monoid-agent-kernel 0.16.1__py3-none-any.whl

monoid_agent_kernel/reference/backend/loop_factory.py

@@ -0,0 +1,253 @@
+from __future__ import annotations
+
+import time
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from monoid_agent_kernel.core.agents import AgentRuntimeConfig, RuntimeConfigProvider
+from monoid_agent_kernel.core.capability import CapabilityBroker
+from monoid_agent_kernel.core.checkpoint import CheckpointStore, RunCheckpoint
+from monoid_agent_kernel.core.context import ContextProvider
+from monoid_agent_kernel.core.events import AgentEvent, EventSink
+from monoid_agent_kernel.core.outbox import OutboxSender
+from monoid_agent_kernel.core.output_validator import OutputValidator
+from monoid_agent_kernel.core.spec import AgentRunSpec, ModelConfig, RunLimits
+from monoid_agent_kernel.loop import AgentLoop
+from monoid_agent_kernel.providers.base import ModelAdapter
+from monoid_agent_kernel.providers.gateway import GatewayModelAdapter
+from monoid_agent_kernel.reference._shared.tokens import TokenKind, TokenManager
+from monoid_agent_kernel.reference.backend.ports import MutableRunRecordPort, RunRequestPort
+from monoid_agent_kernel.reference.backend.run_state import BackendRunStateSink
+from monoid_agent_kernel.tools.base import ToolProvider
+from monoid_agent_kernel.web import WebGatewayClient
+
+ModelAdapterFactory = Callable[[AgentRunSpec, str], ModelAdapter]
+
+
+@dataclass(frozen=True)
+class BackendLoopBuild:
+    spec: AgentRunSpec
+    model_adapter: ModelAdapter
+    web_gateway_client: WebGatewayClient | None
+    runtime_config_provider: RuntimeConfigProvider
+    capability_broker: CapabilityBroker | None
+    outbox_sender: OutboxSender | None
+    loop: AgentLoop
+
+
+@dataclass(frozen=True)
+class BackendLoopFactoryContext:
+    run_root_provider: Callable[[], Path]
+    llm_gateway_url_provider: Callable[[], str]
+    web_gateway_url_provider: Callable[[], str | None]
+    model_adapter_factory_provider: Callable[[], ModelAdapterFactory | None]
+    token_manager_provider: Callable[[], TokenManager]
+    llm_gateway_token_ttl_s_provider: Callable[[], int]
+    checkpoint_store_provider: Callable[[], CheckpointStore | None]
+    emit_output_deltas_provider: Callable[[], bool]
+    extra_event_sink_factories_provider: Callable[[], tuple[Callable[[], EventSink], ...]]
+    subagent_definitions_provider: Callable[[], Mapping[str, Any] | None]
+    tool_providers_provider: Callable[[], tuple[ToolProvider, ...]]
+    context_providers_provider: Callable[[], tuple[ContextProvider, ...]]
+    output_validators_provider: Callable[[], tuple[OutputValidator, ...]]
+    capability_broker_factory_provider: Callable[[], Callable[[RunRequestPort], CapabilityBroker | None] | None]
+    outbox_sender_factory_provider: Callable[[], Callable[[RunRequestPort], OutboxSender | None] | None]
+    current_runtime_config: Callable[[str], AgentRuntimeConfig | None]
+    record: Callable[[str], MutableRunRecordPort]
+    record_event: Callable[[str, AgentEvent], None]
+    persist_checkpoint_payload: Callable[[MutableRunRecordPort, RunCheckpoint, Mapping[str, bytes]], None]
+
+
+class BackendRuntimeConfigProvider(RuntimeConfigProvider):
+    def __init__(
+        self,
+        current_runtime_config: Callable[[str], AgentRuntimeConfig | None],
+        run_id: str,
+    ) -> None:
+        self._current_runtime_config = current_runtime_config
+        self._run_id = run_id
+
+    def current_config(self, run_id: str) -> AgentRuntimeConfig | None:
+        del run_id
+        return self._current_runtime_config(self._run_id)
+
+
+@dataclass
+class _GatewayTokenSource:
+    """Callable gateway token source that re-mints shortly before expiry for long runs."""
+
+    token_manager: TokenManager
+    kind: TokenKind
+    audience: str
+    run_id: str
+    tenant_id: str
+    user_id: str
+    ttl_s: int
+    metadata: dict[str, Any] = field(default_factory=dict)
+    refresh_skew_s: int = 300
+    _token: str = ""
+    _expires_at: float = 0.0
+
+    def __call__(self) -> str:
+        now = time.time()
+        skew = min(self.refresh_skew_s, self.ttl_s // 2)
+        if not self._token or now >= self._expires_at - skew:
+            self._token = self.token_manager.issue(
+                kind=self.kind,
+                audience=self.audience,
+                run_id=self.run_id,
+                tenant_id=self.tenant_id,
+                user_id=self.user_id,
+                ttl_s=self.ttl_s,
+                metadata=dict(self.metadata),
+            )
+            self._expires_at = now + self.ttl_s
+        return self._token
+
+
+class BackendLoopFactory:
+    """Builds AgentLoop instances for the Reference backend facade."""
+
+    def __init__(self, context: BackendLoopFactoryContext) -> None:
+        self._context = context
+
+    def build(
+        self,
+        run_id: str,
+        request: RunRequestPort,
+        workspace_root: Path,
+        llm_gateway_token: str,
+        web_gateway_token: str,
+        *,
+        include_outbox_sender: bool = True,
+    ) -> BackendLoopBuild:
+        spec = self.run_spec_for_request(run_id, request, workspace_root)
+        runtime_config = self._context.current_runtime_config(run_id)
+        model_adapter = self.build_model_adapter(
+            spec,
+            llm_gateway_token,
+            runtime_config.model if runtime_config is not None else None,
+            token_provider=self.llm_token_source(run_id, request, runtime_config),
+        )
+        web_gateway_client = self.web_gateway_client(web_gateway_token)
+        runtime_config_provider = BackendRuntimeConfigProvider(
+            self._context.current_runtime_config,
+            run_id,
+        )
+        capability_broker = self.capability_broker_for(request)
+        outbox_sender = self.outbox_sender_for(request) if include_outbox_sender else None
+        record = self._context.record(run_id)
+        loop = AgentLoop(
+            spec=spec,
+            model_adapter=model_adapter,
+            event_sinks=(
+                BackendRunStateSink(self._context.record_event, run_id),
+                *(make() for make in self._context.extra_event_sink_factories_provider()),
+            ),
+            permission_policy=request.permission_policy,
+            cancellation_token=record.cancellation_token,
+            shell_approval_provider=None,
+            web_gateway_client=web_gateway_client,
+            runtime_config_provider=runtime_config_provider,
+            checkpoint_store=self._context.checkpoint_store_provider(),
+            emit_output_deltas=self._context.emit_output_deltas_provider(),
+            subagent_definitions=self._context.subagent_definitions_provider(),
+            tool_providers=self._context.tool_providers_provider(),
+            context_providers=self._context.context_providers_provider(),
+            output_validators=self._context.output_validators_provider(),
+            capability_broker=capability_broker,
+            checkpoint_persist_callback=lambda checkpoint, blobs: self._context.persist_checkpoint_payload(
+                self._context.record(run_id),
+                checkpoint,
+                blobs,
+            ),
+        )
+        return BackendLoopBuild(
+            spec=spec,
+            model_adapter=model_adapter,
+            web_gateway_client=web_gateway_client,
+            runtime_config_provider=runtime_config_provider,
+            capability_broker=capability_broker,
+            outbox_sender=outbox_sender,
+            loop=loop,
+        )
+
+    def run_spec_for_request(
+        self,
+        run_id: str,
+        request: RunRequestPort,
+        workspace_root: Path,
+    ) -> AgentRunSpec:
+        return AgentRunSpec(
+            workspace_root=workspace_root,
+            run_root=self._context.run_root_provider(),
+            run_id=run_id,
+            mode=request.mode,
+            workspace_backend=request.workspace_backend,
+            limits=RunLimits(
+                max_steps=request.max_steps,
+                max_tool_calls=request.max_tool_calls,
+                max_bytes_read=request.max_bytes_read,
+                max_duration_s=request.max_duration_s,
+            ),
+            permission_policy=request.permission_policy,
+            metadata={
+                **request.metadata,
+                "tenant_id": request.tenant_id,
+                "user_id": request.user_id,
+            },
+        )
+
+    def capability_broker_for(self, request: RunRequestPort) -> CapabilityBroker | None:
+        factory = self._context.capability_broker_factory_provider()
+        if factory is None:
+            return None
+        return factory(request)
+
+    def outbox_sender_for(self, request: RunRequestPort) -> OutboxSender | None:
+        factory = self._context.outbox_sender_factory_provider()
+        if factory is None:
+            return None
+        return factory(request)
+
+    def llm_token_source(
+        self,
+        run_id: str,
+        request: RunRequestPort,
+        runtime_config: AgentRuntimeConfig | None,
+    ) -> _GatewayTokenSource:
+        return _GatewayTokenSource(
+            token_manager=self._context.token_manager_provider(),
+            kind="llm_gateway",
+            audience="csp.llm-gateway",
+            run_id=run_id,
+            tenant_id=request.tenant_id,
+            user_id=request.user_id,
+            ttl_s=self._context.llm_gateway_token_ttl_s_provider(),
+            metadata={"agent_config_hash": runtime_config.config_hash} if runtime_config is not None else {},
+        )
+
+    def build_model_adapter(
+        self,
+        spec: AgentRunSpec,
+        llm_gateway_token: str,
+        model_config: ModelConfig | None,
+        token_provider: Callable[[], str | None] | None = None,
+    ) -> ModelAdapter:
+        factory = self._context.model_adapter_factory_provider()
+        if factory is not None:
+            return factory(spec, llm_gateway_token)
+        return GatewayModelAdapter(
+            model_config or ModelConfig(),
+            gateway_url=self._context.llm_gateway_url_provider(),
+            token=llm_gateway_token,
+            token_provider=token_provider,
+        )
+
+    def web_gateway_client(self, token: str) -> WebGatewayClient | None:
+        gateway_url = self._context.web_gateway_url_provider()
+        if not token or not gateway_url:
+            return None
+        return WebGatewayClient(gateway_url, token=token)

monoid_agent_kernel/reference/backend/outbox_dispatch.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import random
+import time
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from monoid_agent_kernel.core.checkpoint import CheckpointStore, RunCheckpoint
+from monoid_agent_kernel.core.inbox import InboxMessage
+from monoid_agent_kernel.core.outbox import OutboxReceipt, OutboxRequest
+from monoid_agent_kernel.core.trace_context import new_traceparent
+from monoid_agent_kernel.reference.backend.ports import MutableRunRecordPort, queued_message_snapshot
+
+
+class OutboxLoopPort(Protocol):
+    def due_outbox(self, now: float) -> list[OutboxRequest]: ...
+
+    def record_outbox_result(
+        self,
+        request_id: str,
+        receipt: OutboxReceipt,
+        *,
+        max_attempts: int = 5,
+        next_attempt_at: float = 0.0,
+    ) -> str: ...
+
+    def snapshot(self) -> RunCheckpoint | None: ...
+
+    def collect_checkpoint_blobs(self) -> Mapping[str, bytes]: ...
+
+
+@dataclass(frozen=True)
+class OutboxRetryPolicy:
+    max_attempts: int
+    base_s: float
+    factor: float
+    cap_s: float
+
+
+@dataclass(frozen=True)
+class OutboxDispatchContext:
+    retry_policy_provider: Callable[[], OutboxRetryPolicy]
+    max_message_queue_depth_provider: Callable[[], int]
+    checkpoint_store_provider: Callable[[], CheckpointStore]
+    rng_provider: Callable[[], random.Random]
+    live_outbox_runs: Callable[[], list[tuple[MutableRunRecordPort, OutboxLoopPort]]]
+    call_soon: Callable[..., None]
+    record_terminal: Callable[[MutableRunRecordPort], bool]
+
+
+class OutboxDispatchService:
+    """Reference backend edge dispatcher for staged outbox side effects."""
+
+    def __init__(self, context: OutboxDispatchContext) -> None:
+        self._context = context
+
+    def backoff_delay(self, attempts: int) -> float:
+        """Capped exponential backoff with full jitter."""
+        policy = self._context.retry_policy_provider()
+        ceiling = min(policy.cap_s, policy.base_s * (policy.factor**attempts))
+        return self._context.rng_provider().uniform(0.0, max(0.0, ceiling))
+
+    def drain_outbox(self, record: MutableRunRecordPort, loop: OutboxLoopPort) -> None:
+        """Dispatch due staged outbox requests and persist the resulting checkpoint state."""
+        sender = record.outbox_sender
+        now = time.time()
+        due = loop.due_outbox(now)
+        if sender is None or not due:
+            return
+        changed = False
+        policy = self._context.retry_policy_provider()
+        for request in due:
+            if not request.traceparent:
+                request.traceparent = new_traceparent()
+            try:
+                receipt = sender.send(request)
+            except Exception as exc:  # a sender raising is a retryable transport failure
+                receipt = OutboxReceipt(ok=False, error=str(exc), retryable=True)
+            next_attempt_at = now + self.backoff_delay(request.attempts + 1)
+            status = loop.record_outbox_result(
+                request.id,
+                receipt,
+                max_attempts=policy.max_attempts,
+                next_attempt_at=next_attempt_at,
+            )
+            changed = True
+            if request.expect_ack and status in {"dispatched", "failed"}:
+                self.stage_outbox_ack(record, request, status, receipt)
+        if changed:
+            checkpoint = loop.snapshot()
+            if checkpoint is not None:
+                checkpoint.queued_messages = queued_message_snapshot(record.message_queue)
+                checkpoint.inbox_seen_ids = sorted(record.seen_inbox_ids)
+                self._context.checkpoint_store_provider().put(checkpoint, loop.collect_checkpoint_blobs())
+
+    def stage_outbox_ack(
+        self, record: MutableRunRecordPort, request: Any, status: str, receipt: OutboxReceipt
+    ) -> None:
+        """Deliver an outbox send receipt back into the run inbox as a correlated message."""
+        ack_id = f"ack_{request.id}"
+        if self._context.record_terminal(record):
+            return
+        if (
+            ack_id in record.seen_inbox_ids
+            or record.message_queue.qsize() >= self._context.max_message_queue_depth_provider()
+        ):
+            return
+        summary = f"[outbox-ack] request {request.id} to {request.destination!r}: {status}"
+        if receipt.reference:
+            summary += f" (ref={receipt.reference})"
+        if receipt.error:
+            summary += f" (error={receipt.error})"
+        envelope = InboxMessage(
+            content=summary,
+            id=ack_id,
+            source="outbox",
+            type="outbox_ack",
+            run_id=record.run_id,
+            correlation_id=request.correlation_id or request.id,
+            causation_id=request.id,
+            traceparent=request.traceparent,
+            tracestate=request.tracestate,
+        )
+        record.message_queue.put_nowait(envelope.to_json())
+
+    def redrive_outbox(self) -> None:
+        """Schedule due outbox drains for live runs with senders."""
+        for record, loop in self._context.live_outbox_runs():
+            self._context.call_soon(self.drain_outbox, record, loop)

monoid_agent_kernel/reference/backend/ports.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Awaitable, Mapping
+from pathlib import Path
+from typing import Any, Protocol
+
+from monoid_agent_kernel.core.agents import AgentRuntimeConfig
+from monoid_agent_kernel.core.checkpoint import RunCheckpoint
+from monoid_agent_kernel.core.lifecycle import SessionState
+from monoid_agent_kernel.core.outbox import OutboxSender
+from monoid_agent_kernel.core.result import AgentRunResult, Suspension
+
+
+class TokenClaimsPort(Protocol):
+    tenant_id: str
+    user_id: str
+    run_id: str
+
+
+class CancellationTokenPort(Protocol):
+    requested: bool
+
+
+class MessageQueuePort(Protocol):
+    def qsize(self) -> int: ...
+
+    def put_nowait(self, item: Any) -> None: ...
+
+    def get(self) -> Awaitable[Any]: ...
+
+
+def queued_message_snapshot(queue: MessageQueuePort) -> list[str | list[Any] | dict[str, Any]]:
+    raw_queue = getattr(queue, "_queue", ())
+    return [message for message in list(raw_queue) if isinstance(message, (str, list, dict))]
+
+
+class LoopPort(Protocol):
+    def wait_for_pending_tasks(self, timeout_s: float) -> bool: ...
+
+    def has_pending_tasks(self) -> bool: ...
+
+    def emit_external_event(
+        self,
+        event_type: str,
+        *,
+        data: dict[str, Any] | None = None,
+        level: str = "info",
+    ) -> bool: ...
+
+    async def arun_until_suspended(self, user_input: Any | None = None) -> Suspension: ...
+
+    def fail_recoverable(self, error: str, *, error_code: str) -> None: ...
+
+    def await_user_input(self) -> None: ...
+
+    async def aclose(self) -> AgentRunResult: ...
+
+    def snapshot(self) -> RunCheckpoint | None: ...
+
+    def collect_checkpoint_blobs(self) -> Mapping[str, bytes]: ...
+
+    def interrupt_turn(self) -> None: ...
+
+    def pause_turn(self) -> None: ...
+
+    def revoke_capability(
+        self,
+        *,
+        capability: str | None = None,
+        lease_id: str | None = None,
+        before: float | None = None,
+        reason: str = "",
+    ) -> dict[str, Any]: ...
+
+    def report_task_result(
+        self,
+        task_id: str,
+        result: dict[str, Any],
+        *,
+        status: str,
+        persist_checkpoint: bool,
+    ) -> dict[str, Any]: ...
+
+    def create_task(self, kind: str, request: dict[str, Any]) -> str: ...
+
+    def restore(self, checkpoint: RunCheckpoint, *, blobs: Mapping[str, bytes]) -> None: ...
+
+
+class RunStreamPort(Protocol):
+    suspension: Suspension | None
+
+    async def __aenter__(self) -> RunStreamPort: ...
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> bool | None: ...
+
+    def __aiter__(self) -> AsyncIterator[Any]: ...
+
+
+class RunExecutionLoopPort(LoopPort, Protocol):
+    async def aopen(self) -> None: ...
+
+    def astream(self, user_input: Any) -> RunStreamPort: ...
+
+
+class LoopBuildPort(Protocol):
+    loop: RunExecutionLoopPort
+    outbox_sender: OutboxSender | None
+
+
+class RunRecordPort(Protocol):
+    run_id: str
+    tenant_id: str
+    user_id: str
+    workspace_root: Path
+    run_dir: Path
+    state: SessionState
+    terminal: bool
+    created_at: float
+    started_at: float
+    finished_at: float
+    last_event_seq: int
+    last_event_type: str
+    error: str
+    error_code: str
+    result: AgentRunResult | None
+    last_final_output: Any
+    runtime_config: AgentRuntimeConfig | None
+    runtime_config_issuer: str
+    runtime_config_reason: str
+    runtime_config_committed_at: float
+
+
+class MutableRunRecordPort(RunRecordPort, Protocol):
+    message_queue: MessageQueuePort
+    loop: LoopPort | None
+    cancellation_token: CancellationTokenPort
+    seen_inbox_ids: set[str]
+    outbox_sender: OutboxSender | None
+
+
+class PreparedRunPort(Protocol):
+    run_id: str
+    record: MutableRunRecordPort
+    workspace_root: Path
+    run_token: str
+    llm_gateway_token: str
+    web_gateway_token: str
+
+
+class RunRequestPort(Protocol):
+    tenant_id: str
+    user_id: str
+    workspace_root: Path
+    instruction: str
+    input_parts: tuple[Any, ...]
+    mode: str
+    workspace_backend: str
+    max_steps: int
+    max_tool_calls: int
+    max_bytes_read: int
+    max_duration_s: int | None
+    permission_policy: Any
+    runtime_config: AgentRuntimeConfig | None
+    multi_turn: bool
+    metadata: dict[str, Any]
+
+
+class LeaseStorePort(Protocol):
+    def candidate_run_ids(self) -> list[str]: ...
+
+    def is_stale(self, run_id: str) -> bool: ...
+
+    def try_claim(self, run_id: str, worker_id: str, ttl_s: float) -> bool: ...
+
+    def release(self, run_id: str) -> None: ...
+
+
+class DriveOpenSessionPort(Protocol):
+    def __call__(
+        self,
+        record: MutableRunRecordPort,
+        request: RunRequestPort,
+        loop: LoopPort,
+        suspension: Suspension,
+        *,
+        started: float,
+        turns: int,
+    ) -> Awaitable[AgentRunResult]: ...