Skip to content

Commit e1af355

Browse files
feat(session): add upgrade API
Hand-written change from the internal SDK repository. The original commit message is intentionally not mirrored; see the source repository for details. Sync-Source-Commit: 113ec45e73722a1b2dea4ce439c1758f10263277 Hand-Written-Reason: No Ark-APIs provenance marker; treated as a hand-written source commit. Release-Version: 0.9.0
1 parent 0a3aa10 commit e1af355

2 files changed

Lines changed: 209 additions & 0 deletions

File tree

‎src/arkruntime/resources/sessions/sessions.py‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,13 @@
88
from ..._resource import AsyncAPIResource, SyncAPIResource
99
from ..._types import NOT_GIVEN, NotGiven
1010
from ...types.session.agent_identifier import AgentIdentifier
11+
from ...types.session.agent_with_upgrades import AgentWithUpgrades
1112
from ...types.session.checkpoint_restore_config import CheckpointRestoreConfig
1213
from ...types.session.delete_session_response import DeleteSessionResponse
14+
from ...types.session.environment_with_upgrades import EnvironmentWithUpgrades
1315
from ...types.session.list_sessions_order import ListSessionsOrder
1416
from ...types.session.list_sessions_response import ListSessionsResponse
17+
from ...types.session.managed_agents_event_params import ManagedAgentsEventParams
1518
from ...types.session.session import Session
1619
from ...types.session.session_resource import SessionResource
1720
from ...types.session.session_status import SessionStatus
@@ -192,6 +195,48 @@ def update(
192195
cast_to=Session,
193196
)
194197

198+
def upgrade(
199+
self,
200+
session_id: str,
201+
*,
202+
agent: Optional[AgentWithUpgrades] | NotGiven = NOT_GIVEN,
203+
environment: Optional[EnvironmentWithUpgrades] | NotGiven = NOT_GIVEN,
204+
vault_ids: Optional[List[str]] | NotGiven = NOT_GIVEN,
205+
initial_events: Optional[List[ManagedAgentsEventParams]] | NotGiven = NOT_GIVEN,
206+
extra_headers=None,
207+
extra_query=None,
208+
extra_body=None,
209+
timeout=None,
210+
) -> Session:
211+
"""Upgrade an existing Session's runtime configuration.
212+
213+
At least one of ``agent``, ``environment``, or ``vault_ids`` must be
214+
provided. Agent and environment IDs must match the current Session.
215+
Agent array fields use whole-list replacement semantics: an explicit
216+
``[]`` clears the list, while an omitted field keeps its current value.
217+
``vault_ids`` may only refresh the currently bound set; it cannot
218+
change that set. An omitted or zero agent version keeps the current
219+
snapshot rather than selecting the latest version. Self-hosted
220+
environments reject runtime config overrides other than ``config.type``.
221+
"""
222+
if not session_id:
223+
raise ValueError("session_id is required")
224+
return self._post(
225+
f"{_PREFIX}/{session_id}/upgrades",
226+
body=dump_body(
227+
{
228+
"agent": agent,
229+
"environment": environment,
230+
"vault_ids": vault_ids,
231+
"initial_events": initial_events,
232+
}
233+
),
234+
options=make_request_options(
235+
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
236+
),
237+
cast_to=Session,
238+
)
239+
195240
def delete(
196241
self, session_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None
197242
) -> DeleteSessionResponse:
@@ -333,6 +378,48 @@ async def update(
333378
cast_to=Session,
334379
)
335380

381+
async def upgrade(
382+
self,
383+
session_id: str,
384+
*,
385+
agent: Optional[AgentWithUpgrades] | NotGiven = NOT_GIVEN,
386+
environment: Optional[EnvironmentWithUpgrades] | NotGiven = NOT_GIVEN,
387+
vault_ids: Optional[List[str]] | NotGiven = NOT_GIVEN,
388+
initial_events: Optional[List[ManagedAgentsEventParams]] | NotGiven = NOT_GIVEN,
389+
extra_headers=None,
390+
extra_query=None,
391+
extra_body=None,
392+
timeout=None,
393+
) -> Session:
394+
"""Upgrade an existing Session's runtime configuration.
395+
396+
At least one of ``agent``, ``environment``, or ``vault_ids`` must be
397+
provided. Agent and environment IDs must match the current Session.
398+
Agent array fields use whole-list replacement semantics: an explicit
399+
``[]`` clears the list, while an omitted field keeps its current value.
400+
``vault_ids`` may only refresh the currently bound set; it cannot
401+
change that set. An omitted or zero agent version keeps the current
402+
snapshot rather than selecting the latest version. Self-hosted
403+
environments reject runtime config overrides other than ``config.type``.
404+
"""
405+
if not session_id:
406+
raise ValueError("session_id is required")
407+
return await self._post(
408+
f"{_PREFIX}/{session_id}/upgrades",
409+
body=dump_body(
410+
{
411+
"agent": agent,
412+
"environment": environment,
413+
"vault_ids": vault_ids,
414+
"initial_events": initial_events,
415+
}
416+
),
417+
options=make_request_options(
418+
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
419+
),
420+
cast_to=Session,
421+
)
422+
336423
async def delete(
337424
self, session_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None
338425
) -> DeleteSessionResponse:

‎tests/test_session_upgrades.py‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# Copyright (c) 2026 ByteDance Ltd. and/or its affiliates.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
from __future__ import annotations
5+
6+
import pytest
7+
8+
from arkruntime.resources.sessions.sessions import AsyncSessions, Sessions
9+
from arkruntime.types.session.agent_with_upgrades import AgentWithUpgrades
10+
from arkruntime.types.session.environment_with_upgrades import EnvironmentWithUpgrades
11+
from arkruntime.types.session.managed_agents_text_block import ManagedAgentsTextBlock
12+
from arkruntime.types.session.managed_agents_user_message_event_params import (
13+
ManagedAgentsUserMessageEventParams,
14+
)
15+
16+
17+
class _SyncClient:
18+
def __init__(self) -> None:
19+
self.request: dict[str, object] | None = None
20+
21+
def post(self, *args: object, **kwargs: object) -> dict[str, object]:
22+
self.request = {"args": args, **kwargs}
23+
return {}
24+
25+
def post_without_retry(self, *args: object, **kwargs: object) -> dict[str, object]:
26+
return self.post(*args, **kwargs)
27+
28+
def get(self, *args: object, **kwargs: object) -> dict[str, object]:
29+
return {}
30+
31+
def delete(self, *args: object, **kwargs: object) -> dict[str, object]:
32+
return {}
33+
34+
def get_api_list(self, *args: object, **kwargs: object) -> list[object]:
35+
return []
36+
37+
38+
class _AsyncClient:
39+
def __init__(self) -> None:
40+
self.request: dict[str, object] | None = None
41+
42+
async def post(self, *args: object, **kwargs: object) -> dict[str, object]:
43+
self.request = {"args": args, **kwargs}
44+
return {}
45+
46+
async def post_without_retry(self, *args: object, **kwargs: object) -> dict[str, object]:
47+
return await self.post(*args, **kwargs)
48+
49+
async def get(self, *args: object, **kwargs: object) -> dict[str, object]:
50+
return {}
51+
52+
async def delete(self, *args: object, **kwargs: object) -> dict[str, object]:
53+
return {}
54+
55+
async def get_api_list(self, *args: object, **kwargs: object) -> list[object]:
56+
return []
57+
58+
59+
def test_upgrade_session_posts_upgrade_request() -> None:
60+
client = _SyncClient()
61+
Sessions(client).upgrade(
62+
"sess-1",
63+
agent=AgentWithUpgrades(type="agent_with_upgrades", id="agent-1", tools=[]),
64+
vault_ids=[],
65+
)
66+
67+
assert client.request is not None
68+
assert client.request["args"] == ("/sessions/sess-1/upgrades",)
69+
assert client.request["body"] == {
70+
"agent": {
71+
"type": "agent_with_upgrades",
72+
"id": "agent-1",
73+
"tools": [],
74+
},
75+
"vault_ids": [],
76+
}
77+
78+
79+
def test_upgrade_session_requires_session_id() -> None:
80+
with pytest.raises(ValueError, match="session_id"):
81+
Sessions(_SyncClient()).upgrade("")
82+
83+
84+
def test_upgrade_session_serializes_environment_and_initial_events() -> None:
85+
client = _SyncClient()
86+
Sessions(client).upgrade(
87+
"sess-1",
88+
environment=EnvironmentWithUpgrades(type="environment_with_upgrades", id="env-1"),
89+
initial_events=[
90+
ManagedAgentsUserMessageEventParams(
91+
type="user.message",
92+
content=[ManagedAgentsTextBlock(type="text", text="/compact")],
93+
)
94+
],
95+
)
96+
97+
assert client.request is not None
98+
assert client.request["body"] == {
99+
"environment": {"type": "environment_with_upgrades", "id": "env-1"},
100+
"initial_events": [
101+
{
102+
"type": "user.message",
103+
"content": [{"type": "text", "text": "/compact"}],
104+
}
105+
],
106+
}
107+
108+
109+
@pytest.mark.asyncio
110+
async def test_async_upgrade_session_posts_upgrade_request() -> None:
111+
client = _AsyncClient()
112+
await AsyncSessions(client).upgrade("sess-1", vault_ids=[])
113+
114+
assert client.request is not None
115+
assert client.request["args"] == ("/sessions/sess-1/upgrades",)
116+
assert client.request["body"] == {"vault_ids": []}
117+
118+
119+
@pytest.mark.asyncio
120+
async def test_async_upgrade_session_requires_session_id() -> None:
121+
with pytest.raises(ValueError, match="session_id"):
122+
await AsyncSessions(_AsyncClient()).upgrade("")

0 commit comments

Comments
 (0)