feat(api): a per-run hook for callbacks and context

docs/CONSUMING.md

@@ -755,6 +755,31 @@ attribute rather than unpacked because `BuiltAgent` is a `NamedTuple` and a
 consumer's equivalent orders the five differently — unpacking one as the other
 yields `required` where `withheld` belongs, and nothing complains.
 
+**`turn_context` is where what wraps a run goes** — tracing callbacks, a
+correlation id, per-request metadata. A context manager, entered inside the
+stream and given the request and both ids; whatever it yields becomes the
+turn's runnable config:
+
+```python
+@contextmanager
+def traced(request, thread_id, run_id):
+    trace = uuid.uuid4().hex
+    with correlation_id(trace):  # a ContextVar your httpx hook reads
+        yield {"callbacks": [handler(trace)], "metadata": {"thread": thread_id}}
+
+
+app.include_router(create_router(provider, turn_context=traced))
+```
+
+A context manager rather than a config factory because the two things a host
+wants here differ in kind: a config is a *value* handed to the turn, while a
+correlation id stamped onto outgoing MCP calls is a **context variable**, which
+has to be set for the duration. Both need to be in force *while the turn runs*,
+not while the handler is on the stack — by the time the first tool is called,
+the handler has long returned. That is why this is entered beside
+`user_credentials` rather than around the route. `create_app` takes the same
+argument and passes it straight through.
+
 ### 5c. The routes
 
 | | |

src/mcp_agent_api/__init__.py

@@ -28,7 +28,7 @@
     agui_events,
     state_metadata,
 )
-from mcp_agent_api.routes import Built, RunRequest, create_router
+from mcp_agent_api.routes import Built, RunRequest, TurnContext, create_router
 
 __all__ = [
     "ANSWER_CITATIONS",
@@ -38,6 +38,7 @@
     "TOOLS_WITHHELD",
     "Built",
     "RunRequest",
+    "TurnContext",
     "agui_events",
     "create_router",
     "state_metadata",

src/mcp_agent_api/app.py

@@ -41,7 +41,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 
 from mcp_agent.main import AgentSettings, Checkpointing, build_agent
-from mcp_agent_api.routes import Built, create_router
+from mcp_agent_api.routes import Built, TurnContext, create_router
 
 #: Comma-separated origins a browser client may call this API from. Empty (the
 #: default) adds no CORS middleware at all, which is right for an API behind
@@ -90,6 +90,7 @@ def create_app(
     origins: Sequence[str] | None = None,
     prefix: str = "",
     checkpoint: str | None = None,
+    turn_context: TurnContext | None = None,
 ) -> FastAPI:
     """The API as an application, with the agent built during startup.
 
@@ -103,7 +104,9 @@ def create_app(
     same-origin deployment carries none of it.
 
     ``prefix`` is passed to :func:`~mcp_agent_api.routes.create_router`, for
-    mounting the whole service under a path.
+    mounting the whole service under a path, and so is ``turn_context`` — a
+    deployment that wants its runs traced needs that seam whether or not it
+    owns the application around them.
     """
     checkpointing = Checkpointing(checkpoint)
 
@@ -143,7 +146,9 @@ def provider() -> Built:
             raise RuntimeError("the agent is still connecting")
         return built
 
-    app.include_router(create_router(provider, prefix=prefix))
+    app.include_router(
+        create_router(provider, prefix=prefix, turn_context=turn_context)
+    )
     return app
 
 

src/mcp_agent_api/routes.py

@@ -36,6 +36,13 @@
 carry those five fields in different orders, and unpacking one as the other
 yields ``required`` where ``withheld`` belongs.
 
+**What wraps a run is the host's.** Tracing callbacks, a correlation id on the
+outbound MCP calls, per-request metadata — none of that belongs here, and all
+of it has to be in force *while the turn runs* rather than while the handler
+is on the stack. ``turn_context`` is the one seam for it: a context manager
+entered inside the stream, yielding the turn's runnable config. Without one,
+nothing changes.
+
 **History is the server's.** Only the trailing user message of a request is
 read; the rest of ``messages`` is ignored, and the checkpointer's transcript is
 the truth. That diverges from AG-UI's client-is-authoritative convention on
@@ -52,6 +59,7 @@
 import asyncio
 import json
 from collections.abc import AsyncIterator, Callable, Iterable, Mapping, Sequence
+from contextlib import AbstractContextManager, nullcontext
 from typing import Any, Protocol, cast
 
 from ag_ui.core import (
@@ -114,6 +122,14 @@ def required(self) -> dict[str, list[str]] | None: ...
 #: "not built yet"; see :func:`create_router`.
 Provider = Callable[[], Built]
 
+#: Wraps one run, given the request and the ids it was assigned. Entered
+#: *inside* the SSE generator, so context variables it sets are in force while
+#: the turn runs; what it yields becomes ``stream_turn``'s runnable config.
+#: See :func:`create_router`.
+TurnContext = Callable[
+    [Request, str, str], AbstractContextManager[dict[str, Any] | None]
+]
+
 
 class RunRequest(BaseModel):
     """What a client posts to ``/runs``.
@@ -256,14 +272,30 @@ def credentials_for(
     return resolve_credentials(required, supplied)
 
 
-def create_router(provider: Provider, *, prefix: str = "") -> APIRouter:
+def create_router(
+    provider: Provider,
+    *,
+    prefix: str = "",
+    turn_context: TurnContext | None = None,
+) -> APIRouter:
     """Routes serving the agent ``provider`` returns.
 
     ``provider`` is called on each request rather than once here, so an agent
     rebuilt behind it — on a model change, on a reconnect — is picked up
     without remounting. It may raise to signal that the agent is not ready;
     that surfaces as ``503``, which is the honest answer while a lifespan is
     still connecting to MCP servers.
+
+    ``turn_context`` is where a host puts what it needs around each run:
+    tracing callbacks, a correlation id, per-request metadata. It is a context
+    manager rather than a plain config factory because the two things a host
+    wants here differ in kind — a config is a *value* passed to the turn, while
+    a correlation id read by an httpx hook at request time is a **context
+    variable**, and that has to be set for the duration rather than handed
+    over. Entered beside :func:`~mcp_agent.main.user_credentials` and for the
+    same reason: by the time the first tool is called, the request handler has
+    long returned, so anything scoped to the handler's stack is already gone.
+    Whatever it yields — ``None`` is fine — becomes the turn's runnable config.
     """
     router = APIRouter(prefix=prefix)
     views = ViewCache(provider)
@@ -313,11 +345,17 @@ async def create_run(body: RunRequest, request: Request) -> StreamingResponse:
         encoder = EventEncoder(accept=request.headers.get("accept", ""))
 
         async def frames() -> AsyncIterator[str]:
-            # Set inside the generator, so it is in force while the turn runs
-            # rather than only while the handler is on the stack — by the time
-            # the first tool is called, this handler has long returned.
-            with user_credentials(credentials or None):
-                turn = stream_turn(agent.agent, question, thread_id)
+            # Both are entered inside the generator, so they are in force while
+            # the turn runs rather than only while the handler is on the stack
+            # — by the time the first tool is called, this handler has long
+            # returned.
+            around: AbstractContextManager[dict[str, Any] | None] = (
+                turn_context(request, thread_id, run_id)
+                if turn_context
+                else nullcontext(None)
+            )
+            with user_credentials(credentials or None), around as config:
+                turn = stream_turn(agent.agent, question, thread_id, config)
                 async for event in agui_events(
                     turn,
                     thread_id=thread_id,

tests/mcp_agent_api/test_routes.py

@@ -8,11 +8,15 @@
 
 import asyncio
 import json
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
 from typing import Any
 
 import httpx
 import pytest
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
+from langchain_core.callbacks import BaseCallbackHandler
 from langchain_core.messages import AIMessage
 from langchain_core.tools import StructuredTool
 from langgraph.checkpoint.memory import InMemorySaver
@@ -270,6 +274,100 @@ async def call() -> str:
     assert seen == [{"x-cds-token": "the user's"}]
 
 
+async def test_the_turn_context_wraps_the_run_and_is_told_its_ids():
+    """It is handed the request and both ids, and it is left on the way out —
+    a wrapper that never exits leaks whatever it opened, once per run."""
+    seen: list[tuple[str, str, str]] = []
+    exited: list[bool] = []
+
+    @contextmanager
+    def around(request: Request, thread_id: str, run_id: str) -> Iterator[None]:
+        seen.append((request.url.path, thread_id, run_id))
+        try:
+            yield None
+        finally:
+            exited.append(True)
+
+    async with _client(turn_context=around) as client:
+        events = await _run(client, threadId="t9")
+
+    assert seen == [("/runs", "t9", events[0]["runId"])]
+    assert exited == [True]
+
+
+async def test_the_turn_context_yields_the_turn_its_config():
+    """What a host traces with is a callback in the runnable config, so the
+    yielded value has to reach the graph run rather than be dropped."""
+    started: list[str] = []
+
+    class Recorder(BaseCallbackHandler):
+        def on_chain_start(
+            self, serialized: dict[str, Any], inputs: dict[str, Any], **kwargs: Any
+        ) -> None:
+            started.append("chain")
+
+    @contextmanager
+    def around(
+        request: Request, thread_id: str, run_id: str
+    ) -> Iterator[dict[str, Any]]:
+        yield {"callbacks": [Recorder()]}
+
+    async with _client(turn_context=around) as client:
+        await _run(client)
+
+    assert started
+
+
+async def test_the_turn_context_is_in_force_while_the_tool_runs():
+    """The reason this is a context manager and not a config factory: a
+    correlation id read by an httpx hook at request time is a context
+    variable, and the tool call that reads it happens long after the handler
+    has returned. Entered where ``user_credentials`` is, for that reason."""
+    correlation: ContextVar[str | None] = ContextVar("correlation", default=None)
+    seen: list[str | None] = []
+
+    async def call() -> str:
+        seen.append(correlation.get())
+        return "ok"
+
+    peek = StructuredTool(
+        name="peek",
+        description="peek",
+        args_schema={"type": "object", "properties": {}},
+        coroutine=call,
+    )
+    agent, _ = with_session_state(
+        StreamingScriptedModel(
+            script=[
+                AIMessage(
+                    content="",
+                    tool_calls=[
+                        {"name": "peek", "args": {}, "id": "c1", "type": "tool_call"}
+                    ],
+                ),
+                AIMessage(content="done"),
+            ]
+        ),
+        [peek],
+        InMemorySaver(),
+    )
+
+    @contextmanager
+    def around(request: Request, thread_id: str, run_id: str) -> Iterator[None]:
+        token = correlation.set(f"trace-{run_id}")
+        try:
+            yield None
+        finally:
+            correlation.reset(token)
+
+    async with _client(
+        _built(agent=agent, tools=[peek]), turn_context=around
+    ) as client:
+        events = await _run(client)
+
+    assert seen == [f"trace-{events[0]['runId']}"]
+
+
 async def test_the_client_does_not_get_to_write_history():
     """Only the trailing user message is read. A client posting a transcript of
     its own would otherwise be able to put words in the thread."""