diff --git a/src/arkruntime/resources/sessions/sessions.py b/src/arkruntime/resources/sessions/sessions.py index 707cab6..be1e73b 100644 --- a/src/arkruntime/resources/sessions/sessions.py +++ b/src/arkruntime/resources/sessions/sessions.py @@ -8,10 +8,13 @@ from ..._resource import AsyncAPIResource, SyncAPIResource from ..._types import NOT_GIVEN, NotGiven from ...types.session.agent_identifier import AgentIdentifier +from ...types.session.agent_with_upgrades import AgentWithUpgrades from ...types.session.checkpoint_restore_config import CheckpointRestoreConfig from ...types.session.delete_session_response import DeleteSessionResponse +from ...types.session.environment_with_upgrades import EnvironmentWithUpgrades from ...types.session.list_sessions_order import ListSessionsOrder from ...types.session.list_sessions_response import ListSessionsResponse +from ...types.session.managed_agents_event_params import ManagedAgentsEventParams from ...types.session.session import Session from ...types.session.session_resource import SessionResource from ...types.session.session_status import SessionStatus @@ -192,6 +195,48 @@ def update( cast_to=Session, ) + def upgrade( + self, + session_id: str, + *, + agent: Optional[AgentWithUpgrades] | NotGiven = NOT_GIVEN, + environment: Optional[EnvironmentWithUpgrades] | NotGiven = NOT_GIVEN, + vault_ids: Optional[List[str]] | NotGiven = NOT_GIVEN, + initial_events: Optional[List[ManagedAgentsEventParams]] | NotGiven = NOT_GIVEN, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> Session: + """Upgrade an existing Session's runtime configuration. + + At least one of ``agent``, ``environment``, or ``vault_ids`` must be + provided. Agent and environment IDs must match the current Session. + Agent array fields use whole-list replacement semantics: an explicit + ``[]`` clears the list, while an omitted field keeps its current value. + ``vault_ids`` may only refresh the currently bound set; it cannot + change that set. An omitted or zero agent version keeps the current + snapshot rather than selecting the latest version. Self-hosted + environments reject runtime config overrides other than ``config.type``. + """ + if not session_id: + raise ValueError("session_id is required") + return self._post( + f"{_PREFIX}/{session_id}/upgrades", + body=dump_body( + { + "agent": agent, + "environment": environment, + "vault_ids": vault_ids, + "initial_events": initial_events, + } + ), + options=make_request_options( + extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout + ), + cast_to=Session, + ) + def delete( self, session_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None ) -> DeleteSessionResponse: @@ -333,6 +378,48 @@ async def update( cast_to=Session, ) + async def upgrade( + self, + session_id: str, + *, + agent: Optional[AgentWithUpgrades] | NotGiven = NOT_GIVEN, + environment: Optional[EnvironmentWithUpgrades] | NotGiven = NOT_GIVEN, + vault_ids: Optional[List[str]] | NotGiven = NOT_GIVEN, + initial_events: Optional[List[ManagedAgentsEventParams]] | NotGiven = NOT_GIVEN, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> Session: + """Upgrade an existing Session's runtime configuration. + + At least one of ``agent``, ``environment``, or ``vault_ids`` must be + provided. Agent and environment IDs must match the current Session. + Agent array fields use whole-list replacement semantics: an explicit + ``[]`` clears the list, while an omitted field keeps its current value. + ``vault_ids`` may only refresh the currently bound set; it cannot + change that set. An omitted or zero agent version keeps the current + snapshot rather than selecting the latest version. Self-hosted + environments reject runtime config overrides other than ``config.type``. + """ + if not session_id: + raise ValueError("session_id is required") + return await self._post( + f"{_PREFIX}/{session_id}/upgrades", + body=dump_body( + { + "agent": agent, + "environment": environment, + "vault_ids": vault_ids, + "initial_events": initial_events, + } + ), + options=make_request_options( + extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout + ), + cast_to=Session, + ) + async def delete( self, session_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None ) -> DeleteSessionResponse: diff --git a/tests/test_session_upgrades.py b/tests/test_session_upgrades.py new file mode 100644 index 0000000..3943e0e --- /dev/null +++ b/tests/test_session_upgrades.py @@ -0,0 +1,122 @@ +# Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from arkruntime.resources.sessions.sessions import AsyncSessions, Sessions +from arkruntime.types.session.agent_with_upgrades import AgentWithUpgrades +from arkruntime.types.session.environment_with_upgrades import EnvironmentWithUpgrades +from arkruntime.types.session.managed_agents_text_block import ManagedAgentsTextBlock +from arkruntime.types.session.managed_agents_user_message_event_params import ( + ManagedAgentsUserMessageEventParams, +) + + +class _SyncClient: + def __init__(self) -> None: + self.request: dict[str, object] | None = None + + def post(self, *args: object, **kwargs: object) -> dict[str, object]: + self.request = {"args": args, **kwargs} + return {} + + def post_without_retry(self, *args: object, **kwargs: object) -> dict[str, object]: + return self.post(*args, **kwargs) + + def get(self, *args: object, **kwargs: object) -> dict[str, object]: + return {} + + def delete(self, *args: object, **kwargs: object) -> dict[str, object]: + return {} + + def get_api_list(self, *args: object, **kwargs: object) -> list[object]: + return [] + + +class _AsyncClient: + def __init__(self) -> None: + self.request: dict[str, object] | None = None + + async def post(self, *args: object, **kwargs: object) -> dict[str, object]: + self.request = {"args": args, **kwargs} + return {} + + async def post_without_retry(self, *args: object, **kwargs: object) -> dict[str, object]: + return await self.post(*args, **kwargs) + + async def get(self, *args: object, **kwargs: object) -> dict[str, object]: + return {} + + async def delete(self, *args: object, **kwargs: object) -> dict[str, object]: + return {} + + async def get_api_list(self, *args: object, **kwargs: object) -> list[object]: + return [] + + +def test_upgrade_session_posts_upgrade_request() -> None: + client = _SyncClient() + Sessions(client).upgrade( + "sess-1", + agent=AgentWithUpgrades(type="agent_with_upgrades", id="agent-1", tools=[]), + vault_ids=[], + ) + + assert client.request is not None + assert client.request["args"] == ("/sessions/sess-1/upgrades",) + assert client.request["body"] == { + "agent": { + "type": "agent_with_upgrades", + "id": "agent-1", + "tools": [], + }, + "vault_ids": [], + } + + +def test_upgrade_session_requires_session_id() -> None: + with pytest.raises(ValueError, match="session_id"): + Sessions(_SyncClient()).upgrade("") + + +def test_upgrade_session_serializes_environment_and_initial_events() -> None: + client = _SyncClient() + Sessions(client).upgrade( + "sess-1", + environment=EnvironmentWithUpgrades(type="environment_with_upgrades", id="env-1"), + initial_events=[ + ManagedAgentsUserMessageEventParams( + type="user.message", + content=[ManagedAgentsTextBlock(type="text", text="/compact")], + ) + ], + ) + + assert client.request is not None + assert client.request["body"] == { + "environment": {"type": "environment_with_upgrades", "id": "env-1"}, + "initial_events": [ + { + "type": "user.message", + "content": [{"type": "text", "text": "/compact"}], + } + ], + } + + +@pytest.mark.asyncio +async def test_async_upgrade_session_posts_upgrade_request() -> None: + client = _AsyncClient() + await AsyncSessions(client).upgrade("sess-1", vault_ids=[]) + + assert client.request is not None + assert client.request["args"] == ("/sessions/sess-1/upgrades",) + assert client.request["body"] == {"vault_ids": []} + + +@pytest.mark.asyncio +async def test_async_upgrade_session_requires_session_id() -> None: + with pytest.raises(ValueError, match="session_id"): + await AsyncSessions(_AsyncClient()).upgrade("")