dograh/api/mcp_server/tools/save_workflow.py

"""MCP tool that accepts LLM-authored SDK TypeScript and saves it as a draft.

Execution flow:
    1. Parse via the Node TS validator — AST-only, never executes the code.
       Returns either a workflow JSON or per-location parse/validate errors.
    2. Pydantic validation via `ReactFlowDTO.model_validate` (defence in
       depth; the parser is already spec-driven, but the DTO layer is the
       authoritative wire-format gate).
    3. Graph validation via `WorkflowGraph`.
    4. Save as a new draft via `db_client.save_workflow_draft` — the
       published version stays intact, so edits are rollback-safe.

Error codes surfaced to the LLM:
    parse_error       — TS parse failed or a disallowed construct was used
    validation_error  — node data failed spec validation (unknown field,
                        missing required, wrong type, option out of range)
    schema_validation — ReactFlowDTO Pydantic rejection (rare; parser bug)
    graph_validation  — semantic graph rule broken (e.g. no start node)
    bridge_error      — Node subprocess failed before returning JSON

All LLM-facing errors include file:line:column where available so the
LLM can correct its code directly.
"""

from __future__ import annotations

from typing import Any

from fastapi import HTTPException
from loguru import logger
from pydantic import ValidationError as PydanticValidationError

from api.db import db_client
from api.mcp_server.auth import authenticate_mcp_request
from api.mcp_server.server import mcp
from api.mcp_server.tracing import traced_tool
from api.mcp_server.ts_bridge import TsBridgeError, parse_code
from api.services.workflow.dto import ReactFlowDTO
from api.services.workflow.layout import reconcile_positions
from api.services.workflow.workflow import WorkflowGraph


async def _previous_workflow_json(workflow: Any) -> dict[str, Any] | None:
    """Same selection priority as `get_workflow_code` — the version the
    LLM saw is the version we reconcile against.

    `current_definition` (is_current=True) is the published row, so the
    draft must be fetched explicitly. If no draft exists (e.g. the last
    draft was just published), fall through to `released_definition`.
    """
    draft = await db_client.get_draft_version(workflow.id)
    if draft is not None and draft.workflow_json:
        return draft.workflow_json
    released = workflow.released_definition
    if released is not None and released.workflow_json:
        return released.workflow_json
    return workflow.workflow_definition or None


def _error_result(code: str, message: str, **extra: Any) -> dict[str, Any]:
    return {"saved": False, "error_code": code, "error": message, **extra}


def _format_errors(errors: list[dict[str, Any]]) -> str:
    parts: list[str] = []
    for e in errors:
        loc = ""
        line = e.get("line")
        col = e.get("column")
        if line is not None:
            loc = f" (line {line}" + (f", col {col}" if col is not None else "") + ")"
        parts.append(f"{e.get('message', '')}{loc}")
    return "\n".join(parts)


@mcp.tool
@traced_tool
async def save_workflow(workflow_id: int, code: str) -> dict[str, Any]:
    """Parse SDK TypeScript and save the resulting workflow as a draft.

    `code` is TypeScript source using `@dograh/sdk`. Fetch the current
    code first via `get_workflow_code(workflow_id)`, edit it, then pass
    the full updated source here.

    Example code:
        import { Workflow } from "@dograh/sdk";
        import { startCall, endCall } from "@dograh/sdk/typed";

        const wf = new Workflow({ name: "lead_qualification" });
        const greeting = wf.addTyped(startCall({ name: "Greeting", prompt: "Hi!" }));
        const done     = wf.addTyped(endCall({ name: "Done", prompt: "Bye." }));
        wf.edge(greeting, done, { label: "done", condition: "conversation complete" });

    On success the draft version is saved; the published version is
    untouched.
    """
    user = await authenticate_mcp_request()

    workflow = await db_client.get_workflow(
        workflow_id, organization_id=user.selected_organization_id
    )
    if not workflow:
        raise HTTPException(status_code=404, detail=f"Workflow {workflow_id} not found")

    # 1. Parse + spec-validate via the Node TS validator.
    try:
        parsed = await parse_code(code)
    except TsBridgeError as e:
        logger.warning(f"ts_bridge failure: {e}")
        return _error_result("bridge_error", str(e))

    if not parsed.get("ok"):
        stage = parsed.get("stage", "parse")
        errs = parsed.get("errors") or []
        code_key = "parse_error" if stage == "parse" else "validation_error"
        return _error_result(code_key, _format_errors(errs), errors=errs)

    payload = parsed["workflow"]
    new_name = (parsed.get("workflowName") or "").strip()

    # 1b. Reconcile node positions against the previously-stored workflow.
    # The parser drops positions by design (LLMs don't place nodes well);
    # here we fill them back in from what was there before, and pick
    # approximate placements for newly-introduced nodes.
    payload = reconcile_positions(payload, await _previous_workflow_json(workflow))

    # 2. Pydantic shape check (defence in depth — parser is spec-driven).
    try:
        dto = ReactFlowDTO.model_validate(payload)
    except PydanticValidationError as e:
        return _error_result("schema_validation", str(e))

    # 3. Graph-level semantic validation (start-node count, edge shape).
    try:
        WorkflowGraph(dto)
    except (ValueError, Exception) as e:  # WorkflowGraph raises ValueError
        return _error_result("graph_validation", str(e))

    # 4a. If the `new Workflow({ name })` in the edited source differs from
    # the stored name, rename the workflow. Name is a workflow-level field
    # (not versioned), so this takes effect immediately.
    name_changed = bool(new_name) and new_name != workflow.name
    if name_changed:
        await db_client.update_workflow(
            workflow_id=workflow_id,
            name=new_name,
            workflow_definition=None,
            template_context_variables=None,
            workflow_configurations=None,
            organization_id=user.selected_organization_id,
        )

    # 4b. Save as a new draft (existing published version stays intact).
    draft = await db_client.save_workflow_draft(
        workflow_id=workflow_id,
        workflow_definition=payload,
    )

    return {
        "saved": True,
        "workflow_id": workflow_id,
        "version_number": draft.version_number,
        "status": draft.status,
        "node_count": len(payload["nodes"]),
        "edge_count": len(payload["edges"]),
        "name": new_name or workflow.name,
        "renamed": name_changed,
    }