DocsGPT/docsgpt/agents/workflow_agent.py

import logging
from datetime import datetime, timezone
from typing import Any, Dict, Generator, List, Optional, Tuple

from docsgpt.agents.base import BaseAgent
from docsgpt.guardrails.config import (
    DEFAULT_BLOCK_MESSAGE as GUARDRAIL_DEFAULT_MESSAGE,
)
from docsgpt.agents.workflows.schemas import (
    ExecutionStatus,
    Workflow,
    WorkflowEdge,
    WorkflowGraph,
    WorkflowNode,
    WorkflowRun,
)
from docsgpt.agents.workflows.workflow_engine import WorkflowEngine
from docsgpt.core.settings import settings
from docsgpt.logging import LogContext, log_activity
from docsgpt.sandbox.artifacts_capture import QuotaExceeded
from docsgpt.storage.db.base_repository import looks_like_uuid
from docsgpt.storage.db.repositories.workflow_edges import WorkflowEdgesRepository
from docsgpt.storage.db.repositories.workflow_nodes import WorkflowNodesRepository
from docsgpt.storage.db.repositories.workflow_runs import WorkflowRunsRepository
from docsgpt.storage.db.repositories.workflows import WorkflowsRepository
from docsgpt.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]:
        # This override skips BaseAgent.gen, so the input stage has to be run
        # here or a workflow agent would accept guardrail config in the builder
        # and silently enforce none of it.
        self.bind_guardrail_log_context(log_context)
        query, decision = self.apply_input_guardrails(query)
        if decision is not None and decision.blocked:
            yield self._guardrail_block_event(
                decision, decision.block_message or GUARDRAIL_DEFAULT_MESSAGE
            )
            self.flush_guardrail_audit()
            return
        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 docsgpt.sandbox.artifacts_capture import persist_new_artifact
        from docsgpt.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"],
                }
            )
            self._record_preextracted_text(ref["artifact_id"], attachment)
        return refs, dropped

    def _record_preextracted_text(self, artifact_id: str, attachment: Dict[str, Any]) -> None:
        """Stage an attachment's upload-time text on the engine so nodes skip re-parsing it."""
        if self._engine is None:
            return
        text = self._usable_attachment_text(attachment)
        if text:
            self._engine.preextracted_text[str(artifact_id)] = text

    @staticmethod
    def _usable_attachment_text(attachment: Dict[str, Any]) -> Optional[str]:
        """Return an attachment's stored text when it is safe to reuse verbatim, else None.

        The attachment was already parsed (and possibly OCR'd) by ``store_attachment``;
        reusing that text is what keeps a scanned document from being OCR'd once per
        node. Reuse requires the stored text to be the WHOLE document: a truncated row
        was cut head-first at ``ATTACHMENT_MAX_TOKENS`` and would silently lose the tail
        that the workflow's own head+tail bounding keeps, so it falls back to a parse.
        A row with no ``extraction`` metadata at all (legacy, pre-dating that field) is
        trusted when its ``content`` is non-empty: nothing records otherwise, and the
        non-workflow chat path already feeds that same text to the model.
        """
        content = attachment.get("content")
        if not isinstance(content, str) or not content.strip():
            return None
        metadata = attachment.get("metadata")
        extraction = metadata.get("extraction") if isinstance(metadata, dict) else None
        if extraction is None:
            return content
        if not isinstance(extraction, dict):
            return None
        if extraction.get("status") != "ok" or extraction.get("truncated"):
            return None
        return content

    @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)