diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py index 0bfef25eb..ec7d6551c 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py @@ -13,6 +13,7 @@ from app.tasks.chat.streaming.handlers.tools import ( ) from app.tasks.chat.streaming.helpers.tool_output import tool_output_has_error from app.tasks.chat.streaming.relay.state import AgentEventRelayState +from app.tasks.chat.streaming.relay.task_span import clear_task_span_if_delegating_task_ended from app.tasks.chat.streaming.relay.thinking_step_sse import emit_thinking_step_frame @@ -91,6 +92,7 @@ def iter_tool_end_frames( title=title, status="completed", items=completed_items, + metadata=state.span_metadata_if_active(), ) state.just_finished_tool = True @@ -108,5 +110,10 @@ def iter_tool_end_frames( stream_result=result, langgraph_config=config, staged_workspace_file_path=staged_file_path, + tool_metadata=state.span_metadata_if_active(), ) yield from iter_tool_completion_emission_frames(emission_ctx) + + clear_task_span_if_delegating_task_ended( + state, tool_name=tool_name, run_id=run_id + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_output_frame.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_output_frame.py index 07244364c..4cd8e3274 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_output_frame.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_output_frame.py @@ -12,13 +12,18 @@ def emit_tool_output_available_frame( langchain_id_holder: dict[str, str | None], call_id: str, output: Any, + tool_metadata: dict[str, Any] | None = None, ) -> str: if content_builder is not None: content_builder.on_tool_output_available( - call_id, output, langchain_id_holder["value"] + call_id, + output, + langchain_id_holder["value"], + metadata=tool_metadata, ) return streaming_service.format_tool_output_available( call_id, output, langchain_tool_call_id=langchain_id_holder["value"], + metadata=tool_metadata, ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/emission_context.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/emission_context.py index d9ff796c0..baa1d7336 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/emission_context.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/emission_context.py @@ -23,6 +23,7 @@ class ToolCompletionEmissionContext: stream_result: Any langgraph_config: dict[str, Any] staged_workspace_file_path: str | None + tool_metadata: dict[str, Any] | None = None def emit_tool_output_card(self, payload: Any) -> str: return emit_tool_output_available_frame( @@ -31,4 +32,5 @@ class ToolCompletionEmissionContext: langchain_id_holder=self.langchain_tool_call_id_holder, call_id=self.tool_call_id, output=payload, + tool_metadata=self.tool_metadata, )