diff --git a/sdk-python/copilotkit/langgraph_agui_agent.py b/sdk-python/copilotkit/langgraph_agui_agent.py index d1c10fc637..e9bec5cef4 100644 --- a/sdk-python/copilotkit/langgraph_agui_agent.py +++ b/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 diff --git a/sdk-python/tests/test_agui_agent.py b/sdk-python/tests/test_agui_agent.py index 3f9927a949..c3b61bb49b 100644 --- a/sdk-python/tests/test_agui_agent.py +++ b/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 ---------- diff --git a/sdk-python/tests/test_copilotkit_lg_middleware.py b/sdk-python/tests/test_copilotkit_lg_middleware.py index 52d23c30fa..f4227e6f65 100644 --- a/sdk-python/tests/test_copilotkit_lg_middleware.py +++ b/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)