From 6c4ede5f9a9580618e24a85200809c34006043ce Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 02:00:34 +0200 Subject: [PATCH 01/96] fix(test): pin enable_kb_planner_runnable=false for KB-search planner tests --- .../tests/unit/middleware/test_knowledge_search.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/surfsense_backend/tests/unit/middleware/test_knowledge_search.py b/surfsense_backend/tests/unit/middleware/test_knowledge_search.py index 2933a0504..3529a946b 100644 --- a/surfsense_backend/tests/unit/middleware/test_knowledge_search.py +++ b/surfsense_backend/tests/unit/middleware/test_knowledge_search.py @@ -202,6 +202,15 @@ class FakeBudgetLLM: class TestKnowledgeBaseSearchMiddlewarePlanner: + @pytest.fixture(autouse=True) + def _disable_planner_runnable(self, monkeypatch): + # ``FakeLLM`` is a duck-typed mock; ``create_agent`` (used when the + # planner Runnable path is enabled) calls ``.bind()`` on the LLM, + # which the mock does not implement. Pin the flag off so the + # planner falls through to the legacy ``self.llm.ainvoke`` path + # these tests assert against (``llm.calls[0]["config"]``). + monkeypatch.setenv("SURFSENSE_ENABLE_KB_PLANNER_RUNNABLE", "false") + def test_render_recent_conversation_prefers_latest_messages_under_budget(self): messages = [ HumanMessage(content="old user context " * 40), From 2f3a33c9d518bb3ca8c1ac3535fc633375886078 Mon Sep 17 00:00:00 2001 From: guangyang1206 Date: Tue, 5 May 2026 12:48:04 +0800 Subject: [PATCH 02/96] feat(chunker): add table-aware chunk_text_hybrid to prevent mid-row table splits Document_chunker currently splits Markdown tables mid-row when the table is larger than a single chunk window, producing garbled rows that are useless for RAG retrieval (issue #1334). Changes: - document_chunker.py: add chunk_text_hybrid() that detects Markdown table blocks with a regex, emits each table as an indivisible single chunk, and feeds the surrounding prose through the normal chunk_text() chunker. - indexing_pipeline_service.py: route normal (non-code) documents through chunk_text_hybrid instead of chunk_text so tables are protected by default. Fixes #1334 --- .../app/indexing_pipeline/document_chunker.py | 50 +++++++++++++++++++ .../indexing_pipeline_service.py | 20 +++++--- 2 files changed, 64 insertions(+), 6 deletions(-) diff --git a/surfsense_backend/app/indexing_pipeline/document_chunker.py b/surfsense_backend/app/indexing_pipeline/document_chunker.py index 4f3c698ef..6ae81b7a8 100644 --- a/surfsense_backend/app/indexing_pipeline/document_chunker.py +++ b/surfsense_backend/app/indexing_pipeline/document_chunker.py @@ -1,5 +1,15 @@ +import re + from app.config import config +# Regex that matches a Markdown table block (header + separator + one or more rows) +# A table block starts with a | at the beginning of a line and ends when a +# non-table line (or end of string) is encountered. +_TABLE_BLOCK_RE = re.compile( + r"(?:(?:^|\n)(?=[ \t]*\|)(?:[ \t]*\|[^\n]*\n)+)", + re.MULTILINE, +) + def chunk_text(text: str, use_code_chunker: bool = False) -> list[str]: """Chunk a text string using the configured chunker and return the chunk texts.""" @@ -7,3 +17,43 @@ def chunk_text(text: str, use_code_chunker: bool = False) -> list[str]: config.code_chunker_instance if use_code_chunker else config.chunker_instance ) return [c.text for c in chunker.chunk(text)] + + +def chunk_text_hybrid(text: str) -> list[str]: + """Table-aware chunker that prevents Markdown tables from being split mid-row. + + Algorithm: + 1. Scan the document for Markdown table blocks. + 2. Each table block is emitted as a single, unmodified chunk so that its + header, separator row, and data rows always stay together. + 3. The non-table prose segments between (and around) tables are passed through + the normal ``chunk_text`` chunker and their sub-chunks are interleaved in + document order. + + This ensures that table data is never sliced in the middle by the token-based + chunker, which would otherwise produce garbled rows that are useless for RAG. + + Fixes #1334. + """ + chunks: list[str] = [] + cursor = 0 + + for match in _TABLE_BLOCK_RE.finditer(text): + # Prose before this table + prose = text[cursor : match.start()].strip() + if prose: + chunks.extend(chunk_text(prose)) + + # The table itself is kept as one indivisible chunk + table_block = match.group(0).strip() + if table_block: + chunks.append(table_block) + + cursor = match.end() + + # Remaining prose after the last table (or entire text if no tables) + trailing = text[cursor:].strip() + if trailing: + chunks.extend(chunk_text(trailing)) + + return chunks diff --git a/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py b/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py index e6b2458f3..2339647ea 100644 --- a/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py +++ b/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py @@ -19,7 +19,7 @@ from app.db import ( DocumentType, ) from app.indexing_pipeline.connector_document import ConnectorDocument -from app.indexing_pipeline.document_chunker import chunk_text +from app.indexing_pipeline.document_chunker import chunk_text, chunk_text_hybrid from app.indexing_pipeline.document_embedder import embed_texts from app.indexing_pipeline.document_hashing import ( compute_content_hash, @@ -387,11 +387,19 @@ class IndexingPipelineService: ) t_step = time.perf_counter() - chunk_texts = await asyncio.to_thread( - chunk_text, - connector_doc.source_markdown, - use_code_chunker=connector_doc.should_use_code_chunker, - ) + if connector_doc.should_use_code_chunker: + chunk_texts = await asyncio.to_thread( + chunk_text, + connector_doc.source_markdown, + use_code_chunker=True, + ) + else: + # Use the table-aware hybrid chunker so Markdown tables are not + # split mid-row (see issue #1334). + chunk_texts = await asyncio.to_thread( + chunk_text_hybrid, + connector_doc.source_markdown, + ) texts_to_embed = [content, *chunk_texts] embeddings = await asyncio.to_thread(embed_texts, texts_to_embed) From 102f77ab7f22e69d83b66a7201b4218b495199de Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 17:04:19 +0200 Subject: [PATCH 03/96] test(hitl): regression net for surfsense_resume_value pop-on-read --- .../unit/agents/multi_agent_chat/__init__.py | 0 .../multi_agent_chat/middleware/__init__.py | 0 .../__init__.py | 0 .../test_resume_helpers.py | 84 +++++++++++++++++++ 4 files changed, 84 insertions(+) create mode 100644 surfsense_backend/tests/unit/agents/multi_agent_chat/__init__.py create mode 100644 surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/__init__.py create mode 100644 surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/__init__.py create mode 100644 surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/__init__.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/__init__.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/__init__.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py new file mode 100644 index 000000000..2060051a2 --- /dev/null +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py @@ -0,0 +1,84 @@ +"""Pure-function tests for the HITL resume side-channel helpers. + +Tests the invariant that backs the bridge: a queued resume value must be +read exactly once per turn. A second read returns ``None`` so the +parent ``task`` tool falls through to its fail-loud guard rather than +replaying the same resume payload (which would re-fire the interrupt). +""" + +from __future__ import annotations + +from langchain.tools import ToolRuntime + +from app.agents.multi_agent_chat.main_agent.graph.middleware.checkpointed_subagent_middleware.config import ( + consume_surfsense_resume, + has_surfsense_resume, +) + + +def _runtime_with_config(config: dict) -> ToolRuntime: + """Real ToolRuntime; only ``.config`` is exercised by the helpers.""" + return ToolRuntime( + state=None, + context=None, + config=config, + stream_writer=None, + tool_call_id="tcid-test", + store=None, + ) + + +class TestConsumeSurfsenseResume: + def test_pops_value_on_first_call(self): + runtime = _runtime_with_config( + {"configurable": {"surfsense_resume_value": {"decisions": ["approve"]}}} + ) + + assert consume_surfsense_resume(runtime) == {"decisions": ["approve"]} + + def test_second_call_returns_none(self): + # Regression guard: a second read must not replay the queued + # resume. If it did, the subagent would re-invoke with the + # same Command and the user-facing interrupt would fire twice. + configurable: dict = {"surfsense_resume_value": {"decisions": ["approve"]}} + runtime = _runtime_with_config({"configurable": configurable}) + + consume_surfsense_resume(runtime) + + assert consume_surfsense_resume(runtime) is None + assert "surfsense_resume_value" not in configurable + + def test_returns_none_when_no_payload_queued(self): + runtime = _runtime_with_config({"configurable": {}}) + + assert consume_surfsense_resume(runtime) is None + + def test_returns_none_when_configurable_missing(self): + runtime = _runtime_with_config({}) + + assert consume_surfsense_resume(runtime) is None + + +class TestHasSurfsenseResume: + def test_true_when_payload_queued(self): + runtime = _runtime_with_config( + {"configurable": {"surfsense_resume_value": "approve"}} + ) + + assert has_surfsense_resume(runtime) is True + + def test_does_not_consume_payload(self): + # The fail-loud guard in ``task_tool`` calls ``has_surfsense_resume`` + # *before* deciding to consume; the check itself must leave the + # payload queued for the matching ``consume_surfsense_resume`` call. + configurable = {"surfsense_resume_value": "approve"} + runtime = _runtime_with_config({"configurable": configurable}) + + has_surfsense_resume(runtime) + + assert configurable == {"surfsense_resume_value": "approve"} + + def test_false_when_payload_absent(self): + runtime = _runtime_with_config({"configurable": {}}) + + assert has_surfsense_resume(runtime) is False From 681895d28df0aa100159a981b38188e7dcdb4dd3 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 17:13:13 +0200 Subject: [PATCH 04/96] test(hitl): regression net for end-to-end resume bridge dispatch --- .../test_hitl_bridge.py | 132 ++++++++++++++++++ 1 file changed, 132 insertions(+) create mode 100644 surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py new file mode 100644 index 000000000..8e9f0a96c --- /dev/null +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py @@ -0,0 +1,132 @@ +"""End-to-end resume-bridge tests against a real LangGraph subagent. + +Builds a minimal Pregel subagent that calls ``interrupt(...)`` and drives the +``task`` tool directly with a hand-crafted ``ToolRuntime``. Exercises the only +runtime contract we own: parent stashes a decision in +``config["configurable"]["surfsense_resume_value"]`` -> bridge forwards it as +``Command(resume={interrupt_id: value})`` -> subagent completes -> return value +reflects the decision. + +We pause the subagent **outside** the parent task tool (calling +``subagent.ainvoke`` directly) to skip the ``_lg_interrupt`` re-raise path, +which requires a parent runnable context. The bridge logic under test is the +*resume* dispatch, not the propagation; propagation is exercised separately in +its own module's tests. +""" + +from __future__ import annotations + +import pytest +from langchain.tools import ToolRuntime +from langchain_core.messages import HumanMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph import END, START, StateGraph +from langgraph.types import Command, interrupt +from typing_extensions import TypedDict + +from app.agents.multi_agent_chat.main_agent.graph.middleware.checkpointed_subagent_middleware.task_tool import ( + build_task_tool_with_parent_config, +) + + +class _SubagentState(TypedDict, total=False): + messages: list + decision_text: str + + +def _build_single_interrupt_subagent(): + """Subagent that interrupts once, then echoes the resume decision into state.""" + + def approve_node(state): + from langchain_core.messages import AIMessage + + decision = interrupt( + { + "action_requests": [ + { + "name": "do_thing", + "args": {"x": 1}, + "description": "test action", + } + ], + "review_configs": [{}], + } + ) + # Capture the resume payload verbatim so the test can assert the + # bridge forwarded it intact (no reshape, no scalar broadcast). + return { + "messages": [AIMessage(content="done")], + "decision_text": repr(decision), + } + + graph = StateGraph(_SubagentState) + graph.add_node("approve", approve_node) + graph.add_edge(START, "approve") + graph.add_edge("approve", END) + return graph.compile(checkpointer=InMemorySaver()) + + +def _make_runtime(config: dict) -> ToolRuntime: + return ToolRuntime( + state={"messages": [HumanMessage(content="seed")]}, + context=None, + config=config, + stream_writer=None, + tool_call_id="parent-tcid-1", + store=None, + ) + + +@pytest.mark.asyncio +async def test_resume_bridge_dispatches_decision_into_pending_subagent(): + """Side-channel decision -> targeted Command(resume) -> subagent completes.""" + subagent = _build_single_interrupt_subagent() + task_tool = build_task_tool_with_parent_config( + [ + { + "name": "approver", + "description": "approves things", + "runnable": subagent, + } + ] + ) + + # 1. Pause the subagent directly so we can test only the resume path. + parent_config: dict = { + "configurable": {"thread_id": "shared-thread"}, + "recursion_limit": 100, + } + await subagent.ainvoke({"messages": [HumanMessage(content="seed")]}, parent_config) + snap = await subagent.aget_state(parent_config) + assert snap.tasks and snap.tasks[0].interrupts, ( + "fixture broken: subagent should be paused on its interrupt" + ) + + # 2. Stash the user's decision on the side-channel — this is what + # ``stream_resume_chat`` does in production. + parent_config["configurable"]["surfsense_resume_value"] = { + "decisions": ["APPROVED"] + } + runtime = _make_runtime(parent_config) + + # 3. Drive the bridge. Subagent has no remaining interrupt after resume, + # so propagation will not call ``_lg_interrupt`` (no parent ctx needed). + result = await task_tool.coroutine( + description="please approve", + subagent_type="approver", + runtime=runtime, + ) + + assert isinstance(result, Command) + update = result.update + # Bridge forwards the side-channel payload **verbatim** to the + # subagent's ``interrupt()``. A scalar broadcast or accidental + # unwrap would change this shape and we want to catch that. + assert update["decision_text"] == repr({"decisions": ["APPROVED"]}) + + # 4. Side-channel was consumed; a stale replay would re-prompt the user. + assert "surfsense_resume_value" not in parent_config["configurable"] + + # 5. Subagent moved past the interrupt (no pending tasks remain). + final = await subagent.aget_state(parent_config) + assert not final.tasks or all(not t.interrupts for t in final.tasks) From 0c556356032d536dc57d83057595dfd669b90bb3 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 17:15:28 +0200 Subject: [PATCH 05/96] test(hitl): regression net for fail-loud guard on missing resume value --- .../test_hitl_bridge.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py index 8e9f0a96c..2ba95dfef 100644 --- a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py @@ -130,3 +130,41 @@ async def test_resume_bridge_dispatches_decision_into_pending_subagent(): # 5. Subagent moved past the interrupt (no pending tasks remain). final = await subagent.aget_state(parent_config) assert not final.tasks or all(not t.interrupts for t in final.tasks) + + +@pytest.mark.asyncio +async def test_pending_interrupt_without_resume_value_raises_runtime_error(): + """Bridge must fail loud if a paused subagent has no decision queued. + + The fail-open alternative (silently re-invoking) would re-fire the + same interrupt to the user. The error surfaces a real broken bridge + instead of confusing duplicate approval cards. + """ + subagent = _build_single_interrupt_subagent() + task_tool = build_task_tool_with_parent_config( + [ + { + "name": "approver", + "description": "approves things", + "runnable": subagent, + } + ] + ) + + parent_config: dict = { + "configurable": {"thread_id": "guard-thread"}, + "recursion_limit": 100, + } + await subagent.ainvoke({"messages": [HumanMessage(content="seed")]}, parent_config) + snap = await subagent.aget_state(parent_config) + assert snap.tasks and snap.tasks[0].interrupts, "fixture broken" + + # No surfsense_resume_value injected — bridge must refuse to proceed. + runtime = _make_runtime(parent_config) + + with pytest.raises(RuntimeError, match="resume bridge is broken"): + await task_tool.coroutine( + description="please approve", + subagent_type="approver", + runtime=runtime, + ) From fa6f3015a996640fb74fd7b483171d1f747c4302 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 17:19:39 +0200 Subject: [PATCH 06/96] test(hitl): regression net for bundle decision shape preservation --- .../test_hitl_bridge.py | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py index 2ba95dfef..188224074 100644 --- a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py @@ -16,6 +16,8 @@ its own module's tests. from __future__ import annotations +import ast + import pytest from langchain.tools import ToolRuntime from langchain_core.messages import HumanMessage @@ -168,3 +170,84 @@ async def test_pending_interrupt_without_resume_value_raises_runtime_error(): subagent_type="approver", runtime=runtime, ) + + +def _build_bundle_subagent(): + """Subagent that raises a 3-action HITL bundle on its only node.""" + + def bundle_node(state): + from langchain_core.messages import AIMessage + + decision = interrupt( + { + "action_requests": [ + {"name": "create_a", "args": {}, "description": ""}, + {"name": "create_b", "args": {}, "description": ""}, + {"name": "create_c", "args": {}, "description": ""}, + ], + "review_configs": [{}, {}, {}], + } + ) + return { + "messages": [AIMessage(content="bundle-done")], + "decision_text": repr(decision), + } + + graph = StateGraph(_SubagentState) + graph.add_node("bundle", bundle_node) + graph.add_edge(START, "bundle") + graph.add_edge("bundle", END) + return graph.compile(checkpointer=InMemorySaver()) + + +@pytest.mark.asyncio +async def test_bundle_three_mixed_decisions_arrive_in_order(): + """Approve / edit / reject for a 3-action bundle land at ordinals 0/1/2. + + Catches reshape regressions: truncation, decision collapse, order + scrambling, and the legacy single-decision broadcast that would + fan-out one verdict to every action. + """ + subagent = _build_bundle_subagent() + task_tool = build_task_tool_with_parent_config( + [ + { + "name": "bundler", + "description": "creates a bundle", + "runnable": subagent, + } + ] + ) + + parent_config: dict = { + "configurable": {"thread_id": "bundle-thread"}, + "recursion_limit": 100, + } + await subagent.ainvoke({"messages": [HumanMessage(content="seed")]}, parent_config) + + decisions_payload = { + "decisions": [ + {"type": "approve", "args": {}}, + {"type": "edit", "args": {"args": {"name": "edited-b"}}}, + {"type": "reject", "args": {"message": "no thanks"}}, + ] + } + parent_config["configurable"]["surfsense_resume_value"] = decisions_payload + runtime = _make_runtime(parent_config) + + result = await task_tool.coroutine( + description="run bundle", + subagent_type="bundler", + runtime=runtime, + ) + + assert isinstance(result, Command) + decision_text = result.update["decision_text"] + received = ast.literal_eval(decision_text) + assert received == decisions_payload, "bundle decisions must arrive verbatim" + # Cross-checks for the regressions this test exists to catch. + assert len(received["decisions"]) == 3 + assert received["decisions"][0]["type"] == "approve" + assert received["decisions"][1]["type"] == "edit" + assert received["decisions"][1]["args"] == {"args": {"name": "edited-b"}} + assert received["decisions"][2]["type"] == "reject" From f695298d30f6e6d50bb701e3e613fb207c650d30 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 17:39:27 +0200 Subject: [PATCH 07/96] feat(multi-agent): wire model fallback and retry into subagent middleware --- .../graph/middleware/deepagent_stack.py | 39 ++++--- .../multi_agent_chat/subagents/__init__.py | 0 .../subagents/shared/__init__.py | 0 .../subagents/shared/test_subagent_builder.py | 105 ++++++++++++++++++ 4 files changed, 129 insertions(+), 15 deletions(-) create mode 100644 surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/__init__.py create mode 100644 surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/__init__.py create mode 100644 surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py index 74e47cfab..8b7e3d0b0 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py @@ -208,6 +208,26 @@ def build_main_agent_deepagent_middleware( ) gp_middleware.insert(_patch_idx, subagent_deny_permission_mw) + # Defined here (instead of further down with the other ``wrap_model_call`` + # middlewares) so subagents share the same instances as the parent — + # otherwise a connector subagent would die on the first provider hiccup + # while the parent stays resilient. + retry_mw = ( + RetryAfterMiddleware(max_retries=3) + if flags.enable_retry_after and not flags.disable_new_agent_stack + else None + ) + fallback_mw: ModelFallbackMiddleware | None = None + if flags.enable_model_fallback and not flags.disable_new_agent_stack: + try: + fallback_mw = ModelFallbackMiddleware( + "openai:gpt-4o-mini", + "anthropic:claude-3-5-haiku-20241022", + ) + except Exception: + logging.warning("ModelFallbackMiddleware init failed; skipping.") + fallback_mw = None + registry_subagents: list[SubAgent] = [] try: subagent_extra_middleware: list[Any] = [ @@ -222,6 +242,10 @@ def build_main_agent_deepagent_middleware( ] if subagent_deny_permission_mw is not None: subagent_extra_middleware.append(subagent_deny_permission_mw) + if retry_mw is not None: + subagent_extra_middleware.append(retry_mw) + if fallback_mw is not None: + subagent_extra_middleware.append(fallback_mw) registry_subagents = build_subagents( dependencies=subagent_dependencies, model=llm, @@ -268,21 +292,6 @@ def build_main_agent_deepagent_middleware( backend_resolver=backend_resolver, ) - retry_mw = ( - RetryAfterMiddleware(max_retries=3) - if flags.enable_retry_after and not flags.disable_new_agent_stack - else None - ) - fallback_mw: ModelFallbackMiddleware | None = None - if flags.enable_model_fallback and not flags.disable_new_agent_stack: - try: - fallback_mw = ModelFallbackMiddleware( - "openai:gpt-4o-mini", - "anthropic:claude-3-5-haiku-20241022", - ) - except Exception: - logging.warning("ModelFallbackMiddleware init failed; skipping.") - fallback_mw = None model_call_limit_mw = ( ModelCallLimitMiddleware( thread_limit=120, diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/__init__.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/__init__.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py new file mode 100644 index 000000000..82f66891a --- /dev/null +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py @@ -0,0 +1,105 @@ +"""Resilience contract for subagents built via ``pack_subagent``. + +Subagents (jira, linear, notion, ...) run on the same LLM as the parent. When +the provider rate-limits or returns an empty stream, a single hiccup must not +abort the user's HITL flow — the connector subagent has to keep moving. This +relies on ``ModelFallbackMiddleware`` being usable as a subagent +``extra_middleware`` so the production builder can wire it in. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Iterator +from typing import Any + +import pytest +from langchain.agents import create_agent +from langchain.agents.middleware import ModelFallbackMiddleware +from langchain_core.callbacks import ( + AsyncCallbackManagerForLLMRun, + CallbackManagerForLLMRun, +) +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, +) +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage +from langchain_core.outputs import ChatGeneration, ChatResult + +from app.agents.multi_agent_chat.subagents.shared.subagent_builder import ( + pack_subagent, +) + + +class _AlwaysFailingChatModel(BaseChatModel): + """Mimics a provider hard-failing on every call (rate limit / empty stream). + + ``ModelFallbackMiddleware`` triggers on any ``Exception``, so the exact + error type doesn't matter for the contract under test. + """ + + @property + def _llm_type(self) -> str: + return "always-failing-test-model" + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + msg = "primary llm exploded" + raise RuntimeError(msg) + + async def _agenerate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: AsyncCallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + msg = "primary llm exploded" + raise RuntimeError(msg) + + def _stream(self, *args: Any, **kwargs: Any) -> Iterator[ChatGeneration]: + msg = "primary llm exploded" + raise RuntimeError(msg) + + async def _astream( + self, *args: Any, **kwargs: Any + ) -> AsyncIterator[ChatGeneration]: + msg = "primary llm exploded" + raise RuntimeError(msg) + yield # pragma: no cover - unreachable, satisfies async generator typing + + +@pytest.mark.asyncio +async def test_subagent_recovers_when_primary_llm_fails(): + """Primary blows up → fallback in extra_middleware finishes the turn.""" + primary = _AlwaysFailingChatModel() + fallback = FakeMessagesListChatModel( + responses=[AIMessage(content="recovered via fallback")] + ) + + spec = pack_subagent( + name="resilience_test", + description="test subagent", + system_prompt="be helpful", + tools=[], + model=primary, + extra_middleware=[ModelFallbackMiddleware(fallback)], + ) + + agent = create_agent( + model=spec["model"], + tools=spec["tools"], + middleware=spec["middleware"], + system_prompt=spec["system_prompt"], + ) + + result = await agent.ainvoke({"messages": [HumanMessage(content="hi")]}) + + final = result["messages"][-1] + assert isinstance(final, AIMessage) + assert final.content == "recovered via fallback" From 1745d7dccf250fe489c1f9ff422491ca11661d4d Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 18:04:47 +0200 Subject: [PATCH 08/96] feat(middleware): scope model fallback to provider/network errors only --- .../graph/middleware/deepagent_stack.py | 10 +- .../app/agents/new_chat/chat_deepagent.py | 10 +- .../middleware/scoped_model_fallback.py | 106 +++++++++++++ .../subagents/shared/test_subagent_builder.py | 18 +-- .../agents/new_chat/middleware/__init__.py | 0 .../middleware/test_scoped_model_fallback.py | 148 ++++++++++++++++++ 6 files changed, 275 insertions(+), 17 deletions(-) create mode 100644 surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py create mode 100644 surfsense_backend/tests/unit/agents/new_chat/middleware/__init__.py create mode 100644 surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py index 8b7e3d0b0..e490b6b47 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py @@ -14,7 +14,6 @@ from deepagents.middleware.subagents import GENERAL_PURPOSE_SUBAGENT from langchain.agents.middleware import ( LLMToolSelectorMiddleware, ModelCallLimitMiddleware, - ModelFallbackMiddleware, TodoListMiddleware, ToolCallLimitMiddleware, ) @@ -56,6 +55,9 @@ from app.agents.new_chat.middleware import ( create_surfsense_compaction_middleware, default_skills_sources, ) +from app.agents.new_chat.middleware.scoped_model_fallback import ( + ScopedModelFallbackMiddleware, +) from app.agents.new_chat.permissions import Rule, Ruleset from app.agents.new_chat.plugin_loader import ( PluginContext, @@ -217,15 +219,15 @@ def build_main_agent_deepagent_middleware( if flags.enable_retry_after and not flags.disable_new_agent_stack else None ) - fallback_mw: ModelFallbackMiddleware | None = None + fallback_mw: ScopedModelFallbackMiddleware | None = None if flags.enable_model_fallback and not flags.disable_new_agent_stack: try: - fallback_mw = ModelFallbackMiddleware( + fallback_mw = ScopedModelFallbackMiddleware( "openai:gpt-4o-mini", "anthropic:claude-3-5-haiku-20241022", ) except Exception: - logging.warning("ModelFallbackMiddleware init failed; skipping.") + logging.warning("ScopedModelFallbackMiddleware init failed; skipping.") fallback_mw = None registry_subagents: list[SubAgent] = [] diff --git a/surfsense_backend/app/agents/new_chat/chat_deepagent.py b/surfsense_backend/app/agents/new_chat/chat_deepagent.py index 1f4024d9d..605c31416 100644 --- a/surfsense_backend/app/agents/new_chat/chat_deepagent.py +++ b/surfsense_backend/app/agents/new_chat/chat_deepagent.py @@ -31,7 +31,6 @@ from langchain.agents import create_agent from langchain.agents.middleware import ( LLMToolSelectorMiddleware, ModelCallLimitMiddleware, - ModelFallbackMiddleware, TodoListMiddleware, ToolCallLimitMiddleware, ) @@ -77,6 +76,9 @@ from app.agents.new_chat.middleware import ( create_surfsense_compaction_middleware, default_skills_sources, ) +from app.agents.new_chat.middleware.scoped_model_fallback import ( + ScopedModelFallbackMiddleware, +) from app.agents.new_chat.permissions import Rule, Ruleset from app.agents.new_chat.plugin_loader import ( PluginContext, @@ -792,15 +794,15 @@ def _build_compiled_agent_blocking( # Fallback chain — primary is the agent's own model; we add cheap # alternatives. Off by default; only the first call site that # configures the chain via env should enable it. - fallback_mw: ModelFallbackMiddleware | None = None + fallback_mw: ScopedModelFallbackMiddleware | None = None if flags.enable_model_fallback and not flags.disable_new_agent_stack: try: - fallback_mw = ModelFallbackMiddleware( + fallback_mw = ScopedModelFallbackMiddleware( "openai:gpt-4o-mini", "anthropic:claude-3-5-haiku-20241022", ) except Exception: - logging.warning("ModelFallbackMiddleware init failed; skipping.") + logging.warning("ScopedModelFallbackMiddleware init failed; skipping.") fallback_mw = None model_call_limit_mw = ( ModelCallLimitMiddleware( diff --git a/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py b/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py new file mode 100644 index 000000000..de367fda9 --- /dev/null +++ b/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py @@ -0,0 +1,106 @@ +"""Fallback only on provider/network errors; let programming bugs raise. + +Upstream :class:`langchain.agents.middleware.ModelFallbackMiddleware` catches +every ``Exception``. With a non-provider bug (``KeyError``, ``TypeError``, +``AttributeError`` from middleware/state), every fallback model in the chain +hits the same bug — burning latency and tokens before the real cause finally +surfaces. Scoping the catch to provider-style exception types lets bugs fail +fast with clean tracebacks. + +Class-name matching (instead of ``isinstance`` against imported provider +types) keeps the dependency surface flat: openai, anthropic, google, +mistral, etc. all ship their own ``RateLimitError`` and we don't want to +import them all. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from langchain.agents.middleware import ModelFallbackMiddleware + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from langchain.agents.middleware.types import ModelRequest, ModelResponse + from langchain_core.messages import AIMessage + + +_FALLBACK_ELIGIBLE_NAMES: frozenset[str] = frozenset( + { + # Rate / quota + "RateLimitError", + # Server-side + "APIStatusError", + "InternalServerError", + "ServiceUnavailableError", + "BadGatewayError", + "GatewayTimeoutError", + # Network + "APIConnectionError", + "APITimeoutError", + "ConnectError", + "ConnectTimeout", + "ReadTimeout", + "RemoteProtocolError", + "TimeoutError", + "TimeoutException", + } +) + + +def _is_fallback_eligible(exc: BaseException) -> bool: + """Eligible if the exception or any base in its MRO matches by class name.""" + return any(cls.__name__ in _FALLBACK_ELIGIBLE_NAMES for cls in type(exc).__mro__) + + +class ScopedModelFallbackMiddleware(ModelFallbackMiddleware): + """``ModelFallbackMiddleware`` that re-raises non-provider exceptions.""" + + def wrap_model_call( # type: ignore[override] + self, + request: ModelRequest[Any], + handler: Callable[[ModelRequest[Any]], ModelResponse[Any]], + ) -> ModelResponse[Any] | AIMessage: + last_exception: Exception + try: + return handler(request) + except Exception as e: + if not _is_fallback_eligible(e): + raise + last_exception = e + + for fallback_model in self.models: + try: + return handler(request.override(model=fallback_model)) + except Exception as e: + if not _is_fallback_eligible(e): + raise + last_exception = e + continue + + raise last_exception + + async def awrap_model_call( # type: ignore[override] + self, + request: ModelRequest[Any], + handler: Callable[[ModelRequest[Any]], Awaitable[ModelResponse[Any]]], + ) -> ModelResponse[Any] | AIMessage: + last_exception: Exception + try: + return await handler(request) + except Exception as e: + if not _is_fallback_eligible(e): + raise + last_exception = e + + for fallback_model in self.models: + try: + return await handler(request.override(model=fallback_model)) + except Exception as e: + if not _is_fallback_eligible(e): + raise + last_exception = e + continue + + raise last_exception diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py index 82f66891a..859833f1c 100644 --- a/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py @@ -31,12 +31,12 @@ from app.agents.multi_agent_chat.subagents.shared.subagent_builder import ( ) -class _AlwaysFailingChatModel(BaseChatModel): - """Mimics a provider hard-failing on every call (rate limit / empty stream). +class RateLimitError(Exception): + """Provider-style 429; matches the scoped-fallback eligibility allowlist by name.""" - ``ModelFallbackMiddleware`` triggers on any ``Exception``, so the exact - error type doesn't matter for the contract under test. - """ + +class _AlwaysFailingChatModel(BaseChatModel): + """Mimics a provider hard-failing on every call (rate limit / empty stream).""" @property def _llm_type(self) -> str: @@ -50,7 +50,7 @@ class _AlwaysFailingChatModel(BaseChatModel): **kwargs: Any, ) -> ChatResult: msg = "primary llm exploded" - raise RuntimeError(msg) + raise RateLimitError(msg) async def _agenerate( self, @@ -60,17 +60,17 @@ class _AlwaysFailingChatModel(BaseChatModel): **kwargs: Any, ) -> ChatResult: msg = "primary llm exploded" - raise RuntimeError(msg) + raise RateLimitError(msg) def _stream(self, *args: Any, **kwargs: Any) -> Iterator[ChatGeneration]: msg = "primary llm exploded" - raise RuntimeError(msg) + raise RateLimitError(msg) async def _astream( self, *args: Any, **kwargs: Any ) -> AsyncIterator[ChatGeneration]: msg = "primary llm exploded" - raise RuntimeError(msg) + raise RateLimitError(msg) yield # pragma: no cover - unreachable, satisfies async generator typing diff --git a/surfsense_backend/tests/unit/agents/new_chat/middleware/__init__.py b/surfsense_backend/tests/unit/agents/new_chat/middleware/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py b/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py new file mode 100644 index 000000000..af464d1dc --- /dev/null +++ b/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py @@ -0,0 +1,148 @@ +"""Exception-scope contract for ``ScopedModelFallbackMiddleware``. + +Upstream ``ModelFallbackMiddleware`` catches every ``Exception`` and walks +the fallback chain. That means a programming bug (``KeyError`` from a +botched tool config, ``TypeError`` from middleware, ...) burns 1+N model +round-trips and ~Nx tokens before its real cause surfaces. The scoped +variant only falls back on provider/network exception types so bugs fail +fast, with clean tracebacks. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Iterator +from typing import Any + +import pytest +from langchain_core.callbacks import ( + AsyncCallbackManagerForLLMRun, + CallbackManagerForLLMRun, +) +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.messages import AIMessage, BaseMessage +from langchain_core.outputs import ChatGeneration, ChatResult + + +class _RaisingChatModel(BaseChatModel): + """LLM that raises a configurable exception on every invocation.""" + + exc_to_raise: Any + + @property + def _llm_type(self) -> str: + return "raising-test-model" + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + raise self.exc_to_raise + + async def _agenerate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: AsyncCallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + raise self.exc_to_raise + + def _stream(self, *args: Any, **kwargs: Any) -> Iterator[ChatGeneration]: + raise self.exc_to_raise + + async def _astream( + self, *args: Any, **kwargs: Any + ) -> AsyncIterator[ChatGeneration]: + raise self.exc_to_raise + yield # pragma: no cover - unreachable + + +class _RecordingChatModel(BaseChatModel): + """Returns a fixed message and counts how often it was called.""" + + response_text: str = "fallback-ok" + call_count: int = 0 + + @property + def _llm_type(self) -> str: + return "recording-test-model" + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + self.call_count += 1 + return ChatResult( + generations=[ + ChatGeneration(message=AIMessage(content=self.response_text)) + ] + ) + + async def _agenerate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: AsyncCallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + return self._generate(messages, stop, None, **kwargs) + + +# Locally defined provider-style error: importing openai/anthropic/etc. +# would couple the test to provider SDKs the contract intentionally avoids. +class RateLimitError(Exception): + """Mimics ``openai.RateLimitError`` for name-based eligibility.""" + + +def _build_agent(primary: BaseChatModel, fallback: BaseChatModel): + """Compile a no-tools agent with the scoped fallback wired in.""" + from langchain.agents import create_agent + + from app.agents.new_chat.middleware.scoped_model_fallback import ( + ScopedModelFallbackMiddleware, + ) + + return create_agent( + model=primary, + tools=[], + middleware=[ScopedModelFallbackMiddleware(fallback)], + system_prompt="be helpful", + ) + + +@pytest.mark.asyncio +async def test_provider_errors_trigger_fallback(): + """Class names matching the provider allowlist drive the fallback chain.""" + primary = _RaisingChatModel(exc_to_raise=RateLimitError("429 from provider")) + fallback = _RecordingChatModel(response_text="recovered") + + agent = _build_agent(primary, fallback) + result = await agent.ainvoke({"messages": [("user", "hi")]}) + + final = result["messages"][-1] + assert isinstance(final, AIMessage) + assert final.content == "recovered" + assert fallback.call_count == 1 + + +@pytest.mark.asyncio +async def test_programming_errors_propagate_without_invoking_fallback(): + """``KeyError`` from agent-side bugs must surface immediately, no fallback retry.""" + primary = _RaisingChatModel(exc_to_raise=KeyError("missing_state_field")) + fallback = _RecordingChatModel(response_text="should-never-arrive") + + agent = _build_agent(primary, fallback) + + with pytest.raises(KeyError, match="missing_state_field"): + await agent.ainvoke({"messages": [("user", "hi")]}) + + assert fallback.call_count == 0, ( + "fallback was invoked for a programming error; " + "scoping rule is broken" + ) From b394dc71c717ad42d73fea4aa874ee0dace866aa Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 18:12:38 +0200 Subject: [PATCH 09/96] feat(multi-agent): extend model fallback to general-purpose subagent --- .../graph/middleware/deepagent_stack.py | 16 ++++++++++++++++ .../new_chat/middleware/scoped_model_fallback.py | 1 + 2 files changed, 17 insertions(+) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py index e490b6b47..af7fceffa 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py @@ -230,6 +230,22 @@ def build_main_agent_deepagent_middleware( logging.warning("ScopedModelFallbackMiddleware init failed; skipping.") fallback_mw = None + # Mirror the parent's ordering: retry / fallback wrap caching, which wraps + # the model. ``gp_middleware`` is held by reference inside + # ``general_purpose_spec`` so this insertion propagates into the spec. + _gp_resilience: list[Any] = [m for m in (retry_mw, fallback_mw) if m is not None] + if _gp_resilience: + _cache_idx = next( + ( + i + for i, m in enumerate(gp_middleware) + if isinstance(m, AnthropicPromptCachingMiddleware) + ), + len(gp_middleware), + ) + for offset, mw in enumerate(_gp_resilience): + gp_middleware.insert(_cache_idx + offset, mw) + registry_subagents: list[SubAgent] = [] try: subagent_extra_middleware: list[Any] = [ diff --git a/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py b/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py index de367fda9..bde8edeba 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py +++ b/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py @@ -45,6 +45,7 @@ _FALLBACK_ELIGIBLE_NAMES: frozenset[str] = frozenset( "RemoteProtocolError", "TimeoutError", "TimeoutException", + # Can be extended to other exceptions in the future } ) From 309c69553123483a2f1a23b6c6cb544f38699203 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 18:16:00 +0200 Subject: [PATCH 10/96] feat(multi-agent): cap subagent model and tool call counts --- .../graph/middleware/deepagent_stack.py | 59 +++++++++++-------- 1 file changed, 35 insertions(+), 24 deletions(-) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py index af7fceffa..b76b54c27 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py @@ -230,10 +230,34 @@ def build_main_agent_deepagent_middleware( logging.warning("ScopedModelFallbackMiddleware init failed; skipping.") fallback_mw = None - # Mirror the parent's ordering: retry / fallback wrap caching, which wraps - # the model. ``gp_middleware`` is held by reference inside + # Cost / loop ceiling shared with subagents. ``state_schema`` of these + # middlewares is per-agent; counts are not summed across parent + sub — + # the cap acts as a safety net per agent, not a global budget. + model_call_limit_mw = ( + ModelCallLimitMiddleware( + thread_limit=120, + run_limit=80, + exit_behavior="end", + ) + if flags.enable_model_call_limit and not flags.disable_new_agent_stack + else None + ) + tool_call_limit_mw = ( + ToolCallLimitMiddleware( + thread_limit=300, run_limit=80, exit_behavior="continue" + ) + if flags.enable_tool_call_limit and not flags.disable_new_agent_stack + else None + ) + + # Mirror the parent's ordering: retry / fallback / limits wrap caching, + # which wraps the model. ``gp_middleware`` is held by reference inside # ``general_purpose_spec`` so this insertion propagates into the spec. - _gp_resilience: list[Any] = [m for m in (retry_mw, fallback_mw) if m is not None] + _gp_resilience: list[Any] = [ + m + for m in (retry_mw, fallback_mw, model_call_limit_mw, tool_call_limit_mw) + if m is not None + ] if _gp_resilience: _cache_idx = next( ( @@ -260,10 +284,14 @@ def build_main_agent_deepagent_middleware( ] if subagent_deny_permission_mw is not None: subagent_extra_middleware.append(subagent_deny_permission_mw) - if retry_mw is not None: - subagent_extra_middleware.append(retry_mw) - if fallback_mw is not None: - subagent_extra_middleware.append(fallback_mw) + for _resilience_mw in ( + retry_mw, + fallback_mw, + model_call_limit_mw, + tool_call_limit_mw, + ): + if _resilience_mw is not None: + subagent_extra_middleware.append(_resilience_mw) registry_subagents = build_subagents( dependencies=subagent_dependencies, model=llm, @@ -310,23 +338,6 @@ def build_main_agent_deepagent_middleware( backend_resolver=backend_resolver, ) - model_call_limit_mw = ( - ModelCallLimitMiddleware( - thread_limit=120, - run_limit=80, - exit_behavior="end", - ) - if flags.enable_model_call_limit and not flags.disable_new_agent_stack - else None - ) - tool_call_limit_mw = ( - ToolCallLimitMiddleware( - thread_limit=300, run_limit=80, exit_behavior="continue" - ) - if flags.enable_tool_call_limit and not flags.disable_new_agent_stack - else None - ) - noop_mw = ( NoopInjectionMiddleware() if flags.enable_compaction_v2 and not flags.disable_new_agent_stack From 9a4ee5d16bedff12f76810c6b432d76c1a9f615d Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 18:27:46 +0200 Subject: [PATCH 11/96] chore: trim narrative comments and docstrings --- .../graph/middleware/deepagent_stack.py | 15 ++--- .../middleware/scoped_model_fallback.py | 24 ++------ .../test_hitl_bridge.py | 57 ++----------------- .../test_resume_helpers.py | 15 +---- .../subagents/shared/test_subagent_builder.py | 14 +---- .../middleware/test_scoped_model_fallback.py | 28 ++------- 6 files changed, 24 insertions(+), 129 deletions(-) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py index b76b54c27..1d6a8763e 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py @@ -210,10 +210,8 @@ def build_main_agent_deepagent_middleware( ) gp_middleware.insert(_patch_idx, subagent_deny_permission_mw) - # Defined here (instead of further down with the other ``wrap_model_call`` - # middlewares) so subagents share the same instances as the parent — - # otherwise a connector subagent would die on the first provider hiccup - # while the parent stays resilient. + # Defined early so the same instances reach both gp_middleware and + # subagent_extra_middleware below. retry_mw = ( RetryAfterMiddleware(max_retries=3) if flags.enable_retry_after and not flags.disable_new_agent_stack @@ -230,9 +228,7 @@ def build_main_agent_deepagent_middleware( logging.warning("ScopedModelFallbackMiddleware init failed; skipping.") fallback_mw = None - # Cost / loop ceiling shared with subagents. ``state_schema`` of these - # middlewares is per-agent; counts are not summed across parent + sub — - # the cap acts as a safety net per agent, not a global budget. + # Per-agent caps; counts are not summed across parent + subagents. model_call_limit_mw = ( ModelCallLimitMiddleware( thread_limit=120, @@ -250,9 +246,8 @@ def build_main_agent_deepagent_middleware( else None ) - # Mirror the parent's ordering: retry / fallback / limits wrap caching, - # which wraps the model. ``gp_middleware`` is held by reference inside - # ``general_purpose_spec`` so this insertion propagates into the spec. + # gp_middleware is held by reference inside general_purpose_spec, so + # mutating it here propagates into the spec. _gp_resilience: list[Any] = [ m for m in (retry_mw, fallback_mw, model_call_limit_mw, tool_call_limit_mw) diff --git a/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py b/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py index bde8edeba..99eb2d74a 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py +++ b/surfsense_backend/app/agents/new_chat/middleware/scoped_model_fallback.py @@ -1,17 +1,4 @@ -"""Fallback only on provider/network errors; let programming bugs raise. - -Upstream :class:`langchain.agents.middleware.ModelFallbackMiddleware` catches -every ``Exception``. With a non-provider bug (``KeyError``, ``TypeError``, -``AttributeError`` from middleware/state), every fallback model in the chain -hits the same bug — burning latency and tokens before the real cause finally -surfaces. Scoping the catch to provider-style exception types lets bugs fail -fast with clean tracebacks. - -Class-name matching (instead of ``isinstance`` against imported provider -types) keeps the dependency surface flat: openai, anthropic, google, -mistral, etc. all ship their own ``RateLimitError`` and we don't want to -import them all. -""" +"""Fallback only on provider/network errors; let programming bugs raise.""" from __future__ import annotations @@ -26,17 +13,16 @@ if TYPE_CHECKING: from langchain_core.messages import AIMessage +# Matched by class name across the MRO so we don't have to import every +# provider SDK (openai/anthropic/google/...). Extend as new providers ship. _FALLBACK_ELIGIBLE_NAMES: frozenset[str] = frozenset( { - # Rate / quota "RateLimitError", - # Server-side "APIStatusError", "InternalServerError", "ServiceUnavailableError", "BadGatewayError", "GatewayTimeoutError", - # Network "APIConnectionError", "APITimeoutError", "ConnectError", @@ -45,18 +31,16 @@ _FALLBACK_ELIGIBLE_NAMES: frozenset[str] = frozenset( "RemoteProtocolError", "TimeoutError", "TimeoutException", - # Can be extended to other exceptions in the future } ) def _is_fallback_eligible(exc: BaseException) -> bool: - """Eligible if the exception or any base in its MRO matches by class name.""" return any(cls.__name__ in _FALLBACK_ELIGIBLE_NAMES for cls in type(exc).__mro__) class ScopedModelFallbackMiddleware(ModelFallbackMiddleware): - """``ModelFallbackMiddleware`` that re-raises non-provider exceptions.""" + """Re-raise non-provider exceptions instead of walking the fallback chain.""" def wrap_model_call( # type: ignore[override] self, diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py index 188224074..ab6a644b5 100644 --- a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py @@ -1,18 +1,4 @@ -"""End-to-end resume-bridge tests against a real LangGraph subagent. - -Builds a minimal Pregel subagent that calls ``interrupt(...)`` and drives the -``task`` tool directly with a hand-crafted ``ToolRuntime``. Exercises the only -runtime contract we own: parent stashes a decision in -``config["configurable"]["surfsense_resume_value"]`` -> bridge forwards it as -``Command(resume={interrupt_id: value})`` -> subagent completes -> return value -reflects the decision. - -We pause the subagent **outside** the parent task tool (calling -``subagent.ainvoke`` directly) to skip the ``_lg_interrupt`` re-raise path, -which requires a parent runnable context. The bridge logic under test is the -*resume* dispatch, not the propagation; propagation is exercised separately in -its own module's tests. -""" +"""End-to-end resume-bridge tests against a real LangGraph subagent.""" from __future__ import annotations @@ -37,8 +23,6 @@ class _SubagentState(TypedDict, total=False): def _build_single_interrupt_subagent(): - """Subagent that interrupts once, then echoes the resume decision into state.""" - def approve_node(state): from langchain_core.messages import AIMessage @@ -54,8 +38,6 @@ def _build_single_interrupt_subagent(): "review_configs": [{}], } ) - # Capture the resume payload verbatim so the test can assert the - # bridge forwarded it intact (no reshape, no scalar broadcast). return { "messages": [AIMessage(content="done")], "decision_text": repr(decision), @@ -81,7 +63,7 @@ def _make_runtime(config: dict) -> ToolRuntime: @pytest.mark.asyncio async def test_resume_bridge_dispatches_decision_into_pending_subagent(): - """Side-channel decision -> targeted Command(resume) -> subagent completes.""" + """Side-channel decision must reach the subagent's pending interrupt verbatim.""" subagent = _build_single_interrupt_subagent() task_tool = build_task_tool_with_parent_config( [ @@ -93,7 +75,6 @@ async def test_resume_bridge_dispatches_decision_into_pending_subagent(): ] ) - # 1. Pause the subagent directly so we can test only the resume path. parent_config: dict = { "configurable": {"thread_id": "shared-thread"}, "recursion_limit": 100, @@ -104,15 +85,11 @@ async def test_resume_bridge_dispatches_decision_into_pending_subagent(): "fixture broken: subagent should be paused on its interrupt" ) - # 2. Stash the user's decision on the side-channel — this is what - # ``stream_resume_chat`` does in production. parent_config["configurable"]["surfsense_resume_value"] = { "decisions": ["APPROVED"] } runtime = _make_runtime(parent_config) - # 3. Drive the bridge. Subagent has no remaining interrupt after resume, - # so propagation will not call ``_lg_interrupt`` (no parent ctx needed). result = await task_tool.coroutine( description="please approve", subagent_type="approver", @@ -121,27 +98,16 @@ async def test_resume_bridge_dispatches_decision_into_pending_subagent(): assert isinstance(result, Command) update = result.update - # Bridge forwards the side-channel payload **verbatim** to the - # subagent's ``interrupt()``. A scalar broadcast or accidental - # unwrap would change this shape and we want to catch that. assert update["decision_text"] == repr({"decisions": ["APPROVED"]}) - - # 4. Side-channel was consumed; a stale replay would re-prompt the user. assert "surfsense_resume_value" not in parent_config["configurable"] - # 5. Subagent moved past the interrupt (no pending tasks remain). final = await subagent.aget_state(parent_config) assert not final.tasks or all(not t.interrupts for t in final.tasks) @pytest.mark.asyncio async def test_pending_interrupt_without_resume_value_raises_runtime_error(): - """Bridge must fail loud if a paused subagent has no decision queued. - - The fail-open alternative (silently re-invoking) would re-fire the - same interrupt to the user. The error surfaces a real broken bridge - instead of confusing duplicate approval cards. - """ + """Bridge must fail loud rather than silently replay the user's interrupt.""" subagent = _build_single_interrupt_subagent() task_tool = build_task_tool_with_parent_config( [ @@ -161,7 +127,6 @@ async def test_pending_interrupt_without_resume_value_raises_runtime_error(): snap = await subagent.aget_state(parent_config) assert snap.tasks and snap.tasks[0].interrupts, "fixture broken" - # No surfsense_resume_value injected — bridge must refuse to proceed. runtime = _make_runtime(parent_config) with pytest.raises(RuntimeError, match="resume bridge is broken"): @@ -173,8 +138,6 @@ async def test_pending_interrupt_without_resume_value_raises_runtime_error(): def _build_bundle_subagent(): - """Subagent that raises a 3-action HITL bundle on its only node.""" - def bundle_node(state): from langchain_core.messages import AIMessage @@ -202,12 +165,7 @@ def _build_bundle_subagent(): @pytest.mark.asyncio async def test_bundle_three_mixed_decisions_arrive_in_order(): - """Approve / edit / reject for a 3-action bundle land at ordinals 0/1/2. - - Catches reshape regressions: truncation, decision collapse, order - scrambling, and the legacy single-decision broadcast that would - fan-out one verdict to every action. - """ + """Approve / edit / reject for a 3-action bundle must land at ordinals 0/1/2.""" subagent = _build_bundle_subagent() task_tool = build_task_tool_with_parent_config( [ @@ -242,11 +200,8 @@ async def test_bundle_three_mixed_decisions_arrive_in_order(): ) assert isinstance(result, Command) - decision_text = result.update["decision_text"] - received = ast.literal_eval(decision_text) - assert received == decisions_payload, "bundle decisions must arrive verbatim" - # Cross-checks for the regressions this test exists to catch. - assert len(received["decisions"]) == 3 + received = ast.literal_eval(result.update["decision_text"]) + assert received == decisions_payload assert received["decisions"][0]["type"] == "approve" assert received["decisions"][1]["type"] == "edit" assert received["decisions"][1]["args"] == {"args": {"name": "edited-b"}} diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py index 2060051a2..e73fb2823 100644 --- a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py @@ -1,10 +1,4 @@ -"""Pure-function tests for the HITL resume side-channel helpers. - -Tests the invariant that backs the bridge: a queued resume value must be -read exactly once per turn. A second read returns ``None`` so the -parent ``task`` tool falls through to its fail-loud guard rather than -replaying the same resume payload (which would re-fire the interrupt). -""" +"""Resume side-channel must be read exactly once per turn.""" from __future__ import annotations @@ -17,7 +11,6 @@ from app.agents.multi_agent_chat.main_agent.graph.middleware.checkpointed_subage def _runtime_with_config(config: dict) -> ToolRuntime: - """Real ToolRuntime; only ``.config`` is exercised by the helpers.""" return ToolRuntime( state=None, context=None, @@ -37,9 +30,6 @@ class TestConsumeSurfsenseResume: assert consume_surfsense_resume(runtime) == {"decisions": ["approve"]} def test_second_call_returns_none(self): - # Regression guard: a second read must not replay the queued - # resume. If it did, the subagent would re-invoke with the - # same Command and the user-facing interrupt would fire twice. configurable: dict = {"surfsense_resume_value": {"decisions": ["approve"]}} runtime = _runtime_with_config({"configurable": configurable}) @@ -68,9 +58,6 @@ class TestHasSurfsenseResume: assert has_surfsense_resume(runtime) is True def test_does_not_consume_payload(self): - # The fail-loud guard in ``task_tool`` calls ``has_surfsense_resume`` - # *before* deciding to consume; the check itself must leave the - # payload queued for the matching ``consume_surfsense_resume`` call. configurable = {"surfsense_resume_value": "approve"} runtime = _runtime_with_config({"configurable": configurable}) diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py index 859833f1c..5cd62ed36 100644 --- a/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py @@ -1,11 +1,4 @@ -"""Resilience contract for subagents built via ``pack_subagent``. - -Subagents (jira, linear, notion, ...) run on the same LLM as the parent. When -the provider rate-limits or returns an empty stream, a single hiccup must not -abort the user's HITL flow — the connector subagent has to keep moving. This -relies on ``ModelFallbackMiddleware`` being usable as a subagent -``extra_middleware`` so the production builder can wire it in. -""" +"""Subagent resilience contract: ``extra_middleware`` reaches the agent chain.""" from __future__ import annotations @@ -32,11 +25,10 @@ from app.agents.multi_agent_chat.subagents.shared.subagent_builder import ( class RateLimitError(Exception): - """Provider-style 429; matches the scoped-fallback eligibility allowlist by name.""" + """Name matches the scoped-fallback eligibility allowlist.""" class _AlwaysFailingChatModel(BaseChatModel): - """Mimics a provider hard-failing on every call (rate limit / empty stream).""" @property def _llm_type(self) -> str: @@ -76,7 +68,7 @@ class _AlwaysFailingChatModel(BaseChatModel): @pytest.mark.asyncio async def test_subagent_recovers_when_primary_llm_fails(): - """Primary blows up → fallback in extra_middleware finishes the turn.""" + """Fallback in ``extra_middleware`` must finish the turn when primary raises.""" primary = _AlwaysFailingChatModel() fallback = FakeMessagesListChatModel( responses=[AIMessage(content="recovered via fallback")] diff --git a/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py b/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py index af464d1dc..69f6fe6b7 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py +++ b/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py @@ -1,12 +1,4 @@ -"""Exception-scope contract for ``ScopedModelFallbackMiddleware``. - -Upstream ``ModelFallbackMiddleware`` catches every ``Exception`` and walks -the fallback chain. That means a programming bug (``KeyError`` from a -botched tool config, ``TypeError`` from middleware, ...) burns 1+N model -round-trips and ~Nx tokens before its real cause surfaces. The scoped -variant only falls back on provider/network exception types so bugs fail -fast, with clean tracebacks. -""" +"""``ScopedModelFallbackMiddleware`` triggers fallback only on provider errors.""" from __future__ import annotations @@ -24,8 +16,6 @@ from langchain_core.outputs import ChatGeneration, ChatResult class _RaisingChatModel(BaseChatModel): - """LLM that raises a configurable exception on every invocation.""" - exc_to_raise: Any @property @@ -61,8 +51,6 @@ class _RaisingChatModel(BaseChatModel): class _RecordingChatModel(BaseChatModel): - """Returns a fixed message and counts how often it was called.""" - response_text: str = "fallback-ok" call_count: int = 0 @@ -94,14 +82,11 @@ class _RecordingChatModel(BaseChatModel): return self._generate(messages, stop, None, **kwargs) -# Locally defined provider-style error: importing openai/anthropic/etc. -# would couple the test to provider SDKs the contract intentionally avoids. class RateLimitError(Exception): - """Mimics ``openai.RateLimitError`` for name-based eligibility.""" + """Name matches the scoped-fallback eligibility allowlist.""" def _build_agent(primary: BaseChatModel, fallback: BaseChatModel): - """Compile a no-tools agent with the scoped fallback wired in.""" from langchain.agents import create_agent from app.agents.new_chat.middleware.scoped_model_fallback import ( @@ -118,7 +103,7 @@ def _build_agent(primary: BaseChatModel, fallback: BaseChatModel): @pytest.mark.asyncio async def test_provider_errors_trigger_fallback(): - """Class names matching the provider allowlist drive the fallback chain.""" + """Eligible exception names must drive the fallback chain.""" primary = _RaisingChatModel(exc_to_raise=RateLimitError("429 from provider")) fallback = _RecordingChatModel(response_text="recovered") @@ -133,7 +118,7 @@ async def test_provider_errors_trigger_fallback(): @pytest.mark.asyncio async def test_programming_errors_propagate_without_invoking_fallback(): - """``KeyError`` from agent-side bugs must surface immediately, no fallback retry.""" + """Non-eligible exceptions must propagate; fallback must not be invoked.""" primary = _RaisingChatModel(exc_to_raise=KeyError("missing_state_field")) fallback = _RecordingChatModel(response_text="should-never-arrive") @@ -142,7 +127,4 @@ async def test_programming_errors_propagate_without_invoking_fallback(): with pytest.raises(KeyError, match="missing_state_field"): await agent.ainvoke({"messages": [("user", "hi")]}) - assert fallback.call_count == 0, ( - "fallback was invoked for a programming error; " - "scoping rule is broken" - ) + assert fallback.call_count == 0 From a6df944247219a1e2ad6f1d3f0ee06dd6c17ae44 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:49:47 +0200 Subject: [PATCH 12/96] refactor(multi-agent): introduce shared flags helper and permissions package --- .../middleware/shared/__init__.py | 0 .../middleware/shared/flags.py | 10 ++ .../middleware/shared/permissions/__init__.py | 12 ++ .../middleware/shared/permissions/context.py | 109 ++++++++++++++++++ .../shared/permissions/middleware.py | 10 ++ 5 files changed, 141 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/__init__.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/flags.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/__init__.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/context.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/middleware.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/flags.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/flags.py new file mode 100644 index 000000000..69994ae00 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/flags.py @@ -0,0 +1,10 @@ +"""Single source of truth for the feature-flag predicate.""" + +from __future__ import annotations + +from app.agents.new_chat.feature_flags import AgentFeatureFlags + + +def enabled(flags: AgentFeatureFlags, attr: str) -> bool: + """``flags.`` is on AND the new-agent-stack kill switch is off.""" + return getattr(flags, attr) and not flags.disable_new_agent_stack diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/__init__.py new file mode 100644 index 000000000..4f2228170 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/__init__.py @@ -0,0 +1,12 @@ +"""Permission rulesets fanned out to parent / general-purpose / subagent stacks.""" + +from __future__ import annotations + +from .context import PermissionContext, build_permission_context +from .middleware import build_full_permission_mw + +__all__ = [ + "PermissionContext", + "build_full_permission_mw", + "build_permission_context", +] diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/context.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/context.py new file mode 100644 index 000000000..f14d52714 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/context.py @@ -0,0 +1,109 @@ +"""Derive shared permission context once; fan out to all three stack layers. + +The context carries: +- ``rulesets``: full ask/deny/allow rules for the main-agent permission middleware. +- ``general_purpose_interrupt_on``: ``ask`` rules mirrored as deepagents + ``interrupt_on`` so HITL still triggers from inside ``task`` runs (subagents + bypass the main-agent permission middleware). +- ``subagent_deny_mw``: a deny-only ``PermissionMiddleware`` instance shared + across the general-purpose and registry subagent stacks. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +from langchain_core.tools import BaseTool + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.agents.new_chat.middleware import PermissionMiddleware +from app.agents.new_chat.permissions import Rule, Ruleset +from app.agents.new_chat.tools.registry import BUILTIN_TOOLS + +from ..flags import enabled + + +@dataclass(frozen=True) +class PermissionContext: + rulesets: list[Ruleset] + general_purpose_interrupt_on: dict[str, bool] + subagent_deny_mw: PermissionMiddleware | None + + +def build_permission_context( + *, + flags: AgentFeatureFlags, + filesystem_mode: FilesystemMode, + tools: Sequence[BaseTool], + available_connectors: list[str] | None, +) -> PermissionContext: + is_desktop_fs = filesystem_mode == FilesystemMode.DESKTOP_LOCAL_FOLDER + permission_enabled = enabled(flags, "enable_permission") + + rulesets: list[Ruleset] = [] + if permission_enabled or is_desktop_fs: + rulesets.append( + Ruleset( + rules=[Rule(permission="*", pattern="*", action="allow")], + origin="surfsense_defaults", + ) + ) + if is_desktop_fs: + rulesets.append( + Ruleset( + rules=[ + Rule(permission="rm", pattern="*", action="ask"), + Rule(permission="rmdir", pattern="*", action="ask"), + Rule(permission="move_file", pattern="*", action="ask"), + Rule(permission="edit_file", pattern="*", action="ask"), + Rule(permission="write_file", pattern="*", action="ask"), + ], + origin="desktop_safety", + ) + ) + + tool_names_in_use = {t.name for t in tools} + + if permission_enabled: + available_set = set(available_connectors or []) + synthesized: list[Rule] = [] + for tool_def in BUILTIN_TOOLS: + if tool_def.name not in tool_names_in_use: + continue + rc = tool_def.required_connector + if rc and rc not in available_set: + synthesized.append( + Rule(permission=tool_def.name, pattern="*", action="deny") + ) + if synthesized: + rulesets.append( + Ruleset(rules=synthesized, origin="connector_synthesized") + ) + + general_purpose_interrupt_on: dict[str, bool] = { + rule.permission: True + for rs in rulesets + for rule in rs.rules + if rule.action == "ask" and rule.permission in tool_names_in_use + } + + deny_rulesets = [ + Ruleset( + rules=[r for r in rs.rules if r.action == "deny"], + origin=rs.origin, + ) + for rs in rulesets + ] + deny_rulesets = [rs for rs in deny_rulesets if rs.rules] + + subagent_deny_mw: PermissionMiddleware | None = ( + PermissionMiddleware(rulesets=deny_rulesets) if deny_rulesets else None + ) + + return PermissionContext( + rulesets=rulesets, + general_purpose_interrupt_on=general_purpose_interrupt_on, + subagent_deny_mw=subagent_deny_mw, + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/middleware.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/middleware.py new file mode 100644 index 000000000..704a26fb3 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/middleware.py @@ -0,0 +1,10 @@ +"""Main-agent permission middleware (full ask/deny/allow rules).""" + +from __future__ import annotations + +from app.agents.new_chat.middleware import PermissionMiddleware +from app.agents.new_chat.permissions import Ruleset + + +def build_full_permission_mw(rulesets: list[Ruleset]) -> PermissionMiddleware | None: + return PermissionMiddleware(rulesets=rulesets) if rulesets else None From 91701bb49afe93578c99f25cb5b019b2b261ceb8 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:51:02 +0200 Subject: [PATCH 13/96] refactor(multi-agent): split shared resilience bundle into per-concept files --- .../middleware/shared/resilience/__init__.py | 7 +++ .../middleware/shared/resilience/bundle.py | 51 +++++++++++++++++++ .../middleware/shared/resilience/fallback.py | 27 ++++++++++ .../shared/resilience/model_call_limit.py | 21 ++++++++ .../middleware/shared/resilience/retry.py | 16 ++++++ .../shared/resilience/tool_call_limit.py | 21 ++++++++ 6 files changed, 143 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/__init__.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/bundle.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/fallback.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/model_call_limit.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/retry.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/tool_call_limit.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/__init__.py new file mode 100644 index 000000000..92596b771 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/__init__.py @@ -0,0 +1,7 @@ +"""Resilience middleware shared as the same instances across parent / general-purpose / registry.""" + +from __future__ import annotations + +from .bundle import ResilienceBundle, build_resilience_bundle + +__all__ = ["ResilienceBundle", "build_resilience_bundle"] diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/bundle.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/bundle.py new file mode 100644 index 000000000..45f76a6f3 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/bundle.py @@ -0,0 +1,51 @@ +"""Construct each resilience middleware once; same instances flow into every consumer.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from langchain.agents.middleware import ( + ModelCallLimitMiddleware, + ToolCallLimitMiddleware, +) + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import RetryAfterMiddleware +from app.agents.new_chat.middleware.scoped_model_fallback import ( + ScopedModelFallbackMiddleware, +) + +from .fallback import build_fallback_mw +from .model_call_limit import build_model_call_limit_mw +from .retry import build_retry_mw +from .tool_call_limit import build_tool_call_limit_mw + + +@dataclass(frozen=True) +class ResilienceBundle: + retry: RetryAfterMiddleware | None + fallback: ScopedModelFallbackMiddleware | None + model_call_limit: ModelCallLimitMiddleware | None + tool_call_limit: ToolCallLimitMiddleware | None + + def as_list(self) -> list[Any]: + return [ + m + for m in ( + self.retry, + self.fallback, + self.model_call_limit, + self.tool_call_limit, + ) + if m is not None + ] + + +def build_resilience_bundle(flags: AgentFeatureFlags) -> ResilienceBundle: + return ResilienceBundle( + retry=build_retry_mw(flags), + fallback=build_fallback_mw(flags), + model_call_limit=build_model_call_limit_mw(flags), + tool_call_limit=build_tool_call_limit_mw(flags), + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/fallback.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/fallback.py new file mode 100644 index 000000000..ea68a764e --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/fallback.py @@ -0,0 +1,27 @@ +"""Switch to a fallback model on provider/network errors only.""" + +from __future__ import annotations + +import logging + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware.scoped_model_fallback import ( + ScopedModelFallbackMiddleware, +) + +from ..flags import enabled + + +def build_fallback_mw( + flags: AgentFeatureFlags, +) -> ScopedModelFallbackMiddleware | None: + if not enabled(flags, "enable_model_fallback"): + return None + try: + return ScopedModelFallbackMiddleware( + "openai:gpt-4o-mini", + "anthropic:claude-3-5-haiku-20241022", + ) + except Exception: + logging.warning("ScopedModelFallbackMiddleware init failed; skipping.") + return None diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/model_call_limit.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/model_call_limit.py new file mode 100644 index 000000000..85707a385 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/model_call_limit.py @@ -0,0 +1,21 @@ +"""Cap model calls per thread / per run to prevent runaway cost.""" + +from __future__ import annotations + +from langchain.agents.middleware import ModelCallLimitMiddleware + +from app.agents.new_chat.feature_flags import AgentFeatureFlags + +from ..flags import enabled + + +def build_model_call_limit_mw( + flags: AgentFeatureFlags, +) -> ModelCallLimitMiddleware | None: + if not enabled(flags, "enable_model_call_limit"): + return None + return ModelCallLimitMiddleware( + thread_limit=120, + run_limit=80, + exit_behavior="end", + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/retry.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/retry.py new file mode 100644 index 000000000..c98fc4083 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/retry.py @@ -0,0 +1,16 @@ +"""Retry on transient model errors (e.g. Retry-After-bearing 429s).""" + +from __future__ import annotations + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import RetryAfterMiddleware + +from ..flags import enabled + + +def build_retry_mw(flags: AgentFeatureFlags) -> RetryAfterMiddleware | None: + return ( + RetryAfterMiddleware(max_retries=3) + if enabled(flags, "enable_retry_after") + else None + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/tool_call_limit.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/tool_call_limit.py new file mode 100644 index 000000000..dcde81f37 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/resilience/tool_call_limit.py @@ -0,0 +1,21 @@ +"""Cap tool calls per thread / per run to bound infinite-loop blast radius.""" + +from __future__ import annotations + +from langchain.agents.middleware import ToolCallLimitMiddleware + +from app.agents.new_chat.feature_flags import AgentFeatureFlags + +from ..flags import enabled + + +def build_tool_call_limit_mw( + flags: AgentFeatureFlags, +) -> ToolCallLimitMiddleware | None: + if not enabled(flags, "enable_tool_call_limit"): + return None + return ToolCallLimitMiddleware( + thread_limit=300, + run_limit=80, + exit_behavior="continue", + ) From 67036448f9ae42ae788780e2893e13eb86e2f842 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:51:17 +0200 Subject: [PATCH 14/96] refactor(multi-agent): add shared middleware factory per concept --- .../middleware/shared/anthropic_cache.py | 9 +++++++ .../middleware/shared/compaction.py | 14 +++++++++++ .../middleware/shared/file_intent.py | 11 ++++++++ .../middleware/shared/filesystem.py | 25 +++++++++++++++++++ .../middleware/shared/memory.py | 19 ++++++++++++++ .../middleware/shared/patch_tool_calls.py | 9 +++++++ .../middleware/shared/todos.py | 9 +++++++ 7 files changed, 96 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/anthropic_cache.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/compaction.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/file_intent.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/memory.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/patch_tool_calls.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/shared/todos.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/anthropic_cache.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/anthropic_cache.py new file mode 100644 index 000000000..f99fb9c7f --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/anthropic_cache.py @@ -0,0 +1,9 @@ +"""Anthropic prompt caching annotations on system/tool/message blocks.""" + +from __future__ import annotations + +from langchain_anthropic.middleware import AnthropicPromptCachingMiddleware + + +def build_anthropic_cache_mw() -> AnthropicPromptCachingMiddleware: + return AnthropicPromptCachingMiddleware(unsupported_model_behavior="ignore") diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/compaction.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/compaction.py new file mode 100644 index 000000000..b59e7d2c4 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/compaction.py @@ -0,0 +1,14 @@ +"""Context-window summarization with SurfSense protected sections.""" + +from __future__ import annotations + +from typing import Any + +from deepagents.backends import StateBackend +from langchain_core.language_models import BaseChatModel + +from app.agents.new_chat.middleware import create_surfsense_compaction_middleware + + +def build_compaction_mw(llm: BaseChatModel) -> Any: + return create_surfsense_compaction_middleware(llm, StateBackend) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/file_intent.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/file_intent.py new file mode 100644 index 000000000..5ff65aa12 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/file_intent.py @@ -0,0 +1,11 @@ +"""File-intent classifier that gates strict write contracts.""" + +from __future__ import annotations + +from langchain_core.language_models import BaseChatModel + +from app.agents.new_chat.middleware import FileIntentMiddleware + + +def build_file_intent_mw(llm: BaseChatModel) -> FileIntentMiddleware: + return FileIntentMiddleware(llm=llm) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem.py new file mode 100644 index 000000000..9481f5167 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/filesystem.py @@ -0,0 +1,25 @@ +"""SurfSense filesystem tools/middleware.""" + +from __future__ import annotations + +from typing import Any + +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.agents.new_chat.middleware import SurfSenseFilesystemMiddleware + + +def build_filesystem_mw( + *, + backend_resolver: Any, + filesystem_mode: FilesystemMode, + search_space_id: int, + user_id: str | None, + thread_id: int | None, +) -> SurfSenseFilesystemMiddleware: + return SurfSenseFilesystemMiddleware( + backend=backend_resolver, + filesystem_mode=filesystem_mode, + search_space_id=search_space_id, + created_by_id=user_id, + thread_id=thread_id, + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/memory.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/memory.py new file mode 100644 index 000000000..9316b3e21 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/memory.py @@ -0,0 +1,19 @@ +"""User/team memory injection prepended to the conversation.""" + +from __future__ import annotations + +from app.agents.new_chat.middleware import MemoryInjectionMiddleware +from app.db import ChatVisibility + + +def build_memory_mw( + *, + user_id: str | None, + search_space_id: int, + visibility: ChatVisibility, +) -> MemoryInjectionMiddleware: + return MemoryInjectionMiddleware( + user_id=user_id, + search_space_id=search_space_id, + thread_visibility=visibility, + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/patch_tool_calls.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/patch_tool_calls.py new file mode 100644 index 000000000..50036dbbe --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/patch_tool_calls.py @@ -0,0 +1,9 @@ +"""Repair dangling tool-call sequences before each agent turn.""" + +from __future__ import annotations + +from deepagents.middleware.patch_tool_calls import PatchToolCallsMiddleware + + +def build_patch_tool_calls_mw() -> PatchToolCallsMiddleware: + return PatchToolCallsMiddleware() diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/todos.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/todos.py new file mode 100644 index 000000000..ea9173a1d --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/todos.py @@ -0,0 +1,9 @@ +"""Todo-list middleware (each consumer needs its own instance).""" + +from __future__ import annotations + +from langchain.agents.middleware import TodoListMiddleware + + +def build_todos_mw() -> TodoListMiddleware: + return TodoListMiddleware() From 6a4dacda72bce83c6ce072e07ed928bdfe851240 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:53:49 +0200 Subject: [PATCH 15/96] refactor(multi-agent): add main-agent observability and lifecycle middleware factories --- .../middleware/main_agent/__init__.py | 0 .../middleware/main_agent/action_log.py | 36 ++++++++++++++ .../middleware/main_agent/anonymous_doc.py | 16 ++++++ .../middleware/main_agent/busy_mutex.py | 12 +++++ .../middleware/main_agent/otel.py | 12 +++++ .../middleware/main_agent/plugins.py | 49 +++++++++++++++++++ 6 files changed, 125 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/__init__.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/action_log.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/anonymous_doc.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/busy_mutex.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/otel.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/plugins.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/action_log.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/action_log.py new file mode 100644 index 000000000..c9f893d97 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/action_log.py @@ -0,0 +1,36 @@ +"""Audit row per tool call (reversibility metadata).""" + +from __future__ import annotations + +import logging + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import ActionLogMiddleware +from app.agents.new_chat.tools.registry import BUILTIN_TOOLS + +from ..shared.flags import enabled + + +def build_action_log_mw( + *, + flags: AgentFeatureFlags, + thread_id: int | None, + search_space_id: int, + user_id: str | None, +) -> ActionLogMiddleware | None: + if not enabled(flags, "enable_action_log") or thread_id is None: + return None + try: + tool_defs_by_name = {td.name: td for td in BUILTIN_TOOLS} + return ActionLogMiddleware( + thread_id=thread_id, + search_space_id=search_space_id, + user_id=user_id, + tool_definitions=tool_defs_by_name, + ) + except Exception: # pragma: no cover - defensive + logging.warning( + "ActionLogMiddleware init failed; running without it.", + exc_info=True, + ) + return None diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/anonymous_doc.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/anonymous_doc.py new file mode 100644 index 000000000..afd54a2d3 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/anonymous_doc.py @@ -0,0 +1,16 @@ +"""Anonymous document hydration from Redis (cloud only).""" + +from __future__ import annotations + +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.agents.new_chat.middleware import AnonymousDocumentMiddleware + + +def build_anonymous_doc_mw( + *, + filesystem_mode: FilesystemMode, + anon_session_id: str | None, +) -> AnonymousDocumentMiddleware | None: + if filesystem_mode != FilesystemMode.CLOUD: + return None + return AnonymousDocumentMiddleware(anon_session_id=anon_session_id) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/busy_mutex.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/busy_mutex.py new file mode 100644 index 000000000..0ea53bf16 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/busy_mutex.py @@ -0,0 +1,12 @@ +"""Per-thread cooperative lock around the whole turn.""" + +from __future__ import annotations + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import BusyMutexMiddleware + +from ..shared.flags import enabled + + +def build_busy_mutex_mw(flags: AgentFeatureFlags) -> BusyMutexMiddleware | None: + return BusyMutexMiddleware() if enabled(flags, "enable_busy_mutex") else None diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/otel.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/otel.py new file mode 100644 index 000000000..bd7516e65 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/otel.py @@ -0,0 +1,12 @@ +"""OTel spans on model and tool calls.""" + +from __future__ import annotations + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import OtelSpanMiddleware + +from ..shared.flags import enabled + + +def build_otel_mw(flags: AgentFeatureFlags) -> OtelSpanMiddleware | None: + return OtelSpanMiddleware() if enabled(flags, "enable_otel") else None diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/plugins.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/plugins.py new file mode 100644 index 000000000..4418e3806 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/plugins.py @@ -0,0 +1,49 @@ +"""Tail-of-stack plugin slot driven by env allowlist.""" + +from __future__ import annotations + +import logging +from typing import Any + +from langchain_core.language_models import BaseChatModel + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.plugin_loader import ( + PluginContext, + load_allowed_plugin_names_from_env, + load_plugin_middlewares, +) +from app.db import ChatVisibility + +from ..shared.flags import enabled + + +def build_plugin_middlewares( + *, + flags: AgentFeatureFlags, + search_space_id: int, + user_id: str | None, + visibility: ChatVisibility, + llm: BaseChatModel, +) -> list[Any]: + if not enabled(flags, "enable_plugin_loader"): + return [] + try: + allowed_names = load_allowed_plugin_names_from_env() + if not allowed_names: + return [] + return load_plugin_middlewares( + PluginContext.build( + search_space_id=search_space_id, + user_id=user_id, + thread_visibility=visibility, + llm=llm, + ), + allowed_plugin_names=allowed_names, + ) + except Exception: # pragma: no cover - defensive + logging.warning( + "Plugin loader failed; continuing without plugins.", + exc_info=True, + ) + return [] From 390dc9307fac6733df8848b043cac7d7fabd1365 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:53:59 +0200 Subject: [PATCH 16/96] refactor(multi-agent): add main-agent knowledge middleware factories --- .../middleware/main_agent/kb_persistence.py | 23 ++++++++++++++++ .../main_agent/knowledge_priority.py | 27 +++++++++++++++++++ .../middleware/main_agent/knowledge_tree.py | 23 ++++++++++++++++ 3 files changed, 73 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/kb_persistence.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_priority.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_tree.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/kb_persistence.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/kb_persistence.py new file mode 100644 index 000000000..4b27581e7 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/kb_persistence.py @@ -0,0 +1,23 @@ +"""Commit staged cloud filesystem mutations to Postgres at end of turn.""" + +from __future__ import annotations + +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.agents.new_chat.middleware import KnowledgeBasePersistenceMiddleware + + +def build_kb_persistence_mw( + *, + filesystem_mode: FilesystemMode, + search_space_id: int, + user_id: str | None, + thread_id: int | None, +) -> KnowledgeBasePersistenceMiddleware | None: + if filesystem_mode != FilesystemMode.CLOUD: + return None + return KnowledgeBasePersistenceMiddleware( + search_space_id=search_space_id, + created_by_id=user_id, + filesystem_mode=filesystem_mode, + thread_id=thread_id, + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_priority.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_priority.py new file mode 100644 index 000000000..395d2a7af --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_priority.py @@ -0,0 +1,27 @@ +"""KB priority planner: injection.""" + +from __future__ import annotations + +from langchain_core.language_models import BaseChatModel + +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.agents.new_chat.middleware import KnowledgePriorityMiddleware + + +def build_knowledge_priority_mw( + *, + llm: BaseChatModel, + search_space_id: int, + filesystem_mode: FilesystemMode, + available_connectors: list[str] | None, + available_document_types: list[str] | None, + mentioned_document_ids: list[int] | None, +) -> KnowledgePriorityMiddleware: + return KnowledgePriorityMiddleware( + llm=llm, + search_space_id=search_space_id, + filesystem_mode=filesystem_mode, + available_connectors=available_connectors, + available_document_types=available_document_types, + mentioned_document_ids=mentioned_document_ids, + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_tree.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_tree.py new file mode 100644 index 000000000..404082401 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/knowledge_tree.py @@ -0,0 +1,23 @@ +""" injection (cloud only).""" + +from __future__ import annotations + +from langchain_core.language_models import BaseChatModel + +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.agents.new_chat.middleware import KnowledgeTreeMiddleware + + +def build_knowledge_tree_mw( + *, + filesystem_mode: FilesystemMode, + search_space_id: int, + llm: BaseChatModel, +) -> KnowledgeTreeMiddleware | None: + if filesystem_mode != FilesystemMode.CLOUD: + return None + return KnowledgeTreeMiddleware( + search_space_id=search_space_id, + filesystem_mode=filesystem_mode, + llm=llm, + ) From b0ee44b2f15adfdf379c2b7f1967733a66a8e090 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:54:13 +0200 Subject: [PATCH 17/96] refactor(multi-agent): add main-agent safety and llm-shaping middleware factories --- .../middleware/main_agent/context_editing.py | 50 +++++++++++++++++++ .../middleware/main_agent/dedup_hitl.py | 13 +++++ .../middleware/main_agent/doom_loop.py | 12 +++++ .../middleware/main_agent/noop_injection.py | 12 +++++ .../middleware/main_agent/repair.py | 50 +++++++++++++++++++ .../middleware/main_agent/selector.py | 39 +++++++++++++++ .../middleware/main_agent/skills.py | 39 +++++++++++++++ 7 files changed, 215 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/context_editing.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/dedup_hitl.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/noop_injection.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/repair.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/selector.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/skills.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/context_editing.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/context_editing.py new file mode 100644 index 000000000..e8f99933e --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/context_editing.py @@ -0,0 +1,50 @@ +"""Spill + clear-tool-uses passes to keep payloads under budget.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from langchain_core.tools import BaseTool + +from app.agents.multi_agent_chat.main_agent.context_prune.prune_tool_names import ( + safe_exclude_tools, +) +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import ( + ClearToolUsesEdit, + SpillingContextEditingMiddleware, + SpillToBackendEdit, +) + +from ..shared.flags import enabled + + +def build_context_editing_mw( + *, + flags: AgentFeatureFlags, + max_input_tokens: int | None, + tools: Sequence[BaseTool], + backend_resolver: Any, +) -> SpillingContextEditingMiddleware | None: + if not enabled(flags, "enable_context_editing") or not max_input_tokens: + return None + spill_edit = SpillToBackendEdit( + trigger=int(max_input_tokens * 0.55), + clear_at_least=int(max_input_tokens * 0.15), + keep=5, + exclude_tools=safe_exclude_tools(tools), + clear_tool_inputs=True, + ) + clear_edit = ClearToolUsesEdit( + trigger=int(max_input_tokens * 0.55), + clear_at_least=int(max_input_tokens * 0.15), + keep=5, + exclude_tools=safe_exclude_tools(tools), + clear_tool_inputs=True, + placeholder="[cleared - older tool output trimmed for context]", + ) + return SpillingContextEditingMiddleware( + edits=[spill_edit, clear_edit], + backend_resolver=backend_resolver, + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/dedup_hitl.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/dedup_hitl.py new file mode 100644 index 000000000..66cae300b --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/dedup_hitl.py @@ -0,0 +1,13 @@ +"""Drop duplicate HITL tool calls before execution.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from langchain_core.tools import BaseTool + +from app.agents.new_chat.middleware import DedupHITLToolCallsMiddleware + + +def build_dedup_hitl_mw(tools: Sequence[BaseTool]) -> DedupHITLToolCallsMiddleware: + return DedupHITLToolCallsMiddleware(agent_tools=list(tools)) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py new file mode 100644 index 000000000..a0b294092 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py @@ -0,0 +1,12 @@ +"""Stop N identical tool calls in a row via interrupt.""" + +from __future__ import annotations + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import DoomLoopMiddleware + +from ..shared.flags import enabled + + +def build_doom_loop_mw(flags: AgentFeatureFlags) -> DoomLoopMiddleware | None: + return DoomLoopMiddleware(threshold=3) if enabled(flags, "enable_doom_loop") else None diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/noop_injection.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/noop_injection.py new file mode 100644 index 000000000..6e6467ad0 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/noop_injection.py @@ -0,0 +1,12 @@ +"""Provider-compat: append a `_noop` tool when tools=[] but history has tool calls.""" + +from __future__ import annotations + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import NoopInjectionMiddleware + +from ..shared.flags import enabled + + +def build_noop_injection_mw(flags: AgentFeatureFlags) -> NoopInjectionMiddleware | None: + return NoopInjectionMiddleware() if enabled(flags, "enable_compaction_v2") else None diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/repair.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/repair.py new file mode 100644 index 000000000..378b61be1 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/repair.py @@ -0,0 +1,50 @@ +"""Repair miscased / unknown tool names to the registered set or invalid_tool.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from langchain_core.tools import BaseTool + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.middleware import ToolCallNameRepairMiddleware + +from ..shared.flags import enabled + +# deepagents-built-in tool names the repair pass treats as known. +_DEEPAGENT_BUILTIN_TOOL_NAMES: frozenset[str] = frozenset( + { + "write_todos", + "ls", + "read_file", + "write_file", + "edit_file", + "glob", + "grep", + "execute", + "task", + "mkdir", + "cd", + "pwd", + "move_file", + "rm", + "rmdir", + "list_tree", + "execute_code", + } +) + + +def build_repair_mw( + *, + flags: AgentFeatureFlags, + tools: Sequence[BaseTool], +) -> ToolCallNameRepairMiddleware | None: + if not enabled(flags, "enable_tool_call_repair"): + return None + registered_names: set[str] = {t.name for t in tools} + registered_names |= _DEEPAGENT_BUILTIN_TOOL_NAMES + return ToolCallNameRepairMiddleware( + registered_tool_names=registered_names, + fuzzy_match_threshold=None, + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/selector.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/selector.py new file mode 100644 index 000000000..8e7a32be8 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/selector.py @@ -0,0 +1,39 @@ +"""LLM-based tool subset selection (only when >30 tools).""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence + +from langchain.agents.middleware import LLMToolSelectorMiddleware +from langchain_core.tools import BaseTool + +from app.agents.new_chat.feature_flags import AgentFeatureFlags + +from ..shared.flags import enabled + + +def build_selector_mw( + *, + flags: AgentFeatureFlags, + tools: Sequence[BaseTool], +) -> LLMToolSelectorMiddleware | None: + if not enabled(flags, "enable_llm_tool_selector") or len(tools) <= 30: + return None + try: + return LLMToolSelectorMiddleware( + model="openai:gpt-4o-mini", + max_tools=12, + always_include=[ + name + for name in ( + "update_memory", + "get_connected_accounts", + "scrape_webpage", + ) + if name in {t.name for t in tools} + ], + ) + except Exception: + logging.warning("LLMToolSelectorMiddleware init failed; skipping.") + return None diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/skills.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/skills.py new file mode 100644 index 000000000..63a57c5a0 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/skills.py @@ -0,0 +1,39 @@ +"""Skill discovery + injection.""" + +from __future__ import annotations + +import logging + +from deepagents.middleware.skills import SkillsMiddleware + +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.agents.new_chat.middleware import ( + build_skills_backend_factory, + default_skills_sources, +) + +from ..shared.flags import enabled + + +def build_skills_mw( + *, + flags: AgentFeatureFlags, + filesystem_mode: FilesystemMode, + search_space_id: int, +) -> SkillsMiddleware | None: + if not enabled(flags, "enable_skills"): + return None + try: + skills_factory = build_skills_backend_factory( + search_space_id=search_space_id + if filesystem_mode == FilesystemMode.CLOUD + else None, + ) + return SkillsMiddleware( + backend=skills_factory, + sources=default_skills_sources(), + ) + except Exception as exc: # pragma: no cover - defensive + logging.warning("SkillsMiddleware init failed; skipping: %s", exc) + return None From 611fef8666990b4994d8f990b8b796a36027c6fd Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:54:27 +0200 Subject: [PATCH 18/96] refactor(multi-agent): add subagent extras builder and drop filesystem from registry subagents --- .../middleware/subagent/__init__.py | 0 .../middleware/subagent/extras.py | 28 +++++++++++++++++++ 2 files changed, 28 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/subagent/__init__.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/subagent/extras.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/subagent/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/subagent/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/subagent/extras.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/subagent/extras.py new file mode 100644 index 000000000..46dca8a81 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/subagent/extras.py @@ -0,0 +1,28 @@ +"""Extra middleware threaded into every registry subagent's stack. + +Registry subagents are scoped to one domain (deliverables, research, memory, +connectors, MCP) and never read or write the SurfSense filesystem — that +capability belongs to the main agent and is delegated to the general-purpose +subagent as an escape hatch. Keeping FS off the registry stacks avoids +polluting their tool surface with FS tools they never act on. +""" + +from __future__ import annotations + +from typing import Any + +from ..shared.permissions import PermissionContext +from ..shared.resilience import ResilienceBundle +from ..shared.todos import build_todos_mw + + +def build_subagent_extras( + *, + permissions: PermissionContext, + resilience: ResilienceBundle, +) -> list[Any]: + extras: list[Any] = [build_todos_mw()] + if permissions.subagent_deny_mw is not None: + extras.append(permissions.subagent_deny_mw) + extras.extend(resilience.as_list()) + return extras From 7690e8b278e525cae61291ac9bec6cfa773b06db Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:54:45 +0200 Subject: [PATCH 19/96] refactor(multi-agent): relocate general-purpose subagent next to other builtins --- .../builtins/general_purpose/__init__.py | 0 .../builtins/general_purpose/agent.py | 105 ++++++++++++++++++ 2 files changed, 105 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/general_purpose/__init__.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/general_purpose/agent.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/general_purpose/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/general_purpose/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/general_purpose/agent.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/general_purpose/agent.py new file mode 100644 index 000000000..1c3c44f12 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/general_purpose/agent.py @@ -0,0 +1,105 @@ +"""General-purpose subagent for the multi-agent main agent.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any, cast + +from deepagents import SubAgent +from deepagents.middleware.patch_tool_calls import PatchToolCallsMiddleware +from deepagents.middleware.subagents import GENERAL_PURPOSE_SUBAGENT +from langchain_anthropic.middleware import AnthropicPromptCachingMiddleware +from langchain_core.language_models import BaseChatModel +from langchain_core.tools import BaseTool + +from app.agents.multi_agent_chat.middleware.shared.anthropic_cache import ( + build_anthropic_cache_mw, +) +from app.agents.multi_agent_chat.middleware.shared.compaction import ( + build_compaction_mw, +) +from app.agents.multi_agent_chat.middleware.shared.file_intent import ( + build_file_intent_mw, +) +from app.agents.multi_agent_chat.middleware.shared.filesystem import ( + build_filesystem_mw, +) +from app.agents.multi_agent_chat.middleware.shared.patch_tool_calls import ( + build_patch_tool_calls_mw, +) +from app.agents.multi_agent_chat.middleware.shared.permissions import ( + PermissionContext, +) +from app.agents.multi_agent_chat.middleware.shared.resilience import ( + ResilienceBundle, +) +from app.agents.multi_agent_chat.middleware.shared.todos import build_todos_mw +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.agents.new_chat.middleware import MemoryInjectionMiddleware + +NAME = "general-purpose" + + +def build_subagent( + *, + llm: BaseChatModel, + tools: Sequence[BaseTool], + backend_resolver: Any, + filesystem_mode: FilesystemMode, + search_space_id: int, + user_id: str | None, + thread_id: int | None, + permissions: PermissionContext, + resilience: ResilienceBundle, + memory_mw: MemoryInjectionMiddleware, +) -> SubAgent: + """Deny + resilience inserts encapsulated here so the orchestrator never mutates the list.""" + middleware: list[Any] = [ + build_todos_mw(), + memory_mw, + build_file_intent_mw(llm), + build_filesystem_mw( + backend_resolver=backend_resolver, + filesystem_mode=filesystem_mode, + search_space_id=search_space_id, + user_id=user_id, + thread_id=thread_id, + ), + build_compaction_mw(llm), + build_patch_tool_calls_mw(), + build_anthropic_cache_mw(), + ] + + if permissions.subagent_deny_mw is not None: + patch_idx = next( + ( + i + for i, m in enumerate(middleware) + if isinstance(m, PatchToolCallsMiddleware) + ), + len(middleware), + ) + middleware.insert(patch_idx, permissions.subagent_deny_mw) + + resilience_mws = resilience.as_list() + if resilience_mws: + cache_idx = next( + ( + i + for i, m in enumerate(middleware) + if isinstance(m, AnthropicPromptCachingMiddleware) + ), + len(middleware), + ) + for offset, mw in enumerate(resilience_mws): + middleware.insert(cache_idx + offset, mw) + + spec: dict[str, Any] = { + **GENERAL_PURPOSE_SUBAGENT, + "model": llm, + "tools": tools, + "middleware": middleware, + } + if permissions.general_purpose_interrupt_on: + spec["interrupt_on"] = permissions.general_purpose_interrupt_on + return cast(SubAgent, spec) From 5abae09435eaae8885978d12f48898eed83751dc Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:55:03 +0200 Subject: [PATCH 20/96] refactor(multi-agent): add slimmed orchestrator at middleware/stack.py --- .../multi_agent_chat/middleware/__init__.py | 7 + .../multi_agent_chat/middleware/stack.py | 212 ++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/__init__.py create mode 100644 surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/__init__.py new file mode 100644 index 000000000..e6eed9fbe --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/__init__.py @@ -0,0 +1,7 @@ +"""Multi-agent middleware stack assembly.""" + +from __future__ import annotations + +from .stack import build_main_agent_deepagent_middleware + +__all__ = ["build_main_agent_deepagent_middleware"] diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py new file mode 100644 index 000000000..563332986 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py @@ -0,0 +1,212 @@ +"""Main-agent middleware list assembly: one line per slot.""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from typing import Any + +from deepagents import SubAgent +from deepagents.backends import StateBackend +from langchain_core.language_models import BaseChatModel +from langchain_core.tools import BaseTool +from langgraph.types import Checkpointer + +from app.agents.multi_agent_chat.subagents import ( + build_subagents, + get_subagents_to_exclude, +) +from app.agents.multi_agent_chat.subagents.builtins.general_purpose.agent import ( + build_subagent as build_general_purpose_subagent, +) +from app.agents.multi_agent_chat.subagents.shared.permissions import ToolsPermissions +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.db import ChatVisibility + +from .main_agent.action_log import build_action_log_mw +from .main_agent.anonymous_doc import build_anonymous_doc_mw +from .main_agent.busy_mutex import build_busy_mutex_mw +from .main_agent.checkpointed_subagent_middleware import ( + SurfSenseCheckpointedSubAgentMiddleware, +) +from .main_agent.context_editing import build_context_editing_mw +from .main_agent.dedup_hitl import build_dedup_hitl_mw +from .main_agent.doom_loop import build_doom_loop_mw +from .main_agent.kb_persistence import build_kb_persistence_mw +from .main_agent.knowledge_priority import build_knowledge_priority_mw +from .main_agent.knowledge_tree import build_knowledge_tree_mw +from .main_agent.noop_injection import build_noop_injection_mw +from .main_agent.otel import build_otel_mw +from .main_agent.plugins import build_plugin_middlewares +from .main_agent.repair import build_repair_mw +from .main_agent.selector import build_selector_mw +from .main_agent.skills import build_skills_mw +from .shared.anthropic_cache import build_anthropic_cache_mw +from .shared.compaction import build_compaction_mw +from .shared.file_intent import build_file_intent_mw +from .shared.filesystem import build_filesystem_mw +from .shared.memory import build_memory_mw +from .shared.patch_tool_calls import build_patch_tool_calls_mw +from .shared.permissions import ( + build_full_permission_mw, + build_permission_context, +) +from .shared.resilience import build_resilience_bundle +from .shared.todos import build_todos_mw +from .subagent.extras import build_subagent_extras + + +def build_main_agent_deepagent_middleware( + *, + llm: BaseChatModel, + tools: Sequence[BaseTool], + backend_resolver: Any, + filesystem_mode: FilesystemMode, + search_space_id: int, + user_id: str | None, + thread_id: int | None, + visibility: ChatVisibility, + anon_session_id: str | None, + available_connectors: list[str] | None, + available_document_types: list[str] | None, + mentioned_document_ids: list[int] | None, + max_input_tokens: int | None, + flags: AgentFeatureFlags, + subagent_dependencies: dict[str, Any], + checkpointer: Checkpointer, + mcp_tools_by_agent: dict[str, ToolsPermissions] | None = None, + disabled_tools: list[str] | None = None, +) -> list[Any]: + """Ordered middleware for ``create_agent`` (None entries already stripped).""" + permissions = build_permission_context( + flags=flags, + filesystem_mode=filesystem_mode, + tools=tools, + available_connectors=available_connectors, + ) + resilience = build_resilience_bundle(flags) + + # Single instance threaded into both the main-agent stack and the general-purpose subagent. + memory_mw = build_memory_mw( + user_id=user_id, + search_space_id=search_space_id, + visibility=visibility, + ) + + general_purpose_subagent = build_general_purpose_subagent( + llm=llm, + tools=tools, + backend_resolver=backend_resolver, + filesystem_mode=filesystem_mode, + search_space_id=search_space_id, + user_id=user_id, + thread_id=thread_id, + permissions=permissions, + resilience=resilience, + memory_mw=memory_mw, + ) + + subagents_registry: list[SubAgent] = [] + try: + subagent_extras = build_subagent_extras( + permissions=permissions, + resilience=resilience, + ) + subagents_registry = build_subagents( + dependencies=subagent_dependencies, + model=llm, + extra_middleware=subagent_extras, + mcp_tools_by_agent=mcp_tools_by_agent or {}, + exclude=get_subagents_to_exclude(available_connectors), + disabled_tools=disabled_tools, + ) + logging.info( + "Subagents registry: %s", + [s["name"] for s in subagents_registry], + ) + except Exception: + logging.exception("Subagents registry build failed") + raise + + subagents: list[SubAgent] = [general_purpose_subagent, *subagents_registry] + + stack: list[Any] = [ + build_busy_mutex_mw(flags), + build_otel_mw(flags), + build_todos_mw(), + memory_mw, + build_anonymous_doc_mw( + filesystem_mode=filesystem_mode, anon_session_id=anon_session_id + ), + build_knowledge_tree_mw( + filesystem_mode=filesystem_mode, + search_space_id=search_space_id, + llm=llm, + ), + build_knowledge_priority_mw( + llm=llm, + search_space_id=search_space_id, + filesystem_mode=filesystem_mode, + available_connectors=available_connectors, + available_document_types=available_document_types, + mentioned_document_ids=mentioned_document_ids, + ), + build_file_intent_mw(llm), + build_filesystem_mw( + backend_resolver=backend_resolver, + filesystem_mode=filesystem_mode, + search_space_id=search_space_id, + user_id=user_id, + thread_id=thread_id, + ), + build_kb_persistence_mw( + filesystem_mode=filesystem_mode, + search_space_id=search_space_id, + user_id=user_id, + thread_id=thread_id, + ), + build_skills_mw( + flags=flags, + filesystem_mode=filesystem_mode, + search_space_id=search_space_id, + ), + SurfSenseCheckpointedSubAgentMiddleware( + checkpointer=checkpointer, + backend=StateBackend, + subagents=subagents, + ), + build_selector_mw(flags=flags, tools=tools), + resilience.model_call_limit, + resilience.tool_call_limit, + build_context_editing_mw( + flags=flags, + max_input_tokens=max_input_tokens, + tools=tools, + backend_resolver=backend_resolver, + ), + build_compaction_mw(llm), + build_noop_injection_mw(flags), + resilience.retry, + resilience.fallback, + build_repair_mw(flags=flags, tools=tools), + build_full_permission_mw(permissions.rulesets), + build_doom_loop_mw(flags), + build_action_log_mw( + flags=flags, + thread_id=thread_id, + search_space_id=search_space_id, + user_id=user_id, + ), + build_patch_tool_calls_mw(), + build_dedup_hitl_mw(tools), + *build_plugin_middlewares( + flags=flags, + search_space_id=search_space_id, + user_id=user_id, + visibility=visibility, + llm=llm, + ), + build_anthropic_cache_mw(), + ] + return [m for m in stack if m is not None] From 73272ce348e5da1dd92ae4c301a5dff8eb215fef Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 20:55:38 +0200 Subject: [PATCH 21/96] refactor(multi-agent): switch compile graph to new orchestrator and drop deepagent_stack --- .../main_agent/graph/compile_graph_sync.py | 5 +- .../main_agent/graph/middleware/__init__.py | 7 - .../graph/middleware/deepagent_stack.py | 539 ------------------ .../__init__.py | 0 .../config.py | 0 .../constants.py | 0 .../middleware.py | 0 .../propagation.py | 0 .../resume.py | 0 .../task_tool.py | 0 .../test_hitl_bridge.py | 2 +- .../test_resume_helpers.py | 2 +- 12 files changed, 5 insertions(+), 550 deletions(-) delete mode 100644 surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/__init__.py delete mode 100644 surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py rename surfsense_backend/app/agents/multi_agent_chat/{main_agent/graph/middleware => middleware/main_agent}/checkpointed_subagent_middleware/__init__.py (100%) rename surfsense_backend/app/agents/multi_agent_chat/{main_agent/graph/middleware => middleware/main_agent}/checkpointed_subagent_middleware/config.py (100%) rename surfsense_backend/app/agents/multi_agent_chat/{main_agent/graph/middleware => middleware/main_agent}/checkpointed_subagent_middleware/constants.py (100%) rename surfsense_backend/app/agents/multi_agent_chat/{main_agent/graph/middleware => middleware/main_agent}/checkpointed_subagent_middleware/middleware.py (100%) rename surfsense_backend/app/agents/multi_agent_chat/{main_agent/graph/middleware => middleware/main_agent}/checkpointed_subagent_middleware/propagation.py (100%) rename surfsense_backend/app/agents/multi_agent_chat/{main_agent/graph/middleware => middleware/main_agent}/checkpointed_subagent_middleware/resume.py (100%) rename surfsense_backend/app/agents/multi_agent_chat/{main_agent/graph/middleware => middleware/main_agent}/checkpointed_subagent_middleware/task_tool.py (100%) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/compile_graph_sync.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/compile_graph_sync.py index 7afa30a31..4ed94bf7b 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/compile_graph_sync.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/compile_graph_sync.py @@ -11,6 +11,9 @@ from langchain_core.language_models import BaseChatModel from langchain_core.tools import BaseTool from langgraph.types import Checkpointer +from app.agents.multi_agent_chat.middleware import ( + build_main_agent_deepagent_middleware, +) from app.agents.multi_agent_chat.subagents.shared.permissions import ( ToolsPermissions, ) @@ -19,8 +22,6 @@ from app.agents.new_chat.feature_flags import AgentFeatureFlags from app.agents.new_chat.filesystem_selection import FilesystemMode from app.db import ChatVisibility -from .middleware import build_main_agent_deepagent_middleware - def build_compiled_agent_graph_sync( *, diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/__init__.py deleted file mode 100644 index 757ee02f8..000000000 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""Main-agent graph middleware assembly (SurfSense + LangChain + deepagents).""" - -from __future__ import annotations - -from .deepagent_stack import build_main_agent_deepagent_middleware - -__all__ = ["build_main_agent_deepagent_middleware"] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py deleted file mode 100644 index 1d6a8763e..000000000 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/deepagent_stack.py +++ /dev/null @@ -1,539 +0,0 @@ -"""Assemble the main-agent deep-agent middleware list (LangChain + SurfSense + deepagents).""" - -from __future__ import annotations - -import logging -from collections.abc import Sequence -from typing import Any - -from deepagents import SubAgent -from deepagents.backends import StateBackend -from deepagents.middleware.patch_tool_calls import PatchToolCallsMiddleware -from deepagents.middleware.skills import SkillsMiddleware -from deepagents.middleware.subagents import GENERAL_PURPOSE_SUBAGENT -from langchain.agents.middleware import ( - LLMToolSelectorMiddleware, - ModelCallLimitMiddleware, - TodoListMiddleware, - ToolCallLimitMiddleware, -) -from langchain_anthropic.middleware import AnthropicPromptCachingMiddleware -from langchain_core.language_models import BaseChatModel -from langchain_core.tools import BaseTool -from langgraph.types import Checkpointer - -from app.agents.multi_agent_chat.subagents import ( - build_subagents, - get_subagents_to_exclude, -) -from app.agents.multi_agent_chat.subagents.shared.permissions import ( - ToolsPermissions, -) -from app.agents.new_chat.feature_flags import AgentFeatureFlags -from app.agents.new_chat.filesystem_selection import FilesystemMode -from app.agents.new_chat.middleware import ( - ActionLogMiddleware, - AnonymousDocumentMiddleware, - BusyMutexMiddleware, - ClearToolUsesEdit, - DedupHITLToolCallsMiddleware, - DoomLoopMiddleware, - FileIntentMiddleware, - KnowledgeBasePersistenceMiddleware, - KnowledgePriorityMiddleware, - KnowledgeTreeMiddleware, - MemoryInjectionMiddleware, - NoopInjectionMiddleware, - OtelSpanMiddleware, - PermissionMiddleware, - RetryAfterMiddleware, - SpillingContextEditingMiddleware, - SpillToBackendEdit, - SurfSenseFilesystemMiddleware, - ToolCallNameRepairMiddleware, - build_skills_backend_factory, - create_surfsense_compaction_middleware, - default_skills_sources, -) -from app.agents.new_chat.middleware.scoped_model_fallback import ( - ScopedModelFallbackMiddleware, -) -from app.agents.new_chat.permissions import Rule, Ruleset -from app.agents.new_chat.plugin_loader import ( - PluginContext, - load_allowed_plugin_names_from_env, - load_plugin_middlewares, -) -from app.agents.new_chat.tools.registry import BUILTIN_TOOLS -from app.db import ChatVisibility - -from ...context_prune.prune_tool_names import safe_exclude_tools -from .checkpointed_subagent_middleware import SurfSenseCheckpointedSubAgentMiddleware - - -def build_main_agent_deepagent_middleware( - *, - llm: BaseChatModel, - tools: Sequence[BaseTool], - backend_resolver: Any, - filesystem_mode: FilesystemMode, - search_space_id: int, - user_id: str | None, - thread_id: int | None, - visibility: ChatVisibility, - anon_session_id: str | None, - available_connectors: list[str] | None, - available_document_types: list[str] | None, - mentioned_document_ids: list[int] | None, - max_input_tokens: int | None, - flags: AgentFeatureFlags, - subagent_dependencies: dict[str, Any], - checkpointer: Checkpointer, - mcp_tools_by_agent: dict[str, ToolsPermissions] | None = None, - disabled_tools: list[str] | None = None, -) -> list[Any]: - """Build ordered middleware for ``create_agent`` (Nones already stripped).""" - _memory_middleware = MemoryInjectionMiddleware( - user_id=user_id, - search_space_id=search_space_id, - thread_visibility=visibility, - ) - - gp_middleware = [ - TodoListMiddleware(), - _memory_middleware, - FileIntentMiddleware(llm=llm), - SurfSenseFilesystemMiddleware( - backend=backend_resolver, - filesystem_mode=filesystem_mode, - search_space_id=search_space_id, - created_by_id=user_id, - thread_id=thread_id, - ), - create_surfsense_compaction_middleware(llm, StateBackend), - PatchToolCallsMiddleware(), - AnthropicPromptCachingMiddleware(unsupported_model_behavior="ignore"), - ] - - # Build permission rulesets up front so the GP subagent can mirror ``ask`` - # rules into ``interrupt_on``: tool calls emitted from within ``task`` runs - # never reach the parent's ``PermissionMiddleware``. - is_desktop_fs = filesystem_mode == FilesystemMode.DESKTOP_LOCAL_FOLDER - permission_enabled = flags.enable_permission and not flags.disable_new_agent_stack - permission_rulesets: list[Ruleset] = [] - if permission_enabled or is_desktop_fs: - permission_rulesets.append( - Ruleset( - rules=[Rule(permission="*", pattern="*", action="allow")], - origin="surfsense_defaults", - ) - ) - if is_desktop_fs: - permission_rulesets.append( - Ruleset( - rules=[ - Rule(permission="rm", pattern="*", action="ask"), - Rule(permission="rmdir", pattern="*", action="ask"), - Rule(permission="move_file", pattern="*", action="ask"), - Rule(permission="edit_file", pattern="*", action="ask"), - Rule(permission="write_file", pattern="*", action="ask"), - ], - origin="desktop_safety", - ) - ) - - # Tools that self-prompt via ``request_approval`` must not also appear - # as ``ask`` rules — that would double-prompt the user for one call. - _tool_names_in_use = {t.name for t in tools} - - # Deny parent-bound tools whose ``required_connector`` is missing. - # No-op today (connector subagents are pruned upstream); guards future - # additions to the parent's tool list. - if permission_enabled: - _available_set = set(available_connectors or []) - _synthesized: list[Rule] = [] - for tool_def in BUILTIN_TOOLS: - if tool_def.name not in _tool_names_in_use: - continue - rc = tool_def.required_connector - if rc and rc not in _available_set: - _synthesized.append( - Rule(permission=tool_def.name, pattern="*", action="deny") - ) - if _synthesized: - permission_rulesets.append( - Ruleset(rules=_synthesized, origin="connector_synthesized") - ) - gp_interrupt_on: dict[str, bool] = { - rule.permission: True - for rs in permission_rulesets - for rule in rs.rules - if rule.action == "ask" and rule.permission in _tool_names_in_use - } - - general_purpose_spec: SubAgent = { # type: ignore[typeddict-unknown-key] - **GENERAL_PURPOSE_SUBAGENT, - "model": llm, - "tools": tools, - "middleware": gp_middleware, - } - if gp_interrupt_on: - general_purpose_spec["interrupt_on"] = gp_interrupt_on - - # Deny-only on subagents: ``task`` runs bypass the parent's - # PermissionMiddleware, while bucket-based ask gates own the ask path. - subagent_deny_rulesets: list[Ruleset] = [ - Ruleset( - rules=[r for r in rs.rules if r.action == "deny"], - origin=rs.origin, - ) - for rs in permission_rulesets - ] - subagent_deny_rulesets = [rs for rs in subagent_deny_rulesets if rs.rules] - - subagent_deny_permission_mw: PermissionMiddleware | None = ( - PermissionMiddleware(rulesets=subagent_deny_rulesets) - if subagent_deny_rulesets - else None - ) - - if subagent_deny_permission_mw is not None: - # Run deny check on already-repaired tool calls; insert before - # PatchToolCallsMiddleware (append if the slot moves). - _patch_idx = next( - ( - i - for i, m in enumerate(gp_middleware) - if isinstance(m, PatchToolCallsMiddleware) - ), - len(gp_middleware), - ) - gp_middleware.insert(_patch_idx, subagent_deny_permission_mw) - - # Defined early so the same instances reach both gp_middleware and - # subagent_extra_middleware below. - retry_mw = ( - RetryAfterMiddleware(max_retries=3) - if flags.enable_retry_after and not flags.disable_new_agent_stack - else None - ) - fallback_mw: ScopedModelFallbackMiddleware | None = None - if flags.enable_model_fallback and not flags.disable_new_agent_stack: - try: - fallback_mw = ScopedModelFallbackMiddleware( - "openai:gpt-4o-mini", - "anthropic:claude-3-5-haiku-20241022", - ) - except Exception: - logging.warning("ScopedModelFallbackMiddleware init failed; skipping.") - fallback_mw = None - - # Per-agent caps; counts are not summed across parent + subagents. - model_call_limit_mw = ( - ModelCallLimitMiddleware( - thread_limit=120, - run_limit=80, - exit_behavior="end", - ) - if flags.enable_model_call_limit and not flags.disable_new_agent_stack - else None - ) - tool_call_limit_mw = ( - ToolCallLimitMiddleware( - thread_limit=300, run_limit=80, exit_behavior="continue" - ) - if flags.enable_tool_call_limit and not flags.disable_new_agent_stack - else None - ) - - # gp_middleware is held by reference inside general_purpose_spec, so - # mutating it here propagates into the spec. - _gp_resilience: list[Any] = [ - m - for m in (retry_mw, fallback_mw, model_call_limit_mw, tool_call_limit_mw) - if m is not None - ] - if _gp_resilience: - _cache_idx = next( - ( - i - for i, m in enumerate(gp_middleware) - if isinstance(m, AnthropicPromptCachingMiddleware) - ), - len(gp_middleware), - ) - for offset, mw in enumerate(_gp_resilience): - gp_middleware.insert(_cache_idx + offset, mw) - - registry_subagents: list[SubAgent] = [] - try: - subagent_extra_middleware: list[Any] = [ - TodoListMiddleware(), - SurfSenseFilesystemMiddleware( - backend=backend_resolver, - filesystem_mode=filesystem_mode, - search_space_id=search_space_id, - created_by_id=user_id, - thread_id=thread_id, - ), - ] - if subagent_deny_permission_mw is not None: - subagent_extra_middleware.append(subagent_deny_permission_mw) - for _resilience_mw in ( - retry_mw, - fallback_mw, - model_call_limit_mw, - tool_call_limit_mw, - ): - if _resilience_mw is not None: - subagent_extra_middleware.append(_resilience_mw) - registry_subagents = build_subagents( - dependencies=subagent_dependencies, - model=llm, - extra_middleware=subagent_extra_middleware, - mcp_tools_by_agent=mcp_tools_by_agent or {}, - exclude=get_subagents_to_exclude(available_connectors), - disabled_tools=disabled_tools, - ) - logging.info( - "Registry subagents: %s", - [s["name"] for s in registry_subagents], - ) - except Exception: - logging.exception("Registry subagent build failed") - raise - - subagent_specs: list[SubAgent] = [general_purpose_spec, *registry_subagents] - - summarization_mw = create_surfsense_compaction_middleware(llm, StateBackend) - - context_edit_mw = None - if ( - flags.enable_context_editing - and not flags.disable_new_agent_stack - and max_input_tokens - ): - spill_edit = SpillToBackendEdit( - trigger=int(max_input_tokens * 0.55), - clear_at_least=int(max_input_tokens * 0.15), - keep=5, - exclude_tools=safe_exclude_tools(tools), - clear_tool_inputs=True, - ) - clear_edit = ClearToolUsesEdit( - trigger=int(max_input_tokens * 0.55), - clear_at_least=int(max_input_tokens * 0.15), - keep=5, - exclude_tools=safe_exclude_tools(tools), - clear_tool_inputs=True, - placeholder="[cleared - older tool output trimmed for context]", - ) - context_edit_mw = SpillingContextEditingMiddleware( - edits=[spill_edit, clear_edit], - backend_resolver=backend_resolver, - ) - - noop_mw = ( - NoopInjectionMiddleware() - if flags.enable_compaction_v2 and not flags.disable_new_agent_stack - else None - ) - - repair_mw = None - if flags.enable_tool_call_repair and not flags.disable_new_agent_stack: - registered_names: set[str] = {t.name for t in tools} - registered_names |= { - "write_todos", - "ls", - "read_file", - "write_file", - "edit_file", - "glob", - "grep", - "execute", - "task", - "mkdir", - "cd", - "pwd", - "move_file", - "rm", - "rmdir", - "list_tree", - "execute_code", - } - repair_mw = ToolCallNameRepairMiddleware( - registered_tool_names=registered_names, - fuzzy_match_threshold=None, - ) - - doom_loop_mw = ( - DoomLoopMiddleware(threshold=3) - if flags.enable_doom_loop and not flags.disable_new_agent_stack - else None - ) - - permission_mw: PermissionMiddleware | None = ( - PermissionMiddleware(rulesets=permission_rulesets) - if permission_rulesets - else None - ) - - action_log_mw: ActionLogMiddleware | None = None - if ( - flags.enable_action_log - and not flags.disable_new_agent_stack - and thread_id is not None - ): - try: - tool_defs_by_name = {td.name: td for td in BUILTIN_TOOLS} - action_log_mw = ActionLogMiddleware( - thread_id=thread_id, - search_space_id=search_space_id, - user_id=user_id, - tool_definitions=tool_defs_by_name, - ) - except Exception: # pragma: no cover - defensive - logging.warning( - "ActionLogMiddleware init failed; running without it.", - exc_info=True, - ) - action_log_mw = None - - busy_mutex_mw: BusyMutexMiddleware | None = ( - BusyMutexMiddleware() - if flags.enable_busy_mutex and not flags.disable_new_agent_stack - else None - ) - - otel_mw: OtelSpanMiddleware | None = ( - OtelSpanMiddleware() - if flags.enable_otel and not flags.disable_new_agent_stack - else None - ) - - plugin_middlewares: list[Any] = [] - if flags.enable_plugin_loader and not flags.disable_new_agent_stack: - try: - allowed_names = load_allowed_plugin_names_from_env() - if allowed_names: - plugin_middlewares = load_plugin_middlewares( - PluginContext.build( - search_space_id=search_space_id, - user_id=user_id, - thread_visibility=visibility, - llm=llm, - ), - allowed_plugin_names=allowed_names, - ) - except Exception: # pragma: no cover - defensive - logging.warning( - "Plugin loader failed; continuing without plugins.", - exc_info=True, - ) - plugin_middlewares = [] - - skills_mw: SkillsMiddleware | None = None - if flags.enable_skills and not flags.disable_new_agent_stack: - try: - skills_factory = build_skills_backend_factory( - search_space_id=search_space_id - if filesystem_mode == FilesystemMode.CLOUD - else None, - ) - skills_mw = SkillsMiddleware( - backend=skills_factory, - sources=default_skills_sources(), - ) - except Exception as exc: # pragma: no cover - defensive - logging.warning("SkillsMiddleware init failed; skipping: %s", exc) - skills_mw = None - - selector_mw: LLMToolSelectorMiddleware | None = None - if ( - flags.enable_llm_tool_selector - and not flags.disable_new_agent_stack - and len(tools) > 30 - ): - try: - selector_mw = LLMToolSelectorMiddleware( - model="openai:gpt-4o-mini", - max_tools=12, - always_include=[ - name - for name in ( - "update_memory", - "get_connected_accounts", - "scrape_webpage", - ) - if name in {t.name for t in tools} - ], - ) - except Exception: - logging.warning("LLMToolSelectorMiddleware init failed; skipping.") - selector_mw = None - - deepagent_middleware = [ - busy_mutex_mw, - otel_mw, - TodoListMiddleware(), - _memory_middleware, - AnonymousDocumentMiddleware( - anon_session_id=anon_session_id, - ) - if filesystem_mode == FilesystemMode.CLOUD - else None, - KnowledgeTreeMiddleware( - search_space_id=search_space_id, - filesystem_mode=filesystem_mode, - llm=llm, - ) - if filesystem_mode == FilesystemMode.CLOUD - else None, - KnowledgePriorityMiddleware( - llm=llm, - search_space_id=search_space_id, - filesystem_mode=filesystem_mode, - available_connectors=available_connectors, - available_document_types=available_document_types, - mentioned_document_ids=mentioned_document_ids, - ), - FileIntentMiddleware(llm=llm), - SurfSenseFilesystemMiddleware( - backend=backend_resolver, - filesystem_mode=filesystem_mode, - search_space_id=search_space_id, - created_by_id=user_id, - thread_id=thread_id, - ), - KnowledgeBasePersistenceMiddleware( - search_space_id=search_space_id, - created_by_id=user_id, - filesystem_mode=filesystem_mode, - thread_id=thread_id, - ) - if filesystem_mode == FilesystemMode.CLOUD - else None, - skills_mw, - SurfSenseCheckpointedSubAgentMiddleware( - checkpointer=checkpointer, - backend=StateBackend, - subagents=subagent_specs, - ), - selector_mw, - model_call_limit_mw, - tool_call_limit_mw, - context_edit_mw, - summarization_mw, - noop_mw, - retry_mw, - fallback_mw, - repair_mw, - permission_mw, - doom_loop_mw, - action_log_mw, - PatchToolCallsMiddleware(), - DedupHITLToolCallsMiddleware(agent_tools=list(tools)), - *plugin_middlewares, - AnthropicPromptCachingMiddleware(unsupported_model_behavior="ignore"), - ] - return [m for m in deepagent_middleware if m is not None] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/__init__.py similarity index 100% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/__init__.py rename to surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/__init__.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/config.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/config.py similarity index 100% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/config.py rename to surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/config.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/constants.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/constants.py similarity index 100% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/constants.py rename to surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/constants.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/middleware.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/middleware.py similarity index 100% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/middleware.py rename to surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/middleware.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/propagation.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/propagation.py similarity index 100% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/propagation.py rename to surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/propagation.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/resume.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/resume.py similarity index 100% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/resume.py rename to surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/resume.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/task_tool.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/task_tool.py similarity index 100% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/graph/middleware/checkpointed_subagent_middleware/task_tool.py rename to surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/task_tool.py diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py index ab6a644b5..dbc2c9c00 100644 --- a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_hitl_bridge.py @@ -12,7 +12,7 @@ from langgraph.graph import END, START, StateGraph from langgraph.types import Command, interrupt from typing_extensions import TypedDict -from app.agents.multi_agent_chat.main_agent.graph.middleware.checkpointed_subagent_middleware.task_tool import ( +from app.agents.multi_agent_chat.middleware.main_agent.checkpointed_subagent_middleware.task_tool import ( build_task_tool_with_parent_config, ) diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py index e73fb2823..347b32dbd 100644 --- a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_resume_helpers.py @@ -4,7 +4,7 @@ from __future__ import annotations from langchain.tools import ToolRuntime -from app.agents.multi_agent_chat.main_agent.graph.middleware.checkpointed_subagent_middleware.config import ( +from app.agents.multi_agent_chat.middleware.main_agent.checkpointed_subagent_middleware.config import ( consume_surfsense_resume, has_surfsense_resume, ) From bba5fb1db830aa4a95e2011d237726f45d69e9b4 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 22:22:57 +0200 Subject: [PATCH 22/96] fix(multi-agent): fail closed when connector discovery raises --- .../multi_agent_chat/main_agent/runtime/factory.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py index 6a6fd39b7..81123d450 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py @@ -85,7 +85,18 @@ async def create_surfsense_deep_agent( ) except Exception as e: - logging.warning("Failed to discover available connectors/document types: %s", e) + logging.warning( + "Connector/doc-type discovery failed; excluding connector subagents this turn: %s", + e, + ) + + # Fail closed: a None list short-circuits ``get_subagents_to_exclude`` to "exclude + # nothing", which would silently advertise every connector specialist on a flaky + # discovery call. Empty list excludes connector-gated subagents while keeping builtins. + if available_connectors is None: + available_connectors = [] + if available_document_types is None: + available_document_types = [] _perf_log.info( "[create_agent] Connector/doc-type discovery in %.3fs", time.perf_counter() - _t0, From 744ad92971d179af234b06fec0e57fde435e4cf2 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 22:23:17 +0200 Subject: [PATCH 23/96] chore(multi-agent): demote subagent registry listing log to debug --- .../app/agents/multi_agent_chat/middleware/stack.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py index 563332986..f894acc7e 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py @@ -121,7 +121,7 @@ def build_main_agent_deepagent_middleware( exclude=get_subagents_to_exclude(available_connectors), disabled_tools=disabled_tools, ) - logging.info( + logging.debug( "Subagents registry: %s", [s["name"] for s in subagents_registry], ) From c8ed70a26c5726339fe5b5451efdc6f520e1a1a8 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 22:30:08 +0200 Subject: [PATCH 24/96] fix(multi-agent): layer per-thread prompt cache key onto LLM at agent build --- .../app/agents/multi_agent_chat/main_agent/runtime/factory.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py index 81123d450..86da53a1a 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py @@ -26,6 +26,7 @@ from app.agents.new_chat.feature_flags import AgentFeatureFlags, get_flags from app.agents.new_chat.filesystem_backends import build_backend_resolver from app.agents.new_chat.filesystem_selection import FilesystemMode, FilesystemSelection from app.agents.new_chat.llm_config import AgentConfig +from app.agents.new_chat.prompt_caching import apply_litellm_prompt_caching from app.agents.new_chat.tools.invalid_tool import INVALID_TOOL_NAME, invalid_tool from app.agents.new_chat.tools.registry import build_tools_async from app.db import ChatVisibility @@ -62,6 +63,9 @@ async def create_surfsense_deep_agent( ): """Deep agent with SurfSense tools/middleware; registry route subagents behind ``task`` when enabled.""" _t_agent_total = time.perf_counter() + + apply_litellm_prompt_caching(llm, agent_config=agent_config, thread_id=thread_id) + filesystem_selection = filesystem_selection or FilesystemSelection() backend_resolver = build_backend_resolver( filesystem_selection, From 07a84d1a41051f3a3c6068b8deb3c5cf674b7de0 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 22:50:02 +0200 Subject: [PATCH 25/96] fix(multi-agent): cache compiled agent graph keyed on per-request inputs --- .../main_agent/runtime/agent_cache.py | 117 ++++++++++++++++++ .../main_agent/runtime/factory.py | 9 +- 2 files changed, 122 insertions(+), 4 deletions(-) create mode 100644 surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/agent_cache.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/agent_cache.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/agent_cache.py new file mode 100644 index 000000000..42f984b79 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/agent_cache.py @@ -0,0 +1,117 @@ +"""Compiled agent graph caching for the multi-agent path.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Sequence +from typing import Any + +from langchain_core.language_models import BaseChatModel +from langchain_core.tools import BaseTool +from langgraph.types import Checkpointer + +from app.agents.multi_agent_chat.subagents.shared.permissions import ToolsPermissions +from app.agents.new_chat.agent_cache import ( + flags_signature, + get_cache, + stable_hash, + system_prompt_hash, + tools_signature, +) +from app.agents.new_chat.feature_flags import AgentFeatureFlags +from app.agents.new_chat.filesystem_selection import FilesystemMode +from app.db import ChatVisibility + +from ..graph.compile_graph_sync import build_compiled_agent_graph_sync + + +def mcp_signature(mcp_tools_by_agent: dict[str, ToolsPermissions]) -> str: + """Hash the per-agent MCP tool surface so a change rotates the cache key.""" + rows = [] + for agent_name in sorted(mcp_tools_by_agent.keys()): + perms = mcp_tools_by_agent[agent_name] + allow_names = sorted(item.get("name", "") for item in perms.get("allow", [])) + ask_names = sorted(item.get("name", "") for item in perms.get("ask", [])) + rows.append((agent_name, allow_names, ask_names)) + return stable_hash(rows) + + +async def build_agent_with_cache( + *, + llm: BaseChatModel, + tools: Sequence[BaseTool], + final_system_prompt: str, + backend_resolver: Any, + filesystem_mode: FilesystemMode, + search_space_id: int, + user_id: str | None, + thread_id: int | None, + visibility: ChatVisibility, + anon_session_id: str | None, + available_connectors: list[str], + available_document_types: list[str], + mentioned_document_ids: list[int] | None, + max_input_tokens: int | None, + flags: AgentFeatureFlags, + checkpointer: Checkpointer, + subagent_dependencies: dict[str, Any], + mcp_tools_by_agent: dict[str, ToolsPermissions], + disabled_tools: list[str] | None, + config_id: str | None, +) -> Any: + """Compile the multi-agent graph, serving from cache when key components are stable.""" + + async def _build() -> Any: + return await asyncio.to_thread( + build_compiled_agent_graph_sync, + llm=llm, + tools=tools, + final_system_prompt=final_system_prompt, + backend_resolver=backend_resolver, + filesystem_mode=filesystem_mode, + search_space_id=search_space_id, + user_id=user_id, + thread_id=thread_id, + visibility=visibility, + anon_session_id=anon_session_id, + available_connectors=available_connectors, + available_document_types=available_document_types, + mentioned_document_ids=mentioned_document_ids, + max_input_tokens=max_input_tokens, + flags=flags, + checkpointer=checkpointer, + subagent_dependencies=subagent_dependencies, + mcp_tools_by_agent=mcp_tools_by_agent, + disabled_tools=disabled_tools, + ) + + if not (flags.enable_agent_cache and not flags.disable_new_agent_stack): + return await _build() + + # Every per-request value any middleware closes over at __init__ must be in + # the key, otherwise a hit will leak state across threads. Bump the schema + # version when the component list changes shape. + cache_key = stable_hash( + "multi-agent-v1", + config_id, + thread_id, + user_id, + search_space_id, + visibility, + filesystem_mode, + anon_session_id, + tools_signature( + tools, + available_connectors=available_connectors, + available_document_types=available_document_types, + ), + mcp_signature(mcp_tools_by_agent), + flags_signature(flags), + system_prompt_hash(final_system_prompt), + max_input_tokens, + sorted(disabled_tools) if disabled_tools else None, + ) + return await get_cache().get_or_build(cache_key, builder=_build) + + +__all__ = ["build_agent_with_cache", "mcp_signature"] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py index 86da53a1a..9e3c8eab4 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import logging import time from collections.abc import Sequence @@ -33,12 +32,12 @@ from app.db import ChatVisibility from app.services.connector_service import ConnectorService from app.utils.perf import get_perf_logger -from ..graph.compile_graph_sync import build_compiled_agent_graph_sync from ..system_prompt import build_main_agent_system_prompt from ..tools import ( MAIN_AGENT_SURFSENSE_TOOL_NAMES, MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED, ) +from .agent_cache import build_agent_with_cache _perf_log = get_perf_logger() @@ -210,9 +209,10 @@ async def create_surfsense_deep_agent( final_system_prompt = system_prompt + "\n\n" + BASE_AGENT_PROMPT + config_id = agent_config.config_id if agent_config is not None else None + _t0 = time.perf_counter() - agent = await asyncio.to_thread( - build_compiled_agent_graph_sync, + agent = await build_agent_with_cache( llm=llm, tools=tools, final_system_prompt=final_system_prompt, @@ -232,6 +232,7 @@ async def create_surfsense_deep_agent( subagent_dependencies=dependencies, mcp_tools_by_agent=mcp_tools_by_agent, disabled_tools=disabled_tools, + config_id=config_id, ) _perf_log.info( "[create_agent] Middleware stack + graph compiled in %.3fs", From 997d86079046ed6f02bedb4c4a16adba5d1c4ed9 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 22:52:03 +0200 Subject: [PATCH 26/96] fix(multi-agent): defensive message extraction on subagent task return --- .../checkpointed_subagent_middleware/task_tool.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/task_tool.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/task_tool.py index d23dc33a9..5668f8ddb 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/task_tool.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/checkpointed_subagent_middleware/task_tool.py @@ -69,9 +69,16 @@ def build_task_tool_with_parent_config( raise ValueError(msg) state_update = {k: v for k, v in result.items() if k not in EXCLUDED_STATE_KEYS} - message_text = ( - result["messages"][-1].text.rstrip() if result["messages"][-1].text else "" - ) + messages = result["messages"] + if not messages: + msg = ( + "CompiledSubAgent returned an empty 'messages' list. " + "Subagents must produce at least one message so the parent has " + "output to forward back to the user." + ) + raise ValueError(msg) + last_text = getattr(messages[-1], "text", None) or "" + message_text = last_text.rstrip() return Command( update={ **state_update, From a3c3db2a18d270f73dbc19e6b502a6d0b01c2431 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 22:57:35 +0200 Subject: [PATCH 27/96] test(multi-agent): pin first-wins assumption on pending subagent interrupts --- .../test_pending_interrupt.py | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_pending_interrupt.py diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_pending_interrupt.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_pending_interrupt.py new file mode 100644 index 000000000..75242689d --- /dev/null +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/middleware/checkpointed_subagent_middleware/test_pending_interrupt.py @@ -0,0 +1,55 @@ +"""Pins the first-wins assumption of ``get_first_pending_subagent_interrupt``. + +The bridge currently relies on at-most-one pending interrupt per snapshot +(sequential tool nodes). If parallel tool calls are ever enabled, the bridge +needs an id-aware lookup; these tests will need to be revisited at that point. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +from app.agents.multi_agent_chat.middleware.main_agent.checkpointed_subagent_middleware.resume import ( + get_first_pending_subagent_interrupt, +) + + +class TestGetFirstPendingSubagentInterrupt: + def test_returns_first_when_multiple_top_level_interrupts_pending(self): + first = SimpleNamespace(id="i-1", value={"decision": "approve"}) + second = SimpleNamespace(id="i-2", value={"decision": "reject"}) + state = SimpleNamespace(interrupts=(first, second), tasks=()) + + assert get_first_pending_subagent_interrupt(state) == ( + "i-1", + {"decision": "approve"}, + ) + + def test_returns_first_when_multiple_subtask_interrupts_pending(self): + first = SimpleNamespace(id="i-A", value="approve") + second = SimpleNamespace(id="i-B", value="reject") + sub_task = SimpleNamespace(interrupts=(first, second)) + state = SimpleNamespace(interrupts=(), tasks=(sub_task,)) + + assert get_first_pending_subagent_interrupt(state) == ("i-A", "approve") + + def test_returns_none_when_no_interrupts(self): + state = SimpleNamespace(interrupts=(), tasks=()) + + assert get_first_pending_subagent_interrupt(state) == (None, None) + + def test_returns_none_when_state_is_none(self): + assert get_first_pending_subagent_interrupt(None) == (None, None) + + def test_skips_interrupts_with_none_value(self): + empty = SimpleNamespace(id="i-empty", value=None) + real = SimpleNamespace(id="i-real", value="approve") + state = SimpleNamespace(interrupts=(empty, real), tasks=()) + + assert get_first_pending_subagent_interrupt(state) == ("i-real", "approve") + + def test_normalizes_non_string_id_to_none(self): + interrupt = SimpleNamespace(id=12345, value="approve") + state = SimpleNamespace(interrupts=(interrupt,), tasks=()) + + assert get_first_pending_subagent_interrupt(state) == (None, "approve") From 657c31fdf47ecbd141ea1d64ab71b7e3e0727388 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 23:01:24 +0200 Subject: [PATCH 28/96] refactor(stream): rename multi-agent factory alias for clarity --- surfsense_backend/app/tasks/chat/stream_new_chat.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 3ba3912eb..2ee3f075a 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -29,7 +29,7 @@ from sqlalchemy.future import select from sqlalchemy.orm import selectinload from app.agents.multi_agent_chat import ( - create_surfsense_deep_agent as create_registry_deep_agent, + create_surfsense_deep_agent as create_multi_agent_chat, ) from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent from app.agents.new_chat.checkpointer import get_checkpointer @@ -2767,7 +2767,7 @@ async def stream_new_chat( _t0 = time.perf_counter() agent_factory = ( - create_registry_deep_agent + create_multi_agent_chat if use_multi_agent else create_surfsense_deep_agent ) @@ -4130,7 +4130,7 @@ async def stream_resume_chat( _t0 = time.perf_counter() agent_factory = ( - create_registry_deep_agent + create_multi_agent_chat if _app_config.MULTI_AGENT_CHAT_ENABLED else create_surfsense_deep_agent ) From 3cb2c3056ea2df059122250f7084ff7258930603 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 23:35:23 +0200 Subject: [PATCH 29/96] fix(stream): route every agent (re)build through one helper to prevent factory drift --- .../app/agents/multi_agent_chat/__init__.py | 4 +- .../multi_agent_chat/main_agent/__init__.py | 4 +- .../main_agent/runtime/__init__.py | 4 +- .../main_agent/runtime/factory.py | 2 +- .../app/tasks/chat/stream_new_chat.py | 65 +++++++++++++++---- 5 files changed, 60 insertions(+), 19 deletions(-) diff --git a/surfsense_backend/app/agents/multi_agent_chat/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/__init__.py index f568dc6b2..6c7d79eb8 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/__init__.py +++ b/surfsense_backend/app/agents/multi_agent_chat/__init__.py @@ -2,6 +2,6 @@ from __future__ import annotations -from .main_agent import create_surfsense_deep_agent +from .main_agent import create_multi_agent_chat_deep_agent -__all__ = ["create_surfsense_deep_agent"] +__all__ = ["create_multi_agent_chat_deep_agent"] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/__init__.py index b9a18fe53..f74ca0cd0 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/__init__.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/__init__.py @@ -2,6 +2,6 @@ from __future__ import annotations -from .runtime import create_surfsense_deep_agent +from .runtime import create_multi_agent_chat_deep_agent -__all__ = ["create_surfsense_deep_agent"] +__all__ = ["create_multi_agent_chat_deep_agent"] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/__init__.py index 3d4ae977d..593e8da20 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/__init__.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/__init__.py @@ -2,6 +2,6 @@ from __future__ import annotations -from .factory import create_surfsense_deep_agent +from .factory import create_multi_agent_chat_deep_agent -__all__ = ["create_surfsense_deep_agent"] +__all__ = ["create_multi_agent_chat_deep_agent"] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py index 9e3c8eab4..630455694 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py @@ -42,7 +42,7 @@ from .agent_cache import build_agent_with_cache _perf_log = get_perf_logger() -async def create_surfsense_deep_agent( +async def create_multi_agent_chat_deep_agent( llm: BaseChatModel, search_space_id: int, db_session: AsyncSession, diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 2ee3f075a..1a2f38077 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -28,9 +28,7 @@ from langchain_core.messages import HumanMessage from sqlalchemy.future import select from sqlalchemy.orm import selectinload -from app.agents.multi_agent_chat import ( - create_surfsense_deep_agent as create_multi_agent_chat, -) +from app.agents.multi_agent_chat import create_multi_agent_chat_deep_agent from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent from app.agents.new_chat.checkpointer import get_checkpointer from app.agents.new_chat.context import SurfSenseContextSchema @@ -577,6 +575,43 @@ async def _preflight_llm(llm: Any) -> None: ) +async def _build_main_agent_for_thread( + agent_factory: Any, + *, + llm: Any, + search_space_id: int, + db_session: Any, + connector_service: ConnectorService, + checkpointer: Any, + user_id: str | None, + thread_id: int | None, + agent_config: AgentConfig | None, + firecrawl_api_key: str | None, + thread_visibility: ChatVisibility | None, + filesystem_selection: FilesystemSelection | None, + disabled_tools: list[str] | None = None, + mentioned_document_ids: list[int] | None = None, +) -> Any: + """Single (re)build path so the agent factory cannot drift across + initial build, preflight repin, and mid-stream 429 recovery for one + ``thread_id``: a graph swap mid-turn would corrupt checkpointer state.""" + return await agent_factory( + llm=llm, + search_space_id=search_space_id, + db_session=db_session, + connector_service=connector_service, + checkpointer=checkpointer, + user_id=user_id, + thread_id=thread_id, + agent_config=agent_config, + firecrawl_api_key=firecrawl_api_key, + thread_visibility=thread_visibility, + filesystem_selection=filesystem_selection, + disabled_tools=disabled_tools, + mentioned_document_ids=mentioned_document_ids, + ) + + async def _settle_speculative_agent_build(task: asyncio.Task[Any]) -> None: """Wait for a discarded speculative agent build to release shared state. @@ -2767,7 +2802,7 @@ async def stream_new_chat( _t0 = time.perf_counter() agent_factory = ( - create_multi_agent_chat + create_multi_agent_chat_deep_agent if use_multi_agent else create_surfsense_deep_agent ) @@ -2776,7 +2811,8 @@ async def stream_new_chat( # if preflight reports 429 we will discard this future and rebuild # against the freshly pinned config below. agent_build_task = asyncio.create_task( - agent_factory( + _build_main_agent_for_thread( + agent_factory, llm=llm, search_space_id=search_space_id, db_session=session, @@ -2787,9 +2823,9 @@ async def stream_new_chat( agent_config=agent_config, firecrawl_api_key=firecrawl_api_key, thread_visibility=visibility, + filesystem_selection=filesystem_selection, disabled_tools=disabled_tools, mentioned_document_ids=mentioned_document_ids, - filesystem_selection=filesystem_selection, ), name="agent_build:stream_new_chat", ) @@ -3466,7 +3502,8 @@ async def stream_new_chat( title_task = None _t0 = time.perf_counter() - agent = await create_surfsense_deep_agent( + agent = await _build_main_agent_for_thread( + agent_factory, llm=llm, search_space_id=search_space_id, db_session=session, @@ -3477,9 +3514,9 @@ async def stream_new_chat( agent_config=agent_config, firecrawl_api_key=firecrawl_api_key, thread_visibility=visibility, + filesystem_selection=filesystem_selection, disabled_tools=disabled_tools, mentioned_document_ids=mentioned_document_ids, - filesystem_selection=filesystem_selection, ) _perf_log.info( "[stream_new_chat] Runtime rate-limit recovery repinned " @@ -4130,12 +4167,13 @@ async def stream_resume_chat( _t0 = time.perf_counter() agent_factory = ( - create_multi_agent_chat + create_multi_agent_chat_deep_agent if _app_config.MULTI_AGENT_CHAT_ENABLED else create_surfsense_deep_agent ) agent_build_task = asyncio.create_task( - agent_factory( + _build_main_agent_for_thread( + agent_factory, llm=llm, search_space_id=search_space_id, db_session=session, @@ -4224,7 +4262,8 @@ async def stream_resume_chat( "fallback_config_id": llm_config_id, }, ) - agent = await agent_factory( + agent = await _build_main_agent_for_thread( + agent_factory, llm=llm, search_space_id=search_space_id, db_session=session, @@ -4409,7 +4448,8 @@ async def stream_resume_chat( raise stream_exc _t0 = time.perf_counter() - agent = await create_surfsense_deep_agent( + agent = await _build_main_agent_for_thread( + agent_factory, llm=llm, search_space_id=search_space_id, db_session=session, @@ -4421,6 +4461,7 @@ async def stream_resume_chat( firecrawl_api_key=firecrawl_api_key, thread_visibility=visibility, filesystem_selection=filesystem_selection, + disabled_tools=disabled_tools, ) _perf_log.info( "[stream_resume] Runtime rate-limit recovery repinned " From a421e7d792ca7dcdd77c6c2ebdc4944042bb0646 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Tue, 5 May 2026 23:42:11 +0200 Subject: [PATCH 30/96] fix(multi-agent): degrade to builtins-only when MCP or subagent registry build fails --- .../multi_agent_chat/main_agent/runtime/factory.py | 11 ++++++++++- .../app/agents/multi_agent_chat/middleware/stack.py | 8 ++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py index 630455694..d0354aca3 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py @@ -129,7 +129,16 @@ async def create_multi_agent_chat_deep_agent( } _t0 = time.perf_counter() - mcp_tools_by_agent = await load_mcp_tools_by_connector(db_session, search_space_id) + try: + mcp_tools_by_agent = await load_mcp_tools_by_connector(db_session, search_space_id) + except Exception as e: + # Degrade to builtins-only rather than aborting the turn: a transient + # DB or MCP-server hiccup should not deny the user a response. + logging.warning( + "MCP tool discovery failed; subagents will run without MCP tools this turn: %s", + e, + ) + mcp_tools_by_agent = {} _perf_log.info( "[create_agent] load_mcp_tools_by_connector in %.3fs (%d buckets)", time.perf_counter() - _t0, diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py index f894acc7e..6d8faa3f4 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/stack.py @@ -126,8 +126,12 @@ def build_main_agent_deepagent_middleware( [s["name"] for s in subagents_registry], ) except Exception: - logging.exception("Subagents registry build failed") - raise + # Degrade to general-purpose-only rather than aborting the turn: + # one bad subagent dep should not deny the user a response. + logging.exception( + "Subagents registry build failed; falling back to general-purpose only" + ) + subagents_registry = [] subagents: list[SubAgent] = [general_purpose_subagent, *subagents_registry] From 4e174f17f2b4dfd2446408cadd1150978594625b Mon Sep 17 00:00:00 2001 From: "DESKTOP-RTLN3BA\\$punk" Date: Tue, 5 May 2026 17:08:34 -0700 Subject: [PATCH 31/96] chore: linting --- .../agents/multi_agent_chat/main_agent/runtime/factory.py | 4 +++- .../multi_agent_chat/middleware/main_agent/doom_loop.py | 4 +++- .../middleware/shared/permissions/context.py | 4 +--- .../subagents/shared/test_subagent_builder.py | 1 - .../new_chat/middleware/test_scoped_model_fallback.py | 4 +--- .../unit/agents/new_chat/test_memory_response_content.py | 4 +++- .../unit/agents/new_chat/test_permission_middleware.py | 4 +--- .../dashboard/[search_space_id]/purchase-success/page.tsx | 6 +++--- 8 files changed, 15 insertions(+), 16 deletions(-) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py index d0354aca3..cb6410acb 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/runtime/factory.py @@ -130,7 +130,9 @@ async def create_multi_agent_chat_deep_agent( _t0 = time.perf_counter() try: - mcp_tools_by_agent = await load_mcp_tools_by_connector(db_session, search_space_id) + mcp_tools_by_agent = await load_mcp_tools_by_connector( + db_session, search_space_id + ) except Exception as e: # Degrade to builtins-only rather than aborting the turn: a transient # DB or MCP-server hiccup should not deny the user a response. diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py index a0b294092..d67b8d518 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/main_agent/doom_loop.py @@ -9,4 +9,6 @@ from ..shared.flags import enabled def build_doom_loop_mw(flags: AgentFeatureFlags) -> DoomLoopMiddleware | None: - return DoomLoopMiddleware(threshold=3) if enabled(flags, "enable_doom_loop") else None + return ( + DoomLoopMiddleware(threshold=3) if enabled(flags, "enable_doom_loop") else None + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/context.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/context.py index f14d52714..e121421a0 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/context.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/shared/permissions/context.py @@ -78,9 +78,7 @@ def build_permission_context( Rule(permission=tool_def.name, pattern="*", action="deny") ) if synthesized: - rulesets.append( - Ruleset(rules=synthesized, origin="connector_synthesized") - ) + rulesets.append(Ruleset(rules=synthesized, origin="connector_synthesized")) general_purpose_interrupt_on: dict[str, bool] = { rule.permission: True diff --git a/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py b/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py index 5cd62ed36..648e52115 100644 --- a/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py +++ b/surfsense_backend/tests/unit/agents/multi_agent_chat/subagents/shared/test_subagent_builder.py @@ -29,7 +29,6 @@ class RateLimitError(Exception): class _AlwaysFailingChatModel(BaseChatModel): - @property def _llm_type(self) -> str: return "always-failing-test-model" diff --git a/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py b/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py index 69f6fe6b7..80b9862e7 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py +++ b/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py @@ -67,9 +67,7 @@ class _RecordingChatModel(BaseChatModel): ) -> ChatResult: self.call_count += 1 return ChatResult( - generations=[ - ChatGeneration(message=AIMessage(content=self.response_text)) - ] + generations=[ChatGeneration(message=AIMessage(content=self.response_text))] ) async def _agenerate( diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_memory_response_content.py b/surfsense_backend/tests/unit/agents/new_chat/test_memory_response_content.py index 535e4e940..1f338ee3e 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_memory_response_content.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_memory_response_content.py @@ -47,7 +47,9 @@ def test_extract_text_content_ignores_thinking_blocks_and_keeps_markdown_text() assert extract_text_content(content).strip() == markdown.strip() -def test_extract_text_content_returns_empty_when_only_thinking_blocks_are_present() -> None: +def test_extract_text_content_returns_empty_when_only_thinking_blocks_are_present() -> ( + None +): content = [ {"type": "thinking", "thinking": "No durable fact."}, {"type": "thinking", "thinking": "Return no update."}, diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_permission_middleware.py b/surfsense_backend/tests/unit/agents/new_chat/test_permission_middleware.py index eda5be150..47059ade6 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_permission_middleware.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_permission_middleware.py @@ -141,9 +141,7 @@ class TestNormalizeDecision: assert _normalize_permission_decision(decision) == {"decision_type": "reject"} def test_lc_envelope_reject_with_message_carries_feedback(self) -> None: - decision = { - "decisions": [{"type": "reject", "message": "wrong recipient"}] - } + decision = {"decisions": [{"type": "reject", "message": "wrong recipient"}]} out = _normalize_permission_decision(decision) assert out == {"decision_type": "reject", "feedback": "wrong recipient"} diff --git a/surfsense_web/app/dashboard/[search_space_id]/purchase-success/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/purchase-success/page.tsx index b3d504ed5..8eaec3e5a 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/purchase-success/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/purchase-success/page.tsx @@ -119,8 +119,7 @@ export default function PurchaseSuccessPage() { "Stripe reported the checkout as failed or expired. Your card was not charged."} {state.kind === "error" && "Don't worry — if your card was charged, your purchase will still apply within a minute or two."} - {state.kind === "no_session" && - "Your purchase is being applied to your account."} + {state.kind === "no_session" && "Your purchase is being applied to your account."} @@ -134,7 +133,8 @@ export default function PurchaseSuccessPage() { )} {state.kind === "completed" && state.data.purchase_type === "premium_tokens" && (

- New premium credit balance: {formatCredit(state.data.premium_credit_micros_limit ?? 0)} + New premium credit balance:{" "} + {formatCredit(state.data.premium_credit_micros_limit ?? 0)}

)} {state.kind === "error" && ( From 5e87a7a251c5a614ab94f28c28bcea1457b67e47 Mon Sep 17 00:00:00 2001 From: "DESKTOP-RTLN3BA\\$punk" Date: Tue, 5 May 2026 18:57:10 -0700 Subject: [PATCH 32/96] fix: composio tool calls in composio connectors --- .../connectors/calendar/tools/create_event.py | 149 +++--- .../connectors/calendar/tools/delete_event.py | 94 ++-- .../calendar/tools/search_events.py | 59 ++- .../connectors/calendar/tools/update_event.py | 188 ++++--- .../connectors/gmail/tools/create_draft.py | 133 +++-- .../connectors/gmail/tools/read_email.py | 51 +- .../connectors/gmail/tools/search_emails.py | 96 ++-- .../connectors/gmail/tools/send_email.py | 121 +++-- .../connectors/gmail/tools/trash_email.py | 90 ++-- .../connectors/gmail/tools/update_draft.py | 195 ++++--- .../google_drive/tools/create_file.py | 113 ++-- .../google_drive/tools/trash_file.py | 109 ++-- .../app/services/composio_service.py | 499 ++++++++++++++++++ 13 files changed, 1347 insertions(+), 550 deletions(-) diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/create_event.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/create_event.py index 37bcf083e..a8183314a 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/create_event.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/create_event.py @@ -168,20 +168,46 @@ def create_create_calendar_event_tool( f"Creating calendar event: summary='{final_summary}', connector={actual_connector_id}" ) + tz = context.get("timezone", "UTC") + if ( connector.connector_type == SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR ): - from app.utils.google_credentials import build_composio_credentials - cca_id = connector.config.get("composio_connected_account_id") - if cca_id: - creds = build_composio_credentials(cca_id) - else: + if not cca_id: return { "status": "error", "message": "Composio connected account ID not found for this connector.", } + + from app.services.composio_service import ComposioService + + ( + event_id, + html_link, + error, + ) = await ComposioService().create_calendar_event( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + summary=final_summary, + start_datetime=final_start_datetime, + end_datetime=final_end_datetime, + timezone=tz, + description=final_description, + location=final_location, + attendees=final_attendees, + ) + if error: + return {"status": "error", "message": error} + created = { + "id": event_id, + "summary": final_summary, + "htmlLink": html_link, + } + logger.info( + f"Calendar event created via Composio: id={event_id}, summary={final_summary}" + ) else: config_data = dict(connector.config) @@ -211,70 +237,69 @@ def create_create_calendar_event_tool( expiry=datetime.fromisoformat(exp) if exp else None, ) - service = await asyncio.get_event_loop().run_in_executor( - None, lambda: build("calendar", "v3", credentials=creds) - ) - - tz = context.get("timezone", "UTC") - event_body: dict[str, Any] = { - "summary": final_summary, - "start": {"dateTime": final_start_datetime, "timeZone": tz}, - "end": {"dateTime": final_end_datetime, "timeZone": tz}, - } - if final_description: - event_body["description"] = final_description - if final_location: - event_body["location"] = final_location - if final_attendees: - event_body["attendees"] = [ - {"email": e.strip()} for e in final_attendees if e.strip() - ] - - try: - created = await asyncio.get_event_loop().run_in_executor( - None, - lambda: ( - service.events() - .insert(calendarId="primary", body=event_body) - .execute() - ), + service = await asyncio.get_event_loop().run_in_executor( + None, lambda: build("calendar", "v3", credentials=creds) ) - except Exception as api_err: - from googleapiclient.errors import HttpError - if isinstance(api_err, HttpError) and api_err.resp.status == 403: - logger.warning( - f"Insufficient permissions for connector {actual_connector_id}: {api_err}" + event_body: dict[str, Any] = { + "summary": final_summary, + "start": {"dateTime": final_start_datetime, "timeZone": tz}, + "end": {"dateTime": final_end_datetime, "timeZone": tz}, + } + if final_description: + event_body["description"] = final_description + if final_location: + event_body["location"] = final_location + if final_attendees: + event_body["attendees"] = [ + {"email": e.strip()} for e in final_attendees if e.strip() + ] + + try: + created = await asyncio.get_event_loop().run_in_executor( + None, + lambda: ( + service.events() + .insert(calendarId="primary", body=event_body) + .execute() + ), ) - try: - from sqlalchemy.orm.attributes import flag_modified + except Exception as api_err: + from googleapiclient.errors import HttpError - _res = await db_session.execute( - select(SearchSourceConnector).where( - SearchSourceConnector.id == actual_connector_id - ) - ) - _conn = _res.scalar_one_or_none() - if _conn and not _conn.config.get("auth_expired"): - _conn.config = {**_conn.config, "auth_expired": True} - flag_modified(_conn, "config") - await db_session.commit() - except Exception: + if isinstance(api_err, HttpError) and api_err.resp.status == 403: logger.warning( - "Failed to persist auth_expired for connector %s", - actual_connector_id, - exc_info=True, + f"Insufficient permissions for connector {actual_connector_id}: {api_err}" ) - return { - "status": "insufficient_permissions", - "connector_id": actual_connector_id, - "message": "This Google Calendar account needs additional permissions. Please re-authenticate in connector settings.", - } - raise + try: + from sqlalchemy.orm.attributes import flag_modified - logger.info( - f"Calendar event created: id={created.get('id')}, summary={created.get('summary')}" - ) + _res = await db_session.execute( + select(SearchSourceConnector).where( + SearchSourceConnector.id == actual_connector_id + ) + ) + _conn = _res.scalar_one_or_none() + if _conn and not _conn.config.get("auth_expired"): + _conn.config = {**_conn.config, "auth_expired": True} + flag_modified(_conn, "config") + await db_session.commit() + except Exception: + logger.warning( + "Failed to persist auth_expired for connector %s", + actual_connector_id, + exc_info=True, + ) + return { + "status": "insufficient_permissions", + "connector_id": actual_connector_id, + "message": "This Google Calendar account needs additional permissions. Please re-authenticate in connector settings.", + } + raise + + logger.info( + f"Calendar event created via Google API: id={created.get('id')}, summary={created.get('summary')}" + ) kb_message_suffix = "" try: diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/delete_event.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/delete_event.py index 4d9d69b4b..3d160e669 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/delete_event.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/delete_event.py @@ -163,16 +163,22 @@ def create_delete_calendar_event_tool( connector.connector_type == SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR ): - from app.utils.google_credentials import build_composio_credentials - cca_id = connector.config.get("composio_connected_account_id") - if cca_id: - creds = build_composio_credentials(cca_id) - else: + if not cca_id: return { "status": "error", "message": "Composio connected account ID not found for this connector.", } + + from app.services.composio_service import ComposioService + + error = await ComposioService().delete_calendar_event( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + event_id=final_event_id, + ) + if error: + return {"status": "error", "message": error} else: config_data = dict(connector.config) @@ -202,51 +208,51 @@ def create_delete_calendar_event_tool( expiry=datetime.fromisoformat(exp) if exp else None, ) - service = await asyncio.get_event_loop().run_in_executor( - None, lambda: build("calendar", "v3", credentials=creds) - ) - - try: - await asyncio.get_event_loop().run_in_executor( - None, - lambda: ( - service.events() - .delete(calendarId="primary", eventId=final_event_id) - .execute() - ), + service = await asyncio.get_event_loop().run_in_executor( + None, lambda: build("calendar", "v3", credentials=creds) ) - except Exception as api_err: - from googleapiclient.errors import HttpError - if isinstance(api_err, HttpError) and api_err.resp.status == 403: - logger.warning( - f"Insufficient permissions for connector {actual_connector_id}: {api_err}" + try: + await asyncio.get_event_loop().run_in_executor( + None, + lambda: ( + service.events() + .delete(calendarId="primary", eventId=final_event_id) + .execute() + ), ) - try: - from sqlalchemy.orm.attributes import flag_modified + except Exception as api_err: + from googleapiclient.errors import HttpError - _res = await db_session.execute( - select(SearchSourceConnector).where( - SearchSourceConnector.id == actual_connector_id - ) - ) - _conn = _res.scalar_one_or_none() - if _conn and not _conn.config.get("auth_expired"): - _conn.config = {**_conn.config, "auth_expired": True} - flag_modified(_conn, "config") - await db_session.commit() - except Exception: + if isinstance(api_err, HttpError) and api_err.resp.status == 403: logger.warning( - "Failed to persist auth_expired for connector %s", - actual_connector_id, - exc_info=True, + f"Insufficient permissions for connector {actual_connector_id}: {api_err}" ) - return { - "status": "insufficient_permissions", - "connector_id": actual_connector_id, - "message": "This Google Calendar account needs additional permissions. Please re-authenticate in connector settings.", - } - raise + try: + from sqlalchemy.orm.attributes import flag_modified + + _res = await db_session.execute( + select(SearchSourceConnector).where( + SearchSourceConnector.id == actual_connector_id + ) + ) + _conn = _res.scalar_one_or_none() + if _conn and not _conn.config.get("auth_expired"): + _conn.config = {**_conn.config, "auth_expired": True} + flag_modified(_conn, "config") + await db_session.commit() + except Exception: + logger.warning( + "Failed to persist auth_expired for connector %s", + actual_connector_id, + exc_info=True, + ) + return { + "status": "insufficient_permissions", + "connector_id": actual_connector_id, + "message": "This Google Calendar account needs additional permissions. Please re-authenticate in connector settings.", + } + raise logger.info(f"Calendar event deleted: event_id={final_event_id}") diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/search_events.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/search_events.py index dc6adb822..6772d5a1e 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/search_events.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/search_events.py @@ -16,6 +16,14 @@ _CALENDAR_TYPES = [ ] +def _to_calendar_boundary(value: str, *, is_end: bool) -> str: + """Promote a bare YYYY-MM-DD to RFC3339 with a day-edge time, leave full datetimes alone.""" + if "T" in value: + return value + time = "23:59:59" if is_end else "00:00:00" + return f"{value}T{time}Z" + + def create_search_calendar_events_tool( db_session: AsyncSession | None = None, search_space_id: int | None = None, @@ -61,22 +69,47 @@ def create_search_calendar_events_tool( "message": "No Google Calendar connector found. Please connect Google Calendar in your workspace settings.", } - creds = _build_credentials(connector) + if ( + connector.connector_type + == SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR + ): + cca_id = connector.config.get("composio_connected_account_id") + if not cca_id: + return { + "status": "error", + "message": "Composio connected account ID not found for this connector.", + } - from app.connectors.google_calendar_connector import GoogleCalendarConnector + from app.services.composio_service import ComposioService - cal = GoogleCalendarConnector( - credentials=creds, - session=db_session, - user_id=user_id, - connector_id=connector.id, - ) + events_raw, error = await ComposioService().get_calendar_events( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + time_min=_to_calendar_boundary(start_date, is_end=False), + time_max=_to_calendar_boundary(end_date, is_end=True), + max_results=max_results, + ) + if not events_raw and not error: + error = "No events found in the specified date range." + else: + creds = _build_credentials(connector) - events_raw, error = await cal.get_all_primary_calendar_events( - start_date=start_date, - end_date=end_date, - max_results=max_results, - ) + from app.connectors.google_calendar_connector import ( + GoogleCalendarConnector, + ) + + cal = GoogleCalendarConnector( + credentials=creds, + session=db_session, + user_id=user_id, + connector_id=connector.id, + ) + + events_raw, error = await cal.get_all_primary_calendar_events( + start_date=start_date, + end_date=end_date, + max_results=max_results, + ) if error: if ( diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/update_event.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/update_event.py index 259f52bba..a74979484 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/update_event.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/calendar/tools/update_event.py @@ -192,20 +192,62 @@ def create_update_calendar_event_tool( f"Updating calendar event: event_id='{final_event_id}', connector={actual_connector_id}" ) + has_changes = any( + v is not None + for v in ( + final_new_summary, + final_new_start_datetime, + final_new_end_datetime, + final_new_description, + final_new_location, + final_new_attendees, + ) + ) + if not has_changes: + return { + "status": "error", + "message": "No changes specified. Please provide at least one field to update.", + } + if ( connector.connector_type == SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR ): - from app.utils.google_credentials import build_composio_credentials - cca_id = connector.config.get("composio_connected_account_id") - if cca_id: - creds = build_composio_credentials(cca_id) - else: + if not cca_id: return { "status": "error", "message": "Composio connected account ID not found for this connector.", } + + from app.services.composio_service import ComposioService + + tz_for_composio: str | None = None + if final_new_start_datetime is not None and not _is_date_only( + final_new_start_datetime + ): + tz_for_composio = ( + context.get("timezone") if isinstance(context, dict) else None + ) + + _, html_link, error = await ComposioService().update_calendar_event( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + event_id=final_event_id, + summary=final_new_summary, + start_time=final_new_start_datetime, + end_time=final_new_end_datetime, + timezone=tz_for_composio, + description=final_new_description, + location=final_new_location, + attendees=final_new_attendees, + ) + if error: + return {"status": "error", "message": error} + updated = {"htmlLink": html_link} + logger.info( + f"Calendar event updated via Composio: event_id={final_event_id}" + ) else: config_data = dict(connector.config) @@ -235,81 +277,79 @@ def create_update_calendar_event_tool( expiry=datetime.fromisoformat(exp) if exp else None, ) - service = await asyncio.get_event_loop().run_in_executor( - None, lambda: build("calendar", "v3", credentials=creds) - ) - - update_body: dict[str, Any] = {} - if final_new_summary is not None: - update_body["summary"] = final_new_summary - if final_new_start_datetime is not None: - update_body["start"] = _build_time_body( - final_new_start_datetime, context + service = await asyncio.get_event_loop().run_in_executor( + None, lambda: build("calendar", "v3", credentials=creds) ) - if final_new_end_datetime is not None: - update_body["end"] = _build_time_body(final_new_end_datetime, context) - if final_new_description is not None: - update_body["description"] = final_new_description - if final_new_location is not None: - update_body["location"] = final_new_location - if final_new_attendees is not None: - update_body["attendees"] = [ - {"email": e.strip()} for e in final_new_attendees if e.strip() - ] - if not update_body: - return { - "status": "error", - "message": "No changes specified. Please provide at least one field to update.", - } - - try: - updated = await asyncio.get_event_loop().run_in_executor( - None, - lambda: ( - service.events() - .patch( - calendarId="primary", - eventId=final_event_id, - body=update_body, - ) - .execute() - ), - ) - except Exception as api_err: - from googleapiclient.errors import HttpError - - if isinstance(api_err, HttpError) and api_err.resp.status == 403: - logger.warning( - f"Insufficient permissions for connector {actual_connector_id}: {api_err}" + update_body: dict[str, Any] = {} + if final_new_summary is not None: + update_body["summary"] = final_new_summary + if final_new_start_datetime is not None: + update_body["start"] = _build_time_body( + final_new_start_datetime, context ) - try: - from sqlalchemy.orm.attributes import flag_modified + if final_new_end_datetime is not None: + update_body["end"] = _build_time_body( + final_new_end_datetime, context + ) + if final_new_description is not None: + update_body["description"] = final_new_description + if final_new_location is not None: + update_body["location"] = final_new_location + if final_new_attendees is not None: + update_body["attendees"] = [ + {"email": e.strip()} for e in final_new_attendees if e.strip() + ] - _res = await db_session.execute( - select(SearchSourceConnector).where( - SearchSourceConnector.id == actual_connector_id + try: + updated = await asyncio.get_event_loop().run_in_executor( + None, + lambda: ( + service.events() + .patch( + calendarId="primary", + eventId=final_event_id, + body=update_body, ) - ) - _conn = _res.scalar_one_or_none() - if _conn and not _conn.config.get("auth_expired"): - _conn.config = {**_conn.config, "auth_expired": True} - flag_modified(_conn, "config") - await db_session.commit() - except Exception: - logger.warning( - "Failed to persist auth_expired for connector %s", - actual_connector_id, - exc_info=True, - ) - return { - "status": "insufficient_permissions", - "connector_id": actual_connector_id, - "message": "This Google Calendar account needs additional permissions. Please re-authenticate in connector settings.", - } - raise + .execute() + ), + ) + except Exception as api_err: + from googleapiclient.errors import HttpError - logger.info(f"Calendar event updated: event_id={final_event_id}") + if isinstance(api_err, HttpError) and api_err.resp.status == 403: + logger.warning( + f"Insufficient permissions for connector {actual_connector_id}: {api_err}" + ) + try: + from sqlalchemy.orm.attributes import flag_modified + + _res = await db_session.execute( + select(SearchSourceConnector).where( + SearchSourceConnector.id == actual_connector_id + ) + ) + _conn = _res.scalar_one_or_none() + if _conn and not _conn.config.get("auth_expired"): + _conn.config = {**_conn.config, "auth_expired": True} + flag_modified(_conn, "config") + await db_session.commit() + except Exception: + logger.warning( + "Failed to persist auth_expired for connector %s", + actual_connector_id, + exc_info=True, + ) + return { + "status": "insufficient_permissions", + "connector_id": actual_connector_id, + "message": "This Google Calendar account needs additional permissions. Please re-authenticate in connector settings.", + } + raise + + logger.info( + f"Calendar event updated via Google API: event_id={final_event_id}" + ) kb_message_suffix = "" if document_id is not None: diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/create_draft.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/create_draft.py index 0bd044695..59e471097 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/create_draft.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/create_draft.py @@ -161,16 +161,39 @@ def create_create_gmail_draft_tool( connector.connector_type == SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR ): - from app.utils.google_credentials import build_composio_credentials - cca_id = connector.config.get("composio_connected_account_id") - if cca_id: - creds = build_composio_credentials(cca_id) - else: + if not cca_id: return { "status": "error", "message": "Composio connected account ID not found for this Gmail connector.", } + + from app.services.composio_service import ComposioService + + ( + draft_id, + draft_message_id, + draft_thread_id, + error, + ) = await ComposioService().create_gmail_draft( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + to=final_to, + subject=final_subject, + body=final_body, + cc=final_cc, + bcc=final_bcc, + ) + if error: + return {"status": "error", "message": error} + created = { + "id": draft_id, + "message": { + "id": draft_message_id, + "threadId": draft_thread_id, + }, + } + logger.info(f"Gmail draft created via Composio: id={draft_id}") else: from google.oauth2.credentials import Credentials @@ -208,63 +231,65 @@ def create_create_gmail_draft_tool( expiry=datetime.fromisoformat(exp) if exp else None, ) - from googleapiclient.discovery import build + from googleapiclient.discovery import build - gmail_service = build("gmail", "v1", credentials=creds) + gmail_service = build("gmail", "v1", credentials=creds) - message = MIMEText(final_body) - message["to"] = final_to - message["subject"] = final_subject - if final_cc: - message["cc"] = final_cc - if final_bcc: - message["bcc"] = final_bcc - raw = base64.urlsafe_b64encode(message.as_bytes()).decode() + message = MIMEText(final_body) + message["to"] = final_to + message["subject"] = final_subject + if final_cc: + message["cc"] = final_cc + if final_bcc: + message["bcc"] = final_bcc + raw = base64.urlsafe_b64encode(message.as_bytes()).decode() - try: - created = await asyncio.get_event_loop().run_in_executor( - None, - lambda: ( - gmail_service.users() - .drafts() - .create(userId="me", body={"message": {"raw": raw}}) - .execute() - ), - ) - except Exception as api_err: - from googleapiclient.errors import HttpError - - if isinstance(api_err, HttpError) and api_err.resp.status == 403: - logger.warning( - f"Insufficient permissions for connector {actual_connector_id}: {api_err}" + try: + created = await asyncio.get_event_loop().run_in_executor( + None, + lambda: ( + gmail_service.users() + .drafts() + .create(userId="me", body={"message": {"raw": raw}}) + .execute() + ), ) - try: - from sqlalchemy.orm.attributes import flag_modified + except Exception as api_err: + from googleapiclient.errors import HttpError - _res = await db_session.execute( - select(SearchSourceConnector).where( - SearchSourceConnector.id == actual_connector_id - ) - ) - _conn = _res.scalar_one_or_none() - if _conn and not _conn.config.get("auth_expired"): - _conn.config = {**_conn.config, "auth_expired": True} - flag_modified(_conn, "config") - await db_session.commit() - except Exception: + if isinstance(api_err, HttpError) and api_err.resp.status == 403: logger.warning( - "Failed to persist auth_expired for connector %s", - actual_connector_id, - exc_info=True, + f"Insufficient permissions for connector {actual_connector_id}: {api_err}" ) - return { - "status": "insufficient_permissions", - "connector_id": actual_connector_id, - "message": "This Gmail account needs additional permissions. Please re-authenticate in connector settings.", - } - raise + try: + from sqlalchemy.orm.attributes import flag_modified - logger.info(f"Gmail draft created: id={created.get('id')}") + _res = await db_session.execute( + select(SearchSourceConnector).where( + SearchSourceConnector.id == actual_connector_id + ) + ) + _conn = _res.scalar_one_or_none() + if _conn and not _conn.config.get("auth_expired"): + _conn.config = {**_conn.config, "auth_expired": True} + flag_modified(_conn, "config") + await db_session.commit() + except Exception: + logger.warning( + "Failed to persist auth_expired for connector %s", + actual_connector_id, + exc_info=True, + ) + return { + "status": "insufficient_permissions", + "connector_id": actual_connector_id, + "message": "This Gmail account needs additional permissions. Please re-authenticate in connector settings.", + } + raise + + logger.info( + f"Gmail draft created via Google API: id={created.get('id')}" + ) kb_message_suffix = "" try: diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/read_email.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/read_email.py index deec1627c..39526f25e 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/read_email.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/read_email.py @@ -50,7 +50,56 @@ def create_read_gmail_email_tool( "message": "No Gmail connector found. Please connect Gmail in your workspace settings.", } - from app.agents.new_chat.tools.gmail.search_emails import _build_credentials + if ( + connector.connector_type + == SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR + ): + cca_id = connector.config.get("composio_connected_account_id") + if not cca_id: + return { + "status": "error", + "message": "Composio connected account ID not found for this Gmail connector.", + } + + from app.agents.new_chat.tools.gmail.search_emails import ( + _format_gmail_summary, + ) + from app.services.composio_service import ComposioService + + detail, error = await ComposioService().get_gmail_message_detail( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + message_id=message_id, + ) + if error: + return {"status": "error", "message": error} + if not detail: + return { + "status": "not_found", + "message": f"Email with ID '{message_id}' not found.", + } + + summary = _format_gmail_summary(detail) + content = ( + f"# {summary['subject']}\n\n" + f"**From:** {summary['from']}\n" + f"**To:** {summary['to']}\n" + f"**Date:** {summary['date']}\n\n" + f"## Message Content\n\n" + f"{detail.get('messageText') or detail.get('snippet') or ''}\n\n" + f"## Message Details\n\n" + f"- **Message ID:** {summary['message_id']}\n" + f"- **Thread ID:** {summary['thread_id']}\n" + ) + return { + "status": "success", + "message_id": summary["message_id"] or message_id, + "content": content, + } + + from app.agents.new_chat.tools.gmail.search_emails import ( + _build_credentials, + ) creds = _build_credentials(connector) diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/search_emails.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/search_emails.py index 2e363609e..a9d7cdedf 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/search_emails.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/search_emails.py @@ -1,5 +1,4 @@ import logging -from datetime import datetime from typing import Any from langchain_core.tools import tool @@ -15,57 +14,6 @@ _GMAIL_TYPES = [ SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR, ] -_token_encryption_cache: object | None = None - - -def _get_token_encryption(): - global _token_encryption_cache - if _token_encryption_cache is None: - from app.config import config - from app.utils.oauth_security import TokenEncryption - - if not config.SECRET_KEY: - raise RuntimeError("SECRET_KEY not configured for token decryption.") - _token_encryption_cache = TokenEncryption(config.SECRET_KEY) - return _token_encryption_cache - - -def _build_credentials(connector: SearchSourceConnector): - """Build Google OAuth Credentials from a connector's stored config. - - Handles both native OAuth connectors (with encrypted tokens) and - Composio-backed connectors. Shared by Gmail and Calendar tools. - """ - from app.utils.google_credentials import COMPOSIO_GOOGLE_CONNECTOR_TYPES - - if connector.connector_type in COMPOSIO_GOOGLE_CONNECTOR_TYPES: - from app.utils.google_credentials import build_composio_credentials - - cca_id = connector.config.get("composio_connected_account_id") - if not cca_id: - raise ValueError("Composio connected account ID not found.") - return build_composio_credentials(cca_id) - - from google.oauth2.credentials import Credentials - - cfg = dict(connector.config) - if cfg.get("_token_encrypted"): - enc = _get_token_encryption() - for key in ("token", "refresh_token", "client_secret"): - if cfg.get(key): - cfg[key] = enc.decrypt_token(cfg[key]) - - exp = (cfg.get("expiry") or "").replace("Z", "") - return Credentials( - token=cfg.get("token"), - refresh_token=cfg.get("refresh_token"), - token_uri=cfg.get("token_uri"), - client_id=cfg.get("client_id"), - client_secret=cfg.get("client_secret"), - scopes=cfg.get("scopes", []), - expiry=datetime.fromisoformat(exp) if exp else None, - ) - def create_search_gmail_tool( db_session: AsyncSession | None = None, @@ -110,6 +58,50 @@ def create_search_gmail_tool( "message": "No Gmail connector found. Please connect Gmail in your workspace settings.", } + if ( + connector.connector_type + == SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR + ): + cca_id = connector.config.get("composio_connected_account_id") + if not cca_id: + return { + "status": "error", + "message": "Composio connected account ID not found for this Gmail connector.", + } + + from app.agents.new_chat.tools.gmail.search_emails import ( + _format_gmail_summary, + ) + from app.services.composio_service import ComposioService + + ( + messages, + _next, + _estimate, + error, + ) = await ComposioService().get_gmail_messages( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + query=query, + max_results=max_results, + ) + if error: + return {"status": "error", "message": error} + + emails = [_format_gmail_summary(m) for m in messages] + if not emails: + return { + "status": "success", + "emails": [], + "total": 0, + "message": "No emails found.", + } + return {"status": "success", "emails": emails, "total": len(emails)} + + from app.agents.new_chat.tools.gmail.search_emails import ( + _build_credentials, + ) + creds = _build_credentials(connector) from app.connectors.google_gmail_connector import GoogleGmailConnector diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/send_email.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/send_email.py index c3f0999f4..d5de24b62 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/send_email.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/send_email.py @@ -162,16 +162,31 @@ def create_send_gmail_email_tool( connector.connector_type == SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR ): - from app.utils.google_credentials import build_composio_credentials - cca_id = connector.config.get("composio_connected_account_id") - if cca_id: - creds = build_composio_credentials(cca_id) - else: + if not cca_id: return { "status": "error", "message": "Composio connected account ID not found for this Gmail connector.", } + + from app.services.composio_service import ComposioService + + ( + sent_message_id, + sent_thread_id, + error, + ) = await ComposioService().send_gmail_email( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + to=final_to, + subject=final_subject, + body=final_body, + cc=final_cc, + bcc=final_bcc, + ) + if error: + return {"status": "error", "message": error} + sent = {"id": sent_message_id, "threadId": sent_thread_id} else: from google.oauth2.credentials import Credentials @@ -209,61 +224,61 @@ def create_send_gmail_email_tool( expiry=datetime.fromisoformat(exp) if exp else None, ) - from googleapiclient.discovery import build + from googleapiclient.discovery import build - gmail_service = build("gmail", "v1", credentials=creds) + gmail_service = build("gmail", "v1", credentials=creds) - message = MIMEText(final_body) - message["to"] = final_to - message["subject"] = final_subject - if final_cc: - message["cc"] = final_cc - if final_bcc: - message["bcc"] = final_bcc - raw = base64.urlsafe_b64encode(message.as_bytes()).decode() + message = MIMEText(final_body) + message["to"] = final_to + message["subject"] = final_subject + if final_cc: + message["cc"] = final_cc + if final_bcc: + message["bcc"] = final_bcc + raw = base64.urlsafe_b64encode(message.as_bytes()).decode() - try: - sent = await asyncio.get_event_loop().run_in_executor( - None, - lambda: ( - gmail_service.users() - .messages() - .send(userId="me", body={"raw": raw}) - .execute() - ), - ) - except Exception as api_err: - from googleapiclient.errors import HttpError - - if isinstance(api_err, HttpError) and api_err.resp.status == 403: - logger.warning( - f"Insufficient permissions for connector {actual_connector_id}: {api_err}" + try: + sent = await asyncio.get_event_loop().run_in_executor( + None, + lambda: ( + gmail_service.users() + .messages() + .send(userId="me", body={"raw": raw}) + .execute() + ), ) - try: - from sqlalchemy.orm.attributes import flag_modified + except Exception as api_err: + from googleapiclient.errors import HttpError - _res = await db_session.execute( - select(SearchSourceConnector).where( - SearchSourceConnector.id == actual_connector_id - ) - ) - _conn = _res.scalar_one_or_none() - if _conn and not _conn.config.get("auth_expired"): - _conn.config = {**_conn.config, "auth_expired": True} - flag_modified(_conn, "config") - await db_session.commit() - except Exception: + if isinstance(api_err, HttpError) and api_err.resp.status == 403: logger.warning( - "Failed to persist auth_expired for connector %s", - actual_connector_id, - exc_info=True, + f"Insufficient permissions for connector {actual_connector_id}: {api_err}" ) - return { - "status": "insufficient_permissions", - "connector_id": actual_connector_id, - "message": "This Gmail account needs additional permissions. Please re-authenticate in connector settings.", - } - raise + try: + from sqlalchemy.orm.attributes import flag_modified + + _res = await db_session.execute( + select(SearchSourceConnector).where( + SearchSourceConnector.id == actual_connector_id + ) + ) + _conn = _res.scalar_one_or_none() + if _conn and not _conn.config.get("auth_expired"): + _conn.config = {**_conn.config, "auth_expired": True} + flag_modified(_conn, "config") + await db_session.commit() + except Exception: + logger.warning( + "Failed to persist auth_expired for connector %s", + actual_connector_id, + exc_info=True, + ) + return { + "status": "insufficient_permissions", + "connector_id": actual_connector_id, + "message": "This Gmail account needs additional permissions. Please re-authenticate in connector settings.", + } + raise logger.info( f"Gmail email sent: id={sent.get('id')}, threadId={sent.get('threadId')}" diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/trash_email.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/trash_email.py index 1f1f6227a..b78f88934 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/trash_email.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/trash_email.py @@ -162,16 +162,22 @@ def create_trash_gmail_email_tool( connector.connector_type == SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR ): - from app.utils.google_credentials import build_composio_credentials - cca_id = connector.config.get("composio_connected_account_id") - if cca_id: - creds = build_composio_credentials(cca_id) - else: + if not cca_id: return { "status": "error", "message": "Composio connected account ID not found for this Gmail connector.", } + + from app.services.composio_service import ComposioService + + error = await ComposioService().trash_gmail_message( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + message_id=final_message_id, + ) + if error: + return {"status": "error", "message": error} else: from google.oauth2.credentials import Credentials @@ -209,49 +215,49 @@ def create_trash_gmail_email_tool( expiry=datetime.fromisoformat(exp) if exp else None, ) - from googleapiclient.discovery import build + from googleapiclient.discovery import build - gmail_service = build("gmail", "v1", credentials=creds) + gmail_service = build("gmail", "v1", credentials=creds) - try: - await asyncio.get_event_loop().run_in_executor( - None, - lambda: ( - gmail_service.users() - .messages() - .trash(userId="me", id=final_message_id) - .execute() - ), - ) - except Exception as api_err: - from googleapiclient.errors import HttpError - - if isinstance(api_err, HttpError) and api_err.resp.status == 403: - logger.warning( - f"Insufficient permissions for connector {connector.id}: {api_err}" + try: + await asyncio.get_event_loop().run_in_executor( + None, + lambda: ( + gmail_service.users() + .messages() + .trash(userId="me", id=final_message_id) + .execute() + ), ) - try: - from sqlalchemy.orm.attributes import flag_modified + except Exception as api_err: + from googleapiclient.errors import HttpError - if not connector.config.get("auth_expired"): - connector.config = { - **connector.config, - "auth_expired": True, - } - flag_modified(connector, "config") - await db_session.commit() - except Exception: + if isinstance(api_err, HttpError) and api_err.resp.status == 403: logger.warning( - "Failed to persist auth_expired for connector %s", - connector.id, - exc_info=True, + f"Insufficient permissions for connector {connector.id}: {api_err}" ) - return { - "status": "insufficient_permissions", - "connector_id": connector.id, - "message": "This Gmail account needs additional permissions. Please re-authenticate in connector settings.", - } - raise + try: + from sqlalchemy.orm.attributes import flag_modified + + if not connector.config.get("auth_expired"): + connector.config = { + **connector.config, + "auth_expired": True, + } + flag_modified(connector, "config") + await db_session.commit() + except Exception: + logger.warning( + "Failed to persist auth_expired for connector %s", + connector.id, + exc_info=True, + ) + return { + "status": "insufficient_permissions", + "connector_id": connector.id, + "message": "This Gmail account needs additional permissions. Please re-authenticate in connector settings.", + } + raise logger.info(f"Gmail email trashed: message_id={final_message_id}") diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/update_draft.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/update_draft.py index 91178cd21..b6688ac53 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/update_draft.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/gmail/tools/update_draft.py @@ -192,16 +192,51 @@ def create_update_gmail_draft_tool( connector.connector_type == SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR ): - from app.utils.google_credentials import build_composio_credentials - cca_id = connector.config.get("composio_connected_account_id") - if cca_id: - creds = build_composio_credentials(cca_id) - else: + if not cca_id: return { "status": "error", "message": "Composio connected account ID not found for this Gmail connector.", } + + if not final_draft_id: + return { + "status": "error", + "message": ( + "Could not find this draft in Gmail. " + "It may have already been sent or deleted." + ), + } + + from app.services.composio_service import ComposioService + + ( + new_draft_id, + new_message_id, + error, + ) = await ComposioService().update_gmail_draft( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + draft_id=final_draft_id, + to=final_to or None, + subject=final_subject, + body=final_body, + cc=final_cc, + bcc=final_bcc, + ) + if error: + if "not found" in error.lower() or "no longer" in error.lower(): + return { + "status": "error", + "message": "Draft no longer exists in Gmail. It may have been sent or deleted.", + } + return {"status": "error", "message": error} + + updated = { + "id": new_draft_id or final_draft_id, + "message": {"id": new_message_id} if new_message_id else {}, + } + logger.info(f"Gmail draft updated via Composio: id={updated.get('id')}") else: from google.oauth2.credentials import Credentials @@ -239,88 +274,90 @@ def create_update_gmail_draft_tool( expiry=datetime.fromisoformat(exp) if exp else None, ) - from googleapiclient.discovery import build + from googleapiclient.discovery import build - gmail_service = build("gmail", "v1", credentials=creds) + gmail_service = build("gmail", "v1", credentials=creds) - # Resolve draft_id if not already available - if not final_draft_id: - logger.info( - f"draft_id not in metadata, looking up via drafts.list for message_id={message_id}" - ) - final_draft_id = await _find_draft_id_by_message( - gmail_service, message_id - ) - - if not final_draft_id: - return { - "status": "error", - "message": ( - "Could not find this draft in Gmail. " - "It may have already been sent or deleted." - ), - } - - message = MIMEText(final_body) - if final_to: - message["to"] = final_to - message["subject"] = final_subject - if final_cc: - message["cc"] = final_cc - if final_bcc: - message["bcc"] = final_bcc - raw = base64.urlsafe_b64encode(message.as_bytes()).decode() - - try: - updated = await asyncio.get_event_loop().run_in_executor( - None, - lambda: ( - gmail_service.users() - .drafts() - .update( - userId="me", - id=final_draft_id, - body={"message": {"raw": raw}}, - ) - .execute() - ), - ) - except Exception as api_err: - from googleapiclient.errors import HttpError - - if isinstance(api_err, HttpError) and api_err.resp.status == 403: - logger.warning( - f"Insufficient permissions for connector {connector.id}: {api_err}" + # Resolve draft_id if not already available + if not final_draft_id: + logger.info( + f"draft_id not in metadata, looking up via drafts.list for message_id={message_id}" + ) + final_draft_id = await _find_draft_id_by_message( + gmail_service, message_id ) - try: - from sqlalchemy.orm.attributes import flag_modified - if not connector.config.get("auth_expired"): - connector.config = { - **connector.config, - "auth_expired": True, - } - flag_modified(connector, "config") - await db_session.commit() - except Exception: - logger.warning( - "Failed to persist auth_expired for connector %s", - connector.id, - exc_info=True, - ) - return { - "status": "insufficient_permissions", - "connector_id": connector.id, - "message": "This Gmail account needs additional permissions. Please re-authenticate in connector settings.", - } - if isinstance(api_err, HttpError) and api_err.resp.status == 404: + if not final_draft_id: return { "status": "error", - "message": "Draft no longer exists in Gmail. It may have been sent or deleted.", + "message": ( + "Could not find this draft in Gmail. " + "It may have already been sent or deleted." + ), } - raise - logger.info(f"Gmail draft updated: id={updated.get('id')}") + message = MIMEText(final_body) + if final_to: + message["to"] = final_to + message["subject"] = final_subject + if final_cc: + message["cc"] = final_cc + if final_bcc: + message["bcc"] = final_bcc + raw = base64.urlsafe_b64encode(message.as_bytes()).decode() + + try: + updated = await asyncio.get_event_loop().run_in_executor( + None, + lambda: ( + gmail_service.users() + .drafts() + .update( + userId="me", + id=final_draft_id, + body={"message": {"raw": raw}}, + ) + .execute() + ), + ) + except Exception as api_err: + from googleapiclient.errors import HttpError + + if isinstance(api_err, HttpError) and api_err.resp.status == 403: + logger.warning( + f"Insufficient permissions for connector {connector.id}: {api_err}" + ) + try: + from sqlalchemy.orm.attributes import flag_modified + + if not connector.config.get("auth_expired"): + connector.config = { + **connector.config, + "auth_expired": True, + } + flag_modified(connector, "config") + await db_session.commit() + except Exception: + logger.warning( + "Failed to persist auth_expired for connector %s", + connector.id, + exc_info=True, + ) + return { + "status": "insufficient_permissions", + "connector_id": connector.id, + "message": "This Gmail account needs additional permissions. Please re-authenticate in connector settings.", + } + if isinstance(api_err, HttpError) and api_err.resp.status == 404: + return { + "status": "error", + "message": "Draft no longer exists in Gmail. It may have been sent or deleted.", + } + raise + + logger.info( + f"Gmail draft updated via Google API: id={updated.get('id')}" + ) kb_message_suffix = "" if document_id: diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/google_drive/tools/create_file.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/google_drive/tools/create_file.py index f36db8f3f..9e9a30429 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/google_drive/tools/create_file.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/google_drive/tools/create_file.py @@ -179,59 +179,96 @@ def create_create_google_drive_file_tool( f"Creating Google Drive file: name='{final_name}', type='{final_file_type}', connector={actual_connector_id}" ) - pre_built_creds = None + async def _flag_auth_expired() -> None: + try: + from sqlalchemy.orm.attributes import flag_modified + + _res = await db_session.execute( + select(SearchSourceConnector).where( + SearchSourceConnector.id == actual_connector_id + ) + ) + _conn = _res.scalar_one_or_none() + if _conn and not _conn.config.get("auth_expired"): + _conn.config = {**_conn.config, "auth_expired": True} + flag_modified(_conn, "config") + await db_session.commit() + except Exception: + logger.warning( + "Failed to persist auth_expired for connector %s", + actual_connector_id, + exc_info=True, + ) + if ( connector.connector_type == SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR ): - from app.utils.google_credentials import build_composio_credentials - cca_id = connector.config.get("composio_connected_account_id") - if cca_id: - pre_built_creds = build_composio_credentials(cca_id) + if not cca_id: + return { + "status": "error", + "message": "Composio connected account ID not found for this Google Drive connector.", + } - client = GoogleDriveClient( - session=db_session, - connector_id=actual_connector_id, - credentials=pre_built_creds, - ) - try: - created = await client.create_file( + from app.services.composio_service import ComposioService + + created, error = await ComposioService().create_drive_file_from_text( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", name=final_name, mime_type=mime_type, - parent_folder_id=final_parent_folder_id, content=final_content, + parent_id=final_parent_folder_id, ) - except HttpError as http_err: - if http_err.resp.status == 403: - logger.warning( - f"Insufficient permissions for connector {actual_connector_id}: {http_err}" - ) - try: - from sqlalchemy.orm.attributes import flag_modified - _res = await db_session.execute( - select(SearchSourceConnector).where( - SearchSourceConnector.id == actual_connector_id - ) - ) - _conn = _res.scalar_one_or_none() - if _conn and not _conn.config.get("auth_expired"): - _conn.config = {**_conn.config, "auth_expired": True} - flag_modified(_conn, "config") - await db_session.commit() - except Exception: + if error or not created: + err_lower = (error or "").lower() + if ( + "insufficient" in err_lower + or "permission" in err_lower + or "403" in err_lower + ): logger.warning( - "Failed to persist auth_expired for connector %s", - actual_connector_id, - exc_info=True, + f"Insufficient permissions for Composio Drive connector {actual_connector_id}: {error}" ) + await _flag_auth_expired() + return { + "status": "insufficient_permissions", + "connector_id": actual_connector_id, + "message": "This Google Drive account needs additional permissions. Please re-authenticate in connector settings.", + } + logger.error( + f"Composio Drive create_file failed for connector {actual_connector_id}: {error}" + ) return { - "status": "insufficient_permissions", - "connector_id": actual_connector_id, - "message": "This Google Drive account needs additional permissions. Please re-authenticate in connector settings.", + "status": "error", + "message": "Something went wrong while creating the file. Please try again.", } - raise + else: + client = GoogleDriveClient( + session=db_session, + connector_id=actual_connector_id, + ) + try: + created = await client.create_file( + name=final_name, + mime_type=mime_type, + parent_folder_id=final_parent_folder_id, + content=final_content, + ) + except HttpError as http_err: + if http_err.resp.status == 403: + logger.warning( + f"Insufficient permissions for connector {actual_connector_id}: {http_err}" + ) + await _flag_auth_expired() + return { + "status": "insufficient_permissions", + "connector_id": actual_connector_id, + "message": "This Google Drive account needs additional permissions. Please re-authenticate in connector settings.", + } + raise logger.info( f"Google Drive file created: id={created.get('id')}, name={created.get('name')}" diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/google_drive/tools/trash_file.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/google_drive/tools/trash_file.py index 832afff0d..f7531cf3d 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/google_drive/tools/trash_file.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/connectors/google_drive/tools/trash_file.py @@ -158,51 +158,84 @@ def create_delete_google_drive_file_tool( f"Deleting Google Drive file: file_id='{final_file_id}', connector={final_connector_id}" ) - pre_built_creds = None + async def _flag_auth_expired() -> None: + try: + from sqlalchemy.orm.attributes import flag_modified + + if not connector.config.get("auth_expired"): + connector.config = { + **connector.config, + "auth_expired": True, + } + flag_modified(connector, "config") + await db_session.commit() + except Exception: + logger.warning( + "Failed to persist auth_expired for connector %s", + connector.id, + exc_info=True, + ) + if ( connector.connector_type == SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR ): - from app.utils.google_credentials import build_composio_credentials - cca_id = connector.config.get("composio_connected_account_id") - if cca_id: - pre_built_creds = build_composio_credentials(cca_id) - - client = GoogleDriveClient( - session=db_session, - connector_id=connector.id, - credentials=pre_built_creds, - ) - try: - await client.trash_file(file_id=final_file_id) - except HttpError as http_err: - if http_err.resp.status == 403: - logger.warning( - f"Insufficient permissions for connector {connector.id}: {http_err}" - ) - try: - from sqlalchemy.orm.attributes import flag_modified - - if not connector.config.get("auth_expired"): - connector.config = { - **connector.config, - "auth_expired": True, - } - flag_modified(connector, "config") - await db_session.commit() - except Exception: - logger.warning( - "Failed to persist auth_expired for connector %s", - connector.id, - exc_info=True, - ) + if not cca_id: return { - "status": "insufficient_permissions", - "connector_id": connector.id, - "message": "This Google Drive account needs additional permissions. Please re-authenticate in connector settings.", + "status": "error", + "message": "Composio connected account ID not found for this Google Drive connector.", } - raise + + from app.services.composio_service import ComposioService + + error = await ComposioService().trash_drive_file( + connected_account_id=cca_id, + entity_id=f"surfsense_{user_id}", + file_id=final_file_id, + ) + if error: + err_lower = error.lower() + if ( + "insufficient" in err_lower + or "permission" in err_lower + or "403" in err_lower + ): + logger.warning( + f"Insufficient permissions for Composio Drive connector {connector.id}: {error}" + ) + await _flag_auth_expired() + return { + "status": "insufficient_permissions", + "connector_id": connector.id, + "message": "This Google Drive account needs additional permissions. Please re-authenticate in connector settings.", + } + logger.error( + f"Composio Drive trash_file failed for connector {connector.id}: {error}" + ) + return { + "status": "error", + "message": "Something went wrong while trashing the file. Please try again.", + } + else: + client = GoogleDriveClient( + session=db_session, + connector_id=connector.id, + ) + try: + await client.trash_file(file_id=final_file_id) + except HttpError as http_err: + if http_err.resp.status == 403: + logger.warning( + f"Insufficient permissions for connector {connector.id}: {http_err}" + ) + await _flag_auth_expired() + return { + "status": "insufficient_permissions", + "connector_id": connector.id, + "message": "This Google Drive account needs additional permissions. Please re-authenticate in connector settings.", + } + raise logger.info( f"Google Drive file deleted (moved to trash): file_id={final_file_id}" diff --git a/surfsense_backend/app/services/composio_service.py b/surfsense_backend/app/services/composio_service.py index edfab1d15..d73a0d4ce 100644 --- a/surfsense_backend/app/services/composio_service.py +++ b/surfsense_backend/app/services/composio_service.py @@ -1027,6 +1027,505 @@ class ComposioService: logger.error(f"Failed to list Calendar events: {e!s}") return [], str(e) + @staticmethod + def _unwrap_response_data(data: Any) -> Any: + """Composio responses often nest the meaningful payload under + ``data.data.response_data``. Walk that envelope safely and return + whichever inner dict actually has the result keys.""" + if not isinstance(data, dict): + return data + inner = data.get("data", data) + if isinstance(inner, dict): + return inner.get("response_data", inner) + return inner + + @staticmethod + def _split_email_csv(value: str | None) -> list[str] | None: + """Tools accept comma-separated cc/bcc strings; Composio expects an array.""" + if not value: + return None + addrs = [e.strip() for e in value.split(",") if e.strip()] + return addrs or None + + # ===== Gmail write methods ===== + + async def send_gmail_email( + self, + connected_account_id: str, + entity_id: str, + to: str, + subject: str, + body: str, + cc: str | None = None, + bcc: str | None = None, + is_html: bool = False, + ) -> tuple[str | None, str | None, str | None]: + """Send a Gmail message via the Composio ``GMAIL_SEND_EMAIL`` toolkit. + + Returns: + Tuple of (message_id, thread_id, error). On success ``error`` is + None and at least one of the IDs is populated when Composio + returns them; on failure both IDs are None. + """ + try: + params: dict[str, Any] = { + "recipient_email": to, + "subject": subject, + "body": body, + "is_html": is_html, + } + if cc: + cc_list = self._split_email_csv(cc) + if cc_list: + params["cc"] = cc_list + if bcc: + bcc_list = self._split_email_csv(bcc) + if bcc_list: + params["bcc"] = bcc_list + + result = await self.execute_tool( + connected_account_id=connected_account_id, + tool_name="GMAIL_SEND_EMAIL", + params=params, + entity_id=entity_id, + ) + if not result.get("success"): + return None, None, result.get("error", "Unknown error") + + payload = self._unwrap_response_data(result.get("data", {})) + message_id = None + thread_id = None + if isinstance(payload, dict): + message_id = ( + payload.get("id") + or payload.get("message_id") + or payload.get("messageId") + ) + thread_id = payload.get("threadId") or payload.get("thread_id") + return message_id, thread_id, None + except Exception as e: + logger.error(f"Failed to send Gmail email: {e!s}") + return None, None, str(e) + + async def create_gmail_draft( + self, + connected_account_id: str, + entity_id: str, + to: str, + subject: str, + body: str, + cc: str | None = None, + bcc: str | None = None, + is_html: bool = False, + ) -> tuple[str | None, str | None, str | None, str | None]: + """Create a Gmail draft via the Composio ``GMAIL_CREATE_EMAIL_DRAFT`` toolkit. + + Returns: + Tuple of (draft_id, message_id, thread_id, error). On success + ``error`` is None and ``draft_id`` is populated. + """ + try: + params: dict[str, Any] = { + "recipient_email": to, + "subject": subject, + "body": body, + "is_html": is_html, + } + cc_list = self._split_email_csv(cc) + if cc_list: + params["cc"] = cc_list + bcc_list = self._split_email_csv(bcc) + if bcc_list: + params["bcc"] = bcc_list + + result = await self.execute_tool( + connected_account_id=connected_account_id, + tool_name="GMAIL_CREATE_EMAIL_DRAFT", + params=params, + entity_id=entity_id, + ) + if not result.get("success"): + return None, None, None, result.get("error", "Unknown error") + + payload = self._unwrap_response_data(result.get("data", {})) + draft_id = None + message_id = None + thread_id = None + if isinstance(payload, dict): + draft_id = payload.get("id") or payload.get("draft_id") + draft_message = payload.get("message") or {} + if isinstance(draft_message, dict): + message_id = draft_message.get("id") or draft_message.get( + "message_id" + ) + thread_id = draft_message.get("threadId") or draft_message.get( + "thread_id" + ) + if message_id is None: + message_id = payload.get("message_id") or payload.get("messageId") + if thread_id is None: + thread_id = payload.get("thread_id") or payload.get("threadId") + return draft_id, message_id, thread_id, None + except Exception as e: + logger.error(f"Failed to create Gmail draft: {e!s}") + return None, None, None, str(e) + + async def update_gmail_draft( + self, + connected_account_id: str, + entity_id: str, + draft_id: str, + to: str | None = None, + subject: str | None = None, + body: str | None = None, + cc: str | None = None, + bcc: str | None = None, + is_html: bool = False, + ) -> tuple[str | None, str | None, str | None]: + """Update an existing Gmail draft via ``GMAIL_UPDATE_DRAFT``. + + Returns: + Tuple of (draft_id, message_id, error). + """ + try: + params: dict[str, Any] = { + "draft_id": draft_id, + "is_html": is_html, + } + if to: + params["recipient_email"] = to + if subject is not None: + params["subject"] = subject + if body is not None: + params["body"] = body + cc_list = self._split_email_csv(cc) + if cc_list: + params["cc"] = cc_list + bcc_list = self._split_email_csv(bcc) + if bcc_list: + params["bcc"] = bcc_list + + result = await self.execute_tool( + connected_account_id=connected_account_id, + tool_name="GMAIL_UPDATE_DRAFT", + params=params, + entity_id=entity_id, + ) + if not result.get("success"): + return None, None, result.get("error", "Unknown error") + + payload = self._unwrap_response_data(result.get("data", {})) + new_draft_id = draft_id + message_id = None + if isinstance(payload, dict): + new_draft_id = payload.get("id") or payload.get("draft_id") or draft_id + draft_message = payload.get("message") or {} + if isinstance(draft_message, dict): + message_id = draft_message.get("id") or draft_message.get( + "message_id" + ) + if message_id is None: + message_id = payload.get("message_id") or payload.get("messageId") + return new_draft_id, message_id, None + except Exception as e: + logger.error(f"Failed to update Gmail draft: {e!s}") + return None, None, str(e) + + async def trash_gmail_message( + self, + connected_account_id: str, + entity_id: str, + message_id: str, + ) -> str | None: + """Move a Gmail message to trash via ``GMAIL_MOVE_TO_TRASH``. + + Returns the error message on failure, ``None`` on success. + """ + try: + result = await self.execute_tool( + connected_account_id=connected_account_id, + tool_name="GMAIL_MOVE_TO_TRASH", + params={"message_id": message_id}, + entity_id=entity_id, + ) + if not result.get("success"): + return result.get("error", "Unknown error") + return None + except Exception as e: + logger.error(f"Failed to trash Gmail message: {e!s}") + return str(e) + + # ===== Google Calendar write methods ===== + + async def create_calendar_event( + self, + connected_account_id: str, + entity_id: str, + summary: str, + start_datetime: str, + end_datetime: str, + timezone: str | None = None, + description: str | None = None, + location: str | None = None, + attendees: list[str] | None = None, + calendar_id: str = "primary", + ) -> tuple[str | None, str | None, str | None]: + """Create a Google Calendar event via ``GOOGLECALENDAR_CREATE_EVENT``. + + Composio strips trailing timezone info on ``start_datetime`` / + ``end_datetime`` and uses the ``timezone`` field as the IANA name, + so callers may pass ISO 8601 strings with or without offsets. + + Returns: + Tuple of (event_id, html_link, error). + """ + try: + params: dict[str, Any] = { + "summary": summary, + "start_datetime": start_datetime, + "end_datetime": end_datetime, + "calendar_id": calendar_id, + } + if timezone: + params["timezone"] = timezone + if description: + params["description"] = description + if location: + params["location"] = location + if attendees: + params["attendees"] = [a for a in attendees if a] + + result = await self.execute_tool( + connected_account_id=connected_account_id, + tool_name="GOOGLECALENDAR_CREATE_EVENT", + params=params, + entity_id=entity_id, + ) + if not result.get("success"): + return None, None, result.get("error", "Unknown error") + + payload = self._unwrap_response_data(result.get("data", {})) + event_id = None + html_link = None + if isinstance(payload, dict): + event_id = payload.get("id") or payload.get("event_id") + html_link = payload.get("htmlLink") or payload.get("html_link") + return event_id, html_link, None + except Exception as e: + logger.error(f"Failed to create Calendar event: {e!s}") + return None, None, str(e) + + async def update_calendar_event( + self, + connected_account_id: str, + entity_id: str, + event_id: str, + summary: str | None = None, + start_time: str | None = None, + end_time: str | None = None, + timezone: str | None = None, + description: str | None = None, + location: str | None = None, + attendees: list[str] | None = None, + calendar_id: str = "primary", + ) -> tuple[str | None, str | None, str | None]: + """Patch an existing Google Calendar event via ``GOOGLECALENDAR_PATCH_EVENT``. + + Uses PATCH (not PUT) semantics so omitted fields are preserved. + + Returns: + Tuple of (event_id, html_link, error). + """ + try: + params: dict[str, Any] = { + "event_id": event_id, + "calendar_id": calendar_id, + } + if summary is not None: + params["summary"] = summary + if start_time is not None: + params["start_time"] = start_time + if end_time is not None: + params["end_time"] = end_time + if timezone: + params["timezone"] = timezone + if description is not None: + params["description"] = description + if location is not None: + params["location"] = location + if attendees is not None: + params["attendees"] = [a for a in attendees if a] + + result = await self.execute_tool( + connected_account_id=connected_account_id, + tool_name="GOOGLECALENDAR_PATCH_EVENT", + params=params, + entity_id=entity_id, + ) + if not result.get("success"): + return None, None, result.get("error", "Unknown error") + + payload = self._unwrap_response_data(result.get("data", {})) + new_event_id = event_id + html_link = None + if isinstance(payload, dict): + new_event_id = payload.get("id") or payload.get("event_id") or event_id + html_link = payload.get("htmlLink") or payload.get("html_link") + return new_event_id, html_link, None + except Exception as e: + logger.error(f"Failed to patch Calendar event: {e!s}") + return None, None, str(e) + + async def delete_calendar_event( + self, + connected_account_id: str, + entity_id: str, + event_id: str, + calendar_id: str = "primary", + ) -> str | None: + """Delete a Google Calendar event via ``GOOGLECALENDAR_DELETE_EVENT``. + + Returns the error message on failure, ``None`` on success (idempotent + on already-deleted events). + """ + try: + result = await self.execute_tool( + connected_account_id=connected_account_id, + tool_name="GOOGLECALENDAR_DELETE_EVENT", + params={ + "event_id": event_id, + "calendar_id": calendar_id, + }, + entity_id=entity_id, + ) + if not result.get("success"): + return result.get("error", "Unknown error") + return None + except Exception as e: + logger.error(f"Failed to delete Calendar event: {e!s}") + return str(e) + + # ===== Google Drive write methods ===== + + @staticmethod + def _drive_web_view_link(file_id: str, mime_type: str | None) -> str: + """Synthesize a Google Drive ``webViewLink`` from id + mimeType. + + Composio's ``GOOGLEDRIVE_CREATE_FILE_FROM_TEXT`` returns flat + metadata (id, name, mimeType) but does not always include a + ``webViewLink``. We rebuild the canonical UI URL based on the + Workspace MIME type so callers can keep using a single field. + """ + if not file_id: + return "" + mt = (mime_type or "").lower() + if mt == "application/vnd.google-apps.document": + return f"https://docs.google.com/document/d/{file_id}/edit" + if mt == "application/vnd.google-apps.spreadsheet": + return f"https://docs.google.com/spreadsheets/d/{file_id}/edit" + if mt == "application/vnd.google-apps.presentation": + return f"https://docs.google.com/presentation/d/{file_id}/edit" + if mt == "application/vnd.google-apps.folder": + return f"https://drive.google.com/drive/folders/{file_id}" + return f"https://drive.google.com/file/d/{file_id}/view" + + async def create_drive_file_from_text( + self, + connected_account_id: str, + entity_id: str, + name: str, + mime_type: str, + content: str | None = None, + parent_id: str | None = None, + ) -> tuple[dict[str, Any] | None, str | None]: + """Create a Google Drive file from text via ``GOOGLEDRIVE_CREATE_FILE_FROM_TEXT``. + + Composio's tool requires ``text_content`` even for "empty" files; + an empty string is accepted. Native Workspace types (Docs, Sheets) + are produced by setting ``mime_type`` to the Google Apps MIME, and + Drive auto-converts the text payload (e.g. CSV → Sheet). + + Returns: + Tuple of (file_meta, error). ``file_meta`` keys: + ``id``, ``name``, ``mimeType``, ``webViewLink``. + """ + try: + params: dict[str, Any] = { + "file_name": name, + "mime_type": mime_type, + "text_content": content if content is not None else "", + } + if parent_id: + params["parent_id"] = parent_id + + result = await self.execute_tool( + connected_account_id=connected_account_id, + tool_name="GOOGLEDRIVE_CREATE_FILE_FROM_TEXT", + params=params, + entity_id=entity_id, + ) + if not result.get("success"): + return None, result.get("error", "Unknown error") + + payload = self._unwrap_response_data(result.get("data", {})) + file_id: str | None = None + file_name: str | None = name + mime: str | None = mime_type + web_view_link: str | None = None + + if isinstance(payload, dict): + file_id = ( + payload.get("id") or payload.get("file_id") or payload.get("fileId") + ) + file_name = payload.get("name") or payload.get("file_name") or name + mime = payload.get("mimeType") or payload.get("mime_type") or mime_type + web_view_link = payload.get("webViewLink") or payload.get( + "web_view_link" + ) + + if not file_id: + return None, "Composio response did not include a file id" + + if not web_view_link: + web_view_link = self._drive_web_view_link(file_id, mime) + + return ( + { + "id": file_id, + "name": file_name, + "mimeType": mime, + "webViewLink": web_view_link, + }, + None, + ) + except Exception as e: + logger.error(f"Failed to create Drive file: {e!s}") + return None, str(e) + + async def trash_drive_file( + self, + connected_account_id: str, + entity_id: str, + file_id: str, + ) -> str | None: + """Move a Google Drive file to trash via ``GOOGLEDRIVE_TRASH_FILE``. + + Returns the error message on failure, ``None`` on success. + """ + try: + result = await self.execute_tool( + connected_account_id=connected_account_id, + tool_name="GOOGLEDRIVE_TRASH_FILE", + params={"file_id": file_id}, + entity_id=entity_id, + ) + if not result.get("success"): + return result.get("error", "Unknown error") + return None + except Exception as e: + logger.error(f"Failed to trash Drive file: {e!s}") + return str(e) + # ===== User Info Methods ===== async def get_connected_account_email( From 0654662d29c31f8859c6434ad2a69636e9517557 Mon Sep 17 00:00:00 2001 From: "DESKTOP-RTLN3BA\\$punk" Date: Tue, 5 May 2026 19:10:35 -0700 Subject: [PATCH 33/96] refactor(plate-editor): replace markdown deserialization with safeDeserializeMarkdown utility --- .../components/editor/plate-editor.tsx | 28 ++++---- .../editor/utils/safe-deserialize.ts | 64 +++++++++++++++++++ 2 files changed, 79 insertions(+), 13 deletions(-) create mode 100644 surfsense_web/components/editor/utils/safe-deserialize.ts diff --git a/surfsense_web/components/editor/plate-editor.tsx b/surfsense_web/components/editor/plate-editor.tsx index c42cb991e..51ad7d700 100644 --- a/surfsense_web/components/editor/plate-editor.tsx +++ b/surfsense_web/components/editor/plate-editor.tsx @@ -11,6 +11,7 @@ import { EditorSaveContext } from "@/components/editor/editor-save-context"; import { CitationKit, injectCitationNodes } from "@/components/editor/plugins/citation-kit"; import { type EditorPreset, presetMap } from "@/components/editor/presets"; import { escapeMdxExpressions } from "@/components/editor/utils/escape-mdx"; +import { safeDeserializeMarkdown } from "@/components/editor/utils/safe-deserialize"; import { Editor, EditorContainer } from "@/components/ui/editor"; import { preprocessCitationMarkdown } from "@/lib/citations/citation-parser"; @@ -169,15 +170,17 @@ export function PlateEditor({ : markdown ? (editor) => { if (!enableCitations) { - return editor - .getApi(MarkdownPlugin) - .markdown.deserialize(escapeMdxExpressions(markdown)); + return safeDeserializeMarkdown( + editor, + escapeMdxExpressions(markdown) + ) as Value; } const { content: rewritten, urlMap } = preprocessCitationMarkdown(markdown); - const value = editor - .getApi(MarkdownPlugin) - .markdown.deserialize(escapeMdxExpressions(rewritten)); - return injectCitationNodes(value as Descendant[], urlMap) as Value; + const value = safeDeserializeMarkdown( + editor, + escapeMdxExpressions(rewritten) + ); + return injectCitationNodes(value, urlMap) as Value; } : undefined, }); @@ -200,14 +203,13 @@ export function PlateEditor({ let newValue: Descendant[]; if (enableCitations) { const { content: rewritten, urlMap } = preprocessCitationMarkdown(markdown); - const deserialized = editor - .getApi(MarkdownPlugin) - .markdown.deserialize(escapeMdxExpressions(rewritten)) as Descendant[]; + const deserialized = safeDeserializeMarkdown( + editor, + escapeMdxExpressions(rewritten) + ); newValue = injectCitationNodes(deserialized, urlMap); } else { - newValue = editor - .getApi(MarkdownPlugin) - .markdown.deserialize(escapeMdxExpressions(markdown)) as Descendant[]; + newValue = safeDeserializeMarkdown(editor, escapeMdxExpressions(markdown)); } editor.tf.reset(); editor.tf.setValue(newValue as Value); diff --git a/surfsense_web/components/editor/utils/safe-deserialize.ts b/surfsense_web/components/editor/utils/safe-deserialize.ts new file mode 100644 index 000000000..e359a7791 --- /dev/null +++ b/surfsense_web/components/editor/utils/safe-deserialize.ts @@ -0,0 +1,64 @@ +// --------------------------------------------------------------------------- +// Safe markdown deserialization for the Plate editor +// --------------------------------------------------------------------------- +// `remark-mdx` treats any HTML-like tag as JSX, so unbalanced inline HTML +// (very common in GitHub READMEs, web-scraped pages, PDF conversions) makes +// it throw "Expected a closing tag for ``" and crash the editor. +// +// Per the MDX maintainers' guidance (mdx-js/mdx, ipikuka/next-mdx-remote-client +// #14), MDX is the wrong format for untrusted markdown and the recommended +// fix is to fall back to plain markdown parsing. `MarkdownPlugin.deserialize` +// accepts a per-call `remarkPlugins` override, so we can: +// +// 1. Try with `remarkMdx` (rich MDX features, e.g. JSX-style components). +// 2. On failure, retry without `remarkMdx` (lenient HTML, like GitHub). +// 3. As a last resort, render the raw source in a paragraph so the user +// never sees a crashed editor. +// --------------------------------------------------------------------------- + +import { MarkdownPlugin, remarkMdx } from "@platejs/markdown"; +import type { Descendant } from "platejs"; +import remarkGfm from "remark-gfm"; +import remarkMath from "remark-math"; +import type { PlateEditorInstance } from "@/components/editor/plate-editor"; + +const STRICT_PLUGINS = [remarkGfm, remarkMath, remarkMdx]; +const LENIENT_PLUGINS = [remarkGfm, remarkMath]; + +function plainTextFallback(markdown: string): Descendant[] { + return [ + { + type: "p", + children: [{ text: markdown }], + } as unknown as Descendant, + ]; +} + +/** + * Deserialize markdown into a Plate value, gracefully degrading when the + * MDX-strict parser rejects raw HTML. Always returns a renderable value; + * never throws. + */ +export function safeDeserializeMarkdown( + editor: PlateEditorInstance, + markdown: string +): Descendant[] { + const api = editor.getApi(MarkdownPlugin).markdown; + + try { + return api.deserialize(markdown, { remarkPlugins: STRICT_PLUGINS }) as Descendant[]; + } catch (mdxError) { + if (process.env.NODE_ENV !== "production") { + console.warn( + "[plate-editor] MDX parse failed, retrying without remark-mdx:", + mdxError + ); + } + try { + return api.deserialize(markdown, { remarkPlugins: LENIENT_PLUGINS }) as Descendant[]; + } catch (fallbackError) { + console.error("[plate-editor] markdown deserialize failed:", fallbackError); + return plainTextFallback(markdown); + } + } +} From a2ad697a2990b55f292581b03c9f52427fd58e06 Mon Sep 17 00:00:00 2001 From: "DESKTOP-RTLN3BA\\$punk" Date: Tue, 5 May 2026 19:13:38 -0700 Subject: [PATCH 34/96] feat(next.config): enable remote SVG support with enhanced content security policy --- surfsense_web/next.config.ts | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/surfsense_web/next.config.ts b/surfsense_web/next.config.ts index 6cfcb5187..81f349f26 100644 --- a/surfsense_web/next.config.ts +++ b/surfsense_web/next.config.ts @@ -29,6 +29,13 @@ const nextConfig: NextConfig = { hostname: "**", }, ], + // Allow remote SVGs (e.g. README badges from img.shields.io, trendshift.io, + // etc.) which are otherwise blocked by next/image. The CSP below sandboxes + // the SVG and forbids any embedded scripts, which is the mitigation + // recommended by Vercel's NEXTJS_SAFE_SVG_IMAGES conformance rule. + dangerouslyAllowSVG: true, + contentDispositionType: "attachment", + contentSecurityPolicy: "default-src 'self'; script-src 'none'; sandbox;", }, experimental: { optimizePackageImports: [ From 499c6be0997c62b885b1ab61aab74796381f54c6 Mon Sep 17 00:00:00 2001 From: "DESKTOP-RTLN3BA\\$punk" Date: Tue, 5 May 2026 19:21:43 -0700 Subject: [PATCH 35/96] feat: bumped version to 0.0.23 --- VERSION | 2 +- surfsense_backend/pyproject.toml | 2 +- surfsense_backend/uv.lock | 2 +- surfsense_browser_extension/package.json | 2 +- surfsense_desktop/package.json | 2 +- surfsense_web/package.json | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/VERSION b/VERSION index 818944f5b..df5db66fe 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.0.22 +0.0.23 diff --git a/surfsense_backend/pyproject.toml b/surfsense_backend/pyproject.toml index 4235ac962..523a8a1ac 100644 --- a/surfsense_backend/pyproject.toml +++ b/surfsense_backend/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "surf-new-backend" -version = "0.0.22" +version = "0.0.23" description = "SurfSense Backend" requires-python = ">=3.12" dependencies = [ diff --git a/surfsense_backend/uv.lock b/surfsense_backend/uv.lock index 4dd5156e7..812be636a 100644 --- a/surfsense_backend/uv.lock +++ b/surfsense_backend/uv.lock @@ -7947,7 +7947,7 @@ wheels = [ [[package]] name = "surf-new-backend" -version = "0.0.22" +version = "0.0.23" source = { editable = "." } dependencies = [ { name = "alembic" }, diff --git a/surfsense_browser_extension/package.json b/surfsense_browser_extension/package.json index b8b5cb2ec..82c0a349a 100644 --- a/surfsense_browser_extension/package.json +++ b/surfsense_browser_extension/package.json @@ -1,7 +1,7 @@ { "name": "surfsense_browser_extension", "displayName": "Surfsense Browser Extension", - "version": "0.0.22", + "version": "0.0.23", "description": "Extension to collect Browsing History for SurfSense.", "author": "https://github.com/MODSetter", "engines": { diff --git a/surfsense_desktop/package.json b/surfsense_desktop/package.json index 744ab65ab..4ef624760 100644 --- a/surfsense_desktop/package.json +++ b/surfsense_desktop/package.json @@ -1,6 +1,6 @@ { "name": "surfsense-desktop", - "version": "0.0.22", + "version": "0.0.23", "description": "SurfSense Desktop App", "main": "dist/main.js", "scripts": { diff --git a/surfsense_web/package.json b/surfsense_web/package.json index 2adec8638..782409c3c 100644 --- a/surfsense_web/package.json +++ b/surfsense_web/package.json @@ -1,6 +1,6 @@ { "name": "surfsense_web", - "version": "0.0.22", + "version": "0.0.23", "private": true, "description": "SurfSense Frontend", "scripts": { From a9bf7ab7d23d27141da245fbd85d773071203940 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:47 +0200 Subject: [PATCH 36/96] Add SSE envelope helpers under app.services.streaming. --- .../services/streaming/envelope/__init__.py | 23 +++++++++++++++++ .../streaming/envelope/identifiers.py | 25 +++++++++++++++++++ .../app/services/streaming/envelope/sse.py | 25 +++++++++++++++++++ 3 files changed, 73 insertions(+) create mode 100644 surfsense_backend/app/services/streaming/envelope/__init__.py create mode 100644 surfsense_backend/app/services/streaming/envelope/identifiers.py create mode 100644 surfsense_backend/app/services/streaming/envelope/sse.py diff --git a/surfsense_backend/app/services/streaming/envelope/__init__.py b/surfsense_backend/app/services/streaming/envelope/__init__.py new file mode 100644 index 000000000..862e84c8d --- /dev/null +++ b/surfsense_backend/app/services/streaming/envelope/__init__.py @@ -0,0 +1,23 @@ +"""Wire framing layer.""" + +from __future__ import annotations + +from .identifiers import ( + generate_message_id, + generate_reasoning_id, + generate_subagent_run_id, + generate_text_id, + generate_tool_call_id, +) +from .sse import format_done, format_sse, get_response_headers + +__all__ = [ + "format_done", + "format_sse", + "generate_message_id", + "generate_reasoning_id", + "generate_subagent_run_id", + "generate_text_id", + "generate_tool_call_id", + "get_response_headers", +] diff --git a/surfsense_backend/app/services/streaming/envelope/identifiers.py b/surfsense_backend/app/services/streaming/envelope/identifiers.py new file mode 100644 index 000000000..2fdd6ff09 --- /dev/null +++ b/surfsense_backend/app/services/streaming/envelope/identifiers.py @@ -0,0 +1,25 @@ +"""Prefixed UUID generators for stream parts.""" + +from __future__ import annotations + +import uuid + + +def generate_message_id() -> str: + return f"msg_{uuid.uuid4().hex}" + + +def generate_text_id() -> str: + return f"text_{uuid.uuid4().hex}" + + +def generate_reasoning_id() -> str: + return f"reasoning_{uuid.uuid4().hex}" + + +def generate_tool_call_id() -> str: + return f"call_{uuid.uuid4().hex}" + + +def generate_subagent_run_id() -> str: + return f"subagent_{uuid.uuid4().hex}" diff --git a/surfsense_backend/app/services/streaming/envelope/sse.py b/surfsense_backend/app/services/streaming/envelope/sse.py new file mode 100644 index 000000000..508fc1b1c --- /dev/null +++ b/surfsense_backend/app/services/streaming/envelope/sse.py @@ -0,0 +1,25 @@ +"""Server-Sent-Events wire framing.""" + +from __future__ import annotations + +import json +from typing import Any + + +def format_sse(data: Any) -> str: + if isinstance(data, str): + return f"data: {data}\n\n" + return f"data: {json.dumps(data)}\n\n" + + +def format_done() -> str: + return "data: [DONE]\n\n" + + +def get_response_headers() -> dict[str, str]: + return { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "x-vercel-ai-ui-message-stream": "v1", + } From 5510c6c3147e80a4de1b6a25e58bc6479032482e Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:47 +0200 Subject: [PATCH 37/96] Add typed event payload modules for the streaming service. --- .../app/services/streaming/events/__init__.py | 29 +++++ .../services/streaming/events/action_log.py | 24 ++++ .../app/services/streaming/events/data.py | 118 ++++++++++++++++++ .../app/services/streaming/events/error.py | 23 ++++ .../services/streaming/events/interrupt.py | 56 +++++++++ .../services/streaming/events/lifecycle.py | 29 +++++ .../services/streaming/events/reasoning.py | 36 ++++++ .../app/services/streaming/events/source.py | 59 +++++++++ .../streaming/events/subagent_lifecycle.py | 86 +++++++++++++ .../app/services/streaming/events/text.py | 31 +++++ .../app/services/streaming/events/tool.py | 80 ++++++++++++ 11 files changed, 571 insertions(+) create mode 100644 surfsense_backend/app/services/streaming/events/__init__.py create mode 100644 surfsense_backend/app/services/streaming/events/action_log.py create mode 100644 surfsense_backend/app/services/streaming/events/data.py create mode 100644 surfsense_backend/app/services/streaming/events/error.py create mode 100644 surfsense_backend/app/services/streaming/events/interrupt.py create mode 100644 surfsense_backend/app/services/streaming/events/lifecycle.py create mode 100644 surfsense_backend/app/services/streaming/events/reasoning.py create mode 100644 surfsense_backend/app/services/streaming/events/source.py create mode 100644 surfsense_backend/app/services/streaming/events/subagent_lifecycle.py create mode 100644 surfsense_backend/app/services/streaming/events/text.py create mode 100644 surfsense_backend/app/services/streaming/events/tool.py diff --git a/surfsense_backend/app/services/streaming/events/__init__.py b/surfsense_backend/app/services/streaming/events/__init__.py new file mode 100644 index 000000000..91a8ff854 --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/__init__.py @@ -0,0 +1,29 @@ +"""SSE event payload formatters, one module per event family.""" + +from __future__ import annotations + +from . import ( + action_log, + data, + error, + interrupt, + lifecycle, + reasoning, + source, + subagent_lifecycle, + text, + tool, +) + +__all__ = [ + "action_log", + "data", + "error", + "interrupt", + "lifecycle", + "reasoning", + "source", + "subagent_lifecycle", + "text", + "tool", +] diff --git a/surfsense_backend/app/services/streaming/events/action_log.py b/surfsense_backend/app/services/streaming/events/action_log.py new file mode 100644 index 000000000..0a8e46f0a --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/action_log.py @@ -0,0 +1,24 @@ +"""Action-log events relayed from ``ActionLogMiddleware`` custom dispatches.""" + +from __future__ import annotations + +from typing import Any + +from ..emitter import Emitter +from .data import format_data + + +def format_action_log( + payload: dict[str, Any], + *, + emitter: Emitter | None = None, +) -> str: + return format_data("action-log", payload, emitter=emitter) + + +def format_action_log_updated( + payload: dict[str, Any], + *, + emitter: Emitter | None = None, +) -> str: + return format_data("action-log-updated", payload, emitter=emitter) diff --git a/surfsense_backend/app/services/streaming/events/data.py b/surfsense_backend/app/services/streaming/events/data.py new file mode 100644 index 000000000..f6e190578 --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/data.py @@ -0,0 +1,118 @@ +"""Generic ``data-*`` envelopes and SurfSense-specific data parts. + +Inner ``data`` dict fields use snake_case. Legacy ``threadId`` / +``messageId`` keys are preserved where they cross the AI SDK boundary. +""" + +from __future__ import annotations + +from typing import Any + +from ..emitter import Emitter, attach_emitted_by +from ..envelope import format_sse + + +def format_data( + data_type: str, + data: Any, + *, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = {"type": f"data-{data_type}", "data": data} + return format_sse(attach_emitted_by(payload, emitter)) + + +def format_terminal_info( + text: str, + *, + message_type: str = "info", + emitter: Emitter | None = None, +) -> str: + return format_data( + "terminal-info", + {"text": text, "type": message_type}, + emitter=emitter, + ) + + +def format_further_questions( + questions: list[str], + *, + emitter: Emitter | None = None, +) -> str: + return format_data("further-questions", {"questions": questions}, emitter=emitter) + + +def format_thinking_step( + *, + step_id: str, + title: str, + status: str = "in_progress", + items: list[str] | None = None, + emitter: Emitter | None = None, +) -> str: + return format_data( + "thinking-step", + { + "id": step_id, + "title": title, + "status": status, + "items": items or [], + }, + emitter=emitter, + ) + + +def format_thread_title_update( + *, + thread_id: int, + title: str, + emitter: Emitter | None = None, +) -> str: + return format_data( + "thread-title-update", + {"threadId": thread_id, "title": title}, + emitter=emitter, + ) + + +def format_turn_info( + *, + chat_turn_id: str, + emitter: Emitter | None = None, +) -> str: + return format_data("turn-info", {"chat_turn_id": chat_turn_id}, emitter=emitter) + + +def format_turn_status( + *, + status: str, + emitter: Emitter | None = None, +) -> str: + return format_data("turn-status", {"status": status}, emitter=emitter) + + +def format_user_message_id( + *, + message_id: str, + turn_id: str, + emitter: Emitter | None = None, +) -> str: + return format_data( + "user-message-id", + {"message_id": message_id, "turn_id": turn_id}, + emitter=emitter, + ) + + +def format_assistant_message_id( + *, + message_id: str, + turn_id: str, + emitter: Emitter | None = None, +) -> str: + return format_data( + "assistant-message-id", + {"message_id": message_id, "turn_id": turn_id}, + emitter=emitter, + ) diff --git a/surfsense_backend/app/services/streaming/events/error.py b/surfsense_backend/app/services/streaming/events/error.py new file mode 100644 index 000000000..cd190d1f4 --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/error.py @@ -0,0 +1,23 @@ +"""Single terminal error path the orchestrator must route through.""" + +from __future__ import annotations + +from typing import Any + +from ..emitter import Emitter, attach_emitted_by +from ..envelope import format_sse + + +def format_error( + error_text: str, + *, + error_code: str | None = None, + extra: dict[str, Any] | None = None, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = {"type": "error", "errorText": error_text} + if error_code: + payload["errorCode"] = error_code + if extra: + payload.update(extra) + return format_sse(attach_emitted_by(payload, emitter)) diff --git a/surfsense_backend/app/services/streaming/events/interrupt.py b/surfsense_backend/app/services/streaming/events/interrupt.py new file mode 100644 index 000000000..0334b10b3 --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/interrupt.py @@ -0,0 +1,56 @@ +"""Interrupt-request events with a single canonical payload shape.""" + +from __future__ import annotations + +from typing import Any + +from ..emitter import Emitter +from .data import format_data + + +def normalize_interrupt_payload(interrupt_value: dict[str, Any]) -> dict[str, Any]: + if "action_requests" in interrupt_value and "review_configs" in interrupt_value: + return interrupt_value + + interrupt_type = interrupt_value.get("type", "unknown") + message = interrupt_value.get("message") + action = interrupt_value.get("action", {}) or {} + context = interrupt_value.get("context", {}) or {} + + normalized: dict[str, Any] = { + "action_requests": [ + { + "name": action.get("tool", "unknown_tool"), + "args": action.get("params", {}), + } + ], + "review_configs": [ + { + "action_name": action.get("tool", "unknown_tool"), + "allowed_decisions": ["approve", "edit", "reject"], + } + ], + "interrupt_type": interrupt_type, + "context": context, + } + if message: + normalized["message"] = message + return normalized + + +def format_interrupt_request( + interrupt_value: dict[str, Any], + *, + interrupt_id: str | None = None, + pending_interrupt_count: int | None = None, + chat_turn_id: str | None = None, + emitter: Emitter | None = None, +) -> str: + payload = normalize_interrupt_payload(interrupt_value) + if interrupt_id is not None: + payload["interrupt_id"] = interrupt_id + if pending_interrupt_count is not None: + payload["pending_interrupt_count"] = pending_interrupt_count + if chat_turn_id is not None: + payload["chat_turn_id"] = chat_turn_id + return format_data("interrupt-request", payload, emitter=emitter) diff --git a/surfsense_backend/app/services/streaming/events/lifecycle.py b/surfsense_backend/app/services/streaming/events/lifecycle.py new file mode 100644 index 000000000..019718b67 --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/lifecycle.py @@ -0,0 +1,29 @@ +"""High-level message and step lifecycle events. + +Wire verbs are fixed by the AI SDK protocol (``start`` / ``finish`` for +the whole message, ``start-step`` / ``finish-step`` for each step). +Python helpers always read ``format__`` so pairs are +visible at the call site. +""" + +from __future__ import annotations + +from ..emitter import Emitter, attach_emitted_by +from ..envelope import format_sse + + +def format_message_start(message_id: str, *, emitter: Emitter | None = None) -> str: + payload = {"type": "start", "messageId": message_id} + return format_sse(attach_emitted_by(payload, emitter)) + + +def format_message_finish(*, emitter: Emitter | None = None) -> str: + return format_sse(attach_emitted_by({"type": "finish"}, emitter)) + + +def format_step_start(*, emitter: Emitter | None = None) -> str: + return format_sse(attach_emitted_by({"type": "start-step"}, emitter)) + + +def format_step_finish(*, emitter: Emitter | None = None) -> str: + return format_sse(attach_emitted_by({"type": "finish-step"}, emitter)) diff --git a/surfsense_backend/app/services/streaming/events/reasoning.py b/surfsense_backend/app/services/streaming/events/reasoning.py new file mode 100644 index 000000000..5b912d43a --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/reasoning.py @@ -0,0 +1,36 @@ +"""Reasoning-block streaming events.""" + +from __future__ import annotations + +from ..emitter import Emitter, attach_emitted_by +from ..envelope import format_sse + + +def format_reasoning_start( + reasoning_id: str, *, emitter: Emitter | None = None +) -> str: + return format_sse( + attach_emitted_by({"type": "reasoning-start", "id": reasoning_id}, emitter) + ) + + +def format_reasoning_delta( + reasoning_id: str, + delta: str, + *, + emitter: Emitter | None = None, +) -> str: + return format_sse( + attach_emitted_by( + {"type": "reasoning-delta", "id": reasoning_id, "delta": delta}, + emitter, + ) + ) + + +def format_reasoning_end( + reasoning_id: str, *, emitter: Emitter | None = None +) -> str: + return format_sse( + attach_emitted_by({"type": "reasoning-end", "id": reasoning_id}, emitter) + ) diff --git a/surfsense_backend/app/services/streaming/events/source.py b/surfsense_backend/app/services/streaming/events/source.py new file mode 100644 index 000000000..54541e8d2 --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/source.py @@ -0,0 +1,59 @@ +"""Source and file reference events.""" + +from __future__ import annotations + +from typing import Any + +from ..emitter import Emitter, attach_emitted_by +from ..envelope import format_sse + + +def format_source_url( + url: str, + *, + source_id: str | None = None, + title: str | None = None, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = { + "type": "source-url", + "sourceId": source_id or url, + "url": url, + } + if title: + payload["title"] = title + return format_sse(attach_emitted_by(payload, emitter)) + + +def format_source_document( + source_id: str, + *, + media_type: str = "file", + title: str | None = None, + description: str | None = None, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = { + "type": "source-document", + "sourceId": source_id, + "mediaType": media_type, + } + if title: + payload["title"] = title + if description: + payload["description"] = description + return format_sse(attach_emitted_by(payload, emitter)) + + +def format_file( + url: str, + media_type: str, + *, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = { + "type": "file", + "url": url, + "mediaType": media_type, + } + return format_sse(attach_emitted_by(payload, emitter)) diff --git a/surfsense_backend/app/services/streaming/events/subagent_lifecycle.py b/surfsense_backend/app/services/streaming/events/subagent_lifecycle.py new file mode 100644 index 000000000..6dd2d4eab --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/subagent_lifecycle.py @@ -0,0 +1,86 @@ +"""Sub-agent lifecycle events the FE pairs into one timeline lane. + +A sub-agent run is a high-level boundary (a whole agent invocation), +so we use the ``start`` / ``finish`` verb pair, matching how the AI SDK +spells message- and step-level lifecycles. +""" + +from __future__ import annotations + +from typing import Any + +from ..emitter import Emitter +from .data import format_data + + +def format_subagent_start( + *, + subagent_run_id: str, + subagent_type: str, + parent_tool_call_id: str, + chat_turn_id: str | None = None, + description: str | None = None, + started_at: str | None = None, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = { + "subagent_run_id": subagent_run_id, + "subagent_type": subagent_type, + "parent_tool_call_id": parent_tool_call_id, + } + if chat_turn_id is not None: + payload["chat_turn_id"] = chat_turn_id + if description is not None: + payload["description"] = description + if started_at is not None: + payload["started_at"] = started_at + return format_data("subagent-start", payload, emitter=emitter) + + +def format_subagent_finish( + *, + subagent_run_id: str, + subagent_type: str, + parent_tool_call_id: str, + status: str = "completed", + ended_at: str | None = None, + duration_ms: int | None = None, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = { + "subagent_run_id": subagent_run_id, + "subagent_type": subagent_type, + "parent_tool_call_id": parent_tool_call_id, + "status": status, + } + if ended_at is not None: + payload["ended_at"] = ended_at + if duration_ms is not None: + payload["duration_ms"] = duration_ms + return format_data("subagent-finish", payload, emitter=emitter) + + +def format_subagent_error( + *, + subagent_run_id: str, + subagent_type: str, + parent_tool_call_id: str, + error_text: str, + error_type: str | None = None, + ended_at: str | None = None, + duration_ms: int | None = None, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = { + "subagent_run_id": subagent_run_id, + "subagent_type": subagent_type, + "parent_tool_call_id": parent_tool_call_id, + "error_text": error_text, + } + if error_type is not None: + payload["error_type"] = error_type + if ended_at is not None: + payload["ended_at"] = ended_at + if duration_ms is not None: + payload["duration_ms"] = duration_ms + return format_data("subagent-error", payload, emitter=emitter) diff --git a/surfsense_backend/app/services/streaming/events/text.py b/surfsense_backend/app/services/streaming/events/text.py new file mode 100644 index 000000000..3baebdebb --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/text.py @@ -0,0 +1,31 @@ +"""Text-block streaming events.""" + +from __future__ import annotations + +from ..emitter import Emitter, attach_emitted_by +from ..envelope import format_sse + + +def format_text_start(text_id: str, *, emitter: Emitter | None = None) -> str: + return format_sse( + attach_emitted_by({"type": "text-start", "id": text_id}, emitter) + ) + + +def format_text_delta( + text_id: str, + delta: str, + *, + emitter: Emitter | None = None, +) -> str: + return format_sse( + attach_emitted_by( + {"type": "text-delta", "id": text_id, "delta": delta}, emitter + ) + ) + + +def format_text_end(text_id: str, *, emitter: Emitter | None = None) -> str: + return format_sse( + attach_emitted_by({"type": "text-end", "id": text_id}, emitter) + ) diff --git a/surfsense_backend/app/services/streaming/events/tool.py b/surfsense_backend/app/services/streaming/events/tool.py new file mode 100644 index 000000000..c85dc061b --- /dev/null +++ b/surfsense_backend/app/services/streaming/events/tool.py @@ -0,0 +1,80 @@ +"""Tool-call streaming events. + +``toolCallId`` and ``langchainToolCallId`` are AI SDK protocol fields +and stay camelCase. Sub-agent provenance rides on the snake_case +top-level ``emitted_by`` envelope added by :func:`attach_emitted_by`. +""" + +from __future__ import annotations + +from typing import Any + +from ..emitter import Emitter, attach_emitted_by +from ..envelope import format_sse + + +def format_tool_input_start( + tool_call_id: str, + tool_name: str, + *, + langchain_tool_call_id: str | None = None, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = { + "type": "tool-input-start", + "toolCallId": tool_call_id, + "toolName": tool_name, + } + if langchain_tool_call_id: + payload["langchainToolCallId"] = langchain_tool_call_id + return format_sse(attach_emitted_by(payload, emitter)) + + +def format_tool_input_delta( + tool_call_id: str, + input_text_delta: str, + *, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = { + "type": "tool-input-delta", + "toolCallId": tool_call_id, + "inputTextDelta": input_text_delta, + } + return format_sse(attach_emitted_by(payload, emitter)) + + +def format_tool_input_available( + tool_call_id: str, + tool_name: str, + input_data: dict[str, Any], + *, + langchain_tool_call_id: str | None = None, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = { + "type": "tool-input-available", + "toolCallId": tool_call_id, + "toolName": tool_name, + "input": input_data, + } + if langchain_tool_call_id: + payload["langchainToolCallId"] = langchain_tool_call_id + return format_sse(attach_emitted_by(payload, emitter)) + + +def format_tool_output_available( + tool_call_id: str, + output: Any, + *, + langchain_tool_call_id: str | None = None, + emitter: Emitter | None = None, +) -> str: + payload: dict[str, Any] = { + "type": "tool-output-available", + "toolCallId": tool_call_id, + "output": output, + } + if langchain_tool_call_id: + payload["langchainToolCallId"] = langchain_tool_call_id + return format_sse(attach_emitted_by(payload, emitter)) From fc429d87024a6a0a36d520f23dd584bcf7bd8262 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:47 +0200 Subject: [PATCH 38/96] Add streaming emitter and registry for scoped SSE writes. --- .../services/streaming/emitter/__init__.py | 29 +++++++++ .../app/services/streaming/emitter/emitter.py | 61 +++++++++++++++++++ .../services/streaming/emitter/registry.py | 51 ++++++++++++++++ 3 files changed, 141 insertions(+) create mode 100644 surfsense_backend/app/services/streaming/emitter/__init__.py create mode 100644 surfsense_backend/app/services/streaming/emitter/emitter.py create mode 100644 surfsense_backend/app/services/streaming/emitter/registry.py diff --git a/surfsense_backend/app/services/streaming/emitter/__init__.py b/surfsense_backend/app/services/streaming/emitter/__init__.py new file mode 100644 index 000000000..7814894f3 --- /dev/null +++ b/surfsense_backend/app/services/streaming/emitter/__init__.py @@ -0,0 +1,29 @@ +"""Identity of the agent that emitted a streamed event. + +The wire field is ``emitted_by``; the Python identity is :class:`Emitter`. +``EmitterRegistry`` resolves which emitter owns a LangGraph event, with +LangGraph's own namespace metadata as the primary key and a parent_ids +walk as a fallback for cases where context vars don't propagate. +""" + +from __future__ import annotations + +from .emitter import ( + MAIN_EMITTER, + Emitter, + EmitterLevel, + attach_emitted_by, + main_emitter, + subagent_emitter, +) +from .registry import EmitterRegistry + +__all__ = [ + "MAIN_EMITTER", + "Emitter", + "EmitterLevel", + "EmitterRegistry", + "attach_emitted_by", + "main_emitter", + "subagent_emitter", +] diff --git a/surfsense_backend/app/services/streaming/emitter/emitter.py b/surfsense_backend/app/services/streaming/emitter/emitter.py new file mode 100644 index 000000000..08f625a69 --- /dev/null +++ b/surfsense_backend/app/services/streaming/emitter/emitter.py @@ -0,0 +1,61 @@ +"""Identity payload describing which agent produced a stream event.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal + +EmitterLevel = Literal["main", "subagent"] + + +@dataclass(frozen=True) +class Emitter: + level: EmitterLevel + subagent_type: str | None = None + subagent_run_id: str | None = None + parent_tool_call_id: str | None = None + extra: dict[str, Any] = field(default_factory=dict) + + def to_payload(self) -> dict[str, Any]: + payload: dict[str, Any] = {"level": self.level} + if self.subagent_type is not None: + payload["subagent_type"] = self.subagent_type + if self.subagent_run_id is not None: + payload["subagent_run_id"] = self.subagent_run_id + if self.parent_tool_call_id is not None: + payload["parent_tool_call_id"] = self.parent_tool_call_id + if self.extra: + payload.update(self.extra) + return payload + + +MAIN_EMITTER = Emitter(level="main") + + +def main_emitter() -> Emitter: + return MAIN_EMITTER + + +def subagent_emitter( + *, + subagent_type: str, + subagent_run_id: str, + parent_tool_call_id: str | None = None, + extra: dict[str, Any] | None = None, +) -> Emitter: + return Emitter( + level="subagent", + subagent_type=subagent_type, + subagent_run_id=subagent_run_id, + parent_tool_call_id=parent_tool_call_id, + extra=dict(extra or {}), + ) + + +def attach_emitted_by( + payload: dict[str, Any], emitter: Emitter | None +) -> dict[str, Any]: + if emitter is None: + return payload + payload["emitted_by"] = emitter.to_payload() + return payload diff --git a/surfsense_backend/app/services/streaming/emitter/registry.py b/surfsense_backend/app/services/streaming/emitter/registry.py new file mode 100644 index 000000000..cd3e10cdd --- /dev/null +++ b/surfsense_backend/app/services/streaming/emitter/registry.py @@ -0,0 +1,51 @@ +"""Resolve which agent owns a streamed event from its LangGraph run lineage.""" + +from __future__ import annotations + +from collections.abc import Iterable + +from .emitter import Emitter, main_emitter + + +class EmitterRegistry: + def __init__(self) -> None: + self._by_run_id: dict[str, Emitter] = {} + + def register(self, run_id: str, emitter: Emitter) -> None: + if not run_id: + return + self._by_run_id[run_id] = emitter + + def unregister(self, run_id: str) -> Emitter | None: + if not run_id: + return None + return self._by_run_id.pop(run_id, None) + + def get(self, run_id: str | None) -> Emitter | None: + if not run_id: + return None + return self._by_run_id.get(run_id) + + def resolve( + self, + *, + run_id: str | None, + parent_ids: Iterable[str] | None, + ) -> Emitter: + own = self.get(run_id) + if own is not None: + return own + if parent_ids: + for ancestor in reversed(list(parent_ids)): + emitter = self.get(ancestor) + if emitter is not None: + return emitter + return main_emitter() + + def has_active_subagents(self) -> bool: + return any( + emitter.level == "subagent" for emitter in self._by_run_id.values() + ) + + def clear(self) -> None: + self._by_run_id.clear() From fef7621d96c63986e723228c1dcff4a918b443e9 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:47 +0200 Subject: [PATCH 39/96] Add StreamingService and interrupt correlation for chat streams. --- .../app/services/streaming/__init__.py | 20 + .../streaming/interrupt_correlation.py | 84 ++++ .../app/services/streaming/service.py | 414 ++++++++++++++++++ 3 files changed, 518 insertions(+) create mode 100644 surfsense_backend/app/services/streaming/__init__.py create mode 100644 surfsense_backend/app/services/streaming/interrupt_correlation.py create mode 100644 surfsense_backend/app/services/streaming/service.py diff --git a/surfsense_backend/app/services/streaming/__init__.py b/surfsense_backend/app/services/streaming/__init__.py new file mode 100644 index 000000000..287d48a7a --- /dev/null +++ b/surfsense_backend/app/services/streaming/__init__.py @@ -0,0 +1,20 @@ +"""Single-responsibility split of the streaming SSE protocol. + +Layout: +* ``envelope/`` - SSE wire framing + ID generators +* ``emitter/`` - identity of the agent that emitted an event + runtime registry +* ``events/`` - one module per SSE event family +* ``service.py`` - composition root used by the orchestrator +* ``interrupt_correlation.py`` - id-aware lookup over LangGraph state + +Naming on the wire: +* AI SDK protocol fields keep their existing camelCase + (``toolCallId``, ``messageId``, ``inputTextDelta``, ``langchainToolCallId``). +* Every SurfSense-added field uses ``snake_case``, including the + top-level ``emitted_by`` envelope and all inner ``data`` payloads. + +Production keeps using ``app.services.new_streaming_service`` and +``app.tasks.chat.stream_new_chat`` until the cutover phase. +""" + +from __future__ import annotations diff --git a/surfsense_backend/app/services/streaming/interrupt_correlation.py b/surfsense_backend/app/services/streaming/interrupt_correlation.py new file mode 100644 index 000000000..3045dfb6a --- /dev/null +++ b/surfsense_backend/app/services/streaming/interrupt_correlation.py @@ -0,0 +1,84 @@ +"""Id-aware lookup of pending LangGraph interrupts (replaces first-wins).""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class PendingInterrupt: + interrupt_id: str | None + value: dict[str, Any] + source_task_id: str | None = None + + +def list_pending_interrupts(state: Any) -> list[PendingInterrupt]: + out: list[PendingInterrupt] = [] + + for task in getattr(state, "tasks", None) or (): + task_id = _safe_str(getattr(task, "id", None)) + for it in getattr(task, "interrupts", None) or (): + value = _coerce_interrupt_value(it) + if value is None: + continue + interrupt_id = _safe_str(getattr(it, "id", None)) + out.append( + PendingInterrupt( + interrupt_id=interrupt_id, + value=value, + source_task_id=task_id, + ) + ) + + for it in getattr(state, "interrupts", None) or (): + value = _coerce_interrupt_value(it) + if value is None: + continue + interrupt_id = _safe_str(getattr(it, "id", None)) + out.append(PendingInterrupt(interrupt_id=interrupt_id, value=value)) + + return out + + +def get_pending_interrupt_by_id( + state: Any, interrupt_id: str +) -> PendingInterrupt | None: + for pending in list_pending_interrupts(state): + if pending.interrupt_id == interrupt_id: + return pending + return None + + +def get_pending_interrupt_for_tool_call( + state: Any, tool_call_id: str +) -> PendingInterrupt | None: + for pending in list_pending_interrupts(state): + actions = pending.value.get("action_requests") + if not isinstance(actions, list): + continue + for action in actions: + if not isinstance(action, dict): + continue + if action.get("tool_call_id") == tool_call_id: + return pending + return None + + +def first_pending_interrupt(state: Any) -> PendingInterrupt | None: + """Explicit opt-in to legacy first-wins; prefer the id-aware helpers above.""" + pending = list_pending_interrupts(state) + return pending[0] if pending else None + + +def _coerce_interrupt_value(item: Any) -> dict[str, Any] | None: + if isinstance(item, dict): + return item if item else None + value = getattr(item, "value", None) + if isinstance(value, dict): + return value if value else None + return None + + +def _safe_str(value: Any) -> str | None: + return value if isinstance(value, str) and value else None diff --git a/surfsense_backend/app/services/streaming/service.py b/surfsense_backend/app/services/streaming/service.py new file mode 100644 index 000000000..5a75a1b2d --- /dev/null +++ b/surfsense_backend/app/services/streaming/service.py @@ -0,0 +1,414 @@ +"""Composition root: bundles every formatter + a per-invocation emitter registry.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +from . import envelope +from .emitter import Emitter, EmitterRegistry +from .events import ( + action_log, + data, + error, + interrupt, + lifecycle, + reasoning, + source, + subagent_lifecycle, + text, + tool, +) + + +class StreamingService: + def __init__(self) -> None: + self._message_id: str | None = None + self.emitter_registry = EmitterRegistry() + + @property + def message_id(self) -> str | None: + return self._message_id + + def begin_message(self, message_id: str | None = None) -> str: + self._message_id = message_id or envelope.generate_message_id() + return self._message_id + + @staticmethod + def generate_text_id() -> str: + return envelope.generate_text_id() + + @staticmethod + def generate_reasoning_id() -> str: + return envelope.generate_reasoning_id() + + @staticmethod + def generate_tool_call_id() -> str: + return envelope.generate_tool_call_id() + + @staticmethod + def generate_subagent_run_id() -> str: + return envelope.generate_subagent_run_id() + + @staticmethod + def get_response_headers() -> dict[str, str]: + return envelope.get_response_headers() + + @staticmethod + def format_done() -> str: + return envelope.format_done() + + def resolve_emitter( + self, + *, + run_id: str | None, + parent_ids: Iterable[str] | None, + ) -> Emitter: + return self.emitter_registry.resolve(run_id=run_id, parent_ids=parent_ids) + + def format_message_start( + self, + message_id: str | None = None, + *, + emitter: Emitter | None = None, + ) -> str: + chosen = self.begin_message(message_id) + return lifecycle.format_message_start(chosen, emitter=emitter) + + def format_message_finish(self, *, emitter: Emitter | None = None) -> str: + return lifecycle.format_message_finish(emitter=emitter) + + def format_step_start(self, *, emitter: Emitter | None = None) -> str: + return lifecycle.format_step_start(emitter=emitter) + + def format_step_finish(self, *, emitter: Emitter | None = None) -> str: + return lifecycle.format_step_finish(emitter=emitter) + + def format_text_start( + self, text_id: str, *, emitter: Emitter | None = None + ) -> str: + return text.format_text_start(text_id, emitter=emitter) + + def format_text_delta( + self, text_id: str, delta: str, *, emitter: Emitter | None = None + ) -> str: + return text.format_text_delta(text_id, delta, emitter=emitter) + + def format_text_end( + self, text_id: str, *, emitter: Emitter | None = None + ) -> str: + return text.format_text_end(text_id, emitter=emitter) + + def format_reasoning_start( + self, reasoning_id: str, *, emitter: Emitter | None = None + ) -> str: + return reasoning.format_reasoning_start(reasoning_id, emitter=emitter) + + def format_reasoning_delta( + self, + reasoning_id: str, + delta: str, + *, + emitter: Emitter | None = None, + ) -> str: + return reasoning.format_reasoning_delta(reasoning_id, delta, emitter=emitter) + + def format_reasoning_end( + self, reasoning_id: str, *, emitter: Emitter | None = None + ) -> str: + return reasoning.format_reasoning_end(reasoning_id, emitter=emitter) + + def format_tool_input_start( + self, + tool_call_id: str, + tool_name: str, + *, + langchain_tool_call_id: str | None = None, + emitter: Emitter | None = None, + ) -> str: + return tool.format_tool_input_start( + tool_call_id, + tool_name, + langchain_tool_call_id=langchain_tool_call_id, + emitter=emitter, + ) + + def format_tool_input_delta( + self, + tool_call_id: str, + input_text_delta: str, + *, + emitter: Emitter | None = None, + ) -> str: + return tool.format_tool_input_delta( + tool_call_id, input_text_delta, emitter=emitter + ) + + def format_tool_input_available( + self, + tool_call_id: str, + tool_name: str, + input_data: dict[str, Any], + *, + langchain_tool_call_id: str | None = None, + emitter: Emitter | None = None, + ) -> str: + return tool.format_tool_input_available( + tool_call_id, + tool_name, + input_data, + langchain_tool_call_id=langchain_tool_call_id, + emitter=emitter, + ) + + def format_tool_output_available( + self, + tool_call_id: str, + output: Any, + *, + langchain_tool_call_id: str | None = None, + emitter: Emitter | None = None, + ) -> str: + return tool.format_tool_output_available( + tool_call_id, + output, + langchain_tool_call_id=langchain_tool_call_id, + emitter=emitter, + ) + + def format_source_url( + self, + url: str, + *, + source_id: str | None = None, + title: str | None = None, + emitter: Emitter | None = None, + ) -> str: + return source.format_source_url( + url, source_id=source_id, title=title, emitter=emitter + ) + + def format_source_document( + self, + source_id: str, + *, + media_type: str = "file", + title: str | None = None, + description: str | None = None, + emitter: Emitter | None = None, + ) -> str: + return source.format_source_document( + source_id, + media_type=media_type, + title=title, + description=description, + emitter=emitter, + ) + + def format_file( + self, url: str, media_type: str, *, emitter: Emitter | None = None + ) -> str: + return source.format_file(url, media_type, emitter=emitter) + + def format_data( + self, data_type: str, payload: Any, *, emitter: Emitter | None = None + ) -> str: + return data.format_data(data_type, payload, emitter=emitter) + + def format_terminal_info( + self, + text_value: str, + *, + message_type: str = "info", + emitter: Emitter | None = None, + ) -> str: + return data.format_terminal_info( + text_value, message_type=message_type, emitter=emitter + ) + + def format_further_questions( + self, + questions: list[str], + *, + emitter: Emitter | None = None, + ) -> str: + return data.format_further_questions(questions, emitter=emitter) + + def format_thinking_step( + self, + *, + step_id: str, + title: str, + status: str = "in_progress", + items: list[str] | None = None, + emitter: Emitter | None = None, + ) -> str: + return data.format_thinking_step( + step_id=step_id, + title=title, + status=status, + items=items, + emitter=emitter, + ) + + def format_thread_title_update( + self, + *, + thread_id: int, + title: str, + emitter: Emitter | None = None, + ) -> str: + return data.format_thread_title_update( + thread_id=thread_id, title=title, emitter=emitter + ) + + def format_turn_info( + self, + *, + chat_turn_id: str, + emitter: Emitter | None = None, + ) -> str: + return data.format_turn_info(chat_turn_id=chat_turn_id, emitter=emitter) + + def format_turn_status( + self, + *, + status: str, + emitter: Emitter | None = None, + ) -> str: + return data.format_turn_status(status=status, emitter=emitter) + + def format_user_message_id( + self, + *, + message_id: str, + turn_id: str, + emitter: Emitter | None = None, + ) -> str: + return data.format_user_message_id( + message_id=message_id, turn_id=turn_id, emitter=emitter + ) + + def format_assistant_message_id( + self, + *, + message_id: str, + turn_id: str, + emitter: Emitter | None = None, + ) -> str: + return data.format_assistant_message_id( + message_id=message_id, turn_id=turn_id, emitter=emitter + ) + + def format_error( + self, + error_text: str, + *, + error_code: str | None = None, + extra: dict[str, Any] | None = None, + emitter: Emitter | None = None, + ) -> str: + return error.format_error( + error_text, + error_code=error_code, + extra=extra, + emitter=emitter, + ) + + def format_interrupt_request( + self, + interrupt_value: dict[str, Any], + *, + interrupt_id: str | None = None, + pending_interrupt_count: int | None = None, + chat_turn_id: str | None = None, + emitter: Emitter | None = None, + ) -> str: + return interrupt.format_interrupt_request( + interrupt_value, + interrupt_id=interrupt_id, + pending_interrupt_count=pending_interrupt_count, + chat_turn_id=chat_turn_id, + emitter=emitter, + ) + + def format_subagent_start( + self, + *, + subagent_run_id: str, + subagent_type: str, + parent_tool_call_id: str, + chat_turn_id: str | None = None, + description: str | None = None, + started_at: str | None = None, + emitter: Emitter | None = None, + ) -> str: + return subagent_lifecycle.format_subagent_start( + subagent_run_id=subagent_run_id, + subagent_type=subagent_type, + parent_tool_call_id=parent_tool_call_id, + chat_turn_id=chat_turn_id, + description=description, + started_at=started_at, + emitter=emitter, + ) + + def format_subagent_finish( + self, + *, + subagent_run_id: str, + subagent_type: str, + parent_tool_call_id: str, + status: str = "completed", + ended_at: str | None = None, + duration_ms: int | None = None, + emitter: Emitter | None = None, + ) -> str: + return subagent_lifecycle.format_subagent_finish( + subagent_run_id=subagent_run_id, + subagent_type=subagent_type, + parent_tool_call_id=parent_tool_call_id, + status=status, + ended_at=ended_at, + duration_ms=duration_ms, + emitter=emitter, + ) + + def format_subagent_error( + self, + *, + subagent_run_id: str, + subagent_type: str, + parent_tool_call_id: str, + error_text: str, + error_type: str | None = None, + ended_at: str | None = None, + duration_ms: int | None = None, + emitter: Emitter | None = None, + ) -> str: + return subagent_lifecycle.format_subagent_error( + subagent_run_id=subagent_run_id, + subagent_type=subagent_type, + parent_tool_call_id=parent_tool_call_id, + error_text=error_text, + error_type=error_type, + ended_at=ended_at, + duration_ms=duration_ms, + emitter=emitter, + ) + + def format_action_log( + self, + payload: dict[str, Any], + *, + emitter: Emitter | None = None, + ) -> str: + return action_log.format_action_log(payload, emitter=emitter) + + def format_action_log_updated( + self, + payload: dict[str, Any], + *, + emitter: Emitter | None = None, + ) -> str: + return action_log.format_action_log_updated(payload, emitter=emitter) From 3d8c4be369bc887695522e3a18a553cd04954a68 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 40/96] Add unit tests for streaming SSE envelope behavior. --- .../services/streaming/test_sse_envelope.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 surfsense_backend/tests/unit/services/streaming/test_sse_envelope.py diff --git a/surfsense_backend/tests/unit/services/streaming/test_sse_envelope.py b/surfsense_backend/tests/unit/services/streaming/test_sse_envelope.py new file mode 100644 index 000000000..511e4575a --- /dev/null +++ b/surfsense_backend/tests/unit/services/streaming/test_sse_envelope.py @@ -0,0 +1,51 @@ +"""Pin the exact SSE wire bytes the FE parser depends on.""" + +from __future__ import annotations + +import json + +import pytest + +from app.services.streaming.envelope import ( + format_done, + format_sse, + get_response_headers, +) + +pytestmark = pytest.mark.unit + + +class TestFormatSse: + def test_dict_payload_is_json_serialised(self) -> None: + frame = format_sse({"type": "start", "messageId": "msg_1"}) + assert frame.startswith("data: ") + assert frame.endswith("\n\n") + body = frame[len("data: ") : -2] + assert json.loads(body) == {"type": "start", "messageId": "msg_1"} + + def test_string_payload_is_emitted_verbatim(self) -> None: + frame = format_sse('{"already":"json"}') + assert frame == 'data: {"already":"json"}\n\n' + + def test_nested_payload_round_trips(self) -> None: + payload = { + "type": "data-action-log", + "data": {"id": 7, "tool_name": "ls", "reversible": False}, + } + frame = format_sse(payload) + body = frame.removeprefix("data: ").removesuffix("\n\n") + assert json.loads(body) == payload + + +class TestFormatDone: + def test_done_marker_is_literal(self) -> None: + assert format_done() == "data: [DONE]\n\n" + + +class TestResponseHeaders: + def test_headers_pin_ai_sdk_v1_protocol(self) -> None: + headers = get_response_headers() + assert headers["Content-Type"] == "text/event-stream" + assert headers["Cache-Control"] == "no-cache" + assert headers["Connection"] == "keep-alive" + assert headers["x-vercel-ai-ui-message-stream"] == "v1" From 619a8362b7b45034601f35cda15ea349b0f7c701 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 41/96] Add unit tests for streaming emitters and registry wiring. --- .../unit/services/streaming/test_emitter.py | 79 +++++++++++++ .../streaming/test_emitter_registry.py | 111 ++++++++++++++++++ 2 files changed, 190 insertions(+) create mode 100644 surfsense_backend/tests/unit/services/streaming/test_emitter.py create mode 100644 surfsense_backend/tests/unit/services/streaming/test_emitter_registry.py diff --git a/surfsense_backend/tests/unit/services/streaming/test_emitter.py b/surfsense_backend/tests/unit/services/streaming/test_emitter.py new file mode 100644 index 000000000..6c4e1ff58 --- /dev/null +++ b/surfsense_backend/tests/unit/services/streaming/test_emitter.py @@ -0,0 +1,79 @@ +"""Pin the wire compactness rule and the top-level ``emitted_by`` field name.""" + +from __future__ import annotations + +import pytest + +from app.services.streaming.emitter import ( + Emitter, + attach_emitted_by, + main_emitter, + subagent_emitter, +) + +pytestmark = pytest.mark.unit + + +def test_main_emitter_payload_contains_only_level() -> None: + payload = main_emitter().to_payload() + assert payload == {"level": "main"} + + +def test_subagent_emitter_payload_includes_all_set_fields() -> None: + payload = subagent_emitter( + subagent_type="deliverables", + subagent_run_id="subagent_abc", + parent_tool_call_id="call_xyz", + ).to_payload() + assert payload == { + "level": "subagent", + "subagent_type": "deliverables", + "subagent_run_id": "subagent_abc", + "parent_tool_call_id": "call_xyz", + } + + +def test_subagent_emitter_payload_omits_unset_optional_fields() -> None: + """parent_tool_call_id is None when the run is started outside a tool boundary.""" + payload = Emitter( + level="subagent", + subagent_type="email", + subagent_run_id="subagent_1", + ).to_payload() + assert "parent_tool_call_id" not in payload + assert payload["subagent_type"] == "email" + + +def test_extra_fields_merge_into_payload() -> None: + """Future extension fields (e.g. lane colour, label) flow through ``extra``.""" + emitter = subagent_emitter( + subagent_type="search", + subagent_run_id="r1", + extra={"label": "Web Search"}, + ) + assert emitter.to_payload()["label"] == "Web Search" + + +def test_attach_emitted_by_with_none_is_noop() -> None: + payload = {"type": "text-delta", "delta": "hi"} + result = attach_emitted_by(payload, None) + assert "emitted_by" not in result + assert result is payload + + +def test_attach_emitted_by_adds_payload_under_snake_case_top_level_key() -> None: + payload = {"type": "text-delta", "delta": "hi"} + attach_emitted_by( + payload, + subagent_emitter( + subagent_type="x", + subagent_run_id="y", + parent_tool_call_id="z", + ), + ) + assert payload["emitted_by"] == { + "level": "subagent", + "subagent_type": "x", + "subagent_run_id": "y", + "parent_tool_call_id": "z", + } diff --git a/surfsense_backend/tests/unit/services/streaming/test_emitter_registry.py b/surfsense_backend/tests/unit/services/streaming/test_emitter_registry.py new file mode 100644 index 000000000..e459c946a --- /dev/null +++ b/surfsense_backend/tests/unit/services/streaming/test_emitter_registry.py @@ -0,0 +1,111 @@ +"""Pin the parent_ids walk + parallel sub-agent isolation that drives lane attribution.""" + +from __future__ import annotations + +import pytest + +from app.services.streaming.emitter import ( + Emitter, + EmitterRegistry, + main_emitter, + subagent_emitter, +) + +pytestmark = pytest.mark.unit + + +def _sub(run_id: str, kind: str = "deliverables") -> Emitter: + return subagent_emitter( + subagent_type=kind, + subagent_run_id=f"sub_{run_id}", + parent_tool_call_id=f"call_{run_id}", + ) + + +def test_unregistered_event_resolves_to_main_emitter() -> None: + registry = EmitterRegistry() + resolved = registry.resolve(run_id="run_1", parent_ids=["root"]) + assert resolved is main_emitter() + + +def test_event_owned_by_registered_run_id_returns_that_emitter() -> None: + registry = EmitterRegistry() + emitter = _sub("a") + registry.register("run_task_a", emitter) + assert registry.resolve(run_id="run_task_a", parent_ids=[]) is emitter + + +def test_descendant_resolves_via_parent_ids_chain() -> None: + """A model-call event nested under the task tool inherits its sub-agent emitter.""" + registry = EmitterRegistry() + emitter = _sub("a") + registry.register("run_task_a", emitter) + descendant = registry.resolve( + run_id="run_chat_model", + parent_ids=["root", "run_agent", "run_task_a"], + ) + assert descendant is emitter + + +def test_nearest_registered_ancestor_wins_over_distant_ones() -> None: + """Inner sub-agents owe their emitter to the nearest task tool, not the outer one.""" + registry = EmitterRegistry() + outer = _sub("outer", kind="planner") + inner = _sub("inner", kind="email") + registry.register("run_outer", outer) + registry.register("run_inner", inner) + resolved = registry.resolve( + run_id="run_inner_tool", + parent_ids=["root", "run_outer", "run_inner"], + ) + assert resolved is inner + + +def test_parallel_subagents_do_not_bleed_into_each_other() -> None: + """Two concurrent task tools each own their own descendant events.""" + registry = EmitterRegistry() + a = _sub("a", kind="search") + b = _sub("b", kind="email") + registry.register("run_task_a", a) + registry.register("run_task_b", b) + + from_a = registry.resolve(run_id="x", parent_ids=["root", "run_task_a"]) + from_b = registry.resolve(run_id="y", parent_ids=["root", "run_task_b"]) + from_main = registry.resolve(run_id="z", parent_ids=["root"]) + + assert from_a is a + assert from_b is b + assert from_main is main_emitter() + + +def test_unregister_releases_run_id_so_descendants_fall_back_to_main() -> None: + registry = EmitterRegistry() + emitter = _sub("a") + registry.register("run_task_a", emitter) + registry.unregister("run_task_a") + assert registry.resolve(run_id="x", parent_ids=["run_task_a"]) is main_emitter() + + +def test_unregister_returns_the_previously_registered_emitter() -> None: + """Lets callers emit ``data-subagent-finish`` carrying the same emitter they opened with.""" + registry = EmitterRegistry() + emitter = _sub("a") + registry.register("run_task_a", emitter) + assert registry.unregister("run_task_a") is emitter + + +def test_has_active_subagents_tracks_open_lanes() -> None: + registry = EmitterRegistry() + assert not registry.has_active_subagents() + registry.register("run_task_a", _sub("a")) + assert registry.has_active_subagents() + registry.unregister("run_task_a") + assert not registry.has_active_subagents() + + +def test_empty_run_id_and_parent_ids_resolves_to_main() -> None: + """Defensive: events without identifiers always belong to the main lane.""" + registry = EmitterRegistry() + registry.register("run_task_a", _sub("a")) + assert registry.resolve(run_id=None, parent_ids=None) is main_emitter() + assert registry.resolve(run_id="", parent_ids=[]) is main_emitter() From 366122da6e4568289e131e2ac20be63e8bb5bd90 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 42/96] Add unit tests for streaming interrupts and service propagation. --- .../tests/unit/services/streaming/__init__.py | 0 .../streaming/test_interrupt_correlation.py | 164 ++++++++++++++++++ .../streaming/test_interrupt_events.py | 91 ++++++++++ .../test_service_emitter_propagation.py | 142 +++++++++++++++ 4 files changed, 397 insertions(+) create mode 100644 surfsense_backend/tests/unit/services/streaming/__init__.py create mode 100644 surfsense_backend/tests/unit/services/streaming/test_interrupt_correlation.py create mode 100644 surfsense_backend/tests/unit/services/streaming/test_interrupt_events.py create mode 100644 surfsense_backend/tests/unit/services/streaming/test_service_emitter_propagation.py diff --git a/surfsense_backend/tests/unit/services/streaming/__init__.py b/surfsense_backend/tests/unit/services/streaming/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/tests/unit/services/streaming/test_interrupt_correlation.py b/surfsense_backend/tests/unit/services/streaming/test_interrupt_correlation.py new file mode 100644 index 000000000..edf4ecb9a --- /dev/null +++ b/surfsense_backend/tests/unit/services/streaming/test_interrupt_correlation.py @@ -0,0 +1,164 @@ +"""Pin id-aware pending-interrupt lookup that replaces the buggy first-wins.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import pytest + +from app.services.streaming.interrupt_correlation import ( + PendingInterrupt, + first_pending_interrupt, + get_pending_interrupt_by_id, + get_pending_interrupt_for_tool_call, + list_pending_interrupts, +) + +pytestmark = pytest.mark.unit + + +@dataclass +class _Interrupt: + value: dict[str, Any] + id: str | None = None + + +@dataclass +class _Task: + interrupts: tuple[_Interrupt, ...] = () + id: str | None = None + + +@dataclass +class _State: + tasks: tuple[_Task, ...] = () + interrupts: tuple[_Interrupt, ...] = () + + +def _hitl(name: str, tool_call_id: str | None = None) -> dict[str, Any]: + """Minimal LangChain HITLRequest payload for one action.""" + action: dict[str, Any] = {"name": name, "args": {}} + if tool_call_id is not None: + action["tool_call_id"] = tool_call_id + return { + "action_requests": [action], + "review_configs": [{"action_name": name, "allowed_decisions": ["approve"]}], + } + + +def test_empty_state_has_no_pending_interrupts() -> None: + state = _State() + assert list_pending_interrupts(state) == [] + assert first_pending_interrupt(state) is None + + +def test_single_pending_interrupt_in_task_is_returned() -> None: + state = _State( + tasks=( + _Task( + id="task_1", + interrupts=(_Interrupt(value=_hitl("send_email"), id="int_1"),), + ), + ) + ) + pending = list_pending_interrupts(state) + assert len(pending) == 1 + assert pending[0] == PendingInterrupt( + interrupt_id="int_1", + value=_hitl("send_email"), + source_task_id="task_1", + ) + + +def test_pending_interrupts_returned_in_task_then_root_order() -> None: + """Determinism matters: callers iterate in this order to render the UI.""" + state = _State( + tasks=( + _Task( + id="task_a", + interrupts=(_Interrupt(value=_hitl("a"), id="int_a"),), + ), + _Task( + id="task_b", + interrupts=(_Interrupt(value=_hitl("b"), id="int_b"),), + ), + ), + interrupts=(_Interrupt(value=_hitl("c"), id="int_c"),), + ) + pending = list_pending_interrupts(state) + ids = [p.interrupt_id for p in pending] + assert ids == ["int_a", "int_b", "int_c"] + + +def test_get_by_id_finds_the_right_interrupt_under_parallel_load() -> None: + """Replacing first-wins: id-aware lookup MUST pick the requested one.""" + state = _State( + tasks=( + _Task(interrupts=(_Interrupt(value=_hitl("a"), id="int_a"),)), + _Task(interrupts=(_Interrupt(value=_hitl("b"), id="int_b"),)), + _Task(interrupts=(_Interrupt(value=_hitl("c"), id="int_c"),)), + ) + ) + found = get_pending_interrupt_by_id(state, "int_b") + assert found is not None + assert found.value["action_requests"][0]["name"] == "b" + + +def test_get_by_id_returns_none_when_id_is_not_pending() -> None: + state = _State( + tasks=(_Task(interrupts=(_Interrupt(value=_hitl("a"), id="int_a"),)),) + ) + assert get_pending_interrupt_by_id(state, "missing") is None + + +def test_get_by_tool_call_id_matches_action_request_payload() -> None: + """HITLRequest carries ``tool_call_id`` per action; lookup uses that.""" + state = _State( + tasks=( + _Task( + interrupts=( + _Interrupt( + value=_hitl("a", tool_call_id="call_xxx"), id="int_a" + ), + _Interrupt( + value=_hitl("b", tool_call_id="call_yyy"), id="int_b" + ), + ) + ), + ) + ) + found = get_pending_interrupt_for_tool_call(state, "call_yyy") + assert found is not None + assert found.interrupt_id == "int_b" + + +def test_first_pending_interrupt_matches_legacy_first_wins_behaviour() -> None: + """Sequential-turn safety: the explicit shortcut still returns the first.""" + state = _State( + tasks=(_Task(interrupts=(_Interrupt(value=_hitl("first"), id="int_1"),)),), + interrupts=(_Interrupt(value=_hitl("second"), id="int_2"),), + ) + first = first_pending_interrupt(state) + assert first is not None + assert first.interrupt_id == "int_1" + + +def test_interrupt_without_id_falls_back_to_none() -> None: + """Snapshots from older LangGraph versions may omit ``id`` — preserve that.""" + state = _State( + tasks=(_Task(interrupts=(_Interrupt(value=_hitl("a"), id=None),)),) + ) + pending = list_pending_interrupts(state) + assert len(pending) == 1 + assert pending[0].interrupt_id is None + + +def test_non_dict_interrupt_values_are_ignored() -> None: + """Defensive: a non-dict value should not crash the iteration.""" + + class _Raw: + value = "not a dict" + + state = _State(tasks=(_Task(interrupts=(_Raw(),)),)) # type: ignore[arg-type] + assert list_pending_interrupts(state) == [] diff --git a/surfsense_backend/tests/unit/services/streaming/test_interrupt_events.py b/surfsense_backend/tests/unit/services/streaming/test_interrupt_events.py new file mode 100644 index 000000000..dbdd607bf --- /dev/null +++ b/surfsense_backend/tests/unit/services/streaming/test_interrupt_events.py @@ -0,0 +1,91 @@ +"""Pin interrupt-payload normalisation and the optional correlation fields on the wire.""" + +from __future__ import annotations + +import json + +import pytest + +from app.services.streaming.events.interrupt import ( + format_interrupt_request, + normalize_interrupt_payload, +) + +pytestmark = pytest.mark.unit + + +def _decode(frame: str) -> dict: + body = frame.removeprefix("data: ").removesuffix("\n\n") + return json.loads(body) + + +def test_hitlrequest_shape_is_passed_through_unchanged() -> None: + raw = { + "action_requests": [{"name": "send_email", "args": {"to": "a@b"}}], + "review_configs": [ + {"action_name": "send_email", "allowed_decisions": ["approve"]} + ], + } + assert normalize_interrupt_payload(raw) == raw + + +def test_custom_interrupt_primitive_is_converted_to_canonical_shape() -> None: + raw = { + "type": "permission", + "message": "Allow send?", + "action": {"tool": "send_email", "params": {"to": "a@b"}}, + "context": {"reason": "destructive"}, + } + out = normalize_interrupt_payload(raw) + assert out["action_requests"] == [ + {"name": "send_email", "args": {"to": "a@b"}} + ] + assert out["review_configs"] == [ + { + "action_name": "send_email", + "allowed_decisions": ["approve", "edit", "reject"], + } + ] + assert out["interrupt_type"] == "permission" + assert out["message"] == "Allow send?" + assert out["context"] == {"reason": "destructive"} + + +def test_custom_interrupt_without_message_omits_message_key() -> None: + """Optional fields stay optional on the wire; FE does not see ``"message": None``.""" + raw = {"action": {"tool": "send_email"}} + out = normalize_interrupt_payload(raw) + assert "message" not in out + + +def test_custom_interrupt_without_tool_falls_back_to_unknown_tool() -> None: + """Defensive: a malformed ``action`` block must not crash the relay.""" + out = normalize_interrupt_payload({"type": "x", "action": {}}) + assert out["action_requests"][0]["name"] == "unknown_tool" + assert out["review_configs"][0]["action_name"] == "unknown_tool" + + +def test_format_interrupt_request_carries_correlation_fields_on_the_wire() -> None: + frame = format_interrupt_request( + {"action_requests": [], "review_configs": []}, + interrupt_id="int_42", + pending_interrupt_count=3, + chat_turn_id="turn_99", + ) + payload = _decode(frame) + assert payload["type"] == "data-interrupt-request" + inner = payload["data"] + assert inner["interrupt_id"] == "int_42" + assert inner["pending_interrupt_count"] == 3 + assert inner["chat_turn_id"] == "turn_99" + + +def test_format_interrupt_request_omits_correlation_fields_when_unset() -> None: + """Backward compat: legacy single-interrupt callers don't have to supply ids.""" + frame = format_interrupt_request( + {"action_requests": [], "review_configs": []}, + ) + inner = _decode(frame)["data"] + assert "interrupt_id" not in inner + assert "pending_interrupt_count" not in inner + assert "chat_turn_id" not in inner diff --git a/surfsense_backend/tests/unit/services/streaming/test_service_emitter_propagation.py b/surfsense_backend/tests/unit/services/streaming/test_service_emitter_propagation.py new file mode 100644 index 000000000..b381f13bc --- /dev/null +++ b/surfsense_backend/tests/unit/services/streaming/test_service_emitter_propagation.py @@ -0,0 +1,142 @@ +"""Pin that sub-agent emitter reaches every wire event the relay emits.""" + +from __future__ import annotations + +import json + +import pytest + +from app.services.streaming.emitter import subagent_emitter +from app.services.streaming.service import StreamingService + +pytestmark = pytest.mark.unit + + +def _decode(frame: str) -> dict: + body = frame.removeprefix("data: ").removesuffix("\n\n") + return json.loads(body) + + +@pytest.fixture +def service() -> StreamingService: + return StreamingService() + + +@pytest.fixture +def sub_emitter(): + return subagent_emitter( + subagent_type="deliverables", + subagent_run_id="sub_xyz", + parent_tool_call_id="call_parent", + ) + + +def test_text_delta_carries_subagent_emitter_on_the_wire(service, sub_emitter) -> None: + payload = _decode(service.format_text_delta("text_1", "hi", emitter=sub_emitter)) + assert payload["emitted_by"]["subagent_run_id"] == "sub_xyz" + assert payload["delta"] == "hi" + + +def test_reasoning_delta_carries_subagent_emitter_on_the_wire( + service, sub_emitter +) -> None: + payload = _decode( + service.format_reasoning_delta("r_1", "thinking", emitter=sub_emitter) + ) + assert payload["emitted_by"]["subagent_run_id"] == "sub_xyz" + + +def test_tool_input_start_carries_subagent_emitter_and_lc_id( + service, sub_emitter +) -> None: + payload = _decode( + service.format_tool_input_start( + "call_1", + "send_email", + langchain_tool_call_id="lc_1", + emitter=sub_emitter, + ) + ) + assert payload["emitted_by"]["subagent_type"] == "deliverables" + assert payload["langchainToolCallId"] == "lc_1" + assert payload["toolName"] == "send_email" + + +def test_tool_output_available_carries_subagent_emitter(service, sub_emitter) -> None: + payload = _decode( + service.format_tool_output_available( + "call_1", {"ok": True}, emitter=sub_emitter + ) + ) + assert payload["emitted_by"]["subagent_run_id"] == "sub_xyz" + assert payload["output"] == {"ok": True} + + +def test_thinking_step_carries_subagent_emitter(service, sub_emitter) -> None: + payload = _decode( + service.format_thinking_step( + step_id="s1", + title="Sending email", + status="in_progress", + emitter=sub_emitter, + ) + ) + assert payload["type"] == "data-thinking-step" + assert payload["emitted_by"]["subagent_run_id"] == "sub_xyz" + + +def test_action_log_carries_subagent_emitter(service, sub_emitter) -> None: + payload = _decode( + service.format_action_log( + {"id": 1, "tool_name": "send_email", "reversible": False}, + emitter=sub_emitter, + ) + ) + assert payload["emitted_by"]["subagent_run_id"] == "sub_xyz" + assert payload["data"]["tool_name"] == "send_email" + + +def test_subagent_lifecycle_events_share_run_id_for_pairing( + service, sub_emitter +) -> None: + start = _decode( + service.format_subagent_start( + subagent_run_id="sub_xyz", + subagent_type="deliverables", + parent_tool_call_id="call_parent", + emitter=sub_emitter, + ) + ) + finish = _decode( + service.format_subagent_finish( + subagent_run_id="sub_xyz", + subagent_type="deliverables", + parent_tool_call_id="call_parent", + emitter=sub_emitter, + ) + ) + assert start["data"]["subagent_run_id"] == finish["data"]["subagent_run_id"] + assert start["type"] == "data-subagent-start" + assert finish["type"] == "data-subagent-finish" + + +def test_main_emitter_events_omit_emitted_by_field(service) -> None: + payload = _decode(service.format_text_delta("text_1", "hi")) + assert "emitted_by" not in payload + + +def test_resolve_emitter_through_service_uses_registry(service, sub_emitter) -> None: + service.emitter_registry.register("run_task_1", sub_emitter) + resolved = service.resolve_emitter( + run_id="run_chat_model", + parent_ids=["root", "run_task_1"], + ) + assert resolved is sub_emitter + + +def test_message_id_is_assigned_on_message_start_and_reused(service) -> None: + frame = service.format_message_start() + payload = _decode(frame) + assigned = payload["messageId"] + assert assigned.startswith("msg_") + assert service.message_id == assigned From c25b78c30492e53e5f1b69f831f15dda029c9d54 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 43/96] Add chat streaming error classification, helpers, and StreamResult. --- .../app/tasks/chat/streaming/__init__.py | 3 + .../tasks/chat/streaming/errors/__init__.py | 3 + .../tasks/chat/streaming/errors/classifier.py | 187 ++++++++++++++++++ .../tasks/chat/streaming/errors/emitter.py | 38 ++++ .../tasks/chat/streaming/helpers/__init__.py | 3 + .../chat/streaming/helpers/chunk_parts.py | 60 ++++++ .../streaming/helpers/interrupt_inspector.py | 47 +++++ .../streaming/helpers/tool_call_matching.py | 32 +++ .../chat/streaming/helpers/tool_output.py | 43 ++++ .../app/tasks/chat/streaming/stream_result.py | 28 +++ 10 files changed, 444 insertions(+) create mode 100644 surfsense_backend/app/tasks/chat/streaming/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/errors/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/errors/classifier.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/errors/emitter.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/helpers/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/helpers/chunk_parts.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/helpers/interrupt_inspector.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/helpers/tool_call_matching.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/helpers/tool_output.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/stream_result.py diff --git a/surfsense_backend/app/tasks/chat/streaming/__init__.py b/surfsense_backend/app/tasks/chat/streaming/__init__.py new file mode 100644 index 000000000..bb06cc021 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/__init__.py @@ -0,0 +1,3 @@ +"""Chat streaming orchestrator and event relay.""" + +from __future__ import annotations diff --git a/surfsense_backend/app/tasks/chat/streaming/errors/__init__.py b/surfsense_backend/app/tasks/chat/streaming/errors/__init__.py new file mode 100644 index 000000000..02284d4b0 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/errors/__init__.py @@ -0,0 +1,3 @@ +"""Error classification, structured logging, and terminal-error SSE emission.""" + +from __future__ import annotations diff --git a/surfsense_backend/app/tasks/chat/streaming/errors/classifier.py b/surfsense_backend/app/tasks/chat/streaming/errors/classifier.py new file mode 100644 index 000000000..3af2b9f9f --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/errors/classifier.py @@ -0,0 +1,187 @@ +"""Classify stream exceptions for logging and client error payloads.""" + +from __future__ import annotations + +import json +import logging +import time +from typing import Any, Literal + +from app.agents.new_chat.errors import BusyError +from app.agents.new_chat.middleware.busy_mutex import ( + get_cancel_state, + is_cancel_requested, +) + +TURN_CANCELLING_INITIAL_DELAY_MS = 200 +TURN_CANCELLING_BACKOFF_FACTOR = 2 +TURN_CANCELLING_MAX_DELAY_MS = 1500 + + +def compute_turn_cancelling_retry_delay(attempt: int) -> int: + if attempt < 1: + attempt = 1 + delay = TURN_CANCELLING_INITIAL_DELAY_MS * ( + TURN_CANCELLING_BACKOFF_FACTOR ** (attempt - 1) + ) + return min(delay, TURN_CANCELLING_MAX_DELAY_MS) + + +def log_chat_stream_error( + *, + flow: Literal["new", "resume", "regenerate"], + error_kind: str, + error_code: str | None, + severity: Literal["info", "warn", "error"], + is_expected: bool, + request_id: str | None, + thread_id: int | None, + search_space_id: int | None, + user_id: str | None, + message: str, + extra: dict[str, Any] | None = None, +) -> None: + payload: dict[str, Any] = { + "event": "chat_stream_error", + "flow": flow, + "error_kind": error_kind, + "error_code": error_code, + "severity": severity, + "is_expected": is_expected, + "request_id": request_id or "unknown", + "thread_id": thread_id, + "search_space_id": search_space_id, + "user_id": user_id, + "message": message, + } + if extra: + payload.update(extra) + + logger = logging.getLogger(__name__) + rendered = json.dumps(payload, ensure_ascii=False) + if severity == "error": + logger.error("[chat_stream_error] %s", rendered) + elif severity == "warn": + logger.warning("[chat_stream_error] %s", rendered) + else: + logger.info("[chat_stream_error] %s", rendered) + + +def _parse_error_payload(message: str) -> dict[str, Any] | None: + candidates = [message] + first_brace_idx = message.find("{") + if first_brace_idx >= 0: + candidates.append(message[first_brace_idx:]) + + for candidate in candidates: + try: + parsed = json.loads(candidate) + if isinstance(parsed, dict): + return parsed + except Exception: + continue + return None + + +def _extract_provider_error_code(parsed: dict[str, Any] | None) -> int | None: + if not isinstance(parsed, dict): + return None + candidates: list[Any] = [parsed.get("code")] + nested = parsed.get("error") + if isinstance(nested, dict): + candidates.append(nested.get("code")) + for value in candidates: + try: + if value is None: + continue + return int(value) + except Exception: + continue + return None + + +def is_provider_rate_limited(exc: BaseException) -> bool: + """Return True if the exception looks like an upstream HTTP 429 / rate limit.""" + raw = str(exc) + lowered = raw.lower() + if "ratelimit" in type(exc).__name__.lower(): + return True + parsed = _parse_error_payload(raw) + provider_code = _extract_provider_error_code(parsed) + if provider_code == 429: + return True + + provider_error_type = "" + if parsed: + top_type = parsed.get("type") + if isinstance(top_type, str): + provider_error_type = top_type.lower() + nested = parsed.get("error") + if isinstance(nested, dict): + nested_type = nested.get("type") + if isinstance(nested_type, str): + provider_error_type = nested_type.lower() + if provider_error_type == "rate_limit_error": + return True + + return ( + "rate limited" in lowered + or "rate-limited" in lowered + or "temporarily rate-limited upstream" in lowered + ) + + +def classify_stream_exception( + exc: Exception, + *, + flow_label: str, +) -> tuple[ + str, str, Literal["info", "warn", "error"], bool, str, dict[str, Any] | None +]: + """Return kind, code, severity, expected flag, message, and optional extra dict.""" + raw = str(exc) + if isinstance(exc, BusyError) or "Thread is busy with another request" in raw: + busy_thread_id = str(exc.request_id) if isinstance(exc, BusyError) else None + if busy_thread_id and is_cancel_requested(busy_thread_id): + cancel_state = get_cancel_state(busy_thread_id) + attempt = cancel_state[0] if cancel_state else 1 + retry_after_ms = compute_turn_cancelling_retry_delay(attempt) + retry_after_at = int(time.time() * 1000) + retry_after_ms + return ( + "thread_busy", + "TURN_CANCELLING", + "info", + True, + "A previous response is still stopping. Please try again in a moment.", + { + "retry_after_ms": retry_after_ms, + "retry_after_at": retry_after_at, + }, + ) + return ( + "thread_busy", + "THREAD_BUSY", + "warn", + True, + "Another response is still finishing for this thread. Please try again in a moment.", + None, + ) + + if is_provider_rate_limited(exc): + return ( + "rate_limited", + "RATE_LIMITED", + "warn", + True, + "This model is temporarily rate-limited. Please try again in a few seconds or switch models.", + None, + ) + + return ( + "server_error", + "SERVER_ERROR", + "error", + False, + f"Error during {flow_label}: {raw}", + None, + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/errors/emitter.py b/surfsense_backend/app/tasks/chat/streaming/errors/emitter.py new file mode 100644 index 000000000..95806ab87 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/errors/emitter.py @@ -0,0 +1,38 @@ +"""Emit one terminal error SSE frame and log via the stream error classifier.""" + +from __future__ import annotations + +from typing import Any, Literal + +from .classifier import log_chat_stream_error + + +def emit_stream_terminal_error( + *, + streaming_service: Any, + flow: Literal["new", "resume", "regenerate"], + request_id: str | None, + thread_id: int, + search_space_id: int, + user_id: str | None, + message: str, + error_kind: str = "server_error", + error_code: str = "SERVER_ERROR", + severity: Literal["info", "warn", "error"] = "error", + is_expected: bool = False, + extra: dict[str, Any] | None = None, +) -> str: + log_chat_stream_error( + flow=flow, + error_kind=error_kind, + error_code=error_code, + severity=severity, + is_expected=is_expected, + request_id=request_id, + thread_id=thread_id, + search_space_id=search_space_id, + user_id=user_id, + message=message, + extra=extra, + ) + return streaming_service.format_error(message, error_code=error_code, extra=extra) diff --git a/surfsense_backend/app/tasks/chat/streaming/helpers/__init__.py b/surfsense_backend/app/tasks/chat/streaming/helpers/__init__.py new file mode 100644 index 000000000..151dfdaac --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/helpers/__init__.py @@ -0,0 +1,3 @@ +"""Pure helpers for chat streaming.""" + +from __future__ import annotations diff --git a/surfsense_backend/app/tasks/chat/streaming/helpers/chunk_parts.py b/surfsense_backend/app/tasks/chat/streaming/helpers/chunk_parts.py new file mode 100644 index 000000000..48b44fc1d --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/helpers/chunk_parts.py @@ -0,0 +1,60 @@ +"""Split a model chunk into text, reasoning, and tool-call fragment lists.""" + +from __future__ import annotations + +from typing import Any + + +def extract_chunk_parts(chunk: Any) -> dict[str, Any]: + """Return dict with keys text, reasoning, and tool_call_chunks (merged from chunk fields).""" + out: dict[str, Any] = {"text": "", "reasoning": "", "tool_call_chunks": []} + if chunk is None: + return out + + content = getattr(chunk, "content", None) + if isinstance(content, str): + if content: + out["text"] = content + elif isinstance(content, list): + text_parts: list[str] = [] + reasoning_parts: list[str] = [] + for block in content: + if not isinstance(block, dict): + continue + block_type = block.get("type") + if block_type == "text": + value = block.get("text") or block.get("content") or "" + if isinstance(value, str) and value: + text_parts.append(value) + elif block_type == "reasoning": + value = ( + block.get("reasoning") + or block.get("text") + or block.get("content") + or "" + ) + if isinstance(value, str) and value: + reasoning_parts.append(value) + elif block_type in ("tool_call_chunk", "tool_use"): + out["tool_call_chunks"].append(block) + if text_parts: + out["text"] = "".join(text_parts) + if reasoning_parts: + out["reasoning"] = "".join(reasoning_parts) + + additional = getattr(chunk, "additional_kwargs", None) or {} + if isinstance(additional, dict): + extra_reasoning = additional.get("reasoning_content") + if isinstance(extra_reasoning, str) and extra_reasoning: + existing = out["reasoning"] + out["reasoning"] = ( + (existing + extra_reasoning) if existing else extra_reasoning + ) + + extra_tool_chunks = getattr(chunk, "tool_call_chunks", None) + if isinstance(extra_tool_chunks, list): + for tcc in extra_tool_chunks: + if isinstance(tcc, dict): + out["tool_call_chunks"].append(tcc) + + return out diff --git a/surfsense_backend/app/tasks/chat/streaming/helpers/interrupt_inspector.py b/surfsense_backend/app/tasks/chat/streaming/helpers/interrupt_inspector.py new file mode 100644 index 000000000..dca099b3f --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/helpers/interrupt_inspector.py @@ -0,0 +1,47 @@ +"""Read the first interrupt payload from a LangGraph state snapshot.""" + +from __future__ import annotations + +from typing import Any + + +def first_interrupt_value(state: Any) -> dict[str, Any] | None: + """Return the first interrupt payload across all snapshot tasks.""" + + def _extract(candidate: Any) -> dict[str, Any] | None: + if isinstance(candidate, dict): + value = candidate.get("value", candidate) + return value if isinstance(value, dict) else None + value = getattr(candidate, "value", None) + if isinstance(value, dict): + return value + if isinstance(candidate, list | tuple): + for item in candidate: + extracted = _extract(item) + if extracted is not None: + return extracted + return None + + for task in getattr(state, "tasks", ()) or (): + try: + interrupts = getattr(task, "interrupts", ()) or () + except (AttributeError, IndexError, TypeError): + interrupts = () + if not interrupts: + extracted = _extract(task) + if extracted is not None: + return extracted + continue + for interrupt_item in interrupts: + extracted = _extract(interrupt_item) + if extracted is not None: + return extracted + + try: + state_interrupts = getattr(state, "interrupts", ()) or () + except (AttributeError, IndexError, TypeError): + state_interrupts = () + extracted = _extract(state_interrupts) + if extracted is not None: + return extracted + return None diff --git a/surfsense_backend/app/tasks/chat/streaming/helpers/tool_call_matching.py b/surfsense_backend/app/tasks/chat/streaming/helpers/tool_call_matching.py new file mode 100644 index 000000000..fbe4c94b7 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/helpers/tool_call_matching.py @@ -0,0 +1,32 @@ +"""Match buffered model tool-call chunks to a tool start when ids were missing.""" + +from __future__ import annotations + +from typing import Any + + +def match_buffered_langchain_tool_call_id( + pending_tool_call_chunks: list[dict[str, Any]], + tool_name: str, + run_id: str, + lc_tool_call_id_by_run: dict[str, str], +) -> str | None: + matched_idx: int | None = None + for idx, tcc in enumerate(pending_tool_call_chunks): + if tcc.get("name") == tool_name and tcc.get("id"): + matched_idx = idx + break + if matched_idx is None: + for idx, tcc in enumerate(pending_tool_call_chunks): + if tcc.get("id"): + matched_idx = idx + break + if matched_idx is None: + return None + matched = pending_tool_call_chunks.pop(matched_idx) + candidate = matched.get("id") + if isinstance(candidate, str) and candidate: + if run_id: + lc_tool_call_id_by_run[run_id] = candidate + return candidate + return None diff --git a/surfsense_backend/app/tasks/chat/streaming/helpers/tool_output.py b/surfsense_backend/app/tasks/chat/streaming/helpers/tool_output.py new file mode 100644 index 000000000..a7c401dee --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/helpers/tool_output.py @@ -0,0 +1,43 @@ +"""Normalize filesystem tool payloads for SSE cards and messages.""" + +from __future__ import annotations + +import json +from typing import Any + + +def tool_output_to_text(tool_output: Any) -> str: + if isinstance(tool_output, dict): + if isinstance(tool_output.get("result"), str): + return tool_output["result"] + if isinstance(tool_output.get("error"), str): + return tool_output["error"] + return json.dumps(tool_output, ensure_ascii=False) + return str(tool_output) + + +def tool_output_has_error(tool_output: Any) -> bool: + if isinstance(tool_output, dict): + if tool_output.get("error"): + return True + result = tool_output.get("result") + return bool( + isinstance(result, str) and result.strip().lower().startswith("error:") + ) + if isinstance(tool_output, str): + return tool_output.strip().lower().startswith("error:") + return False + + +def extract_resolved_file_path( + *, tool_name: str, tool_output: Any, tool_input: Any | None = None +) -> str | None: + if isinstance(tool_output, dict): + path_value = tool_output.get("path") + if isinstance(path_value, str) and path_value.strip(): + return path_value.strip() + if tool_name in ("write_file", "edit_file") and isinstance(tool_input, dict): + file_path = tool_input.get("file_path") + if isinstance(file_path, str) and file_path.strip(): + return file_path.strip() + return None diff --git a/surfsense_backend/app/tasks/chat/streaming/stream_result.py b/surfsense_backend/app/tasks/chat/streaming/stream_result.py new file mode 100644 index 000000000..8ea3bd295 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/stream_result.py @@ -0,0 +1,28 @@ +"""Mutable facts collected while streaming one agent turn.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class StreamResult: + accumulated_text: str = "" + is_interrupted: bool = False + interrupt_value: dict[str, Any] | None = None + sandbox_files: list[str] = field(default_factory=list) + agent_called_update_memory: bool = False + request_id: str | None = None + turn_id: str = "" + filesystem_mode: str = "cloud" + client_platform: str = "web" + intent_detected: str = "chat_only" + intent_confidence: float = 0.0 + write_attempted: bool = False + write_succeeded: bool = False + verification_succeeded: bool = False + commit_gate_passed: bool = True + commit_gate_reason: str = "" + assistant_message_id: int | None = None + content_builder: Any | None = field(default=None, repr=False) From 7581a7c9c3247bc977ae556daddd25cc185ef2eb Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 44/96] Add chat streaming relay state and thinking-step SSE helpers. --- .../tasks/chat/streaming/relay/__init__.py | 3 + .../app/tasks/chat/streaming/relay/state.py | 55 +++++++++++++++++++ .../relay/thinking_step_completion.py | 31 +++++++++++ .../chat/streaming/relay/thinking_step_sse.py | 24 ++++++++ 4 files changed, 113 insertions(+) create mode 100644 surfsense_backend/app/tasks/chat/streaming/relay/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/relay/state.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py b/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py new file mode 100644 index 000000000..c1a5e7175 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py @@ -0,0 +1,3 @@ +"""Relay state: thinking steps, tool bookkeeping, and stream helpers.""" + +from __future__ import annotations diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/state.py b/surfsense_backend/app/tasks/chat/streaming/relay/state.py new file mode 100644 index 000000000..e8e35d0b2 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/relay/state.py @@ -0,0 +1,55 @@ +"""Mutable counters and maps for one agent stream.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class AgentEventRelayState: + """Tracks text, thinking steps, tool depth, and pending tool-call metadata.""" + + accumulated_text: str = "" + current_text_id: str | None = None + thinking_step_counter: int = 0 + tool_step_ids: dict[str, str] = field(default_factory=dict) + completed_step_ids: set[str] = field(default_factory=set) + last_active_step_id: str | None = None + last_active_step_title: str = "" + last_active_step_items: list[str] = field(default_factory=list) + just_finished_tool: bool = False + active_tool_depth: int = 0 + called_update_memory: bool = False + current_reasoning_id: str | None = None + parity_v2: bool = False + pending_tool_call_chunks: list[dict[str, Any]] = field(default_factory=list) + lc_tool_call_id_by_run: dict[str, str] = field(default_factory=dict) + file_path_by_run: dict[str, str] = field(default_factory=dict) + index_to_meta: dict[int, dict[str, str]] = field(default_factory=dict) + ui_tool_call_id_by_run: dict[str, str] = field(default_factory=dict) + current_lc_tool_call_id: dict[str, str | None] = field( + default_factory=lambda: {"value": None} + ) + + @classmethod + def for_invocation( + cls, + *, + initial_step_id: str | None = None, + initial_step_title: str = "", + initial_step_items: list[str] | None = None, + parity_v2: bool, + ) -> AgentEventRelayState: + counter = 1 if initial_step_id else 0 + return cls( + thinking_step_counter=counter, + last_active_step_id=initial_step_id, + last_active_step_title=initial_step_title, + last_active_step_items=list(initial_step_items or []), + parity_v2=parity_v2, + ) + + def next_thinking_step_id(self, step_prefix: str) -> str: + self.thinking_step_counter += 1 + return f"{step_prefix}-{self.thinking_step_counter}" diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py new file mode 100644 index 000000000..a0be71281 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py @@ -0,0 +1,31 @@ +"""Close the in-progress thinking step with a completed status frame.""" + +from __future__ import annotations + +from typing import Any + +from .thinking_step_sse import emit_thinking_step_frame + + +def complete_active_thinking_step( + *, + streaming_service: Any, + content_builder: Any | None, + last_active_step_id: str | None, + last_active_step_title: str, + last_active_step_items: list[str], + completed_step_ids: set[str], +) -> tuple[str | None, str | None]: + """Emit a completed thinking-step frame once; return (frame or None, next active step id).""" + if last_active_step_id and last_active_step_id not in completed_step_ids: + completed_step_ids.add(last_active_step_id) + event = emit_thinking_step_frame( + streaming_service=streaming_service, + content_builder=content_builder, + step_id=last_active_step_id, + title=last_active_step_title, + status="completed", + items=last_active_step_items if last_active_step_items else None, + ) + return event, None + return None, last_active_step_id diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py new file mode 100644 index 000000000..9e8c08dd5 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py @@ -0,0 +1,24 @@ +"""Thinking-step SSE plus optional content-builder updates.""" + +from __future__ import annotations + +from typing import Any + + +def emit_thinking_step_frame( + *, + streaming_service: Any, + content_builder: Any | None, + step_id: str, + title: str, + status: str = "in_progress", + items: list[str] | None = None, +) -> str: + if content_builder is not None: + content_builder.on_thinking_step(step_id, title, status, items) + return streaming_service.format_thinking_step( + step_id=step_id, + title=title, + status=status, + items=items, + ) From ee16e1d5f96b22d7bb9ed822f1ec983c966b7b91 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 45/96] Add LangGraph handlers for chat model, chain, tool, and custom events. --- .../tasks/chat/streaming/handlers/__init__.py | 3 + .../chat/streaming/handlers/chain_end.py | 23 +++ .../streaming/handlers/chat_model_stream.py | 149 ++++++++++++++++++ .../handlers/custom_event_dispatch.py | 56 +++++++ .../chat/streaming/handlers/custom_events.py | 77 +++++++++ .../tasks/chat/streaming/handlers/tool_end.py | 112 +++++++++++++ .../streaming/handlers/tool_output_frame.py | 24 +++ .../chat/streaming/handlers/tool_start.py | 142 +++++++++++++++++ 8 files changed, 586 insertions(+) create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/chain_end.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/custom_event_dispatch.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tool_output_frame.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/__init__.py new file mode 100644 index 000000000..3e2165932 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/__init__.py @@ -0,0 +1,3 @@ +"""LangGraph stream handlers by event kind.""" + +from __future__ import annotations diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/chain_end.py b/surfsense_backend/app/tasks/chat/streaming/handlers/chain_end.py new file mode 100644 index 000000000..c61058ac7 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/chain_end.py @@ -0,0 +1,23 @@ +"""Close open text when a LangGraph chain or agent node finishes.""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + +from app.tasks.chat.streaming.relay.state import AgentEventRelayState + + +def iter_chain_end_frames( + _event: dict[str, Any], + *, + state: AgentEventRelayState, + streaming_service: Any, + content_builder: Any | None, +) -> Iterator[str]: + """Close the open text stream if one is open.""" + if state.current_text_id is not None: + yield streaming_service.format_text_end(state.current_text_id) + if content_builder is not None: + content_builder.on_text_end(state.current_text_id) + state.current_text_id = None diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py b/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py new file mode 100644 index 000000000..861342b32 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py @@ -0,0 +1,149 @@ +"""Chat model stream: text, reasoning, and tool-call chunk SSE.""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + +from app.tasks.chat.streaming.helpers.chunk_parts import extract_chunk_parts +from app.tasks.chat.streaming.relay.state import AgentEventRelayState +from app.tasks.chat.streaming.relay.thinking_step_completion import ( + complete_active_thinking_step, +) + + +def iter_chat_model_stream_frames( + event: dict[str, Any], + *, + state: AgentEventRelayState, + streaming_service: Any, + content_builder: Any | None, + step_prefix: str, +) -> Iterator[str]: + """SSE frames for one chat-model chunk.""" + if state.active_tool_depth > 0: + return + if "surfsense:internal" in event.get("tags", []): + return + chunk = event.get("data", {}).get("chunk") + if not chunk: + return + parts = extract_chunk_parts(chunk) + + reasoning_delta = parts["reasoning"] + text_delta = parts["text"] + + if state.parity_v2 and reasoning_delta: + if state.current_text_id is not None: + yield streaming_service.format_text_end(state.current_text_id) + if content_builder is not None: + content_builder.on_text_end(state.current_text_id) + state.current_text_id = None + if state.current_reasoning_id is None: + comp, new_active = complete_active_thinking_step( + streaming_service=streaming_service, + content_builder=content_builder, + last_active_step_id=state.last_active_step_id, + last_active_step_title=state.last_active_step_title, + last_active_step_items=state.last_active_step_items, + completed_step_ids=state.completed_step_ids, + ) + if comp: + yield comp + state.last_active_step_id = new_active + if state.just_finished_tool: + state.last_active_step_id = None + state.last_active_step_title = "" + state.last_active_step_items = [] + state.just_finished_tool = False + state.current_reasoning_id = streaming_service.generate_reasoning_id() + yield streaming_service.format_reasoning_start(state.current_reasoning_id) + if content_builder is not None: + content_builder.on_reasoning_start(state.current_reasoning_id) + yield streaming_service.format_reasoning_delta( + state.current_reasoning_id, reasoning_delta + ) + if content_builder is not None: + content_builder.on_reasoning_delta( + state.current_reasoning_id, reasoning_delta + ) + + if text_delta: + if state.current_reasoning_id is not None: + yield streaming_service.format_reasoning_end(state.current_reasoning_id) + if content_builder is not None: + content_builder.on_reasoning_end(state.current_reasoning_id) + state.current_reasoning_id = None + if state.current_text_id is None: + comp, new_active = complete_active_thinking_step( + streaming_service=streaming_service, + content_builder=content_builder, + last_active_step_id=state.last_active_step_id, + last_active_step_title=state.last_active_step_title, + last_active_step_items=state.last_active_step_items, + completed_step_ids=state.completed_step_ids, + ) + if comp: + yield comp + state.last_active_step_id = new_active + if state.just_finished_tool: + state.last_active_step_id = None + state.last_active_step_title = "" + state.last_active_step_items = [] + state.just_finished_tool = False + state.current_text_id = streaming_service.generate_text_id() + yield streaming_service.format_text_start(state.current_text_id) + if content_builder is not None: + content_builder.on_text_start(state.current_text_id) + yield streaming_service.format_text_delta(state.current_text_id, text_delta) + state.accumulated_text += text_delta + if content_builder is not None: + content_builder.on_text_delta(state.current_text_id, text_delta) + + if state.parity_v2 and parts["tool_call_chunks"]: + for tcc in parts["tool_call_chunks"]: + idx = tcc.get("index") + + if idx is not None and idx not in state.index_to_meta: + lc_id = tcc.get("id") + name = tcc.get("name") + if lc_id and name: + ui_id = lc_id + + if state.current_text_id is not None: + yield streaming_service.format_text_end(state.current_text_id) + if content_builder is not None: + content_builder.on_text_end(state.current_text_id) + state.current_text_id = None + if state.current_reasoning_id is not None: + yield streaming_service.format_reasoning_end( + state.current_reasoning_id + ) + if content_builder is not None: + content_builder.on_reasoning_end(state.current_reasoning_id) + state.current_reasoning_id = None + + state.index_to_meta[idx] = { + "ui_id": ui_id, + "lc_id": lc_id, + "name": name, + } + yield streaming_service.format_tool_input_start( + ui_id, + name, + langchain_tool_call_id=lc_id, + ) + if content_builder is not None: + content_builder.on_tool_input_start(ui_id, name, lc_id) + + meta = state.index_to_meta.get(idx) if idx is not None else None + if meta: + args_chunk = tcc.get("args") or "" + if args_chunk: + yield streaming_service.format_tool_input_delta( + meta["ui_id"], args_chunk + ) + if content_builder is not None: + content_builder.on_tool_input_delta(meta["ui_id"], args_chunk) + else: + state.pending_tool_call_chunks.append(tcc) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/custom_event_dispatch.py b/surfsense_backend/app/tasks/chat/streaming/handlers/custom_event_dispatch.py new file mode 100644 index 000000000..b373919cf --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/custom_event_dispatch.py @@ -0,0 +1,56 @@ +"""Custom graph events routed to SSE (documents, action logs, report progress).""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + +from app.tasks.chat.streaming.handlers.custom_events import ( + handle_action_log, + handle_action_log_updated, + handle_document_created, + handle_report_progress, +) +from app.tasks.chat.streaming.relay.state import AgentEventRelayState + + +def iter_custom_event_frames( + event: dict[str, Any], + *, + state: AgentEventRelayState, + streaming_service: Any, + content_builder: Any | None, +) -> Iterator[str]: + """Yield any SSE produced by ad-hoc graph events (documents, action logs, report progress).""" + name = event.get("name") + data = event.get("data", {}) + + if name == "report_progress": + frame, state.last_active_step_items = handle_report_progress( + data, + last_active_step_id=state.last_active_step_id, + last_active_step_title=state.last_active_step_title, + last_active_step_items=state.last_active_step_items, + streaming_service=streaming_service, + content_builder=content_builder, + ) + if frame: + yield frame + return + + if name == "document_created": + frame = handle_document_created(data, streaming_service=streaming_service) + if frame: + yield frame + return + + if name == "action_log": + frame = handle_action_log(data, streaming_service=streaming_service) + if frame: + yield frame + return + + if name == "action_log_updated": + frame = handle_action_log_updated(data, streaming_service=streaming_service) + if frame: + yield frame diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py b/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py new file mode 100644 index 000000000..765f1d790 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py @@ -0,0 +1,77 @@ +"""Custom-event payloads turned into SSE (no model/tool stream handling).""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.relay.thinking_step_sse import emit_thinking_step_frame + + +def handle_report_progress( + data: dict[str, Any], + *, + last_active_step_id: str | None, + last_active_step_title: str, + last_active_step_items: list[str], + streaming_service: Any, + content_builder: Any | None, +) -> tuple[str | None, list[str]]: + """Update report step items; may emit one thinking SSE frame. + + Returns (frame or None, items list after update). + """ + message = data.get("message", "") + if not message or not last_active_step_id: + return None, last_active_step_items + + phase = data.get("phase", "") + topic_items = [ + item for item in last_active_step_items if item.startswith("Topic:") + ] + + if phase in ("revising_section", "adding_section"): + plan_items = [ + item + for item in last_active_step_items + if item.startswith("Topic:") + or item.startswith("Modifying ") + or item.startswith("Adding ") + or item.startswith("Removing ") + ] + plan_items = [item for item in plan_items if not item.endswith("...")] + new_items = [*plan_items, message] + else: + new_items = [*topic_items, message] + + frame = emit_thinking_step_frame( + streaming_service=streaming_service, + content_builder=content_builder, + step_id=last_active_step_id, + title=last_active_step_title, + status="in_progress", + items=new_items, + ) + return frame, new_items + + +def handle_document_created(data: dict[str, Any], *, streaming_service: Any) -> str | None: + if not data.get("id"): + return None + return streaming_service.format_data( + "documents-updated", + {"action": "created", "document": data}, + ) + + +def handle_action_log(data: dict[str, Any], *, streaming_service: Any) -> str | None: + if data.get("id") is None: + return None + return streaming_service.format_data("action-log", data) + + +def handle_action_log_updated( + data: dict[str, Any], *, streaming_service: Any +) -> str | None: + if data.get("id") is None: + return None + return streaming_service.format_data("action-log-updated", data) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py new file mode 100644 index 000000000..0bfef25eb --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py @@ -0,0 +1,112 @@ +"""Tool end: thinking completion, tool output, and terminal SSE.""" + +from __future__ import annotations + +import json +from collections.abc import Iterator +from typing import Any + +from app.tasks.chat.streaming.handlers.tools import ( + ToolCompletionEmissionContext, + iter_tool_completion_emission_frames, + resolve_tool_completed_thinking_step, +) +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.thinking_step_sse import emit_thinking_step_frame + + +def iter_tool_end_frames( + event: dict[str, Any], + *, + state: AgentEventRelayState, + streaming_service: Any, + content_builder: Any | None, + result: Any, + step_prefix: str, + config: dict[str, Any], +) -> Iterator[str]: + """SSE frames when one tool run finishes.""" + state.active_tool_depth = max(0, state.active_tool_depth - 1) + run_id = event.get("run_id", "") + tool_name = event.get("name", "unknown_tool") + raw_output = event.get("data", {}).get("output", "") + staged_file_path = ( + state.file_path_by_run.pop(run_id, None) if run_id else None + ) + + if tool_name == "update_memory": + state.called_update_memory = True + + if hasattr(raw_output, "content"): + content = raw_output.content + if isinstance(content, str): + try: + tool_output = json.loads(content) + except (json.JSONDecodeError, TypeError): + tool_output = {"result": content} + elif isinstance(content, dict): + tool_output = content + else: + tool_output = {"result": str(content)} + elif isinstance(raw_output, dict): + tool_output = raw_output + else: + tool_output = {"result": str(raw_output) if raw_output else "completed"} + + if tool_name in ("write_file", "edit_file"): + if tool_output_has_error(tool_output): + pass + else: + result.write_succeeded = True + result.verification_succeeded = True + + tool_call_id = state.ui_tool_call_id_by_run.get( + run_id, + f"call_{run_id[:32]}" if run_id else "call_unknown", + ) + original_step_id = state.tool_step_ids.get( + run_id, f"{step_prefix}-unknown-{run_id[:8]}" + ) + state.completed_step_ids.add(original_step_id) + + holder = state.current_lc_tool_call_id + holder["value"] = None + authoritative = getattr(raw_output, "tool_call_id", None) + if isinstance(authoritative, str) and authoritative: + holder["value"] = authoritative + if run_id: + state.lc_tool_call_id_by_run[run_id] = authoritative + elif run_id and run_id in state.lc_tool_call_id_by_run: + holder["value"] = state.lc_tool_call_id_by_run[run_id] + + items = state.last_active_step_items + title, completed_items = resolve_tool_completed_thinking_step( + tool_name, tool_output, items + ) + yield emit_thinking_step_frame( + streaming_service=streaming_service, + content_builder=content_builder, + step_id=original_step_id, + title=title, + status="completed", + items=completed_items, + ) + + state.just_finished_tool = True + state.last_active_step_id = None + state.last_active_step_title = "" + state.last_active_step_items = [] + + emission_ctx = ToolCompletionEmissionContext( + tool_name=tool_name, + tool_call_id=tool_call_id, + tool_output=tool_output, + streaming_service=streaming_service, + content_builder=content_builder, + langchain_tool_call_id_holder=holder, + stream_result=result, + langgraph_config=config, + staged_workspace_file_path=staged_file_path, + ) + yield from iter_tool_completion_emission_frames(emission_ctx) 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 new file mode 100644 index 000000000..07244364c --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_output_frame.py @@ -0,0 +1,24 @@ +"""Emit tool-output SSE and optional assistant content updates.""" + +from __future__ import annotations + +from typing import Any + + +def emit_tool_output_available_frame( + *, + streaming_service: Any, + content_builder: Any | None, + langchain_id_holder: dict[str, str | None], + call_id: str, + output: Any, +) -> str: + if content_builder is not None: + content_builder.on_tool_output_available( + call_id, output, langchain_id_holder["value"] + ) + return streaming_service.format_tool_output_available( + call_id, + output, + langchain_tool_call_id=langchain_id_holder["value"], + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py new file mode 100644 index 000000000..c316cc74a --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py @@ -0,0 +1,142 @@ +"""Tool start: thinking-step and tool-input SSE.""" + +from __future__ import annotations + +import json +from collections.abc import Iterator +from typing import Any + +from app.tasks.chat.streaming.handlers.tools import resolve_tool_start_thinking +from app.tasks.chat.streaming.helpers.tool_call_matching import ( + match_buffered_langchain_tool_call_id, +) +from app.tasks.chat.streaming.relay.state import AgentEventRelayState +from app.tasks.chat.streaming.relay.thinking_step_completion import ( + complete_active_thinking_step, +) +from app.tasks.chat.streaming.relay.thinking_step_sse import emit_thinking_step_frame + + +def iter_tool_start_frames( + event: dict[str, Any], + *, + state: AgentEventRelayState, + streaming_service: Any, + content_builder: Any | None, + result: Any, + step_prefix: str, +) -> Iterator[str]: + """SSE frames for the start of one tool run.""" + state.active_tool_depth += 1 + tool_name = event.get("name", "unknown_tool") + run_id = event.get("run_id", "") + tool_input = event.get("data", {}).get("input", {}) + if tool_name in ("write_file", "edit_file"): + result.write_attempted = True + if isinstance(tool_input, dict): + file_path = tool_input.get("file_path") + if isinstance(file_path, str) and file_path.strip() and run_id: + state.file_path_by_run[run_id] = file_path.strip() + + if state.current_text_id is not None: + yield streaming_service.format_text_end(state.current_text_id) + if content_builder is not None: + content_builder.on_text_end(state.current_text_id) + state.current_text_id = None + + if state.last_active_step_title != "Synthesizing response": + comp, new_active = complete_active_thinking_step( + streaming_service=streaming_service, + content_builder=content_builder, + last_active_step_id=state.last_active_step_id, + last_active_step_title=state.last_active_step_title, + last_active_step_items=state.last_active_step_items, + completed_step_ids=state.completed_step_ids, + ) + if comp: + yield comp + state.last_active_step_id = new_active + + state.just_finished_tool = False + tool_step_id = state.next_thinking_step_id(step_prefix) + state.tool_step_ids[run_id] = tool_step_id + state.last_active_step_id = tool_step_id + + thinking = resolve_tool_start_thinking(tool_name, tool_input) + state.last_active_step_title = thinking.title + state.last_active_step_items = thinking.items + frame_kw: dict[str, Any] = { + "streaming_service": streaming_service, + "content_builder": content_builder, + "step_id": tool_step_id, + "title": thinking.title, + "status": "in_progress", + } + if thinking.include_items_on_frame: + frame_kw["items"] = thinking.items + yield emit_thinking_step_frame(**frame_kw) + + matched_meta: dict[str, str] | None = None + if state.parity_v2: + taken_ui_ids = set(state.ui_tool_call_id_by_run.values()) + for meta in state.index_to_meta.values(): + if meta["name"] == tool_name and meta["ui_id"] not in taken_ui_ids: + matched_meta = meta + break + + tool_call_id: str + langchain_tool_call_id: str | None = None + if matched_meta is not None: + tool_call_id = matched_meta["ui_id"] + langchain_tool_call_id = matched_meta["lc_id"] + if run_id: + state.lc_tool_call_id_by_run[run_id] = matched_meta["lc_id"] + else: + tool_call_id = ( + f"call_{run_id[:32]}" + if run_id + else streaming_service.generate_tool_call_id() + ) + if state.parity_v2: + langchain_tool_call_id = match_buffered_langchain_tool_call_id( + state.pending_tool_call_chunks, + tool_name, + run_id, + state.lc_tool_call_id_by_run, + ) + yield streaming_service.format_tool_input_start( + tool_call_id, + tool_name, + langchain_tool_call_id=langchain_tool_call_id, + ) + if content_builder is not None: + content_builder.on_tool_input_start( + tool_call_id, tool_name, langchain_tool_call_id + ) + + if run_id: + state.ui_tool_call_id_by_run[run_id] = tool_call_id + + if isinstance(tool_input, dict): + _safe_input: dict[str, Any] = {} + for _k, _v in tool_input.items(): + try: + json.dumps(_v) + _safe_input[_k] = _v + except (TypeError, ValueError, OverflowError): + pass + else: + _safe_input = {"input": tool_input} + yield streaming_service.format_tool_input_available( + tool_call_id, + tool_name, + _safe_input, + langchain_tool_call_id=langchain_tool_call_id, + ) + if content_builder is not None: + content_builder.on_tool_input_available( + tool_call_id, + tool_name, + _safe_input, + langchain_tool_call_id, + ) From 1392abf5b1d9cbb2309c7780edeacf9c0c72f205 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 46/96] Add chat tool streaming registry with shared, default, and connector tools. --- .../chat/streaming/handlers/tools/__init__.py | 23 +++++ .../handlers/tools/connector/__init__.py | 0 .../tools/connector/shared/__init__.py | 0 .../tools/connector/shared/emission.py | 15 ++++ .../tools/connector/shared/thinking.py | 22 +++++ .../tools/connector/shared/tool_names.py | 31 +++++++ .../handlers/tools/default/__init__.py | 3 + .../handlers/tools/default/emission.py | 24 +++++ .../handlers/tools/default/thinking.py | 23 +++++ .../handlers/tools/emission_context.py | 34 +++++++ .../chat/streaming/handlers/tools/registry.py | 88 +++++++++++++++++++ .../handlers/tools/shared/__init__.py | 0 .../streaming/handlers/tools/shared/model.py | 12 +++ 13 files changed, 275 insertions(+) create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/tool_names.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/emission_context.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/registry.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/shared/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/shared/model.py diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/__init__.py new file mode 100644 index 000000000..4b191c100 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/__init__.py @@ -0,0 +1,23 @@ +"""Per-tool streaming: thinking-step and completion emission.""" + +from __future__ import annotations + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) +from app.tasks.chat.streaming.handlers.tools.registry import ( + iter_tool_completion_emission_frames, + resolve_tool_completed_thinking_step, + resolve_tool_start_thinking, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + +__all__ = [ + "ToolCompletionEmissionContext", + "ToolStartThinking", + "iter_tool_completion_emission_frames", + "resolve_tool_completed_thinking_step", + "resolve_tool_start_thinking", +] diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/emission.py new file mode 100644 index 000000000..8e19dc224 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/emission.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + payload = out if isinstance(out, dict) else {"result": out} + yield ctx.emit_tool_output_card(payload) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/thinking.py new file mode 100644 index 000000000..7e9dd8b96 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/thinking.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.default import ( + thinking as default_thinking, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + return default_thinking.resolve_start_thinking(tool_name, tool_input) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + return default_thinking.resolve_completed_thinking( + tool_name, tool_output, last_items + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/tool_names.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/tool_names.py new file mode 100644 index 000000000..ab698b32d --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/connector/shared/tool_names.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +SHARED_CONNECTOR_TOOLS: frozenset[str] = frozenset( + { + "create_calendar_event", + "create_confluence_page", + "create_dropbox_file", + "create_gmail_draft", + "create_google_drive_file", + "create_jira_issue", + "create_linear_issue", + "create_notion_page", + "create_onedrive_file", + "delete_calendar_event", + "delete_confluence_page", + "delete_dropbox_file", + "delete_google_drive_file", + "delete_jira_issue", + "delete_linear_issue", + "delete_notion_page", + "delete_onedrive_file", + "send_gmail_email", + "trash_gmail_email", + "update_calendar_event", + "update_confluence_page", + "update_gmail_draft", + "update_jira_issue", + "update_linear_issue", + "update_notion_page", + } +) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/__init__.py new file mode 100644 index 000000000..5e84a37f4 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/__init__.py @@ -0,0 +1,3 @@ +"""Fallback tool package.""" + +from __future__ import annotations diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/emission.py new file mode 100644 index 000000000..e24c619a7 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/emission.py @@ -0,0 +1,24 @@ +"""Default tool-output card and a short completion terminal line.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + yield ctx.emit_tool_output_card( + { + "status": "completed", + "result_length": len(str(ctx.tool_output)), + }, + ) + yield ctx.streaming_service.format_terminal_info( + f"Tool {ctx.tool_name} completed", + "success", + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/thinking.py new file mode 100644 index 000000000..46d15a4e7 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/default/thinking.py @@ -0,0 +1,23 @@ +"""Fallback thinking-step copy for unknown tools and connectors without custom UI.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_input + title = tool_name.replace("_", " ").strip().capitalize() or tool_name + return ToolStartThinking(title=title, items=[], include_items_on_frame=False) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str] +) -> tuple[str, list[str]]: + del tool_output + title = tool_name.replace("_", " ").strip().capitalize() or tool_name + return (title, last_items) 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 new file mode 100644 index 000000000..d9ff796c0 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/emission_context.py @@ -0,0 +1,34 @@ +"""Context for one tool-completion emission pass.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from app.tasks.chat.streaming.handlers.tool_output_frame import ( + emit_tool_output_available_frame, +) + + +@dataclass +class ToolCompletionEmissionContext: + """Streaming service, tool output, and ids for completion frames.""" + + tool_name: str + tool_call_id: str + tool_output: Any + streaming_service: Any + content_builder: Any | None + langchain_tool_call_id_holder: dict[str, str | None] + stream_result: Any + langgraph_config: dict[str, Any] + staged_workspace_file_path: str | None + + def emit_tool_output_card(self, payload: Any) -> str: + return emit_tool_output_available_frame( + streaming_service=self.streaming_service, + content_builder=self.content_builder, + langchain_id_holder=self.langchain_tool_call_id_holder, + call_id=self.tool_call_id, + output=payload, + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/registry.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/registry.py new file mode 100644 index 000000000..c0568f870 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/registry.py @@ -0,0 +1,88 @@ +"""Resolve thinking and emission modules by tool name.""" + +from __future__ import annotations + +import importlib +from collections.abc import Iterator +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.connector.shared.tool_names import ( + SHARED_CONNECTOR_TOOLS, +) +from app.tasks.chat.streaming.handlers.tools.deliverables.tool_names import ( + DELIVERABLE_TOOLS, +) +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) +from app.tasks.chat.streaming.handlers.tools.filesystem.tool_names import ( + FILESYSTEM_TOOLS, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + +_BASE = "app.tasks.chat.streaming.handlers.tools" +_CONNECTOR_SHARED = "connector.shared" + +_THINKING_ALIAS: dict[str, str] = { + "execute_code": "filesystem.execute", +} +_EMISSION_ALIAS: dict[str, str] = { + "edit_file": "filesystem.write_file", + "execute_code": "filesystem.execute", +} + + +def _thinking_module(tool_name: str) -> str: + if tool_name in SHARED_CONNECTOR_TOOLS: + return _CONNECTOR_SHARED + if tool_name in FILESYSTEM_TOOLS: + return f"filesystem.{tool_name}" + if tool_name in DELIVERABLE_TOOLS: + return f"deliverables.{tool_name}" + return _THINKING_ALIAS.get(tool_name, tool_name) + + +def _emission_module(tool_name: str) -> str: + if tool_name in _EMISSION_ALIAS: + return _EMISSION_ALIAS[tool_name] + if tool_name in SHARED_CONNECTOR_TOOLS: + return _CONNECTOR_SHARED + if tool_name in DELIVERABLE_TOOLS: + return f"deliverables.{tool_name}" + if tool_name in FILESYSTEM_TOOLS: + return f"filesystem.{tool_name}" + return tool_name + + +def _import_thinking(tool_name: str): + try: + return importlib.import_module(f"{_BASE}.{_thinking_module(tool_name)}.thinking") + except ModuleNotFoundError: + return importlib.import_module(f"{_BASE}.default.thinking") + + +def _import_emission(tool_name: str): + try: + return importlib.import_module(f"{_BASE}.{_emission_module(tool_name)}.emission") + except ModuleNotFoundError: + return importlib.import_module(f"{_BASE}.default.emission") + + +def resolve_tool_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + return _import_thinking(tool_name).resolve_start_thinking(tool_name, tool_input) + + +def resolve_tool_completed_thinking_step( + tool_name: str, tool_output: Any, last_items: list[str] +) -> tuple[str, list[str]]: + return _import_thinking(tool_name).resolve_completed_thinking( + tool_name, tool_output, last_items + ) + + +def iter_tool_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + yield from _import_emission(ctx.tool_name).iter_completion_emission_frames(ctx) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/shared/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/shared/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/shared/model.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/shared/model.py new file mode 100644 index 000000000..047a84374 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/shared/model.py @@ -0,0 +1,12 @@ +"""In-progress thinking-step title and bullet lines.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class ToolStartThinking: + title: str + items: list[str] + include_items_on_frame: bool = True From a322eedaa1a3a0604ffb5f053b9e5381ac539284 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 47/96] Add filesystem tool streaming handlers for chat runs. --- .../handlers/tools/filesystem/__init__.py | 0 .../tools/filesystem/edit_file/__init__.py | 0 .../tools/filesystem/edit_file/thinking.py | 27 +++++++++ .../tools/filesystem/execute/__init__.py | 0 .../tools/filesystem/execute/emission.py | 40 +++++++++++++ .../tools/filesystem/execute/thinking.py | 42 +++++++++++++ .../tools/filesystem/glob/__init__.py | 0 .../tools/filesystem/glob/thinking.py | 27 +++++++++ .../tools/filesystem/grep/__init__.py | 0 .../tools/filesystem/grep/thinking.py | 31 ++++++++++ .../handlers/tools/filesystem/ls/__init__.py | 0 .../handlers/tools/filesystem/ls/thinking.py | 59 +++++++++++++++++++ .../tools/filesystem/mkdir/__init__.py | 0 .../tools/filesystem/mkdir/thinking.py | 27 +++++++++ .../tools/filesystem/move_file/__init__.py | 0 .../tools/filesystem/move_file/thinking.py | 33 +++++++++++ .../tools/filesystem/read_file/__init__.py | 0 .../tools/filesystem/read_file/thinking.py | 27 +++++++++ .../handlers/tools/filesystem/rm/__init__.py | 0 .../handlers/tools/filesystem/rm/thinking.py | 28 +++++++++ .../tools/filesystem/rmdir/__init__.py | 0 .../tools/filesystem/rmdir/thinking.py | 27 +++++++++ .../tools/filesystem/shared/__init__.py | 0 .../tools/filesystem/shared/tool_input.py | 17 ++++++ .../handlers/tools/filesystem/tool_names.py | 18 ++++++ .../tools/filesystem/write_file/__init__.py | 0 .../tools/filesystem/write_file/emission.py | 43 ++++++++++++++ .../tools/filesystem/write_file/thinking.py | 27 +++++++++ .../tools/filesystem/write_todos/__init__.py | 0 .../tools/filesystem/write_todos/thinking.py | 34 +++++++++++ 30 files changed, 507 insertions(+) create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/edit_file/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/edit_file/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/glob/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/glob/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/grep/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/grep/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/ls/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/ls/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/mkdir/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/mkdir/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/move_file/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/move_file/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/read_file/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/read_file/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rm/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rm/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rmdir/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rmdir/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/shared/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/shared/tool_input.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/tool_names.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_todos/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_todos/thinking.py diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/edit_file/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/edit_file/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/edit_file/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/edit_file/thinking.py new file mode 100644 index 000000000..8669107db --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/edit_file/thinking.py @@ -0,0 +1,27 @@ +"""edit_file: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, + truncate_path, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + fp = d.get("file_path", "") if isinstance(tool_input, dict) else str(tool_input) + return ToolStartThinking(title="Editing file", items=[truncate_path(fp)]) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_output, tool_name + return ("Editing file", last_items) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/emission.py new file mode 100644 index 000000000..0ff87a907 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/emission.py @@ -0,0 +1,40 @@ +"""execute: exit code, stdout, sandbox file hints.""" + +from __future__ import annotations + +import re +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + raw_text = out.get("result", "") if isinstance(out, dict) else str(out) + exit_code: int | None = None + output_text = raw_text + m = re.match(r"^Exit code:\s*(\d+)", raw_text) + if m: + exit_code = int(m.group(1)) + om = re.search(r"\nOutput:\n([\s\S]*)", raw_text) + output_text = om.group(1) if om else "" + thread_id_str = ctx.langgraph_config.get("configurable", {}).get("thread_id", "") + + for sf_match in re.finditer( + r"^SANDBOX_FILE:\s*(.+)$", output_text, re.MULTILINE + ): + fpath = sf_match.group(1).strip() + if fpath and fpath not in ctx.stream_result.sandbox_files: + ctx.stream_result.sandbox_files.append(fpath) + + yield ctx.emit_tool_output_card( + { + "exit_code": exit_code, + "output": output_text, + "thread_id": thread_id_str, + }, + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/thinking.py new file mode 100644 index 000000000..2c8aa296b --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/execute/thinking.py @@ -0,0 +1,42 @@ +"""execute: sandbox command thinking + completion lines.""" + +from __future__ import annotations + +import re +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + cmd = d.get("command", "") if isinstance(tool_input, dict) else str(tool_input) + display_cmd = cmd[:80] + ("…" if len(cmd) > 80 else "") + return ToolStartThinking(title="Running command", items=[f"$ {display_cmd}"]) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_name + items = last_items + raw_text = ( + tool_output.get("result", "") + if isinstance(tool_output, dict) + else str(tool_output) + ) + m = re.match(r"^Exit code:\s*(\d+)", raw_text) + exit_code_val = int(m.group(1)) if m else None + if exit_code_val is not None and exit_code_val == 0: + completed = [*items, "Completed successfully"] + elif exit_code_val is not None: + completed = [*items, f"Exit code: {exit_code_val}"] + else: + completed = [*items, "Finished"] + return ("Running command", completed) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/glob/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/glob/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/glob/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/glob/thinking.py new file mode 100644 index 000000000..f5a57beac --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/glob/thinking.py @@ -0,0 +1,27 @@ +"""glob: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + pat = d.get("pattern", "") if isinstance(tool_input, dict) else str(tool_input) + base = d.get("path", "/") if isinstance(tool_input, dict) else "/" + return ToolStartThinking(title="Searching files", items=[f"{pat} in {base}"]) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_output, tool_name + return ("Searching files", last_items) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/grep/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/grep/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/grep/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/grep/thinking.py new file mode 100644 index 000000000..da0864177 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/grep/thinking.py @@ -0,0 +1,31 @@ +"""grep: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + pat = d.get("pattern", "") if isinstance(tool_input, dict) else str(tool_input) + grep_path = d.get("path", "") if isinstance(tool_input, dict) else "" + display_pat = pat[:60] + ("…" if len(pat) > 60 else "") + return ToolStartThinking( + title="Searching content", + items=[f'"{display_pat}"' + (f" in {grep_path}" if grep_path else "")], + ) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_output, tool_name + return ("Searching content", last_items) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/ls/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/ls/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/ls/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/ls/thinking.py new file mode 100644 index 000000000..80c547b5a --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/ls/thinking.py @@ -0,0 +1,59 @@ +"""ls: thinking-step copy for directory listing.""" + +from __future__ import annotations + +import ast +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + if isinstance(tool_input, dict): + path = tool_input.get("path", "/") + else: + path = str(tool_input) + return ToolStartThinking(title="Listing files", items=[path]) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_name + if isinstance(tool_output, dict): + ls_output = tool_output.get("result", "") + elif isinstance(tool_output, str): + ls_output = tool_output + else: + ls_output = str(tool_output) if tool_output else "" + file_names: list[str] = [] + if ls_output: + paths: list[str] = [] + try: + parsed = ast.literal_eval(ls_output) + if isinstance(parsed, list): + paths = [str(p) for p in parsed] + except (ValueError, SyntaxError): + paths = [ + line.strip() + for line in ls_output.strip().split("\n") + if line.strip() + ] + for p in paths: + name = p.rstrip("/").split("/")[-1] + if name and len(name) <= 40: + file_names.append(name) + elif name: + file_names.append(name[:37] + "...") + if file_names: + if len(file_names) <= 5: + completed = [f"[{name}]" for name in file_names] + else: + completed = [f"[{name}]" for name in file_names[:4]] + completed.append(f"(+{len(file_names) - 4} more)") + else: + completed = ["No files found"] + return ("Listing files", completed) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/mkdir/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/mkdir/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/mkdir/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/mkdir/thinking.py new file mode 100644 index 000000000..3a3707698 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/mkdir/thinking.py @@ -0,0 +1,27 @@ +"""mkdir: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + p = d.get("path", "") if isinstance(tool_input, dict) else str(tool_input) + display = p if len(p) <= 80 else "…" + p[-77:] + return ToolStartThinking(title="Creating folder", items=[display] if display else []) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_output, tool_name + return ("Creating folder", last_items) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/move_file/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/move_file/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/move_file/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/move_file/thinking.py new file mode 100644 index 000000000..192a789f4 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/move_file/thinking.py @@ -0,0 +1,33 @@ +"""move_file: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, + truncate_middle, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + src = d.get("source_path", "") if isinstance(tool_input, dict) else "" + dst = d.get("destination_path", "") if isinstance(tool_input, dict) else "" + display_src = truncate_middle(src, max_len=60) + display_dst = truncate_middle(dst, max_len=60) + return ToolStartThinking( + title="Moving file", + items=[f"{display_src} → {display_dst}"] if src or dst else [], + ) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_output, tool_name + return ("Moving file", last_items) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/read_file/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/read_file/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/read_file/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/read_file/thinking.py new file mode 100644 index 000000000..3f4290ad7 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/read_file/thinking.py @@ -0,0 +1,27 @@ +"""read_file: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, + truncate_path, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + fp = d.get("file_path", "") if isinstance(tool_input, dict) else str(tool_input) + return ToolStartThinking(title="Reading file", items=[truncate_path(fp)]) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_output, tool_name + return ("Reading file", last_items) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rm/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rm/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rm/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rm/thinking.py new file mode 100644 index 000000000..a82a44e6f --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rm/thinking.py @@ -0,0 +1,28 @@ +"""rm: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, + truncate_path, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + rm_path = d.get("path", "") if isinstance(tool_input, dict) else str(tool_input) + display = truncate_path(rm_path) + return ToolStartThinking(title="Deleting file", items=[display] if display else []) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_output, tool_name + return ("Deleting file", last_items) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rmdir/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rmdir/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rmdir/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rmdir/thinking.py new file mode 100644 index 000000000..6c97904b7 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/rmdir/thinking.py @@ -0,0 +1,27 @@ +"""rmdir: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + p = d.get("path", "") if isinstance(tool_input, dict) else str(tool_input) + display = p if len(p) <= 80 else "…" + p[-77:] + return ToolStartThinking(title="Deleting folder", items=[display] if display else []) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_output, tool_name + return ("Deleting folder", last_items) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/shared/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/shared/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/shared/tool_input.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/shared/tool_input.py new file mode 100644 index 000000000..507782283 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/shared/tool_input.py @@ -0,0 +1,17 @@ +"""Tool-call args + display truncation for filesystem thinking modules.""" + +from __future__ import annotations + +from typing import Any + + +def as_tool_input_dict(tool_input: Any) -> dict[str, Any]: + return tool_input if isinstance(tool_input, dict) else {} + + +def truncate_path(fp: str, *, max_len: int = 80) -> str: + return fp if len(fp) <= max_len else "…" + fp[-(max_len - 3) :] + + +def truncate_middle(s: str, *, max_len: int = 60) -> str: + return s if len(s) <= max_len else "…" + s[-(max_len - 3) :] diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/tool_names.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/tool_names.py new file mode 100644 index 000000000..e2ad33736 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/tool_names.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +FILESYSTEM_TOOLS: frozenset[str] = frozenset( + { + "read_file", + "glob", + "grep", + "ls", + "mkdir", + "move_file", + "rm", + "rmdir", + "write_todos", + "write_file", + "edit_file", + "execute", + } +) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/emission.py new file mode 100644 index 000000000..820235379 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/emission.py @@ -0,0 +1,43 @@ +"""write_file: path + status envelope on the tool-output card.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) +from app.tasks.chat.streaming.helpers.tool_output import ( + extract_resolved_file_path, + tool_output_has_error, + tool_output_to_text, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + resolved_path = extract_resolved_file_path( + tool_name=ctx.tool_name, + tool_output=ctx.tool_output, + tool_input={"file_path": ctx.staged_workspace_file_path} + if ctx.staged_workspace_file_path + else None, + ) + result_text = tool_output_to_text(ctx.tool_output) + if tool_output_has_error(ctx.tool_output): + yield ctx.emit_tool_output_card( + { + "status": "error", + "error": result_text, + "path": resolved_path, + }, + ) + else: + yield ctx.emit_tool_output_card( + { + "status": "completed", + "path": resolved_path, + "result": result_text, + }, + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/thinking.py new file mode 100644 index 000000000..43bc8a65f --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_file/thinking.py @@ -0,0 +1,27 @@ +"""write_file: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, + truncate_path, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + fp = d.get("file_path", "") if isinstance(tool_input, dict) else str(tool_input) + return ToolStartThinking(title="Writing file", items=[truncate_path(fp)]) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_output, tool_name + return ("Writing file", last_items) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_todos/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_todos/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_todos/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_todos/thinking.py new file mode 100644 index 000000000..43e533daa --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/filesystem/write_todos/thinking.py @@ -0,0 +1,34 @@ +"""write_todos: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.filesystem.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + todos = d.get("todos", []) if isinstance(tool_input, dict) else [] + todo_count = len(todos) if isinstance(todos, list) else 0 + return ToolStartThinking( + title="Planning tasks", + items=( + [f"{todo_count} task{'s' if todo_count != 1 else ''}"] + if todo_count + else [] + ), + ) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_output, tool_name + return ("Planning tasks", last_items) From c8fb4aa5e5b2df6c3c8cca2081f5f8634bd69d16 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 48/96] Add deliverables and web tool streaming handlers for chat runs. --- .../handlers/tools/deliverables/__init__.py | 0 .../deliverables/generate_image/__init__.py | 0 .../deliverables/generate_image/emission.py | 28 +++++++ .../deliverables/generate_image/thinking.py | 39 +++++++++ .../deliverables/generate_podcast/__init__.py | 0 .../deliverables/generate_podcast/emission.py | 37 +++++++++ .../deliverables/generate_podcast/thinking.py | 80 +++++++++++++++++++ .../deliverables/generate_report/__init__.py | 0 .../deliverables/generate_report/emission.py | 33 ++++++++ .../deliverables/generate_report/thinking.py | 77 ++++++++++++++++++ .../deliverables/generate_resume/__init__.py | 0 .../deliverables/generate_resume/emission.py | 32 ++++++++ .../deliverables/generate_resume/thinking.py | 24 ++++++ .../generate_video_presentation/__init__.py | 0 .../generate_video_presentation/emission.py | 28 +++++++ .../generate_video_presentation/thinking.py | 52 ++++++++++++ .../deliverables/save_document/__init__.py | 0 .../deliverables/save_document/emission.py | 16 ++++ .../deliverables/save_document/thinking.py | 38 +++++++++ .../tools/deliverables/shared/__init__.py | 0 .../tools/deliverables/shared/tool_input.py | 9 +++ .../handlers/tools/deliverables/tool_names.py | 12 +++ .../handlers/tools/scrape_webpage/emission.py | 43 ++++++++++ .../tools/scrape_webpage/shared/__init__.py | 0 .../tools/scrape_webpage/shared/tool_input.py | 9 +++ .../handlers/tools/scrape_webpage/thinking.py | 47 +++++++++++ .../handlers/tools/web_search/emission.py | 41 ++++++++++ 27 files changed, 645 insertions(+) create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/shared/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/shared/tool_input.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/tool_names.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/emission.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/shared/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/shared/tool_input.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/thinking.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_search/emission.py diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/emission.py new file mode 100644 index 000000000..762f75cca --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/emission.py @@ -0,0 +1,28 @@ +"""generate_image: tool card + terminal summary.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + payload = out if isinstance(out, dict) else {"result": out} + yield ctx.emit_tool_output_card(payload) + if isinstance(out, dict): + if out.get("error"): + yield ctx.streaming_service.format_terminal_info( + f"Image generation failed: {out['error'][:60]}", + "error", + ) + else: + yield ctx.streaming_service.format_terminal_info( + "Image generated successfully", + "success", + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/thinking.py new file mode 100644 index 000000000..9675cb0f2 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_image/thinking.py @@ -0,0 +1,39 @@ +"""generate_image: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.deliverables.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + prompt = d.get("prompt", "") if isinstance(tool_input, dict) else str(tool_input) + return ToolStartThinking( + title="Generating image", + items=[f"Prompt: {prompt[:80]}{'...' if len(prompt) > 80 else ''}"], + ) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_name + items = last_items + if isinstance(tool_output, dict) and not tool_output.get("error"): + completed = [*items, "Image generated successfully"] + else: + error_msg = ( + tool_output.get("error", "Generation failed") + if isinstance(tool_output, dict) + else "Generation failed" + ) + completed = [*items, f"Error: {error_msg}"] + return ("Generating image", completed) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/emission.py new file mode 100644 index 000000000..f1a1e9c37 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/emission.py @@ -0,0 +1,37 @@ +"""generate_podcast: tool card + queue / success / failure terminal lines.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + payload = out if isinstance(out, dict) else {"result": out} + yield ctx.emit_tool_output_card(payload) + if isinstance(out, dict) and out.get("status") in ( + "pending", + "generating", + "processing", + ): + yield ctx.streaming_service.format_terminal_info( + f"Podcast queued: {out.get('title', 'Podcast')}", + "success", + ) + elif isinstance(out, dict) and out.get("status") in ("ready", "success"): + yield ctx.streaming_service.format_terminal_info( + f"Podcast generated successfully: {out.get('title', 'Podcast')}", + "success", + ) + elif isinstance(out, dict) and out.get("status") in ("failed", "error"): + error_msg = out.get("error", "Unknown error") + yield ctx.streaming_service.format_terminal_info( + f"Podcast generation failed: {error_msg}", + "error", + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/thinking.py new file mode 100644 index 000000000..b92e0c91f --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_podcast/thinking.py @@ -0,0 +1,80 @@ +"""generate_podcast: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.deliverables.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + podcast_title = ( + d.get("podcast_title", "SurfSense Podcast") + if isinstance(tool_input, dict) + else "SurfSense Podcast" + ) + content_len = len( + d.get("source_content", "") if isinstance(tool_input, dict) else "" + ) + return ToolStartThinking( + title="Generating podcast", + items=[ + f"Title: {podcast_title}", + f"Content: {content_len:,} characters", + "Preparing audio generation...", + ], + ) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_name + items = last_items + podcast_status = ( + tool_output.get("status", "unknown") + if isinstance(tool_output, dict) + else "unknown" + ) + podcast_title = ( + tool_output.get("title", "Podcast") + if isinstance(tool_output, dict) + else "Podcast" + ) + if podcast_status in ("pending", "generating", "processing"): + completed = [ + f"Title: {podcast_title}", + "Podcast generation started", + "Processing in background...", + ] + elif podcast_status == "already_generating": + completed = [ + f"Title: {podcast_title}", + "Podcast already in progress", + "Please wait for it to complete", + ] + elif podcast_status in ("failed", "error"): + error_msg = ( + tool_output.get("error", "Unknown error") + if isinstance(tool_output, dict) + else "Unknown error" + ) + completed = [ + f"Title: {podcast_title}", + f"Error: {error_msg[:50]}", + ] + elif podcast_status in ("ready", "success"): + completed = [ + f"Title: {podcast_title}", + "Podcast ready", + ] + else: + completed = items + return ("Generating podcast", completed) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/emission.py new file mode 100644 index 000000000..1c5c71b8b --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/emission.py @@ -0,0 +1,33 @@ +"""generate_report: full payload + terminal line.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + payload = out if isinstance(out, dict) else {"result": out} + yield ctx.emit_tool_output_card(payload) + if isinstance(out, dict) and out.get("status") == "ready": + word_count = out.get("word_count", 0) + yield ctx.streaming_service.format_terminal_info( + f"Report generated: {out.get('title', 'Report')} ({word_count:,} words)", + "success", + ) + else: + error_msg = ( + out.get("error", "Unknown error") + if isinstance(out, dict) + else "Unknown error" + ) + yield ctx.streaming_service.format_terminal_info( + f"Report generation failed: {error_msg}", + "error", + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/thinking.py new file mode 100644 index 000000000..f912350f8 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_report/thinking.py @@ -0,0 +1,77 @@ +"""generate_report: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.deliverables.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + report_topic = ( + d.get("topic", "Report") if isinstance(tool_input, dict) else "Report" + ) + is_revision = bool( + isinstance(tool_input, dict) and tool_input.get("parent_report_id") + ) + step_title = "Revising report" if is_revision else "Generating report" + return ToolStartThinking( + title=step_title, + items=[f"Topic: {report_topic}", "Analyzing source content..."], + ) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_name + items = last_items + report_status = ( + tool_output.get("status", "unknown") + if isinstance(tool_output, dict) + else "unknown" + ) + report_title = ( + tool_output.get("title", "Report") + if isinstance(tool_output, dict) + else "Report" + ) + word_count = ( + tool_output.get("word_count", 0) + if isinstance(tool_output, dict) + else 0 + ) + is_revision = ( + tool_output.get("is_revision", False) + if isinstance(tool_output, dict) + else False + ) + step_title = "Revising report" if is_revision else "Generating report" + + if report_status == "ready": + completed = [ + f"Topic: {report_title}", + f"{word_count:,} words", + "Report ready", + ] + elif report_status == "failed": + error_msg = ( + tool_output.get("error", "Unknown error") + if isinstance(tool_output, dict) + else "Unknown error" + ) + completed = [ + f"Topic: {report_title}", + f"Error: {error_msg[:50]}", + ] + else: + completed = items + + return (step_title, completed) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/emission.py new file mode 100644 index 000000000..dc8d3c7fc --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/emission.py @@ -0,0 +1,32 @@ +"""generate_resume: full payload + terminal line.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + payload = out if isinstance(out, dict) else {"result": out} + yield ctx.emit_tool_output_card(payload) + if isinstance(out, dict) and out.get("status") == "ready": + yield ctx.streaming_service.format_terminal_info( + f"Resume generated: {out.get('title', 'Resume')}", + "success", + ) + else: + error_msg = ( + out.get("error", "Unknown error") + if isinstance(out, dict) + else "Unknown error" + ) + yield ctx.streaming_service.format_terminal_info( + f"Resume generation failed: {error_msg}", + "error", + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/thinking.py new file mode 100644 index 000000000..e81a80679 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_resume/thinking.py @@ -0,0 +1,24 @@ +"""generate_resume: generic thinking titles and items.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.default import ( + thinking as default_thinking, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + return default_thinking.resolve_start_thinking(tool_name, tool_input) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + return default_thinking.resolve_completed_thinking( + tool_name, tool_output, last_items + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/emission.py new file mode 100644 index 000000000..21e27d4c3 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/emission.py @@ -0,0 +1,28 @@ +"""generate_video_presentation: tool card + terminal line.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + payload = out if isinstance(out, dict) else {"result": out} + yield ctx.emit_tool_output_card(payload) + if isinstance(out, dict) and out.get("status") == "pending": + yield ctx.streaming_service.format_terminal_info( + f"Video presentation queued: {out.get('title', 'Presentation')}", + "success", + ) + elif isinstance(out, dict) and out.get("status") == "failed": + error_msg = out.get("error", "Unknown error") + yield ctx.streaming_service.format_terminal_info( + f"Presentation generation failed: {error_msg}", + "error", + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/thinking.py new file mode 100644 index 000000000..5c5aa977d --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/generate_video_presentation/thinking.py @@ -0,0 +1,52 @@ +"""generate_video_presentation: generic in-progress thinking; completion is status-driven.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.default import ( + thinking as default_thinking, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + return default_thinking.resolve_start_thinking(tool_name, tool_input) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_name + items = last_items + vp_status = ( + tool_output.get("status", "unknown") + if isinstance(tool_output, dict) + else "unknown" + ) + vp_title = ( + tool_output.get("title", "Presentation") + if isinstance(tool_output, dict) + else "Presentation" + ) + if vp_status in ("pending", "generating"): + completed = [ + f"Title: {vp_title}", + "Presentation generation started", + "Processing in background...", + ] + elif vp_status == "failed": + error_msg = ( + tool_output.get("error", "Unknown error") + if isinstance(tool_output, dict) + else "Unknown error" + ) + completed = [ + f"Title: {vp_title}", + f"Error: {error_msg[:50]}", + ] + else: + completed = items + return ("Generating video presentation", completed) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/emission.py new file mode 100644 index 000000000..68c93dede --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/emission.py @@ -0,0 +1,16 @@ +"""save_document: default completion card and terminal line.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.default import emission as _default +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + yield from _default.iter_completion_emission_frames(ctx) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/thinking.py new file mode 100644 index 000000000..77059a28c --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/save_document/thinking.py @@ -0,0 +1,38 @@ +"""save_document: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.deliverables.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + doc_title = d.get("title", "") if isinstance(tool_input, dict) else str(tool_input) + display_title = doc_title[:60] + ("…" if len(doc_title) > 60 else "") + return ToolStartThinking(title="Saving document", items=[display_title]) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_name + items = last_items + result_str = ( + tool_output.get("result", "") + if isinstance(tool_output, dict) + else str(tool_output) + ) + is_error = "Error" in result_str + completed = [ + *items, + result_str[:80] if is_error else "Saved to knowledge base", + ] + return ("Saving document", completed) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/shared/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/shared/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/shared/tool_input.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/shared/tool_input.py new file mode 100644 index 000000000..1303cf09f --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/shared/tool_input.py @@ -0,0 +1,9 @@ +"""Tool-call args for deliverable thinking modules.""" + +from __future__ import annotations + +from typing import Any + + +def as_tool_input_dict(tool_input: Any) -> dict[str, Any]: + return tool_input if isinstance(tool_input, dict) else {} diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/tool_names.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/tool_names.py new file mode 100644 index 000000000..5924af196 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/deliverables/tool_names.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +DELIVERABLE_TOOLS: frozenset[str] = frozenset( + { + "generate_image", + "generate_podcast", + "generate_report", + "generate_resume", + "generate_video_presentation", + "save_document", + } +) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/emission.py new file mode 100644 index 000000000..293d2a1e9 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/emission.py @@ -0,0 +1,43 @@ +"""scrape_webpage: redacted payload + terminal summary.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + if isinstance(out, dict): + display_output = {k: v for k, v in out.items() if k != "content"} + if "content" in out: + content = out.get("content", "") + display_output["content_preview"] = ( + content[:500] + "..." if len(content) > 500 else content + ) + yield ctx.emit_tool_output_card(display_output) + else: + yield ctx.emit_tool_output_card({"result": out}) + + if isinstance(out, dict) and "error" not in out: + title = out.get("title", "Webpage") + word_count = out.get("word_count", 0) + yield ctx.streaming_service.format_terminal_info( + f"Scraped: {title[:40]}{'...' if len(title) > 40 else ''} ({word_count:,} words)", + "success", + ) + else: + error_msg = ( + out.get("error", "Failed to scrape") + if isinstance(out, dict) + else "Failed to scrape" + ) + yield ctx.streaming_service.format_terminal_info( + f"Scrape failed: {error_msg}", + "error", + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/shared/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/shared/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/shared/tool_input.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/shared/tool_input.py new file mode 100644 index 000000000..581f0e64a --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/shared/tool_input.py @@ -0,0 +1,9 @@ +"""Tool-call args for scrape_webpage thinking.""" + +from __future__ import annotations + +from typing import Any + + +def as_tool_input_dict(tool_input: Any) -> dict[str, Any]: + return tool_input if isinstance(tool_input, dict) else {} diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/thinking.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/thinking.py new file mode 100644 index 000000000..335cc9703 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/scrape_webpage/thinking.py @@ -0,0 +1,47 @@ +"""scrape_webpage: thinking-step copy.""" + +from __future__ import annotations + +from typing import Any + +from app.tasks.chat.streaming.handlers.tools.scrape_webpage.shared.tool_input import ( + as_tool_input_dict, +) +from app.tasks.chat.streaming.handlers.tools.shared.model import ( + ToolStartThinking, +) + + +def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking: + del tool_name + d = as_tool_input_dict(tool_input) + url = d.get("url", "") if isinstance(tool_input, dict) else str(tool_input) + return ToolStartThinking( + title="Scraping webpage", + items=[f"URL: {url[:80]}{'...' if len(url) > 80 else ''}"], + ) + + +def resolve_completed_thinking( + tool_name: str, tool_output: Any, last_items: list[str], +) -> tuple[str, list[str]]: + del tool_name + items = last_items + if isinstance(tool_output, dict): + title = tool_output.get("title", "Webpage") + word_count = tool_output.get("word_count", 0) + has_error = "error" in tool_output + if has_error: + completed = [ + *items, + f"Error: {tool_output.get('error', 'Failed to scrape')[:50]}", + ] + else: + completed = [ + *items, + f"Title: {title[:50]}{'...' if len(title) > 50 else ''}", + f"Extracted: {word_count:,} words", + ] + else: + completed = [*items, "Content extracted"] + return ("Scraping webpage", completed) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_search/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_search/emission.py new file mode 100644 index 000000000..eccaed708 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_search/emission.py @@ -0,0 +1,41 @@ +"""web_search: citations parsed from provider XML.""" + +from __future__ import annotations + +import re +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + xml = out.get("result", str(out)) if isinstance(out, dict) else str(out) + citations: dict[str, dict[str, str]] = {} + for m in re.finditer( + r"<!\[CDATA\[(.*?)\]\]>\s*", + xml, + ): + title, url = m.group(1).strip(), m.group(2).strip() + if url.startswith("http") and url not in citations: + citations[url] = {"title": title} + for m in re.finditer( + r"", + xml, + ): + chunk_url, content = m.group(1).strip(), m.group(2).strip() + if ( + chunk_url.startswith("http") + and chunk_url in citations + and content + ): + citations[chunk_url]["snippet"] = ( + content[:200] + "…" if len(content) > 200 else content + ) + yield ctx.emit_tool_output_card( + {"status": "completed", "citations": citations}, + ) From ec26ca69a6b4f05471beb8a3c114f689991f5622 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 49/96] Add chat EventRelay and orchestrator stubs for future cutover. --- .../app/tasks/chat/streaming/event_relay.py | 127 ++++++++++++++++++ .../app/tasks/chat/streaming/orchestrator.py | 48 +++++++ 2 files changed, 175 insertions(+) create mode 100644 surfsense_backend/app/tasks/chat/streaming/event_relay.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestrator.py diff --git a/surfsense_backend/app/tasks/chat/streaming/event_relay.py b/surfsense_backend/app/tasks/chat/streaming/event_relay.py new file mode 100644 index 000000000..f86337ad7 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/event_relay.py @@ -0,0 +1,127 @@ +"""Turn LangGraph astream_events into SSE strings via the handler modules.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import dataclass, field +from typing import Any + +from app.services.streaming.emitter import EmitterRegistry +from app.tasks.chat.streaming.handlers.chain_end import iter_chain_end_frames +from app.tasks.chat.streaming.handlers.chat_model_stream import ( + iter_chat_model_stream_frames, +) +from app.tasks.chat.streaming.handlers.custom_event_dispatch import ( + iter_custom_event_frames, +) +from app.tasks.chat.streaming.handlers.tool_end import iter_tool_end_frames +from app.tasks.chat.streaming.handlers.tool_start import iter_tool_start_frames +from app.tasks.chat.streaming.relay.state import AgentEventRelayState +from app.tasks.chat.streaming.relay.thinking_step_completion import ( + complete_active_thinking_step, +) +from app.tasks.chat.streaming.stream_result import StreamResult + + +@dataclass +class EventRelayConfig: + """Optional relay tuning (sub-agent tools, text suppression).""" + + subagent_entry_tool_names: frozenset[str] = field( + default_factory=lambda: frozenset({"task"}) + ) + suppress_main_text_inside_tools: bool = True + + +class EventRelay: + """Dispatches graph events to streaming handlers and optional emitters.""" + + def __init__( + self, + *, + streaming_service: Any, + config: EventRelayConfig | None = None, + ) -> None: + self.streaming_service = streaming_service + self.config = config or EventRelayConfig() + reg = getattr(streaming_service, "emitter_registry", None) + self.emitter_registry = reg if reg is not None else EmitterRegistry() + + async def relay( + self, + events: AsyncIterator[dict[str, Any]], + *, + state: AgentEventRelayState, + result: StreamResult, + step_prefix: str = "thinking", + content_builder: Any | None = None, + config: dict[str, Any] | None = None, + ) -> AsyncIterator[str]: + """Yield SSE for each event from the async iterator, then finalize text/thinking.""" + graph_config = config or {} + async for event in events: + event_type = event.get("event", "") + if event_type == "on_chat_model_stream": + for frame in iter_chat_model_stream_frames( + event, + state=state, + streaming_service=self.streaming_service, + content_builder=content_builder, + step_prefix=step_prefix, + ): + yield frame + elif event_type == "on_tool_start": + for frame in iter_tool_start_frames( + event, + state=state, + streaming_service=self.streaming_service, + content_builder=content_builder, + result=result, + step_prefix=step_prefix, + ): + yield frame + elif event_type == "on_tool_end": + for frame in iter_tool_end_frames( + event, + state=state, + streaming_service=self.streaming_service, + content_builder=content_builder, + result=result, + step_prefix=step_prefix, + config=graph_config, + ): + yield frame + elif event_type == "on_custom_event": + for frame in iter_custom_event_frames( + event, + state=state, + streaming_service=self.streaming_service, + content_builder=content_builder, + ): + yield frame + elif event_type in ("on_chain_end", "on_agent_end"): + for frame in iter_chain_end_frames( + event, + state=state, + streaming_service=self.streaming_service, + content_builder=content_builder, + ): + yield frame + + if state.current_text_id is not None: + yield self.streaming_service.format_text_end(state.current_text_id) + if content_builder is not None: + content_builder.on_text_end(state.current_text_id) + state.current_text_id = None + + completion_event, new_active = complete_active_thinking_step( + streaming_service=self.streaming_service, + content_builder=content_builder, + last_active_step_id=state.last_active_step_id, + last_active_step_title=state.last_active_step_title, + last_active_step_items=state.last_active_step_items, + completed_step_ids=state.completed_step_ids, + ) + if completion_event: + yield completion_event + state.last_active_step_id = new_active diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/orchestrator.py new file mode 100644 index 000000000..1b8558bc6 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/orchestrator.py @@ -0,0 +1,48 @@ +"""Top-level chat streaming entrypoints (stubs until wired).""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator +from typing import Any + + +async def stream_chat( + *, + request: Any, + user: Any, + db_session: Any, +) -> AsyncGenerator[str, None]: # pragma: no cover - orchestrator port in progress + del request, user, db_session + raise NotImplementedError( + "stream_chat: orchestrator not wired yet" + ) + if False: # pragma: no cover + yield "" + + +async def stream_resume( + *, + request: Any, + user: Any, + db_session: Any, +) -> AsyncGenerator[str, None]: # pragma: no cover - orchestrator port in progress + del request, user, db_session + raise NotImplementedError( + "stream_resume: orchestrator not wired yet" + ) + if False: # pragma: no cover + yield "" + + +async def stream_regenerate( + *, + request: Any, + user: Any, + db_session: Any, +) -> AsyncGenerator[str, None]: # pragma: no cover - orchestrator port in progress + del request, user, db_session + raise NotImplementedError( + "stream_regenerate: orchestrator not wired yet" + ) + if False: # pragma: no cover + yield "" From 8b6ffd12b8649bd789a9e780dd90a0a64d04fbac Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:08:48 +0200 Subject: [PATCH 50/96] Add parity unit tests for extracted chat streaming vs legacy. --- .../unit/tasks/chat/streaming/__init__.py | 0 .../chat/streaming/test_stage_1_parity.py | 292 ++++++++++++++++++ .../chat/streaming/test_stage_2_parity.py | 240 ++++++++++++++ 3 files changed, 532 insertions(+) create mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/__init__.py create mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py create mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/__init__.py b/surfsense_backend/tests/unit/tasks/chat/streaming/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py new file mode 100644 index 000000000..9207f37d1 --- /dev/null +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py @@ -0,0 +1,292 @@ +"""Pin Stage 1 extractions as faithful copies of the old helpers. + +The new orchestrator under ``app.tasks.chat.streaming`` is built in +parallel with the production module ``app.tasks.chat.stream_new_chat``. +For each Stage 1 extraction we assert the new function returns the same +output as the old one for a representative input set. The moment the +two diverge - intentionally or otherwise - this file fails loudly so +the divergence is reviewed rather than shipped silently. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import Any + +import pytest + +from app.agents.new_chat.errors import BusyError +from app.agents.new_chat.middleware.busy_mutex import request_cancel, reset_cancel +from app.tasks.chat.stream_new_chat import ( + _classify_stream_exception as old_classify, + _emit_stream_terminal_error as old_emit_terminal_error, + _extract_chunk_parts as old_extract_chunk_parts, + _extract_resolved_file_path as old_extract_resolved_file_path, + _first_interrupt_value as old_first_interrupt_value, + _tool_output_has_error as old_tool_output_has_error, + _tool_output_to_text as old_tool_output_to_text, +) +from app.tasks.chat.streaming.errors.classifier import ( + classify_stream_exception as new_classify, +) +from app.tasks.chat.streaming.errors.emitter import ( + emit_stream_terminal_error as new_emit_terminal_error, +) +from app.tasks.chat.streaming.helpers.chunk_parts import ( + extract_chunk_parts as new_extract_chunk_parts, +) +from app.tasks.chat.streaming.helpers.interrupt_inspector import ( + first_interrupt_value as new_first_interrupt_value, +) +from app.tasks.chat.streaming.helpers.tool_output import ( + extract_resolved_file_path as new_extract_resolved_file_path, + tool_output_has_error as new_tool_output_has_error, + tool_output_to_text as new_tool_output_to_text, +) + +pytestmark = pytest.mark.unit + + +# ---------------------------------------------------------------- chunk parts + + +@dataclass +class _Chunk: + content: Any = "" + additional_kwargs: dict[str, Any] = field(default_factory=dict) + tool_call_chunks: list[dict[str, Any]] = field(default_factory=list) + + +_CHUNK_CASES: list[Any] = [ + None, + _Chunk(content=""), + _Chunk(content="hello"), + _Chunk(content=42), # invalid type, defensively coerced to empty + _Chunk( + content=[ + {"type": "text", "text": "Hello "}, + {"type": "text", "text": "world"}, + ] + ), + _Chunk( + content=[ + {"type": "reasoning", "reasoning": "hmm "}, + {"type": "reasoning", "text": "still"}, + {"type": "text", "text": "answer"}, + ] + ), + _Chunk( + content=[ + {"type": "tool_call_chunk", "id": "c1", "name": "x", "args": "{"}, + {"type": "tool_use", "id": "c2", "name": "y"}, + {"type": "image_url", "url": "ignored"}, + ] + ), + _Chunk( + content="visible", + additional_kwargs={"reasoning_content": "private"}, + ), + _Chunk( + tool_call_chunks=[ + {"id": None, "name": None, "args": '{"a":1}', "index": 0}, + {"id": "c", "name": "n", "args": "}", "index": 0}, + ] + ), + _Chunk( + content=[{"type": "tool_call_chunk", "id": "from-block", "name": "x"}], + tool_call_chunks=[{"id": "from-attr", "name": "y"}], + ), +] + + +@pytest.mark.parametrize("chunk", _CHUNK_CASES) +def test_extract_chunk_parts_matches_old_implementation(chunk: Any) -> None: + assert new_extract_chunk_parts(chunk) == old_extract_chunk_parts(chunk) + + +# ---------------------------------------------------------- interrupt inspector + + +@dataclass +class _Interrupt: + value: dict[str, Any] + + +@dataclass +class _Task: + interrupts: tuple[Any, ...] = () + + +@dataclass +class _State: + tasks: tuple[Any, ...] = () + interrupts: tuple[Any, ...] = () + + +_INTERRUPT_CASES: list[Any] = [ + _State(), + _State(tasks=(_Task(interrupts=(_Interrupt(value={"name": "send"}),)),)), + # Multiple tasks: must return the FIRST one in iteration order. + _State( + tasks=( + _Task(interrupts=(_Interrupt(value={"name": "first"}),)), + _Task(interrupts=(_Interrupt(value={"name": "second"}),)), + ) + ), + # Empty task interrupts -> falls back to root state.interrupts. + _State( + tasks=(_Task(interrupts=()),), + interrupts=(_Interrupt(value={"name": "root"}),), + ), + # Interrupts as plain dicts (not wrapper objects). + _State(interrupts=({"value": {"name": "dict_root"}},)), + # A defective task whose `.interrupts` raises - must be tolerated. + _State(tasks=(object(),)), +] + + +@pytest.mark.parametrize("state", _INTERRUPT_CASES) +def test_first_interrupt_value_matches_old_implementation(state: Any) -> None: + assert new_first_interrupt_value(state) == old_first_interrupt_value(state) + + +# ----------------------------------------------------------- error classifier + + +def _classify_cases() -> list[Exception]: + """Inputs that the FE depends on being mapped to specific error codes.""" + return [ + Exception("totally generic error"), + Exception( + '{"error":{"type":"rate_limit_error","message":"slow down"}}' + ), + Exception( + 'OpenrouterException - {"error":{"message":"Provider returned error",' + '"code":429}}' + ), + BusyError(request_id="thread-busy-parity"), + Exception("Thread is busy with another request"), + ] + + +@pytest.mark.parametrize("exc", _classify_cases()) +def test_classify_stream_exception_matches_old_implementation( + exc: Exception, +) -> None: + new = new_classify(exc, flow_label="parity-test") + old = old_classify(exc, flow_label="parity-test") + # Strip the wall-clock retry timestamp before comparing — both + # implementations call ``time.time()`` independently and the call + # order is enough to differ by 1 ms in practice. Every other field + # in the tuple must match exactly. + new_extra = dict(new[5]) if isinstance(new[5], dict) else new[5] + old_extra = dict(old[5]) if isinstance(old[5], dict) else old[5] + if isinstance(new_extra, dict) and isinstance(old_extra, dict): + new_extra.pop("retry_after_at", None) + old_extra.pop("retry_after_at", None) + assert new[:5] == old[:5] + assert new_extra == old_extra + + +def test_classify_turn_cancelling_branch_parity() -> None: + """The TURN_CANCELLING branch reads cancel state for the busy thread id; + both implementations must agree on retry-window semantics, not just the + plain THREAD_BUSY code.""" + thread_id = "parity-cancelling-thread" + reset_cancel(thread_id) + request_cancel(thread_id) + exc = BusyError(request_id=thread_id) + new = new_classify(exc, flow_label="parity-test") + old = old_classify(exc, flow_label="parity-test") + assert new[0] == old[0] == "thread_busy" + assert new[1] == old[1] == "TURN_CANCELLING" + assert isinstance(new[5], dict) and isinstance(old[5], dict) + assert new[5]["retry_after_ms"] == old[5]["retry_after_ms"] + + +# ------------------------------------------------------------ terminal emitter + + +class _FakeStreamingService: + """Duck-types ``format_error`` for both old and new emitters.""" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def format_error( + self, message: str, *, error_code: str, extra: dict[str, Any] | None = None + ) -> str: + self.calls.append( + {"message": message, "error_code": error_code, "extra": extra} + ) + return f"data: {{\"type\":\"error\",\"errorText\":\"{message}\"}}\n\n" + + +def test_emit_stream_terminal_error_matches_old_output_and_logs(caplog) -> None: + """The new emitter must produce the same SSE frame and log the same + structured payload as the old one for the same arguments.""" + args: dict[str, Any] = { + "flow": "new", + "request_id": "req-parity", + "thread_id": 7, + "search_space_id": 9, + "user_id": "user-parity", + "message": "boom", + "error_kind": "server_error", + "error_code": "SERVER_ERROR", + "severity": "error", + "is_expected": False, + "extra": {"foo": "bar"}, + } + + new_svc = _FakeStreamingService() + old_svc = _FakeStreamingService() + + with caplog.at_level(logging.ERROR): + new_frame = new_emit_terminal_error(streaming_service=new_svc, **args) + old_frame = old_emit_terminal_error(streaming_service=old_svc, **args) + + assert new_frame == old_frame + assert new_svc.calls == old_svc.calls + chat_error_records = [ + r for r in caplog.records if "[chat_stream_error]" in r.message + ] + # One log line per emit call (two emits -> two records). + assert len(chat_error_records) == 2 + + +# ---------------------------------------------------------------- tool output + + +def test_tool_output_helpers_match_old_implementation() -> None: + samples: list[Any] = [ + {"result": "ok"}, + {"error": "bad"}, + {"result": "Error: x"}, + "Error: plain", + "fine", + {"nested": {"a": 1}}, + ] + for s in samples: + assert new_tool_output_to_text(s) == old_tool_output_to_text(s) + assert new_tool_output_has_error(s) == old_tool_output_has_error(s) + + assert new_extract_resolved_file_path( + tool_name="write_file", + tool_output={"path": " /tmp/x "}, + tool_input=None, + ) == old_extract_resolved_file_path( + tool_name="write_file", + tool_output={"path": " /tmp/x "}, + tool_input=None, + ) + assert new_extract_resolved_file_path( + tool_name="write_file", + tool_output={}, + tool_input={"file_path": " /fallback "}, + ) == old_extract_resolved_file_path( + tool_name="write_file", + tool_output={}, + tool_input={"file_path": " /fallback "}, + ) diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py new file mode 100644 index 000000000..892bb7a6a --- /dev/null +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py @@ -0,0 +1,240 @@ +"""Parity tests for Stage 2 extractions (tool matching, thinking step, custom events).""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from app.tasks.chat.stream_new_chat import _legacy_match_lc_id as old_legacy_match +from app.tasks.chat.streaming.handlers.custom_events import ( + handle_action_log, + handle_action_log_updated, + handle_document_created, + handle_report_progress, +) +from app.tasks.chat.streaming.helpers.tool_call_matching import ( + match_buffered_langchain_tool_call_id as new_legacy_match, +) +from app.tasks.chat.streaming.relay.state import AgentEventRelayState +from app.tasks.chat.streaming.relay.thinking_step_completion import ( + complete_active_thinking_step, +) +from app.tasks.chat.streaming.relay.thinking_step_sse import emit_thinking_step_frame + +pytestmark = pytest.mark.unit + + +def _copy_chunk_buffer(raw: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [dict(x) for x in raw] + + +def test_legacy_tool_call_match_matches_old_implementation() -> None: + cases: list[tuple[list[dict[str, Any]], str, str, dict[str, str]]] = [ + ( + [ + {"name": "write_file", "id": "lc-a"}, + {"name": "other", "id": "lc-b"}, + ], + "write_file", + "run-1", + {}, + ), + ( + [{"name": "x", "id": None}, {"name": "y", "id": "lc-fallback"}], + "write_file", + "run-2", + {}, + ), + ([{"name": "no_id"}], "write_file", "run-3", {}), + ] + for chunks_template, tool_name, run_id, lc_map_seed in cases: + old_chunks = _copy_chunk_buffer(chunks_template) + new_chunks = _copy_chunk_buffer(chunks_template) + old_map = dict(lc_map_seed) + new_map = dict(lc_map_seed) + old_out = old_legacy_match(old_chunks, tool_name, run_id, old_map) + new_out = new_legacy_match(new_chunks, tool_name, run_id, new_map) + assert new_out == old_out + assert new_chunks == old_chunks + assert new_map == old_map + + +def test_emit_thinking_step_frame_invokes_builder_before_service() -> None: + order: list[str] = [] + builder = MagicMock() + + def on_ts(*args: Any, **kwargs: Any) -> None: + order.append("builder") + + builder.on_thinking_step.side_effect = on_ts + + svc = MagicMock() + + def fmt(**kwargs: Any) -> str: + order.append("service") + return "frame" + + svc.format_thinking_step.side_effect = fmt + + out = emit_thinking_step_frame( + streaming_service=svc, + content_builder=builder, + step_id="thinking-1", + title="Working", + status="in_progress", + items=["a"], + ) + assert out == "frame" + assert order == ["builder", "service"] + builder.on_thinking_step.assert_called_once() + svc.format_thinking_step.assert_called_once() + + +def test_emit_thinking_step_frame_skips_builder_when_none() -> None: + svc = MagicMock(return_value="x") + svc.format_thinking_step.return_value = "frame" + assert ( + emit_thinking_step_frame( + streaming_service=svc, + content_builder=None, + step_id="s", + title="t", + ) + == "frame" + ) + svc.format_thinking_step.assert_called_once() + + +def test_complete_active_thinking_step_mirrors_closure_semantics() -> None: + svc = MagicMock() + svc.format_thinking_step.return_value = "done-frame" + completed: set[str] = set() + + frame, new_id = complete_active_thinking_step( + streaming_service=svc, + content_builder=None, + last_active_step_id="thinking-1", + last_active_step_title="T", + last_active_step_items=["x"], + completed_step_ids=completed, + ) + assert frame == "done-frame" + assert new_id is None + assert "thinking-1" in completed + + frame2, id2 = complete_active_thinking_step( + streaming_service=svc, + content_builder=None, + last_active_step_id="thinking-1", + last_active_step_title="T", + last_active_step_items=[], + completed_step_ids=completed, + ) + assert frame2 is None + assert id2 == "thinking-1" + + +def test_agent_event_relay_state_factory_matches_counter_rule() -> None: + s0 = AgentEventRelayState.for_invocation(parity_v2=False) + assert s0.thinking_step_counter == 0 + assert s0.last_active_step_id is None + + s1 = AgentEventRelayState.for_invocation( + initial_step_id="thinking-resume-1", + initial_step_title="Inherited", + initial_step_items=["Topic: X"], + parity_v2=True, + ) + assert s1.thinking_step_counter == 1 + assert s1.last_active_step_id == "thinking-resume-1" + assert s1.parity_v2 is True + assert s1.next_thinking_step_id("thinking") == "thinking-2" + + +@pytest.mark.parametrize( + ("phase", "message", "start_items", "expected_tail"), + [ + ( + "revising_section", + "progress line", + ["Topic: Foo", "Modifying bar", "stale..."], + ["Topic: Foo", "Modifying bar", "progress line"], + ), + ( + "other", + "phase msg", + ["Topic: Foo", "old line"], + ["Topic: Foo", "phase msg"], + ), + ], +) +def test_report_progress_items_match_reference( + phase: str, + message: str, + start_items: list[str], + expected_tail: list[str], +) -> None: + svc = MagicMock() + svc.format_thinking_step.return_value = "sse" + + items = list(start_items) + frame, new_items = handle_report_progress( + {"message": message, "phase": phase}, + last_active_step_id="step-1", + last_active_step_title="Report", + last_active_step_items=items, + streaming_service=svc, + content_builder=None, + ) + assert frame == "sse" + assert new_items == expected_tail + kwargs = svc.format_thinking_step.call_args.kwargs + assert kwargs["items"] == expected_tail + + +def test_report_progress_noop_when_missing_message_or_step() -> None: + svc = MagicMock() + items = ["Topic: A"] + f1, i1 = handle_report_progress( + {"message": "", "phase": "x"}, + last_active_step_id="s", + last_active_step_title="t", + last_active_step_items=items, + streaming_service=svc, + content_builder=None, + ) + assert f1 is None and i1 is items + + f2, i2 = handle_report_progress( + {"message": "m", "phase": "x"}, + last_active_step_id=None, + last_active_step_title="t", + last_active_step_items=items, + streaming_service=svc, + content_builder=None, + ) + assert f2 is None and i2 is items + + +def test_document_action_handlers_match_format_data_guards() -> None: + svc = MagicMock() + svc.format_data.return_value = "data-frame" + + assert handle_document_created({}, streaming_service=svc) is None + assert handle_document_created({"id": 0}, streaming_service=svc) is None + handle_document_created({"id": 42, "title": "x"}, streaming_service=svc) + svc.format_data.assert_called_with( + "documents-updated", {"action": "created", "document": {"id": 42, "title": "x"}} + ) + + svc.reset_mock() + assert handle_action_log({"id": None}, streaming_service=svc) is None + handle_action_log({"id": 1}, streaming_service=svc) + svc.format_data.assert_called_once_with("action-log", {"id": 1}) + + svc.reset_mock() + assert handle_action_log_updated({"id": None}, streaming_service=svc) is None + handle_action_log_updated({"id": 2}, streaming_service=svc) + svc.format_data.assert_called_once_with("action-log-updated", {"id": 2}) From 2ec2e82d9d3215734bb367e3cd854035ae078919 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Wed, 6 May 2026 20:10:09 +0200 Subject: [PATCH 51/96] Configure workspace Python interpreter and pyright extra paths. --- .vscode/settings.json | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index 05bd30702..7da4b54f8 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,4 +1,9 @@ { "biome.configurationPath": "./surfsense_web/biome.json", - "deepscan.ignoreConfirmWarning": true + "deepscan.ignoreConfirmWarning": true, + "python.defaultInterpreterPath": "${workspaceFolder}/surfsense_backend/.venv/bin/python", + "basedpyright.analysis.extraPaths": [ + "${workspaceFolder}/surfsense_backend" + ], + "python-envs.pythonProjects": [] } \ No newline at end of file From c0706364d15e25a42d2973a890a106e9df6ccff0 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 7 May 2026 14:44:36 +0200 Subject: [PATCH 52/96] Add a route-level kill switch for streaming orchestrator cutover. --- surfsense_backend/app/config/__init__.py | 6 + .../app/routes/new_chat_routes.py | 33 +++- .../app/tasks/chat/streaming/orchestrator.py | 143 ++++++++++++++---- 3 files changed, 146 insertions(+), 36 deletions(-) diff --git a/surfsense_backend/app/config/__init__.py b/surfsense_backend/app/config/__init__.py index f6f0c7f62..543524456 100644 --- a/surfsense_backend/app/config/__init__.py +++ b/surfsense_backend/app/config/__init__.py @@ -490,6 +490,12 @@ class Config: ENABLE_DESKTOP_LOCAL_FILESYSTEM = ( os.getenv("ENABLE_DESKTOP_LOCAL_FILESYSTEM", "FALSE").upper() == "TRUE" ) + # Streaming entrypoint switch. Keep this at the route layer so orchestrator + # code stays free of legacy fallback branching. + ENABLE_CHAT_STREAM_ORCHESTRATOR = ( + os.getenv("SURFSENSE_ENABLE_CHAT_STREAM_ORCHESTRATOR", "TRUE").upper() + == "TRUE" + ) @classmethod def is_self_hosted(cls) -> bool: diff --git a/surfsense_backend/app/routes/new_chat_routes.py b/surfsense_backend/app/routes/new_chat_routes.py index ad96654f5..7f035daef 100644 --- a/surfsense_backend/app/routes/new_chat_routes.py +++ b/surfsense_backend/app/routes/new_chat_routes.py @@ -71,7 +71,15 @@ from app.schemas.new_chat import ( TokenUsageSummary, TurnStatusResponse, ) -from app.tasks.chat.stream_new_chat import stream_new_chat, stream_resume_chat +from app.tasks.chat.stream_new_chat import ( + stream_new_chat as legacy_stream_new_chat, + stream_resume_chat as legacy_stream_resume_chat, +) +from app.tasks.chat.streaming.orchestrator import ( + stream_chat, + stream_regenerate, + stream_resume, +) from app.users import current_active_user from app.utils.perf import get_perf_logger from app.utils.rbac import check_permission @@ -90,6 +98,10 @@ TURN_CANCELLING_MAX_DELAY_MS = 1500 router = APIRouter() +def _use_streaming_orchestrator() -> bool: + return config.ENABLE_CHAT_STREAM_ORCHESTRATOR + + def _resolve_filesystem_selection( *, mode: str, @@ -1770,7 +1782,11 @@ async def handle_new_chat( ) return StreamingResponse( - stream_new_chat( + ( + stream_chat + if _use_streaming_orchestrator() + else legacy_stream_new_chat + )( user_query=request.user_query, search_space_id=request.search_space_id, chat_id=request.chat_id, @@ -2255,7 +2271,12 @@ async def regenerate_response( else None ) try: - async for chunk in stream_new_chat( + regenerate_fn = ( + stream_regenerate + if _use_streaming_orchestrator() + else legacy_stream_new_chat + ) + async for chunk in regenerate_fn( user_query=str(user_query_to_use), search_space_id=request.search_space_id, chat_id=thread_id, @@ -2387,7 +2408,11 @@ async def resume_chat( await session.close() return StreamingResponse( - stream_resume_chat( + ( + stream_resume + if _use_streaming_orchestrator() + else legacy_stream_resume_chat + )( chat_id=thread_id, search_space_id=request.search_space_id, decisions=decisions, diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/orchestrator.py index 1b8558bc6..e912dd632 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestrator.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestrator.py @@ -1,48 +1,127 @@ -"""Top-level chat streaming entrypoints (stubs until wired).""" +"""Top-level chat streaming entrypoints. + +For now these orchestrator functions are thin compatibility wrappers around the +current ``stream_new_chat`` / ``stream_resume_chat`` implementations. Routing +calls through this module lets us cut over to the fully modular event relay in +one place later without touching API routes again. +""" from __future__ import annotations from collections.abc import AsyncGenerator -from typing import Any +from typing import Any, Literal + +from app.agents.new_chat.filesystem_selection import FilesystemSelection +from app.db import ChatVisibility +from app.tasks.chat.stream_new_chat import stream_new_chat, stream_resume_chat async def stream_chat( *, - request: Any, - user: Any, - db_session: Any, -) -> AsyncGenerator[str, None]: # pragma: no cover - orchestrator port in progress - del request, user, db_session - raise NotImplementedError( - "stream_chat: orchestrator not wired yet" - ) - if False: # pragma: no cover - yield "" + user_query: str, + search_space_id: int, + chat_id: int, + user_id: str | None = None, + llm_config_id: int = -1, + mentioned_document_ids: list[int] | None = None, + mentioned_surfsense_doc_ids: list[int] | None = None, + mentioned_documents: list[dict[str, Any]] | None = None, + checkpoint_id: str | None = None, + needs_history_bootstrap: bool = False, + thread_visibility: ChatVisibility | None = None, + current_user_display_name: str | None = None, + disabled_tools: list[str] | None = None, + filesystem_selection: FilesystemSelection | None = None, + request_id: str | None = None, + user_image_data_urls: list[str] | None = None, +) -> AsyncGenerator[str, None]: + """Stream a new chat turn through the current production pipeline.""" + async for chunk in stream_new_chat( + user_query=user_query, + search_space_id=search_space_id, + chat_id=chat_id, + user_id=user_id, + llm_config_id=llm_config_id, + mentioned_document_ids=mentioned_document_ids, + mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids, + mentioned_documents=mentioned_documents, + checkpoint_id=checkpoint_id, + needs_history_bootstrap=needs_history_bootstrap, + thread_visibility=thread_visibility, + current_user_display_name=current_user_display_name, + disabled_tools=disabled_tools, + filesystem_selection=filesystem_selection, + request_id=request_id, + user_image_data_urls=user_image_data_urls, + ): + yield chunk async def stream_resume( *, - request: Any, - user: Any, - db_session: Any, -) -> AsyncGenerator[str, None]: # pragma: no cover - orchestrator port in progress - del request, user, db_session - raise NotImplementedError( - "stream_resume: orchestrator not wired yet" - ) - if False: # pragma: no cover - yield "" + chat_id: int, + search_space_id: int, + decisions: list[dict], + user_id: str | None = None, + llm_config_id: int = -1, + thread_visibility: ChatVisibility | None = None, + filesystem_selection: FilesystemSelection | None = None, + request_id: str | None = None, + disabled_tools: list[str] | None = None, +) -> AsyncGenerator[str, None]: + """Resume an interrupted chat turn through the current production pipeline.""" + async for chunk in stream_resume_chat( + chat_id=chat_id, + search_space_id=search_space_id, + decisions=decisions, + user_id=user_id, + llm_config_id=llm_config_id, + thread_visibility=thread_visibility, + filesystem_selection=filesystem_selection, + request_id=request_id, + disabled_tools=disabled_tools, + ): + yield chunk async def stream_regenerate( *, - request: Any, - user: Any, - db_session: Any, -) -> AsyncGenerator[str, None]: # pragma: no cover - orchestrator port in progress - del request, user, db_session - raise NotImplementedError( - "stream_regenerate: orchestrator not wired yet" - ) - if False: # pragma: no cover - yield "" + user_query: str, + search_space_id: int, + chat_id: int, + user_id: str | None = None, + llm_config_id: int = -1, + mentioned_document_ids: list[int] | None = None, + mentioned_surfsense_doc_ids: list[int] | None = None, + mentioned_documents: list[dict[str, Any]] | None = None, + checkpoint_id: str | None = None, + needs_history_bootstrap: bool = False, + thread_visibility: ChatVisibility | None = None, + current_user_display_name: str | None = None, + disabled_tools: list[str] | None = None, + filesystem_selection: FilesystemSelection | None = None, + request_id: str | None = None, + user_image_data_urls: list[str] | None = None, + flow: Literal["new", "regenerate"] = "regenerate", +) -> AsyncGenerator[str, None]: + """Regenerate an assistant turn through the current production pipeline.""" + async for chunk in stream_new_chat( + user_query=user_query, + search_space_id=search_space_id, + chat_id=chat_id, + user_id=user_id, + llm_config_id=llm_config_id, + mentioned_document_ids=mentioned_document_ids, + mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids, + mentioned_documents=mentioned_documents, + checkpoint_id=checkpoint_id, + needs_history_bootstrap=needs_history_bootstrap, + thread_visibility=thread_visibility, + current_user_display_name=current_user_display_name, + disabled_tools=disabled_tools, + filesystem_selection=filesystem_selection, + request_id=request_id, + user_image_data_urls=user_image_data_urls, + flow=flow, + ): + yield chunk From 4e664652a86ddcdb81014b92587d5436762cd2fa Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 7 May 2026 15:13:22 +0200 Subject: [PATCH 53/96] Add streaming runtime helpers with behavior-focused unit tests. --- .../app/tasks/chat/streaming/runtime.py | 92 ++++++++++++++ .../streaming/test_orchestrator_runtime.py | 120 ++++++++++++++++++ 2 files changed, 212 insertions(+) create mode 100644 surfsense_backend/app/tasks/chat/streaming/runtime.py create mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_runtime.py diff --git a/surfsense_backend/app/tasks/chat/streaming/runtime.py b/surfsense_backend/app/tasks/chat/streaming/runtime.py new file mode 100644 index 000000000..b45da2789 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/runtime.py @@ -0,0 +1,92 @@ +"""Runtime setup helpers for orchestrated chat streaming.""" + +from __future__ import annotations + +import contextlib +import logging +from collections.abc import Callable +from typing import Any + +_PREFLIGHT_TIMEOUT_SEC: float = 2.5 +_PREFLIGHT_MAX_TOKENS: int = 1 + + +async def preflight_llm( + llm: Any, + *, + is_provider_rate_limited: Callable[[BaseException], bool], +) -> None: + """Issue a minimal completion probe to catch immediate provider 429s.""" + from litellm import acompletion + + model = getattr(llm, "model", None) + if not model or model == "auto": + return + + try: + await acompletion( + model=model, + messages=[{"role": "user", "content": "ping"}], + api_key=getattr(llm, "api_key", None), + api_base=getattr(llm, "api_base", None), + max_tokens=_PREFLIGHT_MAX_TOKENS, + timeout=_PREFLIGHT_TIMEOUT_SEC, + stream=False, + metadata={"tags": ["surfsense:internal", "auto-pin-preflight"]}, + ) + except Exception as exc: + if is_provider_rate_limited(exc): + raise + logging.getLogger(__name__).debug( + "auto_pin_preflight non_rate_limit_error model=%s err=%s", + model, + exc, + ) + + +async def build_main_agent_for_thread( + agent_factory: Any, + *, + llm: Any, + search_space_id: int, + db_session: Any, + connector_service: Any, + checkpointer: Any, + user_id: str | None, + thread_id: int | None, + agent_config: Any, + firecrawl_api_key: str | None, + thread_visibility: Any, + filesystem_selection: Any, + disabled_tools: list[str] | None = None, + mentioned_document_ids: list[int] | None = None, +) -> Any: + """Run one canonical agent-build call for a single thread.""" + return await agent_factory( + llm=llm, + search_space_id=search_space_id, + db_session=db_session, + connector_service=connector_service, + checkpointer=checkpointer, + user_id=user_id, + thread_id=thread_id, + agent_config=agent_config, + firecrawl_api_key=firecrawl_api_key, + thread_visibility=thread_visibility, + filesystem_selection=filesystem_selection, + disabled_tools=disabled_tools, + mentioned_document_ids=mentioned_document_ids, + ) + + +async def settle_speculative_agent_build(task: Any) -> None: + """Wait for a discarded speculative build and swallow its outcome.""" + with contextlib.suppress(BaseException): + await task + + +__all__ = [ + "build_main_agent_for_thread", + "preflight_llm", + "settle_speculative_agent_build", +] diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_runtime.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_runtime.py new file mode 100644 index 000000000..edb05edfa --- /dev/null +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_runtime.py @@ -0,0 +1,120 @@ +"""Behavior tests for streaming runtime helpers.""" + +from __future__ import annotations + +import sys +import types +from typing import Any + +import pytest + +from app.tasks.chat.streaming import runtime + +pytestmark = pytest.mark.unit + + +async def test_preflight_llm_calls_litellm_when_model_present( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: dict[str, Any] = {} + + async def _fake_acompletion(**kwargs: Any): + calls.update(kwargs) + return {"ok": True} + + monkeypatch.setitem( + sys.modules, + "litellm", + types.SimpleNamespace(acompletion=_fake_acompletion), + ) + + llm = types.SimpleNamespace(model="openai/test", api_key="k", api_base="b") + await runtime.preflight_llm(llm, is_provider_rate_limited=lambda _: False) + + assert calls["model"] == "openai/test" + assert calls["max_tokens"] == 1 + assert calls["timeout"] == 2.5 + assert calls["stream"] is False + + +async def test_preflight_llm_rethrows_rate_limited(monkeypatch: pytest.MonkeyPatch) -> None: + class _RateLimitedError(Exception): + pass + + async def _fake_acompletion(**kwargs: Any): + del kwargs + raise _RateLimitedError("rl") + + monkeypatch.setitem( + sys.modules, + "litellm", + types.SimpleNamespace(acompletion=_fake_acompletion), + ) + + with pytest.raises(_RateLimitedError): + await runtime.preflight_llm( + types.SimpleNamespace(model="openai/test"), + is_provider_rate_limited=lambda exc: isinstance(exc, _RateLimitedError), + ) + + +async def test_preflight_llm_skips_probe_for_auto_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + called = {"count": 0} + + async def _fake_acompletion(**kwargs: Any): + del kwargs + called["count"] += 1 + return {"ok": True} + + monkeypatch.setitem( + sys.modules, + "litellm", + types.SimpleNamespace(acompletion=_fake_acompletion), + ) + + await runtime.preflight_llm( + types.SimpleNamespace(model="auto"), + is_provider_rate_limited=lambda _: False, + ) + assert called["count"] == 0 + + +async def test_build_main_agent_for_thread_forwards_arguments() -> None: + seen: dict[str, Any] = {} + + async def _factory(**kwargs: Any): + seen.update(kwargs) + return "agent" + + out = await runtime.build_main_agent_for_thread( + _factory, + llm="llm", + search_space_id=1, + db_session="db", + connector_service="connector", + checkpointer="cp", + user_id="u", + thread_id=10, + agent_config="cfg", + firecrawl_api_key="key", + thread_visibility="vis", + filesystem_selection="fs", + disabled_tools=["a"], + mentioned_document_ids=[5], + ) + assert out == "agent" + assert seen["thread_id"] == 10 + assert seen["mentioned_document_ids"] == [5] + + +async def test_settle_speculative_agent_build_swallows_exceptions() -> None: + async def _boom() -> None: + raise RuntimeError("ignore") + + import asyncio + + task = asyncio.create_task(_boom()) + await runtime.settle_speculative_agent_build(task) + assert task.done() From f8754a9dab480d2cb112f32c2cf1b4c67b4949ea Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 7 May 2026 15:41:33 +0200 Subject: [PATCH 54/96] Rename streaming runtime modules for clearer SRP boundaries. --- .../streaming/{runtime.py => agent_setup.py} | 2 +- .../chat/streaming/orchestration/__init__.py | 5 + .../streaming/orchestration/event_stream.py | 53 +++++++++ ...strator_runtime.py => test_agent_setup.py} | 14 +-- .../test_orchestration_event_stream.py | 107 ++++++++++++++++++ 5 files changed, 173 insertions(+), 8 deletions(-) rename surfsense_backend/app/tasks/chat/streaming/{runtime.py => agent_setup.py} (97%) create mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py rename surfsense_backend/tests/unit/tasks/chat/streaming/{test_orchestrator_runtime.py => test_agent_setup.py} (87%) create mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py diff --git a/surfsense_backend/app/tasks/chat/streaming/runtime.py b/surfsense_backend/app/tasks/chat/streaming/agent_setup.py similarity index 97% rename from surfsense_backend/app/tasks/chat/streaming/runtime.py rename to surfsense_backend/app/tasks/chat/streaming/agent_setup.py index b45da2789..f67c6ad65 100644 --- a/surfsense_backend/app/tasks/chat/streaming/runtime.py +++ b/surfsense_backend/app/tasks/chat/streaming/agent_setup.py @@ -1,4 +1,4 @@ -"""Runtime setup helpers for orchestrated chat streaming.""" +"""Agent setup helpers for orchestrated chat streaming.""" from __future__ import annotations diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py new file mode 100644 index 000000000..8b586f2be --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py @@ -0,0 +1,5 @@ +"""Composable orchestration pieces for chat streaming.""" + +from app.tasks.chat.streaming.orchestration.event_stream import stream_agent_events + +__all__ = ["stream_agent_events"] diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py new file mode 100644 index 000000000..1448cd86a --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py @@ -0,0 +1,53 @@ +"""Run LangGraph event streams through the EventRelay.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import Any + +from app.agents.new_chat.feature_flags import get_flags +from app.tasks.chat.streaming.event_relay import EventRelay +from app.tasks.chat.streaming.relay.state import AgentEventRelayState +from app.tasks.chat.streaming.stream_result import StreamResult + + +async def stream_agent_events( + *, + agent: Any, + config: dict[str, Any], + input_data: Any, + streaming_service: Any, + result: StreamResult, + step_prefix: str = "thinking", + initial_step_id: str | None = None, + initial_step_title: str = "", + initial_step_items: list[str] | None = None, + content_builder: Any | None = None, + runtime_context: Any = None, +) -> AsyncIterator[str]: + """Yield SSE frames from agent ``astream_events`` via ``EventRelay``.""" + state = AgentEventRelayState.for_invocation( + initial_step_id=initial_step_id, + initial_step_title=initial_step_title, + initial_step_items=initial_step_items, + parity_v2=bool(get_flags().enable_stream_parity_v2), + ) + + astream_kwargs: dict[str, Any] = {"config": config, "version": "v2"} + if runtime_context is not None: + astream_kwargs["context"] = runtime_context + + events = agent.astream_events(input_data, **astream_kwargs) + relay = EventRelay(streaming_service=streaming_service) + async for frame in relay.relay( + events, + state=state, + result=result, + step_prefix=step_prefix, + content_builder=content_builder, + config=config, + ): + yield frame + + result.accumulated_text = state.accumulated_text + result.agent_called_update_memory = state.called_update_memory diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_runtime.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_agent_setup.py similarity index 87% rename from surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_runtime.py rename to surfsense_backend/tests/unit/tasks/chat/streaming/test_agent_setup.py index edb05edfa..e1f7dd027 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_runtime.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_agent_setup.py @@ -1,4 +1,4 @@ -"""Behavior tests for streaming runtime helpers.""" +"""Behavior tests for streaming agent setup helpers.""" from __future__ import annotations @@ -8,7 +8,7 @@ from typing import Any import pytest -from app.tasks.chat.streaming import runtime +from app.tasks.chat.streaming import agent_setup pytestmark = pytest.mark.unit @@ -29,7 +29,7 @@ async def test_preflight_llm_calls_litellm_when_model_present( ) llm = types.SimpleNamespace(model="openai/test", api_key="k", api_base="b") - await runtime.preflight_llm(llm, is_provider_rate_limited=lambda _: False) + await agent_setup.preflight_llm(llm, is_provider_rate_limited=lambda _: False) assert calls["model"] == "openai/test" assert calls["max_tokens"] == 1 @@ -52,7 +52,7 @@ async def test_preflight_llm_rethrows_rate_limited(monkeypatch: pytest.MonkeyPat ) with pytest.raises(_RateLimitedError): - await runtime.preflight_llm( + await agent_setup.preflight_llm( types.SimpleNamespace(model="openai/test"), is_provider_rate_limited=lambda exc: isinstance(exc, _RateLimitedError), ) @@ -74,7 +74,7 @@ async def test_preflight_llm_skips_probe_for_auto_model( types.SimpleNamespace(acompletion=_fake_acompletion), ) - await runtime.preflight_llm( + await agent_setup.preflight_llm( types.SimpleNamespace(model="auto"), is_provider_rate_limited=lambda _: False, ) @@ -88,7 +88,7 @@ async def test_build_main_agent_for_thread_forwards_arguments() -> None: seen.update(kwargs) return "agent" - out = await runtime.build_main_agent_for_thread( + out = await agent_setup.build_main_agent_for_thread( _factory, llm="llm", search_space_id=1, @@ -116,5 +116,5 @@ async def test_settle_speculative_agent_build_swallows_exceptions() -> None: import asyncio task = asyncio.create_task(_boom()) - await runtime.settle_speculative_agent_build(task) + await agent_setup.settle_speculative_agent_build(task) assert task.done() diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py new file mode 100644 index 000000000..e12283a75 --- /dev/null +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py @@ -0,0 +1,107 @@ +"""Behavior tests for orchestration event-stream execution.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import pytest + +from app.tasks.chat.streaming.orchestration import stream_agent_events +from app.tasks.chat.streaming.stream_result import StreamResult + +pytestmark = pytest.mark.unit + + +@dataclass +class _Chunk: + content: Any = "" + additional_kwargs: dict[str, Any] = field(default_factory=dict) + tool_call_chunks: list[dict[str, Any]] = field(default_factory=list) + + +class _StreamingService: + def __init__(self) -> None: + self._text_idx = 0 + + def generate_text_id(self) -> str: + self._text_idx += 1 + return f"text-{self._text_idx}" + + def format_text_start(self, text_id: str) -> str: + return f"text_start:{text_id}" + + def format_text_delta(self, text_id: str, text: str) -> str: + return f"text_delta:{text_id}:{text}" + + def format_text_end(self, text_id: str) -> str: + return f"text_end:{text_id}" + + +class _Agent: + def __init__(self, events: list[dict[str, Any]]) -> None: + self.events = list(events) + self.calls: list[tuple[Any, dict[str, Any]]] = [] + + async def astream_events(self, input_data: Any, **kwargs: Any): + self.calls.append((input_data, kwargs)) + for event in self.events: + yield event + + +async def _collect(stream: Any) -> list[str]: + out: list[str] = [] + async for x in stream: + out.append(x) + return out + + +async def test_stream_agent_events_emits_text_lifecycle_and_updates_result() -> None: + service = _StreamingService() + agent = _Agent( + [ + {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(content="Hello")}}, + {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(content=" world")}}, + ] + ) + result = StreamResult() + + frames = await _collect( + stream_agent_events( + agent=agent, + config={"configurable": {"thread_id": "t-1"}}, + input_data={"messages": []}, + streaming_service=service, + result=result, + ) + ) + + assert frames == [ + "text_start:text-1", + "text_delta:text-1:Hello", + "text_delta:text-1: world", + "text_end:text-1", + ] + assert result.accumulated_text == "Hello world" + assert result.agent_called_update_memory is False + assert agent.calls[0][1]["version"] == "v2" + + +async def test_stream_agent_events_passes_runtime_context_to_agent() -> None: + service = _StreamingService() + agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("x")}}]) + result = StreamResult() + + _ = await _collect( + stream_agent_events( + agent=agent, + config={"configurable": {"thread_id": "t-2"}}, + input_data={"messages": []}, + streaming_service=service, + result=result, + runtime_context={"mentioned_document_ids": [1, 2]}, + ) + ) + + assert agent.calls + assert agent.calls[0][1]["context"] == {"mentioned_document_ids": [1, 2]} From 52593d88dbf2dfa5d81587048f4d2bb0ea0d5cad Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 7 May 2026 16:00:15 +0200 Subject: [PATCH 55/96] Reorganize streaming orchestration modules into relay and orchestration folders. --- .../app/routes/new_chat_routes.py | 2 +- .../chat/streaming/orchestration/__init__.py | 8 +- .../streaming/orchestration/event_stream.py | 6 +- .../chat/streaming/orchestration/input.py | 23 +++++ .../{ => orchestration}/orchestrator.py | 33 +++++++ .../output.py} | 8 +- .../tasks/chat/streaming/relay/__init__.py | 4 + .../chat/streaming/{ => relay}/event_relay.py | 4 +- .../test_orchestration_event_stream.py | 6 +- .../test_orchestrator_stream_chat.py | 88 +++++++++++++++++++ 10 files changed, 170 insertions(+), 12 deletions(-) create mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/input.py rename surfsense_backend/app/tasks/chat/streaming/{ => orchestration}/orchestrator.py (76%) rename surfsense_backend/app/tasks/chat/streaming/{stream_result.py => orchestration/output.py} (82%) rename surfsense_backend/app/tasks/chat/streaming/{ => relay}/event_relay.py (97%) create mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py diff --git a/surfsense_backend/app/routes/new_chat_routes.py b/surfsense_backend/app/routes/new_chat_routes.py index 7f035daef..e54497f93 100644 --- a/surfsense_backend/app/routes/new_chat_routes.py +++ b/surfsense_backend/app/routes/new_chat_routes.py @@ -75,7 +75,7 @@ from app.tasks.chat.stream_new_chat import ( stream_new_chat as legacy_stream_new_chat, stream_resume_chat as legacy_stream_resume_chat, ) -from app.tasks.chat.streaming.orchestrator import ( +from app.tasks.chat.streaming.orchestration.orchestrator import ( stream_chat, stream_regenerate, stream_resume, diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py index 8b586f2be..6f683a410 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py @@ -1,5 +1,11 @@ """Composable orchestration pieces for chat streaming.""" from app.tasks.chat.streaming.orchestration.event_stream import stream_agent_events +from app.tasks.chat.streaming.orchestration.input import StreamExecutionInput +from app.tasks.chat.streaming.orchestration.output import StreamOutput -__all__ = ["stream_agent_events"] +__all__ = [ + "StreamExecutionInput", + "StreamOutput", + "stream_agent_events", +] diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py index 1448cd86a..369883c3a 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py @@ -6,9 +6,9 @@ from collections.abc import AsyncIterator from typing import Any from app.agents.new_chat.feature_flags import get_flags -from app.tasks.chat.streaming.event_relay import EventRelay +from app.tasks.chat.streaming.orchestration.output import StreamOutput +from app.tasks.chat.streaming.relay.event_relay import EventRelay from app.tasks.chat.streaming.relay.state import AgentEventRelayState -from app.tasks.chat.streaming.stream_result import StreamResult async def stream_agent_events( @@ -17,7 +17,7 @@ async def stream_agent_events( config: dict[str, Any], input_data: Any, streaming_service: Any, - result: StreamResult, + result: StreamOutput, step_prefix: str = "thinking", initial_step_id: str | None = None, initial_step_title: str = "", diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/input.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/input.py new file mode 100644 index 000000000..13d43b612 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/input.py @@ -0,0 +1,23 @@ +"""Inputs for orchestrator-owned streaming execution.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class StreamExecutionInput: + """Container for dependencies required by ``stream_agent_events``.""" + + agent: Any + config: dict[str, Any] + input_data: Any + streaming_service: Any + step_prefix: str = "thinking" + initial_step_id: str | None = None + initial_step_title: str = "" + initial_step_items: list[str] | None = None + content_builder: Any | None = None + runtime_context: Any = None + diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py similarity index 76% rename from surfsense_backend/app/tasks/chat/streaming/orchestrator.py rename to surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py index e912dd632..ac7abc6f4 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestrator.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py @@ -14,6 +14,9 @@ from typing import Any, Literal from app.agents.new_chat.filesystem_selection import FilesystemSelection from app.db import ChatVisibility from app.tasks.chat.stream_new_chat import stream_new_chat, stream_resume_chat +from app.tasks.chat.streaming.orchestration.event_stream import stream_agent_events +from app.tasks.chat.streaming.orchestration.input import StreamExecutionInput +from app.tasks.chat.streaming.orchestration.output import StreamOutput async def stream_chat( @@ -34,8 +37,38 @@ async def stream_chat( filesystem_selection: FilesystemSelection | None = None, request_id: str | None = None, user_image_data_urls: list[str] | None = None, + orchestration_input: StreamExecutionInput | None = None, ) -> AsyncGenerator[str, None]: """Stream a new chat turn through the current production pipeline.""" + if orchestration_input is not None: + result = StreamOutput( + request_id=request_id, + turn_id=f"{chat_id}:orchestrator", + filesystem_mode=( + filesystem_selection.mode.value if filesystem_selection else "cloud" + ), + client_platform=( + filesystem_selection.client_platform.value + if filesystem_selection + else "web" + ), + ) + async for frame in stream_agent_events( + agent=orchestration_input.agent, + config=orchestration_input.config, + input_data=orchestration_input.input_data, + streaming_service=orchestration_input.streaming_service, + result=result, + step_prefix=orchestration_input.step_prefix, + initial_step_id=orchestration_input.initial_step_id, + initial_step_title=orchestration_input.initial_step_title, + initial_step_items=orchestration_input.initial_step_items, + content_builder=orchestration_input.content_builder, + runtime_context=orchestration_input.runtime_context, + ): + yield frame + return + async for chunk in stream_new_chat( user_query=user_query, search_space_id=search_space_id, diff --git a/surfsense_backend/app/tasks/chat/streaming/stream_result.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/output.py similarity index 82% rename from surfsense_backend/app/tasks/chat/streaming/stream_result.py rename to surfsense_backend/app/tasks/chat/streaming/orchestration/output.py index 8ea3bd295..0c4870ec4 100644 --- a/surfsense_backend/app/tasks/chat/streaming/stream_result.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/output.py @@ -1,4 +1,4 @@ -"""Mutable facts collected while streaming one agent turn.""" +"""Output facts collected while streaming one orchestrated agent turn.""" from __future__ import annotations @@ -7,7 +7,7 @@ from typing import Any @dataclass -class StreamResult: +class StreamOutput: accumulated_text: str = "" is_interrupted: bool = False interrupt_value: dict[str, Any] | None = None @@ -26,3 +26,7 @@ class StreamResult: commit_gate_reason: str = "" assistant_message_id: int | None = None content_builder: Any | None = field(default=None, repr=False) + + +# Backwards-compatible alias while imports migrate. +StreamResult = StreamOutput diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py b/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py index c1a5e7175..351e878a8 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py @@ -1,3 +1,7 @@ """Relay state: thinking steps, tool bookkeeping, and stream helpers.""" from __future__ import annotations + +from app.tasks.chat.streaming.relay.event_relay import EventRelay, EventRelayConfig + +__all__ = ["EventRelay", "EventRelayConfig"] diff --git a/surfsense_backend/app/tasks/chat/streaming/event_relay.py b/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py similarity index 97% rename from surfsense_backend/app/tasks/chat/streaming/event_relay.py rename to surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py index f86337ad7..072baac72 100644 --- a/surfsense_backend/app/tasks/chat/streaming/event_relay.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py @@ -16,11 +16,11 @@ from app.tasks.chat.streaming.handlers.custom_event_dispatch import ( ) from app.tasks.chat.streaming.handlers.tool_end import iter_tool_end_frames from app.tasks.chat.streaming.handlers.tool_start import iter_tool_start_frames +from app.tasks.chat.streaming.orchestration.output import StreamOutput from app.tasks.chat.streaming.relay.state import AgentEventRelayState from app.tasks.chat.streaming.relay.thinking_step_completion import ( complete_active_thinking_step, ) -from app.tasks.chat.streaming.stream_result import StreamResult @dataclass @@ -52,7 +52,7 @@ class EventRelay: events: AsyncIterator[dict[str, Any]], *, state: AgentEventRelayState, - result: StreamResult, + result: StreamOutput, step_prefix: str = "thinking", content_builder: Any | None = None, config: dict[str, Any] | None = None, diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py index e12283a75..e0a1877a8 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py @@ -8,7 +8,7 @@ from typing import Any import pytest from app.tasks.chat.streaming.orchestration import stream_agent_events -from app.tasks.chat.streaming.stream_result import StreamResult +from app.tasks.chat.streaming.orchestration.output import StreamOutput pytestmark = pytest.mark.unit @@ -64,7 +64,7 @@ async def test_stream_agent_events_emits_text_lifecycle_and_updates_result() -> {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(content=" world")}}, ] ) - result = StreamResult() + result = StreamOutput() frames = await _collect( stream_agent_events( @@ -90,7 +90,7 @@ async def test_stream_agent_events_emits_text_lifecycle_and_updates_result() -> async def test_stream_agent_events_passes_runtime_context_to_agent() -> None: service = _StreamingService() agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("x")}}]) - result = StreamResult() + result = StreamOutput() _ = await _collect( stream_agent_events( diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py new file mode 100644 index 000000000..cf54fdab0 --- /dev/null +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py @@ -0,0 +1,88 @@ +"""Behavior tests for orchestrator ``stream_chat`` public API.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import pytest + +from app.tasks.chat.streaming.orchestration import StreamExecutionInput +from app.tasks.chat.streaming.orchestration.orchestrator import stream_chat + +pytestmark = pytest.mark.unit + + +@dataclass +class _Chunk: + content: Any = "" + additional_kwargs: dict[str, Any] = field(default_factory=dict) + tool_call_chunks: list[dict[str, Any]] = field(default_factory=list) + + +class _StreamingService: + def __init__(self) -> None: + self._text_idx = 0 + + def generate_text_id(self) -> str: + self._text_idx += 1 + return f"text-{self._text_idx}" + + def format_text_start(self, text_id: str) -> str: + return f"text_start:{text_id}" + + def format_text_delta(self, text_id: str, text: str) -> str: + return f"text_delta:{text_id}:{text}" + + def format_text_end(self, text_id: str) -> str: + return f"text_end:{text_id}" + + +class _Agent: + def __init__(self, events: list[dict[str, Any]]) -> None: + self.events = list(events) + self.calls: list[tuple[Any, dict[str, Any]]] = [] + + async def astream_events(self, input_data: Any, **kwargs: Any): + self.calls.append((input_data, kwargs)) + for event in self.events: + yield event + + +async def _collect(stream: Any) -> list[str]: + out: list[str] = [] + async for x in stream: + out.append(x) + return out + + +async def test_stream_chat_uses_orchestration_input_path() -> None: + service = _StreamingService() + agent = _Agent( + [ + {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(content="hello")}}, + {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(content="!")}}, + ] + ) + frames = await _collect( + stream_chat( + user_query="ignored-here", + search_space_id=1, + chat_id=77, + orchestration_input=StreamExecutionInput( + agent=agent, + config={"configurable": {"thread_id": "thread-1"}}, + input_data={"messages": []}, + streaming_service=service, + ), + ) + ) + + assert frames == [ + "text_start:text-1", + "text_delta:text-1:hello", + "text_delta:text-1:!", + "text_end:text-1", + ] + assert agent.calls + assert agent.calls[0][1]["version"] == "v2" From 0f40279d9581515dddbd847318ef37589972630c Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 7 May 2026 16:18:29 +0200 Subject: [PATCH 56/96] Expand orchestration gate coverage to resume and regenerate flows. --- .../streaming/orchestration/orchestrator.py | 60 +++++++++++++++++++ .../test_orchestration_event_stream.py | 18 ++++-- .../test_orchestrator_stream_chat.py | 58 +++++++++++++++++- 3 files changed, 128 insertions(+), 8 deletions(-) diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py index ac7abc6f4..1e32e7f5a 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py @@ -101,8 +101,38 @@ async def stream_resume( filesystem_selection: FilesystemSelection | None = None, request_id: str | None = None, disabled_tools: list[str] | None = None, + orchestration_input: StreamExecutionInput | None = None, ) -> AsyncGenerator[str, None]: """Resume an interrupted chat turn through the current production pipeline.""" + if orchestration_input is not None: + result = StreamOutput( + request_id=request_id, + turn_id=f"{chat_id}:orchestrator-resume", + filesystem_mode=( + filesystem_selection.mode.value if filesystem_selection else "cloud" + ), + client_platform=( + filesystem_selection.client_platform.value + if filesystem_selection + else "web" + ), + ) + async for frame in stream_agent_events( + agent=orchestration_input.agent, + config=orchestration_input.config, + input_data=orchestration_input.input_data, + streaming_service=orchestration_input.streaming_service, + result=result, + step_prefix=orchestration_input.step_prefix, + initial_step_id=orchestration_input.initial_step_id, + initial_step_title=orchestration_input.initial_step_title, + initial_step_items=orchestration_input.initial_step_items, + content_builder=orchestration_input.content_builder, + runtime_context=orchestration_input.runtime_context, + ): + yield frame + return + async for chunk in stream_resume_chat( chat_id=chat_id, search_space_id=search_space_id, @@ -136,8 +166,38 @@ async def stream_regenerate( request_id: str | None = None, user_image_data_urls: list[str] | None = None, flow: Literal["new", "regenerate"] = "regenerate", + orchestration_input: StreamExecutionInput | None = None, ) -> AsyncGenerator[str, None]: """Regenerate an assistant turn through the current production pipeline.""" + if orchestration_input is not None: + result = StreamOutput( + request_id=request_id, + turn_id=f"{chat_id}:orchestrator-regenerate", + filesystem_mode=( + filesystem_selection.mode.value if filesystem_selection else "cloud" + ), + client_platform=( + filesystem_selection.client_platform.value + if filesystem_selection + else "web" + ), + ) + async for frame in stream_agent_events( + agent=orchestration_input.agent, + config=orchestration_input.config, + input_data=orchestration_input.input_data, + streaming_service=orchestration_input.streaming_service, + result=result, + step_prefix=orchestration_input.step_prefix, + initial_step_id=orchestration_input.initial_step_id, + initial_step_title=orchestration_input.initial_step_title, + initial_step_items=orchestration_input.initial_step_items, + content_builder=orchestration_input.content_builder, + runtime_context=orchestration_input.runtime_context, + ): + yield frame + return + async for chunk in stream_new_chat( user_query=user_query, search_space_id=search_space_id, diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py index e0a1877a8..bd154e6a0 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py @@ -84,15 +84,20 @@ async def test_stream_agent_events_emits_text_lifecycle_and_updates_result() -> ] assert result.accumulated_text == "Hello world" assert result.agent_called_update_memory is False - assert agent.calls[0][1]["version"] == "v2" async def test_stream_agent_events_passes_runtime_context_to_agent() -> None: service = _StreamingService() - agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("x")}}]) + class _ContextAwareAgent: + async def astream_events(self, input_data: Any, **kwargs: Any): + del input_data + text = "ctx-ok" if kwargs.get("context") else "ctx-missing" + yield {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(text)}} + + agent = _ContextAwareAgent() result = StreamOutput() - _ = await _collect( + frames = await _collect( stream_agent_events( agent=agent, config={"configurable": {"thread_id": "t-2"}}, @@ -103,5 +108,8 @@ async def test_stream_agent_events_passes_runtime_context_to_agent() -> None: ) ) - assert agent.calls - assert agent.calls[0][1]["context"] == {"mentioned_document_ids": [1, 2]} + assert frames == [ + "text_start:text-1", + "text_delta:text-1:ctx-ok", + "text_end:text-1", + ] diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py index cf54fdab0..d9cd7951f 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py @@ -8,7 +8,11 @@ from typing import Any import pytest from app.tasks.chat.streaming.orchestration import StreamExecutionInput -from app.tasks.chat.streaming.orchestration.orchestrator import stream_chat +from app.tasks.chat.streaming.orchestration.orchestrator import ( + stream_chat, + stream_regenerate, + stream_resume, +) pytestmark = pytest.mark.unit @@ -84,5 +88,53 @@ async def test_stream_chat_uses_orchestration_input_path() -> None: "text_delta:text-1:!", "text_end:text-1", ] - assert agent.calls - assert agent.calls[0][1]["version"] == "v2" + + +async def test_stream_resume_uses_orchestration_input_path() -> None: + service = _StreamingService() + agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("r")}}]) + + frames = await _collect( + stream_resume( + chat_id=9, + search_space_id=1, + decisions=[], + orchestration_input=StreamExecutionInput( + agent=agent, + config={"configurable": {"thread_id": "thread-r"}}, + input_data={"messages": []}, + streaming_service=service, + ), + ) + ) + + assert frames == [ + "text_start:text-1", + "text_delta:text-1:r", + "text_end:text-1", + ] + + +async def test_stream_regenerate_uses_orchestration_input_path() -> None: + service = _StreamingService() + agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("g")}}]) + + frames = await _collect( + stream_regenerate( + user_query="q", + search_space_id=1, + chat_id=2, + orchestration_input=StreamExecutionInput( + agent=agent, + config={"configurable": {"thread_id": "thread-g"}}, + input_data={"messages": []}, + streaming_service=service, + ), + ) + ) + + assert frames == [ + "text_start:text-1", + "text_delta:text-1:g", + "text_end:text-1", + ] From a04b2e88bdaeea624c0d192b3259bbd3482bc717 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 7 May 2026 17:06:17 +0200 Subject: [PATCH 57/96] wire orchestrator streaming context path and align event relay outputs --- .../chat/streaming/orchestration/__init__.py | 12 +- .../streaming/orchestration/event_stream.py | 6 +- .../chat/streaming/orchestration/input.py | 4 +- .../streaming/orchestration/orchestrator.py | 142 ++++++++---------- .../chat/streaming/orchestration/output.py | 5 +- .../tasks/chat/streaming/relay/event_relay.py | 4 +- .../test_orchestration_event_stream.py | 16 +- .../test_orchestrator_stream_chat.py | 14 +- 8 files changed, 94 insertions(+), 109 deletions(-) diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py index 6f683a410..b1a201fd3 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py @@ -1,11 +1,11 @@ """Composable orchestration pieces for chat streaming.""" -from app.tasks.chat.streaming.orchestration.event_stream import stream_agent_events -from app.tasks.chat.streaming.orchestration.input import StreamExecutionInput -from app.tasks.chat.streaming.orchestration.output import StreamOutput +from app.tasks.chat.streaming.orchestration.event_stream import stream_output +from app.tasks.chat.streaming.orchestration.input import StreamingContext +from app.tasks.chat.streaming.orchestration.output import StreamingResult __all__ = [ - "StreamExecutionInput", - "StreamOutput", - "stream_agent_events", + "StreamingContext", + "StreamingResult", + "stream_output", ] diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py index 369883c3a..fc8c13027 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py @@ -6,18 +6,18 @@ from collections.abc import AsyncIterator from typing import Any from app.agents.new_chat.feature_flags import get_flags -from app.tasks.chat.streaming.orchestration.output import StreamOutput +from app.tasks.chat.streaming.orchestration.output import StreamingResult from app.tasks.chat.streaming.relay.event_relay import EventRelay from app.tasks.chat.streaming.relay.state import AgentEventRelayState -async def stream_agent_events( +async def stream_output( *, agent: Any, config: dict[str, Any], input_data: Any, streaming_service: Any, - result: StreamOutput, + result: StreamingResult, step_prefix: str = "thinking", initial_step_id: str | None = None, initial_step_title: str = "", diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/input.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/input.py index 13d43b612..45a33d435 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/input.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/input.py @@ -7,8 +7,8 @@ from typing import Any @dataclass(frozen=True) -class StreamExecutionInput: - """Container for dependencies required by ``stream_agent_events``.""" +class StreamingContext: + """Container for dependencies required by ``stream_output``.""" agent: Any config: dict[str, Any] diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py index 1e32e7f5a..b40083f42 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py @@ -1,9 +1,4 @@ """Top-level chat streaming entrypoints. - -For now these orchestrator functions are thin compatibility wrappers around the -current ``stream_new_chat`` / ``stream_resume_chat`` implementations. Routing -calls through this module lets us cut over to the fully modular event relay in -one place later without touching API routes again. """ from __future__ import annotations @@ -14,9 +9,47 @@ from typing import Any, Literal from app.agents.new_chat.filesystem_selection import FilesystemSelection from app.db import ChatVisibility from app.tasks.chat.stream_new_chat import stream_new_chat, stream_resume_chat -from app.tasks.chat.streaming.orchestration.event_stream import stream_agent_events -from app.tasks.chat.streaming.orchestration.input import StreamExecutionInput -from app.tasks.chat.streaming.orchestration.output import StreamOutput +from app.tasks.chat.streaming.orchestration.event_stream import stream_output +from app.tasks.chat.streaming.orchestration.input import StreamingContext +from app.tasks.chat.streaming.orchestration.output import StreamingResult + + +def _build_streaming_result( + *, + chat_id: int, + request_id: str | None, + filesystem_selection: FilesystemSelection | None, + suffix: str, +) -> StreamingResult: + return StreamingResult( + request_id=request_id, + turn_id=f"{chat_id}:{suffix}", + filesystem_mode=(filesystem_selection.mode.value if filesystem_selection else "cloud"), + client_platform=( + filesystem_selection.client_platform.value if filesystem_selection else "web" + ), + ) + + +async def _stream_output_with_streaming_context( + *, + streaming_context: StreamingContext, + result: StreamingResult, +) -> AsyncGenerator[str, None]: + async for frame in stream_output( + agent=streaming_context.agent, + config=streaming_context.config, + input_data=streaming_context.input_data, + streaming_service=streaming_context.streaming_service, + result=result, + step_prefix=streaming_context.step_prefix, + initial_step_id=streaming_context.initial_step_id, + initial_step_title=streaming_context.initial_step_title, + initial_step_items=streaming_context.initial_step_items, + content_builder=streaming_context.content_builder, + runtime_context=streaming_context.runtime_context, + ): + yield frame async def stream_chat( @@ -37,34 +70,19 @@ async def stream_chat( filesystem_selection: FilesystemSelection | None = None, request_id: str | None = None, user_image_data_urls: list[str] | None = None, - orchestration_input: StreamExecutionInput | None = None, + streaming_context: StreamingContext | None = None, ) -> AsyncGenerator[str, None]: """Stream a new chat turn through the current production pipeline.""" - if orchestration_input is not None: - result = StreamOutput( + if streaming_context is not None: + result = _build_streaming_result( + chat_id=chat_id, request_id=request_id, - turn_id=f"{chat_id}:orchestrator", - filesystem_mode=( - filesystem_selection.mode.value if filesystem_selection else "cloud" - ), - client_platform=( - filesystem_selection.client_platform.value - if filesystem_selection - else "web" - ), + filesystem_selection=filesystem_selection, + suffix="orchestrator", ) - async for frame in stream_agent_events( - agent=orchestration_input.agent, - config=orchestration_input.config, - input_data=orchestration_input.input_data, - streaming_service=orchestration_input.streaming_service, + async for frame in _stream_output_with_streaming_context( + streaming_context=streaming_context, result=result, - step_prefix=orchestration_input.step_prefix, - initial_step_id=orchestration_input.initial_step_id, - initial_step_title=orchestration_input.initial_step_title, - initial_step_items=orchestration_input.initial_step_items, - content_builder=orchestration_input.content_builder, - runtime_context=orchestration_input.runtime_context, ): yield frame return @@ -101,34 +119,19 @@ async def stream_resume( filesystem_selection: FilesystemSelection | None = None, request_id: str | None = None, disabled_tools: list[str] | None = None, - orchestration_input: StreamExecutionInput | None = None, + streaming_context: StreamingContext | None = None, ) -> AsyncGenerator[str, None]: """Resume an interrupted chat turn through the current production pipeline.""" - if orchestration_input is not None: - result = StreamOutput( + if streaming_context is not None: + result = _build_streaming_result( + chat_id=chat_id, request_id=request_id, - turn_id=f"{chat_id}:orchestrator-resume", - filesystem_mode=( - filesystem_selection.mode.value if filesystem_selection else "cloud" - ), - client_platform=( - filesystem_selection.client_platform.value - if filesystem_selection - else "web" - ), + filesystem_selection=filesystem_selection, + suffix="orchestrator-resume", ) - async for frame in stream_agent_events( - agent=orchestration_input.agent, - config=orchestration_input.config, - input_data=orchestration_input.input_data, - streaming_service=orchestration_input.streaming_service, + async for frame in _stream_output_with_streaming_context( + streaming_context=streaming_context, result=result, - step_prefix=orchestration_input.step_prefix, - initial_step_id=orchestration_input.initial_step_id, - initial_step_title=orchestration_input.initial_step_title, - initial_step_items=orchestration_input.initial_step_items, - content_builder=orchestration_input.content_builder, - runtime_context=orchestration_input.runtime_context, ): yield frame return @@ -166,34 +169,19 @@ async def stream_regenerate( request_id: str | None = None, user_image_data_urls: list[str] | None = None, flow: Literal["new", "regenerate"] = "regenerate", - orchestration_input: StreamExecutionInput | None = None, + streaming_context: StreamingContext | None = None, ) -> AsyncGenerator[str, None]: """Regenerate an assistant turn through the current production pipeline.""" - if orchestration_input is not None: - result = StreamOutput( + if streaming_context is not None: + result = _build_streaming_result( + chat_id=chat_id, request_id=request_id, - turn_id=f"{chat_id}:orchestrator-regenerate", - filesystem_mode=( - filesystem_selection.mode.value if filesystem_selection else "cloud" - ), - client_platform=( - filesystem_selection.client_platform.value - if filesystem_selection - else "web" - ), + filesystem_selection=filesystem_selection, + suffix="orchestrator-regenerate", ) - async for frame in stream_agent_events( - agent=orchestration_input.agent, - config=orchestration_input.config, - input_data=orchestration_input.input_data, - streaming_service=orchestration_input.streaming_service, + async for frame in _stream_output_with_streaming_context( + streaming_context=streaming_context, result=result, - step_prefix=orchestration_input.step_prefix, - initial_step_id=orchestration_input.initial_step_id, - initial_step_title=orchestration_input.initial_step_title, - initial_step_items=orchestration_input.initial_step_items, - content_builder=orchestration_input.content_builder, - runtime_context=orchestration_input.runtime_context, ): yield frame return diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/output.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/output.py index 0c4870ec4..60f8ee6ee 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/output.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/output.py @@ -7,7 +7,7 @@ from typing import Any @dataclass -class StreamOutput: +class StreamingResult: accumulated_text: str = "" is_interrupted: bool = False interrupt_value: dict[str, Any] | None = None @@ -27,6 +27,3 @@ class StreamOutput: assistant_message_id: int | None = None content_builder: Any | None = field(default=None, repr=False) - -# Backwards-compatible alias while imports migrate. -StreamResult = StreamOutput diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py b/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py index 072baac72..c8aebd99c 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py @@ -16,7 +16,7 @@ from app.tasks.chat.streaming.handlers.custom_event_dispatch import ( ) from app.tasks.chat.streaming.handlers.tool_end import iter_tool_end_frames from app.tasks.chat.streaming.handlers.tool_start import iter_tool_start_frames -from app.tasks.chat.streaming.orchestration.output import StreamOutput +from app.tasks.chat.streaming.orchestration.output import StreamingResult from app.tasks.chat.streaming.relay.state import AgentEventRelayState from app.tasks.chat.streaming.relay.thinking_step_completion import ( complete_active_thinking_step, @@ -52,7 +52,7 @@ class EventRelay: events: AsyncIterator[dict[str, Any]], *, state: AgentEventRelayState, - result: StreamOutput, + result: StreamingResult, step_prefix: str = "thinking", content_builder: Any | None = None, config: dict[str, Any] | None = None, diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py index bd154e6a0..b17d82293 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py @@ -7,8 +7,8 @@ from typing import Any import pytest -from app.tasks.chat.streaming.orchestration import stream_agent_events -from app.tasks.chat.streaming.orchestration.output import StreamOutput +from app.tasks.chat.streaming.orchestration import stream_output +from app.tasks.chat.streaming.orchestration.output import StreamingResult pytestmark = pytest.mark.unit @@ -56,7 +56,7 @@ async def _collect(stream: Any) -> list[str]: return out -async def test_stream_agent_events_emits_text_lifecycle_and_updates_result() -> None: +async def test_stream_output_emits_text_lifecycle_and_updates_result() -> None: service = _StreamingService() agent = _Agent( [ @@ -64,10 +64,10 @@ async def test_stream_agent_events_emits_text_lifecycle_and_updates_result() -> {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(content=" world")}}, ] ) - result = StreamOutput() + result = StreamingResult() frames = await _collect( - stream_agent_events( + stream_output( agent=agent, config={"configurable": {"thread_id": "t-1"}}, input_data={"messages": []}, @@ -86,7 +86,7 @@ async def test_stream_agent_events_emits_text_lifecycle_and_updates_result() -> assert result.agent_called_update_memory is False -async def test_stream_agent_events_passes_runtime_context_to_agent() -> None: +async def test_stream_output_passes_runtime_context_to_agent() -> None: service = _StreamingService() class _ContextAwareAgent: async def astream_events(self, input_data: Any, **kwargs: Any): @@ -95,10 +95,10 @@ async def test_stream_agent_events_passes_runtime_context_to_agent() -> None: yield {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(text)}} agent = _ContextAwareAgent() - result = StreamOutput() + result = StreamingResult() frames = await _collect( - stream_agent_events( + stream_output( agent=agent, config={"configurable": {"thread_id": "t-2"}}, input_data={"messages": []}, diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py index d9cd7951f..b84193cb7 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py @@ -7,7 +7,7 @@ from typing import Any import pytest -from app.tasks.chat.streaming.orchestration import StreamExecutionInput +from app.tasks.chat.streaming.orchestration import StreamingContext from app.tasks.chat.streaming.orchestration.orchestrator import ( stream_chat, stream_regenerate, @@ -60,7 +60,7 @@ async def _collect(stream: Any) -> list[str]: return out -async def test_stream_chat_uses_orchestration_input_path() -> None: +async def test_stream_chat_uses_streaming_context_path() -> None: service = _StreamingService() agent = _Agent( [ @@ -73,7 +73,7 @@ async def test_stream_chat_uses_orchestration_input_path() -> None: user_query="ignored-here", search_space_id=1, chat_id=77, - orchestration_input=StreamExecutionInput( + streaming_context=StreamingContext( agent=agent, config={"configurable": {"thread_id": "thread-1"}}, input_data={"messages": []}, @@ -90,7 +90,7 @@ async def test_stream_chat_uses_orchestration_input_path() -> None: ] -async def test_stream_resume_uses_orchestration_input_path() -> None: +async def test_stream_resume_uses_streaming_context_path() -> None: service = _StreamingService() agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("r")}}]) @@ -99,7 +99,7 @@ async def test_stream_resume_uses_orchestration_input_path() -> None: chat_id=9, search_space_id=1, decisions=[], - orchestration_input=StreamExecutionInput( + streaming_context=StreamingContext( agent=agent, config={"configurable": {"thread_id": "thread-r"}}, input_data={"messages": []}, @@ -115,7 +115,7 @@ async def test_stream_resume_uses_orchestration_input_path() -> None: ] -async def test_stream_regenerate_uses_orchestration_input_path() -> None: +async def test_stream_regenerate_uses_streaming_context_path() -> None: service = _StreamingService() agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("g")}}]) @@ -124,7 +124,7 @@ async def test_stream_regenerate_uses_orchestration_input_path() -> None: user_query="q", search_space_id=1, chat_id=2, - orchestration_input=StreamExecutionInput( + streaming_context=StreamingContext( agent=agent, config={"configurable": {"thread_id": "thread-g"}}, input_data={"messages": []}, From 52895e37e9ec86aba93cf4dd3fcf734da2b2e162 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 7 May 2026 17:57:27 +0200 Subject: [PATCH 58/96] build streaming contexts for chat resume and regenerate paths --- .../streaming/orchestration/orchestrator.py | 55 +++- .../streaming_context/__init__.py | 18 ++ .../orchestration/streaming_context/chat.py | 258 ++++++++++++++++++ .../streaming_context/regenerate.py | 49 ++++ .../orchestration/streaming_context/resume.py | 154 +++++++++++ .../test_orchestrator_stream_chat.py | 100 +++++++ 6 files changed, 633 insertions(+), 1 deletion(-) create mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/__init__.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/chat.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/regenerate.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/resume.py diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py index b40083f42..80cae77a2 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py @@ -9,6 +9,11 @@ from typing import Any, Literal from app.agents.new_chat.filesystem_selection import FilesystemSelection from app.db import ChatVisibility from app.tasks.chat.stream_new_chat import stream_new_chat, stream_resume_chat +from app.tasks.chat.streaming.orchestration.streaming_context import ( + build_chat_streaming_context, + build_regenerate_streaming_context, + build_resume_streaming_context, +) from app.tasks.chat.streaming.orchestration.event_stream import stream_output from app.tasks.chat.streaming.orchestration.input import StreamingContext from app.tasks.chat.streaming.orchestration.output import StreamingResult @@ -38,7 +43,7 @@ async def _stream_output_with_streaming_context( ) -> AsyncGenerator[str, None]: async for frame in stream_output( agent=streaming_context.agent, - config=streaming_context.config, + config=streaming_context.config, input_data=streaming_context.input_data, streaming_service=streaming_context.streaming_service, result=result, @@ -73,6 +78,24 @@ async def stream_chat( streaming_context: StreamingContext | None = None, ) -> AsyncGenerator[str, None]: """Stream a new chat turn through the current production pipeline.""" + if streaming_context is None: + streaming_context = await build_chat_streaming_context( + user_query=user_query, + search_space_id=search_space_id, + chat_id=chat_id, + user_id=user_id, + llm_config_id=llm_config_id, + mentioned_document_ids=mentioned_document_ids, + mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids, + checkpoint_id=checkpoint_id, + needs_history_bootstrap=needs_history_bootstrap, + thread_visibility=thread_visibility, + current_user_display_name=current_user_display_name, + disabled_tools=disabled_tools, + filesystem_selection=filesystem_selection, + request_id=request_id, + user_image_data_urls=user_image_data_urls, + ) if streaming_context is not None: result = _build_streaming_result( chat_id=chat_id, @@ -122,6 +145,18 @@ async def stream_resume( streaming_context: StreamingContext | None = None, ) -> AsyncGenerator[str, None]: """Resume an interrupted chat turn through the current production pipeline.""" + if streaming_context is None: + streaming_context = await build_resume_streaming_context( + chat_id=chat_id, + search_space_id=search_space_id, + decisions=decisions, + user_id=user_id, + llm_config_id=llm_config_id, + thread_visibility=thread_visibility, + filesystem_selection=filesystem_selection, + request_id=request_id, + disabled_tools=disabled_tools, + ) if streaming_context is not None: result = _build_streaming_result( chat_id=chat_id, @@ -172,6 +207,24 @@ async def stream_regenerate( streaming_context: StreamingContext | None = None, ) -> AsyncGenerator[str, None]: """Regenerate an assistant turn through the current production pipeline.""" + if streaming_context is None: + streaming_context = await build_regenerate_streaming_context( + user_query=user_query, + search_space_id=search_space_id, + chat_id=chat_id, + user_id=user_id, + llm_config_id=llm_config_id, + mentioned_document_ids=mentioned_document_ids, + mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids, + checkpoint_id=checkpoint_id, + needs_history_bootstrap=needs_history_bootstrap, + thread_visibility=thread_visibility, + current_user_display_name=current_user_display_name, + disabled_tools=disabled_tools, + filesystem_selection=filesystem_selection, + request_id=request_id, + user_image_data_urls=user_image_data_urls, + ) if streaming_context is not None: result = _build_streaming_result( chat_id=chat_id, diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/__init__.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/__init__.py new file mode 100644 index 000000000..1bd3e103d --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/__init__.py @@ -0,0 +1,18 @@ +"""Streaming context builders per orchestrator entrypoint.""" + +from app.tasks.chat.streaming.orchestration.streaming_context.chat import ( + build_chat_streaming_context, +) +from app.tasks.chat.streaming.orchestration.streaming_context.regenerate import ( + build_regenerate_streaming_context, +) +from app.tasks.chat.streaming.orchestration.streaming_context.resume import ( + build_resume_streaming_context, +) + +__all__ = [ + "build_chat_streaming_context", + "build_regenerate_streaming_context", + "build_resume_streaming_context", +] + diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/chat.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/chat.py new file mode 100644 index 000000000..eb459ae5c --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/chat.py @@ -0,0 +1,258 @@ +"""Build ``StreamingContext`` for chat streaming.""" + +from __future__ import annotations + +import logging +import time +from typing import Any + +from langchain_core.messages import HumanMessage +from sqlalchemy.future import select +from sqlalchemy.orm import selectinload + +from app.agents.multi_agent_chat import create_multi_agent_chat_deep_agent +from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent +from app.agents.new_chat.checkpointer import get_checkpointer +from app.agents.new_chat.context import SurfSenseContextSchema +from app.agents.new_chat.filesystem_selection import FilesystemSelection +from app.agents.new_chat.llm_config import ( + AgentConfig, + create_chat_litellm_from_agent_config, + create_chat_litellm_from_config, + load_agent_config, + load_global_llm_config_by_id, +) +from app.db import ( + ChatVisibility, + NewChatThread, + Report, + SearchSourceConnectorType, + SurfsenseDocsDocument, + async_session_maker, +) +from app.services.auto_model_pin_service import resolve_or_get_pinned_llm_config_id +from app.services.connector_service import ConnectorService +from app.services.new_streaming_service import VercelStreamingService +from app.tasks.chat.stream_new_chat import format_mentioned_surfsense_docs_as_context +from app.tasks.chat.streaming.agent_setup import build_main_agent_for_thread +from app.tasks.chat.streaming.orchestration.input import StreamingContext +from app.utils.content_utils import bootstrap_history_from_db +from app.utils.user_message_multimodal import build_human_message_content + +logger = logging.getLogger(__name__) + + +async def build_chat_streaming_context( + *, + user_query: str, + search_space_id: int, + chat_id: int, + user_id: str | None = None, + llm_config_id: int = -1, + mentioned_document_ids: list[int] | None = None, + mentioned_surfsense_doc_ids: list[int] | None = None, + checkpoint_id: str | None = None, + needs_history_bootstrap: bool = False, + thread_visibility: ChatVisibility | None = None, + current_user_display_name: str | None = None, + disabled_tools: list[str] | None = None, + filesystem_selection: FilesystemSelection | None = None, + request_id: str | None = None, + user_image_data_urls: list[str] | None = None, +) -> StreamingContext | None: + """Build context for ``stream_output`` from route-level chat inputs.""" + session = async_session_maker() + try: + requested_llm_config_id = llm_config_id + llm_config_id = ( + await resolve_or_get_pinned_llm_config_id( + session, + thread_id=chat_id, + search_space_id=search_space_id, + user_id=user_id, + selected_llm_config_id=llm_config_id, + requires_image_input=bool(user_image_data_urls), + ) + ).resolved_llm_config_id + + llm: Any + agent_config: AgentConfig | None + if llm_config_id >= 0: + agent_config = await load_agent_config( + session=session, + config_id=llm_config_id, + search_space_id=search_space_id, + ) + if not agent_config: + logger.warning("streaming context build failed: missing config %s", llm_config_id) + return None + llm = create_chat_litellm_from_agent_config(agent_config) + else: + loaded_llm_config = load_global_llm_config_by_id(llm_config_id) + if not loaded_llm_config: + logger.warning( + "streaming context build failed: missing global config %s", + llm_config_id, + ) + return None + llm = create_chat_litellm_from_config(loaded_llm_config) + agent_config = AgentConfig.from_yaml_config(loaded_llm_config) + + connector_service = ConnectorService(session, search_space_id=search_space_id) + firecrawl_api_key = None + webcrawler_connector = await connector_service.get_connector_by_type( + SearchSourceConnectorType.WEBCRAWLER_CONNECTOR, + search_space_id, + ) + if webcrawler_connector and webcrawler_connector.config: + firecrawl_api_key = webcrawler_connector.config.get("FIRECRAWL_API_KEY") + + checkpointer = await get_checkpointer() + visibility = thread_visibility or ChatVisibility.PRIVATE + + from app.config import config as app_config + + agent_factory = ( + create_multi_agent_chat_deep_agent + if bool(app_config.MULTI_AGENT_CHAT_ENABLED) + else create_surfsense_deep_agent + ) + agent = await build_main_agent_for_thread( + agent_factory, + llm=llm, + search_space_id=search_space_id, + db_session=session, + connector_service=connector_service, + checkpointer=checkpointer, + user_id=user_id, + thread_id=chat_id, + agent_config=agent_config, + firecrawl_api_key=firecrawl_api_key, + thread_visibility=visibility, + filesystem_selection=filesystem_selection, + disabled_tools=disabled_tools, + mentioned_document_ids=mentioned_document_ids, + ) + + langchain_messages = [] + if needs_history_bootstrap: + langchain_messages = await bootstrap_history_from_db( + session, + chat_id, + thread_visibility=visibility, + ) + thread_result = await session.execute( + select(NewChatThread).filter(NewChatThread.id == chat_id) + ) + thread = thread_result.scalars().first() + if thread: + thread.needs_history_bootstrap = False + await session.commit() + + mentioned_surfsense_docs: list[SurfsenseDocsDocument] = [] + if mentioned_surfsense_doc_ids: + result = await session.execute( + select(SurfsenseDocsDocument) + .options(selectinload(SurfsenseDocsDocument.chunks)) + .filter(SurfsenseDocsDocument.id.in_(mentioned_surfsense_doc_ids)) + ) + mentioned_surfsense_docs = list(result.scalars().all()) + + recent_reports_result = await session.execute( + select(Report) + .filter(Report.thread_id == chat_id, Report.content.isnot(None)) + .order_by(Report.id.desc()) + .limit(3) + ) + recent_reports = list(recent_reports_result.scalars().all()) + + final_query = user_query + context_parts = [] + if mentioned_surfsense_docs: + context_parts.append( + format_mentioned_surfsense_docs_as_context(mentioned_surfsense_docs) + ) + if recent_reports: + report_lines = [ + f' - report_id={r.id}, title="{r.title}", style="{r.report_style or "detailed"}"' + for r in recent_reports + ] + reports_listing = "\n".join(report_lines) + context_parts.append( + "\n" + "Previously generated reports in this conversation:\n" + f"{reports_listing}\n\n" + "If the user wants to MODIFY, REVISE, UPDATE, or ADD to one of these reports, " + "set parent_report_id to the relevant report_id above.\n" + "If the user wants a completely NEW report on a different topic, " + "leave parent_report_id unset.\n" + "" + ) + if context_parts: + joined_context = "\n\n".join(context_parts) + final_query = f"{joined_context}\n\n{user_query}" + if visibility == ChatVisibility.SEARCH_SPACE and current_user_display_name: + final_query = f"**[{current_user_display_name}]:** {final_query}" + + human_content = build_human_message_content( + final_query, + list(user_image_data_urls or ()), + ) + langchain_messages.append(HumanMessage(content=human_content)) + + turn_id = f"{chat_id}:{int(time.time() * 1000)}" + input_state = { + "messages": langchain_messages, + "search_space_id": search_space_id, + "request_id": request_id or "unknown", + "turn_id": turn_id, + } + configurable = { + "thread_id": str(chat_id), + "request_id": request_id or "unknown", + "turn_id": turn_id, + } + if checkpoint_id: + configurable["checkpoint_id"] = checkpoint_id + config = {"configurable": configurable, "recursion_limit": 10_000} + + initial_title = ( + "Analyzing referenced content" + if mentioned_surfsense_docs + else "Understanding your request" + ) + action_verb = "Analyzing" if mentioned_surfsense_docs else "Processing" + query_excerpt = user_query[:80] + ("..." if len(user_query) > 80 else "") + query_part = query_excerpt if query_excerpt.strip() else "(message)" + initial_items = [f"{action_verb}: {query_part}"] + + runtime_context = SurfSenseContextSchema( + search_space_id=search_space_id, + mentioned_document_ids=list(mentioned_document_ids or []), + request_id=request_id, + turn_id=turn_id, + ) + + await session.commit() + return StreamingContext( + agent=agent, + config=config, + input_data=input_state, + streaming_service=VercelStreamingService(), + step_prefix="thinking", + initial_step_id="thinking-1", + initial_step_title=initial_title, + initial_step_items=initial_items, + content_builder=None, + runtime_context=runtime_context, + ) + except Exception: + logger.exception( + "Failed to build chat streaming context (llm_config_id=%s requested=%s)", + llm_config_id, + requested_llm_config_id, + ) + return None + finally: + await session.close() + diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/regenerate.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/regenerate.py new file mode 100644 index 000000000..02e871a2c --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/regenerate.py @@ -0,0 +1,49 @@ +"""Build ``StreamingContext`` for regenerate streaming.""" + +from __future__ import annotations + +from app.agents.new_chat.filesystem_selection import FilesystemSelection +from app.db import ChatVisibility +from app.tasks.chat.streaming.orchestration.input import StreamingContext +from app.tasks.chat.streaming.orchestration.streaming_context.chat import ( + build_chat_streaming_context, +) + + +async def build_regenerate_streaming_context( + *, + user_query: str, + search_space_id: int, + chat_id: int, + user_id: str | None = None, + llm_config_id: int = -1, + mentioned_document_ids: list[int] | None = None, + mentioned_surfsense_doc_ids: list[int] | None = None, + checkpoint_id: str | None = None, + needs_history_bootstrap: bool = False, + thread_visibility: ChatVisibility | None = None, + current_user_display_name: str | None = None, + disabled_tools: list[str] | None = None, + filesystem_selection: FilesystemSelection | None = None, + request_id: str | None = None, + user_image_data_urls: list[str] | None = None, +) -> StreamingContext | None: + """Build context for ``stream_regenerate`` execution.""" + return await build_chat_streaming_context( + user_query=user_query, + search_space_id=search_space_id, + chat_id=chat_id, + user_id=user_id, + llm_config_id=llm_config_id, + mentioned_document_ids=mentioned_document_ids, + mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids, + checkpoint_id=checkpoint_id, + needs_history_bootstrap=needs_history_bootstrap, + thread_visibility=thread_visibility, + current_user_display_name=current_user_display_name, + disabled_tools=disabled_tools, + filesystem_selection=filesystem_selection, + request_id=request_id, + user_image_data_urls=user_image_data_urls, + ) + diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/resume.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/resume.py new file mode 100644 index 000000000..6d0caea4d --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/resume.py @@ -0,0 +1,154 @@ +"""Build ``StreamingContext`` for resume streaming.""" + +from __future__ import annotations + +import logging +import time +from typing import Any + +from langgraph.types import Command + +from app.agents.multi_agent_chat import create_multi_agent_chat_deep_agent +from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent +from app.agents.new_chat.checkpointer import get_checkpointer +from app.agents.new_chat.context import SurfSenseContextSchema +from app.agents.new_chat.filesystem_selection import FilesystemSelection +from app.agents.new_chat.llm_config import ( + AgentConfig, + create_chat_litellm_from_agent_config, + create_chat_litellm_from_config, + load_agent_config, + load_global_llm_config_by_id, +) +from app.db import ChatVisibility, SearchSourceConnectorType, async_session_maker +from app.services.auto_model_pin_service import resolve_or_get_pinned_llm_config_id +from app.services.connector_service import ConnectorService +from app.services.new_streaming_service import VercelStreamingService +from app.tasks.chat.streaming.agent_setup import build_main_agent_for_thread +from app.tasks.chat.streaming.orchestration.input import StreamingContext + +logger = logging.getLogger(__name__) + + +async def build_resume_streaming_context( + *, + chat_id: int, + search_space_id: int, + decisions: list[dict], + user_id: str | None = None, + llm_config_id: int = -1, + thread_visibility: ChatVisibility | None = None, + filesystem_selection: FilesystemSelection | None = None, + request_id: str | None = None, + disabled_tools: list[str] | None = None, +) -> StreamingContext | None: + """Build context for ``stream_resume`` execution.""" + session = async_session_maker() + try: + llm_config_id = ( + await resolve_or_get_pinned_llm_config_id( + session, + thread_id=chat_id, + search_space_id=search_space_id, + user_id=user_id, + selected_llm_config_id=llm_config_id, + ) + ).resolved_llm_config_id + + llm: Any + agent_config: AgentConfig | None + if llm_config_id >= 0: + agent_config = await load_agent_config( + session=session, + config_id=llm_config_id, + search_space_id=search_space_id, + ) + if not agent_config: + logger.warning("resume context build failed: missing config %s", llm_config_id) + return None + llm = create_chat_litellm_from_agent_config(agent_config) + else: + loaded_llm_config = load_global_llm_config_by_id(llm_config_id) + if not loaded_llm_config: + logger.warning( + "resume context build failed: missing global config %s", + llm_config_id, + ) + return None + llm = create_chat_litellm_from_config(loaded_llm_config) + agent_config = AgentConfig.from_yaml_config(loaded_llm_config) + + connector_service = ConnectorService(session, search_space_id=search_space_id) + firecrawl_api_key = None + webcrawler_connector = await connector_service.get_connector_by_type( + SearchSourceConnectorType.WEBCRAWLER_CONNECTOR, + search_space_id, + ) + if webcrawler_connector and webcrawler_connector.config: + firecrawl_api_key = webcrawler_connector.config.get("FIRECRAWL_API_KEY") + + checkpointer = await get_checkpointer() + visibility = thread_visibility or ChatVisibility.PRIVATE + + from app.config import config as app_config + + agent_factory = ( + create_multi_agent_chat_deep_agent + if bool(app_config.MULTI_AGENT_CHAT_ENABLED) + else create_surfsense_deep_agent + ) + agent = await build_main_agent_for_thread( + agent_factory, + llm=llm, + search_space_id=search_space_id, + db_session=session, + connector_service=connector_service, + checkpointer=checkpointer, + user_id=user_id, + thread_id=chat_id, + agent_config=agent_config, + firecrawl_api_key=firecrawl_api_key, + thread_visibility=visibility, + filesystem_selection=filesystem_selection, + disabled_tools=disabled_tools, + ) + + turn_id = f"{chat_id}:{int(time.time() * 1000)}" + config = { + "configurable": { + "thread_id": str(chat_id), + "request_id": request_id or "unknown", + "turn_id": turn_id, + "surfsense_resume_value": {"decisions": decisions}, + }, + "recursion_limit": 10_000, + } + + runtime_context = SurfSenseContextSchema( + search_space_id=search_space_id, + request_id=request_id, + turn_id=turn_id, + ) + + await session.commit() + return StreamingContext( + agent=agent, + config=config, + input_data=Command(resume={"decisions": decisions}), + streaming_service=VercelStreamingService(), + step_prefix="thinking-resume", + initial_step_id=None, + initial_step_title="", + initial_step_items=None, + content_builder=None, + runtime_context=runtime_context, + ) + except Exception: + logger.exception( + "Failed to build resume streaming context (llm_config_id=%s)", + llm_config_id, + ) + return None + finally: + await session.close() + diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py index b84193cb7..46c61b498 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py @@ -8,6 +8,7 @@ from typing import Any import pytest from app.tasks.chat.streaming.orchestration import StreamingContext +from app.tasks.chat.streaming.orchestration import orchestrator from app.tasks.chat.streaming.orchestration.orchestrator import ( stream_chat, stream_regenerate, @@ -138,3 +139,102 @@ async def test_stream_regenerate_uses_streaming_context_path() -> None: "text_delta:text-1:g", "text_end:text-1", ] + + +async def test_stream_chat_builds_streaming_context_when_not_provided() -> None: + service = _StreamingService() + agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("b")}}]) + + async def _fake_builder(**kwargs: Any) -> StreamingContext: + del kwargs + return StreamingContext( + agent=agent, + config={"configurable": {"thread_id": "thread-b"}}, + input_data={"messages": []}, + streaming_service=service, + ) + + old = orchestrator.build_chat_streaming_context + orchestrator.build_chat_streaming_context = _fake_builder + try: + frames = await _collect( + stream_chat( + user_query="q", + search_space_id=1, + chat_id=3, + ) + ) + finally: + orchestrator.build_chat_streaming_context = old + + assert frames == [ + "text_start:text-1", + "text_delta:text-1:b", + "text_end:text-1", + ] + + +async def test_stream_resume_builds_streaming_context_when_not_provided() -> None: + service = _StreamingService() + agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("u")}}]) + + async def _fake_builder(**kwargs: Any) -> StreamingContext: + del kwargs + return StreamingContext( + agent=agent, + config={"configurable": {"thread_id": "thread-u"}}, + input_data={"messages": []}, + streaming_service=service, + ) + + old = orchestrator.build_resume_streaming_context + orchestrator.build_resume_streaming_context = _fake_builder + try: + frames = await _collect( + stream_resume( + chat_id=9, + search_space_id=1, + decisions=[], + ) + ) + finally: + orchestrator.build_resume_streaming_context = old + + assert frames == [ + "text_start:text-1", + "text_delta:text-1:u", + "text_end:text-1", + ] + + +async def test_stream_regenerate_builds_streaming_context_when_not_provided() -> None: + service = _StreamingService() + agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("x")}}]) + + async def _fake_builder(**kwargs: Any) -> StreamingContext: + del kwargs + return StreamingContext( + agent=agent, + config={"configurable": {"thread_id": "thread-x"}}, + input_data={"messages": []}, + streaming_service=service, + ) + + old = orchestrator.build_regenerate_streaming_context + orchestrator.build_regenerate_streaming_context = _fake_builder + try: + frames = await _collect( + stream_regenerate( + user_query="q", + search_space_id=1, + chat_id=2, + ) + ) + finally: + orchestrator.build_regenerate_streaming_context = old + + assert frames == [ + "text_start:text-1", + "text_delta:text-1:x", + "text_end:text-1", + ] From 7e07092f67d55a81916b74a5d3691dea5461f273 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 7 May 2026 19:25:20 +0200 Subject: [PATCH 59/96] refactor(chat): drop alternate streaming entry path; use graph_stream --- surfsense_backend/app/config/__init__.py | 6 - .../app/routes/new_chat_routes.py | 32 +-- surfsense_backend/app/schemas/new_chat.py | 2 +- .../app/services/streaming/__init__.py | 6 +- .../app/services/streaming/events/error.py | 2 +- .../app/tasks/chat/streaming/__init__.py | 2 +- .../app/tasks/chat/streaming/agent_setup.py | 92 ------ .../chat/streaming/graph_stream/__init__.py | 21 ++ .../event_stream.py | 4 +- .../output.py => graph_stream/result.py} | 3 +- .../chat/streaming/orchestration/__init__.py | 11 - .../chat/streaming/orchestration/input.py | 23 -- .../streaming/orchestration/orchestrator.py | 261 ------------------ .../streaming_context/__init__.py | 18 -- .../orchestration/streaming_context/chat.py | 258 ----------------- .../streaming_context/regenerate.py | 49 ---- .../orchestration/streaming_context/resume.py | 154 ----------- .../tasks/chat/streaming/relay/__init__.py | 22 +- .../tasks/chat/streaming/relay/event_relay.py | 2 +- .../tasks/chat/streaming/test_agent_setup.py | 120 -------- .../test_orchestrator_stream_chat.py | 240 ---------------- .../chat/streaming/test_stage_1_parity.py | 4 +- ..._event_stream.py => test_stream_output.py} | 7 +- 23 files changed, 61 insertions(+), 1278 deletions(-) delete mode 100644 surfsense_backend/app/tasks/chat/streaming/agent_setup.py create mode 100644 surfsense_backend/app/tasks/chat/streaming/graph_stream/__init__.py rename surfsense_backend/app/tasks/chat/streaming/{orchestration => graph_stream}/event_stream.py (92%) rename surfsense_backend/app/tasks/chat/streaming/{orchestration/output.py => graph_stream/result.py} (91%) delete mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py delete mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/input.py delete mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py delete mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/__init__.py delete mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/chat.py delete mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/regenerate.py delete mode 100644 surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/resume.py delete mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/test_agent_setup.py delete mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py rename surfsense_backend/tests/unit/tasks/chat/streaming/{test_orchestration_event_stream.py => test_stream_output.py} (93%) diff --git a/surfsense_backend/app/config/__init__.py b/surfsense_backend/app/config/__init__.py index 543524456..f6f0c7f62 100644 --- a/surfsense_backend/app/config/__init__.py +++ b/surfsense_backend/app/config/__init__.py @@ -490,12 +490,6 @@ class Config: ENABLE_DESKTOP_LOCAL_FILESYSTEM = ( os.getenv("ENABLE_DESKTOP_LOCAL_FILESYSTEM", "FALSE").upper() == "TRUE" ) - # Streaming entrypoint switch. Keep this at the route layer so orchestrator - # code stays free of legacy fallback branching. - ENABLE_CHAT_STREAM_ORCHESTRATOR = ( - os.getenv("SURFSENSE_ENABLE_CHAT_STREAM_ORCHESTRATOR", "TRUE").upper() - == "TRUE" - ) @classmethod def is_self_hosted(cls) -> bool: diff --git a/surfsense_backend/app/routes/new_chat_routes.py b/surfsense_backend/app/routes/new_chat_routes.py index e54497f93..743b5b849 100644 --- a/surfsense_backend/app/routes/new_chat_routes.py +++ b/surfsense_backend/app/routes/new_chat_routes.py @@ -72,13 +72,8 @@ from app.schemas.new_chat import ( TurnStatusResponse, ) from app.tasks.chat.stream_new_chat import ( - stream_new_chat as legacy_stream_new_chat, - stream_resume_chat as legacy_stream_resume_chat, -) -from app.tasks.chat.streaming.orchestration.orchestrator import ( - stream_chat, - stream_regenerate, - stream_resume, + stream_new_chat, + stream_resume_chat, ) from app.users import current_active_user from app.utils.perf import get_perf_logger @@ -98,10 +93,6 @@ TURN_CANCELLING_MAX_DELAY_MS = 1500 router = APIRouter() -def _use_streaming_orchestrator() -> bool: - return config.ENABLE_CHAT_STREAM_ORCHESTRATOR - - def _resolve_filesystem_selection( *, mode: str, @@ -1782,11 +1773,7 @@ async def handle_new_chat( ) return StreamingResponse( - ( - stream_chat - if _use_streaming_orchestrator() - else legacy_stream_new_chat - )( + stream_new_chat( user_query=request.user_query, search_space_id=request.search_space_id, chat_id=request.chat_id, @@ -2271,12 +2258,7 @@ async def regenerate_response( else None ) try: - regenerate_fn = ( - stream_regenerate - if _use_streaming_orchestrator() - else legacy_stream_new_chat - ) - async for chunk in regenerate_fn( + async for chunk in stream_new_chat( user_query=str(user_query_to_use), search_space_id=request.search_space_id, chat_id=thread_id, @@ -2408,11 +2390,7 @@ async def resume_chat( await session.close() return StreamingResponse( - ( - stream_resume - if _use_streaming_orchestrator() - else legacy_stream_resume_chat - )( + stream_resume_chat( chat_id=thread_id, search_space_id=request.search_space_id, decisions=decisions, diff --git a/surfsense_backend/app/schemas/new_chat.py b/surfsense_backend/app/schemas/new_chat.py index 95d183433..fe8dab076 100644 --- a/surfsense_backend/app/schemas/new_chat.py +++ b/surfsense_backend/app/schemas/new_chat.py @@ -380,7 +380,7 @@ class ResumeRequest(BaseModel): "/regenerate. Resume reuses the original interrupted user " "turn so the server does not write a new user message. " "Currently unused but accepted to keep request bodies " - "uniform across the three streaming entrypoints." + "uniform across new-message, regenerate, and resume stream routes." ), ) diff --git a/surfsense_backend/app/services/streaming/__init__.py b/surfsense_backend/app/services/streaming/__init__.py index 287d48a7a..3ec9b9cf1 100644 --- a/surfsense_backend/app/services/streaming/__init__.py +++ b/surfsense_backend/app/services/streaming/__init__.py @@ -4,7 +4,7 @@ Layout: * ``envelope/`` - SSE wire framing + ID generators * ``emitter/`` - identity of the agent that emitted an event + runtime registry * ``events/`` - one module per SSE event family -* ``service.py`` - composition root used by the orchestrator +* ``service.py`` - composition root used when emitting chat SSE * ``interrupt_correlation.py`` - id-aware lookup over LangGraph state Naming on the wire: @@ -13,8 +13,8 @@ Naming on the wire: * Every SurfSense-added field uses ``snake_case``, including the top-level ``emitted_by`` envelope and all inner ``data`` payloads. -Production keeps using ``app.services.new_streaming_service`` and -``app.tasks.chat.stream_new_chat`` until the cutover phase. +Production chat uses ``app.services.new_streaming_service`` from +``app.tasks.chat.stream_new_chat`` and related routes. """ from __future__ import annotations diff --git a/surfsense_backend/app/services/streaming/events/error.py b/surfsense_backend/app/services/streaming/events/error.py index cd190d1f4..a1e8e01ca 100644 --- a/surfsense_backend/app/services/streaming/events/error.py +++ b/surfsense_backend/app/services/streaming/events/error.py @@ -1,4 +1,4 @@ -"""Single terminal error path the orchestrator must route through.""" +"""Single terminal error path chat streaming must route through.""" from __future__ import annotations diff --git a/surfsense_backend/app/tasks/chat/streaming/__init__.py b/surfsense_backend/app/tasks/chat/streaming/__init__.py index bb06cc021..70c99342a 100644 --- a/surfsense_backend/app/tasks/chat/streaming/__init__.py +++ b/surfsense_backend/app/tasks/chat/streaming/__init__.py @@ -1,3 +1,3 @@ -"""Chat streaming orchestrator and event relay.""" +"""Chat streaming helpers (e.g. LangGraph → SSE relay under ``graph_stream``).""" from __future__ import annotations diff --git a/surfsense_backend/app/tasks/chat/streaming/agent_setup.py b/surfsense_backend/app/tasks/chat/streaming/agent_setup.py deleted file mode 100644 index f67c6ad65..000000000 --- a/surfsense_backend/app/tasks/chat/streaming/agent_setup.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Agent setup helpers for orchestrated chat streaming.""" - -from __future__ import annotations - -import contextlib -import logging -from collections.abc import Callable -from typing import Any - -_PREFLIGHT_TIMEOUT_SEC: float = 2.5 -_PREFLIGHT_MAX_TOKENS: int = 1 - - -async def preflight_llm( - llm: Any, - *, - is_provider_rate_limited: Callable[[BaseException], bool], -) -> None: - """Issue a minimal completion probe to catch immediate provider 429s.""" - from litellm import acompletion - - model = getattr(llm, "model", None) - if not model or model == "auto": - return - - try: - await acompletion( - model=model, - messages=[{"role": "user", "content": "ping"}], - api_key=getattr(llm, "api_key", None), - api_base=getattr(llm, "api_base", None), - max_tokens=_PREFLIGHT_MAX_TOKENS, - timeout=_PREFLIGHT_TIMEOUT_SEC, - stream=False, - metadata={"tags": ["surfsense:internal", "auto-pin-preflight"]}, - ) - except Exception as exc: - if is_provider_rate_limited(exc): - raise - logging.getLogger(__name__).debug( - "auto_pin_preflight non_rate_limit_error model=%s err=%s", - model, - exc, - ) - - -async def build_main_agent_for_thread( - agent_factory: Any, - *, - llm: Any, - search_space_id: int, - db_session: Any, - connector_service: Any, - checkpointer: Any, - user_id: str | None, - thread_id: int | None, - agent_config: Any, - firecrawl_api_key: str | None, - thread_visibility: Any, - filesystem_selection: Any, - disabled_tools: list[str] | None = None, - mentioned_document_ids: list[int] | None = None, -) -> Any: - """Run one canonical agent-build call for a single thread.""" - return await agent_factory( - llm=llm, - search_space_id=search_space_id, - db_session=db_session, - connector_service=connector_service, - checkpointer=checkpointer, - user_id=user_id, - thread_id=thread_id, - agent_config=agent_config, - firecrawl_api_key=firecrawl_api_key, - thread_visibility=thread_visibility, - filesystem_selection=filesystem_selection, - disabled_tools=disabled_tools, - mentioned_document_ids=mentioned_document_ids, - ) - - -async def settle_speculative_agent_build(task: Any) -> None: - """Wait for a discarded speculative build and swallow its outcome.""" - with contextlib.suppress(BaseException): - await task - - -__all__ = [ - "build_main_agent_for_thread", - "preflight_llm", - "settle_speculative_agent_build", -] diff --git a/surfsense_backend/app/tasks/chat/streaming/graph_stream/__init__.py b/surfsense_backend/app/tasks/chat/streaming/graph_stream/__init__.py new file mode 100644 index 000000000..e3bf0426c --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/graph_stream/__init__.py @@ -0,0 +1,21 @@ +"""LangGraph ``astream_events`` → SSE (``stream_output`` + ``StreamingResult``). + +Imports are lazy to avoid a circular import with ``relay.event_relay``. +""" + +from __future__ import annotations + +__all__ = ["StreamingResult", "stream_output"] + + +def __getattr__(name: str): + if name == "stream_output": + from app.tasks.chat.streaming.graph_stream.event_stream import stream_output + + return stream_output + if name == "StreamingResult": + from app.tasks.chat.streaming.graph_stream.result import StreamingResult + + return StreamingResult + msg = f"module {__name__!r} has no attribute {name!r}" + raise AttributeError(msg) diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py b/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py similarity index 92% rename from surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py rename to surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py index fc8c13027..9142dd914 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/event_stream.py +++ b/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py @@ -1,4 +1,4 @@ -"""Run LangGraph event streams through the EventRelay.""" +"""Run LangGraph event streams through ``EventRelay``.""" from __future__ import annotations @@ -6,7 +6,7 @@ from collections.abc import AsyncIterator from typing import Any from app.agents.new_chat.feature_flags import get_flags -from app.tasks.chat.streaming.orchestration.output import StreamingResult +from app.tasks.chat.streaming.graph_stream.result import StreamingResult from app.tasks.chat.streaming.relay.event_relay import EventRelay from app.tasks.chat.streaming.relay.state import AgentEventRelayState diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/output.py b/surfsense_backend/app/tasks/chat/streaming/graph_stream/result.py similarity index 91% rename from surfsense_backend/app/tasks/chat/streaming/orchestration/output.py rename to surfsense_backend/app/tasks/chat/streaming/graph_stream/result.py index 60f8ee6ee..40404e9d0 100644 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/output.py +++ b/surfsense_backend/app/tasks/chat/streaming/graph_stream/result.py @@ -1,4 +1,4 @@ -"""Output facts collected while streaming one orchestrated agent turn.""" +"""Mutable facts collected while relaying one agent stream (``stream_output``).""" from __future__ import annotations @@ -26,4 +26,3 @@ class StreamingResult: commit_gate_reason: str = "" assistant_message_id: int | None = None content_builder: Any | None = field(default=None, repr=False) - diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py deleted file mode 100644 index b1a201fd3..000000000 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Composable orchestration pieces for chat streaming.""" - -from app.tasks.chat.streaming.orchestration.event_stream import stream_output -from app.tasks.chat.streaming.orchestration.input import StreamingContext -from app.tasks.chat.streaming.orchestration.output import StreamingResult - -__all__ = [ - "StreamingContext", - "StreamingResult", - "stream_output", -] diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/input.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/input.py deleted file mode 100644 index 45a33d435..000000000 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/input.py +++ /dev/null @@ -1,23 +0,0 @@ -"""Inputs for orchestrator-owned streaming execution.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - - -@dataclass(frozen=True) -class StreamingContext: - """Container for dependencies required by ``stream_output``.""" - - agent: Any - config: dict[str, Any] - input_data: Any - streaming_service: Any - step_prefix: str = "thinking" - initial_step_id: str | None = None - initial_step_title: str = "" - initial_step_items: list[str] | None = None - content_builder: Any | None = None - runtime_context: Any = None - diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py deleted file mode 100644 index 80cae77a2..000000000 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/orchestrator.py +++ /dev/null @@ -1,261 +0,0 @@ -"""Top-level chat streaming entrypoints. -""" - -from __future__ import annotations - -from collections.abc import AsyncGenerator -from typing import Any, Literal - -from app.agents.new_chat.filesystem_selection import FilesystemSelection -from app.db import ChatVisibility -from app.tasks.chat.stream_new_chat import stream_new_chat, stream_resume_chat -from app.tasks.chat.streaming.orchestration.streaming_context import ( - build_chat_streaming_context, - build_regenerate_streaming_context, - build_resume_streaming_context, -) -from app.tasks.chat.streaming.orchestration.event_stream import stream_output -from app.tasks.chat.streaming.orchestration.input import StreamingContext -from app.tasks.chat.streaming.orchestration.output import StreamingResult - - -def _build_streaming_result( - *, - chat_id: int, - request_id: str | None, - filesystem_selection: FilesystemSelection | None, - suffix: str, -) -> StreamingResult: - return StreamingResult( - request_id=request_id, - turn_id=f"{chat_id}:{suffix}", - filesystem_mode=(filesystem_selection.mode.value if filesystem_selection else "cloud"), - client_platform=( - filesystem_selection.client_platform.value if filesystem_selection else "web" - ), - ) - - -async def _stream_output_with_streaming_context( - *, - streaming_context: StreamingContext, - result: StreamingResult, -) -> AsyncGenerator[str, None]: - async for frame in stream_output( - agent=streaming_context.agent, - config=streaming_context.config, - input_data=streaming_context.input_data, - streaming_service=streaming_context.streaming_service, - result=result, - step_prefix=streaming_context.step_prefix, - initial_step_id=streaming_context.initial_step_id, - initial_step_title=streaming_context.initial_step_title, - initial_step_items=streaming_context.initial_step_items, - content_builder=streaming_context.content_builder, - runtime_context=streaming_context.runtime_context, - ): - yield frame - - -async def stream_chat( - *, - user_query: str, - search_space_id: int, - chat_id: int, - user_id: str | None = None, - llm_config_id: int = -1, - mentioned_document_ids: list[int] | None = None, - mentioned_surfsense_doc_ids: list[int] | None = None, - mentioned_documents: list[dict[str, Any]] | None = None, - checkpoint_id: str | None = None, - needs_history_bootstrap: bool = False, - thread_visibility: ChatVisibility | None = None, - current_user_display_name: str | None = None, - disabled_tools: list[str] | None = None, - filesystem_selection: FilesystemSelection | None = None, - request_id: str | None = None, - user_image_data_urls: list[str] | None = None, - streaming_context: StreamingContext | None = None, -) -> AsyncGenerator[str, None]: - """Stream a new chat turn through the current production pipeline.""" - if streaming_context is None: - streaming_context = await build_chat_streaming_context( - user_query=user_query, - search_space_id=search_space_id, - chat_id=chat_id, - user_id=user_id, - llm_config_id=llm_config_id, - mentioned_document_ids=mentioned_document_ids, - mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids, - checkpoint_id=checkpoint_id, - needs_history_bootstrap=needs_history_bootstrap, - thread_visibility=thread_visibility, - current_user_display_name=current_user_display_name, - disabled_tools=disabled_tools, - filesystem_selection=filesystem_selection, - request_id=request_id, - user_image_data_urls=user_image_data_urls, - ) - if streaming_context is not None: - result = _build_streaming_result( - chat_id=chat_id, - request_id=request_id, - filesystem_selection=filesystem_selection, - suffix="orchestrator", - ) - async for frame in _stream_output_with_streaming_context( - streaming_context=streaming_context, - result=result, - ): - yield frame - return - - async for chunk in stream_new_chat( - user_query=user_query, - search_space_id=search_space_id, - chat_id=chat_id, - user_id=user_id, - llm_config_id=llm_config_id, - mentioned_document_ids=mentioned_document_ids, - mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids, - mentioned_documents=mentioned_documents, - checkpoint_id=checkpoint_id, - needs_history_bootstrap=needs_history_bootstrap, - thread_visibility=thread_visibility, - current_user_display_name=current_user_display_name, - disabled_tools=disabled_tools, - filesystem_selection=filesystem_selection, - request_id=request_id, - user_image_data_urls=user_image_data_urls, - ): - yield chunk - - -async def stream_resume( - *, - chat_id: int, - search_space_id: int, - decisions: list[dict], - user_id: str | None = None, - llm_config_id: int = -1, - thread_visibility: ChatVisibility | None = None, - filesystem_selection: FilesystemSelection | None = None, - request_id: str | None = None, - disabled_tools: list[str] | None = None, - streaming_context: StreamingContext | None = None, -) -> AsyncGenerator[str, None]: - """Resume an interrupted chat turn through the current production pipeline.""" - if streaming_context is None: - streaming_context = await build_resume_streaming_context( - chat_id=chat_id, - search_space_id=search_space_id, - decisions=decisions, - user_id=user_id, - llm_config_id=llm_config_id, - thread_visibility=thread_visibility, - filesystem_selection=filesystem_selection, - request_id=request_id, - disabled_tools=disabled_tools, - ) - if streaming_context is not None: - result = _build_streaming_result( - chat_id=chat_id, - request_id=request_id, - filesystem_selection=filesystem_selection, - suffix="orchestrator-resume", - ) - async for frame in _stream_output_with_streaming_context( - streaming_context=streaming_context, - result=result, - ): - yield frame - return - - async for chunk in stream_resume_chat( - chat_id=chat_id, - search_space_id=search_space_id, - decisions=decisions, - user_id=user_id, - llm_config_id=llm_config_id, - thread_visibility=thread_visibility, - filesystem_selection=filesystem_selection, - request_id=request_id, - disabled_tools=disabled_tools, - ): - yield chunk - - -async def stream_regenerate( - *, - user_query: str, - search_space_id: int, - chat_id: int, - user_id: str | None = None, - llm_config_id: int = -1, - mentioned_document_ids: list[int] | None = None, - mentioned_surfsense_doc_ids: list[int] | None = None, - mentioned_documents: list[dict[str, Any]] | None = None, - checkpoint_id: str | None = None, - needs_history_bootstrap: bool = False, - thread_visibility: ChatVisibility | None = None, - current_user_display_name: str | None = None, - disabled_tools: list[str] | None = None, - filesystem_selection: FilesystemSelection | None = None, - request_id: str | None = None, - user_image_data_urls: list[str] | None = None, - flow: Literal["new", "regenerate"] = "regenerate", - streaming_context: StreamingContext | None = None, -) -> AsyncGenerator[str, None]: - """Regenerate an assistant turn through the current production pipeline.""" - if streaming_context is None: - streaming_context = await build_regenerate_streaming_context( - user_query=user_query, - search_space_id=search_space_id, - chat_id=chat_id, - user_id=user_id, - llm_config_id=llm_config_id, - mentioned_document_ids=mentioned_document_ids, - mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids, - checkpoint_id=checkpoint_id, - needs_history_bootstrap=needs_history_bootstrap, - thread_visibility=thread_visibility, - current_user_display_name=current_user_display_name, - disabled_tools=disabled_tools, - filesystem_selection=filesystem_selection, - request_id=request_id, - user_image_data_urls=user_image_data_urls, - ) - if streaming_context is not None: - result = _build_streaming_result( - chat_id=chat_id, - request_id=request_id, - filesystem_selection=filesystem_selection, - suffix="orchestrator-regenerate", - ) - async for frame in _stream_output_with_streaming_context( - streaming_context=streaming_context, - result=result, - ): - yield frame - return - - async for chunk in stream_new_chat( - user_query=user_query, - search_space_id=search_space_id, - chat_id=chat_id, - user_id=user_id, - llm_config_id=llm_config_id, - mentioned_document_ids=mentioned_document_ids, - mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids, - mentioned_documents=mentioned_documents, - checkpoint_id=checkpoint_id, - needs_history_bootstrap=needs_history_bootstrap, - thread_visibility=thread_visibility, - current_user_display_name=current_user_display_name, - disabled_tools=disabled_tools, - filesystem_selection=filesystem_selection, - request_id=request_id, - user_image_data_urls=user_image_data_urls, - flow=flow, - ): - yield chunk diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/__init__.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/__init__.py deleted file mode 100644 index 1bd3e103d..000000000 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Streaming context builders per orchestrator entrypoint.""" - -from app.tasks.chat.streaming.orchestration.streaming_context.chat import ( - build_chat_streaming_context, -) -from app.tasks.chat.streaming.orchestration.streaming_context.regenerate import ( - build_regenerate_streaming_context, -) -from app.tasks.chat.streaming.orchestration.streaming_context.resume import ( - build_resume_streaming_context, -) - -__all__ = [ - "build_chat_streaming_context", - "build_regenerate_streaming_context", - "build_resume_streaming_context", -] - diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/chat.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/chat.py deleted file mode 100644 index eb459ae5c..000000000 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/chat.py +++ /dev/null @@ -1,258 +0,0 @@ -"""Build ``StreamingContext`` for chat streaming.""" - -from __future__ import annotations - -import logging -import time -from typing import Any - -from langchain_core.messages import HumanMessage -from sqlalchemy.future import select -from sqlalchemy.orm import selectinload - -from app.agents.multi_agent_chat import create_multi_agent_chat_deep_agent -from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent -from app.agents.new_chat.checkpointer import get_checkpointer -from app.agents.new_chat.context import SurfSenseContextSchema -from app.agents.new_chat.filesystem_selection import FilesystemSelection -from app.agents.new_chat.llm_config import ( - AgentConfig, - create_chat_litellm_from_agent_config, - create_chat_litellm_from_config, - load_agent_config, - load_global_llm_config_by_id, -) -from app.db import ( - ChatVisibility, - NewChatThread, - Report, - SearchSourceConnectorType, - SurfsenseDocsDocument, - async_session_maker, -) -from app.services.auto_model_pin_service import resolve_or_get_pinned_llm_config_id -from app.services.connector_service import ConnectorService -from app.services.new_streaming_service import VercelStreamingService -from app.tasks.chat.stream_new_chat import format_mentioned_surfsense_docs_as_context -from app.tasks.chat.streaming.agent_setup import build_main_agent_for_thread -from app.tasks.chat.streaming.orchestration.input import StreamingContext -from app.utils.content_utils import bootstrap_history_from_db -from app.utils.user_message_multimodal import build_human_message_content - -logger = logging.getLogger(__name__) - - -async def build_chat_streaming_context( - *, - user_query: str, - search_space_id: int, - chat_id: int, - user_id: str | None = None, - llm_config_id: int = -1, - mentioned_document_ids: list[int] | None = None, - mentioned_surfsense_doc_ids: list[int] | None = None, - checkpoint_id: str | None = None, - needs_history_bootstrap: bool = False, - thread_visibility: ChatVisibility | None = None, - current_user_display_name: str | None = None, - disabled_tools: list[str] | None = None, - filesystem_selection: FilesystemSelection | None = None, - request_id: str | None = None, - user_image_data_urls: list[str] | None = None, -) -> StreamingContext | None: - """Build context for ``stream_output`` from route-level chat inputs.""" - session = async_session_maker() - try: - requested_llm_config_id = llm_config_id - llm_config_id = ( - await resolve_or_get_pinned_llm_config_id( - session, - thread_id=chat_id, - search_space_id=search_space_id, - user_id=user_id, - selected_llm_config_id=llm_config_id, - requires_image_input=bool(user_image_data_urls), - ) - ).resolved_llm_config_id - - llm: Any - agent_config: AgentConfig | None - if llm_config_id >= 0: - agent_config = await load_agent_config( - session=session, - config_id=llm_config_id, - search_space_id=search_space_id, - ) - if not agent_config: - logger.warning("streaming context build failed: missing config %s", llm_config_id) - return None - llm = create_chat_litellm_from_agent_config(agent_config) - else: - loaded_llm_config = load_global_llm_config_by_id(llm_config_id) - if not loaded_llm_config: - logger.warning( - "streaming context build failed: missing global config %s", - llm_config_id, - ) - return None - llm = create_chat_litellm_from_config(loaded_llm_config) - agent_config = AgentConfig.from_yaml_config(loaded_llm_config) - - connector_service = ConnectorService(session, search_space_id=search_space_id) - firecrawl_api_key = None - webcrawler_connector = await connector_service.get_connector_by_type( - SearchSourceConnectorType.WEBCRAWLER_CONNECTOR, - search_space_id, - ) - if webcrawler_connector and webcrawler_connector.config: - firecrawl_api_key = webcrawler_connector.config.get("FIRECRAWL_API_KEY") - - checkpointer = await get_checkpointer() - visibility = thread_visibility or ChatVisibility.PRIVATE - - from app.config import config as app_config - - agent_factory = ( - create_multi_agent_chat_deep_agent - if bool(app_config.MULTI_AGENT_CHAT_ENABLED) - else create_surfsense_deep_agent - ) - agent = await build_main_agent_for_thread( - agent_factory, - llm=llm, - search_space_id=search_space_id, - db_session=session, - connector_service=connector_service, - checkpointer=checkpointer, - user_id=user_id, - thread_id=chat_id, - agent_config=agent_config, - firecrawl_api_key=firecrawl_api_key, - thread_visibility=visibility, - filesystem_selection=filesystem_selection, - disabled_tools=disabled_tools, - mentioned_document_ids=mentioned_document_ids, - ) - - langchain_messages = [] - if needs_history_bootstrap: - langchain_messages = await bootstrap_history_from_db( - session, - chat_id, - thread_visibility=visibility, - ) - thread_result = await session.execute( - select(NewChatThread).filter(NewChatThread.id == chat_id) - ) - thread = thread_result.scalars().first() - if thread: - thread.needs_history_bootstrap = False - await session.commit() - - mentioned_surfsense_docs: list[SurfsenseDocsDocument] = [] - if mentioned_surfsense_doc_ids: - result = await session.execute( - select(SurfsenseDocsDocument) - .options(selectinload(SurfsenseDocsDocument.chunks)) - .filter(SurfsenseDocsDocument.id.in_(mentioned_surfsense_doc_ids)) - ) - mentioned_surfsense_docs = list(result.scalars().all()) - - recent_reports_result = await session.execute( - select(Report) - .filter(Report.thread_id == chat_id, Report.content.isnot(None)) - .order_by(Report.id.desc()) - .limit(3) - ) - recent_reports = list(recent_reports_result.scalars().all()) - - final_query = user_query - context_parts = [] - if mentioned_surfsense_docs: - context_parts.append( - format_mentioned_surfsense_docs_as_context(mentioned_surfsense_docs) - ) - if recent_reports: - report_lines = [ - f' - report_id={r.id}, title="{r.title}", style="{r.report_style or "detailed"}"' - for r in recent_reports - ] - reports_listing = "\n".join(report_lines) - context_parts.append( - "\n" - "Previously generated reports in this conversation:\n" - f"{reports_listing}\n\n" - "If the user wants to MODIFY, REVISE, UPDATE, or ADD to one of these reports, " - "set parent_report_id to the relevant report_id above.\n" - "If the user wants a completely NEW report on a different topic, " - "leave parent_report_id unset.\n" - "" - ) - if context_parts: - joined_context = "\n\n".join(context_parts) - final_query = f"{joined_context}\n\n{user_query}" - if visibility == ChatVisibility.SEARCH_SPACE and current_user_display_name: - final_query = f"**[{current_user_display_name}]:** {final_query}" - - human_content = build_human_message_content( - final_query, - list(user_image_data_urls or ()), - ) - langchain_messages.append(HumanMessage(content=human_content)) - - turn_id = f"{chat_id}:{int(time.time() * 1000)}" - input_state = { - "messages": langchain_messages, - "search_space_id": search_space_id, - "request_id": request_id or "unknown", - "turn_id": turn_id, - } - configurable = { - "thread_id": str(chat_id), - "request_id": request_id or "unknown", - "turn_id": turn_id, - } - if checkpoint_id: - configurable["checkpoint_id"] = checkpoint_id - config = {"configurable": configurable, "recursion_limit": 10_000} - - initial_title = ( - "Analyzing referenced content" - if mentioned_surfsense_docs - else "Understanding your request" - ) - action_verb = "Analyzing" if mentioned_surfsense_docs else "Processing" - query_excerpt = user_query[:80] + ("..." if len(user_query) > 80 else "") - query_part = query_excerpt if query_excerpt.strip() else "(message)" - initial_items = [f"{action_verb}: {query_part}"] - - runtime_context = SurfSenseContextSchema( - search_space_id=search_space_id, - mentioned_document_ids=list(mentioned_document_ids or []), - request_id=request_id, - turn_id=turn_id, - ) - - await session.commit() - return StreamingContext( - agent=agent, - config=config, - input_data=input_state, - streaming_service=VercelStreamingService(), - step_prefix="thinking", - initial_step_id="thinking-1", - initial_step_title=initial_title, - initial_step_items=initial_items, - content_builder=None, - runtime_context=runtime_context, - ) - except Exception: - logger.exception( - "Failed to build chat streaming context (llm_config_id=%s requested=%s)", - llm_config_id, - requested_llm_config_id, - ) - return None - finally: - await session.close() - diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/regenerate.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/regenerate.py deleted file mode 100644 index 02e871a2c..000000000 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/regenerate.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Build ``StreamingContext`` for regenerate streaming.""" - -from __future__ import annotations - -from app.agents.new_chat.filesystem_selection import FilesystemSelection -from app.db import ChatVisibility -from app.tasks.chat.streaming.orchestration.input import StreamingContext -from app.tasks.chat.streaming.orchestration.streaming_context.chat import ( - build_chat_streaming_context, -) - - -async def build_regenerate_streaming_context( - *, - user_query: str, - search_space_id: int, - chat_id: int, - user_id: str | None = None, - llm_config_id: int = -1, - mentioned_document_ids: list[int] | None = None, - mentioned_surfsense_doc_ids: list[int] | None = None, - checkpoint_id: str | None = None, - needs_history_bootstrap: bool = False, - thread_visibility: ChatVisibility | None = None, - current_user_display_name: str | None = None, - disabled_tools: list[str] | None = None, - filesystem_selection: FilesystemSelection | None = None, - request_id: str | None = None, - user_image_data_urls: list[str] | None = None, -) -> StreamingContext | None: - """Build context for ``stream_regenerate`` execution.""" - return await build_chat_streaming_context( - user_query=user_query, - search_space_id=search_space_id, - chat_id=chat_id, - user_id=user_id, - llm_config_id=llm_config_id, - mentioned_document_ids=mentioned_document_ids, - mentioned_surfsense_doc_ids=mentioned_surfsense_doc_ids, - checkpoint_id=checkpoint_id, - needs_history_bootstrap=needs_history_bootstrap, - thread_visibility=thread_visibility, - current_user_display_name=current_user_display_name, - disabled_tools=disabled_tools, - filesystem_selection=filesystem_selection, - request_id=request_id, - user_image_data_urls=user_image_data_urls, - ) - diff --git a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/resume.py b/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/resume.py deleted file mode 100644 index 6d0caea4d..000000000 --- a/surfsense_backend/app/tasks/chat/streaming/orchestration/streaming_context/resume.py +++ /dev/null @@ -1,154 +0,0 @@ -"""Build ``StreamingContext`` for resume streaming.""" - -from __future__ import annotations - -import logging -import time -from typing import Any - -from langgraph.types import Command - -from app.agents.multi_agent_chat import create_multi_agent_chat_deep_agent -from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent -from app.agents.new_chat.checkpointer import get_checkpointer -from app.agents.new_chat.context import SurfSenseContextSchema -from app.agents.new_chat.filesystem_selection import FilesystemSelection -from app.agents.new_chat.llm_config import ( - AgentConfig, - create_chat_litellm_from_agent_config, - create_chat_litellm_from_config, - load_agent_config, - load_global_llm_config_by_id, -) -from app.db import ChatVisibility, SearchSourceConnectorType, async_session_maker -from app.services.auto_model_pin_service import resolve_or_get_pinned_llm_config_id -from app.services.connector_service import ConnectorService -from app.services.new_streaming_service import VercelStreamingService -from app.tasks.chat.streaming.agent_setup import build_main_agent_for_thread -from app.tasks.chat.streaming.orchestration.input import StreamingContext - -logger = logging.getLogger(__name__) - - -async def build_resume_streaming_context( - *, - chat_id: int, - search_space_id: int, - decisions: list[dict], - user_id: str | None = None, - llm_config_id: int = -1, - thread_visibility: ChatVisibility | None = None, - filesystem_selection: FilesystemSelection | None = None, - request_id: str | None = None, - disabled_tools: list[str] | None = None, -) -> StreamingContext | None: - """Build context for ``stream_resume`` execution.""" - session = async_session_maker() - try: - llm_config_id = ( - await resolve_or_get_pinned_llm_config_id( - session, - thread_id=chat_id, - search_space_id=search_space_id, - user_id=user_id, - selected_llm_config_id=llm_config_id, - ) - ).resolved_llm_config_id - - llm: Any - agent_config: AgentConfig | None - if llm_config_id >= 0: - agent_config = await load_agent_config( - session=session, - config_id=llm_config_id, - search_space_id=search_space_id, - ) - if not agent_config: - logger.warning("resume context build failed: missing config %s", llm_config_id) - return None - llm = create_chat_litellm_from_agent_config(agent_config) - else: - loaded_llm_config = load_global_llm_config_by_id(llm_config_id) - if not loaded_llm_config: - logger.warning( - "resume context build failed: missing global config %s", - llm_config_id, - ) - return None - llm = create_chat_litellm_from_config(loaded_llm_config) - agent_config = AgentConfig.from_yaml_config(loaded_llm_config) - - connector_service = ConnectorService(session, search_space_id=search_space_id) - firecrawl_api_key = None - webcrawler_connector = await connector_service.get_connector_by_type( - SearchSourceConnectorType.WEBCRAWLER_CONNECTOR, - search_space_id, - ) - if webcrawler_connector and webcrawler_connector.config: - firecrawl_api_key = webcrawler_connector.config.get("FIRECRAWL_API_KEY") - - checkpointer = await get_checkpointer() - visibility = thread_visibility or ChatVisibility.PRIVATE - - from app.config import config as app_config - - agent_factory = ( - create_multi_agent_chat_deep_agent - if bool(app_config.MULTI_AGENT_CHAT_ENABLED) - else create_surfsense_deep_agent - ) - agent = await build_main_agent_for_thread( - agent_factory, - llm=llm, - search_space_id=search_space_id, - db_session=session, - connector_service=connector_service, - checkpointer=checkpointer, - user_id=user_id, - thread_id=chat_id, - agent_config=agent_config, - firecrawl_api_key=firecrawl_api_key, - thread_visibility=visibility, - filesystem_selection=filesystem_selection, - disabled_tools=disabled_tools, - ) - - turn_id = f"{chat_id}:{int(time.time() * 1000)}" - config = { - "configurable": { - "thread_id": str(chat_id), - "request_id": request_id or "unknown", - "turn_id": turn_id, - "surfsense_resume_value": {"decisions": decisions}, - }, - "recursion_limit": 10_000, - } - - runtime_context = SurfSenseContextSchema( - search_space_id=search_space_id, - request_id=request_id, - turn_id=turn_id, - ) - - await session.commit() - return StreamingContext( - agent=agent, - config=config, - input_data=Command(resume={"decisions": decisions}), - streaming_service=VercelStreamingService(), - step_prefix="thinking-resume", - initial_step_id=None, - initial_step_title="", - initial_step_items=None, - content_builder=None, - runtime_context=runtime_context, - ) - except Exception: - logger.exception( - "Failed to build resume streaming context (llm_config_id=%s)", - llm_config_id, - ) - return None - finally: - await session.close() - diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py b/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py index 351e878a8..18eda9a6d 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/__init__.py @@ -1,7 +1,23 @@ -"""Relay state: thinking steps, tool bookkeeping, and stream helpers.""" +"""Relay: thinking steps, tool bookkeeping, and ``EventRelay``. + +Package imports are lazy so ``relay.thinking_step_sse`` (and siblings) can load +without pulling in ``event_relay`` (which imports handler modules that may +import those siblings). +""" from __future__ import annotations -from app.tasks.chat.streaming.relay.event_relay import EventRelay, EventRelayConfig - __all__ = ["EventRelay", "EventRelayConfig"] + + +def __getattr__(name: str): + if name == "EventRelay": + from app.tasks.chat.streaming.relay.event_relay import EventRelay + + return EventRelay + if name == "EventRelayConfig": + from app.tasks.chat.streaming.relay.event_relay import EventRelayConfig + + return EventRelayConfig + msg = f"module {__name__!r} has no attribute {name!r}" + raise AttributeError(msg) diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py b/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py index c8aebd99c..872998926 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from typing import Any from app.services.streaming.emitter import EmitterRegistry +from app.tasks.chat.streaming.graph_stream.result import StreamingResult from app.tasks.chat.streaming.handlers.chain_end import iter_chain_end_frames from app.tasks.chat.streaming.handlers.chat_model_stream import ( iter_chat_model_stream_frames, @@ -16,7 +17,6 @@ from app.tasks.chat.streaming.handlers.custom_event_dispatch import ( ) from app.tasks.chat.streaming.handlers.tool_end import iter_tool_end_frames from app.tasks.chat.streaming.handlers.tool_start import iter_tool_start_frames -from app.tasks.chat.streaming.orchestration.output import StreamingResult from app.tasks.chat.streaming.relay.state import AgentEventRelayState from app.tasks.chat.streaming.relay.thinking_step_completion import ( complete_active_thinking_step, diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_agent_setup.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_agent_setup.py deleted file mode 100644 index e1f7dd027..000000000 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_agent_setup.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Behavior tests for streaming agent setup helpers.""" - -from __future__ import annotations - -import sys -import types -from typing import Any - -import pytest - -from app.tasks.chat.streaming import agent_setup - -pytestmark = pytest.mark.unit - - -async def test_preflight_llm_calls_litellm_when_model_present( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: dict[str, Any] = {} - - async def _fake_acompletion(**kwargs: Any): - calls.update(kwargs) - return {"ok": True} - - monkeypatch.setitem( - sys.modules, - "litellm", - types.SimpleNamespace(acompletion=_fake_acompletion), - ) - - llm = types.SimpleNamespace(model="openai/test", api_key="k", api_base="b") - await agent_setup.preflight_llm(llm, is_provider_rate_limited=lambda _: False) - - assert calls["model"] == "openai/test" - assert calls["max_tokens"] == 1 - assert calls["timeout"] == 2.5 - assert calls["stream"] is False - - -async def test_preflight_llm_rethrows_rate_limited(monkeypatch: pytest.MonkeyPatch) -> None: - class _RateLimitedError(Exception): - pass - - async def _fake_acompletion(**kwargs: Any): - del kwargs - raise _RateLimitedError("rl") - - monkeypatch.setitem( - sys.modules, - "litellm", - types.SimpleNamespace(acompletion=_fake_acompletion), - ) - - with pytest.raises(_RateLimitedError): - await agent_setup.preflight_llm( - types.SimpleNamespace(model="openai/test"), - is_provider_rate_limited=lambda exc: isinstance(exc, _RateLimitedError), - ) - - -async def test_preflight_llm_skips_probe_for_auto_model( - monkeypatch: pytest.MonkeyPatch, -) -> None: - called = {"count": 0} - - async def _fake_acompletion(**kwargs: Any): - del kwargs - called["count"] += 1 - return {"ok": True} - - monkeypatch.setitem( - sys.modules, - "litellm", - types.SimpleNamespace(acompletion=_fake_acompletion), - ) - - await agent_setup.preflight_llm( - types.SimpleNamespace(model="auto"), - is_provider_rate_limited=lambda _: False, - ) - assert called["count"] == 0 - - -async def test_build_main_agent_for_thread_forwards_arguments() -> None: - seen: dict[str, Any] = {} - - async def _factory(**kwargs: Any): - seen.update(kwargs) - return "agent" - - out = await agent_setup.build_main_agent_for_thread( - _factory, - llm="llm", - search_space_id=1, - db_session="db", - connector_service="connector", - checkpointer="cp", - user_id="u", - thread_id=10, - agent_config="cfg", - firecrawl_api_key="key", - thread_visibility="vis", - filesystem_selection="fs", - disabled_tools=["a"], - mentioned_document_ids=[5], - ) - assert out == "agent" - assert seen["thread_id"] == 10 - assert seen["mentioned_document_ids"] == [5] - - -async def test_settle_speculative_agent_build_swallows_exceptions() -> None: - async def _boom() -> None: - raise RuntimeError("ignore") - - import asyncio - - task = asyncio.create_task(_boom()) - await agent_setup.settle_speculative_agent_build(task) - assert task.done() diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py deleted file mode 100644 index 46c61b498..000000000 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py +++ /dev/null @@ -1,240 +0,0 @@ -"""Behavior tests for orchestrator ``stream_chat`` public API.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - -import pytest - -from app.tasks.chat.streaming.orchestration import StreamingContext -from app.tasks.chat.streaming.orchestration import orchestrator -from app.tasks.chat.streaming.orchestration.orchestrator import ( - stream_chat, - stream_regenerate, - stream_resume, -) - -pytestmark = pytest.mark.unit - - -@dataclass -class _Chunk: - content: Any = "" - additional_kwargs: dict[str, Any] = field(default_factory=dict) - tool_call_chunks: list[dict[str, Any]] = field(default_factory=list) - - -class _StreamingService: - def __init__(self) -> None: - self._text_idx = 0 - - def generate_text_id(self) -> str: - self._text_idx += 1 - return f"text-{self._text_idx}" - - def format_text_start(self, text_id: str) -> str: - return f"text_start:{text_id}" - - def format_text_delta(self, text_id: str, text: str) -> str: - return f"text_delta:{text_id}:{text}" - - def format_text_end(self, text_id: str) -> str: - return f"text_end:{text_id}" - - -class _Agent: - def __init__(self, events: list[dict[str, Any]]) -> None: - self.events = list(events) - self.calls: list[tuple[Any, dict[str, Any]]] = [] - - async def astream_events(self, input_data: Any, **kwargs: Any): - self.calls.append((input_data, kwargs)) - for event in self.events: - yield event - - -async def _collect(stream: Any) -> list[str]: - out: list[str] = [] - async for x in stream: - out.append(x) - return out - - -async def test_stream_chat_uses_streaming_context_path() -> None: - service = _StreamingService() - agent = _Agent( - [ - {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(content="hello")}}, - {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(content="!")}}, - ] - ) - frames = await _collect( - stream_chat( - user_query="ignored-here", - search_space_id=1, - chat_id=77, - streaming_context=StreamingContext( - agent=agent, - config={"configurable": {"thread_id": "thread-1"}}, - input_data={"messages": []}, - streaming_service=service, - ), - ) - ) - - assert frames == [ - "text_start:text-1", - "text_delta:text-1:hello", - "text_delta:text-1:!", - "text_end:text-1", - ] - - -async def test_stream_resume_uses_streaming_context_path() -> None: - service = _StreamingService() - agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("r")}}]) - - frames = await _collect( - stream_resume( - chat_id=9, - search_space_id=1, - decisions=[], - streaming_context=StreamingContext( - agent=agent, - config={"configurable": {"thread_id": "thread-r"}}, - input_data={"messages": []}, - streaming_service=service, - ), - ) - ) - - assert frames == [ - "text_start:text-1", - "text_delta:text-1:r", - "text_end:text-1", - ] - - -async def test_stream_regenerate_uses_streaming_context_path() -> None: - service = _StreamingService() - agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("g")}}]) - - frames = await _collect( - stream_regenerate( - user_query="q", - search_space_id=1, - chat_id=2, - streaming_context=StreamingContext( - agent=agent, - config={"configurable": {"thread_id": "thread-g"}}, - input_data={"messages": []}, - streaming_service=service, - ), - ) - ) - - assert frames == [ - "text_start:text-1", - "text_delta:text-1:g", - "text_end:text-1", - ] - - -async def test_stream_chat_builds_streaming_context_when_not_provided() -> None: - service = _StreamingService() - agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("b")}}]) - - async def _fake_builder(**kwargs: Any) -> StreamingContext: - del kwargs - return StreamingContext( - agent=agent, - config={"configurable": {"thread_id": "thread-b"}}, - input_data={"messages": []}, - streaming_service=service, - ) - - old = orchestrator.build_chat_streaming_context - orchestrator.build_chat_streaming_context = _fake_builder - try: - frames = await _collect( - stream_chat( - user_query="q", - search_space_id=1, - chat_id=3, - ) - ) - finally: - orchestrator.build_chat_streaming_context = old - - assert frames == [ - "text_start:text-1", - "text_delta:text-1:b", - "text_end:text-1", - ] - - -async def test_stream_resume_builds_streaming_context_when_not_provided() -> None: - service = _StreamingService() - agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("u")}}]) - - async def _fake_builder(**kwargs: Any) -> StreamingContext: - del kwargs - return StreamingContext( - agent=agent, - config={"configurable": {"thread_id": "thread-u"}}, - input_data={"messages": []}, - streaming_service=service, - ) - - old = orchestrator.build_resume_streaming_context - orchestrator.build_resume_streaming_context = _fake_builder - try: - frames = await _collect( - stream_resume( - chat_id=9, - search_space_id=1, - decisions=[], - ) - ) - finally: - orchestrator.build_resume_streaming_context = old - - assert frames == [ - "text_start:text-1", - "text_delta:text-1:u", - "text_end:text-1", - ] - - -async def test_stream_regenerate_builds_streaming_context_when_not_provided() -> None: - service = _StreamingService() - agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("x")}}]) - - async def _fake_builder(**kwargs: Any) -> StreamingContext: - del kwargs - return StreamingContext( - agent=agent, - config={"configurable": {"thread_id": "thread-x"}}, - input_data={"messages": []}, - streaming_service=service, - ) - - old = orchestrator.build_regenerate_streaming_context - orchestrator.build_regenerate_streaming_context = _fake_builder - try: - frames = await _collect( - stream_regenerate( - user_query="q", - search_space_id=1, - chat_id=2, - ) - ) - finally: - orchestrator.build_regenerate_streaming_context = old - - assert frames == [ - "text_start:text-1", - "text_delta:text-1:x", - "text_end:text-1", - ] diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py index 9207f37d1..023c8b999 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_1_parity.py @@ -1,7 +1,7 @@ """Pin Stage 1 extractions as faithful copies of the old helpers. -The new orchestrator under ``app.tasks.chat.streaming`` is built in -parallel with the production module ``app.tasks.chat.stream_new_chat``. +Extractions under ``app.tasks.chat.streaming`` are compared to +``app.tasks.chat.stream_new_chat`` helpers. For each Stage 1 extraction we assert the new function returns the same output as the old one for a representative input set. The moment the two diverge - intentionally or otherwise - this file fails loudly so diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stream_output.py similarity index 93% rename from surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py rename to surfsense_backend/tests/unit/tasks/chat/streaming/test_stream_output.py index b17d82293..9fb876dd7 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestration_event_stream.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stream_output.py @@ -1,4 +1,4 @@ -"""Behavior tests for orchestration event-stream execution.""" +"""Tests for ``stream_output`` (LangGraph events → SSE).""" from __future__ import annotations @@ -7,8 +7,8 @@ from typing import Any import pytest -from app.tasks.chat.streaming.orchestration import stream_output -from app.tasks.chat.streaming.orchestration.output import StreamingResult +from app.tasks.chat.streaming.graph_stream import stream_output +from app.tasks.chat.streaming.graph_stream.result import StreamingResult pytestmark = pytest.mark.unit @@ -88,6 +88,7 @@ async def test_stream_output_emits_text_lifecycle_and_updates_result() -> None: async def test_stream_output_passes_runtime_context_to_agent() -> None: service = _StreamingService() + class _ContextAwareAgent: async def astream_events(self, input_data: Any, **kwargs: Any): del input_data From 78f4747382cead46c0c72040002562fe56bc35e4 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 7 May 2026 19:40:10 +0200 Subject: [PATCH 60/96] refactor(chat): stream agent events via stream_output and remove parity v2 flag --- docker/.env.example | 1 - surfsense_backend/.env.example | 8 - .../app/agents/new_chat/feature_flags.py | 15 - .../app/services/new_streaming_service.py | 17 +- .../app/tasks/chat/content_builder.py | 8 +- .../app/tasks/chat/stream_new_chat.py | 1524 +---------------- .../streaming/graph_stream/event_stream.py | 2 - .../streaming/handlers/chat_model_stream.py | 4 +- .../chat/streaming/handlers/tool_start.py | 24 +- .../app/tasks/chat/streaming/relay/state.py | 3 - .../agents/new_chat/test_feature_flags.py | 3 - .../chat/streaming/test_stage_2_parity.py | 4 +- .../unit/tasks/chat/test_content_builder.py | 4 +- .../tasks/chat/test_tool_input_streaming.py | 112 +- .../assistant-ui/reasoning-message-part.tsx | 4 +- .../components/assistant-ui/tool-fallback.tsx | 14 +- surfsense_web/lib/chat/streaming-state.ts | 5 +- 17 files changed, 76 insertions(+), 1676 deletions(-) diff --git a/docker/.env.example b/docker/.env.example index fd56bdccc..aba15f13f 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -324,7 +324,6 @@ SURFSENSE_ENABLE_ACTION_LOG=true SURFSENSE_ENABLE_REVERT_ROUTE=true SURFSENSE_ENABLE_PERMISSION=true SURFSENSE_ENABLE_DOOM_LOOP=true -SURFSENSE_ENABLE_STREAM_PARITY_V2=true # Periodic connector sync interval (default: 5m) # SCHEDULE_CHECKER_INTERVAL=5m diff --git a/surfsense_backend/.env.example b/surfsense_backend/.env.example index ba89059c8..3d442973c 100644 --- a/surfsense_backend/.env.example +++ b/surfsense_backend/.env.example @@ -315,14 +315,6 @@ LANGSMITH_PROJECT=surfsense # SURFSENSE_ENABLE_ACTION_LOG=false # SURFSENSE_ENABLE_REVERT_ROUTE=false # Backend-only; flip when UI ships -# Streaming parity v2 — opt in to LangChain's structured AIMessageChunk -# content (typed reasoning blocks, tool-input deltas) and propagate the -# real tool_call_id to the SSE layer. When OFF, the stream falls back to -# the str-only text path and synthetic "call_" tool-call ids. -# Schema migrations 135/136 ship unconditionally because they are -# forward-compatible. -# SURFSENSE_ENABLE_STREAM_PARITY_V2=false - # Plugins # SURFSENSE_ENABLE_PLUGIN_LOADER=false # Comma-separated allowlist of plugin entry-point names diff --git a/surfsense_backend/app/agents/new_chat/feature_flags.py b/surfsense_backend/app/agents/new_chat/feature_flags.py index b3dc0fa82..3cea051ef 100644 --- a/surfsense_backend/app/agents/new_chat/feature_flags.py +++ b/surfsense_backend/app/agents/new_chat/feature_flags.py @@ -28,7 +28,6 @@ Defaults: SURFSENSE_ENABLE_PERMISSION=true SURFSENSE_ENABLE_DOOM_LOOP=true SURFSENSE_ENABLE_LLM_TOOL_SELECTOR=false # adds a per-turn LLM call - SURFSENSE_ENABLE_STREAM_PARITY_V2=true Master kill-switch (overrides everything else): @@ -88,15 +87,6 @@ class AgentFeatureFlags: enable_action_log: bool = True enable_revert_route: bool = True - # Streaming parity v2 — opt in to LangChain's structured - # ``AIMessageChunk`` content (typed reasoning blocks, tool-input - # deltas) and propagate the real ``tool_call_id`` to the SSE layer. - # When OFF the ``stream_new_chat`` task falls back to the str-only - # text path and the synthetic ``call_`` tool-call id (no - # ``langchainToolCallId`` propagation). Schema migrations 135/136 - # ship unconditionally because they're forward-compatible. - enable_stream_parity_v2: bool = True - # Plugins enable_plugin_loader: bool = False @@ -169,7 +159,6 @@ class AgentFeatureFlags: enable_kb_planner_runnable=False, enable_action_log=False, enable_revert_route=False, - enable_stream_parity_v2=False, enable_plugin_loader=False, enable_otel=False, enable_agent_cache=False, @@ -208,10 +197,6 @@ class AgentFeatureFlags: # Snapshot / revert enable_action_log=_env_bool("SURFSENSE_ENABLE_ACTION_LOG", True), enable_revert_route=_env_bool("SURFSENSE_ENABLE_REVERT_ROUTE", True), - # Streaming parity v2 - enable_stream_parity_v2=_env_bool( - "SURFSENSE_ENABLE_STREAM_PARITY_V2", True - ), # Plugins enable_plugin_loader=_env_bool("SURFSENSE_ENABLE_PLUGIN_LOADER", False), # Observability diff --git a/surfsense_backend/app/services/new_streaming_service.py b/surfsense_backend/app/services/new_streaming_service.py index 55129668c..cec0c8a5e 100644 --- a/surfsense_backend/app/services/new_streaming_service.py +++ b/surfsense_backend/app/services/new_streaming_service.py @@ -608,15 +608,14 @@ class VercelStreamingService: Args: tool_call_id: The unique tool call identifier. May be EITHER the synthetic ``call_`` id derived from LangGraph - ``run_id`` (legacy / ``SURFSENSE_ENABLE_STREAM_PARITY_V2`` - OFF, or the unmatched-fallback path under parity_v2) OR - the authoritative LangChain ``tool_call.id`` (parity_v2 - path: when the provider streams ``tool_call_chunks`` we - register the ``index`` and reuse the lc-id as the card - id so live ``tool-input-delta`` events can be routed - without a downstream join). Either way, the same id is - preserved across ``tool-input-start`` / ``-delta`` / - ``-available`` / ``tool-output-available`` for one call. + ``run_id`` (unmatched chunk fallback when no ``index`` was + registered) OR the authoritative LangChain ``tool_call.id`` + (when the provider streams ``tool_call_chunks`` we register + the ``index`` and reuse the lc-id as the card id so live + ``tool-input-delta`` events route without a downstream join). + Either way, the same id is preserved across + ``tool-input-start`` / ``-delta`` / ``-available`` / + ``tool-output-available`` for one call. tool_name: The name of the tool being called. langchain_tool_call_id: Optional authoritative LangChain ``tool_call.id``. When set, surfaces as diff --git a/surfsense_backend/app/tasks/chat/content_builder.py b/surfsense_backend/app/tasks/chat/content_builder.py index 041cab286..32b49e6b5 100644 --- a/surfsense_backend/app/tasks/chat/content_builder.py +++ b/surfsense_backend/app/tasks/chat/content_builder.py @@ -85,8 +85,8 @@ class AssistantContentBuilder: self._current_text_idx: int = -1 self._current_reasoning_idx: int = -1 # ``ui_id``-keyed indexes for tool-call parts. ``ui_id`` is the - # synthetic ``call_`` (legacy) or the LangChain - # ``tool_call.id`` (parity_v2) — same key the streaming layer + # synthetic ``call_`` (chunk fallback) or the LangChain + # ``tool_call.id`` (indexed chunk path) — same key the streaming layer # threads through every ``tool-input-*`` / ``tool-output-*`` event. self._tool_call_idx_by_ui_id: dict[str, int] = {} # Live argsText accumulator (concatenated ``tool-input-delta`` chunks) @@ -181,7 +181,7 @@ class AssistantContentBuilder: """Register a tool-call card. Args are filled in by later events.""" if not ui_id: return - # Skip duplicate registration: parity_v2 may emit + # Skip duplicate registration: the stream may emit # ``tool-input-start`` from both ``on_chat_model_stream`` # (when tool_call_chunks register a name) and ``on_tool_start`` # (the canonical path). The FE de-dupes via ``toolCallIndices``; @@ -243,7 +243,7 @@ class AssistantContentBuilder: pretty-printed JSON, sets the full ``args`` dict, and backfills ``langchainToolCallId`` if it wasn't known at ``tool-input-start`` time. Also creates the card if no prior ``tool-input-start`` registered it - (legacy parity_v2-OFF / late-registration paths). + (late-registration when no prior ``tool-input-start``). """ if not ui_id: return diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 1a2f38077..8e135179a 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -9,13 +9,11 @@ Supports loading LLM configurations from: - NewLLMConfig database table (positive IDs for user-created configs with prompt settings) """ -import ast import asyncio import contextlib import gc import json import logging -import re import time from collections.abc import AsyncGenerator from dataclasses import dataclass, field @@ -33,7 +31,6 @@ from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent from app.agents.new_chat.checkpointer import get_checkpointer from app.agents.new_chat.context import SurfSenseContextSchema from app.agents.new_chat.errors import BusyError -from app.agents.new_chat.feature_flags import get_flags from app.agents.new_chat.filesystem_selection import FilesystemMode, FilesystemSelection from app.agents.new_chat.llm_config import ( AgentConfig, @@ -77,6 +74,7 @@ from app.services.chat_session_state_service import ( ) from app.services.connector_service import ConnectorService from app.services.new_streaming_service import VercelStreamingService +from app.tasks.chat.streaming.graph_stream.event_stream import stream_output from app.utils.content_utils import bootstrap_history_from_db from app.utils.perf import get_perf_logger, log_system_snapshot, trim_native_heap from app.utils.user_message_multimodal import build_human_message_content @@ -729,9 +727,9 @@ def _legacy_match_lc_id( ) -> str | None: """Best-effort match a buffered ``tool_call_chunk`` to a tool name. - Pure extract of the legacy in-line match used at ``on_tool_start`` for - parity_v2-OFF and unmatched (chunk path didn't register an index for - this call) tools. Pops the next id-bearing chunk whose ``name`` + Pure extract of the in-line match used at ``on_tool_start`` when the + chunk path didn't register an index for this call. Pops the next + id-bearing chunk whose ``name`` matches ``tool_name`` (or any id-bearing chunk as a fallback) and returns its id. Mutates ``pending_tool_call_chunks`` and ``lc_tool_call_id_by_run`` in place. @@ -803,1505 +801,22 @@ async def _stream_agent_events( Yields: SSE-formatted strings for each event. """ - accumulated_text = "" - current_text_id: str | None = None - thinking_step_counter = 1 if initial_step_id else 0 - tool_step_ids: dict[str, str] = {} - completed_step_ids: set[str] = set() - last_active_step_id: str | None = initial_step_id - last_active_step_title: str = initial_step_title - last_active_step_items: list[str] = initial_step_items or [] - just_finished_tool: bool = False - active_tool_depth: int = 0 # Track nesting: >0 means we're inside a tool - called_update_memory: bool = False + async for sse in stream_output( + agent=agent, + config=config, + input_data=input_data, + streaming_service=streaming_service, + result=result, + step_prefix=step_prefix, + initial_step_id=initial_step_id, + initial_step_title=initial_step_title, + initial_step_items=initial_step_items, + content_builder=content_builder, + runtime_context=runtime_context, + ): + yield sse - # Reasoning-block streaming. We open a reasoning block on the - # first reasoning delta of a step, append deltas as they arrive, and - # close it when text starts (the model has switched to writing its - # answer) or ``on_chat_model_end`` fires for the model node. Reuses - # the same Vercel format-helpers as text-start/delta/end. - current_reasoning_id: str | None = None - - # Streaming-parity v2 feature flag. When OFF we keep the legacy - # shape: str-only content, no reasoning blocks, no - # ``langchainToolCallId`` propagation. The schema migrations - # (135 / 136) ship unconditionally because they're forward-compatible. - parity_v2 = bool(get_flags().enable_stream_parity_v2) - - # Best-effort attach of LangChain ``tool_call_id`` to the synthetic - # ``call_`` card id we already emit. We accumulate - # ``tool_call_chunks`` from ``on_chat_model_stream``, key them by - # name, and pop the next unconsumed entry at ``on_tool_start``. The - # authoritative id is later filled in at ``on_tool_end`` from - # ``ToolMessage.tool_call_id``. Under parity_v2 we ALSO short-circuit - # this list for chunks that already registered into ``index_to_meta`` - # below — so this list is reserved for the parity_v2-OFF / unmatched - # fallback path only and never re-pops a chunk we already streamed. - pending_tool_call_chunks: list[dict[str, Any]] = [] - lc_tool_call_id_by_run: dict[str, str] = {} - file_path_by_run: dict[str, str] = {} - - # parity_v2 only: live tool-call argument streaming. ``index_to_meta`` - # is keyed by the chunk's ``index`` field — LangChain - # ``ToolCallChunk``s for the same call share an index but only the - # first chunk carries id+name (subsequent ones are id=None, - # name=None, args=""). We register an index when both id and - # name are observed on a chunk (per ToolCallChunk semantics they - # arrive together on the first chunk), then route every later chunk - # at that index to the same ``ui_id`` as a ``tool-input-delta``. - # ``ui_tool_call_id_by_run`` maps LangGraph ``run_id`` to the - # ``ui_id`` used for that call's ``tool-input-start`` so the matching - # ``tool-output-available`` (emitted from ``on_tool_end``) lands on - # the same card. - index_to_meta: dict[int, dict[str, str]] = {} - ui_tool_call_id_by_run: dict[str, str] = {} - - # Per-tool-end mutable cache for the LangChain tool_call_id resolved - # at ``on_tool_end``. ``_emit_tool_output`` reads this so every - # ``format_tool_output_available`` call automatically carries the - # authoritative id without duplicating the kwarg at every call site. - current_lc_tool_call_id: dict[str, str | None] = {"value": None} - - def _emit_tool_output(call_id: str, output: Any) -> str: - # Drive the builder before formatting the SSE so the in-memory - # ContentPart[] mirror sees the result attached to the same - # card the FE will render. Builder method is a no-op when - # ``content_builder`` is None (anonymous / legacy paths). - if content_builder is not None: - content_builder.on_tool_output_available( - call_id, output, current_lc_tool_call_id["value"] - ) - return streaming_service.format_tool_output_available( - call_id, - output, - langchain_tool_call_id=current_lc_tool_call_id["value"], - ) - - def _emit_thinking_step( - *, - step_id: str, - title: str, - status: str = "in_progress", - items: list[str] | None = None, - ) -> str: - """Format a thinking-step SSE event and notify the builder. - - Single helper used at every ``format_thinking_step`` yield site - in this generator. Drives ``AssistantContentBuilder.on_thinking_step`` - first so the FE-mirror state lands the update before the SSE - carrying the same data leaves the wire — order matches the FE - pipeline (``processSharedStreamEvent`` updates state, then - flushes). Builder call is a no-op when ``content_builder`` is - None (anonymous / legacy paths). - """ - if content_builder is not None: - content_builder.on_thinking_step(step_id, title, status, items) - return streaming_service.format_thinking_step( - step_id=step_id, - title=title, - status=status, - items=items, - ) - - def next_thinking_step_id() -> str: - nonlocal thinking_step_counter - thinking_step_counter += 1 - return f"{step_prefix}-{thinking_step_counter}" - - def complete_current_step() -> str | None: - nonlocal last_active_step_id - if last_active_step_id and last_active_step_id not in completed_step_ids: - completed_step_ids.add(last_active_step_id) - event = _emit_thinking_step( - step_id=last_active_step_id, - title=last_active_step_title, - status="completed", - items=last_active_step_items if last_active_step_items else None, - ) - last_active_step_id = None - return event - return None - - # Per-invocation runtime context (Phase 1.5). When supplied, - # ``KnowledgePriorityMiddleware`` reads ``mentioned_document_ids`` - # from ``runtime.context`` instead of its constructor closure — the - # prerequisite that lets the compiled-agent cache (Phase 1) reuse a - # single graph across turns. Astream_events_kwargs stays empty when - # callers leave ``runtime_context`` as ``None`` to preserve the - # legacy code path bit-for-bit. - astream_kwargs: dict[str, Any] = {"config": config, "version": "v2"} - if runtime_context is not None: - astream_kwargs["context"] = runtime_context - - async for event in agent.astream_events(input_data, **astream_kwargs): - event_type = event.get("event", "") - - if event_type == "on_chat_model_stream": - if active_tool_depth > 0: - continue # Suppress inner-tool LLM tokens from leaking into chat - if "surfsense:internal" in event.get("tags", []): - continue # Suppress middleware-internal LLM tokens (e.g. KB search classification) - chunk = event.get("data", {}).get("chunk") - if not chunk: - continue - parts = _extract_chunk_parts(chunk) - - reasoning_delta = parts["reasoning"] - text_delta = parts["text"] - - # Reasoning streaming. Open a reasoning block on first - # delta; append every subsequent delta until text begins. - # When text starts we close the reasoning block first so the - # frontend sees the natural hand-off. Gated behind the - # parity-v2 flag so legacy deployments keep today's shape. - if parity_v2 and reasoning_delta: - if current_text_id is not None: - yield streaming_service.format_text_end(current_text_id) - if content_builder is not None: - content_builder.on_text_end(current_text_id) - current_text_id = None - if current_reasoning_id is None: - completion_event = complete_current_step() - if completion_event: - yield completion_event - if just_finished_tool: - last_active_step_id = None - last_active_step_title = "" - last_active_step_items = [] - just_finished_tool = False - current_reasoning_id = streaming_service.generate_reasoning_id() - yield streaming_service.format_reasoning_start(current_reasoning_id) - if content_builder is not None: - content_builder.on_reasoning_start(current_reasoning_id) - yield streaming_service.format_reasoning_delta( - current_reasoning_id, reasoning_delta - ) - if content_builder is not None: - content_builder.on_reasoning_delta( - current_reasoning_id, reasoning_delta - ) - - if text_delta: - if current_reasoning_id is not None: - yield streaming_service.format_reasoning_end(current_reasoning_id) - if content_builder is not None: - content_builder.on_reasoning_end(current_reasoning_id) - current_reasoning_id = None - if current_text_id is None: - completion_event = complete_current_step() - if completion_event: - yield completion_event - if just_finished_tool: - last_active_step_id = None - last_active_step_title = "" - last_active_step_items = [] - just_finished_tool = False - current_text_id = streaming_service.generate_text_id() - yield streaming_service.format_text_start(current_text_id) - if content_builder is not None: - content_builder.on_text_start(current_text_id) - yield streaming_service.format_text_delta(current_text_id, text_delta) - accumulated_text += text_delta - if content_builder is not None: - content_builder.on_text_delta(current_text_id, text_delta) - - # Live tool-call argument streaming. Runs AFTER text/reasoning - # processing so chunks containing both stay in their natural - # wire order (text → text-end → tool-input-start). Active - # text/reasoning are closed inside the registration branch - # before ``tool-input-start`` so the frontend sees a clean - # part boundary even when providers interleave. - if parity_v2 and parts["tool_call_chunks"]: - for tcc in parts["tool_call_chunks"]: - idx = tcc.get("index") - - # Register this index when we first see id+name - # TOGETHER. Per LangChain ToolCallChunk semantics the - # first chunk for a tool call carries both fields - # together; later chunks have id=None, name=None and - # only ``args``. Requiring BOTH keeps wire - # ``tool-input-start`` always carrying a real - # toolName (assistant-ui's typed tool-part dispatch - # keys off it). - if idx is not None and idx not in index_to_meta: - lc_id = tcc.get("id") - name = tcc.get("name") - if lc_id and name: - ui_id = lc_id - - # Close active text/reasoning so wire - # ordering stays clean even on providers - # that interleave text and tool-call chunks - # within the same stream window. - if current_text_id is not None: - yield streaming_service.format_text_end(current_text_id) - if content_builder is not None: - content_builder.on_text_end(current_text_id) - current_text_id = None - if current_reasoning_id is not None: - yield streaming_service.format_reasoning_end( - current_reasoning_id - ) - if content_builder is not None: - content_builder.on_reasoning_end( - current_reasoning_id - ) - current_reasoning_id = None - - index_to_meta[idx] = { - "ui_id": ui_id, - "lc_id": lc_id, - "name": name, - } - yield streaming_service.format_tool_input_start( - ui_id, - name, - langchain_tool_call_id=lc_id, - ) - if content_builder is not None: - content_builder.on_tool_input_start(ui_id, name, lc_id) - - # Emit args delta for any chunk at a registered - # index (including idless continuations). Once an - # index is owned by ``index_to_meta`` we DO NOT - # append to ``pending_tool_call_chunks`` — that list - # is reserved for the parity_v2-OFF / unmatched - # fallback path so it never re-pops chunks already - # consumed here (skip-append). - meta = index_to_meta.get(idx) if idx is not None else None - if meta: - args_chunk = tcc.get("args") or "" - if args_chunk: - yield streaming_service.format_tool_input_delta( - meta["ui_id"], args_chunk - ) - if content_builder is not None: - content_builder.on_tool_input_delta( - meta["ui_id"], args_chunk - ) - else: - pending_tool_call_chunks.append(tcc) - - elif event_type == "on_tool_start": - active_tool_depth += 1 - tool_name = event.get("name", "unknown_tool") - run_id = event.get("run_id", "") - tool_input = event.get("data", {}).get("input", {}) - if tool_name in ("write_file", "edit_file"): - result.write_attempted = True - if isinstance(tool_input, dict): - file_path = tool_input.get("file_path") - if isinstance(file_path, str) and file_path.strip() and run_id: - file_path_by_run[run_id] = file_path.strip() - - if current_text_id is not None: - yield streaming_service.format_text_end(current_text_id) - if content_builder is not None: - content_builder.on_text_end(current_text_id) - current_text_id = None - - if last_active_step_title != "Synthesizing response": - completion_event = complete_current_step() - if completion_event: - yield completion_event - - just_finished_tool = False - tool_step_id = next_thinking_step_id() - tool_step_ids[run_id] = tool_step_id - last_active_step_id = tool_step_id - - if tool_name == "ls": - ls_path = ( - tool_input.get("path", "/") - if isinstance(tool_input, dict) - else str(tool_input) - ) - last_active_step_title = "Listing files" - last_active_step_items = [ls_path] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Listing files", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "read_file": - fp = ( - tool_input.get("file_path", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - display_fp = fp if len(fp) <= 80 else "…" + fp[-77:] - last_active_step_title = "Reading file" - last_active_step_items = [display_fp] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Reading file", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "write_file": - fp = ( - tool_input.get("file_path", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - display_fp = fp if len(fp) <= 80 else "…" + fp[-77:] - last_active_step_title = "Writing file" - last_active_step_items = [display_fp] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Writing file", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "edit_file": - fp = ( - tool_input.get("file_path", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - display_fp = fp if len(fp) <= 80 else "…" + fp[-77:] - last_active_step_title = "Editing file" - last_active_step_items = [display_fp] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Editing file", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "glob": - pat = ( - tool_input.get("pattern", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - base_path = ( - tool_input.get("path", "/") if isinstance(tool_input, dict) else "/" - ) - last_active_step_title = "Searching files" - last_active_step_items = [f"{pat} in {base_path}"] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Searching files", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "grep": - pat = ( - tool_input.get("pattern", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - grep_path = ( - tool_input.get("path", "") if isinstance(tool_input, dict) else "" - ) - display_pat = pat[:60] + ("…" if len(pat) > 60 else "") - last_active_step_title = "Searching content" - last_active_step_items = [ - f'"{display_pat}"' + (f" in {grep_path}" if grep_path else "") - ] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Searching content", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "rm": - rm_path = ( - tool_input.get("path", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - display_path = rm_path if len(rm_path) <= 80 else "…" + rm_path[-77:] - last_active_step_title = "Deleting file" - last_active_step_items = [display_path] if display_path else [] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Deleting file", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "rmdir": - rmdir_path = ( - tool_input.get("path", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - display_path = ( - rmdir_path if len(rmdir_path) <= 80 else "…" + rmdir_path[-77:] - ) - last_active_step_title = "Deleting folder" - last_active_step_items = [display_path] if display_path else [] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Deleting folder", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "mkdir": - mkdir_path = ( - tool_input.get("path", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - display_path = ( - mkdir_path if len(mkdir_path) <= 80 else "…" + mkdir_path[-77:] - ) - last_active_step_title = "Creating folder" - last_active_step_items = [display_path] if display_path else [] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Creating folder", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "move_file": - src = ( - tool_input.get("source_path", "") - if isinstance(tool_input, dict) - else "" - ) - dst = ( - tool_input.get("destination_path", "") - if isinstance(tool_input, dict) - else "" - ) - display_src = src if len(src) <= 60 else "…" + src[-57:] - display_dst = dst if len(dst) <= 60 else "…" + dst[-57:] - last_active_step_title = "Moving file" - last_active_step_items = ( - [f"{display_src} → {display_dst}"] if src or dst else [] - ) - yield _emit_thinking_step( - step_id=tool_step_id, - title="Moving file", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "write_todos": - todos = ( - tool_input.get("todos", []) if isinstance(tool_input, dict) else [] - ) - todo_count = len(todos) if isinstance(todos, list) else 0 - last_active_step_title = "Planning tasks" - last_active_step_items = ( - [f"{todo_count} task{'s' if todo_count != 1 else ''}"] - if todo_count - else [] - ) - yield _emit_thinking_step( - step_id=tool_step_id, - title="Planning tasks", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "save_document": - doc_title = ( - tool_input.get("title", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - display_title = doc_title[:60] + ("…" if len(doc_title) > 60 else "") - last_active_step_title = "Saving document" - last_active_step_items = [display_title] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Saving document", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "generate_image": - prompt = ( - tool_input.get("prompt", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - last_active_step_title = "Generating image" - last_active_step_items = [ - f"Prompt: {prompt[:80]}{'...' if len(prompt) > 80 else ''}" - ] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Generating image", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "scrape_webpage": - url = ( - tool_input.get("url", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - last_active_step_title = "Scraping webpage" - last_active_step_items = [ - f"URL: {url[:80]}{'...' if len(url) > 80 else ''}" - ] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Scraping webpage", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "generate_podcast": - podcast_title = ( - tool_input.get("podcast_title", "SurfSense Podcast") - if isinstance(tool_input, dict) - else "SurfSense Podcast" - ) - content_len = len( - tool_input.get("source_content", "") - if isinstance(tool_input, dict) - else "" - ) - last_active_step_title = "Generating podcast" - last_active_step_items = [ - f"Title: {podcast_title}", - f"Content: {content_len:,} characters", - "Preparing audio generation...", - ] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Generating podcast", - status="in_progress", - items=last_active_step_items, - ) - elif tool_name == "generate_report": - report_topic = ( - tool_input.get("topic", "Report") - if isinstance(tool_input, dict) - else "Report" - ) - is_revision = bool( - isinstance(tool_input, dict) and tool_input.get("parent_report_id") - ) - step_title = "Revising report" if is_revision else "Generating report" - last_active_step_title = step_title - last_active_step_items = [ - f"Topic: {report_topic}", - "Analyzing source content...", - ] - yield _emit_thinking_step( - step_id=tool_step_id, - title=step_title, - status="in_progress", - items=last_active_step_items, - ) - elif tool_name in ("execute", "execute_code"): - cmd = ( - tool_input.get("command", "") - if isinstance(tool_input, dict) - else str(tool_input) - ) - display_cmd = cmd[:80] + ("…" if len(cmd) > 80 else "") - last_active_step_title = "Running command" - last_active_step_items = [f"$ {display_cmd}"] - yield _emit_thinking_step( - step_id=tool_step_id, - title="Running command", - status="in_progress", - items=last_active_step_items, - ) - else: - # Fallback for tools without a curated thinking-step title - # (typically connector tools, MCP-registered tools, or - # newly added tools that haven't been wired up here yet). - # Render the snake_cased name as a sentence-cased phrase - # so non-technical users see e.g. "Send gmail email" - # rather than the raw identifier "send_gmail_email". - last_active_step_title = ( - tool_name.replace("_", " ").strip().capitalize() or tool_name - ) - last_active_step_items = [] - yield _emit_thinking_step( - step_id=tool_step_id, - title=last_active_step_title, - status="in_progress", - ) - - # Resolve the card identity. If the chunk-emission loop - # already registered an ``index`` for this tool call (parity_v2 - # path), reuse the same ui_id so the card sees: - # tool-input-start → deltas… → tool-input-available → - # tool-output-available all keyed by lc_id. Otherwise fall - # back to the synthetic ``call_`` id and the legacy - # best-effort match against ``pending_tool_call_chunks``. - matched_meta: dict[str, str] | None = None - if parity_v2: - # FIFO over indices 0,1,2…; first unassigned same-name - # match wins. Handles parallel same-name calls (e.g. two - # write_file calls) deterministically as long as the - # model interleaves on_tool_start in the same order it - # streamed the args. - taken_ui_ids = set(ui_tool_call_id_by_run.values()) - for meta in index_to_meta.values(): - if meta["name"] == tool_name and meta["ui_id"] not in taken_ui_ids: - matched_meta = meta - break - - tool_call_id: str - langchain_tool_call_id: str | None = None - if matched_meta is not None: - tool_call_id = matched_meta["ui_id"] - langchain_tool_call_id = matched_meta["lc_id"] - # ``tool-input-start`` already fired during chunk - # emission — skip the duplicate. No pruning is needed - # because the chunk-emission loop intentionally never - # appends registered-index chunks to - # ``pending_tool_call_chunks`` (skip-append). - if run_id: - lc_tool_call_id_by_run[run_id] = matched_meta["lc_id"] - else: - tool_call_id = ( - f"call_{run_id[:32]}" - if run_id - else streaming_service.generate_tool_call_id() - ) - # Legacy fallback: parity_v2 OFF, or parity_v2 ON but the - # provider didn't stream tool_call_chunks for this call - # (no index registered). Run the existing best-effort - # match BEFORE emitting start so we still attach an - # authoritative ``langchainToolCallId`` when possible. - if parity_v2: - langchain_tool_call_id = _legacy_match_lc_id( - pending_tool_call_chunks, - tool_name, - run_id, - lc_tool_call_id_by_run, - ) - yield streaming_service.format_tool_input_start( - tool_call_id, - tool_name, - langchain_tool_call_id=langchain_tool_call_id, - ) - if content_builder is not None: - content_builder.on_tool_input_start( - tool_call_id, tool_name, langchain_tool_call_id - ) - - if run_id: - ui_tool_call_id_by_run[run_id] = tool_call_id - - # Sanitize tool_input: strip runtime-injected non-serializable - # values (e.g. LangChain ToolRuntime) before sending over SSE. - if isinstance(tool_input, dict): - _safe_input: dict[str, Any] = {} - for _k, _v in tool_input.items(): - try: - json.dumps(_v) - _safe_input[_k] = _v - except (TypeError, ValueError, OverflowError): - pass - else: - _safe_input = {"input": tool_input} - yield streaming_service.format_tool_input_available( - tool_call_id, - tool_name, - _safe_input, - langchain_tool_call_id=langchain_tool_call_id, - ) - if content_builder is not None: - content_builder.on_tool_input_available( - tool_call_id, - tool_name, - _safe_input, - langchain_tool_call_id, - ) - - elif event_type == "on_tool_end": - active_tool_depth = max(0, active_tool_depth - 1) - run_id = event.get("run_id", "") - tool_name = event.get("name", "unknown_tool") - raw_output = event.get("data", {}).get("output", "") - staged_file_path = file_path_by_run.pop(run_id, None) if run_id else None - - if tool_name == "update_memory": - called_update_memory = True - - if hasattr(raw_output, "content"): - content = raw_output.content - if isinstance(content, str): - try: - tool_output = json.loads(content) - except (json.JSONDecodeError, TypeError): - tool_output = {"result": content} - elif isinstance(content, dict): - tool_output = content - else: - tool_output = {"result": str(content)} - elif isinstance(raw_output, dict): - tool_output = raw_output - else: - tool_output = {"result": str(raw_output) if raw_output else "completed"} - - if tool_name in ("write_file", "edit_file"): - if _tool_output_has_error(tool_output): - # Keep successful evidence if a previous write/edit in this turn succeeded. - pass - else: - result.write_succeeded = True - result.verification_succeeded = True - - # Look up the SAME card id used at on_tool_start (either the - # parity_v2 lc-id-derived ui_id or the legacy synthetic - # ``call_``) so the output event always lands on the - # same card as start/delta/available. Fallback preserves the - # legacy synthetic shape for parity_v2-OFF / unknown-run paths. - tool_call_id = ui_tool_call_id_by_run.get( - run_id, - f"call_{run_id[:32]}" if run_id else "call_unknown", - ) - original_step_id = tool_step_ids.get( - run_id, f"{step_prefix}-unknown-{run_id[:8]}" - ) - completed_step_ids.add(original_step_id) - - # Authoritative LangChain tool_call_id from the returned - # ``ToolMessage``. Falls back to whatever we matched - # at ``on_tool_start`` time (kept in ``lc_tool_call_id_by_run``) - # if the output isn't a ToolMessage. The value is stored in - # ``current_lc_tool_call_id`` so ``_emit_tool_output`` - # picks it up for every output emit below. - # - # Emitted in BOTH parity_v2 and legacy modes: the chat tool - # card needs the LangChain id to match against the - # ``data-action-log`` SSE event (keyed by ``lc_tool_call_id``) - # so the inline Revert button can light up. Reading - # ``raw_output.tool_call_id`` is a cheap, non-mutating attribute - # access that is safe regardless of feature-flag state. - current_lc_tool_call_id["value"] = None - authoritative = getattr(raw_output, "tool_call_id", None) - if isinstance(authoritative, str) and authoritative: - current_lc_tool_call_id["value"] = authoritative - if run_id: - lc_tool_call_id_by_run[run_id] = authoritative - elif run_id and run_id in lc_tool_call_id_by_run: - current_lc_tool_call_id["value"] = lc_tool_call_id_by_run[run_id] - - if tool_name == "read_file": - yield _emit_thinking_step( - step_id=original_step_id, - title="Reading file", - status="completed", - items=last_active_step_items, - ) - elif tool_name == "write_file": - yield _emit_thinking_step( - step_id=original_step_id, - title="Writing file", - status="completed", - items=last_active_step_items, - ) - elif tool_name == "edit_file": - yield _emit_thinking_step( - step_id=original_step_id, - title="Editing file", - status="completed", - items=last_active_step_items, - ) - elif tool_name == "glob": - yield _emit_thinking_step( - step_id=original_step_id, - title="Searching files", - status="completed", - items=last_active_step_items, - ) - elif tool_name == "grep": - yield _emit_thinking_step( - step_id=original_step_id, - title="Searching content", - status="completed", - items=last_active_step_items, - ) - elif tool_name == "rm": - yield _emit_thinking_step( - step_id=original_step_id, - title="Deleting file", - status="completed", - items=last_active_step_items, - ) - elif tool_name == "rmdir": - yield _emit_thinking_step( - step_id=original_step_id, - title="Deleting folder", - status="completed", - items=last_active_step_items, - ) - elif tool_name == "mkdir": - yield _emit_thinking_step( - step_id=original_step_id, - title="Creating folder", - status="completed", - items=last_active_step_items, - ) - elif tool_name == "move_file": - yield _emit_thinking_step( - step_id=original_step_id, - title="Moving file", - status="completed", - items=last_active_step_items, - ) - elif tool_name == "write_todos": - yield _emit_thinking_step( - step_id=original_step_id, - title="Planning tasks", - status="completed", - items=last_active_step_items, - ) - elif tool_name == "save_document": - result_str = ( - tool_output.get("result", "") - if isinstance(tool_output, dict) - else str(tool_output) - ) - is_error = "Error" in result_str - completed_items = [ - *last_active_step_items, - result_str[:80] if is_error else "Saved to knowledge base", - ] - yield _emit_thinking_step( - step_id=original_step_id, - title="Saving document", - status="completed", - items=completed_items, - ) - elif tool_name == "generate_image": - if isinstance(tool_output, dict) and not tool_output.get("error"): - completed_items = [ - *last_active_step_items, - "Image generated successfully", - ] - else: - error_msg = ( - tool_output.get("error", "Generation failed") - if isinstance(tool_output, dict) - else "Generation failed" - ) - completed_items = [*last_active_step_items, f"Error: {error_msg}"] - yield _emit_thinking_step( - step_id=original_step_id, - title="Generating image", - status="completed", - items=completed_items, - ) - elif tool_name == "scrape_webpage": - if isinstance(tool_output, dict): - title = tool_output.get("title", "Webpage") - word_count = tool_output.get("word_count", 0) - has_error = "error" in tool_output - if has_error: - completed_items = [ - *last_active_step_items, - f"Error: {tool_output.get('error', 'Failed to scrape')[:50]}", - ] - else: - completed_items = [ - *last_active_step_items, - f"Title: {title[:50]}{'...' if len(title) > 50 else ''}", - f"Extracted: {word_count:,} words", - ] - else: - completed_items = [*last_active_step_items, "Content extracted"] - yield _emit_thinking_step( - step_id=original_step_id, - title="Scraping webpage", - status="completed", - items=completed_items, - ) - elif tool_name == "generate_podcast": - podcast_status = ( - tool_output.get("status", "unknown") - if isinstance(tool_output, dict) - else "unknown" - ) - podcast_title = ( - tool_output.get("title", "Podcast") - if isinstance(tool_output, dict) - else "Podcast" - ) - if podcast_status in ("pending", "generating", "processing"): - completed_items = [ - f"Title: {podcast_title}", - "Podcast generation started", - "Processing in background...", - ] - elif podcast_status == "already_generating": - completed_items = [ - f"Title: {podcast_title}", - "Podcast already in progress", - "Please wait for it to complete", - ] - elif podcast_status in ("failed", "error"): - error_msg = ( - tool_output.get("error", "Unknown error") - if isinstance(tool_output, dict) - else "Unknown error" - ) - completed_items = [ - f"Title: {podcast_title}", - f"Error: {error_msg[:50]}", - ] - elif podcast_status in ("ready", "success"): - completed_items = [ - f"Title: {podcast_title}", - "Podcast ready", - ] - else: - completed_items = last_active_step_items - yield _emit_thinking_step( - step_id=original_step_id, - title="Generating podcast", - status="completed", - items=completed_items, - ) - elif tool_name == "generate_video_presentation": - vp_status = ( - tool_output.get("status", "unknown") - if isinstance(tool_output, dict) - else "unknown" - ) - vp_title = ( - tool_output.get("title", "Presentation") - if isinstance(tool_output, dict) - else "Presentation" - ) - if vp_status in ("pending", "generating"): - completed_items = [ - f"Title: {vp_title}", - "Presentation generation started", - "Processing in background...", - ] - elif vp_status == "failed": - error_msg = ( - tool_output.get("error", "Unknown error") - if isinstance(tool_output, dict) - else "Unknown error" - ) - completed_items = [ - f"Title: {vp_title}", - f"Error: {error_msg[:50]}", - ] - else: - completed_items = last_active_step_items - yield _emit_thinking_step( - step_id=original_step_id, - title="Generating video presentation", - status="completed", - items=completed_items, - ) - elif tool_name == "generate_report": - report_status = ( - tool_output.get("status", "unknown") - if isinstance(tool_output, dict) - else "unknown" - ) - report_title = ( - tool_output.get("title", "Report") - if isinstance(tool_output, dict) - else "Report" - ) - word_count = ( - tool_output.get("word_count", 0) - if isinstance(tool_output, dict) - else 0 - ) - is_revision = ( - tool_output.get("is_revision", False) - if isinstance(tool_output, dict) - else False - ) - step_title = "Revising report" if is_revision else "Generating report" - - if report_status == "ready": - completed_items = [ - f"Topic: {report_title}", - f"{word_count:,} words", - "Report ready", - ] - elif report_status == "failed": - error_msg = ( - tool_output.get("error", "Unknown error") - if isinstance(tool_output, dict) - else "Unknown error" - ) - completed_items = [ - f"Topic: {report_title}", - f"Error: {error_msg[:50]}", - ] - else: - completed_items = last_active_step_items - - yield _emit_thinking_step( - step_id=original_step_id, - title=step_title, - status="completed", - items=completed_items, - ) - elif tool_name in ("execute", "execute_code"): - raw_text = ( - tool_output.get("result", "") - if isinstance(tool_output, dict) - else str(tool_output) - ) - m = re.match(r"^Exit code:\s*(\d+)", raw_text) - exit_code_val = int(m.group(1)) if m else None - if exit_code_val is not None and exit_code_val == 0: - completed_items = [ - *last_active_step_items, - "Completed successfully", - ] - elif exit_code_val is not None: - completed_items = [ - *last_active_step_items, - f"Exit code: {exit_code_val}", - ] - else: - completed_items = [*last_active_step_items, "Finished"] - yield _emit_thinking_step( - step_id=original_step_id, - title="Running command", - status="completed", - items=completed_items, - ) - elif tool_name == "ls": - if isinstance(tool_output, dict): - ls_output = tool_output.get("result", "") - elif isinstance(tool_output, str): - ls_output = tool_output - else: - ls_output = str(tool_output) if tool_output else "" - file_names: list[str] = [] - if ls_output: - paths: list[str] = [] - try: - parsed = ast.literal_eval(ls_output) - if isinstance(parsed, list): - paths = [str(p) for p in parsed] - except (ValueError, SyntaxError): - paths = [ - line.strip() - for line in ls_output.strip().split("\n") - if line.strip() - ] - for p in paths: - name = p.rstrip("/").split("/")[-1] - if name and len(name) <= 40: - file_names.append(name) - elif name: - file_names.append(name[:37] + "...") - if file_names: - if len(file_names) <= 5: - completed_items = [f"[{name}]" for name in file_names] - else: - completed_items = [f"[{name}]" for name in file_names[:4]] - completed_items.append(f"(+{len(file_names) - 4} more)") - else: - completed_items = ["No files found"] - yield _emit_thinking_step( - step_id=original_step_id, - title="Listing files", - status="completed", - items=completed_items, - ) - else: - # Fallback completion title — see the matching in-progress - # branch above for the wording rationale. - fallback_title = ( - tool_name.replace("_", " ").strip().capitalize() or tool_name - ) - yield _emit_thinking_step( - step_id=original_step_id, - title=fallback_title, - status="completed", - items=last_active_step_items, - ) - - just_finished_tool = True - last_active_step_id = None - last_active_step_title = "" - last_active_step_items = [] - - if tool_name == "generate_podcast": - yield _emit_tool_output( - tool_call_id, - tool_output - if isinstance(tool_output, dict) - else {"result": tool_output}, - ) - if isinstance(tool_output, dict) and tool_output.get("status") in ( - "pending", - "generating", - "processing", - ): - yield streaming_service.format_terminal_info( - f"Podcast queued: {tool_output.get('title', 'Podcast')}", - "success", - ) - elif isinstance(tool_output, dict) and tool_output.get("status") in ( - "ready", - "success", - ): - yield streaming_service.format_terminal_info( - f"Podcast generated successfully: {tool_output.get('title', 'Podcast')}", - "success", - ) - elif isinstance(tool_output, dict) and tool_output.get("status") in ( - "failed", - "error", - ): - error_msg = tool_output.get("error", "Unknown error") - yield streaming_service.format_terminal_info( - f"Podcast generation failed: {error_msg}", - "error", - ) - elif tool_name == "generate_video_presentation": - yield _emit_tool_output( - tool_call_id, - tool_output - if isinstance(tool_output, dict) - else {"result": tool_output}, - ) - if ( - isinstance(tool_output, dict) - and tool_output.get("status") == "pending" - ): - yield streaming_service.format_terminal_info( - f"Video presentation queued: {tool_output.get('title', 'Presentation')}", - "success", - ) - elif ( - isinstance(tool_output, dict) - and tool_output.get("status") == "failed" - ): - error_msg = ( - tool_output.get("error", "Unknown error") - if isinstance(tool_output, dict) - else "Unknown error" - ) - yield streaming_service.format_terminal_info( - f"Presentation generation failed: {error_msg}", - "error", - ) - elif tool_name == "generate_image": - yield _emit_tool_output( - tool_call_id, - tool_output - if isinstance(tool_output, dict) - else {"result": tool_output}, - ) - if isinstance(tool_output, dict): - if tool_output.get("error"): - yield streaming_service.format_terminal_info( - f"Image generation failed: {tool_output['error'][:60]}", - "error", - ) - else: - yield streaming_service.format_terminal_info( - "Image generated successfully", - "success", - ) - elif tool_name == "scrape_webpage": - if isinstance(tool_output, dict): - display_output = { - k: v for k, v in tool_output.items() if k != "content" - } - if "content" in tool_output: - content = tool_output.get("content", "") - display_output["content_preview"] = ( - content[:500] + "..." if len(content) > 500 else content - ) - yield _emit_tool_output( - tool_call_id, - display_output, - ) - else: - yield _emit_tool_output( - tool_call_id, - {"result": tool_output}, - ) - if isinstance(tool_output, dict) and "error" not in tool_output: - title = tool_output.get("title", "Webpage") - word_count = tool_output.get("word_count", 0) - yield streaming_service.format_terminal_info( - f"Scraped: {title[:40]}{'...' if len(title) > 40 else ''} ({word_count:,} words)", - "success", - ) - else: - error_msg = ( - tool_output.get("error", "Failed to scrape") - if isinstance(tool_output, dict) - else "Failed to scrape" - ) - yield streaming_service.format_terminal_info( - f"Scrape failed: {error_msg}", - "error", - ) - elif tool_name in ("write_file", "edit_file"): - resolved_path = _extract_resolved_file_path( - tool_name=tool_name, - tool_output=tool_output, - tool_input={"file_path": staged_file_path} - if staged_file_path - else None, - ) - result_text = _tool_output_to_text(tool_output) - if _tool_output_has_error(tool_output): - yield _emit_tool_output( - tool_call_id, - { - "status": "error", - "error": result_text, - "path": resolved_path, - }, - ) - else: - yield _emit_tool_output( - tool_call_id, - { - "status": "completed", - "path": resolved_path, - "result": result_text, - }, - ) - elif tool_name == "generate_report": - # Stream the full report result so frontend can render the ReportCard - yield _emit_tool_output( - tool_call_id, - tool_output - if isinstance(tool_output, dict) - else {"result": tool_output}, - ) - # Send appropriate terminal message based on status - if ( - isinstance(tool_output, dict) - and tool_output.get("status") == "ready" - ): - word_count = tool_output.get("word_count", 0) - yield streaming_service.format_terminal_info( - f"Report generated: {tool_output.get('title', 'Report')} ({word_count:,} words)", - "success", - ) - else: - error_msg = ( - tool_output.get("error", "Unknown error") - if isinstance(tool_output, dict) - else "Unknown error" - ) - yield streaming_service.format_terminal_info( - f"Report generation failed: {error_msg}", - "error", - ) - elif tool_name == "generate_resume": - yield _emit_tool_output( - tool_call_id, - tool_output - if isinstance(tool_output, dict) - else {"result": tool_output}, - ) - if ( - isinstance(tool_output, dict) - and tool_output.get("status") == "ready" - ): - yield streaming_service.format_terminal_info( - f"Resume generated: {tool_output.get('title', 'Resume')}", - "success", - ) - else: - error_msg = ( - tool_output.get("error", "Unknown error") - if isinstance(tool_output, dict) - else "Unknown error" - ) - yield streaming_service.format_terminal_info( - f"Resume generation failed: {error_msg}", - "error", - ) - elif tool_name in ( - "create_notion_page", - "update_notion_page", - "delete_notion_page", - "create_linear_issue", - "update_linear_issue", - "delete_linear_issue", - "create_google_drive_file", - "delete_google_drive_file", - "create_onedrive_file", - "delete_onedrive_file", - "create_dropbox_file", - "delete_dropbox_file", - "create_gmail_draft", - "update_gmail_draft", - "send_gmail_email", - "trash_gmail_email", - "create_calendar_event", - "update_calendar_event", - "delete_calendar_event", - "create_jira_issue", - "update_jira_issue", - "delete_jira_issue", - "create_confluence_page", - "update_confluence_page", - "delete_confluence_page", - ): - yield _emit_tool_output( - tool_call_id, - tool_output - if isinstance(tool_output, dict) - else {"result": tool_output}, - ) - elif tool_name in ("execute", "execute_code"): - raw_text = ( - tool_output.get("result", "") - if isinstance(tool_output, dict) - else str(tool_output) - ) - exit_code: int | None = None - output_text = raw_text - m = re.match(r"^Exit code:\s*(\d+)", raw_text) - if m: - exit_code = int(m.group(1)) - om = re.search(r"\nOutput:\n([\s\S]*)", raw_text) - output_text = om.group(1) if om else "" - thread_id_str = config.get("configurable", {}).get("thread_id", "") - - for sf_match in re.finditer( - r"^SANDBOX_FILE:\s*(.+)$", output_text, re.MULTILINE - ): - fpath = sf_match.group(1).strip() - if fpath and fpath not in result.sandbox_files: - result.sandbox_files.append(fpath) - - yield _emit_tool_output( - tool_call_id, - { - "exit_code": exit_code, - "output": output_text, - "thread_id": thread_id_str, - }, - ) - elif tool_name == "web_search": - xml = ( - tool_output.get("result", str(tool_output)) - if isinstance(tool_output, dict) - else str(tool_output) - ) - citations: dict[str, dict[str, str]] = {} - for m in re.finditer( - r"<!\[CDATA\[(.*?)\]\]>\s*", - xml, - ): - title, url = m.group(1).strip(), m.group(2).strip() - if url.startswith("http") and url not in citations: - citations[url] = {"title": title} - for m in re.finditer( - r"", - xml, - ): - chunk_url, content = m.group(1).strip(), m.group(2).strip() - if ( - chunk_url.startswith("http") - and chunk_url in citations - and content - ): - citations[chunk_url]["snippet"] = ( - content[:200] + "…" if len(content) > 200 else content - ) - yield _emit_tool_output( - tool_call_id, - {"status": "completed", "citations": citations}, - ) - else: - yield _emit_tool_output( - tool_call_id, - {"status": "completed", "result_length": len(str(tool_output))}, - ) - yield streaming_service.format_terminal_info( - f"Tool {tool_name} completed", "success" - ) - - elif event_type == "on_custom_event" and event.get("name") == "report_progress": - # Live progress updates from inside the generate_report tool - data = event.get("data", {}) - message = data.get("message", "") - if message and last_active_step_id: - phase = data.get("phase", "") - # Always keep the "Topic: ..." line - topic_items = [ - item for item in last_active_step_items if item.startswith("Topic:") - ] - - if phase in ("revising_section", "adding_section"): - # During section-level ops: keep plan summary + show current op - plan_items = [ - item - for item in last_active_step_items - if item.startswith("Topic:") - or item.startswith("Modifying ") - or item.startswith("Adding ") - or item.startswith("Removing ") - ] - # Only keep plan_items that don't end with "..." (not progress lines) - plan_items = [ - item for item in plan_items if not item.endswith("...") - ] - last_active_step_items = [*plan_items, message] - else: - # Phase transitions: replace everything after topic - last_active_step_items = [*topic_items, message] - - yield _emit_thinking_step( - step_id=last_active_step_id, - title=last_active_step_title, - status="in_progress", - items=last_active_step_items, - ) - - elif ( - event_type == "on_custom_event" and event.get("name") == "document_created" - ): - data = event.get("data", {}) - if data.get("id"): - yield streaming_service.format_data( - "documents-updated", - { - "action": "created", - "document": data, - }, - ) - - elif event_type == "on_custom_event" and event.get("name") == "action_log": - # Surface a freshly committed AgentActionLog row so the chat - # tool card can render its Revert button immediately. - data = event.get("data", {}) - if data.get("id") is not None: - yield streaming_service.format_data("action-log", data) - - elif ( - event_type == "on_custom_event" - and event.get("name") == "action_log_updated" - ): - # Reversibility flipped in kb_persistence after the SAVEPOINT - # for a destructive op (rm/rmdir/move/edit/write) committed. - # Frontend uses this to flip the card's Revert - # button on without re-fetching the actions list. - data = event.get("data", {}) - if data.get("id") is not None: - yield streaming_service.format_data("action-log-updated", data) - - elif event_type in ("on_chain_end", "on_agent_end"): - if current_text_id is not None: - yield streaming_service.format_text_end(current_text_id) - if content_builder is not None: - content_builder.on_text_end(current_text_id) - current_text_id = None - - if current_text_id is not None: - yield streaming_service.format_text_end(current_text_id) - if content_builder is not None: - content_builder.on_text_end(current_text_id) - - completion_event = complete_current_step() - if completion_event: - yield completion_event + accumulated_text = result.accumulated_text state = await agent.aget_state(config) state_values = getattr(state, "values", {}) or {} @@ -2397,7 +912,6 @@ async def _stream_agent_events( result.commit_gate_reason = "" result.accumulated_text = accumulated_text - result.agent_called_update_memory = called_update_memory _log_file_contract("turn_outcome", result) interrupt_value = _first_interrupt_value(state) diff --git a/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py b/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py index 9142dd914..9a309f9d7 100644 --- a/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py +++ b/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py @@ -5,7 +5,6 @@ from __future__ import annotations from collections.abc import AsyncIterator from typing import Any -from app.agents.new_chat.feature_flags import get_flags from app.tasks.chat.streaming.graph_stream.result import StreamingResult from app.tasks.chat.streaming.relay.event_relay import EventRelay from app.tasks.chat.streaming.relay.state import AgentEventRelayState @@ -30,7 +29,6 @@ async def stream_output( initial_step_id=initial_step_id, initial_step_title=initial_step_title, initial_step_items=initial_step_items, - parity_v2=bool(get_flags().enable_stream_parity_v2), ) astream_kwargs: dict[str, Any] = {"config": config, "version": "v2"} diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py b/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py index 861342b32..ef86dae56 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py @@ -33,7 +33,7 @@ def iter_chat_model_stream_frames( reasoning_delta = parts["reasoning"] text_delta = parts["text"] - if state.parity_v2 and reasoning_delta: + if reasoning_delta: if state.current_text_id is not None: yield streaming_service.format_text_end(state.current_text_id) if content_builder is not None: @@ -100,7 +100,7 @@ def iter_chat_model_stream_frames( if content_builder is not None: content_builder.on_text_delta(state.current_text_id, text_delta) - if state.parity_v2 and parts["tool_call_chunks"]: + if parts["tool_call_chunks"]: for tcc in parts["tool_call_chunks"]: idx = tcc.get("index") diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py index c316cc74a..e7d2d7f78 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py @@ -77,12 +77,11 @@ def iter_tool_start_frames( yield emit_thinking_step_frame(**frame_kw) matched_meta: dict[str, str] | None = None - if state.parity_v2: - taken_ui_ids = set(state.ui_tool_call_id_by_run.values()) - for meta in state.index_to_meta.values(): - if meta["name"] == tool_name and meta["ui_id"] not in taken_ui_ids: - matched_meta = meta - break + taken_ui_ids = set(state.ui_tool_call_id_by_run.values()) + for meta in state.index_to_meta.values(): + if meta["name"] == tool_name and meta["ui_id"] not in taken_ui_ids: + matched_meta = meta + break tool_call_id: str langchain_tool_call_id: str | None = None @@ -97,13 +96,12 @@ def iter_tool_start_frames( if run_id else streaming_service.generate_tool_call_id() ) - if state.parity_v2: - langchain_tool_call_id = match_buffered_langchain_tool_call_id( - state.pending_tool_call_chunks, - tool_name, - run_id, - state.lc_tool_call_id_by_run, - ) + langchain_tool_call_id = match_buffered_langchain_tool_call_id( + state.pending_tool_call_chunks, + tool_name, + run_id, + state.lc_tool_call_id_by_run, + ) yield streaming_service.format_tool_input_start( tool_call_id, tool_name, diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/state.py b/surfsense_backend/app/tasks/chat/streaming/relay/state.py index e8e35d0b2..7bd996606 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/state.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/state.py @@ -22,7 +22,6 @@ class AgentEventRelayState: active_tool_depth: int = 0 called_update_memory: bool = False current_reasoning_id: str | None = None - parity_v2: bool = False pending_tool_call_chunks: list[dict[str, Any]] = field(default_factory=list) lc_tool_call_id_by_run: dict[str, str] = field(default_factory=dict) file_path_by_run: dict[str, str] = field(default_factory=dict) @@ -39,7 +38,6 @@ class AgentEventRelayState: initial_step_id: str | None = None, initial_step_title: str = "", initial_step_items: list[str] | None = None, - parity_v2: bool, ) -> AgentEventRelayState: counter = 1 if initial_step_id else 0 return cls( @@ -47,7 +45,6 @@ class AgentEventRelayState: last_active_step_id=initial_step_id, last_active_step_title=initial_step_title, last_active_step_items=list(initial_step_items or []), - parity_v2=parity_v2, ) def next_thinking_step_id(self, step_prefix: str) -> str: diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_feature_flags.py b/surfsense_backend/tests/unit/agents/new_chat/test_feature_flags.py index 6800be2af..099aea882 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_feature_flags.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_feature_flags.py @@ -31,7 +31,6 @@ def _clear_all(monkeypatch: pytest.MonkeyPatch) -> None: "SURFSENSE_ENABLE_KB_PLANNER_RUNNABLE", "SURFSENSE_ENABLE_ACTION_LOG", "SURFSENSE_ENABLE_REVERT_ROUTE", - "SURFSENSE_ENABLE_STREAM_PARITY_V2", "SURFSENSE_ENABLE_PLUGIN_LOADER", "SURFSENSE_ENABLE_OTEL", "SURFSENSE_ENABLE_AGENT_CACHE", @@ -61,7 +60,6 @@ def test_defaults_match_shipped_agent_stack(monkeypatch: pytest.MonkeyPatch) -> assert flags.enable_kb_planner_runnable is True assert flags.enable_action_log is True assert flags.enable_revert_route is True - assert flags.enable_stream_parity_v2 is True assert flags.enable_plugin_loader is False assert flags.enable_otel is False # Phase 2: agent cache is now default-on (the prerequisite tool @@ -127,7 +125,6 @@ def test_each_flag_can_be_set_independently(monkeypatch: pytest.MonkeyPatch) -> "enable_kb_planner_runnable": "SURFSENSE_ENABLE_KB_PLANNER_RUNNABLE", "enable_action_log": "SURFSENSE_ENABLE_ACTION_LOG", "enable_revert_route": "SURFSENSE_ENABLE_REVERT_ROUTE", - "enable_stream_parity_v2": "SURFSENSE_ENABLE_STREAM_PARITY_V2", "enable_plugin_loader": "SURFSENSE_ENABLE_PLUGIN_LOADER", "enable_otel": "SURFSENSE_ENABLE_OTEL", } diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py index 892bb7a6a..9ae7defec 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py @@ -137,7 +137,7 @@ def test_complete_active_thinking_step_mirrors_closure_semantics() -> None: def test_agent_event_relay_state_factory_matches_counter_rule() -> None: - s0 = AgentEventRelayState.for_invocation(parity_v2=False) + s0 = AgentEventRelayState.for_invocation() assert s0.thinking_step_counter == 0 assert s0.last_active_step_id is None @@ -145,11 +145,9 @@ def test_agent_event_relay_state_factory_matches_counter_rule() -> None: initial_step_id="thinking-resume-1", initial_step_title="Inherited", initial_step_items=["Topic: X"], - parity_v2=True, ) assert s1.thinking_step_counter == 1 assert s1.last_active_step_id == "thinking-resume-1" - assert s1.parity_v2 is True assert s1.next_thinking_step_id("thinking") == "thinking-2" diff --git a/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py b/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py index c317eba20..4b1fadd9c 100644 --- a/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py +++ b/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py @@ -161,7 +161,7 @@ class TestToolHeavyTurn: _assert_jsonb_safe(snap) def test_tool_input_available_without_prior_start_creates_card(self): - # Legacy / parity_v2-OFF path: tool-input-available may be + # Late-registration: tool-input-available may be # emitted without a prior tool-input-start (no streamed # tool_call_chunks). The card should still be created. b = AssistantContentBuilder() @@ -187,7 +187,7 @@ class TestToolHeavyTurn: assert part["result"] == {"matches": 3} def test_tool_input_start_idempotent_for_same_ui_id(self): - # parity_v2: tool-input-start can fire from BOTH the chunk + # tool-input-start can fire from BOTH the chunk # registration path AND the canonical ``on_tool_start`` path. # The second call must not create a duplicate part. b = AssistantContentBuilder() diff --git a/surfsense_backend/tests/unit/tasks/chat/test_tool_input_streaming.py b/surfsense_backend/tests/unit/tasks/chat/test_tool_input_streaming.py index 60750396c..ada32d168 100644 --- a/surfsense_backend/tests/unit/tasks/chat/test_tool_input_streaming.py +++ b/surfsense_backend/tests/unit/tasks/chat/test_tool_input_streaming.py @@ -1,16 +1,13 @@ """Unit tests for live tool-call argument streaming. -Pins the wire format that ``_stream_agent_events`` emits when -``SURFSENSE_ENABLE_STREAM_PARITY_V2=true``: ``tool-input-start`` → -``tool-input-delta``... → ``tool-input-available`` → ``tool-output-available`` -all keyed by the same LangChain ``tool_call.id``. +Pins the wire format that ``_stream_agent_events`` emits: +``tool-input-start`` → ``tool-input-delta``... → ``tool-input-available`` → +``tool-output-available``, keyed consistently with LangChain ``tool_call.id`` +when the model streams indexed chunks. Identity is tracked in ``index_to_meta`` (per-chunk ``index``) and -``ui_tool_call_id_by_run`` (LangGraph ``run_id``); both are private to -``_stream_agent_events`` so we exercise them via the public wire output. - -These tests also lock in the legacy / parity_v2-OFF behaviour so the -synthetic ``call_`` shape stays stable for older clients. +``ui_tool_call_id_by_run`` (LangGraph ``run_id``); both are internal to the +streaming layer so we assert on the public SSE payloads. """ from __future__ import annotations @@ -22,8 +19,6 @@ from typing import Any import pytest -import app.tasks.chat.stream_new_chat as stream_module -from app.agents.new_chat.feature_flags import AgentFeatureFlags from app.services.new_streaming_service import VercelStreamingService from app.tasks.chat.stream_new_chat import ( StreamResult, @@ -164,24 +159,6 @@ def _tool_end( } -@pytest.fixture -def parity_v2_on(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - stream_module, - "get_flags", - lambda: AgentFeatureFlags(enable_stream_parity_v2=True), - ) - - -@pytest.fixture -def parity_v2_off(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - stream_module, - "get_flags", - lambda: AgentFeatureFlags(enable_stream_parity_v2=False), - ) - - async def _drain( events: list[dict[str, Any]], state: _FakeAgentState | None = None ) -> list[dict[str, Any]]: @@ -253,12 +230,12 @@ class TestLegacyMatch: # --------------------------------------------------------------------------- -# parity_v2 wire format tests. +# Tool input streaming wire format # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_idless_chunk_merging_by_index(parity_v2_on: None) -> None: +async def test_idless_chunk_merging_by_index() -> None: """First chunk carries id+name; later idless chunks at the same ``index`` merge into the SAME ``tool-input-start`` ui id and emit one ``tool-input-delta`` per chunk.""" @@ -302,9 +279,7 @@ async def test_idless_chunk_merging_by_index(parity_v2_on: None) -> None: @pytest.mark.asyncio -async def test_two_interleaved_tool_calls_route_by_index( - parity_v2_on: None, -) -> None: +async def test_two_interleaved_tool_calls_route_by_index() -> None: """Two same-name calls with distinct indices keep their deltas routed to the right card.""" events = [ @@ -344,7 +319,7 @@ async def test_two_interleaved_tool_calls_route_by_index( @pytest.mark.asyncio -async def test_identity_stable_across_lifecycle(parity_v2_on: None) -> None: +async def test_identity_stable_across_lifecycle() -> None: """Whatever id ``tool-input-start`` chose must be the SAME id used on ``tool-input-available`` AND ``tool-output-available``.""" events = [ @@ -367,7 +342,7 @@ async def test_identity_stable_across_lifecycle(parity_v2_on: None) -> None: @pytest.mark.asyncio -async def test_no_duplicate_tool_input_start(parity_v2_on: None) -> None: +async def test_no_duplicate_tool_input_start() -> None: """When the chunk-emission loop already fired ``tool-input-start`` for this run, ``on_tool_start`` MUST NOT emit a second one.""" events = [ @@ -386,9 +361,7 @@ async def test_no_duplicate_tool_input_start(parity_v2_on: None) -> None: @pytest.mark.asyncio -async def test_active_text_closes_before_early_tool_input_start( - parity_v2_on: None, -) -> None: +async def test_active_text_closes_before_early_tool_input_start() -> None: """Streaming a text-delta then a tool-call chunk in subsequent chunks: the wire MUST contain ``text-end`` before the FIRST ``tool-input-start`` (clean part boundary on the frontend).""" @@ -409,9 +382,7 @@ async def test_active_text_closes_before_early_tool_input_start( @pytest.mark.asyncio -async def test_mixed_text_and_tool_chunk_preserve_order( - parity_v2_on: None, -) -> None: +async def test_mixed_text_and_tool_chunk_preserve_order() -> None: """One AIMessageChunk that carries BOTH ``text`` content AND ``tool_call_chunks`` should emit the text delta FIRST, then close text, then ``tool-input-start``+``tool-input-delta``.""" @@ -441,45 +412,7 @@ async def test_mixed_text_and_tool_chunk_preserve_order( @pytest.mark.asyncio -async def test_parity_v2_off_preserves_legacy_shape( - parity_v2_off: None, -) -> None: - """When the flag is OFF, no deltas are emitted and the ``toolCallId`` - is ``call_`` (NOT the lc id).""" - events = [ - _model_stream( - tool_call_chunks=[ - {"id": "lc-1", "name": "ls", "args": '{"path":"/"}', "index": 0} - ] - ), - _tool_start(name="ls", run_id="run-A", input_payload={"path": "/"}), - _tool_end(name="ls", run_id="run-A", tool_call_id="lc-1"), - ] - payloads = await _drain(events) - - assert _of_type(payloads, "tool-input-delta") == [] - starts = _of_type(payloads, "tool-input-start") - assert len(starts) == 1 - assert starts[0]["toolCallId"].startswith("call_run-A") - # No ``langchainToolCallId`` propagation on ``tool-input-start`` in - # legacy mode (the start event fires before the ToolMessage is - # available, so we can't extract the authoritative LangChain id yet). - assert "langchainToolCallId" not in starts[0] - output = _of_type(payloads, "tool-output-available") - assert output[0]["toolCallId"].startswith("call_run-A") - # ``tool-output-available`` MUST carry ``langchainToolCallId`` even - # in legacy mode: the chat tool card uses it to backfill the - # LangChain id and join against the ``data-action-log`` SSE event - # (keyed by ``lc_tool_call_id``) so the inline Revert button can - # light up. Sourced from the returned ``ToolMessage.tool_call_id``, - # which is populated regardless of feature-flag state. - assert output[0]["langchainToolCallId"] == "lc-1" - - -@pytest.mark.asyncio -async def test_skip_append_prevents_stale_id_reuse( - parity_v2_on: None, -) -> None: +async def test_skip_append_prevents_stale_id_reuse() -> None: """Two same-name tools: the SECOND tool's ``langchainToolCallId`` must NOT come from the first tool's chunk (``pending_tool_call_chunks`` must stay empty for indexed-registered chunks).""" @@ -506,9 +439,7 @@ async def test_skip_append_prevents_stale_id_reuse( @pytest.mark.asyncio -async def test_registration_waits_for_both_id_and_name( - parity_v2_on: None, -) -> None: +async def test_registration_waits_for_both_id_and_name() -> None: """An id-only chunk (no name yet) must NOT emit ``tool-input-start``.""" events = [ _model_stream( @@ -520,12 +451,9 @@ async def test_registration_waits_for_both_id_and_name( @pytest.mark.asyncio -async def test_unmatched_fallback_still_attaches_lc_id( - parity_v2_on: None, -) -> None: - """parity_v2 ON, but the provider didn't include an ``index``: the - legacy fallback path must still emit ``tool-input-start`` with the - matching ``langchainToolCallId``.""" +async def test_unmatched_fallback_still_attaches_lc_id() -> None: + """When the provider omits chunk ``index``, buffered chunks still get a + ``tool-input-start`` with the matching ``langchainToolCallId``.""" events = [ # No index on the chunk → not registered into index_to_meta; # falls through to ``pending_tool_call_chunks`` so the legacy @@ -542,9 +470,7 @@ async def test_unmatched_fallback_still_attaches_lc_id( @pytest.mark.asyncio -async def test_interrupt_request_uses_task_that_contains_interrupt( - parity_v2_on: None, -) -> None: +async def test_interrupt_request_uses_task_that_contains_interrupt() -> None: interrupt_payload = { "type": "calendar_event_create", "action": { diff --git a/surfsense_web/components/assistant-ui/reasoning-message-part.tsx b/surfsense_web/components/assistant-ui/reasoning-message-part.tsx index 70636eab8..6e7aaf048 100644 --- a/surfsense_web/components/assistant-ui/reasoning-message-part.tsx +++ b/surfsense_web/components/assistant-ui/reasoning-message-part.tsx @@ -7,8 +7,8 @@ import { TextShimmerLoader } from "@/components/prompt-kit/loader"; import { cn } from "@/lib/utils"; /** - * Renders the structured `reasoning` part emitted by the backend's - * stream-parity v2 path (A1). + * Renders the structured `reasoning` part emitted by the backend stream + * (typed reasoning deltas from the chat model). * * Behaviour mirrors the existing `ThinkingStepsDisplay`: * - collapsed by default; diff --git a/surfsense_web/components/assistant-ui/tool-fallback.tsx b/surfsense_web/components/assistant-ui/tool-fallback.tsx index 06082c9c7..ba58f4158 100644 --- a/surfsense_web/components/assistant-ui/tool-fallback.tsx +++ b/surfsense_web/components/assistant-ui/tool-fallback.tsx @@ -48,13 +48,11 @@ import { cn } from "@/lib/utils"; * stream, post-stream reversibility flip, and explicit revert clicks. * * Match key (in priority order): - * 1. ``a.tool_call_id === toolCallId`` — direct hit in parity_v2 when - * the model streamed ``tool_call_chunks`` so the card's synthetic - * id IS the LangChain id. - * 2. ``a.tool_call_id === langchainToolCallId`` — legacy mode (or - * parity_v2 with provider-side chunk emission) where the card's - * synthetic id is ``call_`` and the LangChain id is - * backfilled onto the part by ``tool-output-available``. + * 1. ``a.tool_call_id === toolCallId`` — direct hit when the model + * streamed ``tool_call_chunks`` so the card id matches the LangChain id. + * 2. ``a.tool_call_id === langchainToolCallId`` — synthetic card id is + * ``call_`` and the LangChain id is backfilled by + * ``tool-output-available``. * 3. ``(chat_turn_id, tool_name, position-within-turn)`` — fallback * for cards whose synthetic id is ``call_`` AND whose * ``langchainToolCallId`` never got backfilled (provider emitted @@ -116,7 +114,7 @@ function ToolCardRevertButton({ const action = useMemo(() => { // Tier 1 + 2: O(1) Map-backed direct id match. Covers - // ~all parity_v2 streams and any legacy stream that backfilled + // Indexed chunk streams and any stream that backfilled // ``langchainToolCallId`` via ``tool-output-available``. const direct = findByToolCallId(toolCallId) ?? findByToolCallId(langchainToolCallId); if (direct) return direct; diff --git a/surfsense_web/lib/chat/streaming-state.ts b/surfsense_web/lib/chat/streaming-state.ts index 27047ecfe..809e214d1 100644 --- a/surfsense_web/lib/chat/streaming-state.ts +++ b/surfsense_web/lib/chat/streaming-state.ts @@ -421,9 +421,8 @@ export type SSEEvent = /** * Live tool-call argument delta. Concatenated into * ``argsText`` on the matching ``tool-call`` content part - * by ``appendToolInputDelta``. parity_v2 only — the legacy - * code path emits ``tool-input-available`` without prior - * deltas. + * by ``appendToolInputDelta``. Some providers emit + * ``tool-input-available`` without prior deltas. */ type: "tool-input-delta"; toolCallId: string; From f0f87107f2a9922bb6e4f8d2cfb48764dc098e5a Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 22:46:58 +0200 Subject: [PATCH 61/96] Track active task span id on the agent event relay state. --- .../app/tasks/chat/streaming/relay/state.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/state.py b/surfsense_backend/app/tasks/chat/streaming/relay/state.py index 7bd996606..82525a52f 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/state.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/state.py @@ -8,7 +8,13 @@ from typing import Any @dataclass class AgentEventRelayState: - """Tracks text, thinking steps, tool depth, and pending tool-call metadata.""" + """Tracks text, thinking steps, tool depth, and pending tool-call metadata. + + ``active_span_id`` groups steps/tools for one open ``task`` episode. + ``active_task_run_id`` is the LangGraph ``run_id`` of that ``task`` so we + only clear the span when that run ends (not when child tools end). Handlers + will set/clear these via ``task_span`` helpers in a later change. + """ accumulated_text: str = "" current_text_id: str | None = None @@ -30,6 +36,17 @@ class AgentEventRelayState: current_lc_tool_call_id: dict[str, str | None] = field( default_factory=lambda: {"value": None} ) + # Open ``task`` delegation span (one id shared by nested activity); unset outside. + active_span_id: str | None = None + active_task_run_id: str | None = None + # Span id minted when a ``task`` tool_call_chunk registers (before ``on_tool_start``). + pending_task_span_by_lc: dict[str, str] = field(default_factory=dict) + + def span_metadata_if_active(self) -> dict[str, Any] | None: + """``{"spanId": ...}`` when a span is active; ``None`` otherwise.""" + if self.active_span_id: + return {"spanId": self.active_span_id} + return None @classmethod def for_invocation( From f944cdacb753369af8e117e34cdaca39c3c48c73 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 22:47:03 +0200 Subject: [PATCH 62/96] Add helpers to open and close task delegation span ids. --- .../tasks/chat/streaming/relay/task_span.py | 74 +++++++++++++++++++ .../tasks/chat/streaming/test_task_span.py | 69 +++++++++++++++++ 2 files changed, 143 insertions(+) create mode 100644 surfsense_backend/app/tasks/chat/streaming/relay/task_span.py create mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/test_task_span.py diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/task_span.py b/surfsense_backend/app/tasks/chat/streaming/relay/task_span.py new file mode 100644 index 000000000..c4cdf24ba --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/relay/task_span.py @@ -0,0 +1,74 @@ +"""Open/close ``active_span_id`` around a delegating ``task`` tool run.""" + +from __future__ import annotations + +import uuid + +from app.tasks.chat.streaming.relay.state import AgentEventRelayState + + +def new_span_id() -> str: + """One delegation-episode id (shared by activity under an open ``task``).""" + return f"spn_{uuid.uuid4().hex}" + + +def _run_key(run_id: str) -> str: + return (run_id or "").strip() + + +def _lc_key(langchain_tool_call_id: str | None) -> str: + return (langchain_tool_call_id or "").strip() + + +def ensure_pending_task_span_for_lc(state: AgentEventRelayState, lc_id: str) -> str: + """Return span id for this LangChain tool call id, storing it in ``pending`` if new. + + Used from ``chat_model_stream`` when the first ``task`` chunk registers so + early ``tool-input-start`` can carry ``metadata.spanId`` before ``on_tool_start``. + """ + key = _lc_key(lc_id) + if not key: + return new_span_id() + existing = state.pending_task_span_by_lc.get(key) + if existing: + return existing + sid = new_span_id() + state.pending_task_span_by_lc[key] = sid + return sid + + +def open_task_span( + state: AgentEventRelayState, + *, + run_id: str, + langchain_tool_call_id: str | None = None, +) -> str: + """Set ``active_span_id`` from pending (same lc) or mint; remember ``active_task_run_id``. + + Call when the ``task`` tool **starts**. Nested ``task`` is not supported: + a second call replaces the previous span without restoring it. + """ + key = _lc_key(langchain_tool_call_id) + sid: str | None = state.pending_task_span_by_lc.pop(key, None) if key else None + if not sid: + sid = new_span_id() + state.active_span_id = sid + state.active_task_run_id = _run_key(run_id) or None + return sid + + +def clear_task_span_if_delegating_task_ended( + state: AgentEventRelayState, + *, + tool_name: str, + run_id: str, +) -> None: + """Clear span state only when this event is the end of the opening ``task`` run.""" + if tool_name != "task": + return + if state.active_task_run_id is None: + return + if state.active_task_run_id != _run_key(run_id): + return + state.active_span_id = None + state.active_task_run_id = None diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_task_span.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_task_span.py new file mode 100644 index 000000000..349c9879c --- /dev/null +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_task_span.py @@ -0,0 +1,69 @@ +"""Unit tests for ``task_span`` open/close helpers.""" + +from __future__ import annotations + +import pytest + +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, + ensure_pending_task_span_for_lc, + open_task_span, +) + +pytestmark = pytest.mark.unit + + +def test_open_task_span_sets_span_and_run_id() -> None: + state = AgentEventRelayState.for_invocation() + sid = open_task_span(state, run_id="run-abc") + assert sid.startswith("spn_") + assert state.active_span_id == sid + assert state.active_task_run_id == "run-abc" + assert state.span_metadata_if_active() == {"spanId": sid} + + +def test_clear_ignored_for_non_task_tool() -> None: + state = AgentEventRelayState.for_invocation() + open_task_span(state, run_id="run-1") + sid = state.active_span_id + clear_task_span_if_delegating_task_ended( + state, tool_name="web_search", run_id="run-1" + ) + assert state.active_span_id == sid + assert state.active_task_run_id == "run-1" + + +def test_clear_ignored_when_task_run_id_mismatches() -> None: + state = AgentEventRelayState.for_invocation() + open_task_span(state, run_id="run-open") + clear_task_span_if_delegating_task_ended(state, tool_name="task", run_id="run-other") + assert state.active_span_id is not None + assert state.active_task_run_id == "run-open" + + +def test_clear_on_matching_task_end() -> None: + state = AgentEventRelayState.for_invocation() + open_task_span(state, run_id="run-x") + clear_task_span_if_delegating_task_ended(state, tool_name="task", run_id="run-x") + assert state.active_span_id is None + assert state.active_task_run_id is None + assert state.span_metadata_if_active() is None + + +def test_clear_noop_when_no_open_span() -> None: + state = AgentEventRelayState.for_invocation() + clear_task_span_if_delegating_task_ended(state, tool_name="task", run_id="run-x") + assert state.active_span_id is None + + +def test_pending_then_open_reuses_same_span_id() -> None: + state = AgentEventRelayState.for_invocation() + sid_pending = ensure_pending_task_span_for_lc(state, "lc-task-1") + assert state.pending_task_span_by_lc["lc-task-1"] == sid_pending + sid_active = open_task_span( + state, run_id="run-1", langchain_tool_call_id="lc-task-1" + ) + assert sid_active == sid_pending + assert state.active_span_id == sid_pending + assert "lc-task-1" not in state.pending_task_span_by_lc From 695f9ded2c6401f23e386019eff346df39a3a6f9 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 22:47:08 +0200 Subject: [PATCH 63/96] Mint pending span id when the task tool registers from chunks. --- .../chat/streaming/handlers/chat_model_stream.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py b/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py index ef86dae56..c3f6d6d59 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py @@ -7,6 +7,7 @@ from typing import Any from app.tasks.chat.streaming.helpers.chunk_parts import extract_chunk_parts from app.tasks.chat.streaming.relay.state import AgentEventRelayState +from app.tasks.chat.streaming.relay.task_span import ensure_pending_task_span_for_lc from app.tasks.chat.streaming.relay.thinking_step_completion import ( complete_active_thinking_step, ) @@ -41,6 +42,7 @@ def iter_chat_model_stream_frames( state.current_text_id = None if state.current_reasoning_id is None: comp, new_active = complete_active_thinking_step( + state=state, streaming_service=streaming_service, content_builder=content_builder, last_active_step_id=state.last_active_step_id, @@ -76,6 +78,7 @@ def iter_chat_model_stream_frames( state.current_reasoning_id = None if state.current_text_id is None: comp, new_active = complete_active_thinking_step( + state=state, streaming_service=streaming_service, content_builder=content_builder, last_active_step_id=state.last_active_step_id, @@ -109,6 +112,10 @@ def iter_chat_model_stream_frames( name = tcc.get("name") if lc_id and name: ui_id = lc_id + tool_input_metadata: dict[str, Any] | None = None + if name == "task": + sid = ensure_pending_task_span_for_lc(state, str(lc_id)) + tool_input_metadata = {"spanId": sid} if state.current_text_id is not None: yield streaming_service.format_text_end(state.current_text_id) @@ -132,9 +139,12 @@ def iter_chat_model_stream_frames( ui_id, name, langchain_tool_call_id=lc_id, + metadata=tool_input_metadata, ) if content_builder is not None: - content_builder.on_tool_input_start(ui_id, name, lc_id) + content_builder.on_tool_input_start( + ui_id, name, lc_id, metadata=tool_input_metadata + ) meta = state.index_to_meta.get(idx) if idx is not None else None if meta: From 2c1b219c6cb32573e08bc0852dd2236fc3e325cc Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 22:47:32 +0200 Subject: [PATCH 64/96] Open task spans at tool start and tag unmatched tool-input SSE. --- .../chat/streaming/handlers/tool_start.py | 50 +++++++++++++------ 1 file changed, 35 insertions(+), 15 deletions(-) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py index e7d2d7f78..3b3537567 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py @@ -11,6 +11,7 @@ from app.tasks.chat.streaming.helpers.tool_call_matching import ( match_buffered_langchain_tool_call_id, ) from app.tasks.chat.streaming.relay.state import AgentEventRelayState +from app.tasks.chat.streaming.relay.task_span import open_task_span from app.tasks.chat.streaming.relay.thinking_step_completion import ( complete_active_thinking_step, ) @@ -46,6 +47,7 @@ def iter_tool_start_frames( if state.last_active_step_title != "Synthesizing response": comp, new_active = complete_active_thinking_step( + state=state, streaming_service=streaming_service, content_builder=content_builder, last_active_step_id=state.last_active_step_id, @@ -62,20 +64,6 @@ def iter_tool_start_frames( state.tool_step_ids[run_id] = tool_step_id state.last_active_step_id = tool_step_id - thinking = resolve_tool_start_thinking(tool_name, tool_input) - state.last_active_step_title = thinking.title - state.last_active_step_items = thinking.items - frame_kw: dict[str, Any] = { - "streaming_service": streaming_service, - "content_builder": content_builder, - "step_id": tool_step_id, - "title": thinking.title, - "status": "in_progress", - } - if thinking.include_items_on_frame: - frame_kw["items"] = thinking.items - yield emit_thinking_step_frame(**frame_kw) - matched_meta: dict[str, str] | None = None taken_ui_ids = set(state.ui_tool_call_id_by_run.values()) for meta in state.index_to_meta.values(): @@ -102,16 +90,46 @@ def iter_tool_start_frames( run_id, state.lc_tool_call_id_by_run, ) + + if tool_name == "task": + open_task_span( + state, + run_id=run_id, + langchain_tool_call_id=langchain_tool_call_id, + ) + + span_md = state.span_metadata_if_active() + + if matched_meta is None: yield streaming_service.format_tool_input_start( tool_call_id, tool_name, langchain_tool_call_id=langchain_tool_call_id, + metadata=span_md, ) if content_builder is not None: content_builder.on_tool_input_start( - tool_call_id, tool_name, langchain_tool_call_id + tool_call_id, + tool_name, + langchain_tool_call_id, + metadata=span_md, ) + thinking = resolve_tool_start_thinking(tool_name, tool_input) + state.last_active_step_title = thinking.title + state.last_active_step_items = thinking.items + frame_kw: dict[str, Any] = { + "streaming_service": streaming_service, + "content_builder": content_builder, + "step_id": tool_step_id, + "title": thinking.title, + "status": "in_progress", + "metadata": span_md, + } + if thinking.include_items_on_frame: + frame_kw["items"] = thinking.items + yield emit_thinking_step_frame(**frame_kw) + if run_id: state.ui_tool_call_id_by_run[run_id] = tool_call_id @@ -130,6 +148,7 @@ def iter_tool_start_frames( tool_name, _safe_input, langchain_tool_call_id=langchain_tool_call_id, + metadata=span_md, ) if content_builder is not None: content_builder.on_tool_input_available( @@ -137,4 +156,5 @@ def iter_tool_start_frames( tool_name, _safe_input, langchain_tool_call_id, + metadata=span_md, ) From 3ed09bdd90f114bfb370c614bf3e697244288965 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 22:47:38 +0200 Subject: [PATCH 65/96] Clear spans after task completion and pass span id on tool output. --- .../app/tasks/chat/streaming/handlers/tool_end.py | 7 +++++++ .../app/tasks/chat/streaming/handlers/tool_output_frame.py | 7 ++++++- .../chat/streaming/handlers/tools/emission_context.py | 2 ++ 3 files changed, 15 insertions(+), 1 deletion(-) 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, ) From 1dcb08e925f3d08bcd7936e6ce52b18efd9f9fd8 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 22:47:46 +0200 Subject: [PATCH 66/96] Attach active span metadata to thinking-step SSE and completion. --- .../app/tasks/chat/streaming/relay/event_relay.py | 1 + .../tasks/chat/streaming/relay/thinking_step_completion.py | 3 +++ .../app/tasks/chat/streaming/relay/thinking_step_sse.py | 6 +++++- 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py b/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py index 872998926..03d6a66e6 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py @@ -115,6 +115,7 @@ class EventRelay: state.current_text_id = None completion_event, new_active = complete_active_thinking_step( + state=state, streaming_service=self.streaming_service, content_builder=content_builder, last_active_step_id=state.last_active_step_id, diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py index a0be71281..ad0930341 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py @@ -4,11 +4,13 @@ from __future__ import annotations from typing import Any +from .state import AgentEventRelayState from .thinking_step_sse import emit_thinking_step_frame def complete_active_thinking_step( *, + state: AgentEventRelayState, streaming_service: Any, content_builder: Any | None, last_active_step_id: str | None, @@ -26,6 +28,7 @@ def complete_active_thinking_step( title=last_active_step_title, status="completed", items=last_active_step_items if last_active_step_items else None, + metadata=state.span_metadata_if_active(), ) return event, None return None, last_active_step_id diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py index 9e8c08dd5..6737f536b 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py @@ -13,12 +13,16 @@ def emit_thinking_step_frame( title: str, status: str = "in_progress", items: list[str] | None = None, + metadata: dict[str, Any] | None = None, ) -> str: if content_builder is not None: - content_builder.on_thinking_step(step_id, title, status, items) + content_builder.on_thinking_step( + step_id, title, status, items, metadata=metadata + ) return streaming_service.format_thinking_step( step_id=step_id, title=title, status=status, items=items, + metadata=metadata, ) From f1d80ffe5d9ca9005b4aa4769cd34a418a6fc345 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 22:47:50 +0200 Subject: [PATCH 67/96] Forward span metadata from report_progress thinking updates. --- .../app/tasks/chat/streaming/handlers/custom_event_dispatch.py | 1 + .../app/tasks/chat/streaming/handlers/custom_events.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/custom_event_dispatch.py b/surfsense_backend/app/tasks/chat/streaming/handlers/custom_event_dispatch.py index b373919cf..69f4b8a24 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/custom_event_dispatch.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/custom_event_dispatch.py @@ -33,6 +33,7 @@ def iter_custom_event_frames( last_active_step_items=state.last_active_step_items, streaming_service=streaming_service, content_builder=content_builder, + thinking_metadata=state.span_metadata_if_active(), ) if frame: yield frame diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py b/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py index 765f1d790..e48e2c493 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py @@ -15,6 +15,7 @@ def handle_report_progress( last_active_step_items: list[str], streaming_service: Any, content_builder: Any | None, + thinking_metadata: dict[str, Any] | None = None, ) -> tuple[str | None, list[str]]: """Update report step items; may emit one thinking SSE frame. @@ -50,6 +51,7 @@ def handle_report_progress( title=last_active_step_title, status="in_progress", items=new_items, + metadata=thinking_metadata, ) return frame, new_items From e802de233330c8aeb913e02ab922fa4a65dc7a00 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 22:47:58 +0200 Subject: [PATCH 68/96] Include optional metadata on tool and thinking-step SSE payloads. --- .../app/services/new_streaming_service.py | 29 +++++++++++++------ 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/surfsense_backend/app/services/new_streaming_service.py b/surfsense_backend/app/services/new_streaming_service.py index cec0c8a5e..ba0cb8753 100644 --- a/surfsense_backend/app/services/new_streaming_service.py +++ b/surfsense_backend/app/services/new_streaming_service.py @@ -456,6 +456,8 @@ class VercelStreamingService: title: str, status: str = "in_progress", items: list[str] | None = None, + *, + metadata: dict[str, Any] | None = None, ) -> str: """ Format a thinking step for chain-of-thought display (SurfSense specific). @@ -469,15 +471,15 @@ class VercelStreamingService: Returns: str: SSE formatted thinking step data part """ - return self.format_data( - "thinking-step", - { - "id": step_id, - "title": title, - "status": status, - "items": items or [], - }, - ) + payload: dict[str, Any] = { + "id": step_id, + "title": title, + "status": status, + "items": items or [], + } + if metadata: + payload["metadata"] = metadata + return self.format_data("thinking-step", payload) def format_thread_title_update(self, thread_id: int, title: str) -> str: """ @@ -601,6 +603,7 @@ class VercelStreamingService: tool_name: str, *, langchain_tool_call_id: str | None = None, + metadata: dict[str, Any] | None = None, ) -> str: """ Format the start of tool input streaming. @@ -635,6 +638,8 @@ class VercelStreamingService: } if langchain_tool_call_id: payload["langchainToolCallId"] = langchain_tool_call_id + if metadata: + payload["metadata"] = metadata return self._format_sse(payload) def format_tool_input_delta(self, tool_call_id: str, input_text_delta: str) -> str: @@ -666,6 +671,7 @@ class VercelStreamingService: input_data: dict[str, Any], *, langchain_tool_call_id: str | None = None, + metadata: dict[str, Any] | None = None, ) -> str: """ Format the completion of tool input. @@ -691,6 +697,8 @@ class VercelStreamingService: } if langchain_tool_call_id: payload["langchainToolCallId"] = langchain_tool_call_id + if metadata: + payload["metadata"] = metadata return self._format_sse(payload) def format_tool_output_available( @@ -699,6 +707,7 @@ class VercelStreamingService: output: Any, *, langchain_tool_call_id: str | None = None, + metadata: dict[str, Any] | None = None, ) -> str: """ Format tool execution output. @@ -725,6 +734,8 @@ class VercelStreamingService: } if langchain_tool_call_id: payload["langchainToolCallId"] = langchain_tool_call_id + if metadata: + payload["metadata"] = metadata return self._format_sse(payload) # ========================================================================= From 3dbcac4b9d656af92fd7fe376386124411fd53e5 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 22:48:07 +0200 Subject: [PATCH 69/96] Merge span metadata into persisted tool-call and thinking parts. --- .../app/tasks/chat/content_builder.py | 38 ++++++++++++++++--- 1 file changed, 32 insertions(+), 6 deletions(-) diff --git a/surfsense_backend/app/tasks/chat/content_builder.py b/surfsense_backend/app/tasks/chat/content_builder.py index 32b49e6b5..2b6dbe649 100644 --- a/surfsense_backend/app/tasks/chat/content_builder.py +++ b/surfsense_backend/app/tasks/chat/content_builder.py @@ -51,6 +51,15 @@ logger = logging.getLogger(__name__) _MEANINGFUL_PART_TYPES: frozenset[str] = frozenset({"text", "reasoning", "tool-call"}) +def _merge_tool_part_metadata(part: dict[str, Any], metadata: dict[str, Any] | None) -> None: + if not metadata: + return + md = part.setdefault("metadata", {}) + for k, v in metadata.items(): + if k not in md: + md[k] = v + + class AssistantContentBuilder: """Server-side projection of ``surfsense_web/lib/chat/streaming-state.ts``. @@ -177,6 +186,8 @@ class AssistantContentBuilder: ui_id: str, tool_name: str, langchain_tool_call_id: str | None, + *, + metadata: dict[str, Any] | None = None, ) -> None: """Register a tool-call card. Args are filled in by later events.""" if not ui_id: @@ -187,11 +198,11 @@ class AssistantContentBuilder: # (the canonical path). The FE de-dupes via ``toolCallIndices``; # we mirror that here. if ui_id in self._tool_call_idx_by_ui_id: - if langchain_tool_call_id: - idx = self._tool_call_idx_by_ui_id[ui_id] - part = self.parts[idx] - if not part.get("langchainToolCallId"): - part["langchainToolCallId"] = langchain_tool_call_id + idx = self._tool_call_idx_by_ui_id[ui_id] + part = self.parts[idx] + if langchain_tool_call_id and not part.get("langchainToolCallId"): + part["langchainToolCallId"] = langchain_tool_call_id + _merge_tool_part_metadata(part, metadata) return part: dict[str, Any] = { @@ -202,6 +213,8 @@ class AssistantContentBuilder: } if langchain_tool_call_id: part["langchainToolCallId"] = langchain_tool_call_id + if metadata: + part["metadata"] = dict(metadata) self.parts.append(part) self._tool_call_idx_by_ui_id[ui_id] = len(self.parts) - 1 @@ -235,6 +248,8 @@ class AssistantContentBuilder: tool_name: str, args: dict[str, Any], langchain_tool_call_id: str | None, + *, + metadata: dict[str, Any] | None = None, ) -> None: """Finalize the tool-call card's input. @@ -264,6 +279,7 @@ class AssistantContentBuilder: part["argsText"] = final_args_text if langchain_tool_call_id and not part.get("langchainToolCallId"): part["langchainToolCallId"] = langchain_tool_call_id + _merge_tool_part_metadata(part, metadata) return # No prior tool-input-start: register the card now. @@ -276,6 +292,7 @@ class AssistantContentBuilder: } if langchain_tool_call_id: new_part["langchainToolCallId"] = langchain_tool_call_id + _merge_tool_part_metadata(new_part, metadata) self.parts.append(new_part) self._tool_call_idx_by_ui_id[ui_id] = len(self.parts) - 1 @@ -287,6 +304,8 @@ class AssistantContentBuilder: ui_id: str, output: Any, langchain_tool_call_id: str | None, + *, + metadata: dict[str, Any] | None = None, ) -> None: """Attach the tool's output (``result``) to the matching card. @@ -305,6 +324,7 @@ class AssistantContentBuilder: part["result"] = output if langchain_tool_call_id and not part.get("langchainToolCallId"): part["langchainToolCallId"] = langchain_tool_call_id + _merge_tool_part_metadata(part, metadata) # ------------------------------------------------------------------ # Thinking steps & step separators @@ -316,6 +336,8 @@ class AssistantContentBuilder: title: str, status: str, items: list[str] | None, + *, + metadata: dict[str, Any] | None = None, ) -> None: """Update / insert the singleton ``data-thinking-steps`` part. @@ -328,12 +350,14 @@ class AssistantContentBuilder: if not step_id: return - new_step = { + new_step: dict[str, Any] = { "id": step_id, "title": title or "", "status": status or "in_progress", "items": list(items) if items else [], } + if metadata: + new_step["metadata"] = dict(metadata) # Find existing data-thinking-steps part. existing_idx = -1 @@ -347,6 +371,8 @@ class AssistantContentBuilder: replaced = False for i, step in enumerate(current_steps): if step.get("id") == step_id: + if not metadata and step.get("metadata"): + new_step["metadata"] = dict(step["metadata"]) current_steps[i] = new_step replaced = True break From d136fcd054d7bbd4d51ecd9c4e5368037cc43e0e Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 23:16:44 +0200 Subject: [PATCH 70/96] Add tool_activity_metadata to merge spanId and thinkingStepId for tools. --- .../app/tasks/chat/streaming/relay/state.py | 37 +++++++++++++++++-- 1 file changed, 33 insertions(+), 4 deletions(-) diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/state.py b/surfsense_backend/app/tasks/chat/streaming/relay/state.py index 82525a52f..27898403d 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/state.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/state.py @@ -10,10 +10,16 @@ from typing import Any class AgentEventRelayState: """Tracks text, thinking steps, tool depth, and pending tool-call metadata. - ``active_span_id`` groups steps/tools for one open ``task`` episode. - ``active_task_run_id`` is the LangGraph ``run_id`` of that ``task`` so we - only clear the span when that run ends (not when child tools end). Handlers - will set/clear these via ``task_span`` helpers in a later change. + **Task span (`spanId`)** — ``active_span_id`` groups steps and tools for one + open delegating ``task`` episode. ``active_task_run_id`` is the LangGraph + ``run_id`` of that ``task`` so the span clears only when that run ends, not + when child tools end. Open/close uses ``relay.task_span`` helpers. + + **Tool ↔ thinking link (`thinkingStepId`)** — Each tool run gets a thinking-row + id (``tool_step_ids[run_id]``, emitted as ``data-thinking-step`` ``data.id``). + ``tool_activity_metadata`` supplies ``metadata`` for ``tool-input-start`` / + ``tool-input-available`` (``handlers.tool_start``) and + ``tool-output-available`` (``handlers.tool_end``). """ accumulated_text: str = "" @@ -48,6 +54,29 @@ class AgentEventRelayState: return {"spanId": self.active_span_id} return None + def tool_activity_metadata( + self, *, thinking_step_id: str | None + ) -> dict[str, Any] | None: + """Build ``metadata`` for tool SSE and ``tool-call`` persistence. + + Contract (keys omitted when not applicable): + + - ``spanId`` (str): present while a task-delegation span is active + (same value as ``span_metadata_if_active()``). + - ``thinkingStepId`` (str): equals the thinking-step row ``id`` for this + tool (``data-thinking-step`` payload ``data.id`` on the wire). + + Returns ``None`` if neither applies. Whitespace-only + ``thinking_step_id`` is ignored. + """ + out: dict[str, Any] = {} + if self.active_span_id: + out["spanId"] = self.active_span_id + tid = (thinking_step_id or "").strip() + if tid: + out["thinkingStepId"] = tid + return out if out else None + @classmethod def for_invocation( cls, From 007a0a30ec7c2850316f5503729a5f8c67312074 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 23:16:56 +0200 Subject: [PATCH 71/96] Cover tool_activity_metadata for span-only, step-only, and combined cases. --- .../streaming/test_tool_activity_metadata.py | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/test_tool_activity_metadata.py diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_tool_activity_metadata.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_tool_activity_metadata.py new file mode 100644 index 000000000..c2e68dacd --- /dev/null +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_tool_activity_metadata.py @@ -0,0 +1,42 @@ +"""Unit tests for ``AgentEventRelayState.tool_activity_metadata``.""" + +from __future__ import annotations + +import pytest + +from app.tasks.chat.streaming.relay.state import AgentEventRelayState +from app.tasks.chat.streaming.relay.task_span import open_task_span + +pytestmark = pytest.mark.unit + + +def test_returns_none_when_no_span_and_no_thinking_step() -> None: + state = AgentEventRelayState.for_invocation() + assert state.tool_activity_metadata(thinking_step_id=None) is None + assert state.tool_activity_metadata(thinking_step_id="") is None + assert state.tool_activity_metadata(thinking_step_id=" ") is None + + +def test_thinking_step_id_only() -> None: + state = AgentEventRelayState.for_invocation() + assert state.tool_activity_metadata(thinking_step_id="thinking-3") == { + "thinkingStepId": "thinking-3", + } + + +def test_span_only_when_active() -> None: + state = AgentEventRelayState.for_invocation() + open_task_span(state, run_id="run-x") + assert state.tool_activity_metadata(thinking_step_id=None) == { + "spanId": state.active_span_id, + } + + +def test_merges_span_and_thinking_step_when_both_set() -> None: + state = AgentEventRelayState.for_invocation() + open_task_span(state, run_id="run-x") + md = state.tool_activity_metadata(thinking_step_id="thinking-7") + assert md == { + "spanId": state.active_span_id, + "thinkingStepId": "thinking-7", + } From a309e830d34dedfb07e3a263bc7e34ec30fa5335 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 23:17:01 +0200 Subject: [PATCH 72/96] Document thinkingStepId on tool-call parts and first-key metadata merge. --- surfsense_backend/app/tasks/chat/content_builder.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/surfsense_backend/app/tasks/chat/content_builder.py b/surfsense_backend/app/tasks/chat/content_builder.py index 2b6dbe649..f0804159a 100644 --- a/surfsense_backend/app/tasks/chat/content_builder.py +++ b/surfsense_backend/app/tasks/chat/content_builder.py @@ -52,6 +52,12 @@ _MEANINGFUL_PART_TYPES: frozenset[str] = frozenset({"text", "reasoning", "tool-c def _merge_tool_part_metadata(part: dict[str, Any], metadata: dict[str, Any] | None) -> None: + """Shallow-merge ``metadata`` into ``part["metadata"]``; first key wins. + + Used for tool-call linkage (``spanId``, ``thinkingStepId``, …): a later + event must not overwrite an existing key so chunk order vs ``on_tool_start`` + stays stable. + """ if not metadata: return md = part.setdefault("metadata", {}) @@ -70,6 +76,7 @@ class AssistantContentBuilder: | { type: "reasoning"; text: string } | { type: "tool-call"; toolCallId: str; toolName: str; args: dict; result?: any; argsText?: str; langchainToolCallId?: str; + metadata?: { spanId?: str; thinkingStepId?: str; ... }; state?: "aborted" } | { type: "data-thinking-steps"; data: { steps: ThinkingStepData[] } } | { type: "data-step-separator"; data: { stepIndex: int } } @@ -189,7 +196,11 @@ class AssistantContentBuilder: *, metadata: dict[str, Any] | None = None, ) -> None: - """Register a tool-call card. Args are filled in by later events.""" + """Register a tool-call card. Args are filled in by later events. + + Optional ``metadata`` (``spanId``, ``thinkingStepId``, …) is stored on the + part; duplicate ``tool-input-start`` calls merge with first-key-wins. + """ if not ui_id: return # Skip duplicate registration: the stream may emit From 32092c0b65914a046a0f6a02e76bfa6a561af71b Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 23:17:05 +0200 Subject: [PATCH 73/96] Pass thinkingStepId through tool-input start and available metadata. --- .../app/tasks/chat/streaming/handlers/tool_start.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py index 3b3537567..e0cac307c 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py @@ -99,20 +99,21 @@ def iter_tool_start_frames( ) span_md = state.span_metadata_if_active() + tool_md = state.tool_activity_metadata(thinking_step_id=tool_step_id) if matched_meta is None: yield streaming_service.format_tool_input_start( tool_call_id, tool_name, langchain_tool_call_id=langchain_tool_call_id, - metadata=span_md, + metadata=tool_md, ) if content_builder is not None: content_builder.on_tool_input_start( tool_call_id, tool_name, langchain_tool_call_id, - metadata=span_md, + metadata=tool_md, ) thinking = resolve_tool_start_thinking(tool_name, tool_input) @@ -148,7 +149,7 @@ def iter_tool_start_frames( tool_name, _safe_input, langchain_tool_call_id=langchain_tool_call_id, - metadata=span_md, + metadata=tool_md, ) if content_builder is not None: content_builder.on_tool_input_available( @@ -156,5 +157,5 @@ def iter_tool_start_frames( tool_name, _safe_input, langchain_tool_call_id, - metadata=span_md, + metadata=tool_md, ) From 1761b60c16397310145cb8739b93cc105c5576e5 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 8 May 2026 23:17:12 +0200 Subject: [PATCH 74/96] Carry thinkingStepId on tool output and extend builder and parity tests. --- .../tasks/chat/streaming/handlers/tool_end.py | 4 +- .../chat/streaming/test_stage_2_parity.py | 3 + .../unit/tasks/chat/test_content_builder.py | 150 ++++++++++++++++++ 3 files changed, 156 insertions(+), 1 deletion(-) 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 ec7d6551c..421c67a6d 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py @@ -110,7 +110,9 @@ 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(), + tool_metadata=state.tool_activity_metadata( + thinking_step_id=original_step_id, + ), ) yield from iter_tool_completion_emission_frames(emission_ctx) diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py index 9ae7defec..3ee1ab622 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stage_2_parity.py @@ -111,8 +111,10 @@ def test_complete_active_thinking_step_mirrors_closure_semantics() -> None: svc = MagicMock() svc.format_thinking_step.return_value = "done-frame" completed: set[str] = set() + relay_state = AgentEventRelayState.for_invocation() frame, new_id = complete_active_thinking_step( + state=relay_state, streaming_service=svc, content_builder=None, last_active_step_id="thinking-1", @@ -125,6 +127,7 @@ def test_complete_active_thinking_step_mirrors_closure_semantics() -> None: assert "thinking-1" in completed frame2, id2 = complete_active_thinking_step( + state=relay_state, streaming_service=svc, content_builder=None, last_active_step_id="thinking-1", diff --git a/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py b/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py index 4b1fadd9c..9d3eb6fa4 100644 --- a/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py +++ b/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py @@ -15,6 +15,7 @@ import json import pytest +from app.services.new_streaming_service import VercelStreamingService from app.tasks.chat.content_builder import AssistantContentBuilder pytestmark = pytest.mark.unit @@ -231,6 +232,155 @@ class TestToolHeavyTurn: ) +# --------------------------------------------------------------------------- +# Task-span metadata on tool-call parts (JSONB persistence) +# --------------------------------------------------------------------------- + + +class TestToolCallSpanMetadata: + def test_input_available_merges_new_metadata_keys_after_start(self): + b = AssistantContentBuilder() + b.on_tool_input_start( + "call_t", "task", "lc_t", metadata={"spanId": "spn_1"} + ) + b.on_tool_input_available( + "call_t", + "task", + {"goal": "x"}, + "lc_t", + metadata={"traceId": "tr_1"}, + ) + part = b.snapshot()[0] + assert part["metadata"]["spanId"] == "spn_1" + assert part["metadata"]["traceId"] == "tr_1" + _assert_jsonb_safe(b.snapshot()) + + def test_input_available_does_not_overwrite_existing_metadata_keys(self): + b = AssistantContentBuilder() + b.on_tool_input_start( + "call_t", "task", "lc_t", metadata={"spanId": "spn_keep"} + ) + b.on_tool_input_available( + "call_t", "task", {}, "lc_t", metadata={"spanId": "spn_other"} + ) + assert b.snapshot()[0]["metadata"]["spanId"] == "spn_keep" + + def test_late_tool_input_available_carries_metadata(self): + b = AssistantContentBuilder() + b.on_tool_input_available( + "call_l", + "grep", + {"pattern": "TODO"}, + None, + metadata={"spanId": "spn_l"}, + ) + part = b.snapshot()[0] + assert part["metadata"] == {"spanId": "spn_l"} + _assert_jsonb_safe(b.snapshot()) + + def test_output_available_merges_without_clobbering_span_id(self): + b = AssistantContentBuilder() + b.on_tool_input_start("call_t", "ls", "lc", metadata={"spanId": "spn_x"}) + b.on_tool_input_available("call_t", "ls", {"path": "/"}, "lc") + b.on_tool_output_available( + "call_t", + {"ok": True}, + "lc", + metadata={"spanId": "spn_y", "extra": 1}, + ) + md = b.snapshot()[0]["metadata"] + assert md["spanId"] == "spn_x" + assert md["extra"] == 1 + + def test_output_available_adds_thinking_step_id_without_clobbering_span(self): + b = AssistantContentBuilder() + b.on_tool_input_start( + "call_t", + "ls", + "lc", + metadata={"spanId": "spn_x", "thinkingStepId": "thinking-3"}, + ) + b.on_tool_input_available("call_t", "ls", {"path": "/"}, "lc") + b.on_tool_output_available( + "call_t", + {"ok": True}, + "lc", + metadata={"spanId": "spn_x", "thinkingStepId": "thinking-3"}, + ) + md = b.snapshot()[0]["metadata"] + assert md["spanId"] == "spn_x" + assert md["thinkingStepId"] == "thinking-3" + + def test_output_available_with_none_metadata_preserves_prior(self): + b = AssistantContentBuilder() + b.on_tool_input_start("c", "ls", "lc", metadata={"spanId": "spn_1"}) + b.on_tool_input_available("c", "ls", {}, "lc") + b.on_tool_output_available("c", {"r": 1}, "lc", metadata=None) + assert b.snapshot()[0]["metadata"] == {"spanId": "spn_1"} + + def test_available_adds_thinking_step_id_after_chunk_only_start(self): + """Mirrors chunk ``tool-input-start`` then ``on_tool_start`` ``available``.""" + b = AssistantContentBuilder() + b.on_tool_input_start("lc_1", "ls", "lc_1", metadata={"spanId": "spn_a"}) + b.on_tool_input_available( + "lc_1", + "ls", + {"path": "/"}, + "lc_1", + metadata={"spanId": "spn_a", "thinkingStepId": "thinking-2"}, + ) + md = b.snapshot()[0]["metadata"] + assert md["spanId"] == "spn_a" + assert md["thinkingStepId"] == "thinking-2" + + +class TestVercelStreamingServiceToolMetadataWire: + """SSE payloads include optional ``metadata`` for FE grouping.""" + + @staticmethod + def _parse_sse_data_line(raw: str) -> dict: + assert raw.startswith("data: ") + payload = raw.split("data: ", 1)[1].split("\n\n", 1)[0].strip() + return json.loads(payload) + + def test_tool_input_available_includes_metadata_when_set(self): + svc = VercelStreamingService() + raw = svc.format_tool_input_available( + "id1", + "task", + {"a": 1}, + langchain_tool_call_id="lc1", + metadata={"spanId": "spn_w", "thinkingStepId": "thinking-4"}, + ) + body = self._parse_sse_data_line(raw) + assert body["type"] == "tool-input-available" + assert body["metadata"] == { + "spanId": "spn_w", + "thinkingStepId": "thinking-4", + } + + def test_tool_output_available_includes_metadata_when_set(self): + svc = VercelStreamingService() + raw = svc.format_tool_output_available( + "id1", + {"status": "completed"}, + langchain_tool_call_id="lc1", + metadata={"spanId": "spn_o", "thinkingStepId": "thinking-9"}, + ) + body = self._parse_sse_data_line(raw) + assert body["type"] == "tool-output-available" + assert body["metadata"] == { + "spanId": "spn_o", + "thinkingStepId": "thinking-9", + } + + def test_tool_input_available_omits_metadata_key_when_none(self): + svc = VercelStreamingService() + raw = svc.format_tool_input_available("id1", "ls", {}) + body = self._parse_sse_data_line(raw) + assert "metadata" not in body + + # --------------------------------------------------------------------------- # Thinking steps & separators # --------------------------------------------------------------------------- From 47e64d1861a85636c8e8bde2268d6f34bf6bee2d Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Sat, 9 May 2026 00:39:43 +0200 Subject: [PATCH 75/96] Persist optional relay metadata on tool-call and thinking-step state. --- surfsense_web/lib/chat/streaming-state.ts | 43 ++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/surfsense_web/lib/chat/streaming-state.ts b/surfsense_web/lib/chat/streaming-state.ts index 809e214d1..ee3160a61 100644 --- a/surfsense_web/lib/chat/streaming-state.ts +++ b/surfsense_web/lib/chat/streaming-state.ts @@ -5,6 +5,11 @@ export interface ThinkingStepData { title: string; status: "pending" | "in_progress" | "completed"; items: string[]; + /** + * Optional relay fields from ``data-thinking-step`` when present on the wire + * (e.g. ``spanId``). Populated in a later slice; equality helpers ignore until wired. + */ + metadata?: Record; } export type ContentPart = @@ -42,6 +47,11 @@ export type ContentPart = * ``data-action-log`` events. */ langchainToolCallId?: string; + /** + * Relay correlation from tool SSE (e.g. ``spanId``, ``thinkingStepId``). + * Merged by ``mergeToolPartMetadata`` when events carry ``metadata``. + */ + metadata?: Record; } | { type: "data-thinking-steps"; @@ -252,6 +262,23 @@ function _toolPasses(gate: ToolUIGate, toolName: string): boolean { return gate === "all" || gate.has(toolName); } +/** + * Shallow-merge relay ``metadata`` into a tool-call part (SSE → content part). + * Keys already set on ``into`` are left unchanged so chunk vs canonical tool + * events cannot reorder or overwrite ``spanId`` / ``thinkingStepId``. + * Matches server ``AssistantContentBuilder`` merge semantics. + */ +function mergeToolPartMetadata( + into: Record, + incoming: Record | undefined +): void { + if (!incoming) return; + for (const [k, v] of Object.entries(incoming)) { + if (k === "__proto__" || k === "constructor") continue; + if (!(k in into)) into[k] = v; + } +} + export function addToolCall( state: ContentPartsState, toolsWithUI: ToolUIGate, @@ -259,15 +286,19 @@ export function addToolCall( toolName: string, args: Record, force = false, - langchainToolCallId?: string + langchainToolCallId?: string, + metadata?: Record ): void { if (force || _toolPasses(toolsWithUI, toolName)) { + const relayMeta: Record = {}; + mergeToolPartMetadata(relayMeta, metadata); state.contentParts.push({ type: "tool-call", toolCallId, toolName, args, ...(langchainToolCallId ? { langchainToolCallId } : {}), + ...(Object.keys(relayMeta).length > 0 ? { metadata: relayMeta } : {}), }); state.toolCallIndices.set(toolCallId, state.contentParts.length - 1); state.currentTextPartIndex = -1; @@ -304,6 +335,7 @@ export function updateToolCall( argsText?: string; result?: unknown; langchainToolCallId?: string; + metadata?: Record; } ): void { const index = state.toolCallIndices.get(toolCallId); @@ -323,6 +355,11 @@ export function updateToolCall( if (update.langchainToolCallId && !tc.langchainToolCallId) { tc.langchainToolCallId = update.langchainToolCallId; } + if (update.metadata && Object.keys(update.metadata).length > 0) { + const md = (tc.metadata ?? {}) as Record; + mergeToolPartMetadata(md, update.metadata); + tc.metadata = md; + } } } @@ -416,6 +453,8 @@ export type SSEEvent = toolName: string; /** Authoritative LangChain ``tool_call.id``. Optional. */ langchainToolCallId?: string; + /** Optional JSON object from tool SSE (same keys as persisted tool-call metadata). */ + metadata?: Record; } | { /** @@ -434,6 +473,7 @@ export type SSEEvent = toolName: string; input: Record; langchainToolCallId?: string; + metadata?: Record; } | { type: "tool-output-available"; @@ -443,6 +483,7 @@ export type SSEEvent = * ``ToolMessage.tool_call_id`` at on_tool_end. Backfills cards * that didn't get the id at tool-input-start time. */ langchainToolCallId?: string; + metadata?: Record; } | { type: "data-thinking-step"; data: ThinkingStepData } | { type: "data-thread-title-update"; data: { threadId: number; title: string } } From 39084b3075ef0dd322d2cfab037fe1b0fbd4dc0f Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Sat, 9 May 2026 00:39:50 +0200 Subject: [PATCH 76/96] Forward tool SSE metadata into the streaming content parts. --- surfsense_web/components/free-chat/free-chat-page.tsx | 8 ++++++-- surfsense_web/lib/chat/stream-pipeline.ts | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/surfsense_web/components/free-chat/free-chat-page.tsx b/surfsense_web/components/free-chat/free-chat-page.tsx index 080d9a2b6..b2f959f2e 100644 --- a/surfsense_web/components/free-chat/free-chat-page.tsx +++ b/surfsense_web/components/free-chat/free-chat-page.tsx @@ -228,7 +228,8 @@ export function FreeChatPage() { parsed.toolName, {}, false, - parsed.langchainToolCallId + parsed.langchainToolCallId, + parsed.metadata ); forceFlush(); break; @@ -245,6 +246,7 @@ export function FreeChatPage() { args: parsed.input || {}, argsText: finalArgsText, langchainToolCallId: parsed.langchainToolCallId, + metadata: parsed.metadata, }); } else { addToolCall( @@ -254,7 +256,8 @@ export function FreeChatPage() { parsed.toolName, parsed.input || {}, false, - parsed.langchainToolCallId + parsed.langchainToolCallId, + parsed.metadata ); updateToolCall(contentPartsState, parsed.toolCallId, { argsText: finalArgsText, @@ -268,6 +271,7 @@ export function FreeChatPage() { updateToolCall(contentPartsState, parsed.toolCallId, { result: parsed.output, langchainToolCallId: parsed.langchainToolCallId, + metadata: parsed.metadata, }); forceFlush(); break; diff --git a/surfsense_web/lib/chat/stream-pipeline.ts b/surfsense_web/lib/chat/stream-pipeline.ts index c76781083..a0f90afeb 100644 --- a/surfsense_web/lib/chat/stream-pipeline.ts +++ b/surfsense_web/lib/chat/stream-pipeline.ts @@ -112,7 +112,8 @@ export function processSharedStreamEvent( parsed.toolName, {}, false, - parsed.langchainToolCallId + parsed.langchainToolCallId, + parsed.metadata ); forceFlush(); return true; @@ -131,6 +132,7 @@ export function processSharedStreamEvent( args: parsed.input || {}, argsText: finalArgsText, langchainToolCallId: parsed.langchainToolCallId, + metadata: parsed.metadata, }); } else { addToolCall( @@ -140,7 +142,8 @@ export function processSharedStreamEvent( parsed.toolName, parsed.input || {}, false, - parsed.langchainToolCallId + parsed.langchainToolCallId, + parsed.metadata ); // addToolCall doesn't accept argsText today; backfill via // updateToolCall so the new card renders pretty-printed JSON. @@ -156,6 +159,7 @@ export function processSharedStreamEvent( updateToolCall(contentPartsState, parsed.toolCallId, { result: parsed.output, langchainToolCallId: parsed.langchainToolCallId, + metadata: parsed.metadata, }); markInterruptsCompleted(contentParts); context.onToolOutputAvailable?.(parsed, { contentPartsState, toolCallIndices }); From e7c5204b0248d7ceb6262478a4419eb7e0fe6f58 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Sat, 9 May 2026 00:39:59 +0200 Subject: [PATCH 77/96] Indent tool cards under an active delegating task span. --- .../assistant-ui/assistant-message.tsx | 88 ++++++++++--------- .../components/assistant-ui/tool-fallback.tsx | 22 +++++ .../components/public-chat/public-thread.tsx | 29 +++--- .../lib/chat/delegation-span-indent.ts | 19 ++++ 4 files changed, 106 insertions(+), 52 deletions(-) create mode 100644 surfsense_web/lib/chat/delegation-span-indent.ts diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx index 7bccc22ee..a21ade74a 100644 --- a/surfsense_web/components/assistant-ui/assistant-message.tsx +++ b/surfsense_web/components/assistant-ui/assistant-message.tsx @@ -4,6 +4,7 @@ import { AuiIf, ErrorPrimitive, MessagePrimitive, + type ToolCallMessagePartComponent, useAui, useAuiState, } from "@assistant-ui/react"; @@ -36,7 +37,7 @@ import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { ReasoningMessagePart } from "@/components/assistant-ui/reasoning-message-part"; import { RevertTurnButton } from "@/components/assistant-ui/revert-turn-button"; import { useTokenUsage } from "@/components/assistant-ui/token-usage-context"; -import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; +import { ToolFallback, withDelegationSpanIndent } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { CommentPanelContainer } from "@/components/chat-comments/comment-panel-container/comment-panel-container"; import { CommentSheet } from "@/components/chat-comments/comment-sheet/comment-sheet"; @@ -505,48 +506,55 @@ const MessageInfoDropdown: FC = () => { // Wrap each tool-ui card with ``withBundleStep`` so multi-card HITL bundles // page through them and stage decisions instead of firing one resume per card. +// ``withDelegationSpanIndent`` wraps every entry (including Fallback) so delegated +// subagent tools don't bypass span indentation via a named ``by_name`` UI. +const bundleTool = (Component: ToolCallMessagePartComponent) => + withBundleStep(withDelegationSpanIndent(Component)); + +const NullToolUi: ToolCallMessagePartComponent = () => null; + const TOOLS_BY_NAME = { - generate_report: withBundleStep(GenerateReportToolUI), - generate_resume: withBundleStep(GenerateResumeToolUI), - generate_podcast: withBundleStep(GeneratePodcastToolUI), - generate_video_presentation: withBundleStep(GenerateVideoPresentationToolUI), - display_image: withBundleStep(GenerateImageToolUI), - generate_image: withBundleStep(GenerateImageToolUI), - update_memory: withBundleStep(UpdateMemoryToolUI), - execute: withBundleStep(SandboxExecuteToolUI), - execute_code: withBundleStep(SandboxExecuteToolUI), - create_notion_page: withBundleStep(CreateNotionPageToolUI), - update_notion_page: withBundleStep(UpdateNotionPageToolUI), - delete_notion_page: withBundleStep(DeleteNotionPageToolUI), - create_linear_issue: withBundleStep(CreateLinearIssueToolUI), - update_linear_issue: withBundleStep(UpdateLinearIssueToolUI), - delete_linear_issue: withBundleStep(DeleteLinearIssueToolUI), - create_google_drive_file: withBundleStep(CreateGoogleDriveFileToolUI), - delete_google_drive_file: withBundleStep(DeleteGoogleDriveFileToolUI), - create_onedrive_file: withBundleStep(CreateOneDriveFileToolUI), - delete_onedrive_file: withBundleStep(DeleteOneDriveFileToolUI), - create_dropbox_file: withBundleStep(CreateDropboxFileToolUI), - delete_dropbox_file: withBundleStep(DeleteDropboxFileToolUI), - create_calendar_event: withBundleStep(CreateCalendarEventToolUI), - update_calendar_event: withBundleStep(UpdateCalendarEventToolUI), - delete_calendar_event: withBundleStep(DeleteCalendarEventToolUI), - create_gmail_draft: withBundleStep(CreateGmailDraftToolUI), - update_gmail_draft: withBundleStep(UpdateGmailDraftToolUI), - send_gmail_email: withBundleStep(SendGmailEmailToolUI), - trash_gmail_email: withBundleStep(TrashGmailEmailToolUI), - create_jira_issue: withBundleStep(CreateJiraIssueToolUI), - update_jira_issue: withBundleStep(UpdateJiraIssueToolUI), - delete_jira_issue: withBundleStep(DeleteJiraIssueToolUI), - create_confluence_page: withBundleStep(CreateConfluencePageToolUI), - update_confluence_page: withBundleStep(UpdateConfluencePageToolUI), - delete_confluence_page: withBundleStep(DeleteConfluencePageToolUI), - web_search: () => null, - link_preview: () => null, - multi_link_preview: () => null, - scrape_webpage: () => null, + generate_report: bundleTool(GenerateReportToolUI), + generate_resume: bundleTool(GenerateResumeToolUI), + generate_podcast: bundleTool(GeneratePodcastToolUI), + generate_video_presentation: bundleTool(GenerateVideoPresentationToolUI), + display_image: bundleTool(GenerateImageToolUI), + generate_image: bundleTool(GenerateImageToolUI), + update_memory: bundleTool(UpdateMemoryToolUI), + execute: bundleTool(SandboxExecuteToolUI), + execute_code: bundleTool(SandboxExecuteToolUI), + create_notion_page: bundleTool(CreateNotionPageToolUI), + update_notion_page: bundleTool(UpdateNotionPageToolUI), + delete_notion_page: bundleTool(DeleteNotionPageToolUI), + create_linear_issue: bundleTool(CreateLinearIssueToolUI), + update_linear_issue: bundleTool(UpdateLinearIssueToolUI), + delete_linear_issue: bundleTool(DeleteLinearIssueToolUI), + create_google_drive_file: bundleTool(CreateGoogleDriveFileToolUI), + delete_google_drive_file: bundleTool(DeleteGoogleDriveFileToolUI), + create_onedrive_file: bundleTool(CreateOneDriveFileToolUI), + delete_onedrive_file: bundleTool(DeleteOneDriveFileToolUI), + create_dropbox_file: bundleTool(CreateDropboxFileToolUI), + delete_dropbox_file: bundleTool(DeleteDropboxFileToolUI), + create_calendar_event: bundleTool(CreateCalendarEventToolUI), + update_calendar_event: bundleTool(UpdateCalendarEventToolUI), + delete_calendar_event: bundleTool(DeleteCalendarEventToolUI), + create_gmail_draft: bundleTool(CreateGmailDraftToolUI), + update_gmail_draft: bundleTool(UpdateGmailDraftToolUI), + send_gmail_email: bundleTool(SendGmailEmailToolUI), + trash_gmail_email: bundleTool(TrashGmailEmailToolUI), + create_jira_issue: bundleTool(CreateJiraIssueToolUI), + update_jira_issue: bundleTool(UpdateJiraIssueToolUI), + delete_jira_issue: bundleTool(DeleteJiraIssueToolUI), + create_confluence_page: bundleTool(CreateConfluencePageToolUI), + update_confluence_page: bundleTool(UpdateConfluencePageToolUI), + delete_confluence_page: bundleTool(DeleteConfluencePageToolUI), + web_search: NullToolUi, + link_preview: NullToolUi, + multi_link_preview: NullToolUi, + scrape_webpage: NullToolUi, } as const; -const TOOLS_FALLBACK = withBundleStep(ToolFallback); +const TOOLS_FALLBACK = bundleTool(ToolFallback); const AssistantMessageInner: FC = () => { const isMobile = !useMediaQuery("(min-width: 768px)"); diff --git a/surfsense_web/components/assistant-ui/tool-fallback.tsx b/surfsense_web/components/assistant-ui/tool-fallback.tsx index ba58f4158..ec93b1018 100644 --- a/surfsense_web/components/assistant-ui/tool-fallback.tsx +++ b/surfsense_web/components/assistant-ui/tool-fallback.tsx @@ -31,6 +31,10 @@ import { Spinner } from "@/components/ui/spinner"; import { getToolDisplayName } from "@/contracts/enums/toolIcons"; import { markActionRevertedInCache, useAgentActionsQuery } from "@/hooks/use-agent-actions-query"; import { agentActionsApiService } from "@/lib/apis/agent-actions-api.service"; +import { + DELEGATION_SPAN_INDENT_CLASS, + shouldIndentToolCallForDelegationSpan, +} from "@/lib/chat/delegation-span-indent"; import { AppError } from "@/lib/error"; import { isInterruptResult } from "@/lib/hitl"; import { cn } from "@/lib/utils"; @@ -499,6 +503,24 @@ const DefaultToolFallbackInner: ToolCallMessagePartComponent = (props) => { ); }; +/** + * Wrap any tool-call UI so cards under an active delegating ``task`` span indent. + * Applied to named tool components as well as ``ToolFallback`` — only ``ToolFallback`` + * would miss delegated tools otherwise. + */ +export function withDelegationSpanIndent( + Component: ToolCallMessagePartComponent +): ToolCallMessagePartComponent { + const Wrapped: ToolCallMessagePartComponent = (props) => { + const metadata = (props as { metadata?: Record }).metadata; + const indent = shouldIndentToolCallForDelegationSpan(props.toolName, metadata); + const inner = ; + return indent ?
{inner}
: inner; + }; + Wrapped.displayName = `withDelegationSpanIndent(${Component.displayName ?? Component.name ?? "ToolUI"})`; + return Wrapped; +} + export const ToolFallback: ToolCallMessagePartComponent = (props) => { if (isInterruptResult(props.result)) { if (isDoomLoopInterrupt(props.result)) { diff --git a/surfsense_web/components/public-chat/public-thread.tsx b/surfsense_web/components/public-chat/public-thread.tsx index 750b7410e..2075d82b8 100644 --- a/surfsense_web/components/public-chat/public-thread.tsx +++ b/surfsense_web/components/public-chat/public-thread.tsx @@ -5,6 +5,7 @@ import { AuiIf, MessagePrimitive, ThreadPrimitive, + type ToolCallMessagePartComponent, useAuiState, } from "@assistant-ui/react"; import { CheckIcon, CopyIcon } from "lucide-react"; @@ -14,7 +15,7 @@ import { type FC, type ReactNode, useState } from "react"; import { CitationMetadataProvider } from "@/components/assistant-ui/citation-metadata-context"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { ReasoningMessagePart } from "@/components/assistant-ui/reasoning-message-part"; -import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; +import { ToolFallback, withDelegationSpanIndent } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { GenerateImageToolUI } from "@/components/tool-ui/generate-image"; import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast"; @@ -29,6 +30,8 @@ const GenerateVideoPresentationToolUI = dynamic( { ssr: false } ); +const NullToolUi: ToolCallMessagePartComponent = () => null; + interface PublicThreadProps { footer?: ReactNode; } @@ -162,18 +165,20 @@ const PublicAssistantMessage: FC = () => { Reasoning: ReasoningMessagePart, tools: { by_name: { - generate_podcast: GeneratePodcastToolUI, - generate_report: GenerateReportToolUI, - generate_resume: GenerateResumeToolUI, - generate_video_presentation: GenerateVideoPresentationToolUI, - display_image: GenerateImageToolUI, - generate_image: GenerateImageToolUI, - web_search: () => null, - link_preview: () => null, - multi_link_preview: () => null, - scrape_webpage: () => null, + generate_podcast: withDelegationSpanIndent(GeneratePodcastToolUI), + generate_report: withDelegationSpanIndent(GenerateReportToolUI), + generate_resume: withDelegationSpanIndent(GenerateResumeToolUI), + generate_video_presentation: withDelegationSpanIndent( + GenerateVideoPresentationToolUI + ), + display_image: withDelegationSpanIndent(GenerateImageToolUI), + generate_image: withDelegationSpanIndent(GenerateImageToolUI), + web_search: NullToolUi, + link_preview: NullToolUi, + multi_link_preview: NullToolUi, + scrape_webpage: NullToolUi, }, - Fallback: ToolFallback, + Fallback: withDelegationSpanIndent(ToolFallback), }, }} /> diff --git a/surfsense_web/lib/chat/delegation-span-indent.ts b/surfsense_web/lib/chat/delegation-span-indent.ts new file mode 100644 index 000000000..99e292eaf --- /dev/null +++ b/surfsense_web/lib/chat/delegation-span-indent.ts @@ -0,0 +1,19 @@ +/** + * Indent tool-call cards that belong to an open delegating ``task`` episode. + * + * The backend only stamps ``metadata.spanId`` on tool SSE / persisted parts + * while a ``task`` is active (see ``AgentEventRelayState.tool_activity_metadata``), + * so its presence is sufficient. The opening ``task`` row itself carries the + * same span id but stays flush — it is the header of the delegation. + */ + +export function shouldIndentToolCallForDelegationSpan( + toolName: string, + metadata: Record | undefined +): boolean { + if (toolName === "task") return false; + const v = metadata?.spanId; + return typeof v === "string" && v.trim().length > 0; +} + +export const DELEGATION_SPAN_INDENT_CLASS = "pl-3 sm:ml-4"; From 4b2c9f07cd4966cd64a2d9a795a8556477cd3b91 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Sat, 9 May 2026 00:40:06 +0200 Subject: [PATCH 78/96] Group delegated thinking steps under their task parent and show subagent name. --- .../assistant-ui/thinking-steps.tsx | 247 +++++++++++++++--- 1 file changed, 213 insertions(+), 34 deletions(-) diff --git a/surfsense_web/components/assistant-ui/thinking-steps.tsx b/surfsense_web/components/assistant-ui/thinking-steps.tsx index df1cef12c..6c3832bff 100644 --- a/surfsense_web/components/assistant-ui/thinking-steps.tsx +++ b/surfsense_web/components/assistant-ui/thinking-steps.tsx @@ -1,7 +1,7 @@ import { makeAssistantDataUI, useAuiState } from "@assistant-ui/react"; import { ChevronRightIcon } from "lucide-react"; import type { FC } from "react"; -import { useCallback, useEffect, useState } from "react"; +import { useCallback, useEffect, useMemo, useState } from "react"; import { ChainOfThoughtItem } from "@/components/prompt-kit/chain-of-thought"; import { TextShimmerLoader } from "@/components/prompt-kit/loader"; import { cn } from "@/lib/utils"; @@ -11,15 +11,170 @@ export interface ThinkingStep { title: string; items: string[]; status: "pending" | "in_progress" | "completed"; + /** + * Optional relay metadata forwarded from ``data-thinking-step`` SSE + * (e.g. ``spanId`` set by ``AgentEventRelayState.span_metadata_if_active``). + * Steps under an open delegating ``task`` carry ``metadata.spanId`` and are + * grouped under the preceding parent (``task`` step) as indented children. + */ + metadata?: Record; } /** - * Chain of thought display component - single collapsible dropdown design + * Per-step info joined from the assistant message ``tool-call`` parts via + * the shared ``metadata.thinkingStepId`` correlation + * (set on the server in ``AgentEventRelayState.tool_activity_metadata``). */ -export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolean }> = ({ - steps, - isThreadRunning = true, -}) => { +interface StepToolInfo { + toolName: string; + args: Record; +} + +export type ThinkingStepToolInfoMap = ReadonlyMap; + +/** + * Build ``thinkingStepId → {toolName, args}`` from message content. Used to + * - identify the opening ``task`` step (parent header, never indents) without + * relying on the human-readable title; + * - render the parent's display title from ``args.subagent_type`` instead of + * the generic "Task" copy. + */ +export function buildThinkingStepToolInfo( + content: readonly unknown[] | undefined +): ThinkingStepToolInfoMap { + const m = new Map(); + if (!content) return m; + for (const part of content) { + if (!part || typeof part !== "object") continue; + const o = part as { + type?: string; + toolName?: string; + args?: Record; + metadata?: Record; + }; + if (o.type !== "tool-call" || !o.toolName) continue; + const tid = o.metadata?.thinkingStepId; + if (typeof tid === "string" && tid.trim().length > 0) { + m.set(tid, { toolName: o.toolName, args: o.args ?? {} }); + } + } + return m; +} + +function asNonEmptyString(v: unknown): string | undefined { + return typeof v === "string" && v.trim().length > 0 ? v.trim() : undefined; +} + +function titleCaseSubagent(raw: string): string { + // "notion" → "Notion", "doc_research" → "Doc Research". + return raw + .split(/[\s_-]+/) + .filter(Boolean) + .map((part) => part.charAt(0).toUpperCase() + part.slice(1)) + .join(" "); +} + +/** + * Display title for a step. For the opening ``task`` step we substitute the + * subagent type from the matching tool-call args (e.g. ``"Notion"`` instead of + * the generic ``"Task"``). Falls back to the step's own title if the tool-call + * hasn't streamed in yet. + */ +function resolveDisplayTitle(step: ThinkingStep, info: StepToolInfo | undefined): string { + if (info?.toolName === "task") { + const subagent = asNonEmptyString(info.args?.subagent_type); + if (subagent) return titleCaseSubagent(subagent); + } + return step.title; +} + +function isDelegatedChild(step: ThinkingStep, info: StepToolInfo | undefined): boolean { + const sid = asNonEmptyString(step.metadata?.spanId); + if (!sid) return false; + // The opening ``task`` step also carries ``spanId`` (it owns the span) but + // must render as the parent header. Prefer the joined ``toolName`` (set by + // ``buildThinkingStepToolInfo`` from ``tool-call.metadata.thinkingStepId``). + // Fall back to the title heuristic when no tool-call is matched — happens + // for messages persisted before ``thinkingStepId`` shipped, and briefly + // during streaming if the ``tool-input-start`` frame hasn't been processed + // yet for some reason. + if (info) return info.toolName !== "task"; + return step.title !== "Task"; +} + +interface StepGroup { + parent: ThinkingStep; + children: ThinkingStep[]; +} + +/** + * Group consecutive delegated child steps under the preceding parent step. + * If the very first step is a child (no parent yet seen), it's promoted to a + * parent so it still renders — defensive only, real flows always start with a + * parent step. + */ +const EMPTY_STEP_TOOL_INFO: ThinkingStepToolInfoMap = new Map(); + +function groupSteps( + steps: readonly ThinkingStep[], + stepToolInfo: ThinkingStepToolInfoMap +): StepGroup[] { + const groups: StepGroup[] = []; + for (const step of steps) { + if (isDelegatedChild(step, stepToolInfo.get(step.id)) && groups.length > 0) { + groups[groups.length - 1].children.push(step); + } else { + groups.push({ parent: step, children: [] }); + } + } + return groups; +} + +const StepBody: FC<{ + step: ThinkingStep; + status: "pending" | "in_progress" | "completed"; + displayTitle: string; +}> = ({ step, status, displayTitle }) => ( +
+
+ {displayTitle} +
+ + {step.items && step.items.length > 0 && ( +
+ {step.items.map((item) => ( + + {item} + + ))} +
+ )} +
+); + +/** + * Chain of thought display component - single collapsible dropdown design. + * + * ``stepToolInfo`` joins each step (by ``thinkingStepId``) to its ``tool-call`` + * part so we can: + * - replace the generic ``"Task"`` title with the real subagent name + * (``args.subagent_type``) on the parent header; + * - decide parent-vs-child purely from the matched ``toolName`` instead of + * relying on the displayed title. + */ +export const ThinkingStepsDisplay: FC<{ + steps: ThinkingStep[]; + isThreadRunning?: boolean; + stepToolInfo?: ThinkingStepToolInfoMap; +}> = ({ steps, isThreadRunning = true, stepToolInfo }) => { + const toolInfo = stepToolInfo ?? EMPTY_STEP_TOOL_INFO; const getEffectiveStatus = useCallback( (step: ThinkingStep): "pending" | "in_progress" | "completed" => { if (step.status === "in_progress" && !isThreadRunning) { @@ -31,6 +186,9 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: ); const inProgressStep = steps.find((s) => getEffectiveStatus(s) === "in_progress"); + const inProgressDisplayTitle = inProgressStep + ? resolveDisplayTitle(inProgressStep, toolInfo.get(inProgressStep.id)) + : undefined; const allCompleted = steps.length > 0 && !isThreadRunning && @@ -49,14 +207,16 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: } }, [allCompleted, isProcessing]); + const groups = useMemo(() => groupSteps(steps, toolInfo), [steps, toolInfo]); + if (steps.length === 0) return null; const getHeaderText = () => { if (allCompleted) { return "Reviewed"; } - if (inProgressStep) { - return inProgressStep.title; + if (inProgressDisplayTitle) { + return inProgressDisplayTitle; } if (isProcessing) { return "Processing"; @@ -94,18 +254,26 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: >
- {steps.map((step, index) => { - const effectiveStatus = getEffectiveStatus(step); - const isLast = index === steps.length - 1; + {groups.map((group, groupIndex) => { + const isLastGroup = groupIndex === groups.length - 1; + const parentStatus = getEffectiveStatus(group.parent); + const parentInfo = toolInfo.get(group.parent.id); + const parentTitle = resolveDisplayTitle(group.parent, parentInfo); + const hasChildren = group.children.length > 0; + // Parent dots are connected by a vertical line that runs through + // any indented children (their column has no dot, so the line + // passes cleanly behind them) and overshoots by ~15px to reach + // the next group's dot center (top-[15px]). + const showParentLine = !isLastGroup; return ( -
-
- {!isLast && ( -
+
+
+ {showParentLine && ( +
)}
- {effectiveStatus === "in_progress" ? ( + {parentStatus === "in_progress" ? ( @@ -117,24 +285,25 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
-
- {step.title} -
+ - {step.items && step.items.length > 0 && ( -
- {step.items.map((item) => ( - - {item} - - ))} + {hasChildren && ( +
+ {group.children.map((child) => { + const childInfo = toolInfo.get(child.id); + return ( + + ); + })}
)}
@@ -158,13 +327,23 @@ function ThinkingStepsDataRenderer({ data }: { name: string; data: unknown }) { const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false); const isMessageStreaming = isThreadRunning && isLastMessage; + const content = useAuiState(({ message }) => message?.content); + + const stepToolInfo = useMemo( + () => buildThinkingStepToolInfo(Array.isArray(content) ? content : undefined), + [content] + ); const steps = (data as { steps: ThinkingStep[] } | null)?.steps ?? []; if (steps.length === 0) return null; return (
- +
); } From c603b46ea43c6e99a092e672885bfcdc70e8e0e2 Mon Sep 17 00:00:00 2001 From: "DESKTOP-RTLN3BA\\$punk" Date: Sat, 9 May 2026 04:31:53 -0700 Subject: [PATCH 79/96] feat: added architecture improvement skill --- .../DEEPENING.md | 37 ++++++++++ .../INTERFACE-DESIGN.md | 44 ++++++++++++ .../improve-codebase-architecture/LANGUAGE.md | 53 ++++++++++++++ .../improve-codebase-architecture/SKILL.md | 71 +++++++++++++++++++ .gitignore | 3 +- skills-lock.json | 6 ++ 6 files changed, 213 insertions(+), 1 deletion(-) create mode 100644 .cursor/skills/improve-codebase-architecture/DEEPENING.md create mode 100644 .cursor/skills/improve-codebase-architecture/INTERFACE-DESIGN.md create mode 100644 .cursor/skills/improve-codebase-architecture/LANGUAGE.md create mode 100644 .cursor/skills/improve-codebase-architecture/SKILL.md diff --git a/.cursor/skills/improve-codebase-architecture/DEEPENING.md b/.cursor/skills/improve-codebase-architecture/DEEPENING.md new file mode 100644 index 000000000..ecaf5d7dc --- /dev/null +++ b/.cursor/skills/improve-codebase-architecture/DEEPENING.md @@ -0,0 +1,37 @@ +# Deepening + +How to deepen a cluster of shallow modules safely, given its dependencies. Assumes the vocabulary in [LANGUAGE.md](LANGUAGE.md) — **module**, **interface**, **seam**, **adapter**. + +## Dependency categories + +When assessing a candidate for deepening, classify its dependencies. The category determines how the deepened module is tested across its seam. + +### 1. In-process + +Pure computation, in-memory state, no I/O. Always deepenable — merge the modules and test through the new interface directly. No adapter needed. + +### 2. Local-substitutable + +Dependencies that have local test stand-ins (PGLite for Postgres, in-memory filesystem). Deepenable if the stand-in exists. The deepened module is tested with the stand-in running in the test suite. The seam is internal; no port at the module's external interface. + +### 3. Remote but owned (Ports & Adapters) + +Your own services across a network boundary (microservices, internal APIs). Define a **port** (interface) at the seam. The deep module owns the logic; the transport is injected as an **adapter**. Tests use an in-memory adapter. Production uses an HTTP/gRPC/queue adapter. + +Recommendation shape: *"Define a port at the seam, implement an HTTP adapter for production and an in-memory adapter for testing, so the logic sits in one deep module even though it's deployed across a network."* + +### 4. True external (Mock) + +Third-party services (Stripe, Twilio, etc.) you don't control. The deepened module takes the external dependency as an injected port; tests provide a mock adapter. + +## Seam discipline + +- **One adapter means a hypothetical seam. Two adapters means a real one.** Don't introduce a port unless at least two adapters are justified (typically production + test). A single-adapter seam is just indirection. +- **Internal seams vs external seams.** A deep module can have internal seams (private to its implementation, used by its own tests) as well as the external seam at its interface. Don't expose internal seams through the interface just because tests use them. + +## Testing strategy: replace, don't layer + +- Old unit tests on shallow modules become waste once tests at the deepened module's interface exist — delete them. +- Write new tests at the deepened module's interface. The **interface is the test surface**. +- Tests assert on observable outcomes through the interface, not internal state. +- Tests should survive internal refactors — they describe behaviour, not implementation. If a test has to change when the implementation changes, it's testing past the interface. diff --git a/.cursor/skills/improve-codebase-architecture/INTERFACE-DESIGN.md b/.cursor/skills/improve-codebase-architecture/INTERFACE-DESIGN.md new file mode 100644 index 000000000..3197723a0 --- /dev/null +++ b/.cursor/skills/improve-codebase-architecture/INTERFACE-DESIGN.md @@ -0,0 +1,44 @@ +# Interface Design + +When the user wants to explore alternative interfaces for a chosen deepening candidate, use this parallel sub-agent pattern. Based on "Design It Twice" (Ousterhout) — your first idea is unlikely to be the best. + +Uses the vocabulary in [LANGUAGE.md](LANGUAGE.md) — **module**, **interface**, **seam**, **adapter**, **leverage**. + +## Process + +### 1. Frame the problem space + +Before spawning sub-agents, write a user-facing explanation of the problem space for the chosen candidate: + +- The constraints any new interface would need to satisfy +- The dependencies it would rely on, and which category they fall into (see [DEEPENING.md](DEEPENING.md)) +- A rough illustrative code sketch to ground the constraints — not a proposal, just a way to make the constraints concrete + +Show this to the user, then immediately proceed to Step 2. The user reads and thinks while the sub-agents work in parallel. + +### 2. Spawn sub-agents + +Spawn 3+ sub-agents in parallel using the Agent tool. Each must produce a **radically different** interface for the deepened module. + +Prompt each sub-agent with a separate technical brief (file paths, coupling details, dependency category from [DEEPENING.md](DEEPENING.md), what sits behind the seam). The brief is independent of the user-facing problem-space explanation in Step 1. Give each agent a different design constraint: + +- Agent 1: "Minimize the interface — aim for 1–3 entry points max. Maximise leverage per entry point." +- Agent 2: "Maximise flexibility — support many use cases and extension." +- Agent 3: "Optimise for the most common caller — make the default case trivial." +- Agent 4 (if applicable): "Design around ports & adapters for cross-seam dependencies." + +Include both [LANGUAGE.md](LANGUAGE.md) vocabulary and CONTEXT.md vocabulary in the brief so each sub-agent names things consistently with the architecture language and the project's domain language. + +Each sub-agent outputs: + +1. Interface (types, methods, params — plus invariants, ordering, error modes) +2. Usage example showing how callers use it +3. What the implementation hides behind the seam +4. Dependency strategy and adapters (see [DEEPENING.md](DEEPENING.md)) +5. Trade-offs — where leverage is high, where it's thin + +### 3. Present and compare + +Present designs sequentially so the user can absorb each one, then compare them in prose. Contrast by **depth** (leverage at the interface), **locality** (where change concentrates), and **seam placement**. + +After comparing, give your own recommendation: which design you think is strongest and why. If elements from different designs would combine well, propose a hybrid. Be opinionated — the user wants a strong read, not a menu. diff --git a/.cursor/skills/improve-codebase-architecture/LANGUAGE.md b/.cursor/skills/improve-codebase-architecture/LANGUAGE.md new file mode 100644 index 000000000..530c27630 --- /dev/null +++ b/.cursor/skills/improve-codebase-architecture/LANGUAGE.md @@ -0,0 +1,53 @@ +# Language + +Shared vocabulary for every suggestion this skill makes. Use these terms exactly — don't substitute "component," "service," "API," or "boundary." Consistent language is the whole point. + +## Terms + +**Module** +Anything with an interface and an implementation. Deliberately scale-agnostic — applies equally to a function, class, package, or tier-spanning slice. +_Avoid_: unit, component, service. + +**Interface** +Everything a caller must know to use the module correctly. Includes the type signature, but also invariants, ordering constraints, error modes, required configuration, and performance characteristics. +_Avoid_: API, signature (too narrow — those refer only to the type-level surface). + +**Implementation** +What's inside a module — its body of code. Distinct from **Adapter**: a thing can be a small adapter with a large implementation (a Postgres repo) or a large adapter with a small implementation (an in-memory fake). Reach for "adapter" when the seam is the topic; "implementation" otherwise. + +**Depth** +Leverage at the interface — the amount of behaviour a caller (or test) can exercise per unit of interface they have to learn. A module is **deep** when a large amount of behaviour sits behind a small interface. A module is **shallow** when the interface is nearly as complex as the implementation. + +**Seam** _(from Michael Feathers)_ +A place where you can alter behaviour without editing in that place. The *location* at which a module's interface lives. Choosing where to put the seam is its own design decision, distinct from what goes behind it. +_Avoid_: boundary (overloaded with DDD's bounded context). + +**Adapter** +A concrete thing that satisfies an interface at a seam. Describes *role* (what slot it fills), not substance (what's inside). + +**Leverage** +What callers get from depth. More capability per unit of interface they have to learn. One implementation pays back across N call sites and M tests. + +**Locality** +What maintainers get from depth. Change, bugs, knowledge, and verification concentrate at one place rather than spreading across callers. Fix once, fixed everywhere. + +## Principles + +- **Depth is a property of the interface, not the implementation.** A deep module can be internally composed of small, mockable, swappable parts — they just aren't part of the interface. A module can have **internal seams** (private to its implementation, used by its own tests) as well as the **external seam** at its interface. +- **The deletion test.** Imagine deleting the module. If complexity vanishes, the module wasn't hiding anything (it was a pass-through). If complexity reappears across N callers, the module was earning its keep. +- **The interface is the test surface.** Callers and tests cross the same seam. If you want to test *past* the interface, the module is probably the wrong shape. +- **One adapter means a hypothetical seam. Two adapters means a real one.** Don't introduce a seam unless something actually varies across it. + +## Relationships + +- A **Module** has exactly one **Interface** (the surface it presents to callers and tests). +- **Depth** is a property of a **Module**, measured against its **Interface**. +- A **Seam** is where a **Module**'s **Interface** lives. +- An **Adapter** sits at a **Seam** and satisfies the **Interface**. +- **Depth** produces **Leverage** for callers and **Locality** for maintainers. + +## Rejected framings + +- **Depth as ratio of implementation-lines to interface-lines** (Ousterhout): rewards padding the implementation. We use depth-as-leverage instead. +- **"Interface" as the TypeScript `interface` keyword or a class's public methods**: too narrow — interface here includes every fact a caller must know. +- **"Boundary"**: overloaded with DDD's bounded context. Say **seam** or **interface**. diff --git a/.cursor/skills/improve-codebase-architecture/SKILL.md b/.cursor/skills/improve-codebase-architecture/SKILL.md new file mode 100644 index 000000000..05984a609 --- /dev/null +++ b/.cursor/skills/improve-codebase-architecture/SKILL.md @@ -0,0 +1,71 @@ +--- +name: improve-codebase-architecture +description: Find deepening opportunities in a codebase, informed by the domain language in CONTEXT.md and the decisions in docs/adr/. Use when the user wants to improve architecture, find refactoring opportunities, consolidate tightly-coupled modules, or make a codebase more testable and AI-navigable. +--- + +# Improve Codebase Architecture + +Surface architectural friction and propose **deepening opportunities** — refactors that turn shallow modules into deep ones. The aim is testability and AI-navigability. + +## Glossary + +Use these terms exactly in every suggestion. Consistent language is the point — don't drift into "component," "service," "API," or "boundary." Full definitions in [LANGUAGE.md](LANGUAGE.md). + +- **Module** — anything with an interface and an implementation (function, class, package, slice). +- **Interface** — everything a caller must know to use the module: types, invariants, error modes, ordering, config. Not just the type signature. +- **Implementation** — the code inside. +- **Depth** — leverage at the interface: a lot of behaviour behind a small interface. **Deep** = high leverage. **Shallow** = interface nearly as complex as the implementation. +- **Seam** — where an interface lives; a place behaviour can be altered without editing in place. (Use this, not "boundary.") +- **Adapter** — a concrete thing satisfying an interface at a seam. +- **Leverage** — what callers get from depth. +- **Locality** — what maintainers get from depth: change, bugs, knowledge concentrated in one place. + +Key principles (see [LANGUAGE.md](LANGUAGE.md) for the full list): + +- **Deletion test**: imagine deleting the module. If complexity vanishes, it was a pass-through. If complexity reappears across N callers, it was earning its keep. +- **The interface is the test surface.** +- **One adapter = hypothetical seam. Two adapters = real seam.** + +This skill is _informed_ by the project's domain model. The domain language gives names to good seams; ADRs record decisions the skill should not re-litigate. + +## Process + +### 1. Explore + +Read the project's domain glossary and any ADRs in the area you're touching first. + +Then use the Agent tool with `subagent_type=Explore` to walk the codebase. Don't follow rigid heuristics — explore organically and note where you experience friction: + +- Where does understanding one concept require bouncing between many small modules? +- Where are modules **shallow** — interface nearly as complex as the implementation? +- Where have pure functions been extracted just for testability, but the real bugs hide in how they're called (no **locality**)? +- Where do tightly-coupled modules leak across their seams? +- Which parts of the codebase are untested, or hard to test through their current interface? + +Apply the **deletion test** to anything you suspect is shallow: would deleting it concentrate complexity, or just move it? A "yes, concentrates" is the signal you want. + +### 2. Present candidates + +Present a numbered list of deepening opportunities. For each candidate: + +- **Files** — which files/modules are involved +- **Problem** — why the current architecture is causing friction +- **Solution** — plain English description of what would change +- **Benefits** — explained in terms of locality and leverage, and also in how tests would improve + +**Use CONTEXT.md vocabulary for the domain, and [LANGUAGE.md](LANGUAGE.md) vocabulary for the architecture.** If `CONTEXT.md` defines "Order," talk about "the Order intake module" — not "the FooBarHandler," and not "the Order service." + +**ADR conflicts**: if a candidate contradicts an existing ADR, only surface it when the friction is real enough to warrant revisiting the ADR. Mark it clearly (e.g. _"contradicts ADR-0007 — but worth reopening because…"_). Don't list every theoretical refactor an ADR forbids. + +Do NOT propose interfaces yet. Ask the user: "Which of these would you like to explore?" + +### 3. Grilling loop + +Once the user picks a candidate, drop into a grilling conversation. Walk the design tree with them — constraints, dependencies, the shape of the deepened module, what sits behind the seam, what tests survive. + +Side effects happen inline as decisions crystallize: + +- **Naming a deepened module after a concept not in `CONTEXT.md`?** Add the term to `CONTEXT.md` — same discipline as `/grill-with-docs` (see [CONTEXT-FORMAT.md](../grill-with-docs/CONTEXT-FORMAT.md)). Create the file lazily if it doesn't exist. +- **Sharpening a fuzzy term during the conversation?** Update `CONTEXT.md` right there. +- **User rejects the candidate with a load-bearing reason?** Offer an ADR, framed as: _"Want me to record this as an ADR so future architecture reviews don't re-suggest it?"_ Only offer when the reason would actually be needed by a future explorer to avoid re-suggesting the same thing — skip ephemeral reasons ("not worth it right now") and self-evident ones. See [ADR-FORMAT.md](../grill-with-docs/ADR-FORMAT.md). +- **Want to explore alternative interfaces for the deepened module?** See [INTERFACE-DESIGN.md](INTERFACE-DESIGN.md). diff --git a/.gitignore b/.gitignore index 2e6ed14e8..f3823a843 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,5 @@ node_modules/ .DS_Store deepagents/ debug.log -opencode/ \ No newline at end of file +opencode/ +hermes-agent/ \ No newline at end of file diff --git a/skills-lock.json b/skills-lock.json index ce251e303..f722ec0d3 100644 --- a/skills-lock.json +++ b/skills-lock.json @@ -46,6 +46,12 @@ "sourceType": "github", "computedHash": "ddd61f32254be1303ce4b7be5d507c932de4af53489a0ebb1309bf61de99018c" }, + "improve-codebase-architecture": { + "source": "mattpocock/skills", + "sourceType": "github", + "skillPath": "skills/engineering/improve-codebase-architecture/SKILL.md", + "computedHash": "2da1d23b8f53cfe67f2e0b68924ab9f4ec400bb6480de097007eeaeb517d1722" + }, "internal-linking-optimizer": { "source": "aaron-he-zhu/seo-geo-claude-skills", "sourceType": "github", From a8417e3c451090da6d05e54ef152b3ff3c72757a Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Sat, 9 May 2026 14:37:06 +0200 Subject: [PATCH 80/96] Render HITL approval cards inline in the thinking-steps timeline. --- .../assistant-ui/assistant-message.tsx | 17 ++- .../assistant-ui/thinking-steps.tsx | 101 ++++++++++++++++-- surfsense_web/lib/hitl/index.ts | 6 ++ surfsense_web/lib/hitl/render-target.tsx | 48 +++++++++ 4 files changed, 158 insertions(+), 14 deletions(-) create mode 100644 surfsense_web/lib/hitl/render-target.tsx diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx index a21ade74a..549141779 100644 --- a/surfsense_web/components/assistant-ui/assistant-message.tsx +++ b/surfsense_web/components/assistant-ui/assistant-message.tsx @@ -59,6 +59,7 @@ import { DropdownMenuLabel } from "@/components/ui/dropdown-menu"; import { useComments } from "@/hooks/use-comments"; import { useMediaQuery } from "@/hooks/use-media-query"; import { useElectronAPI } from "@/hooks/use-platform"; +import { withHitlInTimeline } from "@/lib/hitl"; import { getProviderIcon } from "@/lib/provider-icons"; import { cn } from "@/lib/utils"; @@ -508,12 +509,22 @@ const MessageInfoDropdown: FC = () => { // page through them and stage decisions instead of firing one resume per card. // ``withDelegationSpanIndent`` wraps every entry (including Fallback) so delegated // subagent tools don't bypass span indentation via a named ``by_name`` UI. +// ``withHitlInTimeline`` is the OUTERMOST wrapper so a body render with an +// interrupt result returns ``null`` immediately — no inner wrappers paint +// — while a timeline render (under ``HitlRenderTargetProvider value="timeline"`` +// inside ``ThinkingStepsDisplay``) passes through to the real component. const bundleTool = (Component: ToolCallMessagePartComponent) => - withBundleStep(withDelegationSpanIndent(Component)); + withHitlInTimeline(withBundleStep(withDelegationSpanIndent(Component))); const NullToolUi: ToolCallMessagePartComponent = () => null; -const TOOLS_BY_NAME = { +/** + * Tool-call UI registry. Exported so ``ThinkingStepsDisplay`` can mount + * the SAME wrapped components inline under a step row when the card's + * result is an HITL interrupt. The wrappers handle ``ToolCallIdProvider`` + * and bundle paging consistently across both render targets. + */ +export const TOOLS_BY_NAME = { generate_report: bundleTool(GenerateReportToolUI), generate_resume: bundleTool(GenerateResumeToolUI), generate_podcast: bundleTool(GeneratePodcastToolUI), @@ -554,7 +565,7 @@ const TOOLS_BY_NAME = { scrape_webpage: NullToolUi, } as const; -const TOOLS_FALLBACK = bundleTool(ToolFallback); +export const TOOLS_FALLBACK = bundleTool(ToolFallback); const AssistantMessageInner: FC = () => { const isMobile = !useMediaQuery("(min-width: 768px)"); diff --git a/surfsense_web/components/assistant-ui/thinking-steps.tsx b/surfsense_web/components/assistant-ui/thinking-steps.tsx index 6c3832bff..46d33a9fa 100644 --- a/surfsense_web/components/assistant-ui/thinking-steps.tsx +++ b/surfsense_web/components/assistant-ui/thinking-steps.tsx @@ -1,9 +1,15 @@ -import { makeAssistantDataUI, useAuiState } from "@assistant-ui/react"; +import { + makeAssistantDataUI, + type ToolCallMessagePartComponent, + useAuiState, +} from "@assistant-ui/react"; import { ChevronRightIcon } from "lucide-react"; import type { FC } from "react"; import { useCallback, useEffect, useMemo, useState } from "react"; +import { TOOLS_BY_NAME, TOOLS_FALLBACK } from "@/components/assistant-ui/assistant-message"; import { ChainOfThoughtItem } from "@/components/prompt-kit/chain-of-thought"; import { TextShimmerLoader } from "@/components/prompt-kit/loader"; +import { HitlRenderTargetProvider, isInterruptResult } from "@/lib/hitl"; import { cn } from "@/lib/utils"; export interface ThinkingStep { @@ -24,20 +30,33 @@ export interface ThinkingStep { * Per-step info joined from the assistant message ``tool-call`` parts via * the shared ``metadata.thinkingStepId`` correlation * (set on the server in ``AgentEventRelayState.tool_activity_metadata``). + * + * Carries enough of the part to: + * - identify the opening ``task`` step and substitute the subagent display + * name on the parent header (uses ``toolName`` and ``args``); + * - render the matching tool component inline under the step row when the + * card's result is an HITL interrupt (uses ``toolCallId``, ``argsText``, + * ``result``, ``langchainToolCallId``). */ interface StepToolInfo { + toolCallId: string; toolName: string; args: Record; + argsText?: string; + result?: unknown; + langchainToolCallId?: string; } export type ThinkingStepToolInfoMap = ReadonlyMap; /** - * Build ``thinkingStepId → {toolName, args}`` from message content. Used to + * Build ``thinkingStepId → StepToolInfo`` from message content. Used to * - identify the opening ``task`` step (parent header, never indents) without * relying on the human-readable title; * - render the parent's display title from ``args.subagent_type`` instead of - * the generic "Task" copy. + * the generic "Task" copy; + * - mount the matching tool-call card inline under a step row when the + * result is an HITL interrupt (see ``TimelineHitlCard``). */ export function buildThinkingStepToolInfo( content: readonly unknown[] | undefined @@ -48,14 +67,25 @@ export function buildThinkingStepToolInfo( if (!part || typeof part !== "object") continue; const o = part as { type?: string; + toolCallId?: string; toolName?: string; args?: Record; + argsText?: string; + result?: unknown; + langchainToolCallId?: string; metadata?: Record; }; - if (o.type !== "tool-call" || !o.toolName) continue; + if (o.type !== "tool-call" || !o.toolName || !o.toolCallId) continue; const tid = o.metadata?.thinkingStepId; if (typeof tid === "string" && tid.trim().length > 0) { - m.set(tid, { toolName: o.toolName, args: o.args ?? {} }); + m.set(tid, { + toolCallId: o.toolCallId, + toolName: o.toolName, + args: o.args ?? {}, + argsText: o.argsText, + result: o.result, + langchainToolCallId: o.langchainToolCallId, + }); } } return m; @@ -159,6 +189,47 @@ const StepBody: FC<{
); +/** + * Mount the same tool-call UI used in the message body, but inside the + * chain-of-thought timeline. The body copy returns ``null`` (see + * ``withHitlInTimeline`` in ``lib/hitl/render-target``), so the card + * effectively moves from the body to the timeline for the lifetime of the + * interrupt (pending → processing → complete / rejected). + * + * ``metadata`` is intentionally omitted from the props we forward — the + * step row already provides any indentation it needs, so we don't want + * ``withDelegationSpanIndent`` to add a second indent + border on top. + * + * ``status`` is a placeholder (HITL UIs read only ``args`` + ``result``) + * so we don't need to mirror assistant-ui's runtime status object here. + */ +const TimelineHitlCard: FC<{ info: StepToolInfo }> = ({ info }) => { + const Comp = + (TOOLS_BY_NAME as Record)[info.toolName] ?? + TOOLS_FALLBACK; + const props = { + toolCallId: info.toolCallId, + toolName: info.toolName, + args: info.args, + argsText: info.argsText, + result: info.result, + langchainToolCallId: info.langchainToolCallId, + status: { type: "complete" } as const, + }; + return ( + + {/* biome-ignore lint/suspicious/noExplicitAny: ToolCallMessagePartProps requires + runtime-only fields (addResult, resume, MessagePartState) we don't have when + re-rendering manually; HITL components only read args + result. */} + + + ); +}; + +function hitlInterruptInfo(info: StepToolInfo | undefined): StepToolInfo | undefined { + return info && isInterruptResult(info.result) ? info : undefined; +} + /** * Chain of thought display component - single collapsible dropdown design. * @@ -291,17 +362,25 @@ export const ThinkingStepsDisplay: FC<{ displayTitle={parentTitle} /> + {(() => { + const hitl = hitlInterruptInfo(parentInfo); + return hitl ? : null; + })()} + {hasChildren && (
{group.children.map((child) => { const childInfo = toolInfo.get(child.id); + const childHitl = hitlInterruptInfo(childInfo); return ( - +
+ + {childHitl && } +
); })}
diff --git a/surfsense_web/lib/hitl/index.ts b/surfsense_web/lib/hitl/index.ts index 4bb15e8b5..a2f218d5d 100644 --- a/surfsense_web/lib/hitl/index.ts +++ b/surfsense_web/lib/hitl/index.ts @@ -6,6 +6,12 @@ export { useHitlBundle, useToolCallIdContext, } from "./bundle-context"; +export { + type HitlRenderTarget, + HitlRenderTargetProvider, + useHitlRenderTarget, + withHitlInTimeline, +} from "./render-target"; export type { HitlDecision, InterruptActionRequest, diff --git a/surfsense_web/lib/hitl/render-target.tsx b/surfsense_web/lib/hitl/render-target.tsx new file mode 100644 index 000000000..cbfdbf2be --- /dev/null +++ b/surfsense_web/lib/hitl/render-target.tsx @@ -0,0 +1,48 @@ +"use client"; + +import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; +import { createContext, useContext } from "react"; +import { isInterruptResult } from "./types"; + +/** + * Where this tool-call card is currently rendering. + * + * - ``"body"`` (default) — assistant-ui's ``MessagePrimitive.Parts`` renders + * the card inside the message bubble. + * - ``"timeline"`` — ``ThinkingStepsDisplay`` renders the SAME component + * inline under the matching step row so the HITL approval lives in the + * chain-of-thought instead of as a standalone card in the message body. + * + * The two render targets share one component implementation; the context + * lets the body render skip itself when the timeline copy will show the + * card, avoiding a double-render. + */ +export type HitlRenderTarget = "body" | "timeline"; + +const HitlRenderTargetContext = createContext("body"); + +export const HitlRenderTargetProvider = HitlRenderTargetContext.Provider; + +export function useHitlRenderTarget(): HitlRenderTarget { + return useContext(HitlRenderTargetContext); +} + +/** + * Hide the body render of a tool-call whose result is a HITL interrupt. + * The same component is mounted again inside ``ThinkingStepsDisplay`` + * with ``HitlRenderTargetProvider value="timeline"`` — that copy renders + * normally, so the card "moves" from the message body to the timeline. + * + * Pure pass-through for non-HITL results AND for the timeline render. + */ +export function withHitlInTimeline( + Component: ToolCallMessagePartComponent +): ToolCallMessagePartComponent { + const Wrapped: ToolCallMessagePartComponent = (props) => { + const target = useHitlRenderTarget(); + if (target === "body" && isInterruptResult(props.result)) return null; + return ; + }; + Wrapped.displayName = `withHitlInTimeline(${Component.displayName ?? Component.name ?? "ToolUI"})`; + return Wrapped; +} From 5c1f5edd751e1e6916dff83102226d43f5ee0180 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Sat, 9 May 2026 14:39:44 +0200 Subject: [PATCH 81/96] Add chat-messages feature module architecture doc. --- .../features/chat-messages/ARCHITECTURE.md | 483 ++++++++++++++++++ 1 file changed, 483 insertions(+) create mode 100644 surfsense_web/features/chat-messages/ARCHITECTURE.md diff --git a/surfsense_web/features/chat-messages/ARCHITECTURE.md b/surfsense_web/features/chat-messages/ARCHITECTURE.md new file mode 100644 index 000000000..030374aaf --- /dev/null +++ b/surfsense_web/features/chat-messages/ARCHITECTURE.md @@ -0,0 +1,483 @@ +# `features/chat-messages/` — Architecture + +> **Scope.** This module owns everything between an assistant message +> arriving and its rendering inside the chat UI: the timeline (the +> agent's process — reasoning + every tool call), and the HITL +> primitives that per-tool components compose to render approval views. +> +> It does **NOT** own: the thread shell, the composer, the streaming +> pipeline, the message frame (`assistant-message.tsx`, +> `user-message.tsx`, markdown renderer, citations), the comments +> sidebar, or any of the 63 individual tool-ui integration files +> under `components/tool-ui/`. + +--- + +## 1. Mental model + +Every assistant message has two regions: + +| Region | What it shows | +|---|---| +| **Timeline** | The agent's *process*. Reasoning, every tool call, grouped by delegation `spanId` into a tree. Each tool call is rendered by its registered component, which selects its own view (running, awaiting approval, success, error, etc.) by discriminating its `result` data. | +| **Body** | The agent's *product*. Markdown text, citations, native reasoning blocks, and value-add deliverables (image viewer, chart, canvas). Connector tool cards do NOT render here. | + +**Principle: timeline = process, body = product. No overlap.** + +A tool's UI lives in the body **if and only if** it produces a deliverable +the user wants to interact with — view, scrub, copy, share. If the UI +just shows that the tool ran and what it did, it lives in the timeline. + +``` +┌─ Assistant Message ─────────────────────────────────────────┐ +│ │ +│ ╔═════════════════════════════════════════════════════╗ │ +│ ║ TIMELINE (process) ║ │ +│ ║ ║ │ +│ ║ ▸ task: NotionAgent [running] ║ │ +│ ║ ▸ search_workspace [completed] ║ │ +│ ║ ▸ update_page ← rendered by ║ │ +│ ║ (Notion-styled approval UpdateNotion- ║ │ +│ ║ card OR Notion-styled PageToolUI; ║ │ +│ ║ success/error card, the component ║ │ +│ ║ per its own data picks the view ║ │ +│ ║ discrimination) from result) ║ │ +│ ║ ▸ summarize [completed] ║ │ +│ ╚═════════════════════════════════════════════════════╝ │ +│ │ +│ ╔═════════════════════════════════════════════════════╗ │ +│ ║ BODY (product) ║ │ +│ ║ ║ │ +│ ║ Markdown text, citations, value-add deliverables ║ │ +│ ║ only. Connector tool cards do NOT render here. ║ │ +│ ╚═════════════════════════════════════════════════════╝ │ +│ │ +└─────────────────────────────────────────────────────────────┘ +``` + +--- + +## 2. The single data model + +The timeline reads ONE data structure: a `TimelineItem[]`. There are +no parallel structures for "thinking steps", "tool calls", "HITL +bundles", etc. Every visible piece of agent activity is a `TimelineItem`. + +### 2.1 The discriminated union (outer discrimination) + +Two kinds. The timeline does **outer discrimination** — it chooses +reasoning view vs tool-call mounting based on `kind`. + +```ts +type ItemStatus = + | "pending" | "running" | "completed" | "cancelled" | "error"; + +interface BaseItem { + id: string; + spanId?: string; // groups items into delegation tree (parent task + children) + status: ItemStatus; +} + +interface ReasoningItem extends BaseItem { + kind: "reasoning"; + text: string; +} + +interface ToolCallItem extends BaseItem { + kind: "tool-call"; + toolName: string; + args: Record; + argsText?: string; + result?: unknown; // per-tool component discriminates this internally + langchainToolCallId?: string; +} + +type TimelineItem = ReasoningItem | ToolCallItem; + +interface TimelineGroup { + parent: TimelineItem; + children: TimelineItem[]; +} +``` + +**`ToolCallItem` has no `approval` field, no `phase`, no `view`.** All of +that is derived inside the per-tool component from the result data. + +### 2.2 Inner discrimination (per-tool component) + +Each tool registers a component that receives the tool-call data and +decides what to render based on its own result-shape discriminators: + +```tsx +const UpdateNotionPageToolUI: TimelineToolComponent = (props) => { + if (isInterruptResult(props.result)) return ; + if (isAuthErrorResult(props.result)) return ; + if (isErrorResult(props.result)) return ; + if (isInfoResult(props.result)) return ; + if (isSuccessResult(props.result)) return ; + return ; +}; +``` + +The discriminators (`isInterruptResult`, `isAuthErrorResult`, etc.) +are **types, not centralized infrastructure**. The component owns +the dispatch. The timeline knows none of this. + +### 2.3 The pure builder + +```ts +function buildTimeline( + content: MessageContent[], + thinkingSteps: ThinkingStep[], +): TimelineGroup[] +``` + +Builds the timeline from existing message content + thinking-step data +parts. Pure function. Sets `kind` and `status` on each item; preserves +`result` verbatim for per-tool discrimination. + +### 2.4 The dispatch (timeline-level) + +Two cases. Exhaustive switch. No runtime guards in the timeline renderer. + +```tsx +function TimelineItemView({ item }: { item: TimelineItem }) { + switch (item.kind) { + case "reasoning": return ; + case "tool-call": return ; + } +} + +function ToolCallItemView({ item }: { item: ToolCallItem }) { + const ToolBody = getToolComponent(item.toolName) ?? FallbackToolBody; + return ; +} +``` + +**No card frame, header, body slot, approval area, or result panel +at the timeline level.** Each tool component owns its own visual +presentation. This matches how every existing tool-ui component +already works — they each render their own rounded card with their +own header. + +--- + +## 3. The timeline's tool-component contract + +Tool components mounted by the timeline implement a subset of +assistant-ui's `ToolCallMessagePartProps` — only the fields the +timeline can supply: + +```ts +interface TimelineToolProps { + toolCallId: string; + toolName: string; + args: Record; + argsText?: string; + result?: unknown; + langchainToolCallId?: string; + status: ItemStatus; // simple enum, not assistant-ui's complex status object +} + +type TimelineToolComponent = (props: TimelineToolProps) => ReactNode; +``` + +Notably absent (compared to `ToolCallMessagePartProps`): +- `addResult`, `resume` — runtime-only, not needed; HITL decisions + flow through `useHitlDecision` (a hook) which talks to the runtime + directly. +- The complex `status: ToolCallMessagePartState["status"]` object — + replaced by our simple `ItemStatus` enum. + +The 15 existing HITL-aware tool-ui components only use the subset +above. They are **retyped** to `TimelineToolComponent` in the cutover +commit (mechanical: `ToolCallMessagePartComponent` → `TimelineToolComponent`). + +--- + +## 4. Rendering topology — how the body opts out + +The body uses assistant-ui's `MessagePrimitive.Parts` and registers a +**no-op fallback** for tool calls so they don't render here: + +```tsx + null, // every other tool-call: render nothing in the body + }} +/> +``` + +`BODY_TOOLS` starts empty (no value-add deliverables exist yet) and +grows as we identify them. Every tool not in `BODY_TOOLS` renders +nothing in the body. + +The timeline reads message content via `useAuiState(({ message }) => +message?.content)` and runs `buildTimeline` to produce the items it +renders. Tool-call data IS in the message; the body just chooses not +to render it. + +**Result:** zero dual placement. Zero suppression HOC. Zero +render-target context. Zero coordination. + +--- + +## 5. Slice layout + +``` +features/chat-messages/ +├── ARCHITECTURE.md +│ +├── timeline/ ← the process surface +│ ├── types.ts (TimelineItem union, ToolCallItem, ItemStatus, TimelineGroup) +│ ├── build-timeline.ts (pure: content + thinkingSteps → groups) +│ ├── grouping.ts (pure: group items by spanId) +│ ├── subagent-rename.ts (pure: parent task title from args.subagent_type) +│ ├── tool-registry/ (PRIVATE to timeline; only timeline mounts tools) +│ │ ├── types.ts (TimelineToolComponent, TimelineToolProps) +│ │ ├── registry.ts (TOOLS_BY_NAME from components/tool-ui/*) +│ │ ├── adapt-props.ts (pure: ToolCallItem → TimelineToolProps) +│ │ ├── fallback/ +│ │ │ ├── fallback-tool-body.tsx (TimelineToolComponent for unregistered tools — discriminates internally) +│ │ │ ├── default-fallback-card.tsx (the non-HITL fallback view: status icon + collapsible + JSON) +│ │ │ ├── revert-button.tsx (revert affordance — used by default-fallback-card) +│ │ │ ├── use-tool-action.ts (action lookup hook for revert) +│ │ │ └── index.ts +│ │ └── index.ts +│ ├── items/ +│ │ ├── reasoning-item.tsx (renders kind: "reasoning") +│ │ ├── tool-call-item.tsx (lookup component + mount with adapted props — ~10 lines) +│ │ └── index.ts +│ ├── timeline.tsx (groups + iteration + 2-case dispatch) +│ ├── data-renderer.tsx (assistant-ui adapter; exports TimelineDataUI) +│ └── index.ts +│ +├── hitl/ ← pure HITL primitives +│ ├── types.ts (InterruptResult, HitlPhase, HitlDecision, isInterruptResult) +│ ├── use-hitl-decision.ts (hook: dispatch approve/edit/reject — used by every approval card) +│ ├── use-hitl-phase.ts (hook: tracks pending → processing → approved/rejected/edited) +│ ├── approval-cards/ (the FALLBACK-mounted approval views; per-tool components import from here OR build their own) +│ │ ├── generic-approval.tsx (default approval UI — what FallbackToolBody mounts for interrupt results) +│ │ ├── doom-loop-approval.tsx (special-case approval UI + isDoomLoopInterrupt) +│ │ └── index.ts +│ ├── edit-panel/ +│ │ ├── edit-panel.atom.ts (Jotai atoms for the panel state) +│ │ ├── edit-panel.tsx (root: atom wiring + desktop/mobile switch only) +│ │ ├── fields/ +│ │ │ ├── email-tags-field.tsx (EmailsTagField + parse/format helpers) +│ │ │ ├── calendar-field.tsx (DateTimePickerField + parse/format helpers) +│ │ │ ├── extra-fields.tsx (ExtraField switch renderer) +│ │ │ └── index.ts (private barrel) +│ │ └── index.ts +│ └── index.ts +│ +└── (no body slice yet — body just registers `tools={{ fallback: () => null }}`) +``` + +### 5.1 Notable absences + +| Was | Status | Reason | +|---|---|---| +| `tool-cards/` slice | **Folded into `timeline/`** | Tool-call rendering happens in the timeline; the tool-registry is private to timeline. | +| `bundleTool` composer | **Deleted** | Body opts out via `fallback: () => null`. No HOCs to compose. | +| `withDelegationSpanIndent` HOC | **Deleted** | Tree indent is owned by the timeline's group renderer. | +| `withBundleStep` + `HitlBundleProvider` | **Deleted** | Multi-approval is just N inline renderings; no coordination needed. | +| `withHitlInTimeline` + `HitlRenderTargetProvider` | **Deleted** | Tool cards never render in body; no dual-placement to suppress. | +| `pickApprovalCard` central dispatcher | **Deleted** | Each tool component picks its own view via internal discrimination. The fallback has its OWN internal dispatcher (interrupt → generic-approval; doom-loop → doom-loop-approval). | +| `getHitlToolComponent` registry | **Deleted** | The tool-registry is just a `Record`; lookup is `TOOLS_BY_NAME[name]`. | +| Centralized `approval-area.tsx` in timeline | **Deleted** | The approval is a view the per-tool component renders, not an area the timeline composes. | +| `ApprovalState` on `ToolCallItem` | **Deleted** | Phase is local UI state inside per-tool approval cards (via `useHitlPhase`). The timeline doesn't track it. | +| `ThinkingStepToolInfoMap` Map join | **Deleted** | The unified `TimelineItem` union eliminates the join. | + +--- + +## 6. Public surfaces + +### `timeline/index.ts` + +```ts +export { TimelineDataUI }; // the assistant-ui registration +export { Timeline }; // exposed for tests +export type { TimelineItem, ReasoningItem, ToolCallItem, TimelineGroup, ItemStatus }; +export type { TimelineToolComponent, TimelineToolProps }; +``` + +### `hitl/index.ts` + +```ts +export type { InterruptResult, InterruptActionRequest, InterruptReviewConfig, HitlDecision, HitlPhase }; +export { isInterruptResult }; + +export { useHitlDecision }; +export { useHitlPhase }; + +export { GenericHitlApprovalToolUI }; // for tool-ui integrations that want to compose on top +export { DoomLoopApprovalToolUI, isDoomLoopInterrupt }; + +export { HitlEditPanel, MobileHitlEditPanel }; +export { openHitlEditPanelAtom, closeHitlEditPanelAtom, hitlEditPanelAtom }; +export type { ExtraField }; +``` + +The 63 `components/tool-ui/*` integrations consume `hitl/`'s public +surface (types, hooks, edit-panel atom, optionally the fallback +approval cards). Nothing else. + +--- + +## 7. Layering & SRP rules + +### 7.1 The "what knows about what" rule + +| Component | Knows about | +|---|---| +| `timeline/` | Itself + `hitl/` (via the fallback) + `components/tool-ui/*` (via the registry) | +| `timeline/tool-registry/` | The `TimelineToolComponent` contract, `components/tool-ui/*`, and `hitl/` (for the fallback's approval views) | +| `hitl/` | Itself only — no knowledge of timeline, tool-call types, registry | +| `components/tool-ui/*` | `hitl/` only (for HITL primitives + optional fallback approval cards); never reaches into `timeline/` | +| Body (`assistant-message.tsx`) | The `BODY_TOOLS` registry and `TimelineDataUI` from `timeline/index.ts` | + +`hitl/` does **NOT** import from `timeline/`. The dependency arrow is one-way. + +### 7.2 Render policy belongs to the surface, not the primitive + +- `hitl/` exposes hooks, types, and the fallback approval cards. +- `timeline/` decides WHEN and WHERE tool components mount (inside + `tool-call-item.tsx`). +- A `hitl/` primitive must never assume it's being rendered in the + timeline, the body, or anywhere else. It receives props, renders + UI, returns. No environment sniffing, no context. +- Per-tool components in `components/tool-ui/*` decide WHICH view to + render based on result-shape discriminators. The timeline does not + know these discriminators exist. + +### 7.3 Single Responsibility + +Rules in priority order: + +1. **One responsibility per file.** Need "and" to describe it? Split it. +2. **One responsibility per function.** Same. +3. **Line count is a smell, not a budget.** ~250 lines = pause and + ask "still one responsibility?"; ~500 lines = strong presumption + of split needed unless explicitly justified at the top of the file. + +Notable splits driven by SRP during the port: + +- `hitl-edit-panel.tsx` (current 405 lines, 4 responsibilities) → 5 + files: `edit-panel.tsx` (root + layout switch), `email-tags-field.tsx`, + `calendar-field.tsx`, `extra-fields.tsx`, `edit-panel.atom.ts`. +- `tool-fallback.tsx` (current 533 lines, 3 responsibilities) → split + across `fallback-tool-body.tsx`, `default-fallback-card.tsx`, + `revert-button.tsx`, `use-tool-action.ts`. +- `thinking-steps.tsx` (current 434 lines, 5 responsibilities) → + folded into the new `timeline/` slice across `types.ts`, + `build-timeline.ts`, `grouping.ts`, `subagent-rename.ts`, + `timeline.tsx`, `items/*`, `data-renderer.tsx`. + +--- + +## 8. Tested behaviors + +Unit tests live next to the file they cover (`*.test.ts(x)`). + +- `timeline/build-timeline.test.ts` — content + thinkingSteps → correct items, correct kind, correct status, correct ordering. `result` preserved verbatim. +- `timeline/grouping.test.ts` — items group correctly by spanId; first item with a spanId is the parent; orphaned children are promoted defensively. +- `timeline/subagent-rename.test.ts` — `task` step's display title resolves to `args.subagent_type` (title-cased); falls back to "Task" when subagent type is missing. +- `timeline/tool-registry/registry.test.ts` — `TOOLS_BY_NAME` includes every named tool; `FallbackToolBody` is returned for unknown names; the fallback dispatches correctly (interrupt → generic, doom-loop → doom-loop, otherwise → default fallback). +- `timeline/tool-registry/adapt-props.test.ts` — `ToolCallItem` → `TimelineToolProps` mapping is lossless; status mapping is correct. +- `hitl/use-hitl-phase.test.ts` — phase transitions through pending → processing → approved/rejected/edited correctly. +- `hitl/approval-cards/doom-loop-approval.test.tsx` — `isDoomLoopInterrupt` matches doom-loop-shape interrupts only. + +Smoke test after cutover: +- Assistant message renders; markdown + citations work in body. +- All connector tool calls render in timeline only (none in body). +- Reasoning steps render in timeline. +- Single HITL flow (Notion update): approve, edit, reject — each transitions through the phases correctly. +- Multiple pending HITL cards: each renders inline at its position; deciding one doesn't affect the others. +- Doom-loop approval renders the special card. +- Revert button works on completed default-fallback cards and survives reload. +- Subagent name renaming on `task` parent step. + +--- + +## 9. Migration plan (strangler fig, single atomic cutover) + +### Phase A — Build the new slice in parallel + +In dependency order: `hitl/` first (leaf), then `timeline/`. The +existing code (`thinking-steps.tsx`, `tool-fallback.tsx`, +`assistant-message.tsx`'s tool registry, etc.) remains fully +functional throughout Phase A. + +1. Port `hitl/` primitives. Apply SRP splits (edit panel into 5 files). + `hitl/approval-cards/{generic,doom-loop}-approval.tsx` are ported + as standalone components — they're what the fallback mounts and + what per-tool integrations may compose on top of. +2. Build `timeline/` slice. Implement `buildTimeline` from scratch + (do NOT copy thinking-steps logic verbatim — design the pure + function around the new union). Build the `tool-registry/` with + `TimelineToolComponent` contract; the registry imports from + `components/tool-ui/*` (no file moves yet). +3. Add unit tests as listed in §8. +4. Verify: tsc clean, biome clean, no consumer file touched, no + linter regressions. + +### Phase B — Atomic cutover (single commit) + +| File | Change | +|---|---| +| `components/assistant-ui/assistant-message.tsx` | Replace `TOOLS_BY_NAME`/`TOOLS_FALLBACK` definitions with `BODY_TOOLS` (initially empty) + `tools={{ fallback: () => null }}`. Replace `ThinkingStepsDataUI` registration with `TimelineDataUI`. | +| `components/public-chat/public-thread.tsx` | Same registry + data UI swap. | +| `app/dashboard/.../new-chat/page.tsx` | Switch `ThinkingStepsDataUI` → `TimelineDataUI`. Drop `HitlBundleProvider` (no longer needed). | +| `components/free-chat/free-chat-page.tsx` | Switch `ThinkingStepsDataUI` → `TimelineDataUI`. | +| `components/public-chat/public-chat-view.tsx` | Same. | +| `components/layout/ui/right-panel/RightPanel.tsx` | Switch `HitlEditPanel` import to `@/features/chat-messages/hitl`. | +| The 15 `components/tool-ui/*` HITL-aware integration files | (a) Switch HITL imports from `@/lib/hitl`, `@/hooks/use-hitl-phase`, `@/atoms/chat/hitl-edit-panel.atom` → `@/features/chat-messages/hitl`. (b) Retype from `ToolCallMessagePartComponent` → `TimelineToolComponent` (mechanical type rename). | + +### Phase C — Delete legacy + +After cutover passes smoke tests: + +- `components/assistant-ui/thinking-steps.tsx` +- `components/assistant-ui/tool-fallback.tsx` +- `lib/chat/delegation-span-indent.ts` +- `lib/hitl/` (entire folder) +- `components/hitl-bundle-pager/` (entire folder) +- `components/tool-ui/generic-hitl-approval.tsx` +- `components/tool-ui/doom-loop-approval.tsx` +- `components/hitl-edit-panel/` (entire folder) +- `hooks/use-hitl-phase.ts` +- `atoms/chat/hitl-edit-panel.atom.ts` + +Verify: no orphan files, no dead imports, no test regressions. + +--- + +## 10. Out of scope (and one consumer relationship) + +### 10.1 The 63 `components/tool-ui/*` integrations + +These are **first-class consumers** of `hitl/` and the +`TimelineToolComponent` contract. They are imported by +`timeline/tool-registry/registry.ts` to build `TOOLS_BY_NAME`. They +never reach into `timeline/` themselves. + +They stay where they are. Future option to move them is a separate, +mechanical follow-up refactor. + +### 10.2 Not touched by this refactor + +- The composer (input bar, mention picker, prompt picker, tool toggles). +- The streaming pipeline (`lib/chat/streaming-state.ts`, `stream-pipeline.ts`, `thread-persistence.ts`). +- The chat-comments sidebar. +- The message frame (`assistant-message.tsx`, `user-message.tsx`, `markdown-text.tsx`, `inline-citation.tsx`) beyond swapping the registry imports. + +If any of these become a blocker for the refactor (e.g. the streaming +pipeline needs a metadata field that doesn't exist), surface it +explicitly and decide whether to expand scope before touching it. From 790a6f8c37f64da8982926846cdf09ac8920cb76 Mon Sep 17 00:00:00 2001 From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Date: Sat, 9 May 2026 07:05:57 -0700 Subject: [PATCH 82/96] feat(a11y): add aria-label to clickable media cards The hero carousel video card and use-cases grid image card already had role="button", tabIndex={0}, and onKeyDown handlers. Adds the missing aria-label so screen readers announce what each clickable card does. Both cards now use aria-label={`Expand ${title}`}, matching the example in the issue. Fixes #913 --- surfsense_web/components/homepage/use-cases-grid.tsx | 1 + surfsense_web/components/ui/hero-carousel.tsx | 1 + 2 files changed, 2 insertions(+) diff --git a/surfsense_web/components/homepage/use-cases-grid.tsx b/surfsense_web/components/homepage/use-cases-grid.tsx index 7aa272a5c..b51567700 100644 --- a/surfsense_web/components/homepage/use-cases-grid.tsx +++ b/surfsense_web/components/homepage/use-cases-grid.tsx @@ -68,6 +68,7 @@ function UseCaseCard({
{ diff --git a/surfsense_web/components/ui/hero-carousel.tsx b/surfsense_web/components/ui/hero-carousel.tsx index c93fe5979..2c1a98f46 100644 --- a/surfsense_web/components/ui/hero-carousel.tsx +++ b/surfsense_web/components/ui/hero-carousel.tsx @@ -118,6 +118,7 @@ function HeroCarouselCard({
{ From d9ad9ca5cbd55e178566387e9435ea1e7ea41539 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Sat, 9 May 2026 18:31:16 +0200 Subject: [PATCH 83/96] chat-messages: refresh feature module architecture doc. --- .../features/chat-messages/ARCHITECTURE.md | 120 +++++++++++++++--- 1 file changed, 105 insertions(+), 15 deletions(-) diff --git a/surfsense_web/features/chat-messages/ARCHITECTURE.md b/surfsense_web/features/chat-messages/ARCHITECTURE.md index 030374aaf..5b1dedd3b 100644 --- a/surfsense_web/features/chat-messages/ARCHITECTURE.md +++ b/surfsense_web/features/chat-messages/ARCHITECTURE.md @@ -184,8 +184,10 @@ type TimelineToolComponent = (props: TimelineToolProps) => ReactNode; Notably absent (compared to `ToolCallMessagePartProps`): - `addResult`, `resume` — runtime-only, not needed; HITL decisions - flow through `useHitlDecision` (a hook) which talks to the runtime - directly. + flow through `useHitlDecision`, which either stages in the active + bundle (N≥2) or fires the `hitl-decision` window event the page + listens for (N=1). The hook reads `useToolCallIdContext()` to know + which call is dispatching. - The complex `status: ToolCallMessagePartState["status"]` object — replaced by our simple `ItemStatus` enum. @@ -224,7 +226,81 @@ renders. Tool-call data IS in the message; the body just chooses not to render it. **Result:** zero dual placement. Zero suppression HOC. Zero -render-target context. Zero coordination. +render-target context. Zero pager HOC. + +--- + +## 4a. Multi-approval coordination (the bundle + pager) + +When N HITL interrupts are pending in the same assistant turn (e.g. an +agent fires multiple gated tool calls in parallel), the LangGraph +runtime expects **one resume call with N decisions in order**. Per-card +independent submission isn't possible without backend changes. + +The slice handles this with a single React state container, +`HitlBundleProvider`, mounted once at the thread root by the page that +owns the runtime (currently `app/dashboard/.../new-chat/page.tsx` and +`components/free-chat/free-chat-page.tsx`): + +```tsx + + {/* Thread + Timeline + approval cards mount inside */} + +``` + +Per-card flow: + +1. `tool-call-item.tsx` wraps each mounted tool component in + `` so `useHitlDecision` + knows which call is dispatching. +2. The user clicks approve/edit/reject on a card. +3. `useHitlDecision().dispatch([decision])` runs: + - **Bundle active (N≥2):** stages the decision under this card's + `toolCallId` and fires a `hitl-stage` event so the card's local + result mirror updates immediately (UX continuity — no re-prompt + if user navigates back via the pager). + - **No bundle (N=1):** dispatches the `hitl-decision` event + directly — single-decision fast path. +4. When all N decisions are staged, the user clicks "Submit decisions" + on the pager chrome. `bundle.submit()` dispatches the `hitl-decision` + event with the full ordered array. The page's listener calls + `runtime.resume({ resume: orderedDecisions })` once. + +### Pager UX (kept, not deleted) + +When a bundle is active (N≥2), only ONE approval card is visible at a +time — the current step. Other bundle members are hidden until the user +navigates to them. A small pager chrome (prev/next + "Step X / N" + +"Submit decisions" button) renders once at the end of the timeline. + +Where the responsibilities live: + +- **`tool-call-item.tsx`** (timeline) — checks `useHitlBundle()`. If + the item is in an active bundle but not the current step, returns + null. Otherwise wraps the per-tool component in `ToolCallIdProvider` + and mounts it. +- **`timeline.tsx`** — renders `` once at the bottom, + conditional on `useHitlBundle()` being non-null. +- **`hitl/bundle/pager-chrome.tsx`** — pure presentational component; + reads bundle state, renders nav + Submit. No knowledge of the timeline. + +This is **the only Provider in the slice.** It's a state container, not +a behavior HOC: nothing wraps individual cards. The hide-if-not-current +decision is made at the single mount point (`tool-call-item.tsx`), not +distributed across N HOC wrappers. + +What was deleted vs kept here: + +- **Deleted:** `withBundleStep` HOC. Its two responsibilities (hide + non-current cards; render pager after current card) split into the + two correct places: `tool-call-item.tsx` and `timeline.tsx` + respectively. No HOC to compose. +- **Kept (ported as-is to the slice):** `HitlBundleProvider`, + `useHitlBundle`, `ToolCallIdProvider`, `useToolCallIdContext`, + `BundleSubmit`, `HitlBundleAPI`, `PagerChrome`. --- @@ -258,10 +334,14 @@ features/chat-messages/ │ ├── data-renderer.tsx (assistant-ui adapter; exports TimelineDataUI) │ └── index.ts │ -├── hitl/ ← pure HITL primitives +├── hitl/ ← HITL primitives + bundle state container │ ├── types.ts (InterruptResult, HitlPhase, HitlDecision, isInterruptResult) -│ ├── use-hitl-decision.ts (hook: dispatch approve/edit/reject — used by every approval card) -│ ├── use-hitl-phase.ts (hook: tracks pending → processing → approved/rejected/edited) +│ ├── bundle/ (the ONLY Provider in the slice — coordinates N→1 submission + pager UX) +│ │ ├── bundle-context.tsx (HitlBundleProvider, useHitlBundle, ToolCallIdProvider, useToolCallIdContext, BundleSubmit, HitlBundleAPI) +│ │ ├── pager-chrome.tsx (prev/next/submit chrome — mounted once by timeline.tsx when bundle active) +│ │ └── index.ts +│ ├── use-hitl-decision.ts (hook: stages in bundle when N≥2, direct-dispatches when N=1; used by every approval card) +│ ├── use-hitl-phase.ts (hook: tracks pending → processing → complete/rejected) │ ├── approval-cards/ (the FALLBACK-mounted approval views; per-tool components import from here OR build their own) │ │ ├── generic-approval.tsx (default approval UI — what FallbackToolBody mounts for interrupt results) │ │ ├── doom-loop-approval.tsx (special-case approval UI + isDoomLoopInterrupt) @@ -287,7 +367,8 @@ features/chat-messages/ | `tool-cards/` slice | **Folded into `timeline/`** | Tool-call rendering happens in the timeline; the tool-registry is private to timeline. | | `bundleTool` composer | **Deleted** | Body opts out via `fallback: () => null`. No HOCs to compose. | | `withDelegationSpanIndent` HOC | **Deleted** | Tree indent is owned by the timeline's group renderer. | -| `withBundleStep` + `HitlBundleProvider` | **Deleted** | Multi-approval is just N inline renderings; no coordination needed. | +| `withBundleStep` HOC | **Deleted** | Two responsibilities split into the right places: hide-if-not-current → `tool-call-item.tsx`; render pager after current card → `timeline.tsx`. No HOC. | +| `HitlBundleProvider` + `useHitlBundle` + `PagerChrome` | **Kept** (state container + presentational chrome, not HOCs) | Backend constraint: parallel interrupts need ONE ordered resume call. Provider collects N decisions, pager is the user's submit affordance. | | `withHitlInTimeline` + `HitlRenderTargetProvider` | **Deleted** | Tool cards never render in body; no dual-placement to suppress. | | `pickApprovalCard` central dispatcher | **Deleted** | Each tool component picks its own view via internal discrimination. The fallback has its OWN internal dispatcher (interrupt → generic-approval; doom-loop → doom-loop-approval). | | `getHitlToolComponent` registry | **Deleted** | The tool-registry is just a `Record`; lookup is `TOOLS_BY_NAME[name]`. | @@ -317,6 +398,10 @@ export { isInterruptResult }; export { useHitlDecision }; export { useHitlPhase }; +export { HitlBundleProvider, ToolCallIdProvider, useHitlBundle, useToolCallIdContext }; +export { PagerChrome }; +export type { BundleSubmit, HitlBundleAPI }; + export { GenericHitlApprovalToolUI }; // for tool-ui integrations that want to compose on top export { DoomLoopApprovalToolUI, isDoomLoopInterrupt }; @@ -384,14 +469,19 @@ Notable splits driven by SRP during the port: ## 8. Tested behaviors -Unit tests live next to the file they cover (`*.test.ts(x)`). +> **Status:** No test runner is set up in `surfsense_web` yet. The pure +> functions below are *intended* to be unit-tested but tests are +> deferred to **Phase D** (post-cutover follow-up: install vitest, +> write the suites, update this section). -- `timeline/build-timeline.test.ts` — content + thinkingSteps → correct items, correct kind, correct status, correct ordering. `result` preserved verbatim. -- `timeline/grouping.test.ts` — items group correctly by spanId; first item with a spanId is the parent; orphaned children are promoted defensively. -- `timeline/subagent-rename.test.ts` — `task` step's display title resolves to `args.subagent_type` (title-cased); falls back to "Task" when subagent type is missing. +Planned tests once vitest is in: + +- `timeline/build-timeline.test.ts` — content + thinkingSteps → correct items, correct kind, correct status, correct ordering. `result` preserved verbatim. Orphan tool-calls (no `thinkingStepId`) appended at end. +- `timeline/grouping.test.ts` — items group correctly by spanId; first item with a spanId is the parent; orphaned children become parents defensively. +- `timeline/subagent-rename.test.ts` — `task` tool-call's display title resolves to `args.subagent_type` (title-cased); falls back to `getToolDisplayName("task")` when subagent type is missing. - `timeline/tool-registry/registry.test.ts` — `TOOLS_BY_NAME` includes every named tool; `FallbackToolBody` is returned for unknown names; the fallback dispatches correctly (interrupt → generic, doom-loop → doom-loop, otherwise → default fallback). - `timeline/tool-registry/adapt-props.test.ts` — `ToolCallItem` → `TimelineToolProps` mapping is lossless; status mapping is correct. -- `hitl/use-hitl-phase.test.ts` — phase transitions through pending → processing → approved/rejected/edited correctly. +- `hitl/use-hitl-phase.test.ts` — phase transitions through pending → processing → complete/rejected correctly. - `hitl/approval-cards/doom-loop-approval.test.tsx` — `isDoomLoopInterrupt` matches doom-loop-shape interrupts only. Smoke test after cutover: @@ -434,7 +524,7 @@ functional throughout Phase A. |---|---| | `components/assistant-ui/assistant-message.tsx` | Replace `TOOLS_BY_NAME`/`TOOLS_FALLBACK` definitions with `BODY_TOOLS` (initially empty) + `tools={{ fallback: () => null }}`. Replace `ThinkingStepsDataUI` registration with `TimelineDataUI`. | | `components/public-chat/public-thread.tsx` | Same registry + data UI swap. | -| `app/dashboard/.../new-chat/page.tsx` | Switch `ThinkingStepsDataUI` → `TimelineDataUI`. Drop `HitlBundleProvider` (no longer needed). | +| `app/dashboard/.../new-chat/page.tsx` | Switch `ThinkingStepsDataUI` → `TimelineDataUI`. Switch `HitlBundleProvider` import from `@/lib/hitl` → `@/features/chat-messages/hitl` (keep the wrap; just new path). | | `components/free-chat/free-chat-page.tsx` | Switch `ThinkingStepsDataUI` → `TimelineDataUI`. | | `components/public-chat/public-chat-view.tsx` | Same. | | `components/layout/ui/right-panel/RightPanel.tsx` | Switch `HitlEditPanel` import to `@/features/chat-messages/hitl`. | @@ -447,8 +537,8 @@ After cutover passes smoke tests: - `components/assistant-ui/thinking-steps.tsx` - `components/assistant-ui/tool-fallback.tsx` - `lib/chat/delegation-span-indent.ts` -- `lib/hitl/` (entire folder) -- `components/hitl-bundle-pager/` (entire folder) +- `lib/hitl/` (entire folder — replaced by `features/chat-messages/hitl/{types.ts,bundle/,use-hitl-decision.ts}`) +- `components/hitl-bundle-pager/` (entire folder — `PagerChrome` ported to `hitl/bundle/pager-chrome.tsx`; `withBundleStep` deleted, responsibilities split into `tool-call-item.tsx` + `timeline.tsx`) - `components/tool-ui/generic-hitl-approval.tsx` - `components/tool-ui/doom-loop-approval.tsx` - `components/hitl-edit-panel/` (entire folder) From 9e451a59072619ecfee560b4041d8077178ec681 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Sat, 9 May 2026 18:31:23 +0200 Subject: [PATCH 84/96] chat-messages: add hitl module with types, hooks, bundle, approval cards, and edit panel. --- .../approval-cards/doom-loop-approval.tsx | 191 +++++++++++++ .../hitl/approval-cards/generic-approval.tsx | 261 ++++++++++++++++++ .../hitl/approval-cards/index.ts | 2 + .../hitl/bundle/bundle-context.tsx | 157 +++++++++++ .../chat-messages/hitl/bundle/index.ts | 8 + .../hitl/bundle/pager-chrome.tsx | 65 +++++ .../hitl/edit-panel/edit-panel.atom.ts | 82 ++++++ .../hitl/edit-panel/edit-panel.tsx | 203 ++++++++++++++ .../hitl/edit-panel/fields/calendar-field.tsx | 112 ++++++++ .../edit-panel/fields/email-tags-field.tsx | 86 ++++++ .../hitl/edit-panel/fields/extra-fields.tsx | 74 +++++ .../hitl/edit-panel/fields/index.ts | 3 + .../chat-messages/hitl/edit-panel/index.ts | 7 + .../features/chat-messages/hitl/index.ts | 31 +++ .../features/chat-messages/hitl/types.ts | 51 ++++ .../chat-messages/hitl/use-hitl-decision.ts | 45 +++ .../chat-messages/hitl/use-hitl-phase.ts | 66 +++++ 17 files changed, 1444 insertions(+) create mode 100644 surfsense_web/features/chat-messages/hitl/approval-cards/doom-loop-approval.tsx create mode 100644 surfsense_web/features/chat-messages/hitl/approval-cards/generic-approval.tsx create mode 100644 surfsense_web/features/chat-messages/hitl/approval-cards/index.ts create mode 100644 surfsense_web/features/chat-messages/hitl/bundle/bundle-context.tsx create mode 100644 surfsense_web/features/chat-messages/hitl/bundle/index.ts create mode 100644 surfsense_web/features/chat-messages/hitl/bundle/pager-chrome.tsx create mode 100644 surfsense_web/features/chat-messages/hitl/edit-panel/edit-panel.atom.ts create mode 100644 surfsense_web/features/chat-messages/hitl/edit-panel/edit-panel.tsx create mode 100644 surfsense_web/features/chat-messages/hitl/edit-panel/fields/calendar-field.tsx create mode 100644 surfsense_web/features/chat-messages/hitl/edit-panel/fields/email-tags-field.tsx create mode 100644 surfsense_web/features/chat-messages/hitl/edit-panel/fields/extra-fields.tsx create mode 100644 surfsense_web/features/chat-messages/hitl/edit-panel/fields/index.ts create mode 100644 surfsense_web/features/chat-messages/hitl/edit-panel/index.ts create mode 100644 surfsense_web/features/chat-messages/hitl/index.ts create mode 100644 surfsense_web/features/chat-messages/hitl/types.ts create mode 100644 surfsense_web/features/chat-messages/hitl/use-hitl-decision.ts create mode 100644 surfsense_web/features/chat-messages/hitl/use-hitl-phase.ts diff --git a/surfsense_web/features/chat-messages/hitl/approval-cards/doom-loop-approval.tsx b/surfsense_web/features/chat-messages/hitl/approval-cards/doom-loop-approval.tsx new file mode 100644 index 000000000..5b2b0e385 --- /dev/null +++ b/surfsense_web/features/chat-messages/hitl/approval-cards/doom-loop-approval.tsx @@ -0,0 +1,191 @@ +"use client"; + +import { CornerDownLeftIcon, OctagonAlert } from "lucide-react"; +import { useCallback, useEffect, useMemo } from "react"; +import { TextShimmerLoader } from "@/components/prompt-kit/loader"; +import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Separator } from "@/components/ui/separator"; +import type { HitlApprovalCard, HitlDecision, InterruptResult } from "../types"; +import { isInterruptResult } from "../types"; +import { useHitlDecision } from "../use-hitl-decision"; +import { useHitlPhase } from "../use-hitl-phase"; + +/** + * Specialized HITL card for ``DoomLoopMiddleware`` interrupts. The + * backend signals these by setting ``context.permission === "doom_loop"`` + * on the ``permission_ask`` interrupt. + * + * The card replaces the generic "approve/reject" framing with a + * "continue/stop" affordance that better matches the user's mental + * model: the agent is stuck repeating itself, not asking permission + * for a destructive action. + */ +function DoomLoopCardView({ + toolName, + args, + interruptData, + onDecision, +}: { + toolName: string; + args: Record; + interruptData: InterruptResult; + onDecision: (decision: HitlDecision) => void; +}) { + const { phase, setProcessing, setRejected } = useHitlPhase(interruptData); + + const context = (interruptData.context ?? {}) as Record; + const threshold = typeof context.threshold === "number" ? context.threshold : 3; + const stuckTool = (typeof context.tool === "string" && context.tool) || toolName; + const recentSignatures = Array.isArray(context.recent_signatures) + ? (context.recent_signatures as string[]) + : []; + const displayName = stuckTool.replace(/_/g, " ").replace(/\b\w/g, (c) => c.toUpperCase()); + + const argPreview = useMemo(() => { + if (!args || Object.keys(args).length === 0) return null; + try { + const json = JSON.stringify(args, null, 2); + return json.length > 600 ? `${json.slice(0, 600)}…` : json; + } catch { + return null; + } + }, [args]); + + const handleContinue = useCallback(() => { + if (phase !== "pending") return; + setProcessing(); + onDecision({ type: "approve" }); + }, [phase, setProcessing, onDecision]); + + const handleStop = useCallback(() => { + if (phase !== "pending") return; + setRejected(); + onDecision({ type: "reject", message: "Doom loop: user requested stop." }); + }, [phase, setRejected, onDecision]); + + useEffect(() => { + const handler = (e: KeyboardEvent) => { + if (phase !== "pending") return; + if (e.key === "Enter" && !e.shiftKey && !e.ctrlKey && !e.metaKey) { + e.preventDefault(); + handleStop(); + } + }; + window.addEventListener("keydown", handler); + return () => window.removeEventListener("keydown", handler); + }, [phase, handleStop]); + + const isResolved = phase === "complete" || phase === "rejected"; + + return ( + + + + + {phase === "rejected" + ? "Stopped" + : phase === "processing" + ? "Continuing…" + : phase === "complete" + ? "Continued" + : "I might be stuck"} + + {!isResolved && ( + + doom-loop + + )} + + + {phase === "processing" ? ( + + ) : phase === "rejected" ? ( +

+ I stopped retrying {displayName} as you asked. +

+ ) : phase === "complete" ? ( +

+ Continuing to call {displayName} as you asked. +

+ ) : ( +

+ I called {displayName} {threshold} times in a row + with similar arguments. Should I keep going or stop and rethink? +

+ )} + + {argPreview && phase === "pending" && ( + <> + +
+

+ Last arguments +

+ 
+								{argPreview}
+
+
+ + )} + + {recentSignatures.length > 0 && phase === "pending" && ( +
+ + Show repeated signatures ({recentSignatures.length}) + +
    + {recentSignatures.map((sig) => ( +
  • + {sig} +
    • + ))} +
+
+ )} + + {phase === "pending" && ( +
+ + +
+ )} + + + ); +} + +/** + * Discriminator: returns true iff the result is a ``permission_ask`` + * interrupt with ``context.permission === "doom_loop"``. The fallback + * uses this BEFORE mounting an approval card to choose between + * ``DoomLoopApproval`` and ``GenericHitlApproval``. + */ +export function isDoomLoopInterrupt(result: unknown): boolean { + if (!isInterruptResult(result)) return false; + const ctx = (result.context ?? {}) as Record; + return ctx.permission === "doom_loop"; +} + +/** + * Specialized doom-loop approval mounted by ``FallbackToolBody`` when + * ``isDoomLoopInterrupt(result)`` is true. Caller is responsible for + * the discrimination; this card receives a known ``InterruptResult``. + */ +export const DoomLoopApproval: HitlApprovalCard = ({ toolName, args, result }) => { + const { dispatch } = useHitlDecision(); + return ( + dispatch([decision])} + /> + ); +}; diff --git a/surfsense_web/features/chat-messages/hitl/approval-cards/generic-approval.tsx b/surfsense_web/features/chat-messages/hitl/approval-cards/generic-approval.tsx new file mode 100644 index 000000000..c8b35dbe0 --- /dev/null +++ b/surfsense_web/features/chat-messages/hitl/approval-cards/generic-approval.tsx @@ -0,0 +1,261 @@ +"use client"; + +import { CornerDownLeftIcon, Pencil } from "lucide-react"; +import { useCallback, useEffect, useMemo, useState } from "react"; +import { toast } from "sonner"; +import { TextShimmerLoader } from "@/components/prompt-kit/loader"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Textarea } from "@/components/ui/textarea"; +import { getToolDisplayName } from "@/contracts/enums/toolIcons"; +import { connectorsApiService } from "@/lib/apis/connectors-api.service"; +import type { HitlApprovalCard, HitlDecision, InterruptResult } from "../types"; +import { useHitlDecision } from "../use-hitl-decision"; +import { useHitlPhase } from "../use-hitl-phase"; + +function ParamEditor({ + params, + onChange, + disabled, +}: { + params: Record; + onChange: (updated: Record) => void; + disabled: boolean; +}) { + const entries = Object.entries(params); + if (entries.length === 0) return null; + + return ( +
+ {entries.map(([key, value]) => { + const strValue = value == null ? "" : String(value); + const isLong = strValue.length > 120; + const fieldId = `hitl-param-${key}`; + + return ( +
+ + {isLong ? ( +