runkite-runner 0.1.0__py3-none-any.whl

runkite_runner/__init__.py

@@ -0,0 +1,3 @@
+"""Runkite Python Runner SDK - Execute LangGraph agents against the Runkite control plane."""
+
+__version__ = "0.1.0"

runkite_runner/__main__.py

@@ -0,0 +1,5 @@
+"""Allow running as: python -m runkite_runner --config langgraph.json"""
+
+from .worker import main
+
+main()

runkite_runner/a2a.py

@@ -0,0 +1,133 @@
+"""Agent-to-Agent (A2A) delegation client: agent calls agent via the same
+Agent Protocol API.
+
+`call_agent` is what a running agent's own node code calls to invoke
+another agent as a sub-task -- it POSTs to the control plane's
+`/internal/a2a/runs` endpoint (see internal/api/a2a.go for the full
+server-side design: auth propagation, recursion limits, cost
+attribution via root_run_id).
+
+Deliberately takes the graph node's own `config` dict, not separate
+run_id/user parameters -- every value this needs (the current run_id to
+set as parent_run_id, the authenticated user to forward as
+on_behalf_of) is already there, set by `build_run_config` in worker.py
+for every run. A node function just calls
+``await call_agent(config, "other_agent", {"messages": [...]})``.
+
+Operational note: with ``wait=True`` (the default), the parent run's
+worker slot stays occupied until the child finishes. A runner process
+with ``concurrency=1`` therefore deadlocks on nested A2A -- the child
+job cannot be dequeued until the parent frees its slot. Use
+concurrency >= 2 (or ``wait=False`` + poll) for any graph that
+delegates synchronously.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import httpx
+
+from .tls_utils import httpx_tls_kwargs
+
+
+class A2AError(Exception):
+    """Raised when the control plane rejects a delegation call --
+    e.g. HTTP 400 for a recursion depth limit exceeded (see
+    ErrA2ADepthExceeded in internal/api/server.go), or 404 for an
+    unknown parent_run_id."""
+
+
+async def call_agent(
+    config: dict,
+    agent_id: str,
+    input: dict[str, Any],
+    *,
+    wait: bool = True,
+    thread_id: str | None = None,
+    run_config: dict[str, Any] | None = None,
+    control_plane_url: str | None = None,
+    timeout: float | None = None,
+) -> dict[str, Any]:
+    """Invoke another agent as a sub-task from within a running agent's
+    own node code.
+
+    Args:
+        config: the RunnableConfig LangGraph passes to every node --
+            must be the same `config` the calling node itself received
+            (or its unmodified `configurable` sub-dict), so run_id/user
+            can be forwarded correctly.
+        agent_id: which agent to delegate to.
+        input: the sub-agent's input, same shape as a normal run's input.
+        wait: block until the sub-agent's run reaches a terminal status
+            and return its result (default). False fires the sub-run
+            and returns immediately with just the created run's data --
+            the caller is then responsible for polling/streaming it
+            itself, same as any other async run.
+        thread_id: an existing thread to run on, e.g. to continue a
+            specific sub-conversation. Defaults to a fresh thread per
+            call -- most delegation calls are one-shot sub-tasks, not
+            multi-turn conversations with the sub-agent.
+        run_config: passed through as the sub-run's own `config`
+            (distinct from this function's own `config` argument, which
+            is the CALLER's LangGraph RunnableConfig, not the sub-run's).
+        control_plane_url: defaults to RUNKITE_HTTP_URL, same env var
+            convention as every other control-plane HTTP call in this
+            runner.
+        timeout: httpx request timeout in seconds. None (default) means
+            no timeout -- a `wait=True` call blocks for however long the
+            sub-agent actually takes, which the caller controls via its
+            own input, not an arbitrary client-side cutoff.
+
+    Returns:
+        When wait=True: {"run": {...}, "values": {...}} -- the sub-run's
+        final state and output, same shape as the client-facing
+        `/runs/{id}/wait` response.
+        When wait=False: just the created run object, status "pending".
+
+    Raises:
+        A2AError: the control plane rejected the call (e.g. recursion
+            depth exceeded, unknown parent_run_id).
+        RuntimeError: config has no run_id -- this wasn't called from
+            within an actual graph node's execution (or with a config
+            that build_run_config never touched).
+    """
+    configurable = config.get("configurable", {}) if config else {}
+    parent_run_id = configurable.get("run_id")
+    if not parent_run_id:
+        raise RuntimeError(
+            "call_agent: config has no configurable.run_id -- must be called "
+            "with the RunnableConfig a graph node itself received, not a "
+            "hand-built or empty one"
+        )
+
+    body: dict[str, Any] = {
+        "agent_id": agent_id,
+        "input": input,
+        "parent_run_id": parent_run_id,
+        "wait": wait,
+    }
+    if thread_id:
+        body["thread_id"] = thread_id
+    if run_config:
+        body["config"] = run_config
+
+    user = configurable.get("langgraph_auth_user")
+    if user is not None and hasattr(user, "to_dict"):
+        on_behalf_of = user.to_dict()
+        if on_behalf_of:
+            body["on_behalf_of"] = on_behalf_of
+
+    base_url = control_plane_url or os.environ.get("RUNKITE_HTTP_URL", "http://localhost:2026")
+    headers: dict[str, str] = {}
+    runner_token = os.environ.get("RUNNER_TOKEN")
+    if runner_token:
+        headers["X-Runner-Kind"] = "python-langgraph"
+        headers["X-Runner-Token"] = runner_token
+
+    async with httpx.AsyncClient(timeout=timeout, **httpx_tls_kwargs()) as client:
+        resp = await client.post(f"{base_url}/internal/a2a/runs", json=body, headers=headers)
+        if resp.status_code >= 400:
+            raise A2AError(f"call_agent({agent_id!r}) failed: {resp.status_code} {resp.text}")
+        return resp.json()

runkite_runner/checkpoint.py

@@ -0,0 +1,214 @@
+"""Checkpoint dual mode.
+
+Direct mode (default when POSTGRES_DSN is set -- production, shared DB with
+the control plane): the runner holds its own connection and writes
+checkpoints with LangGraph's native AsyncPostgresSaver. Zero added latency,
+survives runner restarts. Correct only when the control plane also uses
+POSTGRES_DSN against the same database; MySQL/Mongo/SQLite control planes
+must unset POSTGRES_DSN on the runner (see README Checkpoint dual mode).
+
+Local mode (no POSTGRES_DSN -- zero-dependency dev default): falls back to
+LangGraph's in-memory MemorySaver. This is honestly ephemeral -- state does
+NOT survive a runner restart. That's an accepted trade-off for the
+zero-dependency default (same spirit as the control plane's own in-process
+transport), not a hidden gap. Proxy mode (opaque-blob checkpoints via the
+control plane's HTTP API, for non-Python runners or runners without DB
+credentials) is not implemented by this Python runner -- it always has
+direct DB access when Postgres is available, so proxy mode has no benefit
+here; it exists in the protocol for other-language runners.
+
+Connection pooling (runner-side concurrency): a runner process can now
+process multiple jobs at once (see worker.py's --concurrency), so `start`
+takes a `pool_size` and builds an `AsyncConnectionPool` instead of the
+single connection `AsyncPostgresSaver.from_conn_string` opens -- otherwise
+every concurrent job's checkpoint I/O would serialize on that one
+connection's internal lock (correct, but not actually parallel).
+`AsyncPostgresSaver.__init__`'s `conn` parameter accepts a pool directly
+(`Conn = AsyncConnection | AsyncConnectionPool` in langgraph's own
+`checkpoint/postgres/_ainternal.py`) -- confirmed it checks out a
+connection per operation via that same module's `get_connection` helper,
+so this is a supported usage, not a hack.
+
+Concurrent-startup migration race (found live): `AsyncPostgresSaver.setup()` runs `CREATE TABLE IF NOT EXISTS` DDL
+for its own `checkpoint_migrations` table, which is not actually race-free
+on a truly fresh database -- the same class of bug this project's own
+`internal/state/postgres/postgres.go` had and fixed with a session advisory
+lock. When 2+ runner replicas start simultaneously against a fresh
+Postgres, one `setup()` call can lose that race and crash with
+`duplicate key value violates unique constraint "pg_type_typname_nsp_index"`.
+
+Fixed here with an advisory lock too, but `pg_try_advisory_lock` polled in
+a loop, NOT a blocking `pg_advisory_lock` -- tried the blocking version
+first and hit a real deadlock, not just a slower path: `setup()` also runs
+`CREATE INDEX CONCURRENTLY`, which must wait for every *other* backend's
+in-flight statement in the whole database to finish before it can proceed
+(a Postgres-wide barrier, unrelated to which table the other statement
+touches). A losing replica blocked inside a single `SELECT
+pg_advisory_lock(...)` call counts as an in-flight statement -- so the
+winner's `CREATE INDEX CONCURRENTLY` waits on the losers, and the losers
+wait on the winner to finish `setup()` and unlock. Circular wait,
+confirmed live via `pg_stat_activity`/`pg_locks` (losers stuck on
+`Lock/advisory`, winner stuck on `Lock/virtualxid` waiting on them).
+Polling `pg_try_advisory_lock` avoids this because each poll is its own
+complete, instantly-committed statement -- a losing replica is never
+mid-statement between polls, so it never blocks the winner's
+`CREATE INDEX CONCURRENTLY`.
+"""
+
+import asyncio
+import logging
+
+logger = logging.getLogger("runkite.runner")
+
+# Distinct from the Go control plane's own schema-init advisory lock key
+# (894127001, internal/state/postgres/postgres.go) -- unrelated schemas
+# (runner checkpoint tables vs. control-plane state tables), no reason for
+# one to block the other.
+_CHECKPOINT_SETUP_ADVISORY_LOCK_KEY = 894127002
+_LOCK_POLL_INTERVAL_S = 0.2
+_LOCK_POLL_TIMEOUT_S = 60.0
+
+
+class CheckpointerManager:
+    """Owns the runner's single shared checkpointer for its whole lifetime.
+
+    One checkpointer instance is created at worker startup and attached to
+    every loaded graph (overriding whatever checkpointer, if any, the
+    graph module itself compiled with) -- so checkpoint mode is a runner
+    concern, not something agent authors need to configure in their own
+    graph.py. Agent authors get this for free: zero changes to graph code.
+    """
+
+    def __init__(self):
+        self._checkpointer = None
+        self._pool = None  # AsyncConnectionPool, kept open for the runner's lifetime (postgres mode only)
+        self._dsn: str | None = None
+        self._pool_size: int = 4
+        self._attached: list = []  # graphs whose .checkpointer we own
+        self.mode = "none"
+
+    async def start(self, postgres_dsn: str | None, pool_size: int = 4):
+        if postgres_dsn:
+            import psycopg
+
+            self._dsn = postgres_dsn
+            self._pool_size = pool_size
+            await self._open_pool()
+
+            # Serialize setup() across concurrently-starting runner replicas so
+            # its CREATE TABLE IF NOT EXISTS / CREATE INDEX CONCURRENTLY DDL
+            # can't race on a fresh DB (see module docstring). Deliberately a
+            # standalone connection, NOT checked out from self._pool: setup()
+            # itself checks out a connection from that same pool internally,
+            # and with --concurrency 1 (pool max_size=1) holding the lock from
+            # inside the pool would starve setup()'s own checkout of the only
+            # connection available.
+            async with await psycopg.AsyncConnection.connect(postgres_dsn, autocommit=True) as lock_conn:
+                waited = 0.0
+                while True:
+                    row = await (
+                        await lock_conn.execute(
+                            "SELECT pg_try_advisory_lock(%s)", (_CHECKPOINT_SETUP_ADVISORY_LOCK_KEY,)
+                        )
+                    ).fetchone()
+                    if row and row[0]:
+                        break
+                    if waited >= _LOCK_POLL_TIMEOUT_S:
+                        raise TimeoutError(
+                            "timed out waiting for checkpoint setup() advisory lock "
+                            f"(held by another runner replica for over {_LOCK_POLL_TIMEOUT_S}s)"
+                        )
+                    await asyncio.sleep(_LOCK_POLL_INTERVAL_S)
+                    waited += _LOCK_POLL_INTERVAL_S
+                try:
+                    await self._checkpointer.setup()
+                finally:
+                    await lock_conn.execute("SELECT pg_advisory_unlock(%s)", (_CHECKPOINT_SETUP_ADVISORY_LOCK_KEY,))
+            self.mode = "direct-postgres"
+            logger.info(
+                "checkpoint mode: direct (postgres, pool_size=%s) -- LangGraph tables on "
+                "POSTGRES_DSN; requires the control plane to use the same Postgres database "
+                "(Supported profile). If the control plane is MySQL/Mongo/SQLite, unset "
+                "POSTGRES_DSN on this runner and set RUNKITE_HTTP_URL for store proxy mode.",
+                pool_size,
+            )
+        else:
+            from langgraph.checkpoint.memory import MemorySaver
+
+            self._checkpointer = MemorySaver()
+            self.mode = "memory"
+            logger.warning(
+                "checkpoint mode: in-memory (no POSTGRES_DSN set) -- "
+                "thread state will NOT survive a runner restart. "
+                "Set POSTGRES_DSN for production persistence."
+            )
+
+    async def _open_pool(self) -> None:
+        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+        from psycopg.rows import dict_row
+
+        from . import pg_pool
+
+        # Same connection kwargs from_conn_string used
+        # (autocommit/prepare_threshold/row_factory) -- AsyncPostgresSaver
+        # expects dict-row results, not psycopg's tuple default.
+        self._pool = pg_pool.make(
+            self._dsn,
+            max_size=self._pool_size,
+            conn_kwargs={"prepare_threshold": 0, "row_factory": dict_row},
+        )
+        await self._pool.open()
+        self._checkpointer = AsyncPostgresSaver(conn=self._pool)
+        for graph in self._attached:
+            graph.checkpointer = self._checkpointer
+
+    async def recreate_pool(self) -> None:
+        """Drop a wedged pool (idle overnight / laptop sleep) and open a new one.
+
+        Rebinds every previously attach()'d graph so LangGraph does not keep
+        calling through the closed pool.
+        """
+        if self._dsn is None:
+            return
+        old = self._pool
+        self._pool = None
+        if old is not None:
+            try:
+                await old.close()
+            except Exception:
+                logger.exception("error closing wedged checkpoint pool")
+        await self._open_pool()
+
+    async def recover_if_wedged(self) -> None:
+        """Cheap pre-job probe: if getconn hangs/fails, recreate once.
+
+        AsyncPostgresSaver owns the pool reference internally, so unlike
+        store.abatch we cannot catch PoolTimeout inside each checkpoint
+        op -- probe here before astream so an overnight-wedged pool does
+        not burn a full run on a 30s timeout mid-graph.
+        """
+        if self._pool is None:
+            return
+        from psycopg_pool import PoolTimeout
+
+        try:
+            async with self._pool.connection(timeout=5.0) as conn:
+                await conn.execute("SELECT 1")
+        except PoolTimeout:
+            logger.warning("checkpoint pool timed out on health probe; recreating pool")
+            await self.recreate_pool()
+        except Exception:
+            logger.warning("checkpoint pool health probe failed; recreating pool", exc_info=True)
+            await self.recreate_pool()
+
+    async def stop(self):
+        if self._pool is not None:
+            await self._pool.close()
+            self._pool = None
+
+    def attach(self, graph):
+        """Override a compiled graph's checkpointer with the shared one."""
+        if graph not in self._attached:
+            self._attached.append(graph)
+        graph.checkpointer = self._checkpointer
+        return graph

runkite_runner/custom_app.py

@@ -0,0 +1,62 @@
+"""In-runner mode for the Custom routes platform extension.
+
+Loads a user-defined ASGI app (FastAPI, Starlette, or any other ASGI
+framework) from langgraph.json's "custom_app" section and serves it via
+uvicorn, running as a concurrent asyncio task alongside the runner's own
+gRPC worker loop (see run_worker in worker.py) -- same process, same event
+loop, similar to dropping a file into your project.
+
+The control plane reverse-proxies /custom/* to wherever this ends up
+listening (see internal/config.CustomRoutesEntry / cmd/serve.go's
+initCustomRoutesProxy) -- from the control plane's side, in-runner mode and
+a separately-run sidecar are the exact same mechanism, just different
+processes hosting the target URL.
+
+WSGI frameworks (e.g. Flask) aren't directly supported -- uvicorn only
+serves ASGI. Wrap a WSGI app with an adapter (e.g. a2wsgi.WSGIMiddleware)
+if that's the framework of choice; this loader doesn't care what kind of
+object it gets back as long as it's ASGI-callable.
+"""
+
+import importlib.util
+import logging
+import sys
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger("runkite.runner")
+
+
+def load_asgi_app(config_dir: Path, module_ref: str) -> Any:
+    """Loads an ASGI app object from a "path/to/module.py:app_symbol" ref --
+    the same "path:symbol" convention langgraph.json's "graphs" section
+    already uses for agent graphs."""
+    file_path, export_name = module_ref.split(":", 1)
+    abs_path = (config_dir / file_path).resolve()
+
+    spec = importlib.util.spec_from_file_location("runkite_custom_app", str(abs_path))
+    if spec is None or spec.loader is None:
+        raise ValueError(f"Cannot load custom app module: {abs_path}")
+
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+
+    return getattr(module, export_name)
+
+
+async def serve_custom_app(app: Any, host: str, port: int) -> None:
+    """Runs an ASGI app via uvicorn until cancelled. Meant to run as a
+    concurrent asyncio task (asyncio.create_task) alongside the gRPC
+    worker's own poll loop -- both share the process and event loop, so a
+    slow/blocking route handler in the user's app can, in principle, delay
+    the worker's own async work. That's an inherent trade-off of "in-runner,
+    same process" mode; use sidecar mode instead for routes that need
+    independent scaling or isolation.
+    """
+    import uvicorn
+
+    config = uvicorn.Config(app, host=host, port=port, log_level="info", access_log=True)
+    server = uvicorn.Server(config)
+    logger.info(f"Custom routes (in-runner mode): serving on http://{host}:{port}")
+    await server.serve()

runkite_runner/custom_auth.py

@@ -0,0 +1,82 @@
+"""Portable identity helpers for custom-route ASGI apps.
+
+The control plane authenticates the caller, then injects X-Runkite-*
+headers on the reverse-proxied request (see internal/customroutes).
+Custom apps SHOULD trust these headers rather than re-parsing JWT /
+API keys — and MUST NOT trust them on any path that bypasses the
+control plane.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import Any
+
+HEADER_IDENTITY = "x-runkite-identity"
+HEADER_TENANT_ID = "x-runkite-tenant-id"
+HEADER_PERMISSIONS = "x-runkite-permissions"
+HEADER_DISPLAY_NAME = "x-runkite-display-name"
+HEADER_USER_JSON = "x-runkite-user"
+
+
+@dataclass(frozen=True)
+class CustomUser:
+    """Identity forwarded by the control plane into a custom route."""
+
+    identity: str
+    tenant_id: str = "default"
+    permissions: tuple[str, ...] = ()
+    display_name: str = ""
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def is_authenticated(self) -> bool:
+        return bool(self.identity)
+
+
+def _header(headers: Mapping[str, str], name: str) -> str:
+    # ASGI / Starlette headers are lower-case; also accept canonical form.
+    if hasattr(headers, "get"):
+        v = headers.get(name) or headers.get(name.title()) or headers.get(name.upper())
+        if v is not None:
+            return v if isinstance(v, str) else str(v)
+    return ""
+
+
+def user_from_headers(headers: Mapping[str, str]) -> CustomUser | None:
+    """Build a CustomUser from X-Runkite-* headers. Returns None if absent."""
+    raw = _header(headers, HEADER_USER_JSON)
+    if raw:
+        try:
+            data = json.loads(raw)
+            if isinstance(data, dict) and data.get("identity"):
+                perms = data.get("permissions") or []
+                if isinstance(perms, str):
+                    perms = [p for p in perms.split(",") if p]
+                return CustomUser(
+                    identity=str(data["identity"]),
+                    tenant_id=str(data.get("tenant_id") or "default"),
+                    permissions=tuple(str(p) for p in perms),
+                    display_name=str(data.get("display_name") or ""),
+                )
+        except json.JSONDecodeError:
+            pass
+
+    identity = _header(headers, HEADER_IDENTITY)
+    if not identity:
+        return None
+    perms_raw = _header(headers, HEADER_PERMISSIONS)
+    perms = tuple(p for p in (perms_raw.split(",") if perms_raw else []) if p)
+    return CustomUser(
+        identity=identity,
+        tenant_id=_header(headers, HEADER_TENANT_ID) or "default",
+        permissions=perms,
+        display_name=_header(headers, HEADER_DISPLAY_NAME),
+    )
+
+
+def user_from_request(request: Any) -> CustomUser | None:
+    """Starlette/FastAPI convenience: user_from_headers(request.headers)."""
+    return user_from_headers(request.headers)

runkite_runner/custom_helpers.py

@@ -0,0 +1,76 @@
+"""Thin HTTP helpers for custom-route handlers talking back to the CP.
+
+Custom apps are not in-process with the graph runtime. Use these to call
+public Agent Protocol endpoints (store, threads, runs) with the caller's
+Bearer token forwarded from the inbound request.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from urllib.parse import quote
+
+import httpx
+
+
+class ControlPlaneClient:
+    """Minimal sync client for store/run/thread lookups from custom routes."""
+
+    def __init__(self, base_url: str, authorization: str | None = None, timeout: float = 30.0):
+        self.base_url = base_url.rstrip("/")
+        self._headers: dict[str, str] = {}
+        if authorization:
+            self._headers["Authorization"] = authorization
+        self._timeout = timeout
+
+    @classmethod
+    def from_request(cls, request: Any, base_url: str | None = None) -> ControlPlaneClient:
+        """Build a client that reuses the caller's Authorization header.
+
+        base_url defaults to the Host the custom app was reached through
+        (works when the app is reverse-proxied by the same CP). For
+        sidecar mode, pass the control plane URL explicitly.
+        """
+        import os
+
+        auth = None
+        if hasattr(request, "headers"):
+            auth = request.headers.get("authorization") or request.headers.get("Authorization")
+        url = base_url or os.environ.get("RUNKITE_HTTP_URL") or os.environ.get("RUNKITE_URL")
+        if not url and hasattr(request, "base_url"):
+            # Starlette: request.base_url is the public origin as seen by
+            # the app (often the CP origin when StripPrefix-proxied).
+            url = str(request.base_url).rstrip("/")
+        if not url:
+            raise ValueError("control plane base_url required (pass base_url= or set RUNKITE_HTTP_URL)")
+        return cls(url, authorization=auth)
+
+    def _request(self, method: str, path: str, **kwargs: Any) -> Any:
+        with httpx.Client(base_url=self.base_url, headers=self._headers, timeout=self._timeout) as client:
+            resp = client.request(method, path, **kwargs)
+            resp.raise_for_status()
+            if resp.status_code == 204 or not resp.content:
+                return None
+            return resp.json()
+
+    def get_thread(self, thread_id: str) -> dict:
+        return self._request("GET", f"/threads/{quote(thread_id)}")
+
+    def get_run(self, thread_id: str, run_id: str) -> dict:
+        return self._request("GET", f"/threads/{quote(thread_id)}/runs/{quote(run_id)}")
+
+    def store_get(self, namespace: list[str], key: str) -> dict | None:
+        ns = ",".join(namespace)
+        try:
+            return self._request("GET", "/store/items", params={"namespace": ns, "key": key})
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                return None
+            raise
+
+    def store_put(self, namespace: list[str], key: str, value: dict) -> dict:
+        return self._request(
+            "PUT",
+            "/store/items",
+            json={"namespace": namespace, "key": key, "value": value},
+        )