diff --git a/contracts/core-client.v1.json b/contracts/core-client.v1.json new file mode 100644 index 0000000..ebc1359 --- /dev/null +++ b/contracts/core-client.v1.json @@ -0,0 +1,33 @@ +{ + "name": "openlinker-core-client", + "version": "v1", + "package": "openlinker", + "scope": "core", + "rules": { + "single_package_first": true, + "forbidden_path_prefixes": [ + "/api/v1/cloud", + "/api/v1/tasks", + "/api/v1/task-templates", + "/api/v1/workflows", + "/api/v1/workflow-runs", + "/api/v1/billing", + "/api/v1/wallet", + "/api/v1/stripe" + ] + }, + "endpoints": [ + {"client_method":"list_agents","http_method":"GET","path":"/api/v1/agents","auth":"none"}, + {"client_method":"get_agent","http_method":"GET","path":"/api/v1/agents/{slug}","auth":"none"}, + {"client_method":"get_agent_card","http_method":"GET","path":"/api/v1/agents/{slug}/agent-card.json","auth":"none"}, + {"client_method":"get_agent_card","http_method":"GET","path":"/api/v1/agents/{slug}/agent-card.extended.json","auth":"none"}, + {"client_method":"run_agent","http_method":"POST","path":"/api/v1/run","auth":"access_token","required_scope":"agents:run","required_headers":["Idempotency-Key"],"success_statuses":[200,201,202],"response_headers":["Location","Idempotency-Replayed"],"response_fields":["agent_connection_mode","runtime_transport","runtime_transport_reason","runtime_transport_changed_at","dispatch_state","attempt_count","max_attempts","started_at","replayed"]}, + {"client_method":"start_agent_run","http_method":"POST","path":"/api/v1/runs","auth":"access_token","required_scope":"agents:run","required_headers":["Idempotency-Key"],"success_statuses":[200,201,202],"response_headers":["Location","Idempotency-Replayed"],"response_fields":["agent_connection_mode","runtime_transport","runtime_transport_reason","runtime_transport_changed_at","dispatch_state","attempt_count","max_attempts","started_at","replayed"]}, + {"client_method":"get_run","http_method":"GET","path":"/api/v1/runs/{id}","auth":"access_token","required_scope":"runs:read","response_fields":["agent_connection_mode","runtime_transport","runtime_transport_reason","runtime_transport_changed_at","dispatch_state","attempt_count","max_attempts","started_at","replayed"]}, + {"client_method":"list_run_events","http_method":"GET","path":"/api/v1/runs/{id}/events","auth":"access_token","required_scope":"runs:read","response_fields":["items","meta.requested_after_sequence","meta.effective_after_sequence","meta.retained_through_sequence","meta.earliest_available_sequence","meta.latest_available_sequence","meta.retention_gap","meta.terminal","meta.stream_complete"]}, + {"client_method":"list_run_children","http_method":"GET","path":"/api/v1/runs/{id}/children","auth":"access_token","required_scope":"runs:read","response_fields":["parent_run_id","items[].child_run_id","items[].caller_agent_id","items[].target_agent_id","items[].status","items[].children"]}, + {"client_method":"list_run_artifacts","http_method":"GET","path":"/api/v1/runs/{id}/artifacts","auth":"access_token","required_scope":"runs:read"}, + {"client_method":"list_run_messages","http_method":"GET","path":"/api/v1/runs/{id}/messages","auth":"access_token","required_scope":"runs:read"}, + {"client_method":"stream_run_events","http_method":"GET","path":"/api/v1/runs/{id}/stream","auth":"access_token","required_scope":"runs:read"} + ] +} diff --git a/contracts/core-registration.v1.json b/contracts/core-registration.v1.json new file mode 100644 index 0000000..352b163 --- /dev/null +++ b/contracts/core-registration.v1.json @@ -0,0 +1,16 @@ +{ + "name": "openlinker-registration", + "version": "v1", + "scope": "core-registration", + "endpoints": [ + {"client_method":"create_agent","http_method":"POST","path":"/api/v1/creator/agents","auth":"user_token"}, + {"client_method":"list_my_agents","http_method":"GET","path":"/api/v1/creator/agents","auth":"user_token"}, + {"client_method":"get_my_agent","http_method":"GET","path":"/api/v1/creator/agents/{id}","auth":"user_token"}, + {"client_method":"get_my_agent_by_slug","http_method":"GET","path":"/api/v1/creator/agents/by-slug/{slug}","auth":"user_token"}, + {"client_method":"update_agent","http_method":"PATCH","path":"/api/v1/creator/agents/{id}","auth":"user_token"}, + {"client_method":"create_agent_token","http_method":"POST","path":"/api/v1/creator/agent-tokens","auth":"user_token"}, + {"client_method":"list_agent_tokens","http_method":"GET","path":"/api/v1/creator/agent-tokens","auth":"user_token"}, + {"client_method":"revoke_agent_token","http_method":"DELETE","path":"/api/v1/creator/agent-tokens/{id}","auth":"user_token"}, + {"client_method":"register_agent_via_token","http_method":"POST","path":"/api/v1/agent-registration/agents","auth":"agent_token"} + ] +} diff --git a/src/openlinker/__init__.py b/src/openlinker/__init__.py index fcb2625..1f1a630 100644 --- a/src/openlinker/__init__.py +++ b/src/openlinker/__init__.py @@ -1,3 +1,3 @@ -from . import client, runtime +from . import client, registration, runtime -__all__ = ["client", "runtime"] +__all__ = ["client", "registration", "runtime"] diff --git a/src/openlinker/client/__init__.py b/src/openlinker/client/__init__.py index a4a534d..b3fa463 100644 --- a/src/openlinker/client/__init__.py +++ b/src/openlinker/client/__init__.py @@ -3,6 +3,7 @@ import asyncio import inspect import json +import secrets from collections.abc import AsyncIterator, Awaitable from datetime import datetime, timezone, timedelta from email.utils import parsedate_to_datetime @@ -30,6 +31,8 @@ ListRunEventsResponse, MarketListResponse, PlatformCallbackOptions, + RegisterAgentViaTokenRequest, + RegisterAgentViaTokenResponse, RunAgentRequest, RunArtifactResponse, RunMessageResponse, @@ -92,6 +95,24 @@ def _quote(value: str) -> str: return quote(value, safe="") +def _resolve_run_idempotency_key(value: str | None) -> str: + key = secrets.token_hex(32) if value is None else value + if ( + not isinstance(key, str) + or not 1 <= len(key) <= 255 + or any(ord(char) < 0x20 or ord(char) > 0x7E for char in key) + ): + raise ValueError("openlinker: idempotency_key must be 1-255 printable ASCII characters") + return key + + +def _normalize_registration_connection_mode(value: str | None) -> str: + normalized = _clean(value).lower() + if normalized in {"", "runtime", "runtime_ws", "runtime_pull", "agent_node"}: + return "runtime" + return _clean(value) + + def _maybe_await(value: Any) -> Awaitable[Any]: if inspect.isawaitable(value): return value @@ -124,6 +145,24 @@ def _retry_after(headers: httpx.Headers) -> timedelta | None: return None +async def _read_response_body(response: httpx.Response) -> bytes: + declared = response.headers.get("Content-Length") + if declared: + try: + declared_length = int(declared) + except ValueError: + declared_length = 0 + if declared_length > MAX_RESPONSE_BODY_BYTES: + raise ValueError(f"openlinker: response body exceeds {MAX_RESPONSE_BODY_BYTES} bytes") + + body = bytearray() + async for chunk in response.aiter_bytes(): + if len(body) + len(chunk) > MAX_RESPONSE_BODY_BYTES: + raise ValueError(f"openlinker: response body exceeds {MAX_RESPONSE_BODY_BYTES} bytes") + body.extend(chunk) + return bytes(body) + + class Client: def __init__( self, @@ -200,16 +239,20 @@ async def _request( body: Any = None, accept: str = "application/json", token: str = "", + headers: dict[str, str] | None = None, ) -> httpx.Response: content = None if body is not None: content = json.dumps(to_json_value(body)).encode() - return await self._client.request( + request_headers = self._request_headers(accept, token, body is not None) + request_headers.update(headers or {}) + request = self._client.build_request( method, self.endpoint(path, query), content=content, - headers=self._request_headers(accept, token, body is not None), + headers=request_headers, ) + return await self._client.send(request, stream=True) async def _do( self, @@ -219,34 +262,60 @@ async def _do( query: dict[str, Any] | None = None, body: Any = None, out: type[T] | None = None, + headers: dict[str, str] | None = None, + token: str | None = None, ) -> T | dict[str, Any] | None: - response = await self._request( + result, _ = await self._do_with_response( method, path, query=query, body=body, - token=self.user_token, + out=out, + headers=headers, + token=token, ) - if response.status_code < 200 or response.status_code >= 300: - raise self._parse_error(response) - if response.status_code == 204: - return None - raw = response.content - if len(raw) > MAX_RESPONSE_BODY_BYTES: - raise ValueError(f"openlinker: response body exceeds {MAX_RESPONSE_BODY_BYTES} bytes") - if not raw: - return None - data = response.json() - if out is None: - return data - return out.from_dict(data) + return result - def _parse_error(self, response: httpx.Response) -> OpenLinkerError: - raw = response.content[:MAX_RESPONSE_BODY_BYTES] + async def _do_with_response( + self, + method: str, + path: str, + *, + query: dict[str, Any] | None = None, + body: Any = None, + out: type[T] | None = None, + headers: dict[str, str] | None = None, + token: str | None = None, + ) -> tuple[T | dict[str, Any] | None, httpx.Headers]: + response = await self._request( + method, + path, + query=query, + body=body, + token=self.user_token if token is None else token, + headers=headers, + ) + try: + raw = b"" if response.status_code == 204 else await _read_response_body(response) + if response.status_code < 200 or response.status_code >= 300: + raise self._parse_error(response, raw) + response_headers = httpx.Headers(response.headers) + if response.status_code == 204 or not raw: + return None, response_headers + data = json.loads(raw) + if out is None: + return data, response_headers + return out.from_dict(data), response_headers + finally: + await response.aclose() + + def _parse_error(self, response: httpx.Response, raw: bytes | None = None) -> OpenLinkerError: + if raw is None: + raw = response.content[:MAX_RESPONSE_BODY_BYTES] parsed: dict[str, Any] = {} try: - parsed = response.json() - except ValueError: + parsed = json.loads(raw) + except (TypeError, ValueError): pass err = parsed.get("error") if isinstance(parsed, dict) else None err = err if isinstance(err, dict) else {} @@ -280,14 +349,29 @@ async def get_agent_card(self, slug: str, extended: bool = False): return await self._do("GET", f"/agents/{_quote(slug)}/{suffix}", out=AgentCardResponse) async def run_agent(self, req: RunAgentRequest | dict[str, Any]): - return await self._do( - "POST", "/run", body=maybe_model(req, RunAgentRequest), out=RunResponse - ) + return await self._create_run("/run", req) async def start_agent_run(self, req: RunAgentRequest | dict[str, Any]): - return await self._do( - "POST", "/runs", body=maybe_model(req, RunAgentRequest), out=RunResponse + return await self._create_run("/runs", req) + + async def _create_run(self, path: str, req: RunAgentRequest | dict[str, Any]): + request = maybe_model(req, RunAgentRequest) + key = _resolve_run_idempotency_key(request.idempotency_key) + body = request.to_dict() + body.pop("idempotency_key", None) + result, headers = await self._do_with_response( + "POST", + path, + body=body, + out=RunResponse, + headers={"Idempotency-Key": key}, ) + if ( + isinstance(result, RunResponse) + and headers.get("Idempotency-Replayed", "").lower() == "true" + ): + result.replayed = True + return result async def get_run(self, run_id: str): return await self._do("GET", f"/runs/{_quote(run_id)}", out=RunResponse) @@ -330,10 +414,8 @@ async def stream_run_events( headers=self._request_headers("text/event-stream", self.user_token), ) as response: if response.status_code < 200 or response.status_code >= 300: - body = await response.aread() - raise self._parse_error( - httpx.Response(response.status_code, headers=response.headers, content=body) - ) + body = await _read_response_body(response) + raise self._parse_error(response, body) async for event in read_sse(response.aiter_lines()): yield event @@ -342,7 +424,9 @@ async def run_agent_with_callbacks( ): started = await self.start_agent_run(req) await self._stream_platform_callbacks(started.run_id, opts, until_terminal=True) - return await self.get_run(started.run_id) + result = await self.get_run(started.run_id) + result.replayed = result.replayed or started.replayed + return result async def start_agent_run_with_callbacks( self, req: RunAgentRequest | dict[str, Any], opts: PlatformCallbackOptions @@ -431,6 +515,28 @@ async def list_agent_tokens(self, params: ListAgentTokensParams | dict[str, Any] async def revoke_agent_token(self, token_id: str) -> None: await self._do("DELETE", f"/creator/agent-tokens/{_quote(token_id)}") + async def register_agent_via_token( + self, + agent_token: str, + req: RegisterAgentViaTokenRequest | dict[str, Any], + ): + token = _clean(agent_token) + if not token: + raise ValueError("openlinker: Agent Token is required for registration") + request = maybe_model(req, RegisterAgentViaTokenRequest) + body = request.to_dict() + body["visibility"] = request.visibility or "private" + body["connection_mode"] = _normalize_registration_connection_mode( + request.connection_mode + ) + return await self._do( + "POST", + "/agent-registration/agents", + body=body, + out=RegisterAgentViaTokenResponse, + token=token, + ) + def a2a_agent(self, slug: str): from ..a2a import A2AClient @@ -464,6 +570,7 @@ def a2a_agent(self, slug: str): CreateAgentToken = create_agent_token ListAgentTokens = list_agent_tokens RevokeAgentToken = revoke_agent_token + RegisterAgentViaToken = register_agent_via_token A2AAgent = a2a_agent diff --git a/src/openlinker/registration.py b/src/openlinker/registration.py new file mode 100644 index 0000000..1cfea3e --- /dev/null +++ b/src/openlinker/registration.py @@ -0,0 +1,405 @@ +from __future__ import annotations + +import json +import os +import tempfile +import threading +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Protocol, runtime_checkable + +from .client import Client +from .types import CreateAgentTokenRequest, RegisterAgentViaTokenRequest + + +DEFAULT_API_BASE = "https://api.openlinker.ai" +DEFAULT_REGISTRATION_ENV_PATH = ".env" +REGISTER_REUSE_EXISTING = "reuse_existing" +REGISTER_ROTATE_TOKEN = "rotate_token" +REGISTER_FORCE_NEW = "force_new" +REGISTER_VALIDATE_ONLY = "validate_only" +_REGISTER_POLICIES = { + REGISTER_REUSE_EXISTING, + REGISTER_ROTATE_TOKEN, + REGISTER_FORCE_NEW, + REGISTER_VALIDATE_ONLY, +} +_ENV_KEYS = ( + "OPENLINKER_API_BASE", + "OPENLINKER_AGENT_ID", + "OPENLINKER_AGENT_SLUG", + "OPENLINKER_AGENT_NAME", + "OPENLINKER_AGENT_TOKEN", + "OPENLINKER_AGENT_TOKEN_ID", + "OPENLINKER_AGENT_TOKEN_PREFIX", + "OPENLINKER_REGISTERED_AT", + "OPENLINKER_UPDATED_AT", +) + + +@dataclass +class AgentRegistration: + agent_id: str = "" + agent_slug: str = "" + agent_name: str = "" + agent_token: str = "" + token_id: str = "" + token_prefix: str = "" + api_base: str = "" + registered_at: str = "" + updated_at: str = "" + + +@runtime_checkable +class RegistrationStore(Protocol): + def load_agent_registration(self) -> AgentRegistration | None: ... + + def save_agent_registration(self, registration: AgentRegistration) -> None: ... + + +class EnvRegistrationStore: + def __init__(self, path: str | Path = DEFAULT_REGISTRATION_ENV_PATH) -> None: + self.path = Path(path or DEFAULT_REGISTRATION_ENV_PATH) + self._lock = threading.Lock() + + def load_agent_registration(self) -> AgentRegistration | None: + with self._lock: + try: + values = _read_env(self.path) + except FileNotFoundError: + return None + registration = AgentRegistration( + agent_id=values.get("OPENLINKER_AGENT_ID", ""), + agent_slug=values.get("OPENLINKER_AGENT_SLUG", ""), + agent_name=values.get("OPENLINKER_AGENT_NAME", ""), + agent_token=values.get("OPENLINKER_AGENT_TOKEN", ""), + token_id=values.get("OPENLINKER_AGENT_TOKEN_ID", ""), + token_prefix=values.get("OPENLINKER_AGENT_TOKEN_PREFIX", ""), + api_base=values.get("OPENLINKER_API_BASE", ""), + registered_at=values.get("OPENLINKER_REGISTERED_AT", ""), + updated_at=values.get("OPENLINKER_UPDATED_AT", ""), + ) + if not registration.agent_id and not registration.agent_token: + return None + return registration + + def save_agent_registration(self, registration: AgentRegistration) -> None: + if registration is None: + return + with self._lock: + try: + values = _read_env(self.path) + except FileNotFoundError: + values = {} + updates = { + "OPENLINKER_AGENT_ID": registration.agent_id, + "OPENLINKER_AGENT_SLUG": registration.agent_slug, + "OPENLINKER_AGENT_NAME": registration.agent_name, + "OPENLINKER_AGENT_TOKEN": registration.agent_token, + "OPENLINKER_AGENT_TOKEN_ID": registration.token_id, + "OPENLINKER_AGENT_TOKEN_PREFIX": registration.token_prefix, + "OPENLINKER_API_BASE": registration.api_base, + "OPENLINKER_REGISTERED_AT": registration.registered_at, + "OPENLINKER_UPDATED_AT": registration.updated_at, + } + values.update({key: value for key, value in updates.items() if value}) + for key, value in updates.items(): + if not value: + values.pop(key, None) + _write_env(self.path, values) + + +@dataclass +class EnsureAgentRequest: + slug: str = "" + name: str = "" + description: str = "" + endpoint_url: str = "" + endpoint_auth_header: str = "" + price_per_call_cents: int = 0 + tags: list[str] = field(default_factory=list) + skill_ids: list[str] = field(default_factory=list) + visibility: str = "" + connection_mode: str = "" + mcp_tool_name: str = "" + token_name: str = "" + token_scopes: list[str] = field(default_factory=list) + token_expires_in_minutes: int = 0 + policy: str = REGISTER_REUSE_EXISTING + user_token: str = "" + agent_token: str = "" + api_base: str = "" + store: RegistrationStore | None = None + env_path: str | Path = DEFAULT_REGISTRATION_ENV_PATH + + +async def ensure_agent( + request: EnsureAgentRequest, + *, + client: Client | None = None, +) -> AgentRegistration: + if not isinstance(request, EnsureAgentRequest): + raise TypeError("openlinker: ensure_agent requires EnsureAgentRequest") + req = replace(request) + req.tags = list(request.tags) + req.skill_ids = list(request.skill_ids) + req.token_scopes = list(request.token_scopes) + if req.policy not in _REGISTER_POLICIES: + raise ValueError(f"openlinker: unsupported registration policy {req.policy!r}") + + store = req.store or EnvRegistrationStore(req.env_path) + stored = store.load_agent_registration() + req.user_token = _first(req.user_token, os.getenv("OPENLINKER_USER_TOKEN")) + req.agent_token = _first( + req.agent_token, + os.getenv("OPENLINKER_AGENT_TOKEN"), + stored.agent_token if stored else "", + ) + req.api_base = _first( + req.api_base, + os.getenv("OPENLINKER_API_BASE"), + stored.api_base if stored else "", + client.base_url if client else "", + DEFAULT_API_BASE, + ) + req.slug = _first(req.slug, stored.agent_slug if stored else "") + req.name = _first(req.name, stored.agent_name if stored else "") + req.visibility = _first(req.visibility, "private") + req.connection_mode = _normalize_connection_mode(req.connection_mode) + req.token_name = _first(req.token_name, req.name, req.slug, "Python runtime worker") + if not req.token_scopes: + req.token_scopes = ["agent:pull", "agent:call"] + if not req.tags: + req.tags = ["agent", "runtime"] + + if stored and req.policy == REGISTER_REUSE_EXISTING and req.agent_token: + return _effective_registration(stored, req) + + owns_client = client is None + sdk = client or Client(req.api_base, user_token=req.user_token) + try: + if req.policy == REGISTER_VALIDATE_ONLY: + if not stored or not req.agent_token: + raise ValueError( + "openlinker: no stored Agent registration is available to validate" + ) + if not req.user_token and owns_client: + raise ValueError( + "openlinker: OPENLINKER_USER_TOKEN is required to validate registration" + ) + tokens = await sdk.list_agent_tokens( + {"agent_id": stored.agent_id, "limit": 50} + ) + valid = any( + item.id == stored.token_id + and item.status == "active_runtime" + and item.revoked_at is None + for item in tokens.items + ) + if not valid: + raise ValueError("openlinker: no valid stored Agent registration found") + return _effective_registration(stored, req) + + if stored and req.policy == REGISTER_ROTATE_TOKEN: + if not stored.agent_id: + raise ValueError("openlinker: stored Agent ID is required to rotate token") + token = await sdk.create_agent_token( + CreateAgentTokenRequest( + name=req.token_name, + agent_id=stored.agent_id, + scopes=req.token_scopes, + expires_in_minutes=req.token_expires_in_minutes, + ) + ) + if not token.plaintext_token: + raise ValueError("openlinker: platform did not return Agent Token plaintext") + registration = AgentRegistration( + agent_id=stored.agent_id, + agent_slug=_first(req.slug, stored.agent_slug), + agent_name=_first(req.name, stored.agent_name), + agent_token=token.plaintext_token, + token_id=token.id, + token_prefix=token.prefix, + api_base=req.api_base, + registered_at=stored.registered_at or _utc_now(), + updated_at=_utc_now(), + ) + store.save_agent_registration(registration) + return registration + + if stored is None and req.policy == REGISTER_REUSE_EXISTING and req.agent_token: + registered = await sdk.register_agent_via_token( + req.agent_token, _registration_request(req) + ) + registration = _registered_response(req, req.agent_token, registered) + store.save_agent_registration(registration) + return registration + + if not req.user_token and owns_client: + raise ValueError( + "openlinker: OPENLINKER_USER_TOKEN is required to create an Agent" + ) + + if not req.slug or not req.name: + raise ValueError("openlinker: Agent slug and name are required to create an Agent") + pending = await sdk.create_agent_token( + CreateAgentTokenRequest( + name=req.token_name, + scopes=req.token_scopes, + expires_in_minutes=req.token_expires_in_minutes, + ) + ) + if not pending.plaintext_token: + raise ValueError( + "openlinker: platform did not return pending Agent Token plaintext" + ) + registered = await sdk.register_agent_via_token( + pending.plaintext_token, _registration_request(req) + ) + registration = _registered_response(req, pending.plaintext_token, registered) + store.save_agent_registration(registration) + return registration + finally: + if owns_client: + await sdk.aclose() + + +def _registration_request(req: EnsureAgentRequest) -> RegisterAgentViaTokenRequest: + return RegisterAgentViaTokenRequest( + slug=req.slug or None, + name=req.name, + description=req.description or None, + endpoint_url=req.endpoint_url or None, + endpoint_auth_header=req.endpoint_auth_header or None, + price_per_call_cents=req.price_per_call_cents, + tags=req.tags, + ability_tags=req.tags, + skill_ids=req.skill_ids, + visibility=req.visibility, + connection_mode=req.connection_mode, + mcp_tool_name=req.mcp_tool_name or None, + ) + + +def _registered_response(req, agent_token, response) -> AgentRegistration: + now = _utc_now() + return AgentRegistration( + agent_id=response.agent.id, + agent_slug=response.agent.slug, + agent_name=response.agent.name, + agent_token=agent_token, + token_id=response.agent_token.id, + token_prefix=response.agent_token.prefix, + api_base=req.api_base, + registered_at=now, + updated_at=now, + ) + + +def _effective_registration( + stored: AgentRegistration, + req: EnsureAgentRequest, +) -> AgentRegistration: + return replace( + stored, + agent_slug=_first(req.slug, stored.agent_slug), + agent_name=_first(req.name, stored.agent_name), + agent_token=_first(req.agent_token, stored.agent_token), + api_base=_first(req.api_base, stored.api_base), + ) + + +def _normalize_connection_mode(value: str) -> str: + normalized = value.strip().lower() + if normalized in {"", "runtime", "runtime_ws", "runtime_pull", "agent_node"}: + return "runtime" + return value.strip() + + +def _first(*values: str | None) -> str: + for value in values: + if value and value.strip(): + return value.strip() + return "" + + +def _utc_now() -> str: + from datetime import datetime, timezone + + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") + + +def _read_env(path: Path) -> dict[str, str]: + values: dict[str, str] = {} + for raw in path.read_text(encoding="utf-8").splitlines(): + line = raw.strip() + if line.startswith("export "): + line = line[7:].strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + values[key.strip()] = _unquote(value.strip()) + return values + + +def _write_env(path: Path, values: dict[str, str]) -> None: + path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + try: + existing = path.read_text(encoding="utf-8").splitlines() + except FileNotFoundError: + existing = [] + managed = set(_ENV_KEYS) + seen: set[str] = set() + output: list[str] = [] + for raw in existing: + line = raw.strip() + candidate = line[7:].strip() if line.startswith("export ") else line + key = candidate.split("=", 1)[0].strip() if "=" in candidate else "" + if key not in managed: + output.append(raw) + continue + seen.add(key) + if values.get(key): + output.append(f"{key}={json.dumps(values[key])}") + for key in _ENV_KEYS: + if key not in seen and values.get(key): + output.append(f"{key}={json.dumps(values[key])}") + content = "\n".join(output) + ("\n" if output else "") + descriptor, temporary = tempfile.mkstemp( + prefix=".openlinker-registration-", + suffix=".tmp", + dir=path.parent, + ) + try: + os.fchmod(descriptor, 0o600) + with os.fdopen(descriptor, "w", encoding="utf-8", closefd=True) as stream: + descriptor = -1 + stream.write(content) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + os.chmod(path, 0o600) + finally: + if descriptor >= 0: + os.close(descriptor) + try: + os.unlink(temporary) + except FileNotFoundError: + pass + + +def _unquote(value: str) -> str: + if len(value) >= 2 and value[0] == value[-1] == '"': + try: + decoded = json.loads(value) + if isinstance(decoded, str): + return decoded + except ValueError: + pass + if len(value) >= 2 and value[0] == value[-1] == "'": + return value[1:-1] + return value + + +EnsureAgent = ensure_agent +NewEnvRegistrationStore = EnvRegistrationStore diff --git a/src/openlinker/runtime/worker.py b/src/openlinker/runtime/worker.py index b351e87..2c01b2f 100644 --- a/src/openlinker/runtime/worker.py +++ b/src/openlinker/runtime/worker.py @@ -85,6 +85,7 @@ "RUNTIME_SPOOL_CORRUPT", } _LEASE_TERMINAL_CODES = {"STALE_LEASE", "LEASE_EXPIRED", "RUN_ALREADY_TERMINAL"} +_ASSIGNMENT_TERMINAL_CODES = _LEASE_TERMINAL_CODES | {"RUN_CANCEL_REQUESTED"} class _RuntimeStoppedBeforeDrain(RuntimeError): @@ -277,7 +278,8 @@ def __init__( self._active: dict[str, _ActiveAttempt] = {} self._attempt_locks: dict[str, asyncio.Lock] = {} self._spool_permissions: dict[str, tuple[bool, bool]] = {} - self._cancellations: set[str] = set() + self._cancellations: set[tuple[str, str]] = set() + self._terminal_attempts: set[str] = set() self._background: set[asyncio.Task[Any]] = set() self._websocket_probe_task: asyncio.Task[Any] | None = None self._claim_switch_lock = asyncio.Lock() @@ -715,6 +717,9 @@ async def claim() -> ClaimedAssignment | None: async def _handle_assignment(self, claimed: ClaimedAssignment) -> None: assignment = claimed.assignment + attempt_id = assignment.attempt_identity.attempt_id + if attempt_id in self._terminal_attempts: + return local = self._local_identity(assignment.attempt_identity) record = AssignmentRecord( identity=local, @@ -735,15 +740,20 @@ async def _handle_assignment(self, claimed: ClaimedAssignment) -> None: if record.state == ASSIGNMENT_RECEIVED: record = store.advance_assignment(local.assignment_message_id, ASSIGNMENT_ACK_SENT) if record.state == ASSIGNMENT_ACK_SENT: - confirmation = await self._retry_call( + terminal, confirmation = await self._retry_assignment_operation( + assignment.attempt_identity, lambda: self._transport_required().ack_assignment( {"attempt_identity": assignment.attempt_identity.to_dict()}, delivery_id=claimed.delivery_id, - ) + ), ) + if terminal: + return self._validate_confirmation(assignment.attempt_identity, confirmation) + if attempt_id in self._terminal_attempts: + return record = store.advance_assignment(local.assignment_message_id, ASSIGNMENT_CONFIRMED) - self._spool_permissions[assignment.attempt_identity.attempt_id] = (True, True) + self._spool_permissions[attempt_id] = (True, True) await self._start_confirmed_attempt( record, parse_datetime(confirmation["lease_expires_at"]) ) @@ -753,14 +763,17 @@ async def _handle_assignment(self, claimed: ClaimedAssignment) -> None: await self._start_confirmed_attempt(record, None) return if record.state in {ASSIGNMENT_STARTED, ASSIGNMENT_FINISHED}: - confirmation = await self._retry_call( + terminal, confirmation = await self._retry_assignment_operation( + assignment.attempt_identity, lambda: self._transport_required().ack_assignment( {"attempt_identity": assignment.attempt_identity.to_dict()}, delivery_id=claimed.delivery_id, - ) + ), ) + if terminal: + return self._validate_confirmation(assignment.attempt_identity, confirmation) - active = self._active.get(assignment.attempt_identity.attempt_id) + active = self._active.get(attempt_id) if active is not None: active.lease_expires_at = parse_datetime(confirmation["lease_expires_at"]) @@ -772,7 +785,8 @@ async def _reject_assignment(self, record: AssignmentRecord, delivery_id: str) - ) capacity, inflight = self._capacity_snapshot() reason = "NODE_DRAINING" if self._draining else "NODE_AT_CAPACITY" - response = await self._retry_call( + terminal, response = await self._retry_assignment_operation( + record.identity.attempt, lambda: self._transport_required().reject_assignment( { "attempt_identity": record.identity.attempt.to_dict(), @@ -781,8 +795,10 @@ async def _reject_assignment(self, record: AssignmentRecord, delivery_id: str) - "inflight": inflight, }, delivery_id=delivery_id, - ) + ), ) + if terminal: + return _require_response_keys( response, required={"attempt_identity", "outcome", "dispatch_state"}, @@ -795,13 +811,49 @@ async def _reject_assignment(self, record: AssignmentRecord, delivery_id: str) - store.advance_assignment(record.identity.assignment_message_id, ASSIGNMENT_REJECTED) store.delete_assignment(record.identity.assignment_message_id) + async def _retry_assignment_operation( + self, + identity: RuntimeAttemptIdentity, + operation: Callable[[], Awaitable[Any]], + ) -> tuple[bool, Any]: + if identity.attempt_id in self._terminal_attempts: + return True, None + try: + value = await self._retry_call(operation) + except RuntimeRemoteError as exc: + if await self._converge_assignment_terminal_error(identity, exc): + return True, None + raise + if identity.attempt_id in self._terminal_attempts: + return True, None + return False, value + + async def _converge_assignment_terminal_error( + self, + identity: RuntimeAttemptIdentity, + error: RuntimeRemoteError, + ) -> bool: + if error.code not in _ASSIGNMENT_TERMINAL_CODES: + return False + try: + record = self._store_required().assignment_for_attempt(identity.attempt_id) + except RuntimeStoreError as exc: + if str(exc) == "assignment not found": + self._terminal_attempts.add(identity.attempt_id) + return True + raise + if record.identity.attempt != identity: + raise RuntimeProtocolError("Runtime assignment terminal identity mismatch") + await self._revoke_attempt(record) + return True + async def _start_confirmed_attempt( self, record: AssignmentRecord, lease_expires_at: datetime | None, ) -> None: attempt_id = record.identity.attempt.attempt_id - if attempt_id in self._active: + if attempt_id in self._active or attempt_id in self._terminal_attempts: return if record.state != ASSIGNMENT_CONFIRMED: raise RuntimeStoreError("handler requires a confirmed assignment") @@ -1073,10 +1125,14 @@ async def _handle_command(self, command: dict[str, Any]) -> None: parse_datetime(payload["deadline_at"]) except (TypeError, ValueError) as exc: raise RuntimeProtocolError("Runtime cancellation deadline is invalid") from exc - if cancellation_id in self._cancellations: + identity = RuntimeAttemptIdentity.from_dict( + _protocol_object(payload.get("attempt_identity"), "cancel Attempt") + ) + cancellation_key = (cancellation_id, identity.attempt_id) + if cancellation_key in self._cancellations: return - self._cancellations.add(cancellation_id) - self._spawn(self._handle_cancel(payload)) + self._cancellations.add(cancellation_key) + self._spawn(self._handle_cancel(payload, cancellation_key)) elif command_type == "runtime.drain": validate_runtime_drain_payload(payload) self._draining = True @@ -1107,46 +1163,111 @@ async def _handle_command(self, command: dict[str, Any]) -> None: else: raise RuntimeProtocolError(f"unknown Runtime command {command_type!r}") - async def _handle_cancel(self, payload: dict[str, Any]) -> None: + async def _handle_cancel( + self, + payload: dict[str, Any], + cancellation_key: tuple[str, str], + ) -> None: cancellation_id = str(payload.get("cancellation_id", "")) _canonical_protocol_uuid(cancellation_id, "cancellation_id") identity = RuntimeAttemptIdentity.from_dict( _protocol_object(payload.get("attempt_identity"), "cancel Attempt") ) try: - record = self._store_required().assignment_for_attempt(identity.attempt_id) - except RuntimeStoreError as exc: - if str(exc) != "assignment not found": - raise - await self._ack_cancel(payload, "failed", "ATTEMPT_IDENTITY_MISMATCH") - return - if record.identity.attempt != identity: - await self._ack_cancel(payload, "failed", "ATTEMPT_IDENTITY_MISMATCH") - return - await self._ack_cancel(payload, "stopping", "") - active = self._active.get(identity.attempt_id) - if active is not None: - active.cancel_event.set() - if active.task is not None: - active.task.cancel() - async with self._attempt_lock(identity.attempt_id): - self._spool_permissions.pop(identity.attempt_id, None) - if active is not None: deadline = parse_datetime(payload["deadline_at"]) - timeout = max(0.0, (deadline - datetime.now(timezone.utc)).total_seconds()) try: + record = self._store_required().assignment_for_attempt(identity.attempt_id) + except RuntimeStoreError as exc: + if str(exc) != "assignment not found": + raise + await self._ack_cancel_best_effort( + payload, + "failed", + "ATTEMPT_IDENTITY_MISMATCH", + deadline, + ) + return + if record.identity.attempt != identity: + await self._ack_cancel_best_effort( + payload, + "failed", + "ATTEMPT_IDENTITY_MISMATCH", + deadline, + ) + return + + self._terminal_attempts.add(identity.attempt_id) + now = datetime.now(timezone.utc) + initial_remaining = max(0.0, (deadline - now).total_seconds()) + stopping_budget = max(0.001, min(2.0, initial_remaining / 2)) + stopping_deadline = min(deadline, now + timedelta(seconds=stopping_budget)) + stopping_ack = asyncio.create_task( + self._ack_cancel(payload, "stopping", "", deadline=stopping_deadline) + ) + + active = self._active.get(identity.attempt_id) + if active is not None: + active.cancel_event.set() if active.task is not None: - await asyncio.wait_for( - asyncio.gather(active.task, return_exceptions=True), timeout=timeout - ) - except asyncio.TimeoutError: - await self._ack_cancel(payload, "failed", "CANCEL_DEADLINE_EXCEEDED") + active.task.cancel() + async with self._attempt_lock(identity.attempt_id): + self._spool_permissions.pop(identity.attempt_id, None) + + handler_stopped = True + if active is not None and active.task is not None: + timeout = max(0.0, (deadline - datetime.now(timezone.utc)).total_seconds()) + done, _ = await asyncio.wait({active.task}, timeout=timeout) + handler_stopped = active.task in done + + try: + await stopping_ack + except Exception as exc: + self.logger.warning( + "Runtime cancel stopping ACK was not confirmed: %s", exc + ) + + if not handler_stopped: + await self._ack_cancel_best_effort( + payload, + "failed", + "CANCEL_DEADLINE_EXCEEDED", + deadline, + ) return - await self._ack_cancel(payload, "stopped", "") - await self._finalize_revoked_attempt(record) - self.logger.debug("Runtime cancellation %s stopped", cancellation_id) + if not await self._ack_cancel_best_effort( + payload, + "stopped", + "", + deadline, + ): + return + await self._finalize_revoked_attempt(record) + self.logger.debug("Runtime cancellation %s stopped", cancellation_id) + finally: + self._cancellations.discard(cancellation_key) + + async def _ack_cancel_best_effort( + self, + payload: dict[str, Any], + state: str, + error_code: str, + deadline: datetime, + ) -> bool: + try: + await self._ack_cancel(payload, state, error_code, deadline=deadline) + return True + except Exception as exc: + self.logger.warning("Runtime cancel %s ACK was not confirmed: %s", state, exc) + return False - async def _ack_cancel(self, payload: dict[str, Any], state: str, error_code: str) -> None: + async def _ack_cancel( + self, + payload: dict[str, Any], + state: str, + error_code: str, + *, + deadline: datetime | None = None, + ) -> None: request = { "cancellation_id": payload["cancellation_id"], "attempt_identity": payload["attempt_identity"], @@ -1154,9 +1275,18 @@ async def _ack_cancel(self, payload: dict[str, Any], state: str, error_code: str } if error_code: request["error_code"] = error_code + deadline = deadline or parse_datetime(payload["deadline_at"]) + + async def send() -> dict[str, Any]: + remaining = max(0.001, (deadline - datetime.now(timezone.utc)).total_seconds()) + return await asyncio.wait_for( + self._transport_required().ack_cancel(request), + timeout=remaining, + ) + response = await self._retry_call( - lambda: self._transport_required().ack_cancel(request), - deadline=parse_datetime(payload["deadline_at"]), + send, + deadline=deadline, ) if response.get("cancellation_id") != payload["cancellation_id"]: raise RuntimeProtocolError("Runtime cancellation ACK identity mismatch") @@ -1573,6 +1703,7 @@ async def _websocket_probe_loop(self) -> None: async def _revoke_attempt(self, record: AssignmentRecord) -> None: attempt_id = record.identity.attempt.attempt_id + self._terminal_attempts.add(attempt_id) active = self._active.get(attempt_id) if active is not None: active.cancel_event.set() diff --git a/src/openlinker/types.py b/src/openlinker/types.py index bcfe943..3e812a8 100644 --- a/src/openlinker/types.py +++ b/src/openlinker/types.py @@ -163,6 +163,7 @@ class TaskCallbackConfig(Model): class RunAgentRequest(Model): agent_id: str = "" input: Any = None + idempotency_key: str | None = None metadata: Any = None a2a_context: RunA2AContext | None = None task_callback: TaskCallbackConfig | None = None @@ -245,19 +246,71 @@ class RunEventResponse(Model): created_at: str = "" +@dataclass +class RunEventPageMeta(Model): + requested_after_sequence: int = 0 + effective_after_sequence: int = 0 + retained_through_sequence: int = 0 + earliest_available_sequence: int | None = None + latest_available_sequence: int | None = None + retention_gap: bool = False + terminal: bool = False + stream_complete: bool = False + + @dataclass class ListRunEventsResponse(Model): + items: list[RunEventResponse] = jfield(default_factory=list) + meta: RunEventPageMeta = jfield(default_factory=RunEventPageMeta) + # Deprecated compatibility view. Core returns `items`; keep `events` for + # callers of older Python SDK versions. events: list[RunEventResponse] = jfield(default_factory=list) + @classmethod + def from_dict(cls, data: dict[str, Any] | None): + result = super().from_dict(data) + if result.items and not result.events: + result.events = result.items + elif result.events and not result.items: + result.items = result.events + return result + + +@dataclass +class RunSkillRef(Model): + id: str = "" + name: str = "" + @dataclass class RunChildResponse(Model): child_run_id: str = "" + parent_run_id: str = "" + caller_agent_id: str = "" + caller_agent_slug: str = "" + caller_agent_name: str = "" + caller_agent_tags: list[str] = jfield(default_factory=list) + caller_skills: list[RunSkillRef] = jfield(default_factory=list) + target_agent_id: str = "" + target_agent_slug: str = "" + target_agent_name: str = "" + target_agent_tags: list[str] = jfield(default_factory=list) + target_skills: list[RunSkillRef] = jfield(default_factory=list) + reason: str = "" status: str = "" + cost_cents: int = 0 + duration_ms: int | None = None + started_at: str = "" + finished_at: str | None = None + source: str = "" + billing_mode: str = "" + a2a_context: Any = None + children: list[RunChildResponse] = jfield(default_factory=list) @dataclass class ListRunChildrenResponse(Model): + parent_run_id: str = "" items: list[RunChildResponse] = jfield(default_factory=list) @@ -439,3 +492,25 @@ class AgentTokenListResponse(Model): sort_by: str = "" sort_dir: str = "" has_more: bool = False + + +@dataclass +class RegisterAgentViaTokenRequest(Model): + slug: str | None = None + name: str = "" + description: str | None = None + endpoint_url: str | None = None + endpoint_auth_header: str | None = None + price_per_call_cents: int = 0 + tags: list[str] = jfield(default_factory=list) + ability_tags: list[str] = jfield(default_factory=list) + skill_ids: list[str] = jfield(default_factory=list) + visibility: str | None = None + connection_mode: str | None = None + mcp_tool_name: str | None = None + + +@dataclass +class RegisterAgentViaTokenResponse(Model): + agent: AgentResponse = jfield(default_factory=AgentResponse) + agent_token: AgentTokenResponse = jfield(default_factory=AgentTokenResponse) diff --git a/tests/test_client.py b/tests/test_client.py index cf30354..26c8dd9 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -51,8 +51,10 @@ async def test_run_agent_encodes_request_body(): async def handler(request: httpx.Request) -> httpx.Response: seen["path"] = request.url.path seen["body"] = json.loads(request.content) + seen["idempotency_key"] = request.headers.get("Idempotency-Key") return httpx.Response( 200, + headers={"Idempotency-Replayed": "true"}, json={ "run_id": "run-1", "agent_id": "agent-1", @@ -83,6 +85,7 @@ async def handler(request: httpx.Request) -> httpx.Response: openlinker_client.RunAgentRequest( agent_id="agent-1", input={"query": "hello"}, + idempotency_key="logical-run-1", task_callback=openlinker_client.TaskCallbackConfig( url="https://caller.example.com/events", token="caller-token", @@ -102,7 +105,187 @@ async def handler(request: httpx.Request) -> httpx.Response: assert seen["body"]["agent_id"] == "agent-1" assert "agent_connection_mode" not in seen["body"] assert "runtime_transport" not in seen["body"] + assert "idempotency_key" not in seen["body"] assert seen["body"]["task_callback"]["secret"] == "caller-secret" + assert seen["idempotency_key"] == "logical-run-1" + assert resp.replayed is True + + +@pytest.mark.asyncio +async def test_run_agent_generates_unique_idempotency_keys_when_omitted(): + keys = [] + + async def handler(request: httpx.Request) -> httpx.Response: + keys.append(request.headers.get("Idempotency-Key")) + return httpx.Response(201, json={"run_id": f"run-{len(keys)}", "status": "running"}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as hc: + sdk = openlinker_client.Client("https://api.example.test", http_client=hc) + await sdk.run_agent({"agent_id": "agent-1", "input": {}}) + await sdk.run_agent({"agent_id": "agent-1", "input": {}}) + + assert len(keys) == 2 + assert keys[0] != keys[1] + assert all(key and len(key) == 64 for key in keys) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key", ["", "line\nbreak", "é", "x" * 256]) +async def test_run_agent_rejects_invalid_idempotency_keys_before_network(key): + calls = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response(500) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as hc: + sdk = openlinker_client.Client("https://api.example.test", http_client=hc) + with pytest.raises(ValueError) as raised: + await sdk.run_agent( + {"agent_id": "agent-1", "input": {}, "idempotency_key": key} + ) + + assert calls == 0 + if key: + assert key not in str(raised.value) + + +@pytest.mark.asyncio +async def test_list_run_events_decodes_items_meta_and_legacy_events(): + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "items": [ + { + "event_id": "event-1", + "run_id": "run-1", + "sequence": 4, + "event_type": "run.completed", + "payload": {"ok": True}, + "created_at": "2026-07-18T00:00:00Z", + } + ], + "meta": { + "requested_after_sequence": 0, + "effective_after_sequence": 3, + "retained_through_sequence": 3, + "earliest_available_sequence": 4, + "latest_available_sequence": 4, + "retention_gap": True, + "terminal": True, + "stream_complete": True, + }, + }, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as hc: + sdk = openlinker_client.Client("https://api.example.test", http_client=hc) + response = await sdk.list_run_events("run-1") + + assert response.items[0].event_id == "event-1" + assert response.events is response.items + assert response.meta.retention_gap is True + assert response.meta.stream_complete is True + + +@pytest.mark.asyncio +async def test_list_run_children_decodes_complete_recursive_dto(): + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "parent_run_id": "run-parent", + "items": [ + { + "child_run_id": "run-child", + "parent_run_id": "run-parent", + "caller_agent_id": "agent-caller", + "caller_agent_slug": "caller", + "caller_agent_name": "Caller", + "caller_agent_tags": ["orchestrator"], + "caller_skills": [{"id": "skill-plan", "name": "Planning"}], + "target_agent_id": "agent-target", + "target_agent_slug": "target", + "target_agent_name": "Target", + "target_agent_tags": ["research"], + "target_skills": [ + {"id": "skill-research", "name": "Research"} + ], + "reason": "delegate research", + "status": "success", + "cost_cents": 4, + "duration_ms": 12, + "started_at": "2026-07-18T00:00:00Z", + "finished_at": "2026-07-18T00:00:01Z", + "source": "api", + "billing_mode": "caller", + "a2a_context": {"trace_id": "trace-1"}, + "children": [], + } + ], + }, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as hc: + sdk = openlinker_client.Client("https://api.example.test", http_client=hc) + response = await sdk.list_run_children("run-parent") + + assert response.parent_run_id == "run-parent" + assert response.items[0].target_skills[0].name == "Research" + assert response.items[0].a2a_context["trace_id"] == "trace-1" + + +@pytest.mark.asyncio +async def test_register_agent_via_token_uses_agent_auth_and_runtime_defaults(): + seen = {} + + async def handler(request: httpx.Request) -> httpx.Response: + seen["path"] = request.url.path + seen["auth"] = request.headers.get("Authorization") + seen["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "agent": {"id": "agent-1", "slug": "demo", "name": "Demo"}, + "agent_token": { + "id": "token-1", + "prefix": "ol_agent_demo", + "status": "active_runtime", + }, + }, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as hc: + sdk = openlinker_client.Client( + "https://api.example.test", http_client=hc, user_token="ol_user_creator" + ) + response = await sdk.register_agent_via_token( + " ol_agent_pending ", {"name": "Demo", "tags": ["runtime"]} + ) + + assert response.agent.id == "agent-1" + assert response.agent_token.id == "token-1" + assert seen["path"] == "/api/v1/agent-registration/agents" + assert seen["auth"] == "Bearer ol_agent_pending" + assert seen["body"]["visibility"] == "private" + assert seen["body"]["connection_mode"] == "runtime" + + +@pytest.mark.asyncio +async def test_client_rejects_oversized_response_before_reading_body(): + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"Content-Length": str(openlinker_client.MAX_RESPONSE_BODY_BYTES + 1)}, + content=b"{}", + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as hc: + sdk = openlinker_client.Client("https://api.example.test", http_client=hc) + with pytest.raises(ValueError, match="response body exceeds"): + await sdk.list_agents() @pytest.mark.asyncio diff --git a/tests/test_client_contract.py b/tests/test_client_contract.py new file mode 100644 index 0000000..2450824 --- /dev/null +++ b/tests/test_client_contract.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import json +from dataclasses import fields +from pathlib import Path + +from openlinker.client import Client +from openlinker.types import ListRunChildrenResponse, ListRunEventsResponse, RunResponse + + +ROOT = Path(__file__).parents[1] + + +def test_core_client_contract_maps_to_methods_and_response_types(): + contract = _contract("core-client.v1.json") + assert contract["package"] == "openlinker" + forbidden = contract["rules"]["forbidden_path_prefixes"] + for endpoint in contract["endpoints"]: + assert callable(getattr(Client, endpoint["client_method"], None)) + assert endpoint["path"].startswith("/api/v1/") + assert not any(endpoint["path"].startswith(prefix) for prefix in forbidden) + + run_fields = {item.name for item in fields(RunResponse)} + for endpoint in contract["endpoints"]: + if endpoint["path"] in {"/api/v1/run", "/api/v1/runs"}: + assert endpoint["required_headers"] == ["Idempotency-Key"] + assert endpoint["success_statuses"] == [200, 201, 202] + assert set(endpoint["response_fields"]).issubset(run_fields) + + event_fields = {item.name for item in fields(ListRunEventsResponse)} + child_fields = {item.name for item in fields(ListRunChildrenResponse)} + assert {"items", "meta"}.issubset(event_fields) + assert {"parent_run_id", "items"}.issubset(child_fields) + + +def test_core_registration_contract_maps_to_client_methods(): + contract = _contract("core-registration.v1.json") + assert contract["scope"] == "core-registration" + assert len(contract["endpoints"]) == 9 + for endpoint in contract["endpoints"]: + assert callable(getattr(Client, endpoint["client_method"], None)) + assert endpoint["auth"] in {"user_token", "agent_token"} + + +def _contract(name: str): + return json.loads((ROOT / "contracts" / name).read_text(encoding="utf-8")) diff --git a/tests/test_registration.py b/tests/test_registration.py new file mode 100644 index 0000000..890fe91 --- /dev/null +++ b/tests/test_registration.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import json +import stat +from dataclasses import replace + +import httpx +import pytest + +from openlinker.client import Client +from openlinker.registration import ( + AgentRegistration, + EnsureAgentRequest, + EnvRegistrationStore, + REGISTER_REUSE_EXISTING, + ensure_agent, +) + + +class MemoryRegistrationStore: + def __init__(self, registration: AgentRegistration | None = None) -> None: + self.registration = registration + + def load_agent_registration(self) -> AgentRegistration | None: + return replace(self.registration) if self.registration else None + + def save_agent_registration(self, registration: AgentRegistration) -> None: + self.registration = replace(registration) + + +@pytest.mark.asyncio +async def test_ensure_agent_registers_pending_token_then_reuses_store(monkeypatch): + for key in ("OPENLINKER_USER_TOKEN", "OPENLINKER_AGENT_TOKEN", "OPENLINKER_API_BASE"): + monkeypatch.delenv(key, raising=False) + calls = [] + + async def handler(request: httpx.Request) -> httpx.Response: + calls.append((request.url.path, request.headers.get("Authorization"), json.loads(request.content))) + if request.url.path == "/api/v1/creator/agent-tokens": + return httpx.Response( + 200, + json={ + "id": "token-1", + "prefix": "ol_agent_demo", + "status": "pending_registration", + "plaintext_token": "ol_agent_plaintext", + }, + ) + if request.url.path == "/api/v1/agent-registration/agents": + return httpx.Response( + 200, + json={ + "agent": {"id": "agent-1", "slug": "demo", "name": "Demo"}, + "agent_token": { + "id": "token-1", + "prefix": "ol_agent_demo", + "status": "active_runtime", + }, + }, + ) + raise AssertionError(f"unexpected request {request.url.path}") + + store = MemoryRegistrationStore() + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = Client( + "https://api.example.test", + http_client=http_client, + user_token="ol_user_creator", + ) + request = EnsureAgentRequest( + slug="demo", + name="Demo", + api_base="https://api.example.test", + user_token="ol_user_creator", + store=store, + ) + first = await ensure_agent(request, client=client) + second = await ensure_agent( + EnsureAgentRequest(policy=REGISTER_REUSE_EXISTING, store=store), + client=client, + ) + + assert first.agent_id == "agent-1" + assert first.agent_token == "ol_agent_plaintext" + assert second == first + assert [item[0] for item in calls] == [ + "/api/v1/creator/agent-tokens", + "/api/v1/agent-registration/agents", + ] + assert calls[0][1] == "Bearer ol_user_creator" + assert calls[1][1] == "Bearer ol_agent_plaintext" + assert calls[1][2]["connection_mode"] == "runtime" + + +def test_env_registration_store_preserves_unrelated_values_and_forces_private_mode(tmp_path): + path = tmp_path / ".env" + path.write_text("UNRELATED=value\nOPENLINKER_AGENT_TOKEN=old\n", encoding="utf-8") + path.chmod(0o644) + store = EnvRegistrationStore(path) + + store.save_agent_registration( + AgentRegistration( + agent_id="agent-1", + agent_slug="demo", + agent_name="Demo", + agent_token="ol_agent_secret", + token_id="token-1", + api_base="https://api.example.test", + updated_at="2026-07-18T00:00:00Z", + ) + ) + + raw = path.read_text(encoding="utf-8") + loaded = store.load_agent_registration() + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + assert "UNRELATED=value" in raw + assert "OPENLINKER_AGENT_TOKEN=old" not in raw + assert loaded is not None + assert loaded.agent_id == "agent-1" + assert loaded.agent_token == "ol_agent_secret" diff --git a/tests/test_runtime.py b/tests/test_runtime.py index cc6f95f..cdc8132 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -12,6 +12,7 @@ from openlinker import runtime import openlinker.runtime.worker as runtime_worker_module +from openlinker.runtime.store import AssignmentRecord from openlinker.runtime.transport import ( ClaimedAssignment, RuntimeDiscoveryConnection, @@ -85,6 +86,8 @@ def __init__(self, *, kind: str = "pull") -> None: self.create_calls = 0 self.session_conflicts = 0 self.ack_failures = 0 + self.ack_error: Exception | None = None + self.reject_error: Exception | None = None self.event_failures = 0 self.result_failures = 0 self.event_upload_available = True @@ -169,6 +172,8 @@ async def ack_assignment( self.ack_attempts.append(request) self.ack_entered.set() await self.ack_release.wait() + if self.ack_error is not None: + raise self.ack_error if len(self.ack_attempts) <= self.ack_failures: raise ConnectionError("assignment ACK response was lost") return { @@ -181,6 +186,8 @@ async def reject_assignment( self, request: dict[str, Any], *, delivery_id: str = "" ) -> dict[str, Any]: del delivery_id + if self.reject_error is not None: + raise self.reject_error return { "attempt_identity": request["attempt_identity"], "outcome": "offer_rejected", @@ -400,6 +407,94 @@ async def handler(context: runtime.RuntimeContext) -> dict[str, Any]: assert transport.result_attempts[0]["result_id"] == transport.result_attempts[1]["result_id"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "code", + ["RUN_CANCEL_REQUESTED", "STALE_LEASE", "LEASE_EXPIRED", "RUN_ALREADY_TERMINAL"], +) +async def test_assignment_confirmation_terminal_errors_converge_one_attempt(code): + store = runtime.MemoryRuntimeStore() + transport = FakeTransport() + offered = assignment(store) + transport.ack_error = runtime.RuntimeRemoteError(code, "terminal", status_code=409) + handler_calls = 0 + + async def handler(context: runtime.RuntimeContext) -> dict[str, Any]: + nonlocal handler_calls + del context + handler_calls += 1 + return {} + + worker = make_worker(store, transport, handler) + await worker._handle_assignment(ClaimedAssignment(offered, "delivery-terminal")) + + assert handler_calls == 0 + assert store.assignments() == [] + assert offered.attempt_identity.attempt_id in worker._terminal_attempts + assert worker._fatal.empty() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("state", ["started", "finished"]) +async def test_duplicate_owned_assignment_terminal_reack_is_benign(state): + store = runtime.MemoryRuntimeStore() + transport = FakeTransport() + offered = assignment(store) + worker = make_worker(store, transport, lambda _context: {}) + local = worker._local_identity(offered.attempt_identity) + record = store.create_assignment( + AssignmentRecord( + identity=local, + input=offered.input, + metadata=offered.metadata, + node_envelope=offered.node_envelope, + agent_invocation_token=offered.agent_invocation_token, + offer_expires_at=offered.offer_expires_at, + attempt_deadline_at=offered.attempt_deadline_at, + run_deadline_at=offered.run_deadline_at, + ) + ) + for next_state in ("ack_sent", "confirmed", "started"): + record = store.advance_assignment(record.identity.assignment_message_id, next_state) + if state == "finished": + store.store_result( + offered.attempt_identity.attempt_id, + { + "attempt_identity": offered.attempt_identity.to_dict(), + "result_id": CANCELLATION_ID, + "duration_ms": 1, + "final_client_event_seq": 0, + "status": "success", + "output": {}, + }, + ) + transport.ack_error = runtime.RuntimeRemoteError( + "RUN_ALREADY_TERMINAL", "terminal", status_code=409 + ) + + await worker._handle_assignment(ClaimedAssignment(offered, "delivery-duplicate")) + + assert store.assignments() == [] + assert offered.attempt_identity.attempt_id in worker._terminal_attempts + + +@pytest.mark.asyncio +async def test_assignment_rejection_terminal_error_is_benign(): + store = runtime.MemoryRuntimeStore() + transport = FakeTransport() + offered = assignment(store) + transport.reject_error = runtime.RuntimeRemoteError( + "RUN_CANCEL_REQUESTED", "canceled before rejection", status_code=409 + ) + worker = make_worker(store, transport, lambda _context: {}) + worker._draining = True + + await worker._handle_assignment(ClaimedAssignment(offered, "delivery-reject")) + + assert store.assignments() == [] + assert offered.attempt_identity.attempt_id in worker._terminal_attempts + + @pytest.mark.asyncio async def test_finished_handler_keeps_lease_and_capacity_until_result_ack(): store = runtime.MemoryRuntimeStore() @@ -797,7 +892,76 @@ async def handler(context: runtime.RuntimeContext) -> dict[str, Any]: break await asyncio.sleep(0.005) assert transport.cancel_states == ["stopping", "stopped"] + for _ in range(100): + if not worker._cancellations: + break + await asyncio.sleep(0.005) + assert worker._cancellations == set() + finally: + await worker.stop() + await running + + +@pytest.mark.asyncio +async def test_cancel_deadline_failure_does_not_leak_dedupe_or_stop_worker(): + class SlowStoppingAckTransport(FakeTransport): + async def ack_cancel(self, request: dict[str, Any]) -> dict[str, Any]: + self.cancel_states.append(request["cancel_state"]) + if request["cancel_state"] == "stopping": + await asyncio.Event().wait() + return { + "cancellation_id": request["cancellation_id"], + "cancel_state": request["cancel_state"], + "updated_at": datetime.now(timezone.utc).isoformat(), + } + + store = runtime.MemoryRuntimeStore() + transport = SlowStoppingAckTransport() + offered = assignment(store) + transport.assignment = offered + handler_started = asyncio.Event() + ignored_cancel = asyncio.Event() + release_handler = asyncio.Event() + + async def handler(context: runtime.RuntimeContext) -> dict[str, Any]: + del context + handler_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + ignored_cancel.set() + await release_handler.wait() + return {"late": True} + + worker = make_worker(store, transport, handler) + running = asyncio.create_task(worker.run()) + try: + await asyncio.wait_for(handler_started.wait(), timeout=1) + await transport.commands.put( + { + "type": "run.cancel", + "payload": { + "cancellation_id": CANCELLATION_ID, + "attempt_identity": offered.attempt_identity.to_dict(), + "deadline_at": ( + datetime.now(timezone.utc) + timedelta(seconds=0.05) + ).isoformat(), + "reason_code": "caller_requested", + }, + } + ) + await asyncio.wait_for(ignored_cancel.wait(), timeout=1) + for _ in range(200): + if transport.cancel_states == ["stopping", "failed"]: + break + await asyncio.sleep(0.005) + + assert transport.cancel_states == ["stopping", "failed"] + assert worker._cancellations == set() + assert not running.done() + assert store.assignments() finally: + release_handler.set() await worker.stop() await running