diff --git a/python/tests/integrations/deepagents_tests/test_deepagents_integration.py b/python/tests/integrations/deepagents_tests/test_deepagents_integration.py index 4bc13cecf..f9de8ce37 100644 --- a/python/tests/integrations/deepagents_tests/test_deepagents_integration.py +++ b/python/tests/integrations/deepagents_tests/test_deepagents_integration.py @@ -329,6 +329,51 @@ async def async_handler(next_request: ToolCallRequest) -> ToolMessage: assert kwargs["handle"] is parent_handle +def test_skill_load_mark_survives_deepagents_middleware( + subscribed_events: list[nemo_relay.Event], + deepagents_integration_module: types.ModuleType, +): + from langchain.agents.middleware import ToolCallRequest + from langchain_core.messages import ToolMessage + + middleware = deepagents_integration_module.NemoRelayDeepAgentsMiddleware() + request = ToolCallRequest( + tool_call={ + "name": "read_file", + "args": {"path": "/skills/review/SKILL.md"}, + "id": "call-skill", + }, + tool=None, + state={}, + runtime=MagicMock(), + ) + + with nemo_relay.scope.scope("deepagents-skill", nemo_relay.ScopeType.Agent): + response = middleware.wrap_tool_call( + request, + lambda next_request: ToolMessage( + content="loaded", + tool_call_id=next_request.tool_call["id"], + ), + ) + + nemo_relay.subscribers.flush() + assert response.content == "loaded" + marks = _filter_mark_events(subscribed_events) + assert [mark.name for mark in marks] == ["skill.load"] + assert _mark_data(marks[0]) == {"skill_name": "review"} + assert _mark_metadata(marks[0]) == { + "skill_load_source": "structured_read", + "tool_name": "read_file", + } + tool_start = next( + event + for event in subscribed_events + if isinstance(event, nemo_relay.ScopeEvent) and event.name == "read_file" and event.scope_category == "start" + ) + assert marks[0].parent_uuid == tool_start.uuid + + def test_callback_handler_emits_human_in_the_loop_marks( subscribed_events: list[nemo_relay.Event], callback_handler: deepagents_integration.NemoRelayDeepAgentsCallbackHandler, diff --git a/python/tests/integrations/langchain_tests/test_middleware.py b/python/tests/integrations/langchain_tests/test_middleware.py index fffb369b3..18feb6e82 100644 --- a/python/tests/integrations/langchain_tests/test_middleware.py +++ b/python/tests/integrations/langchain_tests/test_middleware.py @@ -395,6 +395,50 @@ def test_awrap_tool_call_routes_through_tool_execute( assert isinstance(kwargs["result_codec"], nemo_relay.typed.BestEffortAnyCodec) +def test_complete_skill_read_emits_mark_through_langchain_middleware( + subscribed_events: list[nemo_relay.Event], + nemo_relay_middleware: NemoRelayMiddleware, +): + from langchain.agents.middleware import ToolCallRequest + from langchain_core.messages import ToolMessage + + request = ToolCallRequest( + tool_call={ + "name": "read_file", + "args": {"path": "/skills/review/SKILL.md"}, + "id": "call-skill", + }, + tool=None, + state={}, + runtime=MagicMock(), + ) + with nemo_relay.scope.scope("langchain-skill", nemo_relay.ScopeType.Agent): + response = nemo_relay_middleware.wrap_tool_call( + request, + lambda next_request: ToolMessage( + content="loaded", + tool_call_id=next_request.tool_call["id"], + ), + ) + nemo_relay.subscribers.flush() + + assert response.content == "loaded" + mark = next( + event for event in subscribed_events if isinstance(event, nemo_relay.MarkEvent) and event.name == "skill.load" + ) + tool_start = next( + event + for event in subscribed_events + if isinstance(event, nemo_relay.ScopeEvent) and event.name == "read_file" and event.scope_category == "start" + ) + assert mark.parent_uuid == tool_start.uuid + assert mark.data == {"skill_name": "review"} + assert mark.metadata == { + "skill_load_source": "structured_read", + "tool_name": "read_file", + } + + @pytest.mark.parametrize("use_async", [False, True]) def test_agent_integration(use_async: bool, nemo_relay_middleware: NemoRelayMiddleware): """An integration test to verify that the middleware correctly wraps a model call end-to-end.""" diff --git a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py index 464c88dfe..cd2decb73 100644 --- a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py +++ b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py @@ -124,6 +124,40 @@ async def test_async( assert _events_to_strings(subscribed_events) == self._expected_events +def test_complete_skill_read_inside_langgraph_emits_mark( + subscribed_events: list[nemo_relay.Event], + callback_handler: NemoRelayCallbackHandler, +): + from langgraph.graph import END, START, StateGraph + + def load_skill(state: State) -> State: + handle = nemo_relay.tools.call("read_file", {"path": "/skills/review/SKILL.md"}) + nemo_relay.tools.call_end(handle, {"loaded": True}) + return state + + builder = StateGraph(cast(Any, State)) + builder.add_node("load_skill", load_skill) + builder.add_edge(START, "load_skill") + builder.add_edge("load_skill", END) + graph = builder.compile() + + with nemo_relay.scope.scope("request", nemo_relay.ScopeType.Agent): + result = graph.invoke({"value": 1}, config={"callbacks": [callback_handler]}) + + nemo_relay.subscribers.flush() + assert result == {"value": 1} + mark = next( + event for event in subscribed_events if isinstance(event, nemo_relay.MarkEvent) and event.name == "skill.load" + ) + tool_start = next( + event + for event in subscribed_events + if isinstance(event, nemo_relay.ScopeEvent) and event.name == "read_file" and event.scope_category == "start" + ) + assert mark.parent_uuid == tool_start.uuid + assert mark.data == {"skill_name": "review"} + + def test_graph_lifecycle_callbacks_emit_marks( subscribed_events: list[nemo_relay.Event], callback_handler: NemoRelayCallbackHandler, diff --git a/python/tests/test_tools.py b/python/tests/test_tools.py index 18a0eaf9b..0fa667afa 100644 --- a/python/tests/test_tools.py +++ b/python/tests/test_tools.py @@ -8,6 +8,7 @@ import pytest from nemo_relay import ( + Event, MarkEvent, PendingMarkSpec, ScopeEvent, @@ -67,6 +68,19 @@ def test_call_with_parent_handle(self): tools.call_end(handle, {}) scope.pop(parent) + def test_complete_skill_read_emits_minimal_eager_mark(self, subscribed_events: list[Event]): + handle = tools.call("read_file", {"path": "/skills/review/SKILL.md"}) + tools.call_end(handle, {"ok": True}) + subscribers.flush() + + start = _tool_event(subscribed_events, "read_file", "start") + mark = next(event for event in subscribed_events if isinstance(event, MarkEvent) and event.name == "skill.load") + end = _tool_event(subscribed_events, "read_file", "end") + assert subscribed_events.index(start) < subscribed_events.index(mark) < subscribed_events.index(end) + assert mark.parent_uuid == start.uuid + assert mark.data == {"skill_name": "review"} + assert mark.metadata == {"skill_load_source": "structured_read", "tool_name": "read_file"} + class TestToolsAsync: async def test_execute_basic(self):