From f0b2c6a44c1b604d2dded8f4e67ce6fbebdd3fd0 Mon Sep 17 00:00:00 2001 From: SSharma-10 Date: Tue, 23 Jun 2026 16:01:21 +0530 Subject: [PATCH 1/8] Integrate Harness APIs --- Makefile | 5 +- examples/agents/async_stream_session.py | 24 ++ examples/agents/create_session.py | 41 ++++ examples/agents/destroy_session.py | 12 + examples/agents/get_session.py | 13 ++ examples/agents/list_sessions.py | 11 + examples/agents/resolve_hitl.py | 19 ++ examples/agents/send_input.py | 14 ++ examples/agents/start_oauth_flow.py | 18 ++ examples/agents/stream_session.py | 31 +++ src/pydo/_patch.py | 10 + src/pydo/agents/__init__.py | 69 ++++++ src/pydo/agents/custom_models.py | 123 ++++++++++ src/pydo/agents/custom_sessions.py | 289 ++++++++++++++++++++++++ src/pydo/aio/_patch.py | 10 + src/pydo/aio/agents/__init__.py | 29 +++ src/pydo/aio/agents/custom_sessions.py | 233 +++++++++++++++++++ tests/agents/test_sessions.py | 259 +++++++++++++++++++++ 18 files changed, 1209 insertions(+), 1 deletion(-) create mode 100644 examples/agents/async_stream_session.py create mode 100644 examples/agents/create_session.py create mode 100644 examples/agents/destroy_session.py create mode 100644 examples/agents/get_session.py create mode 100644 examples/agents/list_sessions.py create mode 100644 examples/agents/resolve_hitl.py create mode 100644 examples/agents/send_input.py create mode 100644 examples/agents/start_oauth_flow.py create mode 100644 examples/agents/stream_session.py create mode 100644 src/pydo/agents/__init__.py create mode 100644 src/pydo/agents/custom_models.py create mode 100644 src/pydo/agents/custom_sessions.py create mode 100644 src/pydo/aio/agents/__init__.py create mode 100644 src/pydo/aio/agents/custom_sessions.py create mode 100644 tests/agents/test_sessions.py diff --git a/Makefile b/Makefile index a5157d7e..51f6b6f1 100644 --- a/Makefile +++ b/Makefile @@ -25,7 +25,10 @@ clean: ## Removes all generated code (except _patch.py files) @printf "=== Cleaning src directory\n" @rm -rf src/pydo/resources @rm -rf src/pydo/types - @find src/pydo -type f ! -name "_patch.py" ! -name "custom_*.py" ! -name "exceptions.py" -exec rm -rf {} + + @find src/pydo -type f \ + ! -name "_patch.py" ! -name "custom_*.py" ! -name "exceptions.py" \ + ! -path "*/agents/__init__.py" ! -path "*/aio/agents/__init__.py" \ + -exec rm -rf {} + .PHONY: download-spec download-spec: ## Download Latest DO Spec diff --git a/examples/agents/async_stream_session.py b/examples/agents/async_stream_session.py new file mode 100644 index 00000000..37b4684e --- /dev/null +++ b/examples/agents/async_stream_session.py @@ -0,0 +1,24 @@ +"""Async stream session events. Set DIGITALOCEAN_TOKEN and SESSION_ID.""" + +import asyncio +import os + +from pydo.aio import Client + +SESSION_ID = os.environ["SESSION_ID"] + + +async def main() -> None: + async with Client(token=os.environ["DIGITALOCEAN_TOKEN"]) as client: + stream = await client.agents.sessions.stream(SESSION_ID) + async with stream as events: + async for event in events: + if getattr(event, "type", None) == "run.token_delta" and event.get("data"): + print(event.data.text, end="", flush=True) + elif "token_chunk" in event: + print(event.token_chunk.text, end="", flush=True) + else: + print(event) + + +asyncio.run(main()) diff --git a/examples/agents/create_session.py b/examples/agents/create_session.py new file mode 100644 index 00000000..6c85c196 --- /dev/null +++ b/examples/agents/create_session.py @@ -0,0 +1,41 @@ +"""Create a session. + +Set DIGITALOCEAN_TOKEN (and PYDO_AGENTS_ENDPOINT for stage2). + +Optional PYDO_AGENT_KIND (Private Beta server support): + CLAUDE_CODE → coding-claude-code (default) + OPENCODE → coding-opencode + CODEX_CLI → coding-codex + NONE → coding-base (requires ops to enable AGENT_KIND_NONE on the server) +""" + +import json +import os + +from pydo import Client +from pydo.agents import AgentKind + +_AGENT_KINDS = { + "CLAUDE_CODE": AgentKind.CLAUDE_CODE, + "OPENCODE": AgentKind.OPENCODE, + "CODEX_CLI": AgentKind.CODEX_CLI, + "NONE": AgentKind.NONE, +} + +client = Client(token=os.environ["DIGITALOCEAN_TOKEN"]) + +kind_name = os.environ.get("PYDO_AGENT_KIND", "CLAUDE_CODE").upper() +try: + agent_kind = _AGENT_KINDS[kind_name] +except KeyError as exc: + raise SystemExit( + f"Unknown PYDO_AGENT_KIND={kind_name!r}; " + f"use one of: {', '.join(_AGENT_KINDS)}" + ) from exc + +resp = client.agents.sessions.create( + agent_kind=agent_kind, + repo_hint="digitalocean/pydo", +) + +print(json.dumps(resp, indent=2, default=str)) diff --git a/examples/agents/destroy_session.py b/examples/agents/destroy_session.py new file mode 100644 index 00000000..ca6b2379 --- /dev/null +++ b/examples/agents/destroy_session.py @@ -0,0 +1,12 @@ +"""Destroy a session. Set DIGITALOCEAN_TOKEN and SESSION_ID.""" + +import os + +from pydo import Client + +SESSION_ID = os.environ["SESSION_ID"] + +client = Client(token=os.environ["DIGITALOCEAN_TOKEN"]) +client.agents.sessions.destroy(SESSION_ID) + +print("ok", SESSION_ID) diff --git a/examples/agents/get_session.py b/examples/agents/get_session.py new file mode 100644 index 00000000..bd135240 --- /dev/null +++ b/examples/agents/get_session.py @@ -0,0 +1,13 @@ +"""Get a session. Set DIGITALOCEAN_TOKEN and SESSION_ID.""" + +import json +import os + +from pydo import Client + +SESSION_ID = os.environ["SESSION_ID"] + +client = Client(token=os.environ["DIGITALOCEAN_TOKEN"]) +resp = client.agents.sessions.get(SESSION_ID) + +print(json.dumps(resp, indent=2, default=str)) diff --git a/examples/agents/list_sessions.py b/examples/agents/list_sessions.py new file mode 100644 index 00000000..1ec10a6a --- /dev/null +++ b/examples/agents/list_sessions.py @@ -0,0 +1,11 @@ +"""List sessions. Set DIGITALOCEAN_TOKEN (and PYDO_AGENTS_ENDPOINT for stage2).""" + +import os + +from pydo import Client + +client = Client(token=os.environ["DIGITALOCEAN_TOKEN"]) + +resp = client.agents.sessions.list(page_size=50) +for session in resp.get("sessions", []): + print(session.session_id, session.status, session.agent_kind) diff --git a/examples/agents/resolve_hitl.py b/examples/agents/resolve_hitl.py new file mode 100644 index 00000000..73b17359 --- /dev/null +++ b/examples/agents/resolve_hitl.py @@ -0,0 +1,19 @@ +"""Resolve HITL. Set DIGITALOCEAN_TOKEN, SESSION_ID, REQUEST_ID, and PYDO_AGENTS_ENDPOINT (stage2).""" + +import os + +from pydo import Client +from pydo.agents import HITLOutcome, ResolutionSource + +SESSION_ID = os.environ["SESSION_ID"] +REQUEST_ID = os.environ["REQUEST_ID"] + +client = Client(token=os.environ["DIGITALOCEAN_TOKEN"]) +client.agents.sessions.resolve_hitl( + SESSION_ID, + REQUEST_ID, + outcome=HITLOutcome.APPROVE, + source=ResolutionSource.OUT_OF_BAND, +) + +print("ok", REQUEST_ID) diff --git a/examples/agents/send_input.py b/examples/agents/send_input.py new file mode 100644 index 00000000..60386bc4 --- /dev/null +++ b/examples/agents/send_input.py @@ -0,0 +1,14 @@ +"""Send input to a session. Set DIGITALOCEAN_TOKEN and SESSION_ID.""" + +import json +import os + +from pydo import Client + +SESSION_ID = os.environ["SESSION_ID"] +TEXT = "Summarise the README in two sentences." + +client = Client(token=os.environ["DIGITALOCEAN_TOKEN"]) +resp = client.agents.sessions.send_input(SESSION_ID, text=TEXT) + +print(json.dumps(resp, indent=2, default=str)) diff --git a/examples/agents/start_oauth_flow.py b/examples/agents/start_oauth_flow.py new file mode 100644 index 00000000..5b01824c --- /dev/null +++ b/examples/agents/start_oauth_flow.py @@ -0,0 +1,18 @@ +"""Start GitHub OAuth. Set DIGITALOCEAN_TOKEN and SESSION_ID.""" + +import json +import os + +from pydo import Client +from pydo.agents import OAuthProvider + +SESSION_ID = os.environ["SESSION_ID"] + +client = Client(token=os.environ["DIGITALOCEAN_TOKEN"]) +resp = client.agents.sessions.start_oauth_flow( + SESSION_ID, + OAuthProvider.GITHUB, + requested_scopes=["repo"], +) + +print(json.dumps(resp, indent=2, default=str)) diff --git a/examples/agents/stream_session.py b/examples/agents/stream_session.py new file mode 100644 index 00000000..3b6454be --- /dev/null +++ b/examples/agents/stream_session.py @@ -0,0 +1,31 @@ +"""Stream session events. + +Required env: + DIGITALOCEAN_TOKEN + SESSION_ID + PYDO_AGENTS_ENDPOINT (stage2: https://api.s2r1.internal.digitalocean.com) + +Tip: open this stream *before* send_input.py so you catch live events. +If the run already finished, the stream may sit idle until the next input. +""" + +import os +import sys + +from pydo import Client + +SESSION_ID = os.environ["SESSION_ID"] + +client = Client(token=os.environ["DIGITALOCEAN_TOKEN"]) +print(f"agents endpoint: {client.agents.base_url}", file=sys.stderr) + +with client.agents.sessions.stream(SESSION_ID) as events: + for event in events: + # SPI wire (harness-api HTTP handler): type + data.text + if getattr(event, "type", None) == "run.token_delta" and event.get("data"): + print(event.data.text, end="", flush=True) + # grpc-gateway proto envelope (legacy) + elif "token_chunk" in event: + print(event.token_chunk.text, end="", flush=True) + else: + print(event) diff --git a/src/pydo/_patch.py b/src/pydo/_patch.py index d979f13d..2fbe1905 100644 --- a/src/pydo/_patch.py +++ b/src/pydo/_patch.py @@ -58,6 +58,8 @@ class Client( # type: ignore subdomain (e.g. ``"https://.agents.do-ai.run"``). Required only when using agent inference endpoints. :paramtype agent_endpoint: str + :keyword agents_endpoint: Hosted Agents API base URL (default + ``api.digitalocean.com``; override via ``PYDO_AGENTS_ENDPOINT``). """ def __init__( @@ -68,6 +70,7 @@ def __init__( timeout: int = 120, inference_endpoint: str = INFERENCE_BASE_URL, agent_endpoint: str = "", + agents_endpoint: Optional[str] = None, **kwargs, ): if token is not None and api_key is not None: @@ -111,6 +114,13 @@ def __init__( self.images.generate = inference_images.generate self.images.generations = inference_images.generations + try: + from pydo.agents import AgentsResources + except ImportError: + self.agents = None + else: + self.agents = AgentsResources(self, agents_endpoint=agents_endpoint) + def _setup_inference_routing( self, inference_endpoint: str, diff --git a/src/pydo/agents/__init__.py b/src/pydo/agents/__init__.py new file mode 100644 index 00000000..659c6aa6 --- /dev/null +++ b/src/pydo/agents/__init__.py @@ -0,0 +1,69 @@ +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""Hosted Agents (Harness) API — hand-written; preserved across ``make generate``.""" +from __future__ import annotations + +import os +from typing import Optional + +from pydo.custom_extensions import _BaseURLProxy + +from .custom_models import ( + AgentKind, + HITLActionKind, + HITLOutcome, + OAuthFlowKind, + OAuthProvider, + ProviderAuthState, + ResolutionSource, + RunFailureCode, + RunState, + SessionStatus, +) +from .custom_sessions import HarnessEventStream, HarnessStreamError, SessionsOperations + +DEFAULT_AGENTS_BASE_URL = "https://api.digitalocean.com" +_ENV_VAR = "PYDO_AGENTS_ENDPOINT" + + +def resolve_agents_base_url(explicit: Optional[str] = None) -> str: + url = explicit or os.environ.get(_ENV_VAR) or DEFAULT_AGENTS_BASE_URL + url = url.rstrip("/") + if "://" not in url: + url = f"https://{url}" + return url + + +class AgentsResources: + def __init__(self, parent_client, *, agents_endpoint: Optional[str] = None): + self._proxy = _BaseURLProxy( + parent_client._client, + resolve_agents_base_url(agents_endpoint), + ) + self.sessions = SessionsOperations(self._proxy) + + @property + def base_url(self) -> str: + return self._proxy._base_url + + +__all__ = [ + "AgentsResources", + "SessionsOperations", + "HarnessEventStream", + "HarnessStreamError", + "DEFAULT_AGENTS_BASE_URL", + "resolve_agents_base_url", + "AgentKind", + "SessionStatus", + "RunState", + "RunFailureCode", + "HITLOutcome", + "HITLActionKind", + "ResolutionSource", + "OAuthProvider", + "OAuthFlowKind", + "ProviderAuthState", +] diff --git a/src/pydo/agents/custom_models.py b/src/pydo/agents/custom_models.py new file mode 100644 index 00000000..940b3b21 --- /dev/null +++ b/src/pydo/agents/custom_models.py @@ -0,0 +1,123 @@ +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""Hosted Agents API enum string constants.""" +from __future__ import annotations + + +class AgentKind: + """Discriminator for what runs inside a session's sandbox.""" + + UNSPECIFIED = "AGENT_KIND_UNSPECIFIED" + CLAUDE_CODE = "AGENT_KIND_CLAUDE_CODE" + OPENCODE = "AGENT_KIND_OPENCODE" + CODEX_CLI = "AGENT_KIND_CODEX_CLI" + CURSOR_CLI = "AGENT_KIND_CURSOR_CLI" + NONE = "AGENT_KIND_NONE" + CUSTOM = "AGENT_KIND_CUSTOM" + + +class SessionStatus: + """Lifecycle states for a session.""" + + UNSPECIFIED = "SESSION_STATUS_UNSPECIFIED" + PROVISIONING = "SESSION_STATUS_PROVISIONING" + READY = "SESSION_STATUS_READY" + DETACHED = "SESSION_STATUS_DETACHED" + DESTROYING = "SESSION_STATUS_DESTROYING" + DESTROYED = "SESSION_STATUS_DESTROYED" + FAILED = "SESSION_STATUS_FAILED" + + +class RunState: + """Finite-state machine for an individual agent run inside a session.""" + + UNSPECIFIED = "RUN_STATE_UNSPECIFIED" + QUEUED = "RUN_STATE_QUEUED" + RUNNING = "RUN_STATE_RUNNING" + AWAITING_HITL = "RUN_STATE_AWAITING_HITL" + PAUSED = "RUN_STATE_PAUSED" + COMPLETED = "RUN_STATE_COMPLETED" + FAILED = "RUN_STATE_FAILED" + + +class RunFailureCode: + """Canonical reasons a run can fail. Closed enum; routable by clients.""" + + UNSPECIFIED = "RUN_FAILURE_CODE_UNSPECIFIED" + MODEL_ERROR = "RUN_FAILURE_CODE_MODEL_ERROR" + MODEL_TIMEOUT = "RUN_FAILURE_CODE_MODEL_TIMEOUT" + TOOL_ERROR = "RUN_FAILURE_CODE_TOOL_ERROR" + SANDBOX_LOST = "RUN_FAILURE_CODE_SANDBOX_LOST" + HITL_REJECTED = "RUN_FAILURE_CODE_HITL_REJECTED" + BUDGET_EXCEEDED = "RUN_FAILURE_CODE_BUDGET_EXCEEDED" + INTERNAL = "RUN_FAILURE_CODE_INTERNAL" + + +class HITLOutcome: + """User decision on a pending HITL request.""" + + UNSPECIFIED = "HITL_OUTCOME_UNSPECIFIED" + APPROVE = "HITL_OUTCOME_APPROVE" + REJECT = "HITL_OUTCOME_REJECT" + DEFER = "HITL_OUTCOME_DEFER" + + +class HITLActionKind: + """Classifies the action the agent is asking permission for.""" + + UNSPECIFIED = "HITL_ACTION_KIND_UNSPECIFIED" + BASH = "HITL_ACTION_BASH" + FILE_WRITE_OUTSIDE_WORKSPACE = "HITL_ACTION_FILE_WRITE_OUTSIDE_WORKSPACE" + GITHUB_COMMIT_PUSH = "HITL_ACTION_GITHUB_COMMIT_PUSH" + GITHUB_CREATE_PR = "HITL_ACTION_GITHUB_CREATE_PR" + GITHUB_BRANCH_DELETE = "HITL_ACTION_GITHUB_BRANCH_DELETE" + GITHUB_FORCE_PUSH = "HITL_ACTION_GITHUB_FORCE_PUSH" + + +class ResolutionSource: + """How a HITL resolution was triggered. Captured for audit / UX analytics.""" + + UNSPECIFIED = "RESOLUTION_SOURCE_UNSPECIFIED" + INLINE_KEYSTROKE = "RESOLUTION_SOURCE_INLINE_KEYSTROKE" + OUT_OF_BAND = "RESOLUTION_SOURCE_OUT_OF_BAND" + + +class OAuthProvider: + """External identity provider a session is linking to.""" + + UNSPECIFIED = "OAUTH_PROVIDER_UNSPECIFIED" + GITHUB = "OAUTH_PROVIDER_GITHUB" + + +class OAuthFlowKind: + """Interaction model the client should drive the developer through.""" + + UNSPECIFIED = "OAUTH_FLOW_KIND_UNSPECIFIED" + WEB_CALLBACK = "OAUTH_FLOW_KIND_WEB_CALLBACK" + DEVICE = "OAUTH_FLOW_KIND_DEVICE" + + +class ProviderAuthState: + """OAuth state for a single provider in ``Session.provider_auth``.""" + + UNSPECIFIED = "PROVIDER_AUTH_STATE_UNSPECIFIED" + NONE = "PROVIDER_AUTH_STATE_NONE" + PENDING = "PROVIDER_AUTH_STATE_PENDING" + AUTHORIZED = "PROVIDER_AUTH_STATE_AUTHORIZED" + EXPIRED = "PROVIDER_AUTH_STATE_EXPIRED" + + +__all__ = [ + "AgentKind", + "SessionStatus", + "RunState", + "RunFailureCode", + "HITLOutcome", + "HITLActionKind", + "ResolutionSource", + "OAuthProvider", + "OAuthFlowKind", + "ProviderAuthState", +] diff --git a/src/pydo/agents/custom_sessions.py b/src/pydo/agents/custom_sessions.py new file mode 100644 index 00000000..0e811a91 --- /dev/null +++ b/src/pydo/agents/custom_sessions.py @@ -0,0 +1,289 @@ +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""Sync Hosted Agents session operations (``/v2/agents/sessions/...``).""" +from __future__ import annotations + +import json as _json +from typing import Any, Dict, Iterator, List, Optional +from urllib.parse import quote + +from azure.core.exceptions import ( + ClientAuthenticationError, + HttpResponseError, + ResourceExistsError, + ResourceNotFoundError, + ResourceNotModifiedError, + map_error, +) +from azure.core.rest import HttpRequest + +from pydo.custom_extensions import SSEStream, _wrap + +_ERROR_MAP = { + 401: ClientAuthenticationError, + 404: ResourceNotFoundError, + 409: ResourceExistsError, + 304: ResourceNotModifiedError, +} + +_BASE_PATH = "/v2/agents/sessions" + + +def _unwrap_harness_sse_chunk(chunk: Dict[str, Any]) -> Optional[Any]: + """Normalize SSE JSON to a harness Event. + + harness-api's HTTP handler emits SPI canonical events + (``event_id``, ``type``, ``data``). grpc-gateway streaming uses a + ``{result, error}`` envelope — accept both. + """ + if chunk.get("result") is not None: + return chunk["result"] + if chunk.get("event_id") and chunk.get("type"): + return chunk + return None + + +def _quote(value: str) -> str: + return quote(str(value), safe="") + + +def _response_body_text(response) -> str: + try: + if hasattr(response, "read"): + try: + response.read() + except Exception: # noqa: BLE001 + pass + body = response.text() if hasattr(response, "text") else response.body() + if isinstance(body, bytes): + body = body.decode("utf-8", errors="replace") + return body or "" + except Exception: # noqa: BLE001 — best-effort error detail for callers + return "" + + +def _raise_agents_http_error(response) -> None: + body = _response_body_text(response) + map_error( + status_code=response.status_code, + response=response, + error_map=_ERROR_MAP, + ) + message = body.strip() or getattr(response, "reason", None) or "request failed" + raise HttpResponseError(message=message, response=response) + + +class HarnessEventStream: + """Unwraps grpc-gateway SSE envelopes ``{result, error}`` into harness Events.""" + + def __init__(self, sse_stream: SSEStream): + self._sse = sse_stream + + def __iter__(self) -> Iterator[Any]: + for chunk in self._sse: + if not isinstance(chunk, dict): + continue + if chunk.get("error"): + err = chunk["error"] + raise HarnessStreamError( + grpc_code=err.get("grpc_code"), + http_code=err.get("http_code"), + message=err.get("message") or "stream error", + http_status=err.get("http_status"), + details=err.get("details") or [], + ) + event = _unwrap_harness_sse_chunk(chunk) + if event is not None: + yield event + + def close(self) -> None: + self._sse.close() + + def __enter__(self) -> "HarnessEventStream": + return self + + def __exit__(self, *args: Any) -> None: + self.close() + + +class HarnessStreamError(RuntimeError): + """SSE stream error frame from harness-api.""" + + def __init__( + self, + *, + grpc_code: Optional[int], + http_code: Optional[int], + message: str, + http_status: Optional[str] = None, + details: Optional[List[Any]] = None, + ): + self.grpc_code = grpc_code + self.http_code = http_code + self.http_status = http_status + self.details = details or [] + super().__init__(message) + + +class SessionsOperations: + """Hosted Agents session REST operations.""" + + def __init__(self, base_url_proxy): + self._client = base_url_proxy + + def _send( + self, + method: str, + path: str, + *, + body: Optional[Dict[str, Any]] = None, + params: Optional[Dict[str, Any]] = None, + stream: bool = False, + ): + headers = {"Accept": "application/json"} + kwargs: Dict[str, Any] = {"headers": headers} + if params: + kwargs["params"] = { + k: v for k, v in params.items() if v is not None and v != "" + } + if body is not None: + headers["Content-Type"] = "application/json" + kwargs["json"] = body + + request = HttpRequest(method, path, **kwargs) + request.url = self._client.format_url(request.url) + pipeline_response = self._client._pipeline.run(request, stream=stream) + response = pipeline_response.http_response + + if response.status_code not in (200, 204): + _raise_agents_http_error(response) + return pipeline_response + + @staticmethod + def _parse_json(pipeline_response) -> Any: + response = pipeline_response.http_response + body = response.text() if hasattr(response, "text") else response.body() + if not body: + return None + if isinstance(body, bytes): + body = body.decode("utf-8") + return _wrap(_json.loads(body)) + + def list( + self, + *, + page_token: Optional[str] = None, + page_size: Optional[int] = None, + status: Optional[str] = None, + ) -> Any: + return self._parse_json( + self._send( + "GET", + _BASE_PATH, + params={ + "page_token": page_token, + "page_size": page_size, + "status": status, + }, + ), + ) + + def create( + self, + *, + agent_kind: str, + repo_hint: Optional[str] = None, + idle_timeout_seconds: Optional[int] = None, + ) -> Any: + body: Dict[str, Any] = {"agent_kind": agent_kind} + if repo_hint is not None: + body["repo_hint"] = repo_hint + if idle_timeout_seconds is not None: + body["idle_timeout_seconds"] = idle_timeout_seconds + return self._parse_json(self._send("POST", _BASE_PATH, body=body)) + + def get(self, session_id: str) -> Any: + return self._parse_json( + self._send("GET", f"{_BASE_PATH}/{_quote(session_id)}"), + ) + + def destroy(self, session_id: str) -> None: + self._send("DELETE", f"{_BASE_PATH}/{_quote(session_id)}") + + def send_input(self, session_id: str, *, text: str) -> Any: + return self._parse_json( + self._send( + "POST", + f"{_BASE_PATH}/{_quote(session_id)}/input", + body={"text": text}, + ), + ) + + def resolve_hitl( + self, + session_id: str, + request_id: str, + *, + outcome: str, + reason: Optional[str] = None, + source: Optional[str] = None, + ) -> None: + body: Dict[str, Any] = {"outcome": outcome} + if reason is not None: + body["reason"] = reason + if source is not None: + body["source"] = source + self._send( + "POST", + f"{_BASE_PATH}/{_quote(session_id)}/hitl/{_quote(request_id)}", + body=body, + ) + + def start_oauth_flow( + self, + session_id: str, + provider: str, + *, + requested_scopes: Optional[List[str]] = None, + ) -> Any: + body: Dict[str, Any] = {} + if requested_scopes is not None: + body["requested_scopes"] = list(requested_scopes) + return self._parse_json( + self._send( + "POST", + f"{_BASE_PATH}/{_quote(session_id)}/oauth/{_quote(provider)}", + body=body, + ), + ) + + def stream( + self, + session_id: str, + *, + replay_from: Optional[str] = None, + replay_only: bool = False, + ) -> HarnessEventStream: + params: Dict[str, Any] = {} + if replay_from: + params["replay_from"] = replay_from + if replay_only: + params["replay_only"] = "true" + + request = HttpRequest( + "GET", + f"{_BASE_PATH}/{_quote(session_id)}/stream", + headers={"Accept": "text/event-stream, application/json"}, + params=params, + ) + request.url = self._client.format_url(request.url) + pipeline_response = self._client._pipeline.run(request, stream=True) + response = pipeline_response.http_response + if response.status_code != 200: + _raise_agents_http_error(response) + return HarnessEventStream(SSEStream(response)) + + +__all__ = ["SessionsOperations", "HarnessEventStream", "HarnessStreamError"] diff --git a/src/pydo/aio/_patch.py b/src/pydo/aio/_patch.py index 1d317f97..e92d6e49 100644 --- a/src/pydo/aio/_patch.py +++ b/src/pydo/aio/_patch.py @@ -64,6 +64,8 @@ class Client( # type: ignore subdomain (e.g. ``"https://.agents.do-ai.run"``). Required only when using agent inference endpoints. :paramtype agent_endpoint: str + :keyword agents_endpoint: Hosted Agents API base URL (default + ``api.digitalocean.com``; override via ``PYDO_AGENTS_ENDPOINT``). """ def __init__( @@ -74,6 +76,7 @@ def __init__( timeout: int = 120, inference_endpoint: str = INFERENCE_BASE_URL, agent_endpoint: str = "", + agents_endpoint: Optional[str] = None, **kwargs, ): if token is not None and api_key is not None: @@ -117,6 +120,13 @@ def __init__( self.images.generate = inference_images.generate self.images.generations = inference_images.generations + try: + from pydo.aio.agents import AsyncAgentsResources + except ImportError: + self.agents = None + else: + self.agents = AsyncAgentsResources(self, agents_endpoint=agents_endpoint) + def _setup_inference_routing( self, inference_endpoint: str, diff --git a/src/pydo/aio/agents/__init__.py b/src/pydo/aio/agents/__init__.py new file mode 100644 index 00000000..dd184182 --- /dev/null +++ b/src/pydo/aio/agents/__init__.py @@ -0,0 +1,29 @@ +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""Async Hosted Agents API — hand-written; preserved across ``make generate``.""" +from __future__ import annotations + +from typing import Optional + +from pydo.agents import resolve_agents_base_url +from pydo.custom_extensions import _BaseURLProxy + +from .custom_sessions import AsyncHarnessEventStream, AsyncSessionsOperations + + +class AsyncAgentsResources: + def __init__(self, parent_client, *, agents_endpoint: Optional[str] = None): + self._proxy = _BaseURLProxy( + parent_client._client, + resolve_agents_base_url(agents_endpoint), + ) + self.sessions = AsyncSessionsOperations(self._proxy) + + @property + def base_url(self) -> str: + return self._proxy._base_url + + +__all__ = ["AsyncAgentsResources", "AsyncSessionsOperations", "AsyncHarnessEventStream"] diff --git a/src/pydo/aio/agents/custom_sessions.py b/src/pydo/aio/agents/custom_sessions.py new file mode 100644 index 00000000..64f1bbde --- /dev/null +++ b/src/pydo/aio/agents/custom_sessions.py @@ -0,0 +1,233 @@ +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""Async Hosted Agents session operations.""" +from __future__ import annotations + +import json as _json +from typing import Any, AsyncIterator, Dict, List, Optional +from urllib.parse import quote + +from azure.core.exceptions import ( + ClientAuthenticationError, + HttpResponseError, + ResourceExistsError, + ResourceNotFoundError, + ResourceNotModifiedError, + map_error, +) +from azure.core.rest import HttpRequest + +from pydo.agents.custom_sessions import HarnessStreamError, _raise_agents_http_error, _unwrap_harness_sse_chunk +from pydo.custom_extensions import AsyncSSEStream, _wrap + +_ERROR_MAP = { + 401: ClientAuthenticationError, + 404: ResourceNotFoundError, + 409: ResourceExistsError, + 304: ResourceNotModifiedError, +} + +_BASE_PATH = "/v2/agents/sessions" + + +def _quote(value: str) -> str: + return quote(str(value), safe="") + + +class AsyncHarnessEventStream: + def __init__(self, sse_stream: AsyncSSEStream): + self._sse = sse_stream + + def __aiter__(self) -> AsyncIterator[Any]: + return self._iter() + + async def _iter(self) -> AsyncIterator[Any]: + async for chunk in self._sse: + if not isinstance(chunk, dict): + continue + if chunk.get("error"): + err = chunk["error"] + raise HarnessStreamError( + grpc_code=err.get("grpc_code"), + http_code=err.get("http_code"), + message=err.get("message") or "stream error", + http_status=err.get("http_status"), + details=err.get("details") or [], + ) + event = _unwrap_harness_sse_chunk(chunk) + if event is not None: + yield event + + async def close(self) -> None: + await self._sse.close() + + async def __aenter__(self) -> "AsyncHarnessEventStream": + return self + + async def __aexit__(self, *args: Any) -> None: + await self.close() + + +class AsyncSessionsOperations: + def __init__(self, base_url_proxy): + self._client = base_url_proxy + + async def _send( + self, + method: str, + path: str, + *, + body: Optional[Dict[str, Any]] = None, + params: Optional[Dict[str, Any]] = None, + stream: bool = False, + ): + headers = {"Accept": "application/json"} + kwargs: Dict[str, Any] = {"headers": headers} + if params: + kwargs["params"] = { + k: v for k, v in params.items() if v is not None and v != "" + } + if body is not None: + headers["Content-Type"] = "application/json" + kwargs["json"] = body + + request = HttpRequest(method, path, **kwargs) + request.url = self._client.format_url(request.url) + pipeline_response = await self._client._pipeline.run(request, stream=stream) + response = pipeline_response.http_response + + if response.status_code not in (200, 204): + await response.read() + _raise_agents_http_error(response) + return pipeline_response + + @staticmethod + async def _parse_json(pipeline_response) -> Any: + body = await pipeline_response.http_response.read() + if not body: + return None + if isinstance(body, bytes): + body = body.decode("utf-8") + return _wrap(_json.loads(body)) + + async def list( + self, + *, + page_token: Optional[str] = None, + page_size: Optional[int] = None, + status: Optional[str] = None, + ) -> Any: + return await self._parse_json( + await self._send( + "GET", + _BASE_PATH, + params={ + "page_token": page_token, + "page_size": page_size, + "status": status, + }, + ), + ) + + async def create( + self, + *, + agent_kind: str, + repo_hint: Optional[str] = None, + idle_timeout_seconds: Optional[int] = None, + ) -> Any: + body: Dict[str, Any] = {"agent_kind": agent_kind} + if repo_hint is not None: + body["repo_hint"] = repo_hint + if idle_timeout_seconds is not None: + body["idle_timeout_seconds"] = idle_timeout_seconds + return await self._parse_json( + await self._send("POST", _BASE_PATH, body=body), + ) + + async def get(self, session_id: str) -> Any: + return await self._parse_json( + await self._send("GET", f"{_BASE_PATH}/{_quote(session_id)}"), + ) + + async def destroy(self, session_id: str) -> None: + await self._send("DELETE", f"{_BASE_PATH}/{_quote(session_id)}") + + async def send_input(self, session_id: str, *, text: str) -> Any: + return await self._parse_json( + await self._send( + "POST", + f"{_BASE_PATH}/{_quote(session_id)}/input", + body={"text": text}, + ), + ) + + async def resolve_hitl( + self, + session_id: str, + request_id: str, + *, + outcome: str, + reason: Optional[str] = None, + source: Optional[str] = None, + ) -> None: + body: Dict[str, Any] = {"outcome": outcome} + if reason is not None: + body["reason"] = reason + if source is not None: + body["source"] = source + await self._send( + "POST", + f"{_BASE_PATH}/{_quote(session_id)}/hitl/{_quote(request_id)}", + body=body, + ) + + async def start_oauth_flow( + self, + session_id: str, + provider: str, + *, + requested_scopes: Optional[List[str]] = None, + ) -> Any: + body: Dict[str, Any] = {} + if requested_scopes is not None: + body["requested_scopes"] = list(requested_scopes) + return await self._parse_json( + await self._send( + "POST", + f"{_BASE_PATH}/{_quote(session_id)}/oauth/{_quote(provider)}", + body=body, + ), + ) + + async def stream( + self, + session_id: str, + *, + replay_from: Optional[str] = None, + replay_only: bool = False, + ) -> AsyncHarnessEventStream: + params: Dict[str, Any] = {} + if replay_from: + params["replay_from"] = replay_from + if replay_only: + params["replay_only"] = "true" + + request = HttpRequest( + "GET", + f"{_BASE_PATH}/{_quote(session_id)}/stream", + headers={"Accept": "text/event-stream, application/json"}, + params=params, + ) + request.url = self._client.format_url(request.url) + pipeline_response = await self._client._pipeline.run(request, stream=True) + response = pipeline_response.http_response + if response.status_code != 200: + await response.read() + _raise_agents_http_error(response) + return AsyncHarnessEventStream(AsyncSSEStream(response)) + + +__all__ = ["AsyncSessionsOperations", "AsyncHarnessEventStream"] diff --git a/tests/agents/test_sessions.py b/tests/agents/test_sessions.py new file mode 100644 index 00000000..0481d3f3 --- /dev/null +++ b/tests/agents/test_sessions.py @@ -0,0 +1,259 @@ +# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""Unit tests for :mod:`pydo.agents.custom_sessions`.""" +from __future__ import annotations + +import json +from types import SimpleNamespace +from typing import Any, List +from unittest.mock import MagicMock + +import pytest + +from pydo.agents import ( + AgentKind, + AgentsResources, + HITLOutcome, + HarnessStreamError, + OAuthProvider, + ResolutionSource, + SessionStatus, + resolve_agents_base_url, +) + + +# --------------------------------------------------------------------------- +# Fake pipeline / response plumbing +# --------------------------------------------------------------------------- + + +class _FakeResponse: + def __init__(self, status_code: int, body: Any = None, *, sse_chunks=None): + self.status_code = status_code + if isinstance(body, (dict, list)): + self._body_bytes = json.dumps(body).encode("utf-8") + elif isinstance(body, str): + self._body_bytes = body.encode("utf-8") + elif isinstance(body, bytes): + self._body_bytes = body + else: + self._body_bytes = b"" + self._sse = sse_chunks + + def text(self) -> str: + return self._body_bytes.decode("utf-8") + + def body(self) -> bytes: + return self._body_bytes + + def read(self) -> bytes: + return self._body_bytes + + def iter_bytes(self): + for chunk in self._sse or []: + yield chunk + + def close(self) -> None: + pass + + +class _FakePipeline: + def __init__(self, responses: List[_FakeResponse]): + self._responses = list(responses) + self.calls: List[Any] = [] + + def run(self, request, *, stream=False): + self.calls.append(SimpleNamespace(request=request, stream=stream)) + response = self._responses.pop(0) + return SimpleNamespace(http_response=response) + + +def _make_resources(responses: List[_FakeResponse]) -> AgentsResources: + parent = MagicMock() + parent._client = MagicMock() + parent._client._pipeline = _FakePipeline(responses) + return AgentsResources(parent, agents_endpoint="https://api.stage2.digitalocean.com") + + +# --------------------------------------------------------------------------- +# CRUD endpoints +# --------------------------------------------------------------------------- + + +def test_create_session_posts_expected_body(): + body = {"session": {"session_id": "abc", "status": SessionStatus.PROVISIONING}} + resources = _make_resources([_FakeResponse(200, body)]) + + resp = resources.sessions.create( + agent_kind=AgentKind.CLAUDE_CODE, + repo_hint="digitalocean/pydo", + idle_timeout_seconds=900, + ) + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.method == "POST" + assert call.request.url.endswith("/v2/agents/sessions") + sent = json.loads(call.request.content) + assert sent == { + "agent_kind": "AGENT_KIND_CLAUDE_CODE", + "repo_hint": "digitalocean/pydo", + "idle_timeout_seconds": 900, + } + assert resp.session.session_id == "abc" + + +def test_get_session_url_encodes_id(): + resources = _make_resources([_FakeResponse(200, {"session": {"session_id": "x/y"}})]) + resources.sessions.get("x/y") + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.method == "GET" + assert call.request.url.endswith("/v2/agents/sessions/x%2Fy") + + +def test_destroy_session(): + resources = _make_resources([_FakeResponse(200, "")]) + resources.sessions.destroy("abc-123") + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.method == "DELETE" + assert call.request.url.endswith("/v2/agents/sessions/abc-123") + + +def test_list_sessions_propagates_query_params(): + resources = _make_resources([_FakeResponse(200, {"sessions": [], "next_page_token": ""})]) + resources.sessions.list(page_token="tok", page_size=10, status=SessionStatus.READY) + + call = resources._proxy._original._pipeline.calls[0] + raw = call.request.url + assert "page_token=tok" in raw + assert "page_size=10" in raw + assert "status=SESSION_STATUS_READY" in raw + + +def test_send_input_body_shape(): + resources = _make_resources([_FakeResponse(200, {"run_id": "r1"})]) + resp = resources.sessions.send_input("s1", text="hello world") + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.method == "POST" + assert call.request.url.endswith("/v2/agents/sessions/s1/input") + assert json.loads(call.request.content) == {"text": "hello world"} + assert resp.run_id == "r1" + + +def test_resolve_hitl_url_and_body(): + resources = _make_resources([_FakeResponse(200, "")]) + resources.sessions.resolve_hitl( + "s1", + "req-9", + outcome=HITLOutcome.APPROVE, + reason="looks safe", + source=ResolutionSource.OUT_OF_BAND, + ) + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.url.endswith("/v2/agents/sessions/s1/hitl/req-9") + assert json.loads(call.request.content) == { + "outcome": "HITL_OUTCOME_APPROVE", + "reason": "looks safe", + "source": "RESOLUTION_SOURCE_OUT_OF_BAND", + } + + +def test_start_oauth_flow(): + body = { + "authorize_url": "https://github.com/login/oauth/authorize?...", + "flow_kind": "OAUTH_FLOW_KIND_WEB_CALLBACK", + } + resources = _make_resources([_FakeResponse(200, body)]) + + resp = resources.sessions.start_oauth_flow( + "s1", + OAuthProvider.GITHUB, + requested_scopes=["repo"], + ) + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.url.endswith( + "/v2/agents/sessions/s1/oauth/OAUTH_PROVIDER_GITHUB" + ) + assert json.loads(call.request.content) == {"requested_scopes": ["repo"]} + assert resp.flow_kind == "OAUTH_FLOW_KIND_WEB_CALLBACK" + + +# --------------------------------------------------------------------------- +# Streaming +# --------------------------------------------------------------------------- + + +def test_stream_unwraps_spi_canonical_envelope(): + sse_payload = ( + b'data: {"event_id":"e1","type":"run.token_delta","data":{"text":"hello "}}\n\n' + b'data: {"event_id":"e2","type":"run.token_delta","data":{"text":"world"}}\n\n' + b'data: {"event_id":"e3","type":"run.completed","data":{"run_cost_micros":1234}}\n\n' + ) + resources = _make_resources( + [_FakeResponse(200, sse_chunks=[sse_payload])] + ) + + events = list(resources.sessions.stream("s1")) + assert events[0].type == "run.token_delta" + assert events[0].data.text == "hello " + assert events[1].data.text == "world" + assert events[2].type == "run.completed" + assert events[2].data.run_cost_micros == 1234 + + +def test_stream_unwraps_result_envelope(): + sse_payload = ( + b'data: {"result":{"event_id":"e1","token_chunk":{"text":"hello "}}}\n\n' + b'data: {"result":{"event_id":"e2","token_chunk":{"text":"world"}}}\n\n' + b'data: {"result":{"event_id":"e3","run_completed":{"run_cost_micros":1234}}}\n\n' + ) + resources = _make_resources( + [_FakeResponse(200, sse_chunks=[sse_payload])] + ) + + events = list(resources.sessions.stream("s1")) + assert events[0].token_chunk.text == "hello " + assert events[1].token_chunk.text == "world" + assert events[2].run_completed.run_cost_micros == 1234 + + +def test_stream_error_envelope_raises(): + sse_payload = ( + b'data: {"error":{"grpc_code":9,"http_code":412,"message":"not ready"}}\n\n' + ) + resources = _make_resources( + [_FakeResponse(200, sse_chunks=[sse_payload])] + ) + + with pytest.raises(HarnessStreamError) as excinfo: + list(resources.sessions.stream("s1")) + + assert excinfo.value.grpc_code == 9 + assert excinfo.value.http_code == 412 + assert "not ready" in str(excinfo.value) + + +def test_stream_passes_replay_query_params(): + sse_payload = b"" + resources = _make_resources( + [_FakeResponse(200, sse_chunks=[sse_payload])] + ) + + stream = resources.sessions.stream("s1", replay_from="evt-42", replay_only=True) + list(stream) + + call = resources._proxy._original._pipeline.calls[0] + assert "replay_from=evt-42" in call.request.url + assert "replay_only=true" in call.request.url + + +def test_resolve_agents_base_url_adds_https_scheme(): + assert resolve_agents_base_url("api.digitalocean.com") == "https://api.digitalocean.com" + assert resolve_agents_base_url("http://127.0.0.1:8080") == "http://127.0.0.1:8080" From 30bfebd98434df32a0e6e23893cc3828a91bd1e5 Mon Sep 17 00:00:00 2001 From: SSharma-10 Date: Tue, 23 Jun 2026 16:21:08 +0530 Subject: [PATCH 2/8] lint fix --- tests/agents/test_sessions.py | 35 ++++++++++++++++++----------------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/tests/agents/test_sessions.py b/tests/agents/test_sessions.py index 0481d3f3..b381c274 100644 --- a/tests/agents/test_sessions.py +++ b/tests/agents/test_sessions.py @@ -1,4 +1,4 @@ -# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access +# pylint: disable=line-too-long,missing-class-docstring,missing-function-docstring,protected-access # ------------------------------------ # Copyright (c) DigitalOcean. # Licensed under the Apache-2.0 License. @@ -75,7 +75,9 @@ def _make_resources(responses: List[_FakeResponse]) -> AgentsResources: parent = MagicMock() parent._client = MagicMock() parent._client._pipeline = _FakePipeline(responses) - return AgentsResources(parent, agents_endpoint="https://api.stage2.digitalocean.com") + return AgentsResources( + parent, agents_endpoint="https://api.stage2.digitalocean.com" + ) # --------------------------------------------------------------------------- @@ -106,7 +108,9 @@ def test_create_session_posts_expected_body(): def test_get_session_url_encodes_id(): - resources = _make_resources([_FakeResponse(200, {"session": {"session_id": "x/y"}})]) + resources = _make_resources( + [_FakeResponse(200, {"session": {"session_id": "x/y"}})] + ) resources.sessions.get("x/y") call = resources._proxy._original._pipeline.calls[0] @@ -124,7 +128,9 @@ def test_destroy_session(): def test_list_sessions_propagates_query_params(): - resources = _make_resources([_FakeResponse(200, {"sessions": [], "next_page_token": ""})]) + resources = _make_resources( + [_FakeResponse(200, {"sessions": [], "next_page_token": ""})] + ) resources.sessions.list(page_token="tok", page_size=10, status=SessionStatus.READY) call = resources._proxy._original._pipeline.calls[0] @@ -196,9 +202,7 @@ def test_stream_unwraps_spi_canonical_envelope(): b'data: {"event_id":"e2","type":"run.token_delta","data":{"text":"world"}}\n\n' b'data: {"event_id":"e3","type":"run.completed","data":{"run_cost_micros":1234}}\n\n' ) - resources = _make_resources( - [_FakeResponse(200, sse_chunks=[sse_payload])] - ) + resources = _make_resources([_FakeResponse(200, sse_chunks=[sse_payload])]) events = list(resources.sessions.stream("s1")) assert events[0].type == "run.token_delta" @@ -214,9 +218,7 @@ def test_stream_unwraps_result_envelope(): b'data: {"result":{"event_id":"e2","token_chunk":{"text":"world"}}}\n\n' b'data: {"result":{"event_id":"e3","run_completed":{"run_cost_micros":1234}}}\n\n' ) - resources = _make_resources( - [_FakeResponse(200, sse_chunks=[sse_payload])] - ) + resources = _make_resources([_FakeResponse(200, sse_chunks=[sse_payload])]) events = list(resources.sessions.stream("s1")) assert events[0].token_chunk.text == "hello " @@ -228,9 +230,7 @@ def test_stream_error_envelope_raises(): sse_payload = ( b'data: {"error":{"grpc_code":9,"http_code":412,"message":"not ready"}}\n\n' ) - resources = _make_resources( - [_FakeResponse(200, sse_chunks=[sse_payload])] - ) + resources = _make_resources([_FakeResponse(200, sse_chunks=[sse_payload])]) with pytest.raises(HarnessStreamError) as excinfo: list(resources.sessions.stream("s1")) @@ -242,9 +242,7 @@ def test_stream_error_envelope_raises(): def test_stream_passes_replay_query_params(): sse_payload = b"" - resources = _make_resources( - [_FakeResponse(200, sse_chunks=[sse_payload])] - ) + resources = _make_resources([_FakeResponse(200, sse_chunks=[sse_payload])]) stream = resources.sessions.stream("s1", replay_from="evt-42", replay_only=True) list(stream) @@ -255,5 +253,8 @@ def test_stream_passes_replay_query_params(): def test_resolve_agents_base_url_adds_https_scheme(): - assert resolve_agents_base_url("api.digitalocean.com") == "https://api.digitalocean.com" + assert ( + resolve_agents_base_url("api.digitalocean.com") + == "https://api.digitalocean.com" + ) assert resolve_agents_base_url("http://127.0.0.1:8080") == "http://127.0.0.1:8080" From 19139223c02ff94e7789653169afca441aadb043 Mon Sep 17 00:00:00 2001 From: Amulya <18429813+logwolvy@users.noreply.github.com> Date: Fri, 26 Jun 2026 17:44:10 +0530 Subject: [PATCH 3/8] Managed Agents: spec-based session creation + high-level AgentSession interface (#682) * Create hosted-agent sessions from an agents.yaml manifest The harness-api backend provisions a session entirely from the agent spec, so replace the legacy JSON create(agent_kind=...) with create_from_manifest(), which uploads the manifest verbatim as Content-Type: application/x-yaml. The spec defines the runtime adapter, sandbox, env vars, and egress; there are no other arguments. - Add create_from_manifest() to the sync and async session clients, sharing manifest-normalization and YAML media-type helpers. - Remove the now-unsupported create(agent_kind=...) path. - Point the create_session.py example at a spec file and add session_e2e.py demonstrating the full create -> attach (send a prompt, stream SSE back) -> destroy flow. - Cover manifest upload in the sync and new async unit tests; drop the legacy create test. * regenerate fix * fix --------- Co-authored-by: SSharma-10 --- Makefile | 1 + examples/agents/attach.py | 53 ++++ examples/agents/create_session.py | 41 +-- examples/agents/run_blocking.py | 31 ++ examples/agents/session_e2e.py | 70 +++++ src/pydo/agents/__init__.py | 32 ++ src/pydo/agents/custom_sessions.py | 59 +++- src/pydo/agents/session.py | 414 +++++++++++++++++++++++++ src/pydo/aio/agents/__init__.py | 27 +- src/pydo/aio/agents/custom_sessions.py | 47 ++- src/pydo/aio/agents/session.py | 245 +++++++++++++++ tests/agents/test_async_sessions.py | 85 +++++ tests/agents/test_session_highlevel.py | 292 +++++++++++++++++ tests/agents/test_sessions.py | 40 ++- 14 files changed, 1365 insertions(+), 72 deletions(-) create mode 100644 examples/agents/attach.py create mode 100644 examples/agents/run_blocking.py create mode 100644 examples/agents/session_e2e.py create mode 100644 src/pydo/agents/session.py create mode 100644 src/pydo/aio/agents/session.py create mode 100644 tests/agents/test_async_sessions.py create mode 100644 tests/agents/test_session_highlevel.py diff --git a/Makefile b/Makefile index 755765eb..a2167026 100644 --- a/Makefile +++ b/Makefile @@ -29,6 +29,7 @@ clean: ## Removes all generated code (except _patch.py files) @find src/pydo -type f \ ! -name "_patch.py" ! -name "custom_*.py" ! -name "exceptions.py" \ ! -path "*/agents/__init__.py" ! -path "*/aio/agents/__init__.py" \ + ! -path "*/agents/session.py" ! -path "*/aio/agents/session.py" \ -exec rm -rf {} + .PHONY: download-spec diff --git a/examples/agents/attach.py b/examples/agents/attach.py new file mode 100644 index 00000000..af74da85 --- /dev/null +++ b/examples/agents/attach.py @@ -0,0 +1,53 @@ +"""Attach to an EXISTING agent session and stream one turn. + +Like ``doctl agents attach``: connect to a session that already exists, send a +prompt, and consume the SSE event feed through the high-level API — no thread, +no raw event-string matching. attach() never destroys the session, so it stays +alive for further turns. + +Get a SESSION_ID first by creating one (the session stays up): + AGENT_SPEC=... python examples/agents/create_session.py +then copy the "session_id" it prints. + +Required env: + DIGITALOCEAN_TOKEN + PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com + SESSION_ID the session to attach to + +Optional env: + PROMPT message to send (default below) +""" + +import os +import sys + +from pydo import Client +from pydo.agents import AgentEventType + +client = Client( + token=os.environ["DIGITALOCEAN_TOKEN"], + agents_endpoint=os.environ.get("PYDO_AGENTS_ENDPOINT"), +) + +agent = client.agents.attach(os.environ["SESSION_ID"]) +prompt = os.environ.get("PROMPT", "In one short sentence, what is DigitalOcean?") + +print(f"[attached {agent.session_id}]\n>>> {prompt}\n", file=sys.stderr) + +stream = agent.run_streamed(prompt) # opens the stream, sends input, yields events +for event in stream: + if event.type == AgentEventType.TOKEN: + print(event.text, end="", flush=True) # live reply -> stdout + elif event.type == AgentEventType.TOOL_CALL: + print(f"\n[tool] {event.tool_name}", file=sys.stderr) + elif event.type == AgentEventType.HITL_REQUESTED: + print(f"\n[hitl auto-approved] {event.request_id}", file=sys.stderr) + +result = stream.result +print( + f"\n\n[{result.status}] captured {len(result.final_output)} chars " + f"(tokens out={result.usage.get('tokens_out')})", + file=sys.stderr, +) +# attach() leaves the session running; destroy it when done: +# SESSION_ID=... python examples/agents/destroy_session.py diff --git a/examples/agents/create_session.py b/examples/agents/create_session.py index 6c85c196..e41e339d 100644 --- a/examples/agents/create_session.py +++ b/examples/agents/create_session.py @@ -1,41 +1,26 @@ -"""Create a session. +"""Create a session from an agent manifest (``agents.yaml``). -Set DIGITALOCEAN_TOKEN (and PYDO_AGENTS_ENDPOINT for stage2). +A session is created entirely from the agent spec — the runtime adapter, +sandbox template, env vars, etc. are all defined in the manifest. The client +uploads it verbatim (``Content-Type: application/x-yaml``); the server parses +and validates it. -Optional PYDO_AGENT_KIND (Private Beta server support): - CLAUDE_CODE → coding-claude-code (default) - OPENCODE → coding-opencode - CODEX_CLI → coding-codex - NONE → coding-base (requires ops to enable AGENT_KIND_NONE on the server) +Required env: + DIGITALOCEAN_TOKEN + PYDO_AGENTS_ENDPOINT (stage2: https://api.s2r1.internal.digitalocean.com) + AGENT_SPEC path to the agent spec YAML (default: agent-spec.yaml) """ import json import os from pydo import Client -from pydo.agents import AgentKind -_AGENT_KINDS = { - "CLAUDE_CODE": AgentKind.CLAUDE_CODE, - "OPENCODE": AgentKind.OPENCODE, - "CODEX_CLI": AgentKind.CODEX_CLI, - "NONE": AgentKind.NONE, -} +spec_path = os.environ.get("AGENT_SPEC", "agent-spec.yaml") +with open(spec_path, "r", encoding="utf-8") as fh: + manifest = fh.read() client = Client(token=os.environ["DIGITALOCEAN_TOKEN"]) - -kind_name = os.environ.get("PYDO_AGENT_KIND", "CLAUDE_CODE").upper() -try: - agent_kind = _AGENT_KINDS[kind_name] -except KeyError as exc: - raise SystemExit( - f"Unknown PYDO_AGENT_KIND={kind_name!r}; " - f"use one of: {', '.join(_AGENT_KINDS)}" - ) from exc - -resp = client.agents.sessions.create( - agent_kind=agent_kind, - repo_hint="digitalocean/pydo", -) +resp = client.agents.sessions.create_from_manifest(manifest) print(json.dumps(resp, indent=2, default=str)) diff --git a/examples/agents/run_blocking.py b/examples/agents/run_blocking.py new file mode 100644 index 00000000..e950a130 --- /dev/null +++ b/examples/agents/run_blocking.py @@ -0,0 +1,31 @@ +"""One-liner blocking run: create from spec -> run -> print -> destroy. + +The high-level API collapses the whole flow into ``agent.run(prompt)``, which +blocks until the run finishes and returns the assembled output. The ``with`` +block auto-destroys the session on exit. + +Required env: + DIGITALOCEAN_TOKEN + PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com + AGENT_SPEC path to the agent spec YAML + +Optional env: + PROMPT +""" + +import os + +from pydo import Client + +client = Client( + token=os.environ["DIGITALOCEAN_TOKEN"], + agents_endpoint=os.environ.get("PYDO_AGENTS_ENDPOINT"), +) + +with open(os.environ.get("AGENT_SPEC", "agent-spec.yaml"), encoding="utf-8") as fh: + manifest = fh.read() + +prompt = os.environ.get("PROMPT", "In one short sentence, what is DigitalOcean?") + +with client.agents.start(manifest) as agent: # auto-destroys on exit + print(agent.run(prompt).final_output) diff --git a/examples/agents/session_e2e.py b/examples/agents/session_e2e.py new file mode 100644 index 00000000..10457af3 --- /dev/null +++ b/examples/agents/session_e2e.py @@ -0,0 +1,70 @@ +"""End-to-end hosted-agents demo: create -> attach -> destroy. + +Uses pydo's high-level agent interface, so consuming the SSE feed needs no +manual wiring — no background thread, no completion event, no dispatching on +raw ``run.*`` event strings, and no explicit teardown: + + * ``client.agents.start(manifest)`` creates the session and returns a handle + that auto-destroys when the ``with`` block exits. + * ``agent.run_streamed(prompt)`` opens the stream, sends the prompt, and + yields normalized, typed events (auto-approving HITL prompts by default). + * ``agent.run(prompt)`` is the blocking one-liner equivalent. + +Required env: + DIGITALOCEAN_TOKEN + PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com + AGENT_SPEC path to the agent spec YAML (agents.yaml manifest) + +Optional env: + PROMPT message to send (default: a short demo prompt) +""" + +import os +import sys + +from pydo import Client +from pydo.agents import AgentEventType + + +def main(): + spec_path = os.environ.get("AGENT_SPEC", "agent-spec.yaml") + with open(spec_path, "r", encoding="utf-8") as fh: + manifest = fh.read() + + prompt = os.environ.get("PROMPT", "In one short sentence, what is DigitalOcean?") + + client = Client( + token=os.environ["DIGITALOCEAN_TOKEN"], + agents_endpoint=os.environ.get("PYDO_AGENTS_ENDPOINT"), + ) + + # create + auto-destroy via the context manager. + with client.agents.start(manifest) as agent: + print(f"[session {agent.session_id} | {agent.status}]", file=sys.stderr) + print(f">>> {prompt}\n", file=sys.stderr) + + # attach: send the prompt and consume typed events as they stream in. + stream = agent.run_streamed(prompt) + for event in stream: + if event.type == AgentEventType.TOKEN: + print(event.text, end="", flush=True) # live reply -> stdout + elif event.type == AgentEventType.TOOL_CALL: + print(f"\n[tool] {event.tool_name}", file=sys.stderr) + elif event.type == AgentEventType.HITL_REQUESTED: + print(f"\n[hitl auto-approved] {event.request_id}", file=sys.stderr) + + result = stream.result + print( + f"\n\n[{result.status}] " + f"tokens in={result.usage.get('tokens_in')} " + f"out={result.usage.get('tokens_out')} " + f"| captured {len(result.final_output)} chars", + file=sys.stderr, + ) + + # session is destroyed here. + return 0 if stream.status == "completed" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/pydo/agents/__init__.py b/src/pydo/agents/__init__.py index 659c6aa6..60ada7c7 100644 --- a/src/pydo/agents/__init__.py +++ b/src/pydo/agents/__init__.py @@ -23,6 +23,14 @@ SessionStatus, ) from .custom_sessions import HarnessEventStream, HarnessStreamError, SessionsOperations +from .session import ( + AgentEvent, + AgentEventType, + AgentSession, + HITLPolicy, + RunResult, + RunStream, +) DEFAULT_AGENTS_BASE_URL = "https://api.digitalocean.com" _ENV_VAR = "PYDO_AGENTS_ENDPOINT" @@ -48,9 +56,33 @@ def __init__(self, parent_client, *, agents_endpoint: Optional[str] = None): def base_url(self) -> str: return self._proxy._base_url + def start(self, manifest: "str | bytes") -> AgentSession: + """Create a session from an ``agents.yaml`` manifest and return a handle. + + Use as a context manager to auto-destroy on exit:: + + with client.agents.start(manifest) as agent: + print(agent.run("hello").final_output) + """ + resp = self.sessions.create_from_manifest(manifest) + get = getattr(resp, "get", None) + info = get("session") if get else None + session_id = (getattr(info or resp, "get", lambda *_: None))("session_id") + return AgentSession(self.sessions, session_id, raw=resp) + + def attach(self, session_id: str) -> AgentSession: + """Return an :class:`AgentSession` handle for an existing session.""" + return AgentSession(self.sessions, session_id) + __all__ = [ "AgentsResources", + "AgentSession", + "AgentEvent", + "AgentEventType", + "RunResult", + "RunStream", + "HITLPolicy", "SessionsOperations", "HarnessEventStream", "HarnessStreamError", diff --git a/src/pydo/agents/custom_sessions.py b/src/pydo/agents/custom_sessions.py index 0e811a91..fbe07752 100644 --- a/src/pydo/agents/custom_sessions.py +++ b/src/pydo/agents/custom_sessions.py @@ -3,10 +3,11 @@ # Licensed under the Apache-2.0 License. # ------------------------------------ """Sync Hosted Agents session operations (``/v2/agents/sessions/...``).""" + from __future__ import annotations import json as _json -from typing import Any, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional, Union from urllib.parse import quote from azure.core.exceptions import ( @@ -30,6 +31,24 @@ _BASE_PATH = "/v2/agents/sessions" +# Private Beta contract: a session is created from an ``agents.yaml`` manifest +# uploaded verbatim. The server routes on this media type (handlers.go: +# isYAMLContentType) and parses the body as the agent spec. +_YAML_MEDIA_TYPE = "application/x-yaml" + + +def _manifest_bytes(manifest: Union[str, bytes]) -> bytes: + """Normalize an agents.yaml manifest to non-empty UTF-8 bytes.""" + if isinstance(manifest, str): + data = manifest.encode("utf-8") + elif isinstance(manifest, (bytes, bytearray)): + data = bytes(manifest) + else: + raise TypeError("manifest must be a str or bytes YAML document") + if not data.strip(): + raise ValueError("manifest is empty") + return data + def _unwrap_harness_sse_chunk(chunk: Dict[str, Any]) -> Optional[Any]: """Normalize SSE JSON to a harness Event. @@ -139,6 +158,8 @@ def _send( path: str, *, body: Optional[Dict[str, Any]] = None, + content: Optional[Union[str, bytes]] = None, + content_type: Optional[str] = None, params: Optional[Dict[str, Any]] = None, stream: bool = False, ): @@ -151,6 +172,10 @@ def _send( if body is not None: headers["Content-Type"] = "application/json" kwargs["json"] = body + elif content is not None: + if content_type: + headers["Content-Type"] = content_type + kwargs["content"] = content request = HttpRequest(method, path, **kwargs) request.url = self._client.format_url(request.url) @@ -190,19 +215,25 @@ def list( ), ) - def create( - self, - *, - agent_kind: str, - repo_hint: Optional[str] = None, - idle_timeout_seconds: Optional[int] = None, - ) -> Any: - body: Dict[str, Any] = {"agent_kind": agent_kind} - if repo_hint is not None: - body["repo_hint"] = repo_hint - if idle_timeout_seconds is not None: - body["idle_timeout_seconds"] = idle_timeout_seconds - return self._parse_json(self._send("POST", _BASE_PATH, body=body)) + def create_from_manifest(self, manifest: Union[str, bytes]) -> Any: + """Create a session from an ``agents.yaml`` manifest. + + This is the supported creation path: the manifest defines everything + about the session (runtime adapter, sandbox, env vars, egress). It is + uploaded verbatim as ``application/x-yaml`` and the server owns parsing + and validation. There are no ``agent_kind``/``repo_hint`` arguments. + + :param manifest: The agent spec as a YAML ``str`` or ``bytes`` document. + """ + data = _manifest_bytes(manifest) + return self._parse_json( + self._send( + "POST", + _BASE_PATH, + content=data, + content_type=_YAML_MEDIA_TYPE, + ), + ) def get(self, session_id: str) -> Any: return self._parse_json( diff --git a/src/pydo/agents/session.py b/src/pydo/agents/session.py new file mode 100644 index 00000000..711212d6 --- /dev/null +++ b/src/pydo/agents/session.py @@ -0,0 +1,414 @@ +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""High-level, self-contained agent interface over the Hosted Agents session API. + +The low-level :class:`~pydo.agents.custom_sessions.SessionsOperations` mirrors the +REST endpoints one-to-one, which means consuming a run requires hand-wiring the +SSE feed: spawn a reader, dispatch on raw ``run.*`` event strings, accumulate +token deltas, resolve HITL prompts, and coordinate completion. + +This module wraps that into an ergonomic surface inspired by +``openai-agents-python``:: + + with client.agents.start(manifest) as agent: # create + auto-destroy + result = agent.run("Summarize the repo") # blocking + print(result.final_output) + + for event in agent.run_streamed("Now add tests"): # streamed, typed + if event.type == AgentEventType.TOKEN: + print(event.text, end="") + +No threads, no raw event-string matching, no manual teardown. +""" +from __future__ import annotations + +import time +from typing import Any, Callable, Dict, List, Optional, Union + +from .custom_models import HITLOutcome, ResolutionSource, SessionStatus + +# Normalized event type -> raw SPI ``type`` it maps from. +_RAW_TO_TYPE = { + "run.started": "run_started", + "run.token_delta": "token", + "run.tool_call_started": "tool_call", + "run.tool_call_completed": "tool_result", + "run.human_input_requested": "hitl_requested", + "run.human_input_received": "hitl_resolved", + "run.completed": "completed", + "run.failed": "failed", +} + +_HITL_OUTCOMES = { + "approve": HITLOutcome.APPROVE, + "reject": HITLOutcome.REJECT, + "defer": HITLOutcome.DEFER, +} + +# A HITL policy is either a fixed decision ("approve"/"reject"/"defer"), an +# explicit HITLOutcome constant, a callable mapping an event -> decision, or +# None to leave prompts unresolved for the caller to handle. +HITLPolicy = Union[str, Callable[["AgentEvent"], Optional[str]], None] + + +class AgentEventType: + """Normalized, friendly event kinds yielded by a run stream.""" + + RUN_STARTED = "run_started" + TOKEN = "token" + TOOL_CALL = "tool_call" + TOOL_RESULT = "tool_result" + HITL_REQUESTED = "hitl_requested" + HITL_RESOLVED = "hitl_resolved" + COMPLETED = "completed" + FAILED = "failed" + OTHER = "other" + + +class AgentEvent: + """A normalized view over a raw harness SSE event. + + Exposes a stable :attr:`type` (see :class:`AgentEventType`) plus typed + accessors, while keeping the underlying event available as :attr:`raw`. + """ + + def __init__(self, raw: Any): + self.raw = raw + get = getattr(raw, "get", None) + self.raw_type: Optional[str] = get("type") if get else None + self.type = _RAW_TO_TYPE.get(self.raw_type or "", AgentEventType.OTHER) + data = get("data") if get else None + self.data = data if data is not None else {} + self.run_id: str = (get("run_id") if get else None) or "" + + def _d(self, key: str, default: Any = None) -> Any: + get = getattr(self.data, "get", None) + return get(key, default) if get else default + + @property + def text(self) -> str: + """Token text for ``TOKEN`` events ("" otherwise).""" + return self._d("text", "") or "" + + @property + def tool_name(self) -> Optional[str]: + """Tool name for ``TOOL_CALL`` events.""" + return self._d("name") + + @property + def request_id(self) -> Optional[str]: + """HITL request id for ``HITL_REQUESTED`` / ``HITL_RESOLVED`` events.""" + return self._d("hitl_id") or self._d("request_id") + + @property + def usage(self) -> Dict[str, Any]: + """Token/cost totals for ``COMPLETED`` events.""" + return { + "tokens_in": self._d("total_tokens_in"), + "tokens_out": self._d("total_tokens_out"), + "cost_micros": self._d("run_cost_micros"), + } + + @property + def error(self) -> Optional[Dict[str, Any]]: + """Failure ``{code, message}`` for ``FAILED`` events.""" + if self.type != AgentEventType.FAILED: + return None + return {"code": self._d("code"), "message": self._d("message")} + + def __repr__(self) -> str: + return f"AgentEvent(type={self.type!r}, run_id={self.run_id!r})" + + +class RunResult: + """The outcome of a single run: assembled output, usage, and raw events.""" + + def __init__( + self, + *, + run_id: Optional[str], + final_output: str, + events: List[AgentEvent], + usage: Dict[str, Any], + status: str, + error: Optional[Dict[str, Any]] = None, + ): + self.run_id = run_id + self.final_output = final_output + self.events = events + self.usage = usage + self.status = status # "completed" | "failed" | "timeout" + self.error = error + + @property + def ok(self) -> bool: + return self.status == "completed" + + def __str__(self) -> str: + return self.final_output + + def __repr__(self) -> str: + return ( + f"RunResult(status={self.status!r}, run_id={self.run_id!r}, " + f"chars={len(self.final_output)})" + ) + + +def _decide_hitl(policy: HITLPolicy, event: AgentEvent) -> Optional[str]: + """Resolve a HITL policy to a concrete outcome constant (or None).""" + decision: Any = policy(event) if callable(policy) else policy + if not decision: + return None + if isinstance(decision, str): + return _HITL_OUTCOMES.get(decision.lower(), decision) + return decision + + +class RunStream: + """Iterable of :class:`AgentEvent` for one run. + + Iterating yields normalized events, auto-resolves HITL prompts per the + configured policy, and accumulates the assembled output. After iteration, + :attr:`final_output`, :attr:`usage`, :attr:`status`, and :attr:`result` + are populated. + """ + + def __init__( + self, + *, + raw_stream: Any, + run_id: Optional[str], + session: "AgentSession", + hitl: HITLPolicy = "approve", + timeout: Optional[float] = None, + ): + self._raw = raw_stream + self.run_id = run_id + self._session = session + self._hitl = hitl + self._timeout = timeout + self._chunks: List[str] = [] + self.events: List[AgentEvent] = [] + self.usage: Dict[str, Any] = {} + self.status = "running" + self.error: Optional[Dict[str, Any]] = None + + def __iter__(self): + deadline = time.monotonic() + self._timeout if self._timeout else None + try: + for raw in self._raw: + event = AgentEvent(raw) + self.events.append(event) + + if event.type == AgentEventType.TOKEN: + self._chunks.append(event.text) + elif event.type == AgentEventType.HITL_REQUESTED: + self._auto_resolve(event) + + yield event + + if self._is_terminal(event): + break + if deadline and time.monotonic() > deadline: + self.status = "timeout" + break + finally: + self.close() + + def _auto_resolve(self, event: AgentEvent) -> None: + outcome = _decide_hitl(self._hitl, event) + if not outcome or not event.request_id: + return + try: + self._session.resolve_hitl( + event.request_id, + outcome=outcome, + source=ResolutionSource.OUT_OF_BAND, + ) + except Exception: # noqa: BLE001 - best-effort; surfaced via the feed + pass + + def _is_terminal(self, event: AgentEvent) -> bool: + if self.run_id and event.run_id and event.run_id != self.run_id: + return False + if event.type == AgentEventType.COMPLETED: + self.usage = event.usage + self.status = "completed" + return True + if event.type == AgentEventType.FAILED: + self.error = event.error + self.status = "failed" + return True + return False + + @property + def final_output(self) -> str: + return "".join(self._chunks).strip() + + @property + def result(self) -> RunResult: + return RunResult( + run_id=self.run_id, + final_output=self.final_output, + events=self.events, + usage=self.usage, + status=self.status, + error=self.error, + ) + + def close(self) -> None: + closer = getattr(self._raw, "close", None) + if closer: + try: + closer() + except Exception: # noqa: BLE001 + pass + + def __enter__(self) -> "RunStream": + return self + + def __exit__(self, *args: Any) -> None: + self.close() + + +class AgentSession: + """A self-managing handle to one hosted-agent session. + + Wraps :class:`~pydo.agents.custom_sessions.SessionsOperations`, binding the + ``session_id`` so callers never re-pass it, and adds :meth:`run` / + :meth:`run_streamed`. Use as a context manager to auto-destroy on exit. + """ + + def __init__(self, sessions: Any, session_id: str, *, raw: Any = None): + self._sessions = sessions + self.session_id = session_id + self._raw = raw + + @property + def sessions(self) -> Any: + """The underlying low-level operations object.""" + return self._sessions + + @property + def info(self) -> Any: + """The most recent session object (``{...}``), if known.""" + raw = self._raw + if raw is None: + return None + get = getattr(raw, "get", None) + inner = get("session") if get else None + return inner if inner is not None else raw + + @property + def status(self) -> Optional[str]: + info = self.info + get = getattr(info, "get", None) + return get("status") if get else None + + def refresh(self) -> Any: + """Fetch and cache the latest session state.""" + self._raw = self._sessions.get(self.session_id) + return self.info + + def wait_until_ready( + self, *, timeout: float = 120.0, poll_interval: float = 2.0 + ) -> "AgentSession": + """Block until the session reports ``READY`` (or raise).""" + deadline = time.monotonic() + timeout + while True: + status = (getattr(self.refresh(), "get", lambda *_: None))("status") + if status == SessionStatus.READY: + return self + if status in (SessionStatus.FAILED, SessionStatus.DESTROYED): + raise RuntimeError(f"session {self.session_id} is {status}") + if time.monotonic() > deadline: + raise TimeoutError( + f"session {self.session_id} not ready after {timeout}s" + ) + time.sleep(poll_interval) + + def run_streamed( + self, + prompt: str, + *, + hitl: HITLPolicy = "approve", + timeout: Optional[float] = 300.0, + ) -> RunStream: + """Send ``prompt`` and return a :class:`RunStream` of typed events. + + The SSE subscription is opened *before* the input is submitted, so no + early events are missed — no caller-managed thread required. + """ + raw_stream = self._sessions.stream(self.session_id) + run = self._sessions.send_input(self.session_id, text=prompt) + run_id = (getattr(run, "get", lambda *_: None))("run_id") + return RunStream( + raw_stream=raw_stream, + run_id=run_id, + session=self, + hitl=hitl, + timeout=timeout, + ) + + def run( + self, + prompt: str, + *, + hitl: HITLPolicy = "approve", + timeout: Optional[float] = 300.0, + ) -> RunResult: + """Send ``prompt`` and block until the run finishes, returning a result.""" + stream = self.run_streamed(prompt, hitl=hitl, timeout=timeout) + for _ in stream: + pass + return stream.result + + # --- thin passthroughs (session id bound) --------------------------- + def send_input(self, text: str) -> Any: + return self._sessions.send_input(self.session_id, text=text) + + def stream(self, **kwargs: Any) -> Any: + return self._sessions.stream(self.session_id, **kwargs) + + def resolve_hitl( + self, + request_id: str, + *, + outcome: str, + reason: Optional[str] = None, + source: Optional[str] = None, + ) -> None: + self._sessions.resolve_hitl( + self.session_id, + request_id, + outcome=outcome, + reason=reason, + source=source, + ) + + def destroy(self) -> None: + self._sessions.destroy(self.session_id) + + def __enter__(self) -> "AgentSession": + return self + + def __exit__(self, *args: Any) -> None: + try: + self.destroy() + except Exception: # noqa: BLE001 - teardown best-effort + pass + + def __repr__(self) -> str: + return f"AgentSession(session_id={self.session_id!r})" + + +__all__ = [ + "AgentEvent", + "AgentEventType", + "AgentSession", + "RunResult", + "RunStream", + "HITLPolicy", +] diff --git a/src/pydo/aio/agents/__init__.py b/src/pydo/aio/agents/__init__.py index dd184182..8c29af31 100644 --- a/src/pydo/aio/agents/__init__.py +++ b/src/pydo/aio/agents/__init__.py @@ -11,6 +11,7 @@ from pydo.custom_extensions import _BaseURLProxy from .custom_sessions import AsyncHarnessEventStream, AsyncSessionsOperations +from .session import AsyncAgentSession, AsyncRunStream class AsyncAgentsResources: @@ -25,5 +26,27 @@ def __init__(self, parent_client, *, agents_endpoint: Optional[str] = None): def base_url(self) -> str: return self._proxy._base_url - -__all__ = ["AsyncAgentsResources", "AsyncSessionsOperations", "AsyncHarnessEventStream"] + async def start(self, manifest: "str | bytes") -> AsyncAgentSession: + """Create a session from an ``agents.yaml`` manifest and return a handle. + + Use as ``async with await client.agents.start(manifest) as agent:`` to + auto-destroy on exit. + """ + resp = await self.sessions.create_from_manifest(manifest) + get = getattr(resp, "get", None) + info = get("session") if get else None + session_id = (getattr(info or resp, "get", lambda *_: None))("session_id") + return AsyncAgentSession(self.sessions, session_id, raw=resp) + + def attach(self, session_id: str) -> AsyncAgentSession: + """Return an :class:`AsyncAgentSession` handle for an existing session.""" + return AsyncAgentSession(self.sessions, session_id) + + +__all__ = [ + "AsyncAgentsResources", + "AsyncAgentSession", + "AsyncRunStream", + "AsyncSessionsOperations", + "AsyncHarnessEventStream", +] diff --git a/src/pydo/aio/agents/custom_sessions.py b/src/pydo/aio/agents/custom_sessions.py index 64f1bbde..a0764cf7 100644 --- a/src/pydo/aio/agents/custom_sessions.py +++ b/src/pydo/aio/agents/custom_sessions.py @@ -3,10 +3,11 @@ # Licensed under the Apache-2.0 License. # ------------------------------------ """Async Hosted Agents session operations.""" + from __future__ import annotations import json as _json -from typing import Any, AsyncIterator, Dict, List, Optional +from typing import Any, AsyncIterator, Dict, List, Optional, Union from urllib.parse import quote from azure.core.exceptions import ( @@ -19,7 +20,13 @@ ) from azure.core.rest import HttpRequest -from pydo.agents.custom_sessions import HarnessStreamError, _raise_agents_http_error, _unwrap_harness_sse_chunk +from pydo.agents.custom_sessions import ( + _YAML_MEDIA_TYPE, + HarnessStreamError, + _manifest_bytes, + _raise_agents_http_error, + _unwrap_harness_sse_chunk, +) from pydo.custom_extensions import AsyncSSEStream, _wrap _ERROR_MAP = { @@ -80,6 +87,8 @@ async def _send( path: str, *, body: Optional[Dict[str, Any]] = None, + content: Optional[Union[str, bytes]] = None, + content_type: Optional[str] = None, params: Optional[Dict[str, Any]] = None, stream: bool = False, ): @@ -92,6 +101,10 @@ async def _send( if body is not None: headers["Content-Type"] = "application/json" kwargs["json"] = body + elif content is not None: + if content_type: + headers["Content-Type"] = content_type + kwargs["content"] = content request = HttpRequest(method, path, **kwargs) request.url = self._client.format_url(request.url) @@ -131,20 +144,24 @@ async def list( ), ) - async def create( - self, - *, - agent_kind: str, - repo_hint: Optional[str] = None, - idle_timeout_seconds: Optional[int] = None, - ) -> Any: - body: Dict[str, Any] = {"agent_kind": agent_kind} - if repo_hint is not None: - body["repo_hint"] = repo_hint - if idle_timeout_seconds is not None: - body["idle_timeout_seconds"] = idle_timeout_seconds + async def create_from_manifest(self, manifest: Union[str, bytes]) -> Any: + """Create a session from an ``agents.yaml`` manifest. + + This is the supported creation path: the manifest defines everything + about the session (runtime adapter, sandbox, env vars, egress). It is + uploaded verbatim as ``application/x-yaml`` and the server owns parsing + and validation. There are no ``agent_kind``/``repo_hint`` arguments. + + :param manifest: The agent spec as a YAML ``str`` or ``bytes`` document. + """ + data = _manifest_bytes(manifest) return await self._parse_json( - await self._send("POST", _BASE_PATH, body=body), + await self._send( + "POST", + _BASE_PATH, + content=data, + content_type=_YAML_MEDIA_TYPE, + ), ) async def get(self, session_id: str) -> Any: diff --git a/src/pydo/aio/agents/session.py b/src/pydo/aio/agents/session.py new file mode 100644 index 00000000..0b3e8139 --- /dev/null +++ b/src/pydo/aio/agents/session.py @@ -0,0 +1,245 @@ +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""Async high-level agent interface (see :mod:`pydo.agents.session`).""" +from __future__ import annotations + +import asyncio +from typing import Any, Dict, List, Optional + +from pydo.agents.custom_models import ResolutionSource, SessionStatus +from pydo.agents.session import ( + AgentEvent, + AgentEventType, + HITLPolicy, + RunResult, + _decide_hitl, +) + + +class AsyncRunStream: + """Async iterable of :class:`~pydo.agents.session.AgentEvent` for one run.""" + + def __init__( + self, + *, + raw_stream: Any, + run_id: Optional[str], + session: "AsyncAgentSession", + hitl: HITLPolicy = "approve", + timeout: Optional[float] = None, + ): + self._raw = raw_stream + self.run_id = run_id + self._session = session + self._hitl = hitl + self._timeout = timeout + self._chunks: List[str] = [] + self.events: List[AgentEvent] = [] + self.usage: Dict[str, Any] = {} + self.status = "running" + self.error: Optional[Dict[str, Any]] = None + + def __aiter__(self): + return self._iter() + + async def _iter(self): + deadline = ( + (asyncio.get_event_loop().time() + self._timeout) if self._timeout else None + ) + try: + async for raw in self._raw: + event = AgentEvent(raw) + self.events.append(event) + + if event.type == AgentEventType.TOKEN: + self._chunks.append(event.text) + elif event.type == AgentEventType.HITL_REQUESTED: + await self._auto_resolve(event) + + yield event + + if self._is_terminal(event): + break + if deadline and asyncio.get_event_loop().time() > deadline: + self.status = "timeout" + break + finally: + await self.close() + + async def _auto_resolve(self, event: AgentEvent) -> None: + outcome = _decide_hitl(self._hitl, event) + if not outcome or not event.request_id: + return + try: + await self._session.resolve_hitl( + event.request_id, + outcome=outcome, + source=ResolutionSource.OUT_OF_BAND, + ) + except Exception: # noqa: BLE001 - best-effort; surfaced via the feed + pass + + def _is_terminal(self, event: AgentEvent) -> bool: + if self.run_id and event.run_id and event.run_id != self.run_id: + return False + if event.type == AgentEventType.COMPLETED: + self.usage = event.usage + self.status = "completed" + return True + if event.type == AgentEventType.FAILED: + self.error = event.error + self.status = "failed" + return True + return False + + @property + def final_output(self) -> str: + return "".join(self._chunks).strip() + + @property + def result(self) -> RunResult: + return RunResult( + run_id=self.run_id, + final_output=self.final_output, + events=self.events, + usage=self.usage, + status=self.status, + error=self.error, + ) + + async def close(self) -> None: + closer = getattr(self._raw, "close", None) + if closer: + try: + await closer() + except Exception: # noqa: BLE001 + pass + + async def __aenter__(self) -> "AsyncRunStream": + return self + + async def __aexit__(self, *args: Any) -> None: + await self.close() + + +class AsyncAgentSession: + """Async self-managing handle to one hosted-agent session.""" + + def __init__(self, sessions: Any, session_id: str, *, raw: Any = None): + self._sessions = sessions + self.session_id = session_id + self._raw = raw + + @property + def sessions(self) -> Any: + return self._sessions + + @property + def info(self) -> Any: + raw = self._raw + if raw is None: + return None + get = getattr(raw, "get", None) + inner = get("session") if get else None + return inner if inner is not None else raw + + @property + def status(self) -> Optional[str]: + info = self.info + get = getattr(info, "get", None) + return get("status") if get else None + + async def refresh(self) -> Any: + self._raw = await self._sessions.get(self.session_id) + return self.info + + async def wait_until_ready( + self, *, timeout: float = 120.0, poll_interval: float = 2.0 + ) -> "AsyncAgentSession": + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while True: + info = await self.refresh() + status = (getattr(info, "get", lambda *_: None))("status") + if status == SessionStatus.READY: + return self + if status in (SessionStatus.FAILED, SessionStatus.DESTROYED): + raise RuntimeError(f"session {self.session_id} is {status}") + if loop.time() > deadline: + raise TimeoutError( + f"session {self.session_id} not ready after {timeout}s" + ) + await asyncio.sleep(poll_interval) + + async def run_streamed( + self, + prompt: str, + *, + hitl: HITLPolicy = "approve", + timeout: Optional[float] = 300.0, + ) -> AsyncRunStream: + raw_stream = await self._sessions.stream(self.session_id) + run = await self._sessions.send_input(self.session_id, text=prompt) + run_id = (getattr(run, "get", lambda *_: None))("run_id") + return AsyncRunStream( + raw_stream=raw_stream, + run_id=run_id, + session=self, + hitl=hitl, + timeout=timeout, + ) + + async def run( + self, + prompt: str, + *, + hitl: HITLPolicy = "approve", + timeout: Optional[float] = 300.0, + ) -> RunResult: + stream = await self.run_streamed(prompt, hitl=hitl, timeout=timeout) + async for _ in stream: + pass + return stream.result + + # --- thin passthroughs (session id bound) --------------------------- + async def send_input(self, text: str) -> Any: + return await self._sessions.send_input(self.session_id, text=text) + + async def stream(self, **kwargs: Any) -> Any: + return await self._sessions.stream(self.session_id, **kwargs) + + async def resolve_hitl( + self, + request_id: str, + *, + outcome: str, + reason: Optional[str] = None, + source: Optional[str] = None, + ) -> None: + await self._sessions.resolve_hitl( + self.session_id, + request_id, + outcome=outcome, + reason=reason, + source=source, + ) + + async def destroy(self) -> None: + await self._sessions.destroy(self.session_id) + + async def __aenter__(self) -> "AsyncAgentSession": + return self + + async def __aexit__(self, *args: Any) -> None: + try: + await self.destroy() + except Exception: # noqa: BLE001 - teardown best-effort + pass + + def __repr__(self) -> str: + return f"AsyncAgentSession(session_id={self.session_id!r})" + + +__all__ = ["AsyncAgentSession", "AsyncRunStream"] diff --git a/tests/agents/test_async_sessions.py b/tests/agents/test_async_sessions.py new file mode 100644 index 00000000..1acb8e6e --- /dev/null +++ b/tests/agents/test_async_sessions.py @@ -0,0 +1,85 @@ +# pylint: disable=line-too-long,missing-class-docstring,missing-function-docstring,protected-access +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""Unit tests for :mod:`pydo.aio.agents.custom_sessions`.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from typing import Any, List +from unittest.mock import MagicMock + +import pytest + +from pydo.aio.agents import AsyncAgentsResources + + +class _FakeAsyncResponse: + def __init__(self, status_code: int, body: Any = None): + self.status_code = status_code + if isinstance(body, (dict, list)): + self._body_bytes = json.dumps(body).encode("utf-8") + elif isinstance(body, str): + self._body_bytes = body.encode("utf-8") + elif isinstance(body, bytes): + self._body_bytes = body + else: + self._body_bytes = b"" + + async def read(self) -> bytes: + return self._body_bytes + + def text(self) -> str: + return self._body_bytes.decode("utf-8") + + def body(self) -> bytes: + return self._body_bytes + + +class _FakeAsyncPipeline: + def __init__(self, responses: List[_FakeAsyncResponse]): + self._responses = list(responses) + self.calls: List[Any] = [] + + async def run(self, request, *, stream=False): + self.calls.append(SimpleNamespace(request=request, stream=stream)) + return SimpleNamespace(http_response=self._responses.pop(0)) + + +def _make_async_resources(responses: List[_FakeAsyncResponse]) -> AsyncAgentsResources: + parent = MagicMock() + parent._client = MagicMock() + parent._client._pipeline = _FakeAsyncPipeline(responses) + return AsyncAgentsResources( + parent, agents_endpoint="https://api.stage2.digitalocean.com" + ) + + +@pytest.mark.asyncio +async def test_async_create_from_manifest_uploads_yaml_verbatim(): + resources = _make_async_resources( + [_FakeAsyncResponse(200, {"session": {"session_id": "abc"}})] + ) + manifest = "apiVersion: agents.digitalocean.com/v1alpha1\nkind: Agent\n" + + resp = await resources.sessions.create_from_manifest(manifest) + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.method == "POST" + assert call.request.url.endswith("/v2/agents/sessions") + assert call.request.headers.get("Content-Type") == "application/x-yaml" + content = call.request.content + if isinstance(content, bytes): + content = content.decode("utf-8") + assert content == manifest + assert resp.session.session_id == "abc" + + +@pytest.mark.asyncio +async def test_async_create_from_manifest_rejects_empty(): + resources = _make_async_resources([]) + with pytest.raises(ValueError): + await resources.sessions.create_from_manifest("") diff --git a/tests/agents/test_session_highlevel.py b/tests/agents/test_session_highlevel.py new file mode 100644 index 00000000..6d51b2d5 --- /dev/null +++ b/tests/agents/test_session_highlevel.py @@ -0,0 +1,292 @@ +# pylint: disable=line-too-long,missing-class-docstring,missing-function-docstring,protected-access +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""Unit tests for the high-level agent interface (:mod:`pydo.agents.session`).""" +from __future__ import annotations + +from typing import Any, List + +import pytest + +from pydo.agents import AgentEvent, AgentEventType, AgentSession, RunResult +from pydo.agents.custom_models import HITLOutcome +from pydo.aio.agents import AsyncAgentSession + + +# --------------------------------------------------------------------------- +# Sync fakes +# --------------------------------------------------------------------------- + + +class _FakeRawStream: + def __init__(self, events: List[dict]): + self._events = events + self.closed = False + + def __iter__(self): + return iter(self._events) + + def close(self): + self.closed = True + + +class _FakeSessions: + """Stands in for SessionsOperations, recording call order.""" + + def __init__(self, events: List[dict]): + self._events = events + self.calls: List[Any] = [] + self.resolved: List[Any] = [] + self.destroyed = False + + def create_from_manifest(self, manifest): + self.calls.append(("create", manifest)) + return {"session": {"session_id": "s1", "status": "SESSION_STATUS_READY"}} + + def stream(self, session_id, **kwargs): + self.calls.append(("stream", session_id)) + return _FakeRawStream(self._events) + + def send_input(self, session_id, *, text): + self.calls.append(("send_input", session_id, text)) + return {"run_id": "r1"} + + def resolve_hitl( + self, session_id, request_id, *, outcome, reason=None, source=None + ): + self.calls.append(("resolve_hitl", session_id, request_id, outcome)) + self.resolved.append((request_id, outcome)) + + def get(self, session_id): + return {"session": {"session_id": session_id, "status": "SESSION_STATUS_READY"}} + + def destroy(self, session_id): + self.calls.append(("destroy", session_id)) + self.destroyed = True + + +# --------------------------------------------------------------------------- +# AgentEvent normalization +# --------------------------------------------------------------------------- + + +def test_agent_event_normalizes_token(): + ev = AgentEvent({"type": "run.token_delta", "data": {"text": "hi"}, "run_id": "r1"}) + assert ev.type == AgentEventType.TOKEN + assert ev.text == "hi" + assert ev.run_id == "r1" + + +def test_agent_event_normalizes_completed_and_failed(): + done = AgentEvent({"type": "run.completed", "data": {"total_tokens_out": 7}}) + assert done.type == AgentEventType.COMPLETED + assert done.usage["tokens_out"] == 7 + + failed = AgentEvent({"type": "run.failed", "data": {"code": 3, "message": "boom"}}) + assert failed.type == AgentEventType.FAILED + assert failed.error == {"code": 3, "message": "boom"} + + +def test_agent_event_unknown_type_is_other(): + assert AgentEvent({"type": "session.updated"}).type == AgentEventType.OTHER + + +# --------------------------------------------------------------------------- +# run / run_streamed +# --------------------------------------------------------------------------- + + +def test_run_collects_final_output_usage_and_order(): + events = [ + {"type": "run.started", "data": {}}, + {"type": "run.token_delta", "data": {"text": "Hello "}}, + {"type": "run.token_delta", "data": {"text": "world"}}, + { + "type": "run.completed", + "data": { + "total_tokens_in": 3, + "total_tokens_out": 5, + "run_cost_micros": 1234, + }, + }, + ] + sessions = _FakeSessions(events) + result = AgentSession(sessions, "s1").run("hi") + + assert isinstance(result, RunResult) + assert result.final_output == "Hello world" + assert result.status == "completed" and result.ok + assert result.usage["tokens_out"] == 5 + assert result.run_id == "r1" + # stream is opened BEFORE input is sent (so no early events are missed) + kinds = [c[0] for c in sessions.calls] + assert kinds.index("stream") < kinds.index("send_input") + + +def test_run_streamed_yields_typed_events(): + events = [ + {"type": "run.token_delta", "data": {"text": "hi"}}, + {"type": "run.completed", "data": {}}, + ] + stream = AgentSession(_FakeSessions(events), "s1").run_streamed("x") + seen = [e.type for e in stream] + assert seen == [AgentEventType.TOKEN, AgentEventType.COMPLETED] + assert stream.final_output == "hi" + assert stream.status == "completed" + + +def test_run_failed_sets_status_and_error(): + events = [ + {"type": "run.token_delta", "data": {"text": "partial"}}, + {"type": "run.failed", "data": {"code": 3, "message": "boom"}}, + ] + result = AgentSession(_FakeSessions(events), "s1").run("x") + assert result.status == "failed" and not result.ok + assert result.error["message"] == "boom" + + +# --------------------------------------------------------------------------- +# HITL policy +# --------------------------------------------------------------------------- + + +def test_run_auto_approves_hitl_by_default(): + events = [ + {"type": "run.human_input_requested", "data": {"hitl_id": "req-1"}}, + {"type": "run.token_delta", "data": {"text": "done"}}, + {"type": "run.completed", "data": {}}, + ] + sessions = _FakeSessions(events) + result = AgentSession(sessions, "s1").run("do it") + assert ("req-1", HITLOutcome.APPROVE) in sessions.resolved + assert result.final_output == "done" + + +def test_hitl_callable_policy_receives_event(): + events = [ + {"type": "run.human_input_requested", "data": {"hitl_id": "req-9"}}, + {"type": "run.completed", "data": {}}, + ] + sessions = _FakeSessions(events) + seen = [] + + def policy(event): + seen.append(event.request_id) + return "reject" + + AgentSession(sessions, "s1").run("x", hitl=policy) + assert seen == ["req-9"] + assert ("req-9", HITLOutcome.REJECT) in sessions.resolved + + +def test_hitl_none_leaves_prompt_unresolved(): + events = [ + {"type": "run.human_input_requested", "data": {"hitl_id": "req-9"}}, + {"type": "run.completed", "data": {}}, + ] + sessions = _FakeSessions(events) + AgentSession(sessions, "s1").run("x", hitl=None) + assert sessions.resolved == [] + + +# --------------------------------------------------------------------------- +# Lifecycle +# --------------------------------------------------------------------------- + + +def test_context_manager_destroys(): + sessions = _FakeSessions([]) + with AgentSession(sessions, "s1") as agent: + assert agent.session_id == "s1" + assert sessions.destroyed + + +def test_run_stream_closes_underlying_stream(): + events = [{"type": "run.completed", "data": {}}] + sessions = _FakeSessions(events) + stream = AgentSession(sessions, "s1").run_streamed("x") + raw = stream._raw + list(stream) + assert raw.closed + + +# --------------------------------------------------------------------------- +# Async parity +# --------------------------------------------------------------------------- + + +class _FakeAsyncRawStream: + def __init__(self, events: List[dict]): + self._events = events + self.closed = False + + def __aiter__(self): + return self._gen() + + async def _gen(self): + for ev in self._events: + yield ev + + async def close(self): + self.closed = True + + +class _FakeAsyncSessions: + def __init__(self, events: List[dict]): + self._events = events + self.calls: List[Any] = [] + self.resolved: List[Any] = [] + self.destroyed = False + + async def stream(self, session_id, **kwargs): + self.calls.append(("stream", session_id)) + return _FakeAsyncRawStream(self._events) + + async def send_input(self, session_id, *, text): + self.calls.append(("send_input", session_id, text)) + return {"run_id": "r1"} + + async def resolve_hitl( + self, session_id, request_id, *, outcome, reason=None, source=None + ): + self.resolved.append((request_id, outcome)) + + async def destroy(self, session_id): + self.destroyed = True + + +@pytest.mark.asyncio +async def test_async_run_collects_output_and_order(): + events = [ + {"type": "run.token_delta", "data": {"text": "Hi "}}, + {"type": "run.token_delta", "data": {"text": "there"}}, + {"type": "run.completed", "data": {"total_tokens_out": 2}}, + ] + sessions = _FakeAsyncSessions(events) + result = await AsyncAgentSession(sessions, "s1").run("x") + assert result.final_output == "Hi there" + assert result.status == "completed" + assert result.usage["tokens_out"] == 2 + kinds = [c[0] for c in sessions.calls] + assert kinds.index("stream") < kinds.index("send_input") + + +@pytest.mark.asyncio +async def test_async_auto_approves_hitl(): + events = [ + {"type": "run.human_input_requested", "data": {"hitl_id": "req-1"}}, + {"type": "run.completed", "data": {}}, + ] + sessions = _FakeAsyncSessions(events) + await AsyncAgentSession(sessions, "s1").run("x") + assert ("req-1", HITLOutcome.APPROVE) in sessions.resolved + + +@pytest.mark.asyncio +async def test_async_context_manager_destroys(): + sessions = _FakeAsyncSessions([]) + async with AsyncAgentSession(sessions, "s1") as agent: + assert agent.session_id == "s1" + assert sessions.destroyed diff --git a/tests/agents/test_sessions.py b/tests/agents/test_sessions.py index b381c274..ccc23406 100644 --- a/tests/agents/test_sessions.py +++ b/tests/agents/test_sessions.py @@ -4,6 +4,7 @@ # Licensed under the Apache-2.0 License. # ------------------------------------ """Unit tests for :mod:`pydo.agents.custom_sessions`.""" + from __future__ import annotations import json @@ -14,7 +15,6 @@ import pytest from pydo.agents import ( - AgentKind, AgentsResources, HITLOutcome, HarnessStreamError, @@ -24,7 +24,6 @@ resolve_agents_base_url, ) - # --------------------------------------------------------------------------- # Fake pipeline / response plumbing # --------------------------------------------------------------------------- @@ -85,28 +84,43 @@ def _make_resources(responses: List[_FakeResponse]) -> AgentsResources: # --------------------------------------------------------------------------- -def test_create_session_posts_expected_body(): +def test_create_from_manifest_uploads_yaml_verbatim(): body = {"session": {"session_id": "abc", "status": SessionStatus.PROVISIONING}} resources = _make_resources([_FakeResponse(200, body)]) - resp = resources.sessions.create( - agent_kind=AgentKind.CLAUDE_CODE, - repo_hint="digitalocean/pydo", - idle_timeout_seconds=900, + manifest = ( + "apiVersion: agents.digitalocean.com/v1alpha1\n" + "kind: Agent\n" + "metadata:\n" + " name: harness-demo\n" ) + resp = resources.sessions.create_from_manifest(manifest) call = resources._proxy._original._pipeline.calls[0] assert call.request.method == "POST" assert call.request.url.endswith("/v2/agents/sessions") - sent = json.loads(call.request.content) - assert sent == { - "agent_kind": "AGENT_KIND_CLAUDE_CODE", - "repo_hint": "digitalocean/pydo", - "idle_timeout_seconds": 900, - } + assert call.request.headers.get("Content-Type") == "application/x-yaml" + content = call.request.content + if isinstance(content, bytes): + content = content.decode("utf-8") + assert content == manifest assert resp.session.session_id == "abc" +def test_create_from_manifest_accepts_bytes(): + resources = _make_resources([_FakeResponse(200, {"session": {"session_id": "z"}})]) + resources.sessions.create_from_manifest(b"kind: Agent\n") + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.headers.get("Content-Type") == "application/x-yaml" + + +def test_create_from_manifest_rejects_empty(): + resources = _make_resources([]) + with pytest.raises(ValueError): + resources.sessions.create_from_manifest(" \n ") + + def test_get_session_url_encodes_id(): resources = _make_resources( [_FakeResponse(200, {"session": {"session_id": "x/y"}})] From 0ea762110c7761313bd10a8651b4f3d503efe725 Mon Sep 17 00:00:00 2001 From: SSharma-10 Date: Tue, 30 Jun 2026 19:29:40 +0530 Subject: [PATCH 4/8] add workspace file upload/download endpoints --- examples/agents/workspace_transfer.py | 116 ++++++ src/pydo/agents/__init__.py | 10 +- src/pydo/agents/custom_sessions.py | 362 ++++++++++++++++++- src/pydo/agents/session.py | 34 ++ src/pydo/aio/agents/__init__.py | 7 +- src/pydo/aio/agents/custom_sessions.py | 181 +++++++++- src/pydo/aio/agents/session.py | 32 ++ tests/agents/test_workspace.py | 476 +++++++++++++++++++++++++ 8 files changed, 1208 insertions(+), 10 deletions(-) create mode 100644 examples/agents/workspace_transfer.py create mode 100644 tests/agents/test_workspace.py diff --git a/examples/agents/workspace_transfer.py b/examples/agents/workspace_transfer.py new file mode 100644 index 00000000..39d47022 --- /dev/null +++ b/examples/agents/workspace_transfer.py @@ -0,0 +1,116 @@ +"""Upload a file to a session's sandbox workspace and download it back. + +Hosted-agents sessions expose a ``/workspace`` sandbox. These two custom REST +endpoints stream arbitrary bytes in and out of it: + + * ``client.agents.sessions.workspace_upload(...)`` POST .../workspace/upload + * ``client.agents.sessions.workspace_download(...)`` GET .../workspace/download + +This script does a full round trip against an *existing* session: + + 1. Upload a local file to ``GUEST_PATH`` inside the workspace. The optional + ``X-Content-Sha256`` header is sent so the guest can verify the upload. + 2. Download it back. The download is chunked and the SHA-256 digest arrives as + an HTTP *trailer* after the body, so the SDK reads the whole body first, + then verifies it. Note: CPython's HTTP stack discards chunked trailers, so + the X-Content-Sha256 trailer is usually unreadable from Python — the SDK + falls back to the server's size hint (X-Workspace-Size-Bytes) to detect + truncation and warns that the checksum could not be confirmed. A real + mismatch (trailer or size) raises ``WorkspaceTransferError``; pass + require_checksum=True for strict mode on a trailer-capable transport. + 3. Confirm the bytes survived the round trip. + +The high-level handle also offers ``agent.upload_file(...)`` / +``agent.download_file(...)`` which bind the session id for you; the equivalent +low-level calls are shown in comments. + +Required env: + DIGITALOCEAN_TOKEN + PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com + SESSION_ID an existing, READY session id + +Optional env: + LOCAL_FILE file to upload (default: a small generated text file) + GUEST_PATH destination inside /workspace (default: uploads/example.txt) + DOWNLOAD_TO local path to write the round-tripped copy (default: a temp file) +""" + +import hashlib +import os +import sys +import tempfile + +from pydo import Client +from pydo.agents import WorkspaceTransferError + + +def _sample_file() -> str: + """Create a small throwaway file to upload when LOCAL_FILE is unset.""" + fd, path = tempfile.mkstemp(prefix="pydo-ws-", suffix=".txt") + with os.fdopen(fd, "w", encoding="utf-8") as fh: + fh.write("hello from pydo workspace transfer\n" * 4) + return path + + +def main() -> int: + session_id = os.environ["SESSION_ID"] + local_file = os.environ.get("LOCAL_FILE") or _sample_file() + guest_path = os.environ.get("GUEST_PATH", "uploads/example.txt") + download_to = os.environ.get("DOWNLOAD_TO") or tempfile.mktemp(prefix="pydo-dl-") + + client = Client( + token=os.environ["DIGITALOCEAN_TOKEN"], + agents_endpoint=os.environ.get("PYDO_AGENTS_ENDPOINT"), + ) + agent = client.agents.attach(session_id) + + with open(local_file, "rb") as fh: + original = fh.read() + sha256 = hashlib.sha256(original).hexdigest() + print( + f"[session {session_id}] uploading {len(original)} bytes " + f"({local_file}) -> /workspace/{guest_path}", + file=sys.stderr, + ) + + # --- upload ----------------------------------------------------------- + # `data` accepts bytes, a filesystem path, or a readable binary stream. + # Passing content_sha256 lets the guest verify what it received. + up = agent.upload_file(path=guest_path, data=local_file, content_sha256=sha256) + # low-level equivalent: + # client.agents.sessions.workspace_upload( + # session_id, path=guest_path, data=local_file, content_sha256=sha256) + print(f"[uploaded] {dict(up)}", file=sys.stderr) + + # --- download --------------------------------------------------------- + # The returned object streams the body; .save()/.read() consume it fully + # and then verify integrity (trailer if readable, else the size hint). On a + # detected corruption/truncation a partial file written by .save() is + # removed and WorkspaceTransferError is raised, so you never keep a corrupt + # download. (A "could not verify trailer" warning is expected on CPython.) + download = agent.download_file(path=guest_path) + # low-level equivalent: + # download = client.agents.sessions.workspace_download( + # session_id, path=guest_path) + try: + written = download.save(download_to) + except WorkspaceTransferError as exc: + print(f"[integrity check FAILED] {exc}", file=sys.stderr) + return 1 + + print( + f"[downloaded] {written} bytes -> {download_to} " + f"(is_archive={download.is_archive}, size_hint={download.size_hint})", + file=sys.stderr, + ) + + # --- verify the round trip ------------------------------------------- + with open(download_to, "rb") as fh: + roundtripped = fh.read() + ok = roundtripped == original + print(f"[round-trip] bytes match: {ok}", file=sys.stderr) + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/pydo/agents/__init__.py b/src/pydo/agents/__init__.py index 60ada7c7..33b76b68 100644 --- a/src/pydo/agents/__init__.py +++ b/src/pydo/agents/__init__.py @@ -22,7 +22,13 @@ RunState, SessionStatus, ) -from .custom_sessions import HarnessEventStream, HarnessStreamError, SessionsOperations +from .custom_sessions import ( + HarnessEventStream, + HarnessStreamError, + SessionsOperations, + WorkspaceDownload, + WorkspaceTransferError, +) from .session import ( AgentEvent, AgentEventType, @@ -86,6 +92,8 @@ def attach(self, session_id: str) -> AgentSession: "SessionsOperations", "HarnessEventStream", "HarnessStreamError", + "WorkspaceDownload", + "WorkspaceTransferError", "DEFAULT_AGENTS_BASE_URL", "resolve_agents_base_url", "AgentKind", diff --git a/src/pydo/agents/custom_sessions.py b/src/pydo/agents/custom_sessions.py index fbe07752..7de32bb3 100644 --- a/src/pydo/agents/custom_sessions.py +++ b/src/pydo/agents/custom_sessions.py @@ -6,8 +6,11 @@ from __future__ import annotations +import hashlib import json as _json -from typing import Any, Dict, Iterator, List, Optional, Union +import os +import warnings +from typing import Any, BinaryIO, Dict, Iterator, List, Optional, Union from urllib.parse import quote from azure.core.exceptions import ( @@ -50,6 +53,175 @@ def _manifest_bytes(manifest: Union[str, bytes]) -> bytes: return data +# Workspace file-transfer contract (custom REST handlers, not grpc-gateway). +_OCTET_STREAM = "application/octet-stream" +_SHA256_HEADER = "X-Content-Sha256" +_IS_ARCHIVE_HEADER = "X-Workspace-Is-Archive" +_SIZE_HINT_HEADER = "X-Workspace-Size-Bytes" +# Per-request upload cap enforced by the server (413 beyond this); guarded +# client-side when the payload size is known to avoid a pointless round trip. +_MAX_UPLOAD_BYTES = 500 * 1024 * 1024 + +UploadData = Union[bytes, bytearray, str, "os.PathLike[str]", BinaryIO] + + +class WorkspaceTransferError(RuntimeError): + """A workspace upload/download failed an integrity check. + + Raised on a download when the ``X-Content-Sha256`` trailer does not match + the received bytes, when the byte count disagrees with the server's size + hint (truncation), or — in strict mode (``require_checksum=True``) — when + the trailer cannot be read at all. The output should be discarded. + """ + + +def _bool_param(value: bool) -> str: + return "true" if value else "false" + + +def _ci_get(mapping: Any, key: str) -> Optional[str]: + """Case-insensitive lookup over a header-like mapping.""" + if not mapping: + return None + getter = getattr(mapping, "get", None) + if getter is not None: + value = getter(key) + if value is not None: + return value + lower = key.lower() + try: + items = mapping.items() + except (AttributeError, TypeError): + return None + for name, value in items: + if isinstance(name, str) and name.lower() == lower: + return value + return None + + +def _extract_trailer(response: Any, name: str) -> Optional[str]: + """Best-effort read of an HTTP trailer that arrives after the body. + + The digest is sent as a chunked-transfer trailer, so it is only available + once the body has been fully consumed. Transports surface trailers + differently (and some not at all), so check the response headers and the + underlying transport's raw response. + """ + value = _ci_get(getattr(response, "headers", None), name) + if value: + return value + internal = getattr(response, "internal_response", None) + for obj in (internal, getattr(internal, "raw", None)): + if obj is None: + continue + value = _ci_get(getattr(obj, "trailers", None), name) + if value: + return value + return None + + +def _download_is_archive(response: Any) -> bool: + value = _ci_get(getattr(response, "headers", None), _IS_ARCHIVE_HEADER) + return str(value or "").strip().lower() == "true" + + +def _download_size_hint(response: Any) -> Optional[int]: + raw = _ci_get(getattr(response, "headers", None), _SIZE_HINT_HEADER) + try: + return int(raw) if raw not in (None, "") else None + except (TypeError, ValueError): + return None + + +def _verify_download( + response: Any, + computed_hex: str, + total_bytes: int, + require_checksum: bool, +) -> None: + """Verify a finished download against whatever integrity signal is readable. + + Per the contract the integrity digest is the ``X-Content-Sha256`` trailer. + When it is readable it is authoritative: a mismatch means corruption. But + CPython's HTTP stack (``http.client`` under both ``requests`` and + ``aiohttp``) reads and *discards* chunked trailers, so on the default + transports the trailer is usually unavailable even when the server sent it. + + The check therefore degrades gracefully: + + * trailer readable -> must match the computed digest, else raise; + * else size hint present -> received byte count must match, else raise + (catches truncation for known-size single files); + * else ``require_checksum`` -> raise (strict mode, for trailer-capable + transports); otherwise warn that the digest could not be verified. + """ + expected = _extract_trailer(response, _SHA256_HEADER) + if expected: + if expected.strip().lower() != computed_hex.lower(): + raise WorkspaceTransferError( + "workspace download integrity check failed: SHA-256 mismatch " + f"(trailer {expected.strip()!r} != computed {computed_hex!r}) — " + "discard the output." + ) + return + + size_hint = _download_size_hint(response) + if size_hint is not None and size_hint != total_bytes: + raise WorkspaceTransferError( + "workspace download is truncated: received " + f"{total_bytes} bytes but the server reported {size_hint} — " + "discard the output." + ) + if require_checksum: + raise WorkspaceTransferError( + "workspace download integrity check failed: the X-Content-Sha256 " + "trailer is missing or could not be read. Python's HTTP stack " + "discards chunked trailers, so strict checksum verification is not " + "possible on this transport; pass require_checksum=False to accept " + "downloads (verified by size when the server provides a size hint)." + ) + warnings.warn( + "workspace download could not verify the X-Content-Sha256 trailer " + "(Python's HTTP stack discards chunked trailers); integrity was " + + ( + "confirmed via the size hint." + if size_hint is not None + else "NOT independently verified." + ), + stacklevel=2, + ) + + +def _coerce_upload_content(data: UploadData) -> "tuple[Any, Optional[int], Any]": + """Normalize an upload payload to ``(content, size_or_None, handle_to_close)``. + + ``content`` is suitable to hand to the transport (raw bytes or a readable + binary stream). For filesystem paths a handle is opened and returned so the + caller can close it after the request. + """ + if isinstance(data, (bytes, bytearray)): + payload = bytes(data) + return payload, len(payload), None + if isinstance(data, (str, os.PathLike)): + path = os.fspath(data) + size = os.path.getsize(path) + handle = open(path, "rb") # pylint: disable=consider-using-with + return handle, size, handle + if hasattr(data, "read"): + size: Optional[int] = None + try: + current = data.tell() + data.seek(0, os.SEEK_END) + size = data.tell() - current + data.seek(current) + except (OSError, AttributeError, ValueError): + size = None + return data, size, None + raise TypeError( + "data must be bytes, a filesystem path, or a readable binary stream" + ) + + def _unwrap_harness_sse_chunk(chunk: Dict[str, Any]) -> Optional[Any]: """Normalize SSE JSON to a harness Event. @@ -158,12 +330,13 @@ def _send( path: str, *, body: Optional[Dict[str, Any]] = None, - content: Optional[Union[str, bytes]] = None, + content: Optional[Any] = None, content_type: Optional[str] = None, params: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, str]] = None, stream: bool = False, ): - headers = {"Accept": "application/json"} + headers = {"Accept": "application/json", **(headers or {})} kwargs: Dict[str, Any] = {"headers": headers} if params: kwargs["params"] = { @@ -316,5 +489,186 @@ def stream( _raise_agents_http_error(response) return HarnessEventStream(SSEStream(response)) + def workspace_upload( + self, + session_id: str, + *, + path: str, + data: UploadData, + is_archive: bool = False, + content_sha256: Optional[str] = None, + ) -> Any: + """Upload raw file (or tar) bytes into a session's sandbox workspace. + + ``POST /v2/agents/sessions/{session_id}/workspace/upload``. + + :param path: Destination path, resolved inside the workspace root + (``/workspace``). Anything escaping the root is rejected with 403. + :param data: Raw bytes, a filesystem path, or a readable binary stream. + :param is_archive: When ``True`` the body is a tar archive to extract at + ``path``. + :param content_sha256: Optional hex digest of the full payload, forwarded + via the ``X-Content-Sha256`` header for the guest to verify. + :returns: ``{"path": ..., "bytes_written": N}``. + """ + if not path: + raise ValueError("path is required") + content, size, handle = _coerce_upload_content(data) + try: + if size is not None and size > _MAX_UPLOAD_BYTES: + raise ValueError( + f"upload of {size} bytes exceeds the 500 MiB per-request limit" + ) + headers = {_SHA256_HEADER: content_sha256} if content_sha256 else None + return self._parse_json( + self._send( + "POST", + f"{_BASE_PATH}/{_quote(session_id)}/workspace/upload", + content=content, + content_type=_OCTET_STREAM, + params={"path": path, "is_archive": _bool_param(is_archive)}, + headers=headers, + ), + ) + finally: + if handle is not None: + handle.close() + + def workspace_download( + self, + session_id: str, + *, + path: str, + as_archive: bool = False, + require_checksum: bool = False, + ) -> "WorkspaceDownload": + """Download a file (or tar-streamed directory) from a session workspace. + + ``GET /v2/agents/sessions/{session_id}/workspace/download``. + + The response is chunked with the SHA-256 digest delivered as an HTTP + trailer after the body. The returned :class:`WorkspaceDownload` streams + the body and verifies integrity once it is fully consumed. A mismatched + trailer (or a byte count that disagrees with the server's size hint) + raises :class:`WorkspaceTransferError` and the output should be + discarded. + + :param path: Source path, resolved inside the workspace root. + :param as_archive: When ``True`` the directory at ``path`` is + tar-streamed. + :param require_checksum: When ``True``, raise if the SHA-256 trailer + cannot be read. The default is ``False`` because CPython's HTTP + stack discards chunked trailers, so strict verification only works + on a trailer-capable transport. + """ + if not path: + raise ValueError("path is required") + request = HttpRequest( + "GET", + f"{_BASE_PATH}/{_quote(session_id)}/workspace/download", + headers={"Accept": _OCTET_STREAM}, + params={"path": path, "as_archive": _bool_param(as_archive)}, + ) + request.url = self._client.format_url(request.url) + pipeline_response = self._client._pipeline.run(request, stream=True) + response = pipeline_response.http_response + if response.status_code != 200: + response.read() + _raise_agents_http_error(response) + return WorkspaceDownload(response, require_checksum=require_checksum) + + +class WorkspaceDownload: + """A streaming workspace download with best-effort integrity verification. + + Iterating yields the raw body chunks while the SHA-256 digest is computed + incrementally. Once the body is fully consumed, integrity is verified (see + :func:`_verify_download`): a mismatched ``X-Content-Sha256`` trailer or a + byte count disagreeing with the size hint raises + :class:`WorkspaceTransferError`. Always consume the body to completion + (via iteration, :meth:`read`, or :meth:`save`) before trusting the output. + """ + + def __init__(self, response: Any, *, require_checksum: bool = False): + self._response = response + self._require_checksum = require_checksum + self.bytes_read = 0 + + @property + def is_archive(self) -> bool: + """Whether the response is a tar stream (``X-Workspace-Is-Archive``).""" + return _download_is_archive(self._response) + + @property + def size_hint(self) -> Optional[int]: + """Size hint for progress UIs (``X-Workspace-Size-Bytes``); not framing.""" + return _download_size_hint(self._response) + + def __iter__(self) -> Iterator[bytes]: + hasher = hashlib.sha256() + total = 0 + for chunk in self._response.iter_bytes(): + if not chunk: + continue + if isinstance(chunk, str): + chunk = chunk.encode("utf-8") + hasher.update(chunk) + total += len(chunk) + yield bytes(chunk) + self.bytes_read = total + _verify_download( + self._response, hasher.hexdigest(), total, self._require_checksum + ) + + def read(self) -> bytes: + """Consume the whole body and return the (verified) bytes.""" + return b"".join(self) + + def save(self, dest: Union[str, "os.PathLike[str]", BinaryIO]) -> int: + """Stream the body to *dest* (a path or writable binary file). + + Returns the number of bytes written. If verification fails, a + file opened from a path is removed before re-raising. + """ + own = isinstance(dest, (str, os.PathLike)) + handle = open(os.fspath(dest), "wb") if own else dest # type: ignore[arg-type] + total = 0 + try: + for chunk in self: + handle.write(chunk) + total += len(chunk) + except WorkspaceTransferError: + if own: + handle.close() + try: + os.remove(os.fspath(dest)) # type: ignore[arg-type] + except OSError: + pass + raise + finally: + if own and not handle.closed: + handle.close() + return total + + def close(self) -> None: + closer = getattr(self._response, "close", None) + if closer is not None: + try: + closer() + except Exception: # noqa: BLE001 — best-effort cleanup + pass + + def __enter__(self) -> "WorkspaceDownload": + return self + + def __exit__(self, *args: Any) -> None: + self.close() + -__all__ = ["SessionsOperations", "HarnessEventStream", "HarnessStreamError"] +__all__ = [ + "SessionsOperations", + "HarnessEventStream", + "HarnessStreamError", + "WorkspaceDownload", + "WorkspaceTransferError", +] diff --git a/src/pydo/agents/session.py b/src/pydo/agents/session.py index 711212d6..d1eb7be8 100644 --- a/src/pydo/agents/session.py +++ b/src/pydo/agents/session.py @@ -388,6 +388,40 @@ def resolve_hitl( source=source, ) + def upload_file( + self, + *, + path: str, + data: Any, + is_archive: bool = False, + content_sha256: Optional[str] = None, + ) -> Any: + """Upload bytes/a tar into the session workspace (see + :meth:`SessionsOperations.workspace_upload`).""" + return self._sessions.workspace_upload( + self.session_id, + path=path, + data=data, + is_archive=is_archive, + content_sha256=content_sha256, + ) + + def download_file( + self, + *, + path: str, + as_archive: bool = False, + require_checksum: bool = False, + ) -> Any: + """Download a workspace file/tar with integrity verification (see + :meth:`SessionsOperations.workspace_download`).""" + return self._sessions.workspace_download( + self.session_id, + path=path, + as_archive=as_archive, + require_checksum=require_checksum, + ) + def destroy(self) -> None: self._sessions.destroy(self.session_id) diff --git a/src/pydo/aio/agents/__init__.py b/src/pydo/aio/agents/__init__.py index 8c29af31..ea08da9c 100644 --- a/src/pydo/aio/agents/__init__.py +++ b/src/pydo/aio/agents/__init__.py @@ -10,7 +10,11 @@ from pydo.agents import resolve_agents_base_url from pydo.custom_extensions import _BaseURLProxy -from .custom_sessions import AsyncHarnessEventStream, AsyncSessionsOperations +from .custom_sessions import ( + AsyncHarnessEventStream, + AsyncSessionsOperations, + AsyncWorkspaceDownload, +) from .session import AsyncAgentSession, AsyncRunStream @@ -49,4 +53,5 @@ def attach(self, session_id: str) -> AsyncAgentSession: "AsyncRunStream", "AsyncSessionsOperations", "AsyncHarnessEventStream", + "AsyncWorkspaceDownload", ] diff --git a/src/pydo/aio/agents/custom_sessions.py b/src/pydo/aio/agents/custom_sessions.py index a0764cf7..5e7fe4fc 100644 --- a/src/pydo/aio/agents/custom_sessions.py +++ b/src/pydo/aio/agents/custom_sessions.py @@ -6,8 +6,10 @@ from __future__ import annotations +import hashlib import json as _json -from typing import Any, AsyncIterator, Dict, List, Optional, Union +import os +from typing import Any, AsyncIterator, BinaryIO, Dict, List, Optional, Union from urllib.parse import quote from azure.core.exceptions import ( @@ -21,11 +23,21 @@ from azure.core.rest import HttpRequest from pydo.agents.custom_sessions import ( + _MAX_UPLOAD_BYTES, + _OCTET_STREAM, + _SHA256_HEADER, _YAML_MEDIA_TYPE, HarnessStreamError, + UploadData, + WorkspaceTransferError, + _bool_param, + _coerce_upload_content, + _download_is_archive, + _download_size_hint, _manifest_bytes, _raise_agents_http_error, _unwrap_harness_sse_chunk, + _verify_download, ) from pydo.custom_extensions import AsyncSSEStream, _wrap @@ -87,12 +99,13 @@ async def _send( path: str, *, body: Optional[Dict[str, Any]] = None, - content: Optional[Union[str, bytes]] = None, + content: Optional[Any] = None, content_type: Optional[str] = None, params: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, str]] = None, stream: bool = False, ): - headers = {"Accept": "application/json"} + headers = {"Accept": "application/json", **(headers or {})} kwargs: Dict[str, Any] = {"headers": headers} if params: kwargs["params"] = { @@ -246,5 +259,165 @@ async def stream( _raise_agents_http_error(response) return AsyncHarnessEventStream(AsyncSSEStream(response)) + async def workspace_upload( + self, + session_id: str, + *, + path: str, + data: UploadData, + is_archive: bool = False, + content_sha256: Optional[str] = None, + ) -> Any: + """Upload raw file (or tar) bytes into a session's sandbox workspace. + + Async counterpart of + :meth:`pydo.agents.custom_sessions.SessionsOperations.workspace_upload`. + """ + if not path: + raise ValueError("path is required") + content, size, handle = _coerce_upload_content(data) + try: + # aiohttp does not reliably stream arbitrary sync file objects, so + # materialize non-bytes payloads before sending. + if hasattr(content, "read"): + content = content.read() + if isinstance(content, str): + content = content.encode("utf-8") + size = len(content) + finally: + if handle is not None: + handle.close() + if size is not None and size > _MAX_UPLOAD_BYTES: + raise ValueError( + f"upload of {size} bytes exceeds the 500 MiB per-request limit" + ) + headers = {_SHA256_HEADER: content_sha256} if content_sha256 else None + return await self._parse_json( + await self._send( + "POST", + f"{_BASE_PATH}/{_quote(session_id)}/workspace/upload", + content=content, + content_type=_OCTET_STREAM, + params={"path": path, "is_archive": _bool_param(is_archive)}, + headers=headers, + ), + ) + + async def workspace_download( + self, + session_id: str, + *, + path: str, + as_archive: bool = False, + require_checksum: bool = False, + ) -> "AsyncWorkspaceDownload": + """Download a file (or tar-streamed directory) from a session workspace. + + Async counterpart of + :meth:`pydo.agents.custom_sessions.SessionsOperations.workspace_download`. + """ + if not path: + raise ValueError("path is required") + request = HttpRequest( + "GET", + f"{_BASE_PATH}/{_quote(session_id)}/workspace/download", + headers={"Accept": _OCTET_STREAM}, + params={"path": path, "as_archive": _bool_param(as_archive)}, + ) + request.url = self._client.format_url(request.url) + pipeline_response = await self._client._pipeline.run(request, stream=True) + response = pipeline_response.http_response + if response.status_code != 200: + await response.read() + _raise_agents_http_error(response) + return AsyncWorkspaceDownload(response, require_checksum=require_checksum) + + +class AsyncWorkspaceDownload: + """Async streaming workspace download with best-effort integrity verification. + + Async counterpart of :class:`pydo.agents.custom_sessions.WorkspaceDownload`; + see that class and :func:`pydo.agents.custom_sessions._verify_download` for + the verification semantics. + """ + + def __init__(self, response: Any, *, require_checksum: bool = False): + self._response = response + self._require_checksum = require_checksum + self.bytes_read = 0 + + @property + def is_archive(self) -> bool: + return _download_is_archive(self._response) + + @property + def size_hint(self) -> Optional[int]: + return _download_size_hint(self._response) + + def __aiter__(self) -> AsyncIterator[bytes]: + return self._iter() + + async def _iter(self) -> AsyncIterator[bytes]: + hasher = hashlib.sha256() + total = 0 + async for chunk in self._response.iter_bytes(): + if not chunk: + continue + if isinstance(chunk, str): + chunk = chunk.encode("utf-8") + hasher.update(chunk) + total += len(chunk) + yield bytes(chunk) + self.bytes_read = total + _verify_download( + self._response, hasher.hexdigest(), total, self._require_checksum + ) + + async def read(self) -> bytes: + chunks = [chunk async for chunk in self] + return b"".join(chunks) + + async def save(self, dest: Union[str, "os.PathLike[str]", BinaryIO]) -> int: + own = isinstance(dest, (str, os.PathLike)) + handle = open(os.fspath(dest), "wb") if own else dest # type: ignore[arg-type] + total = 0 + try: + async for chunk in self: + handle.write(chunk) + total += len(chunk) + except WorkspaceTransferError: + if own: + handle.close() + try: + os.remove(os.fspath(dest)) # type: ignore[arg-type] + except OSError: + pass + raise + finally: + if own and not handle.closed: + handle.close() + return total + + async def close(self) -> None: + closer = getattr(self._response, "close", None) + if closer is None: + return + try: + result = closer() + if hasattr(result, "__await__"): + await result + except Exception: # noqa: BLE001 — best-effort cleanup + pass + + async def __aenter__(self) -> "AsyncWorkspaceDownload": + return self + + async def __aexit__(self, *args: Any) -> None: + await self.close() + -__all__ = ["AsyncSessionsOperations", "AsyncHarnessEventStream"] +__all__ = [ + "AsyncSessionsOperations", + "AsyncHarnessEventStream", + "AsyncWorkspaceDownload", +] diff --git a/src/pydo/aio/agents/session.py b/src/pydo/aio/agents/session.py index 0b3e8139..06da8d8a 100644 --- a/src/pydo/aio/agents/session.py +++ b/src/pydo/aio/agents/session.py @@ -226,6 +226,38 @@ async def resolve_hitl( source=source, ) + async def upload_file( + self, + *, + path: str, + data: Any, + is_archive: bool = False, + content_sha256: Optional[str] = None, + ) -> Any: + """Upload bytes/a tar into the session workspace.""" + return await self._sessions.workspace_upload( + self.session_id, + path=path, + data=data, + is_archive=is_archive, + content_sha256=content_sha256, + ) + + async def download_file( + self, + *, + path: str, + as_archive: bool = False, + require_checksum: bool = False, + ) -> Any: + """Download a workspace file/tar with integrity verification.""" + return await self._sessions.workspace_download( + self.session_id, + path=path, + as_archive=as_archive, + require_checksum=require_checksum, + ) + async def destroy(self) -> None: await self._sessions.destroy(self.session_id) diff --git a/tests/agents/test_workspace.py b/tests/agents/test_workspace.py new file mode 100644 index 00000000..0c1832e7 --- /dev/null +++ b/tests/agents/test_workspace.py @@ -0,0 +1,476 @@ +# pylint: disable=line-too-long,missing-class-docstring,missing-function-docstring,protected-access +# ------------------------------------ +# Copyright (c) DigitalOcean. +# Licensed under the Apache-2.0 License. +# ------------------------------------ +"""Unit tests for workspace upload/download (sync + async).""" + +from __future__ import annotations + +import hashlib +import io +import json +from types import SimpleNamespace +from typing import Any, List, Optional +from unittest.mock import MagicMock + +import pytest + +from pydo.agents import AgentsResources, WorkspaceTransferError +from pydo.aio.agents import AsyncAgentsResources + +# --------------------------------------------------------------------------- +# Sync fakes +# --------------------------------------------------------------------------- + + +class _FakeResponse: + def __init__( + self, + status_code: int, + *, + body: Any = None, + chunks: Optional[List[bytes]] = None, + headers: Optional[dict] = None, + trailer: Optional[str] = None, + ): + self.status_code = status_code + self.headers = dict(headers or {}) + self.reason = None + self._chunks = list(chunks or []) + self._trailer = trailer + self.closed = False + if isinstance(body, (dict, list)): + self._body_bytes = json.dumps(body).encode("utf-8") + elif isinstance(body, str): + self._body_bytes = body.encode("utf-8") + elif isinstance(body, bytes): + self._body_bytes = body + else: + self._body_bytes = b"".join(self._chunks) + + def text(self) -> str: + return self._body_bytes.decode("utf-8") + + def body(self) -> bytes: + return self._body_bytes + + def read(self) -> bytes: + return self._body_bytes + + def iter_bytes(self): + for chunk in self._chunks: + yield chunk + # The integrity digest is a trailer: only visible once the body is done. + if self._trailer is not None: + self.headers["X-Content-Sha256"] = self._trailer + + def close(self) -> None: + self.closed = True + + +class _FakePipeline: + def __init__(self, responses: List[_FakeResponse]): + self._responses = list(responses) + self.calls: List[Any] = [] + + def run(self, request, *, stream=False): + # The real transport reads a streamed request body here, before the + # caller closes any opened file handle; mirror that so tests can inspect + # the bytes that were actually sent. + content_bytes = request.content + if hasattr(content_bytes, "read"): + content_bytes = content_bytes.read() + self.calls.append( + SimpleNamespace(request=request, stream=stream, content_bytes=content_bytes) + ) + return SimpleNamespace(http_response=self._responses.pop(0)) + + +def _make_resources(responses: List[_FakeResponse]) -> AgentsResources: + parent = MagicMock() + parent._client = MagicMock() + parent._client._pipeline = _FakePipeline(responses) + return AgentsResources(parent, agents_endpoint="https://api.stage2.digitalocean.com") + + +def _calls(resources) -> List[Any]: + return resources._proxy._original._pipeline.calls + + +# --------------------------------------------------------------------------- +# Upload +# --------------------------------------------------------------------------- + + +def test_upload_bytes_sets_path_params_and_content_type(): + resources = _make_resources( + [_FakeResponse(200, body={"path": "/workspace/a.txt", "bytes_written": 5})] + ) + + resp = resources.sessions.workspace_upload( + "s1", path="a.txt", data=b"hello" + ) + + call = _calls(resources)[0] + assert call.request.method == "POST" + assert "/v2/agents/sessions/s1/workspace/upload" in call.request.url + assert "path=a.txt" in call.request.url + assert "is_archive=false" in call.request.url + assert call.request.headers.get("Content-Type") == "application/octet-stream" + assert "X-Content-Sha256" not in call.request.headers + assert call.request.content == b"hello" + assert resp.bytes_written == 5 + + +def test_upload_forwards_sha256_header_and_is_archive(): + resources = _make_resources([_FakeResponse(200, body={"bytes_written": 3})]) + + resources.sessions.workspace_upload( + "s1", + path="dir", + data=b"tar", + is_archive=True, + content_sha256="deadbeef", + ) + + call = _calls(resources)[0] + assert "is_archive=true" in call.request.url + assert call.request.headers.get("X-Content-Sha256") == "deadbeef" + + +def test_upload_accepts_filesystem_path(tmp_path): + payload = b"file-on-disk" + src = tmp_path / "input.bin" + src.write_bytes(payload) + resources = _make_resources([_FakeResponse(200, body={"bytes_written": len(payload)})]) + + resources.sessions.workspace_upload("s1", path="dest.bin", data=str(src)) + + assert _calls(resources)[0].content_bytes == payload + + +def test_upload_rejects_payload_over_500_mib(): + class _HugeStream: + def __init__(self): + self._pos = 0 + + def tell(self): + return self._pos + + def seek(self, offset, whence=io.SEEK_SET): + self._pos = (500 * 1024 * 1024 + 1) if whence == io.SEEK_END else offset + return self._pos + + def read(self, *_a, **_k): + return b"" + + resources = _make_resources([]) + with pytest.raises(ValueError, match="500 MiB"): + resources.sessions.workspace_upload("s1", path="big", data=_HugeStream()) + + +def test_upload_requires_path(): + resources = _make_resources([]) + with pytest.raises(ValueError, match="path is required"): + resources.sessions.workspace_upload("s1", path="", data=b"x") + + +# --------------------------------------------------------------------------- +# Download +# --------------------------------------------------------------------------- + + +def test_download_verifies_matching_trailer_and_returns_bytes(): + payload = b"hello workspace" + digest = hashlib.sha256(payload).hexdigest() + resources = _make_resources( + [ + _FakeResponse( + 200, + chunks=[b"hello ", b"workspace"], + headers={"X-Workspace-Size-Bytes": str(len(payload))}, + trailer=digest, + ) + ] + ) + + download = resources.sessions.workspace_download("s1", path="out.txt") + assert download.size_hint == len(payload) + assert download.is_archive is False + + data = download.read() + assert data == payload + assert download.bytes_read == len(payload) + + call = _calls(resources)[0] + assert call.request.method == "GET" + assert "/v2/agents/sessions/s1/workspace/download" in call.request.url + assert "path=out.txt" in call.request.url + assert "as_archive=false" in call.request.url + + +def test_download_archive_flag_and_header(): + payload = b"tarbytes" + resources = _make_resources( + [ + _FakeResponse( + 200, + chunks=[payload], + headers={"X-Workspace-Is-Archive": "true"}, + trailer=hashlib.sha256(payload).hexdigest(), + ) + ] + ) + + download = resources.sessions.workspace_download( + "s1", path="dir", as_archive=True + ) + assert download.read() == payload + assert download.is_archive is True + assert "as_archive=true" in _calls(resources)[0].request.url + + +def test_download_missing_trailer_strict_mode_is_failure(): + payload = b"truncated" + resources = _make_resources([_FakeResponse(200, chunks=[payload], trailer=None)]) + + download = resources.sessions.workspace_download( + "s1", path="x", require_checksum=True + ) + with pytest.raises(WorkspaceTransferError, match="trailer"): + download.read() + + +def test_download_missing_trailer_default_warns_and_returns(): + # Python's HTTP stack discards chunked trailers, so the default tolerates a + # missing trailer (with a warning) rather than failing every real download. + payload = b"no-trailer" + resources = _make_resources([_FakeResponse(200, chunks=[payload], trailer=None)]) + + download = resources.sessions.workspace_download("s1", path="x") + with pytest.warns(UserWarning, match="X-Content-Sha256"): + data = download.read() + assert data == payload + + +def test_download_size_hint_mismatch_is_truncation_failure(): + resources = _make_resources( + [ + _FakeResponse( + 200, + chunks=[b"abc"], + headers={"X-Workspace-Size-Bytes": "99"}, + trailer=None, + ) + ] + ) + download = resources.sessions.workspace_download("s1", path="x") + with pytest.raises(WorkspaceTransferError, match="truncated"): + download.read() + + +def test_download_size_hint_match_is_accepted(): + payload = b"sized" + resources = _make_resources( + [ + _FakeResponse( + 200, + chunks=[payload], + headers={"X-Workspace-Size-Bytes": str(len(payload))}, + trailer=None, + ) + ] + ) + download = resources.sessions.workspace_download("s1", path="x") + assert download.read() == payload + + +def test_download_mismatched_trailer_is_failure(): + resources = _make_resources( + [_FakeResponse(200, chunks=[b"abc"], trailer="0" * 64)] + ) + + download = resources.sessions.workspace_download("s1", path="x") + with pytest.raises(WorkspaceTransferError, match="mismatch"): + download.read() + + +def test_download_require_checksum_false_skips_verification(): + resources = _make_resources([_FakeResponse(200, chunks=[b"abc"], trailer=None)]) + + download = resources.sessions.workspace_download( + "s1", path="x", require_checksum=False + ) + assert download.read() == b"abc" + + +def test_download_save_writes_file_and_discards_on_failure(tmp_path): + good = tmp_path / "good.bin" + payload = b"good-payload" + resources = _make_resources( + [ + _FakeResponse( + 200, chunks=[payload], trailer=hashlib.sha256(payload).hexdigest() + ) + ] + ) + written = resources.sessions.workspace_download("s1", path="g").save(str(good)) + assert written == len(payload) + assert good.read_bytes() == payload + + bad = tmp_path / "bad.bin" + resources = _make_resources([_FakeResponse(200, chunks=[b"abc"], trailer="0" * 64)]) + with pytest.raises(WorkspaceTransferError): + resources.sessions.workspace_download("s1", path="b").save(str(bad)) + assert not bad.exists() + + +def test_download_non_200_raises(): + from azure.core.exceptions import HttpResponseError + + resources = _make_resources( + [_FakeResponse(404, body="path not found")] + ) + with pytest.raises(HttpResponseError): + resources.sessions.workspace_download("s1", path="missing") + + +def test_agent_session_upload_download_passthrough(tmp_path): + payload = b"round-trip" + resources = _make_resources( + [ + _FakeResponse(200, body={"bytes_written": len(payload)}), + _FakeResponse( + 200, chunks=[payload], trailer=hashlib.sha256(payload).hexdigest() + ), + ] + ) + agent = resources.attach("s1") + + up = agent.upload_file(path="f.bin", data=payload) + assert up.bytes_written == len(payload) + assert agent.download_file(path="f.bin").read() == payload + + +# --------------------------------------------------------------------------- +# Async fakes + tests +# --------------------------------------------------------------------------- + + +class _FakeAsyncResponse: + def __init__( + self, + status_code: int, + *, + body: Any = None, + chunks: Optional[List[bytes]] = None, + headers: Optional[dict] = None, + trailer: Optional[str] = None, + ): + self.status_code = status_code + self.headers = dict(headers or {}) + self.reason = None + self._chunks = list(chunks or []) + self._trailer = trailer + if isinstance(body, (dict, list)): + self._body_bytes = json.dumps(body).encode("utf-8") + elif isinstance(body, str): + self._body_bytes = body.encode("utf-8") + elif isinstance(body, bytes): + self._body_bytes = body + else: + self._body_bytes = b"".join(self._chunks) + + async def read(self) -> bytes: + return self._body_bytes + + def text(self) -> str: + return self._body_bytes.decode("utf-8") + + def body(self) -> bytes: + return self._body_bytes + + async def iter_bytes(self): + for chunk in self._chunks: + yield chunk + if self._trailer is not None: + self.headers["X-Content-Sha256"] = self._trailer + + def close(self) -> None: + pass + + +class _FakeAsyncPipeline: + def __init__(self, responses: List[_FakeAsyncResponse]): + self._responses = list(responses) + self.calls: List[Any] = [] + + async def run(self, request, *, stream=False): + self.calls.append(SimpleNamespace(request=request, stream=stream)) + return SimpleNamespace(http_response=self._responses.pop(0)) + + +def _make_async_resources(responses: List[_FakeAsyncResponse]) -> AsyncAgentsResources: + parent = MagicMock() + parent._client = MagicMock() + parent._client._pipeline = _FakeAsyncPipeline(responses) + return AsyncAgentsResources( + parent, agents_endpoint="https://api.stage2.digitalocean.com" + ) + + +@pytest.mark.asyncio +async def test_async_upload_materializes_and_sets_headers(): + resources = _make_async_resources([_FakeAsyncResponse(200, body={"bytes_written": 3})]) + + resp = await resources.sessions.workspace_upload( + "s1", path="a.txt", data=b"abc", content_sha256="cafe" + ) + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.method == "POST" + assert "/v2/agents/sessions/s1/workspace/upload" in call.request.url + assert call.request.headers.get("Content-Type") == "application/octet-stream" + assert call.request.headers.get("X-Content-Sha256") == "cafe" + assert resp.bytes_written == 3 + + +@pytest.mark.asyncio +async def test_async_download_verifies_trailer(): + payload = b"async-bytes" + resources = _make_async_resources( + [ + _FakeAsyncResponse( + 200, chunks=[b"async-", b"bytes"], trailer=hashlib.sha256(payload).hexdigest() + ) + ] + ) + + download = await resources.sessions.workspace_download("s1", path="o") + data = await download.read() + assert data == payload + assert download.bytes_read == len(payload) + + +@pytest.mark.asyncio +async def test_async_download_missing_trailer_strict_fails(): + resources = _make_async_resources([_FakeAsyncResponse(200, chunks=[b"x"], trailer=None)]) + + download = await resources.sessions.workspace_download( + "s1", path="o", require_checksum=True + ) + with pytest.raises(WorkspaceTransferError): + await download.read() + + +@pytest.mark.asyncio +async def test_async_download_save_discards_on_failure(tmp_path): + bad = tmp_path / "bad.bin" + resources = _make_async_resources( + [_FakeAsyncResponse(200, chunks=[b"abc"], trailer="0" * 64)] + ) + download = await resources.sessions.workspace_download("s1", path="b") + with pytest.raises(WorkspaceTransferError): + await download.save(str(bad)) + assert not bad.exists() From 42c74c3f4b7dbce92bd87f01202b87dc145a7f95 Mon Sep 17 00:00:00 2001 From: SSharma-10 Date: Tue, 30 Jun 2026 19:46:39 +0530 Subject: [PATCH 5/8] lint fix --- examples/agents/workspace_transfer.py | 41 ------------ src/pydo/agents/custom_sessions.py | 88 ++++++-------------------- src/pydo/aio/agents/custom_sessions.py | 3 +- tests/agents/test_workspace.py | 36 ++++++----- 4 files changed, 40 insertions(+), 128 deletions(-) diff --git a/examples/agents/workspace_transfer.py b/examples/agents/workspace_transfer.py index 39d47022..42e45d92 100644 --- a/examples/agents/workspace_transfer.py +++ b/examples/agents/workspace_transfer.py @@ -1,29 +1,5 @@ """Upload a file to a session's sandbox workspace and download it back. -Hosted-agents sessions expose a ``/workspace`` sandbox. These two custom REST -endpoints stream arbitrary bytes in and out of it: - - * ``client.agents.sessions.workspace_upload(...)`` POST .../workspace/upload - * ``client.agents.sessions.workspace_download(...)`` GET .../workspace/download - -This script does a full round trip against an *existing* session: - - 1. Upload a local file to ``GUEST_PATH`` inside the workspace. The optional - ``X-Content-Sha256`` header is sent so the guest can verify the upload. - 2. Download it back. The download is chunked and the SHA-256 digest arrives as - an HTTP *trailer* after the body, so the SDK reads the whole body first, - then verifies it. Note: CPython's HTTP stack discards chunked trailers, so - the X-Content-Sha256 trailer is usually unreadable from Python — the SDK - falls back to the server's size hint (X-Workspace-Size-Bytes) to detect - truncation and warns that the checksum could not be confirmed. A real - mismatch (trailer or size) raises ``WorkspaceTransferError``; pass - require_checksum=True for strict mode on a trailer-capable transport. - 3. Confirm the bytes survived the round trip. - -The high-level handle also offers ``agent.upload_file(...)`` / -``agent.download_file(...)`` which bind the session id for you; the equivalent -low-level calls are shown in comments. - Required env: DIGITALOCEAN_TOKEN PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com @@ -45,7 +21,6 @@ def _sample_file() -> str: - """Create a small throwaway file to upload when LOCAL_FILE is unset.""" fd, path = tempfile.mkstemp(prefix="pydo-ws-", suffix=".txt") with os.fdopen(fd, "w", encoding="utf-8") as fh: fh.write("hello from pydo workspace transfer\n" * 4) @@ -73,25 +48,10 @@ def main() -> int: file=sys.stderr, ) - # --- upload ----------------------------------------------------------- - # `data` accepts bytes, a filesystem path, or a readable binary stream. - # Passing content_sha256 lets the guest verify what it received. up = agent.upload_file(path=guest_path, data=local_file, content_sha256=sha256) - # low-level equivalent: - # client.agents.sessions.workspace_upload( - # session_id, path=guest_path, data=local_file, content_sha256=sha256) print(f"[uploaded] {dict(up)}", file=sys.stderr) - # --- download --------------------------------------------------------- - # The returned object streams the body; .save()/.read() consume it fully - # and then verify integrity (trailer if readable, else the size hint). On a - # detected corruption/truncation a partial file written by .save() is - # removed and WorkspaceTransferError is raised, so you never keep a corrupt - # download. (A "could not verify trailer" warning is expected on CPython.) download = agent.download_file(path=guest_path) - # low-level equivalent: - # download = client.agents.sessions.workspace_download( - # session_id, path=guest_path) try: written = download.save(download_to) except WorkspaceTransferError as exc: @@ -104,7 +64,6 @@ def main() -> int: file=sys.stderr, ) - # --- verify the round trip ------------------------------------------- with open(download_to, "rb") as fh: roundtripped = fh.read() ok = roundtripped == original diff --git a/src/pydo/agents/custom_sessions.py b/src/pydo/agents/custom_sessions.py index 7de32bb3..d5f5b167 100644 --- a/src/pydo/agents/custom_sessions.py +++ b/src/pydo/agents/custom_sessions.py @@ -58,21 +58,13 @@ def _manifest_bytes(manifest: Union[str, bytes]) -> bytes: _SHA256_HEADER = "X-Content-Sha256" _IS_ARCHIVE_HEADER = "X-Workspace-Is-Archive" _SIZE_HINT_HEADER = "X-Workspace-Size-Bytes" -# Per-request upload cap enforced by the server (413 beyond this); guarded -# client-side when the payload size is known to avoid a pointless round trip. -_MAX_UPLOAD_BYTES = 500 * 1024 * 1024 +_MAX_UPLOAD_BYTES = 500 * 1024 * 1024 # server returns 413 beyond this UploadData = Union[bytes, bytearray, str, "os.PathLike[str]", BinaryIO] class WorkspaceTransferError(RuntimeError): - """A workspace upload/download failed an integrity check. - - Raised on a download when the ``X-Content-Sha256`` trailer does not match - the received bytes, when the byte count disagrees with the server's size - hint (truncation), or — in strict mode (``require_checksum=True``) — when - the trailer cannot be read at all. The output should be discarded. - """ + """A workspace download failed its integrity check (discard the output).""" def _bool_param(value: bool) -> str: @@ -100,13 +92,7 @@ def _ci_get(mapping: Any, key: str) -> Optional[str]: def _extract_trailer(response: Any, name: str) -> Optional[str]: - """Best-effort read of an HTTP trailer that arrives after the body. - - The digest is sent as a chunked-transfer trailer, so it is only available - once the body has been fully consumed. Transports surface trailers - differently (and some not at all), so check the response headers and the - underlying transport's raw response. - """ + """Best-effort read of a chunked-transfer trailer (only valid post-body).""" value = _ci_get(getattr(response, "headers", None), name) if value: return value @@ -141,19 +127,10 @@ def _verify_download( ) -> None: """Verify a finished download against whatever integrity signal is readable. - Per the contract the integrity digest is the ``X-Content-Sha256`` trailer. - When it is readable it is authoritative: a mismatch means corruption. But - CPython's HTTP stack (``http.client`` under both ``requests`` and - ``aiohttp``) reads and *discards* chunked trailers, so on the default - transports the trailer is usually unavailable even when the server sent it. - - The check therefore degrades gracefully: - - * trailer readable -> must match the computed digest, else raise; - * else size hint present -> received byte count must match, else raise - (catches truncation for known-size single files); - * else ``require_checksum`` -> raise (strict mode, for trailer-capable - transports); otherwise warn that the digest could not be verified. + The ``X-Content-Sha256`` trailer is authoritative when readable, but + CPython's HTTP stack discards chunked trailers, so the check degrades: use + the trailer if present, else the size hint, else warn (or raise under + ``require_checksum``). """ expected = _extract_trailer(response, _SHA256_HEADER) if expected: @@ -193,12 +170,7 @@ def _verify_download( def _coerce_upload_content(data: UploadData) -> "tuple[Any, Optional[int], Any]": - """Normalize an upload payload to ``(content, size_or_None, handle_to_close)``. - - ``content`` is suitable to hand to the transport (raw bytes or a readable - binary stream). For filesystem paths a handle is opened and returned so the - caller can close it after the request. - """ + """Normalize an upload payload to ``(content, size_or_None, handle_to_close)``.""" if isinstance(data, (bytes, bytearray)): payload = bytes(data) return payload, len(payload), None @@ -500,16 +472,10 @@ def workspace_upload( ) -> Any: """Upload raw file (or tar) bytes into a session's sandbox workspace. - ``POST /v2/agents/sessions/{session_id}/workspace/upload``. - - :param path: Destination path, resolved inside the workspace root - (``/workspace``). Anything escaping the root is rejected with 403. - :param data: Raw bytes, a filesystem path, or a readable binary stream. - :param is_archive: When ``True`` the body is a tar archive to extract at - ``path``. - :param content_sha256: Optional hex digest of the full payload, forwarded - via the ``X-Content-Sha256`` header for the guest to verify. - :returns: ``{"path": ..., "bytes_written": N}``. + ``POST /v2/agents/sessions/{session_id}/workspace/upload``. ``data`` is + bytes, a filesystem path, or a readable binary stream. ``is_archive`` + extracts the body as a tar at ``path``; ``content_sha256`` is forwarded + for the guest to verify. Returns ``{"path": ..., "bytes_written": N}``. """ if not path: raise ValueError("path is required") @@ -544,22 +510,11 @@ def workspace_download( ) -> "WorkspaceDownload": """Download a file (or tar-streamed directory) from a session workspace. - ``GET /v2/agents/sessions/{session_id}/workspace/download``. - - The response is chunked with the SHA-256 digest delivered as an HTTP - trailer after the body. The returned :class:`WorkspaceDownload` streams - the body and verifies integrity once it is fully consumed. A mismatched - trailer (or a byte count that disagrees with the server's size hint) - raises :class:`WorkspaceTransferError` and the output should be - discarded. - - :param path: Source path, resolved inside the workspace root. - :param as_archive: When ``True`` the directory at ``path`` is - tar-streamed. - :param require_checksum: When ``True``, raise if the SHA-256 trailer - cannot be read. The default is ``False`` because CPython's HTTP - stack discards chunked trailers, so strict verification only works - on a trailer-capable transport. + ``GET /v2/agents/sessions/{session_id}/workspace/download``. Returns a + :class:`WorkspaceDownload` that streams and verifies the body. + ``as_archive`` tar-streams the directory at ``path``. ``require_checksum`` + raises when the SHA-256 trailer cannot be read (default ``False`` since + CPython's HTTP stack discards chunked trailers). """ if not path: raise ValueError("path is required") @@ -581,12 +536,9 @@ def workspace_download( class WorkspaceDownload: """A streaming workspace download with best-effort integrity verification. - Iterating yields the raw body chunks while the SHA-256 digest is computed - incrementally. Once the body is fully consumed, integrity is verified (see - :func:`_verify_download`): a mismatched ``X-Content-Sha256`` trailer or a - byte count disagreeing with the size hint raises - :class:`WorkspaceTransferError`. Always consume the body to completion - (via iteration, :meth:`read`, or :meth:`save`) before trusting the output. + Iterating yields body chunks while computing the SHA-256; integrity is + verified once the body is fully consumed (see :func:`_verify_download`). + Consume it fully (iteration, :meth:`read`, or :meth:`save`) before trusting. """ def __init__(self, response: Any, *, require_checksum: bool = False): diff --git a/src/pydo/aio/agents/custom_sessions.py b/src/pydo/aio/agents/custom_sessions.py index 5e7fe4fc..4c276749 100644 --- a/src/pydo/aio/agents/custom_sessions.py +++ b/src/pydo/aio/agents/custom_sessions.py @@ -277,8 +277,7 @@ async def workspace_upload( raise ValueError("path is required") content, size, handle = _coerce_upload_content(data) try: - # aiohttp does not reliably stream arbitrary sync file objects, so - # materialize non-bytes payloads before sending. + # aiohttp can't reliably stream sync file objects; materialize them. if hasattr(content, "read"): content = content.read() if isinstance(content, str): diff --git a/tests/agents/test_workspace.py b/tests/agents/test_workspace.py index 0c1832e7..e828b693 100644 --- a/tests/agents/test_workspace.py +++ b/tests/agents/test_workspace.py @@ -91,7 +91,9 @@ def _make_resources(responses: List[_FakeResponse]) -> AgentsResources: parent = MagicMock() parent._client = MagicMock() parent._client._pipeline = _FakePipeline(responses) - return AgentsResources(parent, agents_endpoint="https://api.stage2.digitalocean.com") + return AgentsResources( + parent, agents_endpoint="https://api.stage2.digitalocean.com" + ) def _calls(resources) -> List[Any]: @@ -108,9 +110,7 @@ def test_upload_bytes_sets_path_params_and_content_type(): [_FakeResponse(200, body={"path": "/workspace/a.txt", "bytes_written": 5})] ) - resp = resources.sessions.workspace_upload( - "s1", path="a.txt", data=b"hello" - ) + resp = resources.sessions.workspace_upload("s1", path="a.txt", data=b"hello") call = _calls(resources)[0] assert call.request.method == "POST" @@ -143,7 +143,9 @@ def test_upload_accepts_filesystem_path(tmp_path): payload = b"file-on-disk" src = tmp_path / "input.bin" src.write_bytes(payload) - resources = _make_resources([_FakeResponse(200, body={"bytes_written": len(payload)})]) + resources = _make_resources( + [_FakeResponse(200, body={"bytes_written": len(payload)})] + ) resources.sessions.workspace_upload("s1", path="dest.bin", data=str(src)) @@ -223,9 +225,7 @@ def test_download_archive_flag_and_header(): ] ) - download = resources.sessions.workspace_download( - "s1", path="dir", as_archive=True - ) + download = resources.sessions.workspace_download("s1", path="dir", as_archive=True) assert download.read() == payload assert download.is_archive is True assert "as_archive=true" in _calls(resources)[0].request.url @@ -287,9 +287,7 @@ def test_download_size_hint_match_is_accepted(): def test_download_mismatched_trailer_is_failure(): - resources = _make_resources( - [_FakeResponse(200, chunks=[b"abc"], trailer="0" * 64)] - ) + resources = _make_resources([_FakeResponse(200, chunks=[b"abc"], trailer="0" * 64)]) download = resources.sessions.workspace_download("s1", path="x") with pytest.raises(WorkspaceTransferError, match="mismatch"): @@ -329,9 +327,7 @@ def test_download_save_writes_file_and_discards_on_failure(tmp_path): def test_download_non_200_raises(): from azure.core.exceptions import HttpResponseError - resources = _make_resources( - [_FakeResponse(404, body="path not found")] - ) + resources = _make_resources([_FakeResponse(404, body="path not found")]) with pytest.raises(HttpResponseError): resources.sessions.workspace_download("s1", path="missing") @@ -422,7 +418,9 @@ def _make_async_resources(responses: List[_FakeAsyncResponse]) -> AsyncAgentsRes @pytest.mark.asyncio async def test_async_upload_materializes_and_sets_headers(): - resources = _make_async_resources([_FakeAsyncResponse(200, body={"bytes_written": 3})]) + resources = _make_async_resources( + [_FakeAsyncResponse(200, body={"bytes_written": 3})] + ) resp = await resources.sessions.workspace_upload( "s1", path="a.txt", data=b"abc", content_sha256="cafe" @@ -442,7 +440,9 @@ async def test_async_download_verifies_trailer(): resources = _make_async_resources( [ _FakeAsyncResponse( - 200, chunks=[b"async-", b"bytes"], trailer=hashlib.sha256(payload).hexdigest() + 200, + chunks=[b"async-", b"bytes"], + trailer=hashlib.sha256(payload).hexdigest(), ) ] ) @@ -455,7 +455,9 @@ async def test_async_download_verifies_trailer(): @pytest.mark.asyncio async def test_async_download_missing_trailer_strict_fails(): - resources = _make_async_resources([_FakeAsyncResponse(200, chunks=[b"x"], trailer=None)]) + resources = _make_async_resources( + [_FakeAsyncResponse(200, chunks=[b"x"], trailer=None)] + ) download = await resources.sessions.workspace_download( "s1", path="o", require_checksum=True From 5e5eed6ca9d585ffb310b1d4b96b31215f8951f6 Mon Sep 17 00:00:00 2001 From: SSharma-10 Date: Fri, 3 Jul 2026 13:51:55 +0530 Subject: [PATCH 6/8] add support for session name --- examples/agents/attach_by_name.py | 52 ++++++++++++++++++++++++++ src/pydo/agents/__init__.py | 27 +++++++++++++ src/pydo/agents/custom_sessions.py | 7 ++++ src/pydo/aio/agents/__init__.py | 12 +++++- src/pydo/aio/agents/custom_sessions.py | 7 ++++ tests/agents/test_async_sessions.py | 43 +++++++++++++++++++++ tests/agents/test_sessions.py | 42 +++++++++++++++++++++ 7 files changed, 189 insertions(+), 1 deletion(-) create mode 100644 examples/agents/attach_by_name.py diff --git a/examples/agents/attach_by_name.py b/examples/agents/attach_by_name.py new file mode 100644 index 00000000..b68b18e9 --- /dev/null +++ b/examples/agents/attach_by_name.py @@ -0,0 +1,52 @@ +"""Look up a hosted-agent session by name (instead of by id). + +Sessions can be filtered server-side by name (``GET /v2/agents/sessions?name=``). +This script lists the matches and resolves the name to a session handle via +``client.agents.attach_by_name(...)``. + +Required env: + DIGITALOCEAN_TOKEN + PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com + SESSION_NAME the session name to look up +""" + +import os +import sys + +from pydo import Client + + +def main() -> int: + name = os.environ["SESSION_NAME"] + + client = Client( + token=os.environ["DIGITALOCEAN_TOKEN"], + agents_endpoint=os.environ.get("PYDO_AGENTS_ENDPOINT"), + ) + + resp = client.agents.sessions.list(name=name) + sessions = resp.get("sessions", []) if hasattr(resp, "get") else [] + print(f"[list name={name!r}] {len(sessions)} match(es):", file=sys.stderr) + for s in sessions: + print(f" {s['session_id']} {s['name']} {s['status']}", file=sys.stderr) + + if not sessions: + # Help diagnose a 0-match result: show what names actually exist now. + all_resp = client.agents.sessions.list() + all_sessions = all_resp.get("sessions", []) if hasattr(all_resp, "get") else [] + print( + f"[attach_by_name] no session named {name!r}. " + f"{len(all_sessions)} session(s) currently exist:", + file=sys.stderr, + ) + for s in all_sessions: + print(f" {s['name']} ({s['status']})", file=sys.stderr) + return 1 + + agent = client.agents.attach_by_name(name) + print(f"[attach_by_name] resolved session_id: {agent.session_id}", file=sys.stderr) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/pydo/agents/__init__.py b/src/pydo/agents/__init__.py index 33b76b68..ea3bcb30 100644 --- a/src/pydo/agents/__init__.py +++ b/src/pydo/agents/__init__.py @@ -50,6 +50,21 @@ def resolve_agents_base_url(explicit: Optional[str] = None) -> str: return url +def _select_session_by_name(list_response, name: str): + """Pick the most recently created session from a name-filtered list. + + Raises :class:`LookupError` when the response contains no sessions. + """ + get = getattr(list_response, "get", None) + sessions = (get("sessions") if get else None) or [] + if not sessions: + raise LookupError(f"no session found with name {name!r}") + return max( + sessions, + key=lambda s: (getattr(s, "get", lambda *_: "")("created_at") or ""), + ) + + class AgentsResources: def __init__(self, parent_client, *, agents_endpoint: Optional[str] = None): self._proxy = _BaseURLProxy( @@ -80,6 +95,18 @@ def attach(self, session_id: str) -> AgentSession: """Return an :class:`AgentSession` handle for an existing session.""" return AgentSession(self.sessions, session_id) + def attach_by_name(self, name: str) -> AgentSession: + """Resolve a session by ``name`` and return an :class:`AgentSession`. + + Looks up ``GET /v2/agents/sessions?name=``. If several sessions + share the name (e.g. reused over time), the most recently created one + is chosen. Raises :class:`LookupError` when there is no match. + """ + resp = self.sessions.list(name=name) + session = _select_session_by_name(resp, name) + session_id = (getattr(session, "get", lambda *_: None))("session_id") + return AgentSession(self.sessions, session_id, raw=session) + __all__ = [ "AgentsResources", diff --git a/src/pydo/agents/custom_sessions.py b/src/pydo/agents/custom_sessions.py index d5f5b167..005a06b1 100644 --- a/src/pydo/agents/custom_sessions.py +++ b/src/pydo/agents/custom_sessions.py @@ -347,7 +347,13 @@ def list( page_token: Optional[str] = None, page_size: Optional[int] = None, status: Optional[str] = None, + name: Optional[str] = None, ) -> Any: + """List sessions, optionally filtered by ``status`` and/or ``name``. + + ``name`` filters server-side (``GET /v2/agents/sessions?name=...``) and + may match more than one session (e.g. a name reused over time). + """ return self._parse_json( self._send( "GET", @@ -356,6 +362,7 @@ def list( "page_token": page_token, "page_size": page_size, "status": status, + "name": name, }, ), ) diff --git a/src/pydo/aio/agents/__init__.py b/src/pydo/aio/agents/__init__.py index ea08da9c..7ea22d9c 100644 --- a/src/pydo/aio/agents/__init__.py +++ b/src/pydo/aio/agents/__init__.py @@ -7,7 +7,7 @@ from typing import Optional -from pydo.agents import resolve_agents_base_url +from pydo.agents import _select_session_by_name, resolve_agents_base_url from pydo.custom_extensions import _BaseURLProxy from .custom_sessions import ( @@ -46,6 +46,16 @@ def attach(self, session_id: str) -> AsyncAgentSession: """Return an :class:`AsyncAgentSession` handle for an existing session.""" return AsyncAgentSession(self.sessions, session_id) + async def attach_by_name(self, name: str) -> AsyncAgentSession: + """Resolve a session by ``name`` and return an :class:`AsyncAgentSession`. + + See :meth:`pydo.agents.AgentsResources.attach_by_name`. + """ + resp = await self.sessions.list(name=name) + session = _select_session_by_name(resp, name) + session_id = (getattr(session, "get", lambda *_: None))("session_id") + return AsyncAgentSession(self.sessions, session_id, raw=session) + __all__ = [ "AsyncAgentsResources", diff --git a/src/pydo/aio/agents/custom_sessions.py b/src/pydo/aio/agents/custom_sessions.py index 4c276749..4cdffe9a 100644 --- a/src/pydo/aio/agents/custom_sessions.py +++ b/src/pydo/aio/agents/custom_sessions.py @@ -144,7 +144,13 @@ async def list( page_token: Optional[str] = None, page_size: Optional[int] = None, status: Optional[str] = None, + name: Optional[str] = None, ) -> Any: + """List sessions, optionally filtered by ``status`` and/or ``name``. + + ``name`` filters server-side (``GET /v2/agents/sessions?name=...``) and + may match more than one session (e.g. a name reused over time). + """ return await self._parse_json( await self._send( "GET", @@ -153,6 +159,7 @@ async def list( "page_token": page_token, "page_size": page_size, "status": status, + "name": name, }, ), ) diff --git a/tests/agents/test_async_sessions.py b/tests/agents/test_async_sessions.py index 1acb8e6e..be4f58e2 100644 --- a/tests/agents/test_async_sessions.py +++ b/tests/agents/test_async_sessions.py @@ -83,3 +83,46 @@ async def test_async_create_from_manifest_rejects_empty(): resources = _make_async_resources([]) with pytest.raises(ValueError): await resources.sessions.create_from_manifest("") + + +@pytest.mark.asyncio +async def test_async_list_filters_by_name(): + resources = _make_async_resources([_FakeAsyncResponse(200, {"sessions": []})]) + await resources.sessions.list(name="my-session") + + call = resources._proxy._original._pipeline.calls[0] + assert "name=my-session" in call.request.url + + +@pytest.mark.asyncio +async def test_async_attach_by_name_picks_most_recent_match(): + resources = _make_async_resources( + [ + _FakeAsyncResponse( + 200, + { + "sessions": [ + { + "session_id": "old", + "name": "dup", + "created_at": "2026-01-01T00:00:00Z", + }, + { + "session_id": "new", + "name": "dup", + "created_at": "2026-07-01T00:00:00Z", + }, + ] + }, + ) + ] + ) + agent = await resources.attach_by_name("dup") + assert agent.session_id == "new" + + +@pytest.mark.asyncio +async def test_async_attach_by_name_raises_when_not_found(): + resources = _make_async_resources([_FakeAsyncResponse(200, {"sessions": []})]) + with pytest.raises(LookupError): + await resources.attach_by_name("missing") diff --git a/tests/agents/test_sessions.py b/tests/agents/test_sessions.py index ccc23406..7d94aceb 100644 --- a/tests/agents/test_sessions.py +++ b/tests/agents/test_sessions.py @@ -154,6 +154,48 @@ def test_list_sessions_propagates_query_params(): assert "status=SESSION_STATUS_READY" in raw +def test_list_sessions_filters_by_name(): + resources = _make_resources([_FakeResponse(200, {"sessions": []})]) + resources.sessions.list(name="my-session") + + call = resources._proxy._original._pipeline.calls[0] + assert "name=my-session" in call.request.url + + +def test_attach_by_name_picks_most_recent_match(): + resources = _make_resources( + [ + _FakeResponse( + 200, + { + "sessions": [ + { + "session_id": "old", + "name": "dup", + "created_at": "2026-01-01T00:00:00Z", + }, + { + "session_id": "new", + "name": "dup", + "created_at": "2026-07-01T00:00:00Z", + }, + ] + }, + ) + ] + ) + agent = resources.attach_by_name("dup") + + assert "name=dup" in resources._proxy._original._pipeline.calls[0].request.url + assert agent.session_id == "new" + + +def test_attach_by_name_raises_when_not_found(): + resources = _make_resources([_FakeResponse(200, {"sessions": []})]) + with pytest.raises(LookupError): + resources.attach_by_name("missing") + + def test_send_input_body_shape(): resources = _make_resources([_FakeResponse(200, {"run_id": "r1"})]) resp = resources.sessions.send_input("s1", text="hello world") From d6cc1ffe39b976878d5ff6535996d6a03b8bd84c Mon Sep 17 00:00:00 2001 From: SSharma-10 Date: Fri, 3 Jul 2026 18:34:59 +0530 Subject: [PATCH 7/8] add support for pause/resume --- examples/agents/pause_resume.py | 76 ++++++++++++++++++++++++++ src/pydo/agents/custom_models.py | 1 + src/pydo/agents/custom_sessions.py | 12 ++++ src/pydo/agents/session.py | 6 ++ src/pydo/aio/agents/custom_sessions.py | 12 ++++ src/pydo/aio/agents/session.py | 6 ++ tests/agents/test_async_sessions.py | 24 ++++++++ tests/agents/test_sessions.py | 22 ++++++++ 8 files changed, 159 insertions(+) create mode 100644 examples/agents/pause_resume.py diff --git a/examples/agents/pause_resume.py b/examples/agents/pause_resume.py new file mode 100644 index 00000000..4453dc92 --- /dev/null +++ b/examples/agents/pause_resume.py @@ -0,0 +1,76 @@ +"""Pause and resume a hosted-agent session. + +Attaches to an existing session (by id or by name), pauses it, waits for it to +report ``SESSION_STATUS_PAUSED``, then resumes it and waits for ``READY``. + +Required env: + DIGITALOCEAN_TOKEN + PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com + +Pick a session with one of (SESSION_ID takes precedence): + SESSION_ID the session id to pause/resume + SESSION_NAME resolve the session by name instead +""" + +import os +import sys +import time + +from pydo import Client +from pydo.agents.custom_models import SessionStatus + + +def _wait_for_status(agent, target, *, timeout=120.0, poll_interval=2.0): + """Poll until the session reaches ``target`` (or a terminal/failed state).""" + deadline = time.monotonic() + timeout + while True: + agent.refresh() + status = agent.status + print(f" status: {status}", file=sys.stderr) + if status == target: + return + if status in (SessionStatus.FAILED, SessionStatus.DESTROYED): + raise RuntimeError(f"session {agent.session_id} is {status}") + if time.monotonic() > deadline: + raise TimeoutError( + f"session {agent.session_id} did not reach {target} in {timeout}s " + f"(last status: {status})" + ) + time.sleep(poll_interval) + + +def main() -> int: + client = Client( + token=os.environ["DIGITALOCEAN_TOKEN"], + agents_endpoint=os.environ.get("PYDO_AGENTS_ENDPOINT"), + ) + + session_id = os.environ.get("SESSION_ID") + session_name = os.environ.get("SESSION_NAME") + if session_id: + agent = client.agents.attach(session_id) + elif session_name: + agent = client.agents.attach_by_name(session_name) + else: + print("set SESSION_ID or SESSION_NAME", file=sys.stderr) + return 2 + + print(f"[attach] session_id={agent.session_id}", file=sys.stderr) + agent.refresh() + print(f"[attach] current status: {agent.status}", file=sys.stderr) + + print("[pause] pausing session...", file=sys.stderr) + agent.pause() + _wait_for_status(agent, SessionStatus.PAUSED) + print("[pause] session is PAUSED", file=sys.stderr) + + print("[resume] resuming session...", file=sys.stderr) + agent.resume() + _wait_for_status(agent, SessionStatus.READY) + print("[resume] session is READY", file=sys.stderr) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/pydo/agents/custom_models.py b/src/pydo/agents/custom_models.py index 940b3b21..ec49b279 100644 --- a/src/pydo/agents/custom_models.py +++ b/src/pydo/agents/custom_models.py @@ -25,6 +25,7 @@ class SessionStatus: PROVISIONING = "SESSION_STATUS_PROVISIONING" READY = "SESSION_STATUS_READY" DETACHED = "SESSION_STATUS_DETACHED" + PAUSED = "SESSION_STATUS_PAUSED" DESTROYING = "SESSION_STATUS_DESTROYING" DESTROYED = "SESSION_STATUS_DESTROYED" FAILED = "SESSION_STATUS_FAILED" diff --git a/src/pydo/agents/custom_sessions.py b/src/pydo/agents/custom_sessions.py index 005a06b1..b86b7956 100644 --- a/src/pydo/agents/custom_sessions.py +++ b/src/pydo/agents/custom_sessions.py @@ -395,6 +395,18 @@ def get(self, session_id: str) -> Any: def destroy(self, session_id: str) -> None: self._send("DELETE", f"{_BASE_PATH}/{_quote(session_id)}") + def pause(self, session_id: str) -> Any: + """Pause a running session (``POST .../{session_id}/pause``).""" + return self._parse_json( + self._send("POST", f"{_BASE_PATH}/{_quote(session_id)}/pause"), + ) + + def resume(self, session_id: str) -> Any: + """Resume a paused session (``POST .../{session_id}/resume``).""" + return self._parse_json( + self._send("POST", f"{_BASE_PATH}/{_quote(session_id)}/resume"), + ) + def send_input(self, session_id: str, *, text: str) -> Any: return self._parse_json( self._send( diff --git a/src/pydo/agents/session.py b/src/pydo/agents/session.py index d1eb7be8..4ae48f2f 100644 --- a/src/pydo/agents/session.py +++ b/src/pydo/agents/session.py @@ -369,6 +369,12 @@ def run( def send_input(self, text: str) -> Any: return self._sessions.send_input(self.session_id, text=text) + def pause(self) -> Any: + return self._sessions.pause(self.session_id) + + def resume(self) -> Any: + return self._sessions.resume(self.session_id) + def stream(self, **kwargs: Any) -> Any: return self._sessions.stream(self.session_id, **kwargs) diff --git a/src/pydo/aio/agents/custom_sessions.py b/src/pydo/aio/agents/custom_sessions.py index 4cdffe9a..a2e8a1c3 100644 --- a/src/pydo/aio/agents/custom_sessions.py +++ b/src/pydo/aio/agents/custom_sessions.py @@ -192,6 +192,18 @@ async def get(self, session_id: str) -> Any: async def destroy(self, session_id: str) -> None: await self._send("DELETE", f"{_BASE_PATH}/{_quote(session_id)}") + async def pause(self, session_id: str) -> Any: + """Pause a running session (``POST .../{session_id}/pause``).""" + return await self._parse_json( + await self._send("POST", f"{_BASE_PATH}/{_quote(session_id)}/pause"), + ) + + async def resume(self, session_id: str) -> Any: + """Resume a paused session (``POST .../{session_id}/resume``).""" + return await self._parse_json( + await self._send("POST", f"{_BASE_PATH}/{_quote(session_id)}/resume"), + ) + async def send_input(self, session_id: str, *, text: str) -> Any: return await self._parse_json( await self._send( diff --git a/src/pydo/aio/agents/session.py b/src/pydo/aio/agents/session.py index 06da8d8a..574bfc12 100644 --- a/src/pydo/aio/agents/session.py +++ b/src/pydo/aio/agents/session.py @@ -207,6 +207,12 @@ async def run( async def send_input(self, text: str) -> Any: return await self._sessions.send_input(self.session_id, text=text) + async def pause(self) -> Any: + return await self._sessions.pause(self.session_id) + + async def resume(self) -> Any: + return await self._sessions.resume(self.session_id) + async def stream(self, **kwargs: Any) -> Any: return await self._sessions.stream(self.session_id, **kwargs) diff --git a/tests/agents/test_async_sessions.py b/tests/agents/test_async_sessions.py index be4f58e2..3de3b91f 100644 --- a/tests/agents/test_async_sessions.py +++ b/tests/agents/test_async_sessions.py @@ -85,6 +85,30 @@ async def test_async_create_from_manifest_rejects_empty(): await resources.sessions.create_from_manifest("") +@pytest.mark.asyncio +async def test_async_pause_session(): + resources = _make_async_resources( + [_FakeAsyncResponse(200, {"session": {"session_id": "abc-123"}})] + ) + await resources.sessions.pause("abc-123") + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.method == "POST" + assert call.request.url.endswith("/v2/agents/sessions/abc-123/pause") + + +@pytest.mark.asyncio +async def test_async_resume_session(): + resources = _make_async_resources( + [_FakeAsyncResponse(200, {"session": {"session_id": "abc-123"}})] + ) + await resources.sessions.resume("abc-123") + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.method == "POST" + assert call.request.url.endswith("/v2/agents/sessions/abc-123/resume") + + @pytest.mark.asyncio async def test_async_list_filters_by_name(): resources = _make_async_resources([_FakeAsyncResponse(200, {"sessions": []})]) diff --git a/tests/agents/test_sessions.py b/tests/agents/test_sessions.py index 7d94aceb..439f0bab 100644 --- a/tests/agents/test_sessions.py +++ b/tests/agents/test_sessions.py @@ -141,6 +141,28 @@ def test_destroy_session(): assert call.request.url.endswith("/v2/agents/sessions/abc-123") +def test_pause_session(): + resources = _make_resources( + [_FakeResponse(200, {"session": {"session_id": "abc-123"}})] + ) + resources.sessions.pause("abc-123") + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.method == "POST" + assert call.request.url.endswith("/v2/agents/sessions/abc-123/pause") + + +def test_resume_session(): + resources = _make_resources( + [_FakeResponse(200, {"session": {"session_id": "abc-123"}})] + ) + resources.sessions.resume("abc-123") + + call = resources._proxy._original._pipeline.calls[0] + assert call.request.method == "POST" + assert call.request.url.endswith("/v2/agents/sessions/abc-123/resume") + + def test_list_sessions_propagates_query_params(): resources = _make_resources( [_FakeResponse(200, {"sessions": [], "next_page_token": ""})] From 2cb74b86c1b1c5dcfe30db3982c3ff813eba81f1 Mon Sep 17 00:00:00 2001 From: rahulkumarvh Date: Mon, 6 Jul 2026 12:31:50 -0700 Subject: [PATCH 8/8] =?UTF-8?q?examples/agents:=20add=20policy=20engine=20?= =?UTF-8?q?test=20scripts=20Four=20standalone=20scripts=20exercising=20the?= =?UTF-8?q?=20permissions=20block:=20-=20policy=5Fauto=5Fallow.py=20=20?= =?UTF-8?q?=E2=80=94=20touch=20*=20=E2=86=92=20allow=20runs=20without=20HI?= =?UTF-8?q?TL=20-=20policy=5Fauto=5Fdeny.py=20=20=20=E2=80=94=20ls=20*=20?= =?UTF-8?q?=E2=86=92=20deny=20blocks=20without=20HITL=20prompt=20-=20polic?= =?UTF-8?q?y=5Fhitl=5Fask.py=20=20=20=20=E2=80=94=20defaultAction:=20ask?= =?UTF-8?q?=20gates=20unmatched=20commands;=20=20=20=20=20=20=20=20=20=20?= =?UTF-8?q?=20=20=20=20=20=20=20=20=20=20=20=20=20=20=20=20=20tests=20both?= =?UTF-8?q?=20approve=20and=20reject=20resolution=20paths=20-=20policy=5Fw?= =?UTF-8?q?rite=5Finterception.py=20=E2=80=94=20Write/Edit=20(non-Bash)=20?= =?UTF-8?q?tool=20intercepted=20=20=20=20=20=20=20=20=20=20=20=20=20=20=20?= =?UTF-8?q?=20=20=20=20=20=20=20=20=20=20=20=20=20=20=20=20=20=20=20under?= =?UTF-8?q?=20defaultAction:=20ask=20All=20scripts=20load=20the=20base=20A?= =?UTF-8?q?GENT=5FSPEC,=20inject=20a=20permissions=20block,=20and=20assert?= =?UTF-8?q?=20pass/fail=20without=20external=20YAML=20files.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- examples/agents/policy_auto_allow.py | 99 ++++++++++++++ examples/agents/policy_auto_deny.py | 100 ++++++++++++++ examples/agents/policy_hitl_ask.py | 129 +++++++++++++++++++ examples/agents/policy_write_interception.py | 104 +++++++++++++++ 4 files changed, 432 insertions(+) create mode 100644 examples/agents/policy_auto_allow.py create mode 100644 examples/agents/policy_auto_deny.py create mode 100644 examples/agents/policy_hitl_ask.py create mode 100644 examples/agents/policy_write_interception.py diff --git a/examples/agents/policy_auto_allow.py b/examples/agents/policy_auto_allow.py new file mode 100644 index 00000000..3d9260b5 --- /dev/null +++ b/examples/agents/policy_auto_allow.py @@ -0,0 +1,99 @@ +"""Policy engine test: auto-allow rule. + +Creates a session with ``defaultAction: ask`` and a single ``touch * → allow`` +rule, then asks the agent to run ``touch /workspace/allow_test.txt``. + +Expected: the Bash tool call matches the allow rule and executes WITHOUT any +HITL prompt being raised. The test FAILS if a ``hitl_requested`` event fires. + +Required env: + DIGITALOCEAN_TOKEN + AGENT_SPEC path to the base agents.yaml (default: agent-spec.yaml) + +Optional env: + PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com +""" + +import os +import sys + +import yaml + +from pydo import Client +from pydo.agents import AgentEventType + + +_PERMISSIONS = { + "defaultAction": "ask", + "rules": [ + {"tool": "bash", "match": {"command": "touch *"}, "action": "allow"}, + {"tool": "bash", "match": {"command": "git status"}, "action": "allow"}, + ], +} + +_PROMPT = ( + "Use bash to run exactly this command and nothing else: " + "touch /workspace/allow_test.txt" +) + + +def _load_manifest(name: str, permissions: dict) -> str: + spec_path = os.environ.get("AGENT_SPEC", "agent-spec.yaml") + with open(spec_path, "r", encoding="utf-8") as fh: + spec = yaml.safe_load(fh) + spec["metadata"]["name"] = name + spec["spec"]["permissions"] = permissions + return yaml.dump(spec, default_flow_style=False) + + +def main() -> int: + client = Client( + token=os.environ["DIGITALOCEAN_TOKEN"], + agents_endpoint=os.environ.get("PYDO_AGENTS_ENDPOINT"), + ) + + manifest = _load_manifest("policy-auto-allow-test", _PERMISSIONS) + hitl_events = [] + + def _track_and_approve(event): + hitl_events.append(event) + return "approve" # resolve so the run doesn't hang if HITL fires unexpectedly + + with client.agents.start(manifest) as agent: + print(f"[session {agent.session_id}]", file=sys.stderr) + print(f"[policy] defaultAction=ask touch * → allow", file=sys.stderr) + print(f"[prompt] {_PROMPT}\n", file=sys.stderr) + + stream = agent.run_streamed(_PROMPT, hitl=_track_and_approve, timeout=180) + for event in stream: + if event.type == AgentEventType.TOKEN: + print(event.text, end="", flush=True) + elif event.type == AgentEventType.TOOL_CALL: + print(f"\n[tool_call] {event.tool_name}", file=sys.stderr) + elif event.type == AgentEventType.HITL_REQUESTED: + print(f"\n[hitl_fired] request_id={event.request_id}", file=sys.stderr) + + result = stream.result + print(f"\n[{result.status}]", file=sys.stderr) + + if hitl_events: + print( + f"\nFAIL HITL fired {len(hitl_events)} time(s) — " + "touch should have been auto-allowed without prompting", + file=sys.stderr, + ) + return 1 + + if result.status != "completed": + print( + f"\nFAIL run ended with status={result.status!r}", + file=sys.stderr, + ) + return 1 + + print("\nPASS touch ran without HITL prompt (auto-allow confirmed)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/agents/policy_auto_deny.py b/examples/agents/policy_auto_deny.py new file mode 100644 index 00000000..3ac4c5d6 --- /dev/null +++ b/examples/agents/policy_auto_deny.py @@ -0,0 +1,100 @@ +"""Policy engine test: auto-deny rule. + +Creates a session with ``ls * → deny`` and ``rm -rf * → deny`` rules and asks +the agent to run ``ls /workspace``. + +Expected: the deny rule blocks the Bash call automatically — the agent receives +a "forbidden" response and continues WITHOUT raising a HITL prompt. The test +FAILS if a ``hitl_requested`` event fires. + +Required env: + DIGITALOCEAN_TOKEN + AGENT_SPEC path to the base agents.yaml (default: agent-spec.yaml) + +Optional env: + PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com +""" + +import os +import sys + +import yaml + +from pydo import Client +from pydo.agents import AgentEventType + + +_PERMISSIONS = { + "defaultAction": "ask", + "rules": [ + {"tool": "bash", "match": {"command": "ls *"}, "action": "deny"}, + {"tool": "bash", "match": {"command": "rm -rf *"}, "action": "deny"}, + {"tool": "bash", "match": {"command": "touch *"}, "action": "allow"}, + ], +} + +_PROMPT = ( + "Use bash to run exactly this command and nothing else: ls /workspace" +) + + +def _load_manifest(name: str, permissions: dict) -> str: + spec_path = os.environ.get("AGENT_SPEC", "agent-spec.yaml") + with open(spec_path, "r", encoding="utf-8") as fh: + spec = yaml.safe_load(fh) + spec["metadata"]["name"] = name + spec["spec"]["permissions"] = permissions + return yaml.dump(spec, default_flow_style=False) + + +def main() -> int: + client = Client( + token=os.environ["DIGITALOCEAN_TOKEN"], + agents_endpoint=os.environ.get("PYDO_AGENTS_ENDPOINT"), + ) + + manifest = _load_manifest("policy-auto-deny-test", _PERMISSIONS) + hitl_events = [] + + def _track_and_approve(event): + hitl_events.append(event) + return "approve" # resolve to not hang; we assert this never fires + + with client.agents.start(manifest) as agent: + print(f"[session {agent.session_id}]", file=sys.stderr) + print(f"[policy] ls * → deny rm -rf * → deny defaultAction=ask", file=sys.stderr) + print(f"[prompt] {_PROMPT}\n", file=sys.stderr) + + stream = agent.run_streamed(_PROMPT, hitl=_track_and_approve, timeout=180) + for event in stream: + if event.type == AgentEventType.TOKEN: + print(event.text, end="", flush=True) + elif event.type == AgentEventType.TOOL_CALL: + print(f"\n[tool_call] {event.tool_name}", file=sys.stderr) + elif event.type == AgentEventType.HITL_REQUESTED: + print(f"\n[hitl_fired] request_id={event.request_id}", file=sys.stderr) + + result = stream.result + print(f"\n[{result.status}]", file=sys.stderr) + + if hitl_events: + print( + f"\nFAIL HITL fired {len(hitl_events)} time(s) — " + "ls should have been auto-denied without prompting", + file=sys.stderr, + ) + return 1 + + if result.status != "completed": + print( + f"\nFAIL run ended with status={result.status!r}", + file=sys.stderr, + ) + return 1 + + print("\nPASS ls was auto-denied without HITL prompt") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/agents/policy_hitl_ask.py b/examples/agents/policy_hitl_ask.py new file mode 100644 index 00000000..268af67d --- /dev/null +++ b/examples/agents/policy_hitl_ask.py @@ -0,0 +1,129 @@ +"""Policy engine test: defaultAction: ask → HITL. + +Runs two sub-tests against a session with ``defaultAction: ask`` and no rule +matching ``mkdir``: + + Sub-test approve: + Agent tries ``mkdir /workspace/hitl_dir``. + Expected: HITL prompt fires → we approve → run completes. + + Sub-test reject: + Same prompt on a fresh session. + Expected: HITL prompt fires → we reject → agent acknowledges denial, run + completes (agent continues the conversation after the rejection). + +Both sub-tests FAIL if no ``hitl_requested`` event fires, or if the run ends in +a non-completed state. + +Required env: + DIGITALOCEAN_TOKEN + AGENT_SPEC path to the base agents.yaml (default: agent-spec.yaml) + +Optional env: + PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com +""" + +import os +import sys + +import yaml + +from pydo import Client +from pydo.agents import AgentEventType + + +_PERMISSIONS = { + "defaultAction": "ask", + "rules": [ + # touch is explicitly allowed so the agent can set up; mkdir is unmatched → ask + {"tool": "bash", "match": {"command": "touch *"}, "action": "allow"}, + {"tool": "bash", "match": {"command": "git status"}, "action": "allow"}, + ], +} + +_PROMPT = ( + "Use bash to run exactly this command and nothing else: " + "mkdir /workspace/hitl_test_dir" +) + + +def _load_manifest(name: str, permissions: dict) -> str: + spec_path = os.environ.get("AGENT_SPEC", "agent-spec.yaml") + with open(spec_path, "r", encoding="utf-8") as fh: + spec = yaml.safe_load(fh) + spec["metadata"]["name"] = name + spec["spec"]["permissions"] = permissions + return yaml.dump(spec, default_flow_style=False) + + +def _run_hitl_subtest(client, manifest_name: str, decision: str, label: str) -> int: + """Run one HITL sub-test, resolving HITL with ``decision`` (approve/reject).""" + manifest = _load_manifest(manifest_name, _PERMISSIONS) + hitl_events = [] + + def _track_and_decide(event): + hitl_events.append(event) + print(f"\n[hitl_{decision}] request_id={event.request_id}", file=sys.stderr) + return decision + + with client.agents.start(manifest) as agent: + print(f"\n[{label}] session={agent.session_id}", file=sys.stderr) + print(f"[{label}] prompt: {_PROMPT}", file=sys.stderr) + + stream = agent.run_streamed(_PROMPT, hitl=_track_and_decide, timeout=180) + for event in stream: + if event.type == AgentEventType.TOKEN: + print(event.text, end="", flush=True) + elif event.type == AgentEventType.TOOL_CALL: + print(f"\n[tool_call] {event.tool_name}", file=sys.stderr) + + result = stream.result + print(f"\n[{result.status}]", file=sys.stderr) + + if not hitl_events: + print( + f"\nFAIL [{label}] no HITL event fired — " + "mkdir should have been gated by defaultAction: ask", + file=sys.stderr, + ) + return 1 + + if result.status != "completed": + print( + f"\nFAIL [{label}] run ended with status={result.status!r} " + f"(expected 'completed')", + file=sys.stderr, + ) + return 1 + + print(f"\nPASS [{label}] HITL fired and was {decision}d, run completed") + return 0 + + +def main() -> int: + client = Client( + token=os.environ["DIGITALOCEAN_TOKEN"], + agents_endpoint=os.environ.get("PYDO_AGENTS_ENDPOINT"), + ) + + print("[policy] defaultAction=ask touch * → allow mkdir → unmatched", file=sys.stderr) + + rc_a = _run_hitl_subtest( + client, + manifest_name="policy-hitl-approve-test", + decision="approve", + label="approve", + ) + + rc_b = _run_hitl_subtest( + client, + manifest_name="policy-hitl-reject-test", + decision="reject", + label="reject", + ) + + return max(rc_a, rc_b) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/agents/policy_write_interception.py b/examples/agents/policy_write_interception.py new file mode 100644 index 00000000..d85cb5d3 --- /dev/null +++ b/examples/agents/policy_write_interception.py @@ -0,0 +1,104 @@ +"""Policy engine test: non-Bash tool interception (Write/Edit). + +Creates a session with ``defaultAction: ask`` (no explicit bash rules) and asks +the agent to write a file. Codex uses its Write tool (not bash) for this, so +the interception exercises the non-Bash tool path of the policy engine. + +Expected: a ``hitl_requested`` event fires for the Write/Edit tool call. +The test FAILS if no HITL prompt is raised. + +Required env: + DIGITALOCEAN_TOKEN + AGENT_SPEC path to the base agents.yaml (default: agent-spec.yaml) + +Optional env: + PYDO_AGENTS_ENDPOINT stage2: https://api.s2r1.internal.digitalocean.com +""" + +import os +import sys + +import yaml + +from pydo import Client +from pydo.agents import AgentEventType + + +_PERMISSIONS = { + "defaultAction": "ask", + # No rules → every tool call (Bash, Write, Edit, …) requires approval +} + +_PROMPT = ( + "Write the text 'policy engine write interception test' into the file " + "/workspace/write_intercept_test.txt using your file-write capability " + "(not bash). Do not use bash." +) + + +def _load_manifest(name: str, permissions: dict) -> str: + spec_path = os.environ.get("AGENT_SPEC", "agent-spec.yaml") + with open(spec_path, "r", encoding="utf-8") as fh: + spec = yaml.safe_load(fh) + spec["metadata"]["name"] = name + spec["spec"]["permissions"] = permissions + return yaml.dump(spec, default_flow_style=False) + + +def main() -> int: + client = Client( + token=os.environ["DIGITALOCEAN_TOKEN"], + agents_endpoint=os.environ.get("PYDO_AGENTS_ENDPOINT"), + ) + + manifest = _load_manifest("policy-write-intercept-test", _PERMISSIONS) + hitl_events = [] + tool_calls = [] + + def _track_and_approve(event): + hitl_events.append(event) + print(f"\n[hitl_fired] request_id={event.request_id}", file=sys.stderr) + return "approve" + + with client.agents.start(manifest) as agent: + print(f"[session {agent.session_id}]", file=sys.stderr) + print(f"[policy] defaultAction=ask (no rules — all tools intercepted)", file=sys.stderr) + print(f"[prompt] {_PROMPT}\n", file=sys.stderr) + + stream = agent.run_streamed(_PROMPT, hitl=_track_and_approve, timeout=180) + for event in stream: + if event.type == AgentEventType.TOKEN: + print(event.text, end="", flush=True) + elif event.type == AgentEventType.TOOL_CALL: + tool_calls.append(event.tool_name) + print(f"\n[tool_call] {event.tool_name}", file=sys.stderr) + elif event.type == AgentEventType.HITL_REQUESTED: + pass # already logged in the callback + + result = stream.result + print(f"\n[{result.status}] tool_calls={tool_calls}", file=sys.stderr) + + if not hitl_events: + print( + "\nFAIL no HITL event fired — Write/Edit tool should have been " + "intercepted by defaultAction: ask", + file=sys.stderr, + ) + return 1 + + if result.status != "completed": + print( + f"\nFAIL run ended with status={result.status!r}", + file=sys.stderr, + ) + return 1 + + print( + f"\nPASS Write/Edit tool intercepted by policy engine " + f"({len(hitl_events)} HITL event(s) fired, tool_calls={tool_calls})" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())