Reorganize streaming orchestration modules into relay and orchestration folders.

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",
+]

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 = "",

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
+

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,

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

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"]

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,