From e68b3f953233d5252d99e470d18e7064f5b4fd80 Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Wed, 3 Jun 2026 22:51:55 +0530
Subject: [PATCH] feat(chat): fix support for mentioned connectors in new chat
input state and runtime context, enhancing query rendering and context
management
---
.../streaming/flows/new_chat/input_state.py | 46 ++++++++++++++++++-
.../streaming/flows/new_chat/orchestrator.py | 5 ++
.../flows/new_chat/runtime_context.py | 6 +++
.../builtin/agent_task/test_dependencies.py | 4 +-
.../test_parallel_refactor_parity.py | 4 ++
5 files changed, 62 insertions(+), 3 deletions(-)
diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/input_state.py b/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/input_state.py
index 0c6704bd1..c9ef6edd6 100644
--- a/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/input_state.py
+++ b/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/input_state.py
@@ -62,6 +62,7 @@ async def build_new_chat_input_state(
user_image_data_urls: list[str] | None,
mentioned_document_ids: list[int] | None,
mentioned_folder_ids: list[int] | None,
+ mentioned_connectors: list[dict[str, Any]] | None,
mentioned_documents: list[dict[str, Any]] | None,
needs_history_bootstrap: bool,
thread_visibility: ChatVisibility,
@@ -110,6 +111,7 @@ async def build_new_chat_input_state(
final_query = _render_query_with_context(
agent_user_query=agent_user_query,
+ mentioned_connectors=mentioned_connectors,
recent_reports=recent_reports,
)
@@ -196,11 +198,16 @@ async def _resolve_mentions_for_query(
def _render_query_with_context(
*,
agent_user_query: str,
+ mentioned_connectors: list[dict[str, Any]] | None,
recent_reports: list[Report],
) -> str:
- """Prepend recent-reports XML block to the user query."""
+ """Prepend connector/report XML context blocks to the user query."""
context_parts: list[str] = []
+ connector_context = _render_mentioned_connectors(mentioned_connectors)
+ if connector_context:
+ context_parts.append(connector_context)
+
if recent_reports:
report_lines: list[str] = []
for r in recent_reports:
@@ -225,3 +232,40 @@ def _render_query_with_context(
return f"{context}\n\n{agent_user_query}"
return agent_user_query
+
+
+def _render_mentioned_connectors(
+ mentioned_connectors: list[dict[str, Any]] | None,
+) -> str | None:
+ """Render selected connector account metadata for connector-backed tools."""
+ if not mentioned_connectors:
+ return None
+
+ connector_lines: list[str] = []
+ for connector in mentioned_connectors:
+ if not isinstance(connector, dict):
+ continue
+ connector_id = connector.get("id")
+ connector_type = connector.get("connector_type") or connector.get(
+ "document_type"
+ )
+ account_name = connector.get("account_name") or connector.get("title")
+ if connector_id is None or connector_type is None:
+ continue
+ connector_lines.append(
+ f' - connector_id={connector_id}, connector_type="{connector_type}", '
+ f'account_name="{account_name or ""}"'
+ )
+
+ if not connector_lines:
+ return None
+
+ return (
+ "\n"
+ "The user selected these exact connector accounts with @. "
+ "These entries are selection metadata, not retrieved connector content. "
+ "When a connector-backed tool needs an account, use the matching "
+ "connector_id from this list if the tool supports connector_id:\n"
+ + "\n".join(connector_lines)
+ + "\n"
+ )
diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py
index 1892320d3..9c25218bf 100644
--- a/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py
+++ b/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py
@@ -124,6 +124,8 @@ async def stream_new_chat(
llm_config_id: int = -1,
mentioned_document_ids: list[int] | None = None,
mentioned_folder_ids: list[int] | None = None,
+ mentioned_connector_ids: list[int] | None = None,
+ mentioned_connectors: list[dict[str, Any]] | None = None,
mentioned_documents: list[dict[str, Any]] | None = None,
checkpoint_id: str | None = None,
needs_history_bootstrap: bool = False,
@@ -435,6 +437,7 @@ async def stream_new_chat(
user_image_data_urls=user_image_data_urls,
mentioned_document_ids=mentioned_document_ids,
mentioned_folder_ids=mentioned_folder_ids,
+ mentioned_connectors=mentioned_connectors,
mentioned_documents=mentioned_documents,
needs_history_bootstrap=needs_history_bootstrap,
thread_visibility=visibility,
@@ -588,6 +591,8 @@ async def stream_new_chat(
mentioned_document_ids=mentioned_document_ids,
accepted_folder_ids=accepted_folder_ids,
mentioned_folder_ids=mentioned_folder_ids,
+ mentioned_connector_ids=mentioned_connector_ids,
+ mentioned_connectors=mentioned_connectors,
request_id=request_id,
turn_id=stream_result.turn_id,
)
diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/runtime_context.py b/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/runtime_context.py
index cf1e8c3fb..66233dec8 100644
--- a/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/runtime_context.py
+++ b/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/runtime_context.py
@@ -8,6 +8,8 @@ mention lists / request ids / turn ids without rebuilding the graph.
from __future__ import annotations
+from typing import Any
+
from app.agents.new_chat.context import SurfSenseContextSchema
@@ -17,6 +19,8 @@ def build_new_chat_runtime_context(
mentioned_document_ids: list[int] | None,
accepted_folder_ids: list[int],
mentioned_folder_ids: list[int] | None,
+ mentioned_connector_ids: list[int] | None,
+ mentioned_connectors: list[dict[str, Any]] | None,
request_id: str | None,
turn_id: str,
) -> SurfSenseContextSchema:
@@ -31,6 +35,8 @@ def build_new_chat_runtime_context(
search_space_id=search_space_id,
mentioned_document_ids=list(mentioned_document_ids or []),
mentioned_folder_ids=list(accepted_folder_ids or mentioned_folder_ids or []),
+ mentioned_connector_ids=list(mentioned_connector_ids or []),
+ mentioned_connectors=list(mentioned_connectors or []),
request_id=request_id,
turn_id=turn_id,
)
diff --git a/surfsense_backend/tests/unit/automations/actions/builtin/agent_task/test_dependencies.py b/surfsense_backend/tests/unit/automations/actions/builtin/agent_task/test_dependencies.py
index ac20b2608..79da12933 100644
--- a/surfsense_backend/tests/unit/automations/actions/builtin/agent_task/test_dependencies.py
+++ b/surfsense_backend/tests/unit/automations/actions/builtin/agent_task/test_dependencies.py
@@ -14,8 +14,8 @@ from typing import Any
import pytest
-import app.automations.actions.agent_task.dependencies as deps_mod
-from app.automations.actions.agent_task.dependencies import (
+import app.automations.actions.builtin.agent_task.dependencies as deps_mod
+from app.automations.actions.builtin.agent_task.dependencies import (
DependencyError,
build_dependencies,
)
diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_parallel_refactor_parity.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_parallel_refactor_parity.py
index e014bb911..fa6d8b9e2 100644
--- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_parallel_refactor_parity.py
+++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_parallel_refactor_parity.py
@@ -246,6 +246,8 @@ def test_new_chat_runtime_context_prefers_accepted_folder_ids() -> None:
mentioned_document_ids=[1, 2],
accepted_folder_ids=[10],
mentioned_folder_ids=[20, 30],
+ mentioned_connector_ids=None,
+ mentioned_connectors=None,
request_id="req",
turn_id="t1",
)
@@ -263,6 +265,8 @@ def test_new_chat_runtime_context_falls_back_to_mentioned_folder_ids() -> None:
mentioned_document_ids=None,
accepted_folder_ids=[],
mentioned_folder_ids=[20, 30],
+ mentioned_connector_ids=None,
+ mentioned_connectors=None,
request_id=None,
turn_id="t2",
)