fix(sdk-python): capture subgraph context from run input (#3886)

sdk-python/copilotkit/langgraph_agui_agent.py

@@ -1,3 +1,4 @@
+import inspect
 import json
 import logging
 from enum import Enum
@@ -70,6 +71,7 @@ class LangGraphAGUIAgent(LangGraphAgent):
     ):
         super().__init__(name=name, graph=graph, description=description, config=config)
         self.constant_schema_keys = self.constant_schema_keys + ["copilotkit"]
+        self._copilotkit_runtime_payload: dict[str, Any] | None = None
 
     def _dispatch_event(self, event) -> str:
         """Override the dispatch event method to handle custom CopilotKit events and filtering.
@@ -248,9 +250,15 @@ class LangGraphAGUIAgent(LangGraphAgent):
 
     async def run(self, input):
         """Override run to filter out None events from _dispatch_event filtering."""
-        async for event in super().run(input):
-            if event is not None:
-                yield event
+        self._copilotkit_runtime_payload = self._serialize_copilotkit_runtime_payload(
+            input
+        )
+        try:
+            async for event in super().run(input):
+                if event is not None:
+                    yield event
+        finally:
+            self._copilotkit_runtime_payload = None
 
     async def _handle_single_event(
         self, event: Any, state: State
@@ -295,13 +303,29 @@ class LangGraphAGUIAgent(LangGraphAgent):
         fork: Optional[Any] = None,
     ) -> Dict[str, Any]:
         """Thread CopilotKit payload through LangGraph runtime context for subgraphs."""
+        supports_context = (
+            "context" in inspect.signature(self.graph.astream_events).parameters
+        )
         merged_context = dict(context or {})
-        existing_copilotkit = merged_context.get("copilotkit") or {}
-        merged_context["copilotkit"] = {
-            **existing_copilotkit,
-            **self._serialize_copilotkit_runtime_payload(input),
-        }
-        return super().get_stream_kwargs(
+        captured_payload = self._copilotkit_runtime_payload
+        if captured_payload is not None:
+            if supports_context:
+                existing_copilotkit = merged_context.get("copilotkit") or {}
+                merged_context["copilotkit"] = {
+                    **existing_copilotkit,
+                    **captured_payload,
+                }
+            else:
+                next_config = dict(config or {})
+                configurable = dict(next_config.get("configurable") or {})
+                existing_copilotkit = configurable.get("copilotkit") or {}
+                configurable["copilotkit"] = {
+                    **existing_copilotkit,
+                    **captured_payload,
+                }
+                next_config["configurable"] = configurable
+                config = next_config
+        stream_kwargs = super().get_stream_kwargs(
             input=input,
             subgraphs=subgraphs,
             version=version,
@@ -309,6 +333,7 @@ class LangGraphAGUIAgent(LangGraphAgent):
             context=merged_context,
             fork=fork,
         )
+        return stream_kwargs
 
     def langgraph_default_merge_state(
         self, state: State, messages: List[BaseMessage], input: Any

sdk-python/tests/test_agui_agent.py

@@ -457,8 +457,8 @@ class TestLanggraphDefaultMergeState:
         assert action_names == ["first", "second", "third"]
 
 
-class TestGetStreamKwargs:
-    """get_stream_kwargs threads CopilotKit payload into LangGraph context."""
+class TestRunRuntimeContextBridge:
+    """run captures the request before the base agent derives stream input."""
 
     class _GraphWithContext:
         nodes = {}
@@ -478,7 +478,24 @@ class TestGetStreamKwargs:
             if False:
                 yield input, subgraphs, version, config, context
 
-    def test_copilotkit_payload_added_to_runtime_context(self):
+    class _GraphWithoutContext:
+        nodes = {}
+
+        def get_state(self):
+            return None
+
+        async def astream_events(
+            self,
+            *,
+            input=None,
+            subgraphs=False,
+            version="v2",
+            config=None,
+        ):
+            if False:
+                yield input, subgraphs, version, config
+
+    def test_run_captures_payload_for_derived_stream_input(self):
         agent = LangGraphAGUIAgent(name="test", graph=self._GraphWithContext())
         run_input = RunAgentInput(
             thread_id="t-1",
@@ -500,29 +517,89 @@ class TestGetStreamKwargs:
             forwarded_props={},
         )
 
-        kwargs = agent.get_stream_kwargs(
-            input=run_input,
-            subgraphs=True,
-            version="v2",
-            config={"configurable": {"thread_id": "t-1", "x-trace": "abc"}},
+        captured = {}
+
+        async def fake_base_run(_self, _input):
+            captured.update(
+                _self.get_stream_kwargs(
+                    input={"messages": []},
+                    subgraphs=True,
+                    version="v2",
+                    config={"configurable": {"thread_id": "t-1", "x-trace": "abc"}},
+                )
+            )
+            if False:
+                yield None
+
+        with patch.object(AGUIBase, "run", new=fake_base_run):
+            import asyncio
+
+            asyncio.run(self._consume(agent.run(run_input)))
+
+        assert captured["context"]["thread_id"] == "t-1"
+        assert captured["context"]["x-trace"] == "abc"
+        assert captured["context"]["copilotkit"]["actions"][0]["name"] == (
+            "frontend_lookup"
+        )
+        assert captured["context"]["copilotkit"]["context"] == [
+            {"description": "viewer role", "value": "admin"}
+        ]
+
+    def test_run_does_not_force_context_for_graphs_without_context_support(self):
+        agent = LangGraphAGUIAgent(name="test", graph=self._GraphWithoutContext())
+        run_input = RunAgentInput(
+            thread_id="t-1",
+            run_id="r-1",
+            state={},
+            messages=[],
+            tools=[
+                {
+                    "name": "frontend_lookup",
+                    "description": "frontend tool",
+                }
+            ],
+            context=[
+                {
+                    "description": "viewer role",
+                    "value": "admin",
+                }
+            ],
+            forwarded_props={},
         )
 
-        assert kwargs["context"]["thread_id"] == "t-1"
-        assert kwargs["context"]["x-trace"] == "abc"
-        assert kwargs["context"]["copilotkit"]["actions"] == [
-            {
-                "name": "frontend_lookup",
-                "description": "frontend tool",
-                "parameters": None,
-            }
-        ]
-        assert kwargs["context"]["copilotkit"]["context"] == [
-            {
-                "description": "viewer role",
-                "value": "admin",
-            }
+        captured = {}
+
+        async def fake_base_run(_self, _input):
+            captured.update(
+                _self.get_stream_kwargs(
+                    input={"messages": []},
+                    subgraphs=True,
+                    version="v2",
+                    config={"configurable": {"thread_id": "t-1", "x-trace": "abc"}},
+                )
+            )
+            if False:
+                yield None
+
+        with patch.object(AGUIBase, "run", new=fake_base_run):
+            import asyncio
+
+            asyncio.run(self._consume(agent.run(run_input)))
+
+        assert "context" not in captured
+        assert captured["config"]["configurable"]["thread_id"] == "t-1"
+        assert captured["config"]["configurable"]["x-trace"] == "abc"
+        assert captured["config"]["configurable"]["copilotkit"]["actions"][0][
+            "name"
+        ] == ("frontend_lookup")
+        assert captured["config"]["configurable"]["copilotkit"]["context"] == [
+            {"description": "viewer role", "value": "admin"}
         ]
 
+    @staticmethod
+    async def _consume(events):
+        return [event async for event in events]
+
 
 # ---------- Reasoning content preservation ----------
 

sdk-python/tests/test_copilotkit_lg_middleware.py

@@ -30,6 +30,7 @@ from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
+from ag_ui.core import RunAgentInput, UserMessage
 from pydantic import Field
 from typing_extensions import TypedDict
 from langchain.agents import create_agent
@@ -43,6 +44,7 @@ from langchain_core.messages import (
 )
 from langchain_core.outputs import ChatGeneration, ChatResult
 from langgraph.constants import END, START
+from langgraph.checkpoint.memory import MemorySaver
 from langchain.agents.middleware import ModelRequest
 from langgraph.graph import StateGraph
 
@@ -51,6 +53,7 @@ from copilotkit.copilotkit_lg_middleware import (
     _extract_forwarded_headers_from_config,
 )
 from copilotkit.header_propagation import get_forwarded_headers, set_forwarded_headers
+from copilotkit.langgraph_agui_agent import LangGraphAGUIAgent
 
 
 # ---------------------------------------------------------------------------
@@ -330,40 +333,52 @@ def test_wrap_model_call_injects_frontend_tools_from_context_bridge(monkeypatch)
 
 
 def test_real_subgraph_context_bridge_reaches_child_agent():
-    """A real create_agent subgraph sees frontend tools and app context via runtime context."""
+    """A real AG-UI run carries tools and app context into a child agent."""
     model = _RecordingToolAwareChatModel()
+    middleware = CopilotKitMiddleware()
     child_agent = create_agent(
         model=model,
         tools=[],
-        middleware=[CopilotKitMiddleware()],
+        middleware=[middleware],
+        context_schema=_ParentContext,
     )
     parent = StateGraph(_ParentState, context_schema=_ParentContext)
     parent.add_node("child", child_agent)
     parent.add_edge(START, "child")
     parent.add_edge("child", END)
-    compiled = parent.compile()
-
-    result = compiled.invoke(
-        {"messages": [HumanMessage("hi")]},
-        context={
-            "copilotkit": {
-                "actions": [
-                    {
-                        "name": "frontend_lookup",
-                        "description": "frontend tool",
-                    }
-                ],
-                "context": [
-                    {
-                        "description": "viewer role",
-                        "value": "admin",
-                    }
-                ],
-            }
-        },
+    agent = LangGraphAGUIAgent(
+        name="parent", graph=parent.compile(checkpointer=MemorySaver())
     )
 
-    assert result["messages"][-1].content == "ok"
+    async def consume_run():
+        return [
+            event
+            async for event in agent.run(
+                RunAgentInput(
+                    thread_id="t-1",
+                    run_id="r-1",
+                    state={},
+                    messages=[UserMessage(id="m-1", content="hi")],
+                    tools=[
+                        {
+                            "name": "frontend_lookup",
+                            "description": "frontend tool",
+                        }
+                    ],
+                    context=[
+                        {
+                            "description": "viewer role",
+                            "value": "admin",
+                        }
+                    ],
+                    forwarded_props={},
+                )
+            )
+        ]
+
+    asyncio.run(consume_run())
+
+    assert model.last_messages, "child model should receive the parent run"
     assert [tool.get("name") for tool in model.bound_tools] == ["frontend_lookup"]
     system_messages = [
         msg for msg in model.last_messages if isinstance(msg, SystemMessage)