fix: tests and sources on workflow agent

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") == []