chore: import upstream snapshot with attribution

application/agents/workflow_agent.py

@@ -0,0 +1,496 @@
+import logging
+from datetime import datetime, timezone
+from typing import Any, Dict, Generator, List, Optional, Tuple
+
+from application.agents.base import BaseAgent
+from application.agents.workflows.schemas import (
+    ExecutionStatus,
+    Workflow,
+    WorkflowEdge,
+    WorkflowGraph,
+    WorkflowNode,
+    WorkflowRun,
+)
+from application.agents.workflows.workflow_engine import WorkflowEngine
+from application.core.settings import settings
+from application.logging import LogContext, log_activity
+from application.sandbox.artifacts_capture import QuotaExceeded
+from application.storage.db.base_repository import looks_like_uuid
+from application.storage.db.repositories.workflow_edges import WorkflowEdgesRepository
+from application.storage.db.repositories.workflow_nodes import WorkflowNodesRepository
+from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository
+from application.storage.db.repositories.workflows import WorkflowsRepository
+from application.storage.db.session import db_readonly, db_session
+
+logger = logging.getLogger(__name__)
+
+# Per-run cap on attachments staged as run-scoped artifacts; the remainder is
+# dropped (the per-user artifact quota is only a best-effort soft cap).
+_MAX_INPUT_DOCUMENTS = 25
+
+
+class WorkflowAgent(BaseAgent):
+    """A specialized agent that executes predefined workflows."""
+
+    def __init__(
+        self,
+        *args,
+        workflow_id: Optional[str] = None,
+        workflow: Optional[Dict[str, Any]] = None,
+        workflow_owner: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.workflow_id = workflow_id
+        self.workflow_owner = workflow_owner
+        self._workflow_data = workflow
+        self._engine: Optional[WorkflowEngine] = None
+        self._run_persisted = False
+        # Set to a message when the input-document bridge fails fatally (quota), so
+        # the run is finalized FAILED instead of running with missing documents.
+        self._bridge_error: Optional[str] = None
+
+    @log_activity()
+    def gen(self, query: str, log_context: LogContext = None) -> Generator[Dict[str, str], None, None]:
+        yield from self._gen_inner(query, log_context)
+
+    def _gen_inner(self, query: str, log_context: LogContext) -> Generator[Dict[str, str], None, None]:
+        graph = self._load_workflow_graph()
+        if not graph:
+            yield {"type": "error", "error": "Failed to load workflow configuration."}
+            return
+        self._engine = WorkflowEngine(graph, self)
+
+        # Two distinct identities: the workflow *owner* (A) owns the workflow
+        # definition and is used to resolve the workflow row; the *runner* (B,
+        # the caller) owns the run and its artifacts. They are the same user
+        # except for a shared agent, where B != A. Owning run artifacts by the
+        # runner means quota is charged to the uploader and the caller can read
+        # the outputs of the run they triggered (authz is run.user_id == caller).
+        workflow_owner_id = self._resolve_owner_id()
+        run_user_id = self._resolve_run_user_id(workflow_owner_id)
+        pg_workflow_id = self._precreate_workflow_run(workflow_owner_id, run_user_id, query)
+        self._run_persisted = pg_workflow_id is not None
+
+        try:
+            input_documents, dropped = self._bridge_attachments(run_user_id, persisted=self._run_persisted)
+        except QuotaExceeded as exc:
+            # The run's input documents exceed the uploader's artifact quota. Surface
+            # a clean error and finalize the pre-created RUNNING row as FAILED rather
+            # than executing nodes with silently-missing documents.
+            self._bridge_error = str(exc)
+            yield {
+                "type": "error",
+                "user_facing": True,
+                "error": (
+                    "This run's input documents exceed your artifact storage quota. "
+                    "Delete some artifacts and try again."
+                ),
+            }
+            self._finalize_workflow_run(workflow_owner_id, run_user_id, pg_workflow_id, query)
+            return
+
+        # Non-fatal: some attachments were dropped (oversize / unreadable). Emit a
+        # ``notice`` -- NOT an ``error`` -- so the client surfaces which were dropped
+        # without marking the turn failed or ending the stream (an ``error`` event is
+        # terminal client-side and disables reconnect). The run then still executes
+        # with the documents that did bridge.
+        if dropped:
+            yield {"type": "notice", "notice": " ".join(dropped)}
+
+        self._engine.run_persisted = self._run_persisted
+        interrupted = True
+        try:
+            yield from self._engine.execute({"input_documents": input_documents}, query)
+            interrupted = False
+        finally:
+            self._finalize_workflow_run(
+                workflow_owner_id,
+                run_user_id,
+                pg_workflow_id,
+                query,
+                interrupted=interrupted,
+            )
+
+    def _load_workflow_graph(self) -> Optional[WorkflowGraph]:
+        if self._workflow_data:
+            return self._parse_embedded_workflow()
+        if self.workflow_id:
+            return self._load_from_database()
+        return None
+
+    def _parse_embedded_workflow(self) -> Optional[WorkflowGraph]:
+        try:
+            nodes_data = self._workflow_data.get("nodes", [])
+            edges_data = self._workflow_data.get("edges", [])
+
+            workflow = Workflow(
+                name=self._workflow_data.get("name", "Embedded Workflow"),
+                description=self._workflow_data.get("description"),
+            )
+
+            nodes = []
+            for n in nodes_data:
+                node_config = n.get("data", {})
+                nodes.append(
+                    WorkflowNode(
+                        id=n["id"],
+                        workflow_id=self.workflow_id or "embedded",
+                        type=n["type"],
+                        title=n.get("title", "Node"),
+                        description=n.get("description"),
+                        position=n.get("position", {"x": 0, "y": 0}),
+                        config=node_config,
+                    )
+                )
+            edges = []
+            for e in edges_data:
+                edges.append(
+                    WorkflowEdge(
+                        id=e["id"],
+                        workflow_id=self.workflow_id or "embedded",
+                        source=e.get("source") or e.get("source_id"),
+                        target=e.get("target") or e.get("target_id"),
+                        sourceHandle=e.get("sourceHandle") or e.get("source_handle"),
+                        targetHandle=e.get("targetHandle") or e.get("target_handle"),
+                    )
+                )
+            return WorkflowGraph(workflow=workflow, nodes=nodes, edges=edges)
+        except Exception as e:
+            logger.error(f"Invalid embedded workflow: {e}")
+            return None
+
+    def _load_from_database(self) -> Optional[WorkflowGraph]:
+        try:
+            if not self.workflow_id:
+                logger.error("Missing workflow ID for load")
+                return None
+            owner_id = self.workflow_owner
+            if not owner_id and isinstance(self.decoded_token, dict):
+                owner_id = self.decoded_token.get("sub")
+            if not owner_id:
+                logger.error(f"Workflow owner not available for workflow load: {self.workflow_id}")
+                return None
+
+            with db_readonly() as conn:
+                wf_repo = WorkflowsRepository(conn)
+                if looks_like_uuid(self.workflow_id):
+                    workflow_row = wf_repo.get(self.workflow_id, owner_id)
+                else:
+                    workflow_row = wf_repo.get_by_legacy_id(self.workflow_id, owner_id)
+                if workflow_row is None:
+                    logger.error(f"Workflow {self.workflow_id} not found or inaccessible for user {owner_id}")
+                    return None
+                pg_workflow_id = str(workflow_row["id"])
+                graph_version = workflow_row.get("current_graph_version", 1)
+                try:
+                    graph_version = int(graph_version)
+                    if graph_version <= 0:
+                        graph_version = 1
+                except (ValueError, TypeError):
+                    graph_version = 1
+
+                node_rows = WorkflowNodesRepository(conn).find_by_version(
+                    pg_workflow_id,
+                    graph_version,
+                )
+                edge_rows = WorkflowEdgesRepository(conn).find_by_version(
+                    pg_workflow_id,
+                    graph_version,
+                )
+
+            workflow = Workflow(
+                name=workflow_row.get("name"),
+                description=workflow_row.get("description"),
+            )
+            nodes = [
+                WorkflowNode(
+                    id=n["node_id"],
+                    workflow_id=pg_workflow_id,
+                    type=n["node_type"],
+                    title=n.get("title") or "Node",
+                    description=n.get("description"),
+                    position=n.get("position") or {"x": 0, "y": 0},
+                    config=n.get("config") or {},
+                )
+                for n in node_rows
+            ]
+            edges = [
+                WorkflowEdge(
+                    id=e["edge_id"],
+                    workflow_id=pg_workflow_id,
+                    source=e.get("source_id"),
+                    target=e.get("target_id"),
+                    sourceHandle=e.get("source_handle"),
+                    targetHandle=e.get("target_handle"),
+                )
+                for e in edge_rows
+            ]
+
+            return WorkflowGraph(workflow=workflow, nodes=nodes, edges=edges)
+        except Exception as e:
+            logger.error(f"Failed to load workflow from database: {e}")
+            return None
+
+    def _resolve_owner_id(self) -> Optional[str]:
+        """Resolve the workflow *owner* (used to resolve the owned workflow row)."""
+        owner_id = self.workflow_owner
+        if not owner_id and isinstance(self.decoded_token, dict):
+            owner_id = self.decoded_token.get("sub")
+        return owner_id
+
+    def _resolve_run_user_id(self, workflow_owner_id: Optional[str]) -> Optional[str]:
+        """Resolve the *runner* (caller) who owns the run and its artifacts.
+
+        Equals the owner for a user running their own workflow (and for external
+        API-key calls, where the key owner is the caller); for a shared agent it
+        is the calling user, so their uploads/outputs are owned by and readable
+        to them rather than silently accruing under the agent owner's account.
+        """
+        return getattr(self, "initial_user_id", None) or getattr(self, "user", None) or workflow_owner_id
+
+    def _resolve_owned_workflow_pg_id(self, conn: Any, owner_id: Optional[str]) -> Optional[str]:
+        """Return the owned workflow's PG id, or None for an unowned/draft id."""
+        if not self.workflow_id or not owner_id:
+            return None
+        wf_repo = WorkflowsRepository(conn)
+        if looks_like_uuid(self.workflow_id):
+            workflow_row = wf_repo.get(self.workflow_id, owner_id)
+        else:
+            workflow_row = wf_repo.get_by_legacy_id(self.workflow_id, owner_id)
+        return str(workflow_row["id"]) if workflow_row is not None else None
+
+    def _precreate_workflow_run(
+        self,
+        workflow_owner_id: Optional[str],
+        run_user_id: Optional[str],
+        query: str,
+    ) -> Optional[str]:
+        """Insert the run row up front so run-scoped artifacts are authz-reachable mid-run.
+
+        The workflow row is resolved against its *owner*; the run is owned by the
+        *runner* so artifact access (``run.user_id == caller``) tracks the caller.
+        """
+        if not self._engine or not self.workflow_id or not workflow_owner_id or not run_user_id:
+            return None
+        try:
+            with db_session() as conn:
+                pg_workflow_id = self._resolve_owned_workflow_pg_id(conn, workflow_owner_id)
+                if pg_workflow_id is None:
+                    return None
+                WorkflowRunsRepository(conn).create(
+                    pg_workflow_id,
+                    run_user_id,
+                    ExecutionStatus.RUNNING.value,
+                    run_id=self._engine.workflow_run_id,
+                    inputs={"query": query},
+                    started_at=datetime.now(timezone.utc),
+                )
+            return pg_workflow_id
+        except Exception as e:
+            logger.error(f"Failed to pre-create workflow run: {e}")
+            return None
+
+    def _bridge_attachments(
+        self, run_user_id: Optional[str], *, persisted: bool
+    ) -> Tuple[List[Dict[str, Any]], List[str]]:
+        """Stage uploaded attachments as run-scoped artifacts the nodes can read.
+
+        Bytes are read server-side from each attachment's ``upload_path`` (bounded
+        by ``ARTIFACT_MAX_BYTES``, handle always closed) and re-persisted through
+        ``persist_new_artifact`` (size/sha256/storage key all derived server-side);
+        only the resulting references enter the run state. Artifacts are owned by
+        the *runner* (the uploader), not the workflow owner.
+
+        Returns the bridged references and a list of user-facing notices for
+        attachments that were dropped (oversize / unreadable / unstorable).
+        ``QuotaExceeded`` is NOT swallowed: it propagates to the caller so the run
+        fails cleanly instead of running with silently-missing documents.
+        """
+        if not self._engine or not self.attachments or not run_user_id:
+            return [], []
+        # Without a persisted run row the artifacts would be orphaned (no authz
+        # parent), so skip the bridge for unowned/draft ids.
+        if not persisted:
+            return [], []
+        from application.sandbox.artifacts_capture import persist_new_artifact
+        from application.storage.storage_creator import StorageCreator
+
+        storage = StorageCreator.get_storage()
+        max_bytes = int(getattr(settings, "ARTIFACT_MAX_BYTES", 0) or 0)
+        dropped: List[str] = []
+        if len(self.attachments) > _MAX_INPUT_DOCUMENTS:
+            over = len(self.attachments) - _MAX_INPUT_DOCUMENTS
+            logger.warning(
+                "Workflow run input documents exceed cap (%d); dropping %d attachment(s)",
+                _MAX_INPUT_DOCUMENTS,
+                over,
+            )
+            dropped.append(
+                f"Only the first {_MAX_INPUT_DOCUMENTS} input document(s) were used; "
+                f"{over} additional attachment(s) were dropped."
+            )
+        refs: List[Dict[str, Any]] = []
+        for attachment in self.attachments[:_MAX_INPUT_DOCUMENTS]:
+            upload_path = attachment.get("upload_path") or attachment.get("path")
+            if not upload_path:
+                continue
+            filename = attachment.get("filename") or "attachment"
+            mime_type = attachment.get("mime_type") or "application/octet-stream"
+            # Reject oversize attachments via the authoritative ``size`` column
+            # BEFORE buffering the bytes into worker memory (a memory-DoS guard);
+            # the bounded read below backstops a missing/lying ``size``.
+            declared_size = attachment.get("size")
+            if max_bytes and isinstance(declared_size, (int, float)) and declared_size > max_bytes:
+                dropped.append(f'Document "{filename}" exceeds the artifact size limit and was skipped.')
+                continue
+            data = self._read_attachment_bytes(storage, upload_path, max_bytes)
+            if data is None:
+                dropped.append(f'Document "{filename}" could not be read and was skipped.')
+                continue
+            if max_bytes and len(data) > max_bytes:
+                dropped.append(f'Document "{filename}" exceeds the artifact size limit and was skipped.')
+                continue
+            # QuotaExceeded propagates (fatal); persist_new_artifact returns None on
+            # any other error, which we report as a per-attachment drop.
+            ref = persist_new_artifact(
+                user_id=run_user_id,
+                kind="file",
+                data=data,
+                filename=filename,
+                mime_type=mime_type,
+                title=filename,
+                workflow_run_id=self._engine.workflow_run_id,
+            )
+            if ref is None:
+                dropped.append(f'Document "{filename}" could not be stored and was skipped.')
+                continue
+            refs.append(
+                {
+                    "artifact_id": ref["artifact_id"],
+                    "ref": ref.get("ref"),
+                    "filename": ref["filename"],
+                    "mime_type": ref["mime_type"],
+                }
+            )
+        return refs, dropped
+
+    @staticmethod
+    def _read_attachment_bytes(storage: Any, upload_path: str, max_bytes: int) -> Optional[bytes]:
+        """Read an attachment with a bounded read and a guaranteed handle close; None on failure."""
+        try:
+            file_obj = storage.get_file(upload_path)
+        except Exception as exc:
+            logger.error("Failed to open attachment for workflow run: %s", type(exc).__name__)
+            return None
+        try:
+            return file_obj.read(max_bytes + 1) if max_bytes else file_obj.read()
+        except Exception as exc:
+            logger.error("Failed to read attachment for workflow run: %s", type(exc).__name__)
+            return None
+        finally:
+            close = getattr(file_obj, "close", None)
+            if callable(close):
+                close()
+
+    def _finalize_workflow_run(
+        self,
+        workflow_owner_id: Optional[str],
+        run_user_id: Optional[str],
+        pg_workflow_id: Optional[str],
+        query: str,
+        interrupted: bool = False,
+    ) -> None:
+        """Write the run's terminal status/result; upsert the row if pre-creation was skipped.
+
+        The run is owned by the *runner* (so it stays readable to the caller and
+        matches the pre-created row); the workflow row is resolved by its *owner*.
+        When ``interrupted`` is set (client disconnect / mid-run error), the run is
+        recorded as FAILED regardless of the per-node log, so a partial run is never
+        left looking complete.
+        """
+        if not self._engine:
+            return
+        try:
+            status = ExecutionStatus.FAILED if interrupted else self._determine_run_status()
+            run = WorkflowRun(
+                workflow_id=self.workflow_id or "unknown",
+                user=run_user_id,
+                status=status,
+                inputs={"query": query},
+                outputs=self._serialize_state(self._engine.state),
+                steps=self._engine.get_execution_summary(),
+                created_at=datetime.now(timezone.utc),
+                completed_at=datetime.now(timezone.utc),
+            )
+            steps_json = [step.model_dump(mode="json") for step in run.steps]
+
+            if not self.workflow_id or not workflow_owner_id or not run_user_id:
+                return
+            with db_session() as conn:
+                if pg_workflow_id is None:
+                    pg_workflow_id = self._resolve_owned_workflow_pg_id(conn, workflow_owner_id)
+                    if pg_workflow_id is None:
+                        return
+                runs_repo = WorkflowRunsRepository(conn)
+                updated = False
+                if self._run_persisted:
+                    updated = runs_repo.finalize(
+                        self._engine.workflow_run_id,
+                        run_user_id,
+                        run.status.value,
+                        result=run.outputs,
+                        steps=steps_json,
+                        ended_at=run.completed_at,
+                    )
+                    if not updated:
+                        logger.warning(
+                            "Workflow run %s finalize matched no row; "
+                            "recovering via insert so terminal data is not lost",
+                            self._engine.workflow_run_id,
+                        )
+                if not self._run_persisted or not updated:
+                    runs_repo.create(
+                        pg_workflow_id,
+                        run_user_id,
+                        run.status.value,
+                        run_id=self._engine.workflow_run_id,
+                        inputs=run.inputs,
+                        result=run.outputs,
+                        steps=steps_json,
+                        started_at=run.created_at,
+                        ended_at=run.completed_at,
+                    )
+        except Exception as e:
+            logger.error(f"Failed to save workflow run: {e}")
+
+    def _determine_run_status(self) -> ExecutionStatus:
+        # A fatal input-document bridge failure (quota) means the engine never ran;
+        # the run is FAILED even though there is no per-node failure log entry.
+        if self._bridge_error is not None:
+            return ExecutionStatus.FAILED
+        if not self._engine or not self._engine.execution_log:
+            return ExecutionStatus.COMPLETED
+        for log in self._engine.execution_log:
+            if log.get("status") == ExecutionStatus.FAILED.value:
+                return ExecutionStatus.FAILED
+        return ExecutionStatus.COMPLETED
+
+    def _serialize_state(self, state: Dict[str, Any]) -> Dict[str, Any]:
+        serialized: Dict[str, Any] = {}
+        for key, value in state.items():
+            serialized[key] = self._serialize_state_value(value)
+        return serialized
+
+    def _serialize_state_value(self, value: Any) -> Any:
+        if isinstance(value, dict):
+            return {str(dict_key): self._serialize_state_value(dict_value) for dict_key, dict_value in value.items()}
+        if isinstance(value, list):
+            return [self._serialize_state_value(item) for item in value]
+        if isinstance(value, tuple):
+            return [self._serialize_state_value(item) for item in value]
+        if isinstance(value, datetime):
+            return value.isoformat()
+        if isinstance(value, (str, int, float, bool, type(None))):
+            return value
+        return str(value)