fix: project agent node outputs into draft graph (#37467)

api/controllers/console/app/workflow.py

@@ -2,12 +2,12 @@ import json
 import logging
 from collections.abc import Sequence
 from datetime import datetime
-from typing import Any, NotRequired, TypedDict
+from typing import Any, NotRequired, TypedDict, cast
 
 from flask import abort, request
 from flask_restx import Resource, fields
 from pydantic import AliasChoices, BaseModel, Field, RootModel, ValidationError, field_validator
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 from werkzeug.exceptions import BadRequest, Forbidden, InternalServerError, NotFound
 
 import services
@@ -449,8 +449,16 @@ class DraftWorkflowApi(Resource):
         if not workflow:
             raise DraftWorkflowNotExist()
 
-        # return workflow, if not found, return 404
-        return dump_response(WorkflowResponse, workflow)
+        from services.agent.workflow_publish_service import WorkflowAgentPublishService
+
+        # Return workflow with response-only Agent node job projection so the
+        # front-end can treat draft graph node data as the editing source.
+        response = WorkflowResponse.model_validate(workflow, from_attributes=True).model_dump(mode="json")
+        response["graph"] = WorkflowAgentPublishService.project_draft_bindings_to_graph(
+            session=cast(Session, db.session),
+            draft_workflow=workflow,
+        )
+        return response
 
     @setup_required
     @login_required

api/services/agent/workflow_publish_service.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import copy
 from collections.abc import Mapping
-from typing import Any
+from typing import Any, cast
 
 from pydantic import ValidationError
 from sqlalchemy import select
@@ -21,6 +22,41 @@ class WorkflowAgentPublishService:
     _AGENT_TASK_KEY = "agent_task"
     _AGENT_DECLARED_OUTPUTS_KEY = "agent_declared_outputs"
 
+    @classmethod
+    def project_draft_bindings_to_graph(cls, *, session: Session, draft_workflow: Workflow) -> dict[str, Any]:
+        """Return draft graph with persisted Agent node job config projected into node data.
+
+        Workflow draft graph is the front-end's editing source of truth, while
+        runtime/publish reads WorkflowAgentNodeBinding.node_job_config. This
+        response-only projection keeps reads aligned without writing binding
+        details back into the stored graph JSON.
+        """
+        graph = cast(dict[str, Any], copy.deepcopy(draft_workflow.graph_dict))
+        agent_nodes = dict(WorkflowAgentNodeValidator.iter_agent_v2_nodes(graph))
+        if not agent_nodes:
+            return graph
+
+        bindings = session.scalars(
+            select(WorkflowAgentNodeBinding).where(
+                WorkflowAgentNodeBinding.tenant_id == draft_workflow.tenant_id,
+                WorkflowAgentNodeBinding.app_id == draft_workflow.app_id,
+                WorkflowAgentNodeBinding.workflow_id == draft_workflow.id,
+                WorkflowAgentNodeBinding.workflow_version == cls._DRAFT_WORKFLOW_VERSION,
+                WorkflowAgentNodeBinding.node_id.in_(list(agent_nodes.keys())),
+            )
+        ).all()
+        for binding in bindings:
+            node_data = agent_nodes.get(binding.node_id)
+            if not isinstance(node_data, dict):
+                continue
+            node_job = WorkflowNodeJobConfig.model_validate(binding.node_job_config_dict)
+            if node_job.workflow_prompt is not None:
+                node_data[cls._AGENT_TASK_KEY] = node_job.workflow_prompt
+            node_data[cls._AGENT_DECLARED_OUTPUTS_KEY] = [
+                output.model_dump(mode="json") for output in node_job.declared_outputs
+            ]
+        return graph
+
     @classmethod
     def validate_agent_nodes_for_publish(cls, *, session: Session, draft_workflow: Workflow) -> None:
         WorkflowAgentNodeValidator.validate_published_workflow(session=session, workflow=draft_workflow)

api/tests/unit_tests/controllers/console/app/test_workflow.py

@@ -540,6 +540,58 @@ def test_draft_workflow_get_not_found(monkeypatch: pytest.MonkeyPatch) -> None:
         handler(api, app_model=SimpleNamespace(id="app"))
 
 
+def test_draft_workflow_get_projects_agent_node_job_to_graph(monkeypatch: pytest.MonkeyPatch) -> None:
+    workflow = _make_workflow(
+        graph_dict={
+            "nodes": [
+                {
+                    "id": "agent-node",
+                    "data": {
+                        "type": "agent",
+                        "version": "2",
+                    },
+                }
+            ],
+            "edges": [],
+        }
+    )
+    projected_graph = {
+        "nodes": [
+            {
+                "id": "agent-node",
+                "data": {
+                    "type": "agent",
+                    "version": "2",
+                    "agent_task": "Summarize it.",
+                    "agent_declared_outputs": [{"name": "summary", "type": "string"}],
+                },
+            }
+        ],
+        "edges": [],
+    }
+
+    monkeypatch.setattr(
+        workflow_module,
+        "WorkflowService",
+        lambda: SimpleNamespace(get_draft_workflow=lambda **_k: workflow),
+    )
+
+    from services.agent.workflow_publish_service import WorkflowAgentPublishService
+
+    monkeypatch.setattr(
+        WorkflowAgentPublishService,
+        "project_draft_bindings_to_graph",
+        lambda **_k: projected_graph,
+    )
+
+    api = workflow_module.DraftWorkflowApi()
+    handler = inspect.unwrap(api.get)
+
+    response = handler(api, app_model=SimpleNamespace(id="app"))
+
+    assert response["graph"] == projected_graph
+
+
 def test_advanced_chat_run_conversation_not_exists(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(
         workflow_module.AppGenerateService,

api/tests/unit_tests/services/agent/test_agent_services.py

@@ -1140,6 +1140,62 @@ class TestListWorkflowsReferencingAppAgent:
 
 
 class TestWorkflowAgentDraftBindingSync:
+    def test_projects_binding_declared_outputs_to_draft_graph_response(self):
+        workflow = Workflow(
+            id="workflow-1",
+            tenant_id="tenant-1",
+            app_id="app-1",
+            version=Workflow.VERSION_DRAFT,
+            graph=json.dumps(
+                {
+                    "nodes": [
+                        {
+                            "id": "agent-node",
+                            "data": {
+                                "type": "agent",
+                                "version": "2",
+                                "agent_binding": {
+                                    "binding_type": "roster_agent",
+                                    "agent_id": "agent-1",
+                                },
+                            },
+                        }
+                    ],
+                    "edges": [],
+                }
+            ),
+        )
+        binding = WorkflowAgentNodeBinding(
+            id="binding-1",
+            tenant_id="tenant-1",
+            app_id="app-1",
+            workflow_id="workflow-1",
+            workflow_version=Workflow.VERSION_DRAFT,
+            node_id="agent-node",
+            binding_type=WorkflowAgentBindingType.ROSTER_AGENT,
+            agent_id="agent-1",
+            current_snapshot_id="snapshot-1",
+            node_job_config=WorkflowNodeJobConfig(
+                workflow_prompt="Summarize the upstream result.",
+                declared_outputs=[
+                    DeclaredOutputConfig(name="summary", type=DeclaredOutputType.STRING, description="Short summary")
+                ],
+            ),
+        )
+        session = FakeSession(scalars=[[binding]])
+
+        graph = WorkflowAgentPublishService.project_draft_bindings_to_graph(
+            session=session,
+            draft_workflow=workflow,
+        )
+
+        node_data = graph["nodes"][0]["data"]
+        assert node_data["agent_task"] == "Summarize the upstream result."
+        assert node_data["agent_declared_outputs"][0]["name"] == "summary"
+        assert node_data["agent_declared_outputs"][0]["type"] == "string"
+        assert node_data["agent_declared_outputs"][0]["description"] == "Short summary"
+        assert "agent_declared_outputs" not in workflow.graph_dict["nodes"][0]["data"]
+
     def test_creates_roster_binding_from_agent_node_graph(self):
         workflow = Workflow(
             id="workflow-1",