diff --git a/docs/CONSUMING.md b/docs/CONSUMING.md index 221954b..cc9f2f7 100644 --- a/docs/CONSUMING.md +++ b/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 | | | diff --git a/src/mcp_agent_api/__init__.py b/src/mcp_agent_api/__init__.py index 151bada..184e51d 100644 --- a/src/mcp_agent_api/__init__.py +++ b/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", diff --git a/src/mcp_agent_api/app.py b/src/mcp_agent_api/app.py index 9ad0d30..e35aa93 100644 --- a/src/mcp_agent_api/app.py +++ b/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 diff --git a/src/mcp_agent_api/routes.py b/src/mcp_agent_api/routes.py index 3d5cbe4..0b88d9e 100644 --- a/src/mcp_agent_api/routes.py +++ b/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,7 +272,12 @@ 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 @@ -264,6 +285,17 @@ def create_router(provider: Provider, *, prefix: str = "") -> APIRouter: 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, diff --git a/tests/mcp_agent_api/test_routes.py b/tests/mcp_agent_api/test_routes.py index 28dbca8..7df931d 100644 --- a/tests/mcp_agent_api/test_routes.py +++ b/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."""