induscode 0.1.0__py3-none-any.whl

induscode/runtime_bridge/broker.py

@@ -0,0 +1,397 @@
+"""RuntimeBroker — the single decision point that routes every turn to
+either an external runtime (a spawned child coding-agent driven by a
+:class:`~induscode.runtime_bridge.contract.RuntimeBridge`) or the framework
+network stream (:func:`indusagi.ai.stream_simple` over an HTTP provider)
+(port of TS ``src/runtime-bridge/broker.ts``).
+
+The broker is the only component the product asks "how is this turn
+produced?". It owns three concerns:
+
+1. **Routing.** :meth:`_Broker.route` resolves the model's optional
+   :class:`~induscode.runtime_bridge.contract.ExternalRuntimeSpec` (a
+   ``bridge:<adapter>`` baseUrl decode) and a matching registered bridge. A
+   model with a spec whose adapter is registered routes ``"external"``;
+   everything else routes ``"framework"``.
+2. **Dispatch.** :meth:`RuntimeBrokerRuntime.exchange` acts on that
+   decision: for an external route it builds (or is injected) a
+   :class:`~induscode.runtime_bridge.contract.ChildTransport` and calls
+   :meth:`RuntimeBridge.run_exchange` over it; for a framework route it
+   calls the framework ``stream_simple``. Either way it returns the
+   framework ``AssistantMessageEventStream`` the turn streams into — the two
+   paths are indistinguishable to the caller.
+3. **Resume persistence.** When a bridge surfaces a ``resume`` token (a CLI
+   session id / thread id), the broker persists it through an injected
+   :class:`~induscode.runtime_bridge.contract.RuntimeLinkStore` as a
+   *renamed* custom transcript entry —
+   :data:`~induscode.runtime_bridge.contract.RUNTIME_LINK_ENTRY` =
+   ``"external-runtime-link"``, shape ``{source, bridge, resumeToken, at}``
+   — so a later exchange can reattach the same underlying session. Handle
+   reuse keys on the composite :func:`runtime_source_key`
+   (``source|model|bridge``), not three separate field comparisons.
+
+The ``ChildTransport`` is injectable end to end: a production factory wraps
+a spawned process, a test factory returns a hand-written fake. The broker
+never imports a subprocess module; nothing here spawns a real
+``claude``/``codex`` binary.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, replace
+from datetime import datetime, timezone
+from types import SimpleNamespace
+from typing import Any, Protocol
+
+from indusagi.ai import stream_simple
+
+from induscode.runtime_bridge.bridges import spec_from_model
+from induscode.runtime_bridge.contract import (
+    AssistantMessageEventStream,
+    ChildTransportFactory,
+    Context,
+    ExchangeOptions,
+    ExternalRoute,
+    ExternalRuntimeSpec,
+    FrameworkRoute,
+    FrameworkStream,
+    Model,
+    RuntimeBridge,
+    RuntimeBroker,
+    RuntimeLink,
+    RuntimeLinkStore,
+    RuntimeRoute,
+    TransportContext,
+)
+
+__all__ = [
+    "RuntimeBrokerDeps",
+    "RuntimeBrokerRuntime",
+    "create_runtime_broker",
+    "runtime_source_key",
+]
+
+
+# ---------------------------------------------------------------------------
+# Resume-token reuse key
+# ---------------------------------------------------------------------------
+
+
+def runtime_source_key(source: str, model_id: str, bridge: str) -> str:
+    """Compose the composite reuse key a broker matches a persisted
+    :class:`~induscode.runtime_bridge.contract.RuntimeLink` against (TS
+    ``runtimeSourceKey``). A single ``source|model|bridge`` string rather
+    than three field comparisons — a stored link is reattachable iff its key
+    equals the key of the exchange about to run.
+
+    :param source: the model provider slug
+    :param model_id: the model id
+    :param bridge: the bridge adapter id
+    """
+    return f"{source}|{model_id}|{bridge}"
+
+
+# ---------------------------------------------------------------------------
+# Broker dependencies
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True)
+class RuntimeBrokerDeps:
+    """Construction-time dependencies for :func:`create_runtime_broker` (TS
+    ``RuntimeBrokerDeps``). All optional — an empty broker registers bridges
+    later, routes everything to the framework until a transport factory is
+    wired, and skips persistence when no store is set. Field names keep the
+    TS spelling."""
+
+    #: The seam that builds a ``ChildTransport`` for an external exchange.
+    #: When absent, an external route still *decides* ``"external"`` but
+    #: :meth:`RuntimeBrokerRuntime.exchange` cannot drive a child — so it
+    #: falls through to the framework path. Tests inject a fake-transport
+    #: factory here.
+    transportFactory: ChildTransportFactory | None = None
+    #: Where resume tokens persist; omitted ⇒ tokens are not persisted.
+    linkStore: RuntimeLinkStore | None = None
+    #: Override for the framework network stream (defaults to
+    #: ``stream_simple``). Injected in tests so the framework path is
+    #: observable without a network.
+    frameworkStream: FrameworkStream | None = None
+    #: Bridges to pre-register at construction (equivalent to ``register``).
+    bridges: Sequence[RuntimeBridge] = ()
+
+
+# ---------------------------------------------------------------------------
+# The runtime broker
+# ---------------------------------------------------------------------------
+
+
+class RuntimeBrokerRuntime(RuntimeBroker, Protocol):
+    """The broker the product drives (TS ``RuntimeBrokerRuntime``). Extends
+    the frozen :class:`~induscode.runtime_bridge.contract.RuntimeBroker`
+    routing surface with :meth:`exchange`: the dispatch half that acts on a
+    :meth:`~induscode.runtime_bridge.contract.RuntimeBroker.route` decision
+    and returns the framework stream the turn streams into."""
+
+    def exchange(
+        self,
+        model: Model,
+        context: Context,
+        opts: ExchangeOptions | None = None,
+    ) -> AssistantMessageEventStream:
+        """Produce the turn for ``model``. Routes (via ``route``); on an
+        ``"external"`` route with a wired transport factory it drives the
+        bridge's ``run_exchange`` over a freshly-built ``ChildTransport``
+        (resolving + later persisting the resume token); otherwise it runs
+        the framework stream. Returns the ``AssistantMessageEventStream``
+        synchronously, like ``stream_simple``.
+
+        :param model: the model the turn is bound to
+        :param context: the framework conversation context
+        :param opts: per-exchange options
+        """
+        ...
+
+
+class _Broker:
+    """The concrete broker. Private to the module; consumers obtain one
+    through :func:`create_runtime_broker`."""
+
+    __slots__ = ("_bridges", "_framework", "_link_store", "_transport_factory")
+
+    def __init__(self, deps: RuntimeBrokerDeps | None = None) -> None:
+        deps = deps if deps is not None else RuntimeBrokerDeps()
+        #: Registered bridges, keyed by adapter id (last registration wins).
+        self._bridges: dict[str, RuntimeBridge] = {}
+        self._transport_factory = deps.transportFactory
+        self._link_store = deps.linkStore
+        self._framework: FrameworkStream = (
+            deps.frameworkStream
+            if deps.frameworkStream is not None
+            else (lambda model, context, opts: stream_simple(model, context, opts))
+        )
+        for bridge in deps.bridges:
+            self.register(bridge)
+
+    # ---- registry ----
+
+    def register(self, bridge: RuntimeBridge) -> None:
+        self._bridges[bridge.adapter] = bridge
+
+    # ---- spec / credential resolution ----
+
+    def resolve_spec(self, model: Model) -> ExternalRuntimeSpec | None:
+        return spec_from_model(model)
+
+    def requires_credential(self, spec: ExternalRuntimeSpec) -> bool:
+        # Delegate to the owning bridge when registered; else fall back to the
+        # spec's own auth mode so the predicate answers without a bridge present.
+        bridge = self._bridges.get(spec.adapter)
+        if bridge is not None:
+            return bridge.requires_credential(spec)
+        return spec.authMode == "api-key"
+
+    # ---- routing ----
+
+    def route(self, model: Model, _context: Context, _opts: ExchangeOptions) -> RuntimeRoute:
+        spec = self.resolve_spec(model)
+        if spec is None:
+            return FrameworkRoute()
+        bridge = self._bridges.get(spec.adapter)
+        if bridge is None:
+            return FrameworkRoute()
+        return ExternalRoute(bridge=bridge, spec=spec)
+
+    # ---- dispatch ----
+
+    def exchange(
+        self,
+        model: Model,
+        context: Context,
+        opts: ExchangeOptions | None = None,
+    ) -> AssistantMessageEventStream:
+        opts = opts if opts is not None else ExchangeOptions()
+        route = self.route(model, context, opts)
+
+        # No external runtime, or no transport seam to drive a child: the
+        # framework network path produces the turn unchanged.
+        if route.target == "framework" or self._transport_factory is None:
+            return self._framework(model, context, opts)
+
+        bridge, spec = route.bridge, route.spec
+        source = str(_model_field(model, "provider"))
+        model_id = str(_model_field(model, "id"))
+        source_key = runtime_source_key(source, model_id, spec.adapter)
+
+        # Resolve a persisted resume token for this reuse key.
+        #
+        # # parity: SYNC FAST-PATH (TS `typeof persisted === "string"`) — only
+        # a *synchronously-available* resume token reattaches this dispatch.
+        # An async `find` (a coroutine/awaitable result) simply means no
+        # reattach on this first dispatch: the TS broker never awaits the
+        # promise, and this port preserves that quirk verbatim (the awaitable
+        # is closed, never awaited). A disk-backed store that wants reattach
+        # must answer `find` synchronously.
+        persisted = self._link_store.find(source_key) if self._link_store is not None else None
+        if inspect.iscoroutine(persisted):
+            persisted.close()  # drop, don't await — see the parity note above
+        resume = persisted if isinstance(persisted, str) else opts.resume
+        exchange_opts = replace(opts, resume=resume) if resume is not None else opts
+
+        transport = self._transport_factory(
+            TransportContext(spec=spec, model=model, opts=exchange_opts, resume=resume)
+        )
+
+        # Tap inbound resume tokens off the transport so they persist as the
+        # renamed link entry. The tap is independent of the bridge's own
+        # parsing — both read the same messages; the bridge maps them to the
+        # event stream, this records the reattach token. Best-effort: a
+        # persistence failure never breaks a turn.
+        tap = _make_resume_tap(self._link_store, source, model_id, spec.adapter)
+        if tap is not None:
+            dispose = transport.on_message(tap.on_message)
+            tap.on_close(dispose)
+
+        return bridge.run_exchange(model, context, exchange_opts, transport)
+
+
+def _model_field(model: Any, key: str) -> Any:
+    """Tolerant model-field read (dataclass attribute or mapping key)."""
+    try:
+        return getattr(model, key)
+    except AttributeError:
+        return model.get(key) if hasattr(model, "get") else None
+
+
+# ---------------------------------------------------------------------------
+# Resume tap
+# ---------------------------------------------------------------------------
+
+#: Live fire-and-forget persistence tasks (kept referenced; pruned on settle).
+_PENDING_SAVES: set[asyncio.Task[None]] = set()
+
+
+async def _swallow(awaitable: Any) -> None:
+    try:
+        await awaitable
+    except Exception:
+        # best-effort: a persist failure must not break the turn.
+        pass
+
+
+def _settle_quietly(result: Any) -> None:
+    """Settle an (optionally) async ``save`` result without ever letting it
+    break the turn — the Python analogue of the TS ``r.catch(() => {})``."""
+    if not inspect.isawaitable(result):
+        return
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        if inspect.iscoroutine(result):
+            result.close()
+        return
+    task = loop.create_task(_swallow(result))
+    _PENDING_SAVES.add(task)
+    task.add_done_callback(_PENDING_SAVES.discard)
+
+
+def _make_resume_tap(
+    store: RuntimeLinkStore | None,
+    source: str,
+    model_id: str,
+    bridge: str,
+) -> Any | None:
+    """Build a best-effort resume-token tap that persists a
+    :class:`~induscode.runtime_bridge.contract.RuntimeLink` when a child
+    message carries one (TS ``makeResumeTap``). Returns ``None`` when no
+    store is wired (nothing to persist). The tap recognizes the
+    cross-dialect token shapes the bridges also surface as a ``resume``
+    normalized event: a CLI ``session_id`` (Anthropic ``system/init``), a
+    ``thread_id`` (OpenAI ``thread.started``), and an explicit
+    ``resumeToken`` (peer ``session/resume``). It is intentionally
+    permissive — it only records a token; misreads cost nothing because
+    reattach is keyed on :func:`runtime_source_key` and a stale token simply
+    yields a fresh session."""
+    if store is None:
+        return None
+    saved = False
+    disposer: Callable[[], None] | None = None
+
+    def persist(resume_token: str) -> None:
+        nonlocal saved
+        if saved:
+            return
+        saved = True
+        link = RuntimeLink(
+            source=source,
+            bridge=bridge,
+            resumeToken=resume_token,
+            at=datetime.now(timezone.utc).isoformat(),
+        )
+        try:
+            _settle_quietly(store.save(link))
+        except Exception:
+            # best-effort
+            pass
+        # One token per exchange is enough; drop the subscription once captured.
+        if disposer is not None:
+            disposer()
+
+    def on_message(message: Any) -> None:
+        token = _resume_token_of(message.payload)
+        if token is not None:
+            persist(token)
+
+    def on_close(dispose: Callable[[], None]) -> None:
+        nonlocal disposer
+        disposer = dispose
+
+    return SimpleNamespace(on_message=on_message, on_close=on_close)
+
+
+def _resume_token_of(payload: Any) -> str | None:
+    """Extract a resume token from an opaque child payload across the wire
+    dialects (TS ``resumeTokenOf``)."""
+    if not isinstance(payload, dict):
+        return None
+
+    # Peer JSON-RPC: { method: "session/resume", params: { resumeToken } }.
+    if payload.get("method") == "session/resume":
+        params = payload.get("params")
+        if isinstance(params, dict):
+            token = params.get("resumeToken")
+            if isinstance(token, str):
+                return token
+        return None
+
+    # Anthropic stream-json: { type: "system", subtype: "init", session_id }.
+    if payload.get("type") == "system" and payload.get("subtype") == "init":
+        session_id = payload.get("session_id")
+        if isinstance(session_id, str):
+            return session_id
+
+    # OpenAI --json: { type: "thread.started", thread_id }.
+    if payload.get("type") == "thread.started":
+        thread_id = payload.get("thread_id")
+        if isinstance(thread_id, str):
+            return thread_id
+
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Factory
+# ---------------------------------------------------------------------------
+
+
+def create_runtime_broker(deps: RuntimeBrokerDeps | None = None) -> RuntimeBrokerRuntime:
+    """Construct a :class:`RuntimeBrokerRuntime` (TS ``createRuntimeBroker``).
+    The single sanctioned way to obtain a broker: build one with optional
+    dependencies (a transport factory, a resume link store, a
+    framework-stream override, pre-registered bridges), then ``register``
+    further bridges and drive turns through ``exchange`` / ``route``.
+
+    :param deps: optional construction-time dependencies
+    """
+    return _Broker(deps)