From 47dcbcb74b1e8233ff36580bc4fe7624b9edfa0c Mon Sep 17 00:00:00 2001
From: Alex <a@tushynski.me>
Date: Sat, 28 Mar 2026 12:03:16 +0000
Subject: [PATCH] fix: tests and sources on workflow agent

---
 .../src/agents/workflow/WorkflowBuilder.tsx   |  40 ++-
 frontend/src/components/ui/multi-select.tsx   |   8 +-
 pytest.ini                                    |   3 +-
 tests/integration/test_workflows.py           | 253 ++++++++++++++++++
 4 files changed, 296 insertions(+), 8 deletions(-)

diff --git a/frontend/src/agents/workflow/WorkflowBuilder.tsx b/frontend/src/agents/workflow/WorkflowBuilder.tsx
index 7caf9453..ad4fec28 100644
--- a/frontend/src/agents/workflow/WorkflowBuilder.tsx
+++ b/frontend/src/agents/workflow/WorkflowBuilder.tsx
@@ -54,7 +54,10 @@ import { FileUpload } from '../../components/FileUpload';
 import AgentDetailsModal from '../../modals/AgentDetailsModal';
 import ConfirmationModal from '../../modals/ConfirmationModal';
 import { ActiveState } from '../../models/misc';
-import { selectToken } from '../../preferences/preferenceSlice';
+import {
+  selectSourceDocs,
+  selectToken,
+} from '../../preferences/preferenceSlice';
 import { getToolDisplayName } from '../../utils/toolUtils';
 import { Agent } from '../types';
 import { ConditionCase, WorkflowNode } from '../types/workflow';
@@ -300,6 +303,7 @@ function createWorkflowPayload(
 function WorkflowBuilderInner() {
   const navigate = useNavigate();
   const token = useSelector(selectToken);
+  const sourceDocs = useSelector(selectSourceDocs);
   const { agentId } = useParams<{ agentId?: string }>();
   const [searchParams] = useSearchParams();
   const folderId = searchParams.get('folder_id');
@@ -341,6 +345,14 @@ function WorkflowBuilderInner() {
   const [availableModels, setAvailableModels] = useState<Model[]>([]);
   const [defaultAgentModelId, setDefaultAgentModelId] = useState('');
   const [availableTools, setAvailableTools] = useState<UserTool[]>([]);
+  const sourceOptions = useMemo(
+    () =>
+      (sourceDocs ?? []).map((doc) => ({
+        value: doc.id ?? 'default',
+        label: doc.name,
+      })),
+    [sourceDocs],
+  );
   const [agentJsonSchemaDrafts, setAgentJsonSchemaDrafts] = useState<
     Record<string, string>
   >({});
@@ -1279,8 +1291,8 @@ function WorkflowBuilderInner() {
 
   const handlePrimaryAction = useCallback(() => {
     if (isPrimaryActionDisabled) return;
-    void persistWorkflow(!canManageAgent);
-  }, [isPrimaryActionDisabled, persistWorkflow, canManageAgent]);
+    void persistWorkflow(false);
+  }, [isPrimaryActionDisabled, persistWorkflow]);
 
   const agentForDetails = useMemo<Agent>(
     () => ({
@@ -1918,6 +1930,28 @@ function WorkflowBuilderInner() {
                                     emptyText="No tools available"
                                   />
                                 </div>
+                                <div>
+                                  <label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
+                                    Sources
+                                  </label>
+                                  <MultiSelect
+                                    options={sourceOptions}
+                                    selected={
+                                      selectedNode.data.config?.sources || []
+                                    }
+                                    onChange={(newSources) =>
+                                      handleUpdateNodeData({
+                                        config: {
+                                          ...(selectedNode.data.config || {}),
+                                          sources: newSources,
+                                        },
+                                      })
+                                    }
+                                    placeholder="Select sources..."
+                                    searchPlaceholder="Search sources..."
+                                    emptyText="No sources available"
+                                  />
+                                </div>
                                 <div>
                                   <label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
                                     Structured Output (JSON Schema)
diff --git a/frontend/src/components/ui/multi-select.tsx b/frontend/src/components/ui/multi-select.tsx
index 0a07d0c0..42ec75b0 100644
--- a/frontend/src/components/ui/multi-select.tsx
+++ b/frontend/src/components/ui/multi-select.tsx
@@ -70,12 +70,12 @@ export function MultiSelect({
           role="combobox"
           aria-expanded={open}
           className={cn(
-            'w-full justify-between border-[#E5E5E5] bg-white hover:bg-gray-50 dark:border-[#3A3A3A] dark:bg-[#2C2C2C] dark:hover:bg-[#383838]',
+            'h-auto min-h-[2.5rem] w-full justify-between border-[#E5E5E5] bg-white py-1.5 hover:bg-gray-50 dark:border-[#3A3A3A] dark:bg-[#2C2C2C] dark:hover:bg-[#383838]',
             !selected.length && 'text-gray-500 dark:text-gray-400',
             className,
           )}
         >
-          <div className="flex flex-wrap gap-1">
+          <div className="flex min-w-0 flex-wrap gap-1">
             {selected.length === 0 ? (
               placeholder
             ) : (
@@ -85,9 +85,9 @@ export function MultiSelect({
                   return (
                     <span
                       key={option?.value || label}
-                      className="dark:bg-purple-30/30 bg-violets-are-blue/20 inline-flex items-center gap-1 rounded-md px-2 py-0.5 text-xs font-medium text-purple-700 dark:text-purple-300"
+                      className="dark:bg-purple-30/30 bg-violets-are-blue/20 inline-flex max-w-[calc(100%-1rem)] items-center gap-1 rounded-md px-2 py-0.5 text-xs font-medium text-purple-700 dark:text-purple-300"
                     >
-                      {label}
+                      <span className="truncate">{label}</span>
                       <span
                         role="button"
                         tabIndex={0}
diff --git a/pytest.ini b/pytest.ini
index 30b76ad3..bfe52165 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -3,7 +3,7 @@ testpaths = tests
 python_files = test_*.py
 python_classes = Test*
 python_functions = test_*
-addopts = 
+addopts =
     -v
     --strict-markers
     --tb=short
@@ -11,6 +11,7 @@ addopts =
     --cov-report=html
     --cov-report=term-missing
     --cov-report=xml
+    --ignore=tests/integration
 markers =
     unit: Unit tests
     integration: Integration tests
diff --git a/tests/integration/test_workflows.py b/tests/integration/test_workflows.py
index a28d34d5..9e41cb30 100644
--- a/tests/integration/test_workflows.py
+++ b/tests/integration/test_workflows.py
@@ -322,3 +322,256 @@ class TestMultiInputEndNode:
         assert len(edges_to_end) >= 2, (
             f"Expected >=2 edges to end after update, got {len(edges_to_end)}"
         )
+
+
+# ---------------------------------------------------------------------------
+# Source-aware payload helpers
+# ---------------------------------------------------------------------------
+
+def workflow_with_sources(sources, suffix=""):
+    """Start -> Agent (with sources) -> End."""
+    return {
+        "name": f"Source WF {int(time.time())}{suffix}",
+        "description": "integration test with sources",
+        "nodes": [
+            {"id": "start_1", "type": "start", "title": "Start",
+             "position": {"x": 0, "y": 0}, "data": {}},
+            {"id": "agent_1", "type": "agent", "title": "Agent",
+             "position": {"x": 200, "y": 0}, "data": {
+                 "agent_type": "classic",
+                 "system_prompt": "You are helpful.",
+                 "prompt_template": "",
+                 "stream_to_user": False,
+                 "sources": sources,
+                 "tools": [],
+             }},
+            {"id": "end_1", "type": "end", "title": "End",
+             "position": {"x": 400, "y": 0}, "data": {}},
+        ],
+        "edges": [
+            {"id": "edge_1", "source": "start_1", "target": "agent_1"},
+            {"id": "edge_2", "source": "agent_1", "target": "end_1"},
+        ],
+    }
+
+
+def workflow_multi_agent_sources(suffix=""):
+    """Start -> Agent A (sources A) -> Agent B (sources B) -> End."""
+    return {
+        "name": f"Multi-Agent Sources {int(time.time())}{suffix}",
+        "description": "two agents with different sources",
+        "nodes": [
+            {"id": "start_1", "type": "start", "title": "Start",
+             "position": {"x": 0, "y": 0}, "data": {}},
+            {"id": "agent_a", "type": "agent", "title": "Agent A",
+             "position": {"x": 200, "y": 0}, "data": {
+                 "agent_type": "agentic",
+                 "system_prompt": "Agent A prompt",
+                 "prompt_template": "",
+                 "stream_to_user": False,
+                 "sources": ["src_alpha", "src_beta"],
+                 "tools": [],
+             }},
+            {"id": "agent_b", "type": "agent", "title": "Agent B",
+             "position": {"x": 400, "y": 0}, "data": {
+                 "agent_type": "classic",
+                 "system_prompt": "Agent B prompt",
+                 "prompt_template": "",
+                 "stream_to_user": True,
+                 "sources": ["src_gamma"],
+                 "tools": [],
+             }},
+            {"id": "end_1", "type": "end", "title": "End",
+             "position": {"x": 600, "y": 0}, "data": {}},
+        ],
+        "edges": [
+            {"id": "e1", "source": "start_1", "target": "agent_a"},
+            {"id": "e2", "source": "agent_a", "target": "agent_b"},
+            {"id": "e3", "source": "agent_b", "target": "end_1"},
+        ],
+    }
+
+
+def _find_agent_node(nodes, node_id):
+    """Find a specific node by id."""
+    return next((n for n in nodes if n["id"] == node_id), None)
+
+
+# ===========================================================================
+# Workflow integration tests
+# ===========================================================================
+
+
+class TestWorkflowIntegration:
+    """Verify end-to-end workflow create → get → update → get round-trips."""
+
+    def test_linear_workflow_round_trip(self, client, created_ids):
+        """Create a linear workflow and verify all nodes/edges survive the round-trip."""
+        payload = linear_workflow(" round-trip")
+        resp = client.post("/api/workflows", json=payload)
+        assert resp.status_code in (200, 201), resp.get_data(as_text=True)
+        wf_id = _extract_id(resp)
+        assert wf_id
+        created_ids.append(wf_id)
+
+        nodes, edges = _get_graph(client, wf_id)
+        assert len(nodes) == 3
+        assert len(edges) == 2
+
+        # Verify node types
+        types = {n["id"]: n["type"] for n in nodes}
+        assert types["start_1"] == "start"
+        assert types["agent_1"] == "agent"
+        assert types["end_1"] == "end"
+
+    def test_agent_config_persisted(self, client, created_ids):
+        """Agent node config (type, prompts, stream_to_user) round-trips correctly."""
+        payload = linear_workflow(" config")
+        resp = client.post("/api/workflows", json=payload)
+        wf_id = _extract_id(resp)
+        created_ids.append(wf_id)
+
+        nodes, _ = _get_graph(client, wf_id)
+        agent = _find_agent_node(nodes, "agent_1")
+        assert agent is not None
+        assert agent["data"]["agent_type"] == "classic"
+        assert agent["data"]["system_prompt"] == "You are helpful."
+        assert agent["data"]["stream_to_user"] is False
+
+    def test_update_workflow_replaces_graph(self, client, created_ids):
+        """Updating a workflow fully replaces nodes and edges."""
+        resp = client.post("/api/workflows", json=simple_workflow(" replace"))
+        wf_id = _extract_id(resp)
+        created_ids.append(wf_id)
+
+        nodes, edges = _get_graph(client, wf_id)
+        assert len(nodes) == 2
+
+        # Update to linear
+        update_resp = client.put(
+            f"/api/workflows/{wf_id}", json=linear_workflow(" replaced")
+        )
+        assert update_resp.status_code == 200
+
+        nodes, edges = _get_graph(client, wf_id)
+        assert len(nodes) == 3
+        assert len(edges) == 2
+
+
+# ===========================================================================
+# Source-specific integration tests
+# ===========================================================================
+
+
+class TestWorkflowSources:
+    """Verify that agent node sources are persisted and retrieved correctly."""
+
+    def test_create_workflow_with_single_source(self, client, created_ids):
+        """A workflow with one source on an agent node persists it."""
+        payload = workflow_with_sources(["default"])
+        resp = client.post("/api/workflows", json=payload)
+        assert resp.status_code in (200, 201), resp.get_data(as_text=True)
+        wf_id = _extract_id(resp)
+        assert wf_id
+        created_ids.append(wf_id)
+
+        nodes, _ = _get_graph(client, wf_id)
+        agent = _find_agent_node(nodes, "agent_1")
+        assert agent is not None, "Agent node not found"
+        assert agent["data"].get("sources") == ["default"], (
+            f"Expected sources=['default'], got {agent['data'].get('sources')}"
+        )
+
+    def test_create_workflow_with_multiple_sources(self, client, created_ids):
+        """Multiple sources on an agent node are all persisted."""
+        sources = ["src_1", "src_2", "src_3"]
+        payload = workflow_with_sources(sources)
+        resp = client.post("/api/workflows", json=payload)
+        assert resp.status_code in (200, 201), resp.get_data(as_text=True)
+        wf_id = _extract_id(resp)
+        created_ids.append(wf_id)
+
+        nodes, _ = _get_graph(client, wf_id)
+        agent = _find_agent_node(nodes, "agent_1")
+        assert agent is not None
+        assert agent["data"].get("sources") == sources
+
+    def test_create_workflow_with_empty_sources(self, client, created_ids):
+        """An agent node with empty sources list is accepted and persisted."""
+        payload = workflow_with_sources([])
+        resp = client.post("/api/workflows", json=payload)
+        assert resp.status_code in (200, 201), resp.get_data(as_text=True)
+        wf_id = _extract_id(resp)
+        assert wf_id
+        created_ids.append(wf_id)
+
+        nodes, _ = _get_graph(client, wf_id)
+        agent = _find_agent_node(nodes, "agent_1")
+        assert agent is not None
+        assert agent["data"].get("sources") == []
+
+    def test_update_workflow_sources(self, client, created_ids):
+        """Updating a workflow replaces agent sources."""
+        # Create with original sources
+        payload = workflow_with_sources(["old_src"])
+        resp = client.post("/api/workflows", json=payload)
+        wf_id = _extract_id(resp)
+        created_ids.append(wf_id)
+
+        # Update with new sources
+        updated_payload = workflow_with_sources(["new_src_1", "new_src_2"], " upd")
+        update_resp = client.put(f"/api/workflows/{wf_id}", json=updated_payload)
+        assert update_resp.status_code == 200, update_resp.get_data(as_text=True)
+
+        nodes, _ = _get_graph(client, wf_id)
+        agent = _find_agent_node(nodes, "agent_1")
+        assert agent is not None
+        assert agent["data"].get("sources") == ["new_src_1", "new_src_2"]
+
+    def test_multi_agent_independent_sources(self, client, created_ids):
+        """Each agent node keeps its own distinct sources list."""
+        payload = workflow_multi_agent_sources()
+        resp = client.post("/api/workflows", json=payload)
+        assert resp.status_code in (200, 201), resp.get_data(as_text=True)
+        wf_id = _extract_id(resp)
+        created_ids.append(wf_id)
+
+        nodes, _ = _get_graph(client, wf_id)
+        agent_a = _find_agent_node(nodes, "agent_a")
+        agent_b = _find_agent_node(nodes, "agent_b")
+
+        assert agent_a is not None, "Agent A not found"
+        assert agent_b is not None, "Agent B not found"
+        assert agent_a["data"].get("sources") == ["src_alpha", "src_beta"]
+        assert agent_b["data"].get("sources") == ["src_gamma"]
+
+    def test_sources_survive_workflow_update(self, client, created_ids):
+        """Sources survive when a workflow is updated without changing sources."""
+        payload = workflow_with_sources(["persistent_src"])
+        resp = client.post("/api/workflows", json=payload)
+        wf_id = _extract_id(resp)
+        created_ids.append(wf_id)
+
+        # Update keeping same sources
+        update_resp = client.put(f"/api/workflows/{wf_id}", json=payload)
+        assert update_resp.status_code == 200
+
+        nodes, _ = _get_graph(client, wf_id)
+        agent = _find_agent_node(nodes, "agent_1")
+        assert agent["data"].get("sources") == ["persistent_src"]
+
+    def test_remove_sources_on_update(self, client, created_ids):
+        """Clearing sources on update results in empty list."""
+        payload = workflow_with_sources(["will_be_removed"])
+        resp = client.post("/api/workflows", json=payload)
+        wf_id = _extract_id(resp)
+        created_ids.append(wf_id)
+
+        # Update with no sources
+        cleared_payload = workflow_with_sources([], " cleared")
+        update_resp = client.put(f"/api/workflows/{wf_id}", json=cleared_payload)
+        assert update_resp.status_code == 200
+
+        nodes, _ = _get_graph(client, wf_id)
+        agent = _find_agent_node(nodes, "agent_1")
+        assert agent["data"].get("sources") == []