diff --git a/src/agent_runtime_kit/__init__.py b/src/agent_runtime_kit/__init__.py index 9825b54..71fb1b4 100644 --- a/src/agent_runtime_kit/__init__.py +++ b/src/agent_runtime_kit/__init__.py @@ -24,6 +24,7 @@ SessionResumeState, ToolCallAudit, Usage, + runtime_kind_value, ) from agent_runtime_kit.events import ( output_delta_event, @@ -61,6 +62,7 @@ "UnsupportedTaskInputError", "Usage", "create_default_registry", + "runtime_kind_value", "output_delta_event", "safe_emit", "task_completed_event", diff --git a/src/agent_runtime_kit/_errors.py b/src/agent_runtime_kit/_errors.py index afc41b8..51466ac 100644 --- a/src/agent_runtime_kit/_errors.py +++ b/src/agent_runtime_kit/_errors.py @@ -2,7 +2,7 @@ from __future__ import annotations -from agent_runtime_kit._types import AgentRuntimeKind +from agent_runtime_kit._types import AgentRuntimeKind, runtime_kind_value class AgentRuntimeError(RuntimeError): @@ -13,7 +13,7 @@ class AgentRuntimeUnavailableError(AgentRuntimeError): """A runtime cannot be constructed or used in the current environment.""" def __init__(self, kind: AgentRuntimeKind | str, message: str) -> None: - self.kind = AgentRuntimeKind.coerce(kind) + self.kind: AgentRuntimeKind | str = AgentRuntimeKind.coerce(kind) super().__init__(message) @@ -21,14 +21,14 @@ class UnsupportedTaskInputError(AgentRuntimeError, ValueError): """A runtime was asked to honor an input it does not support.""" def __init__(self, kind: AgentRuntimeKind | str, field: str, message: str) -> None: - self.kind = AgentRuntimeKind.coerce(kind) + self.kind: AgentRuntimeKind | str = AgentRuntimeKind.coerce(kind) self.field = field - super().__init__(f"{self.kind.value} cannot honor {field}: {message}") + super().__init__(f"{runtime_kind_value(self.kind)} cannot honor {field}: {message}") class RuntimeNotRegisteredError(AgentRuntimeError, LookupError): """No runtime factory is registered for the requested runtime kind.""" def __init__(self, kind: AgentRuntimeKind | str) -> None: - self.kind = AgentRuntimeKind.coerce(kind) - super().__init__(f"No runtime registered for {self.kind.value!r}") + self.kind: AgentRuntimeKind | str = AgentRuntimeKind.coerce(kind) + super().__init__(f"No runtime registered for {runtime_kind_value(self.kind)!r}") diff --git a/src/agent_runtime_kit/_runtime.py b/src/agent_runtime_kit/_runtime.py index a5fb92e..39e63ed 100644 --- a/src/agent_runtime_kit/_runtime.py +++ b/src/agent_runtime_kit/_runtime.py @@ -111,6 +111,15 @@ async def cancel(self, task_id: str) -> None: self.cancelled_task_ids.add(task_id) + async def aclose(self) -> None: + """No-op: the fake runtime owns no vendor process.""" + + async def __aenter__(self) -> FakeAgentRuntime: + return self + + async def __aexit__(self, _exc_type: object, _exc: object, _tb: object) -> None: + await self.aclose() + def _ensure_supported( kind: AgentRuntimeKind, diff --git a/src/agent_runtime_kit/_types.py b/src/agent_runtime_kit/_types.py index 5ea028a..1da287f 100644 --- a/src/agent_runtime_kit/_types.py +++ b/src/agent_runtime_kit/_types.py @@ -19,12 +19,30 @@ class AgentRuntimeKind(str, Enum): ANTIGRAVITY_AGENT_SDK = "antigravity-agent-sdk" @classmethod - def coerce(cls, value: AgentRuntimeKind | str) -> AgentRuntimeKind: - """Normalize a string or enum value into an ``AgentRuntimeKind``.""" + def coerce(cls, value: AgentRuntimeKind | str) -> AgentRuntimeKind | str: + """Normalize a runtime kind, allowing namespaced third-party strings. + + A value matching a built-in member returns that member. Any other + non-empty string is returned as-is so a third party can register and + dispatch a runtime kind (e.g. ``"x-myorg-agent"``) without forking the + enum. Empty/blank values still raise ``ValueError``. + """ if isinstance(value, cls): return value - return cls(value) + try: + return cls(value) + except ValueError: + normalized = str(value).strip() + if not normalized: + raise ValueError("runtime kind must be a non-empty string") from None + return normalized + + +def runtime_kind_value(value: AgentRuntimeKind | str) -> str: + """Return the wire/string form of a runtime kind (enum member or raw string).""" + + return value.value if isinstance(value, AgentRuntimeKind) else str(value) class AvailabilityReason(str, Enum): @@ -81,7 +99,7 @@ class AgentCapabilities: class RuntimeAvailability: """Availability diagnostic for a runtime in the current environment.""" - kind: AgentRuntimeKind + kind: AgentRuntimeKind | str available: bool reason: AvailabilityReason = AvailabilityReason.UNKNOWN message: str = "" @@ -254,7 +272,11 @@ def cost_usd(self) -> float: class AgentRuntime(Protocol): """Async runtime that drives an ``AgentTask`` to completion.""" - kind: AgentRuntimeKind + # Read-only (covariant) so a concrete adapter may narrow it to a specific + # ``AgentRuntimeKind`` member while third-party adapters use a namespaced str. + @property + def kind(self) -> AgentRuntimeKind | str: ... + capabilities: AgentCapabilities def availability(self) -> RuntimeAvailability: @@ -265,3 +287,19 @@ async def run(self, task: AgentTask) -> AgentResult: async def cancel(self, task_id: str) -> None: """Request cancellation for a task if supported.""" + + async def aclose(self) -> None: + """Release any resources (e.g. a reused vendor process) owned by the runtime. + + Stateless runtimes may implement this as a no-op, but every runtime must + expose it so callers can manage lifecycle uniformly without ``getattr``. + """ + + async def __aenter__(self) -> AgentRuntime: + """Enter an async context managing this runtime's lifecycle.""" + + async def __aexit__(self, exc_type: object, exc: object, tb: object, /) -> None: + """Exit the async context, releasing resources via :meth:`aclose`. + + Parameters are positional-only so implementations may name them freely. + """ diff --git a/src/agent_runtime_kit/events.py b/src/agent_runtime_kit/events.py index 49b5feb..75c15e5 100644 --- a/src/agent_runtime_kit/events.py +++ b/src/agent_runtime_kit/events.py @@ -7,7 +7,13 @@ from datetime import datetime, timezone from typing import Any -from agent_runtime_kit._types import AgentResult, AgentRuntimeKind, AgentTask, ToolCallAudit +from agent_runtime_kit._types import ( + AgentResult, + AgentRuntimeKind, + AgentTask, + ToolCallAudit, + runtime_kind_value, +) DEFAULT_PREVIEW_CHARS = 1000 SENSITIVE_KEY_SUBSTRINGS = ( @@ -40,7 +46,7 @@ def task_started_event( attrs = { "task_id": task.task_id, - "runtime_kind": AgentRuntimeKind.coerce(kind).value, + "runtime_kind": runtime_kind_value(kind), "session_id": task.session_id, "task_goal": task.goal, "system_prompt": task.system, @@ -66,7 +72,7 @@ def task_completed_event( attrs = { "task_id": task.task_id, - "runtime_kind": AgentRuntimeKind.coerce(kind).value, + "runtime_kind": runtime_kind_value(kind), "finish_reason": result.finish_reason, "output": result.output, "rounds": result.rounds, @@ -97,7 +103,7 @@ def task_failed_event( attrs = { "task_id": task.task_id, - "runtime_kind": AgentRuntimeKind.coerce(kind).value, + "runtime_kind": runtime_kind_value(kind), "finish_reason": finish_reason, "error": error, } @@ -119,7 +125,7 @@ def output_delta_event( preview, truncated = _preview(text) attrs = { "task_id": task.task_id, - "runtime_kind": AgentRuntimeKind.coerce(kind).value, + "runtime_kind": runtime_kind_value(kind), "text_delta": preview, "text_delta_length": len(text), "text_delta_truncated": truncated, @@ -142,7 +148,7 @@ def tool_requested_event( attrs = { "task_id": task.task_id, - "runtime_kind": AgentRuntimeKind.coerce(kind).value, + "runtime_kind": runtime_kind_value(kind), "tool_name": tool_name, "argument_count": len(arguments or {}), "argument_keys": sorted(str(key) for key in (arguments or {})), @@ -164,7 +170,7 @@ def tool_completed_event( attrs = { "task_id": task.task_id, - "runtime_kind": AgentRuntimeKind.coerce(kind).value, + "runtime_kind": runtime_kind_value(kind), "tool_name": audit.tool_name, "status": audit.status, "duration_ms": audit.duration_ms, @@ -192,7 +198,7 @@ def vendor_turn_event( attrs = { "task_id": task.task_id, - "runtime_kind": AgentRuntimeKind.coerce(kind).value, + "runtime_kind": runtime_kind_value(kind), "turn_index": turn_index, "payload": dict(payload or {}), } diff --git a/src/agent_runtime_kit/registry.py b/src/agent_runtime_kit/registry.py index dbb6d64..b5ed6dd 100644 --- a/src/agent_runtime_kit/registry.py +++ b/src/agent_runtime_kit/registry.py @@ -12,6 +12,7 @@ AgentRuntime, AgentRuntimeKind, RuntimeAvailability, + runtime_kind_value, ) RuntimeFactory = Callable[..., AgentRuntime] @@ -21,7 +22,7 @@ class RuntimeRegistry: """Register and resolve runtime factories by kind.""" def __init__(self) -> None: - self._factories: dict[AgentRuntimeKind, RuntimeFactory] = {} + self._factories: dict[AgentRuntimeKind | str, RuntimeFactory] = {} def register( self, @@ -34,7 +35,9 @@ def register( normalized = AgentRuntimeKind.coerce(kind) if normalized in self._factories and not replace: - raise ValueError(f"Runtime already registered for {normalized.value!r}") + raise ValueError( + f"Runtime already registered for {runtime_kind_value(normalized)!r}" + ) self._factories[normalized] = factory def unregister(self, kind: AgentRuntimeKind | str) -> None: @@ -43,7 +46,7 @@ def unregister(self, kind: AgentRuntimeKind | str) -> None: normalized = AgentRuntimeKind.coerce(kind) self._factories.pop(normalized, None) - def kinds(self) -> tuple[AgentRuntimeKind, ...]: + def kinds(self) -> tuple[AgentRuntimeKind | str, ...]: """Return registered runtime kinds.""" return tuple(self._factories) diff --git a/src/agent_runtime_kit/testing/fakes.py b/src/agent_runtime_kit/testing/fakes.py index 56f1a16..03831bf 100644 --- a/src/agent_runtime_kit/testing/fakes.py +++ b/src/agent_runtime_kit/testing/fakes.py @@ -151,6 +151,15 @@ async def cancel(self, task_id: str) -> None: del task_id + async def aclose(self) -> None: + """No-op: the fake runtime owns no vendor process.""" + + async def __aenter__(self) -> FakeSDKRuntime: + return self + + async def __aexit__(self, _exc_type: object, _exc: object, _tb: object) -> None: + await self.aclose() + class RecordingEventSink: """In-memory event sink for tests.""" diff --git a/tests/test_core.py b/tests/test_core.py index 94cc304..7885377 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -4,12 +4,14 @@ from agent_runtime_kit import ( AgentCapabilities, + AgentRuntime, AgentRuntimeKind, AgentTask, FakeAgentRuntime, RuntimeNotRegisteredError, UnsupportedTaskInputError, create_default_registry, + runtime_kind_value, ) @@ -42,6 +44,37 @@ def test_registry_rejects_missing_runtime() -> None: registry.resolve(AgentRuntimeKind.FAKE) +def test_registry_accepts_namespaced_third_party_kind() -> None: + # A third party can register a runtime under a namespaced string kind without + # forking the built-in enum. + registry = create_default_registry() + registry.register("x-myorg-agent", lambda **_: FakeAgentRuntime(output="third-party")) + + assert "x-myorg-agent" in registry.kinds() + runtime = registry.resolve("x-myorg-agent") + assert isinstance(runtime, FakeAgentRuntime) + + +def test_coerce_returns_namespaced_string_and_rejects_blank() -> None: + assert AgentRuntimeKind.coerce("claude-agent-sdk") is AgentRuntimeKind.CLAUDE_AGENT_SDK + assert AgentRuntimeKind.coerce("x-myorg-agent") == "x-myorg-agent" + assert runtime_kind_value(AgentRuntimeKind.FAKE) == "fake" + assert runtime_kind_value("x-myorg-agent") == "x-myorg-agent" + with pytest.raises(ValueError): + AgentRuntimeKind.coerce(" ") + + +@pytest.mark.asyncio +async def test_fake_runtime_satisfies_protocol_with_lifecycle() -> None: + runtime = FakeAgentRuntime(output="done") + + assert isinstance(runtime, AgentRuntime) + async with runtime as entered: + result = await entered.run(AgentTask(goal="x")) + assert result.output == "done" + await runtime.aclose() + + @pytest.mark.asyncio async def test_fake_runtime_rejects_unsupported_structured_output() -> None: runtime = FakeAgentRuntime(capabilities=AgentCapabilities(structured_output=False))