diff --git a/src/arkruntime/_client.py b/src/arkruntime/_client.py index 2434277..e2f91f6 100644 --- a/src/arkruntime/_client.py +++ b/src/arkruntime/_client.py @@ -135,6 +135,7 @@ def __init__( self.environments = resources.Environments(self) self.memory_stores = resources.MemoryStores(self) self.sessions = resources.Sessions(self) + self.checkpoints = resources.Checkpoints(self) self.skills = resources.Skills(self) self.vaults = resources.Vaults(self) @@ -289,6 +290,7 @@ def __init__( self.environments = resources.AsyncEnvironments(self) self.memory_stores = resources.AsyncMemoryStores(self) self.sessions = resources.AsyncSessions(self) + self.checkpoints = resources.AsyncCheckpoints(self) self.skills = resources.AsyncSkills(self) self.vaults = resources.AsyncVaults(self) diff --git a/src/arkruntime/resources/__init__.py b/src/arkruntime/resources/__init__.py index e31e22f..fb6f798 100644 --- a/src/arkruntime/resources/__init__.py +++ b/src/arkruntime/resources/__init__.py @@ -5,6 +5,7 @@ from .agents import Agents, AsyncAgents from .batch import AsyncBatch, Batch from .chat import AsyncChat, Chat +from .checkpoints import AsyncCheckpoints, Checkpoints from .content_generation import AsyncContentGeneration, ContentGeneration from .embeddings import AsyncEmbeddings, Embeddings from .environments import AsyncEnvironments, AsyncEnvironmentWork, Environments, EnvironmentWork @@ -31,7 +32,9 @@ from .vaults import AsyncCredentials, AsyncVaults, Credentials, Vaults __all__ = [ + "Checkpoints", "Chat", + "AsyncCheckpoints", "AsyncChat", "Embeddings", "AsyncEmbeddings", diff --git a/src/arkruntime/resources/agents/agents.py b/src/arkruntime/resources/agents/agents.py index 4400a0c..51887e6 100644 --- a/src/arkruntime/resources/agents/agents.py +++ b/src/arkruntime/resources/agents/agents.py @@ -83,13 +83,29 @@ def create( cast_to=Agent, ) - def retrieve(self, agent_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None) -> Agent: + def retrieve( + self, + agent_id: str, + *, + version: Optional[int] | NotGiven = NOT_GIVEN, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> Agent: if not agent_id: raise ValueError("agent_id is required") + query: dict = {} + if version is not NOT_GIVEN and version is not None: + query["version"] = version return self._get( f"{_PREFIX}/{agent_id}", options=make_request_options( - extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout + query=query or None, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, ), cast_to=Agent, ) @@ -248,14 +264,28 @@ async def create( ) async def retrieve( - self, agent_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None + self, + agent_id: str, + *, + version: Optional[int] | NotGiven = NOT_GIVEN, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, ) -> Agent: if not agent_id: raise ValueError("agent_id is required") + query: dict = {} + if version is not NOT_GIVEN and version is not None: + query["version"] = version return await self._get( f"{_PREFIX}/{agent_id}", options=make_request_options( - extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout + query=query or None, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, ), cast_to=Agent, ) diff --git a/src/arkruntime/resources/checkpoints.py b/src/arkruntime/resources/checkpoints.py new file mode 100644 index 0000000..7760ffd --- /dev/null +++ b/src/arkruntime/resources/checkpoints.py @@ -0,0 +1,250 @@ +from __future__ import annotations + +from typing import List, Optional + +from .._base_client import make_request_options +from .._managed_agents_serialize import dump_body +from .._resource import AsyncAPIResource, SyncAPIResource +from .._types import NOT_GIVEN, NotGiven +from ..types.session.checkpoint_status import CheckpointStatus +from ..types.session.create_session_checkpoint_response import ( + CreateSessionCheckpointResponse, +) +from ..types.session.delete_session_checkpoint_response import ( + DeleteSessionCheckpointResponse, +) +from ..types.session.get_session_checkpoint_response import ( + GetSessionCheckpointResponse, +) +from ..types.session.list_session_checkpoints_response import ( + ListSessionCheckpointsResponse, +) + +__all__ = ["Checkpoints", "AsyncCheckpoints"] + +_PREFIX = "/checkpoints" + + +def _list_query( + *, + source_session_id=NOT_GIVEN, + status=NOT_GIVEN, + limit=NOT_GIVEN, + page=NOT_GIVEN, +) -> dict: + q: dict = {} + for k, v in { + "source_session_id": source_session_id, + "status": status, + "limit": limit, + "page": page, + }.items(): + if v is not NOT_GIVEN and v is not None: + q[k] = v + return q + + +class Checkpoints(SyncAPIResource): + def create( + self, + session_id: str, + *, + name: str, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> CreateSessionCheckpointResponse: + """Create a checkpoint of an idle Session. + + The returned checkpoint is initially ``creating`` and becomes + ``ready`` once the asynchronous snapshot completes; only ready + checkpoints can be used as ``checkpoint_id`` in + ``sessions.create``. + """ + if not session_id: + raise ValueError("session_id is required") + return self._post( + f"/sessions/{session_id}/checkpoints", + body=dump_body({"name": name}), + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ), + cast_to=CreateSessionCheckpointResponse, + ) + + def list( + self, + *, + source_session_id: Optional[str] | NotGiven = NOT_GIVEN, + status: Optional[List[CheckpointStatus]] | NotGiven = NOT_GIVEN, + limit: Optional[int] | NotGiven = NOT_GIVEN, + page: Optional[str] | NotGiven = NOT_GIVEN, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> ListSessionCheckpointsResponse: + return self._get( + _PREFIX, + options=make_request_options( + query=_list_query( + source_session_id=source_session_id, + status=status, + limit=limit, + page=page, + ), + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ), + cast_to=ListSessionCheckpointsResponse, + ) + + def retrieve( + self, + checkpoint_id: str, + *, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> GetSessionCheckpointResponse: + if not checkpoint_id: + raise ValueError("checkpoint_id is required") + return self._get( + f"{_PREFIX}/{checkpoint_id}", + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ), + cast_to=GetSessionCheckpointResponse, + ) + + def delete( + self, + checkpoint_id: str, + *, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> DeleteSessionCheckpointResponse: + if not checkpoint_id: + raise ValueError("checkpoint_id is required") + return self._delete( + f"{_PREFIX}/{checkpoint_id}", + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ), + cast_to=DeleteSessionCheckpointResponse, + ) + + +class AsyncCheckpoints(AsyncAPIResource): + async def create( + self, + session_id: str, + *, + name: str, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> CreateSessionCheckpointResponse: + """Async variant of :meth:`Checkpoints.create`.""" + if not session_id: + raise ValueError("session_id is required") + return await self._post( + f"/sessions/{session_id}/checkpoints", + body=dump_body({"name": name}), + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ), + cast_to=CreateSessionCheckpointResponse, + ) + + async def list( + self, + *, + source_session_id: Optional[str] | NotGiven = NOT_GIVEN, + status: Optional[List[CheckpointStatus]] | NotGiven = NOT_GIVEN, + limit: Optional[int] | NotGiven = NOT_GIVEN, + page: Optional[str] | NotGiven = NOT_GIVEN, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> ListSessionCheckpointsResponse: + return await self._get( + _PREFIX, + options=make_request_options( + query=_list_query( + source_session_id=source_session_id, + status=status, + limit=limit, + page=page, + ), + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ), + cast_to=ListSessionCheckpointsResponse, + ) + + async def retrieve( + self, + checkpoint_id: str, + *, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> GetSessionCheckpointResponse: + if not checkpoint_id: + raise ValueError("checkpoint_id is required") + return await self._get( + f"{_PREFIX}/{checkpoint_id}", + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ), + cast_to=GetSessionCheckpointResponse, + ) + + async def delete( + self, + checkpoint_id: str, + *, + extra_headers=None, + extra_query=None, + extra_body=None, + timeout=None, + ) -> DeleteSessionCheckpointResponse: + if not checkpoint_id: + raise ValueError("checkpoint_id is required") + return await self._delete( + f"{_PREFIX}/{checkpoint_id}", + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ), + cast_to=DeleteSessionCheckpointResponse, + ) diff --git a/src/arkruntime/resources/sessions/sessions.py b/src/arkruntime/resources/sessions/sessions.py index 3b5c83a..707cab6 100644 --- a/src/arkruntime/resources/sessions/sessions.py +++ b/src/arkruntime/resources/sessions/sessions.py @@ -8,6 +8,7 @@ from ..._resource import AsyncAPIResource, SyncAPIResource from ..._types import NOT_GIVEN, NotGiven from ...types.session.agent_identifier import AgentIdentifier +from ...types.session.checkpoint_restore_config import CheckpointRestoreConfig from ...types.session.delete_session_response import DeleteSessionResponse from ...types.session.list_sessions_order import ListSessionsOrder from ...types.session.list_sessions_response import ListSessionsResponse @@ -73,17 +74,26 @@ def threads(self) -> SessionThreads: def create( self, *, - agent: AgentIdentifier, - environment_id: str, + agent: Optional[AgentIdentifier] | NotGiven = NOT_GIVEN, + environment_id: Optional[str] | NotGiven = NOT_GIVEN, tags: Optional[List[Tag]] | NotGiven = NOT_GIVEN, resources: Optional[List[SessionResource]] | NotGiven = NOT_GIVEN, title: Optional[str] | NotGiven = NOT_GIVEN, vault_ids: Optional[List[str]] | NotGiven = NOT_GIVEN, + checkpoint_id: Optional[str] | NotGiven = NOT_GIVEN, + checkpoint_restore_config: Optional[CheckpointRestoreConfig] | NotGiven = NOT_GIVEN, extra_headers=None, extra_query=None, extra_body=None, timeout=None, ) -> Session: + """Create a Session, or restore one from a checkpoint. + + When ``checkpoint_id`` is given, ``agent`` / ``environment_id`` may + be omitted and are inherited from the checkpoint snapshot. + ``resources`` must be omitted when restoring; ``vault_ids`` must be + omitted for same-account restore. + """ return self._post( _PREFIX, body=dump_body( @@ -94,6 +104,8 @@ def create( "resources": resources, "title": title, "vault_ids": vault_ids, + "checkpoint_id": checkpoint_id, + "checkpoint_restore_config": checkpoint_restore_config, } ), options=make_request_options( @@ -210,12 +222,14 @@ def threads(self) -> AsyncSessionThreads: async def create( self, *, - agent: AgentIdentifier, - environment_id: str, + agent: Optional[AgentIdentifier] | NotGiven = NOT_GIVEN, + environment_id: Optional[str] | NotGiven = NOT_GIVEN, tags=NOT_GIVEN, resources=NOT_GIVEN, title=NOT_GIVEN, vault_ids=NOT_GIVEN, + checkpoint_id=NOT_GIVEN, + checkpoint_restore_config=NOT_GIVEN, extra_headers=None, extra_query=None, extra_body=None, @@ -231,6 +245,8 @@ async def create( "resources": resources, "title": title, "vault_ids": vault_ids, + "checkpoint_id": checkpoint_id, + "checkpoint_restore_config": checkpoint_restore_config, } ), options=make_request_options( diff --git a/tests/test_checkpoints.py b/tests/test_checkpoints.py new file mode 100644 index 0000000..38a3522 --- /dev/null +++ b/tests/test_checkpoints.py @@ -0,0 +1,269 @@ +# Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the hand-written checkpoint / restore / agent-version call layer. + +The generated ``types/session/*`` models are covered by the ark-apis +codegen contract; these tests pin the request wiring (URL, query +encoding, NOT_GIVEN omission) of the resource wrappers. +""" + +from __future__ import annotations + +import pytest + +from arkruntime.resources.agents.agents import Agents, AsyncAgents +from arkruntime.resources.checkpoints import AsyncCheckpoints, Checkpoints +from arkruntime.resources.sessions.sessions import AsyncSessions, Sessions +from arkruntime.types.session.checkpoint_restore_config import CheckpointRestoreConfig +from arkruntime.types.session.checkpoint_status import CheckpointStatus + + +class _SyncClient: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def _record(self, method: str, *args: object, **kwargs: object) -> dict[str, object]: + call = {"method": method, "args": args, **kwargs} + self.calls.append(call) + return call + + def post(self, *args: object, **kwargs: object) -> dict[str, object]: + return self._record("POST", *args, **kwargs) + + def post_without_retry(self, *args: object, **kwargs: object) -> dict[str, object]: + return self._record("POST", *args, **kwargs) + + def get(self, *args: object, **kwargs: object) -> dict[str, object]: + return self._record("GET", *args, **kwargs) + + def delete(self, *args: object, **kwargs: object) -> dict[str, object]: + return self._record("DELETE", *args, **kwargs) + + def get_api_list(self, *args: object, **kwargs: object) -> list[object]: + return [] + + +class _AsyncClient: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def _record(self, method: str, *args: object, **kwargs: object) -> dict[str, object]: + call = {"method": method, "args": args, **kwargs} + self.calls.append(call) + return call + + async def post(self, *args: object, **kwargs: object) -> dict[str, object]: + return await self._record("POST", *args, **kwargs) + + async def post_without_retry(self, *args: object, **kwargs: object) -> dict[str, object]: + return await self._record("POST", *args, **kwargs) + + async def get(self, *args: object, **kwargs: object) -> dict[str, object]: + return await self._record("GET", *args, **kwargs) + + async def delete(self, *args: object, **kwargs: object) -> dict[str, object]: + return await self._record("DELETE", *args, **kwargs) + + async def get_api_list(self, *args: object, **kwargs: object) -> list[object]: + return [] + + +# --- Checkpoints.create --------------------------------------------------- + + +def test_create_checkpoint_posts_to_session_subresource() -> None: + client = _SyncClient() + Checkpoints(client).create("sess-1", name="nightly") + + call = client.calls[0] + assert call["method"] == "POST" + assert call["args"] == ("/sessions/sess-1/checkpoints",) + assert call["body"] == {"name": "nightly"} + + +def test_create_checkpoint_requires_session_id() -> None: + with pytest.raises(ValueError, match="session_id"): + Checkpoints(_SyncClient()).create("", name="nightly") + + +@pytest.mark.asyncio +async def test_async_create_checkpoint() -> None: + client = _AsyncClient() + await AsyncCheckpoints(client).create("sess-1", name="nightly") + + call = client.calls[0] + assert call["method"] == "POST" + assert call["args"] == ("/sessions/sess-1/checkpoints",) + assert call["body"] == {"name": "nightly"} + + +# --- Checkpoints.list ----------------------------------------------------- + + +def test_list_checkpoints_encodes_filters_as_repeated_query() -> None: + client = _SyncClient() + Checkpoints(client).list( + source_session_id="sess-1", + status=[CheckpointStatus.ready, CheckpointStatus.failed], + limit=10, + page="cursor-1", + ) + + call = client.calls[0] + assert call["method"] == "GET" + assert call["args"] == ("/checkpoints",) + assert call["options"]["params"] == { + "source_session_id": "sess-1", + "status": [CheckpointStatus.ready, CheckpointStatus.failed], + "limit": 10, + "page": "cursor-1", + } + + +def test_list_checkpoints_defaults_to_empty_query() -> None: + client = _SyncClient() + Checkpoints(client).list() + + call = client.calls[0] + assert call["args"] == ("/checkpoints",) + assert call["options"]["params"] == {} + + +def test_list_checkpoints_drops_explicit_none_filters() -> None: + client = _SyncClient() + Checkpoints(client).list(source_session_id=None, status=None, limit=None, page=None) + + assert client.calls[0]["options"]["params"] == {} + + +@pytest.mark.asyncio +async def test_async_list_checkpoints() -> None: + client = _AsyncClient() + await AsyncCheckpoints(client).list(status=[CheckpointStatus.ready]) + + assert client.calls[0]["options"]["params"] == {"status": [CheckpointStatus.ready]} + + +# --- Checkpoints.retrieve / delete ---------------------------------------- + + +def test_retrieve_checkpoint() -> None: + client = _SyncClient() + Checkpoints(client).retrieve("cp-1") + + call = client.calls[0] + assert call["method"] == "GET" + assert call["args"] == ("/checkpoints/cp-1",) + + +def test_retrieve_checkpoint_requires_id() -> None: + with pytest.raises(ValueError, match="checkpoint_id"): + Checkpoints(_SyncClient()).retrieve("") + + +def test_delete_checkpoint() -> None: + client = _SyncClient() + Checkpoints(client).delete("cp-1") + + call = client.calls[0] + assert call["method"] == "DELETE" + assert call["args"] == ("/checkpoints/cp-1",) + + +def test_delete_checkpoint_requires_id() -> None: + with pytest.raises(ValueError, match="checkpoint_id"): + Checkpoints(_SyncClient()).delete("") + + +@pytest.mark.asyncio +async def test_async_retrieve_and_delete_checkpoint() -> None: + client = _AsyncClient() + await AsyncCheckpoints(client).retrieve("cp-1") + await AsyncCheckpoints(client).delete("cp-1") + + assert client.calls[0]["args"] == ("/checkpoints/cp-1",) + assert client.calls[0]["method"] == "GET" + assert client.calls[1]["args"] == ("/checkpoints/cp-1",) + assert client.calls[1]["method"] == "DELETE" + + +# --- Sessions.create restore wiring --------------------------------------- + + +def test_create_session_from_checkpoint_omits_unrelated_fields() -> None: + client = _SyncClient() + Sessions(client).create(checkpoint_id="cp-1") + + call = client.calls[0] + assert call["args"] == ("/sessions",) + # NOT_GIVEN fields (agent / environment_id / resources / ...) must not + # hit the wire — the server rejects resources on restore and inherits + # agent/environment from the snapshot when omitted. + assert call["body"] == {"checkpoint_id": "cp-1"} + + +def test_create_session_from_checkpoint_with_restore_config() -> None: + client = _SyncClient() + Sessions(client).create( + checkpoint_id="cp-1", + checkpoint_restore_config=CheckpointRestoreConfig(tos_bucket="target-bkt"), + ) + + assert client.calls[0]["body"] == { + "checkpoint_id": "cp-1", + "checkpoint_restore_config": {"tos_bucket": "target-bkt"}, + } + + +def test_create_session_normal_path_unchanged() -> None: + client = _SyncClient() + Sessions(client).create(agent="agent-1", environment_id="env-1", title="plain") + + body = client.calls[0]["body"] + assert body["agent"] == "agent-1" + assert body["environment_id"] == "env-1" + assert body["title"] == "plain" + assert "checkpoint_id" not in body + assert "checkpoint_restore_config" not in body + + +@pytest.mark.asyncio +async def test_async_create_session_from_checkpoint() -> None: + client = _AsyncClient() + await AsyncSessions(client).create(checkpoint_id="cp-1") + + assert client.calls[0]["body"] == {"checkpoint_id": "cp-1"} + + +# --- Agents.retrieve version query ---------------------------------------- + + +def test_retrieve_agent_with_version() -> None: + client = _SyncClient() + Agents(client).retrieve("agent-1", version=3) + + call = client.calls[0] + assert call["method"] == "GET" + assert call["args"] == ("/agents/agent-1",) + assert call["options"]["params"] == {"version": 3} + + +def test_retrieve_agent_without_version_has_no_query() -> None: + client = _SyncClient() + Agents(client).retrieve("agent-1") + + assert "params" not in client.calls[0]["options"] + + +def test_retrieve_agent_requires_id() -> None: + with pytest.raises(ValueError, match="agent_id"): + Agents(_SyncClient()).retrieve("", version=1) + + +@pytest.mark.asyncio +async def test_async_retrieve_agent_with_version() -> None: + client = _AsyncClient() + await AsyncAgents(client).retrieve("agent-1", version=7) + + assert client.calls[0]["options"]["params"] == {"version": 7}