193 lines
6.0 KiB
Python
193 lines
6.0 KiB
Python
"""Runtime context handed to skill handlers.
|
|
|
|
The same agent code runs unchanged on local dev, Docker, Kubernetes, and
|
|
hosted runtimes — the runtime provides a concrete :class:`RunContext` that
|
|
implements artifact storage, secret access, streaming, and cancellation.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Generic, Sequence, TypeVar
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from .workspace import WorkspaceClient
|
|
|
|
AuthT = TypeVar("AuthT", bound=BaseModel)
|
|
|
|
|
|
class CancelledByCaller(RuntimeError):
|
|
"""Raised by :meth:`RunContext.check_cancelled` when the caller cancelled."""
|
|
|
|
|
|
class MissingScopes(PermissionError):
|
|
"""Raised by :meth:`RunContext.require_scopes` when caller lacks scopes."""
|
|
|
|
def __init__(self, missing: Sequence[str]) -> None:
|
|
self.missing = tuple(missing)
|
|
super().__init__(f"missing scopes: {sorted(self.missing)}")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ArtifactRef:
|
|
"""Opaque handle to a stored artifact (blob, file, etc.)."""
|
|
|
|
name: str
|
|
uri: str
|
|
mime_type: str
|
|
size_bytes: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AgentEvent:
|
|
"""A structured event emitted during a skill run."""
|
|
|
|
kind: str
|
|
payload: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
class RunContext(ABC, Generic[AuthT]):
|
|
"""Per-invocation context.
|
|
|
|
A new context is constructed by the runtime for every skill call. It
|
|
carries caller identity (``auth``), the task identity, and runtime
|
|
capabilities (artifacts, secrets, streaming, cancellation).
|
|
|
|
Agents must depend only on this abstract interface, never on a concrete
|
|
runtime implementation.
|
|
"""
|
|
|
|
task_id: str
|
|
auth: AuthT
|
|
|
|
@abstractmethod
|
|
async def emit_event(self, event: AgentEvent) -> None:
|
|
"""Publish a structured event to subscribers (UI, logs, traces)."""
|
|
|
|
@abstractmethod
|
|
async def write_artifact(
|
|
self, name: str, data: bytes, mime_type: str
|
|
) -> ArtifactRef:
|
|
"""Persist ``data`` as a named artifact and return a reference."""
|
|
|
|
@abstractmethod
|
|
async def check_cancelled(self) -> None:
|
|
"""Raise :class:`CancelledByCaller` if the caller cancelled."""
|
|
|
|
@abstractmethod
|
|
def secret(self, name: str) -> str:
|
|
"""Look up a runtime-injected secret by logical name."""
|
|
|
|
@property
|
|
@abstractmethod
|
|
def workspace(self) -> WorkspaceClient:
|
|
"""Negotiation surface for workspace access.
|
|
|
|
Raises if the agent's :attr:`A2AAgent.workspace_access` is disabled.
|
|
"""
|
|
|
|
# --- concrete helpers built on emit_event ---
|
|
|
|
async def emit_progress(self, message: str) -> None:
|
|
"""Emit a human-readable progress event."""
|
|
await self.emit_event(AgentEvent(kind="progress", payload={"message": message}))
|
|
|
|
async def emit_text_delta(self, text: str) -> None:
|
|
"""Emit a streamed token chunk (for LLM-style streaming output)."""
|
|
await self.emit_event(AgentEvent(kind="text_delta", payload={"text": text}))
|
|
|
|
async def emit_artifact(self, ref: ArtifactRef) -> None:
|
|
"""Notify subscribers that a new artifact is available."""
|
|
await self.emit_event(
|
|
AgentEvent(
|
|
kind="artifact",
|
|
payload={
|
|
"name": ref.name,
|
|
"uri": ref.uri,
|
|
"mime_type": ref.mime_type,
|
|
"size_bytes": ref.size_bytes,
|
|
},
|
|
)
|
|
)
|
|
|
|
async def emit_error(self, message: str, *, code: str | None = None) -> None:
|
|
"""Emit a structured error event (does not raise)."""
|
|
await self.emit_event(
|
|
AgentEvent(kind="error", payload={"message": message, "code": code})
|
|
)
|
|
|
|
def require_scopes(self, required: Sequence[str]) -> None:
|
|
"""Raise :class:`MissingScopes` if ``self.auth`` lacks any required scope.
|
|
|
|
Auth models without a ``scopes`` attribute (e.g. :class:`NoAuth`) are
|
|
treated as having an empty scope set.
|
|
"""
|
|
if not required:
|
|
return
|
|
auth_scopes = set(getattr(self.auth, "scopes", ()) or ())
|
|
missing = [s for s in required if s not in auth_scopes]
|
|
if missing:
|
|
raise MissingScopes(missing)
|
|
|
|
|
|
class LocalRunContext(RunContext[AuthT]):
|
|
"""In-memory context for local dev and tests.
|
|
|
|
Stores events and artifacts in lists/dicts. Secrets come from a plain
|
|
mapping. Cancellation is driven by an :class:`asyncio.Event`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
auth: AuthT,
|
|
task_id: str = "local-task",
|
|
secrets: dict[str, str] | None = None,
|
|
workspace: WorkspaceClient | None = None,
|
|
) -> None:
|
|
self.task_id = task_id
|
|
self.auth = auth
|
|
self._secrets: dict[str, str] = dict(secrets or {})
|
|
self._workspace = workspace
|
|
self._cancel = asyncio.Event()
|
|
self.events: list[AgentEvent] = []
|
|
self.artifacts: dict[str, bytes] = {}
|
|
|
|
@property
|
|
def workspace(self) -> WorkspaceClient:
|
|
if self._workspace is None:
|
|
raise PermissionError(
|
|
"no workspace bound to this context; agent did not declare "
|
|
"workspace_access or runtime did not provision one"
|
|
)
|
|
return self._workspace
|
|
|
|
async def emit_event(self, event: AgentEvent) -> None:
|
|
self.events.append(event)
|
|
|
|
async def write_artifact(
|
|
self, name: str, data: bytes, mime_type: str
|
|
) -> ArtifactRef:
|
|
self.artifacts[name] = data
|
|
return ArtifactRef(
|
|
name=name,
|
|
uri=f"memory://{self.task_id}/{name}",
|
|
mime_type=mime_type,
|
|
size_bytes=len(data),
|
|
)
|
|
|
|
async def check_cancelled(self) -> None:
|
|
if self._cancel.is_set():
|
|
raise CancelledByCaller(self.task_id)
|
|
|
|
def cancel(self) -> None:
|
|
self._cancel.set()
|
|
|
|
def secret(self, name: str) -> str:
|
|
try:
|
|
return self._secrets[name]
|
|
except KeyError as exc:
|
|
raise KeyError(f"unknown secret: {name!r}") from exc
|