Files
DocsGPT/docsgpt/api/user/agents/portability.py
T
2026-09-29 10:42:39 +04:00

1842 lines
74 KiB
Python

"""Agent YAML export / import.
Agents are mostly references (sources, tools, prompt, models). Export
translates those into portable identifiers and never emits a secret;
import resolves them back against the importing user — matching existing
resources or creating them from the file plus user-supplied tokens.
Re-import is idempotent at every level: the agent is matched by
``metadata.id`` then ``metadata.slug``, tools by ``(type, name)``,
prompts by ``(name, content)``, custom models by
``(display_name, upstream_model_id, base_url)`` — so the same file
applied twice updates rather than duplicating.
Workflow agents additionally carry their graph under ``spec.workflow``.
The graph's agent nodes reference tools, sources and models by raw id in
their config; export rewrites those into the same portable identities the
top-level sections use (``tool-N`` keys into ``spec.tools``, source names,
custom-model display names), and import resolves them back and writes the
graph through the same validation gate as the workflow API.
"""
from __future__ import annotations
import copy
import re
from typing import Any, Optional
import yaml
from flask import current_app, jsonify, make_response, request
from flask_restx import Namespace, Resource
from docsgpt.agents.default_tools import (
default_tool_id,
is_synthesized_tool_id,
synthesize_tool_by_name,
synthesized_tool_name_for_id,
)
from docsgpt.api import api
from docsgpt.api.pat.rules import allowed_ids
from docsgpt.api.user.resource_access import AccessDenied, require
from docsgpt.core.model_utils import validate_model_id
from docsgpt.core.url_validation import SSRFError, validate_url
from docsgpt.security.safe_url import UnsafeUserUrlError, validate_user_base_url
from docsgpt.storage.db.base_repository import looks_like_uuid
from docsgpt.storage.db.repositories.agents import AgentsRepository
from docsgpt.storage.db.repositories.prompts import PromptsRepository
from docsgpt.storage.db.repositories.sources import SourcesRepository
from docsgpt.storage.db.repositories.user_custom_models import (
UserCustomModelsRepository,
)
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
from docsgpt.storage.db.repositories.workflow_edges import WorkflowEdgesRepository
from docsgpt.storage.db.repositories.workflow_nodes import WorkflowNodesRepository
from docsgpt.storage.db.repositories.workflows import WorkflowsRepository
from docsgpt.storage.db.session import db_readonly, db_session
API_VERSION = "docsgpt.arc53.com/v1"
KIND = "Agent"
# Import safety bounds (DoS): cap the document size and the number of
# child resources a single import may reference / create.
MAX_IMPORT_BYTES = 512 * 1024
MAX_LIST_ITEMS = 100
MAX_WORKFLOW_NODES = 200
MAX_WORKFLOW_EDGES = 400
class AgentImportError(Exception):
"""Raised when an agent YAML document is malformed or unsupported."""
class AgentExportError(Exception):
"""Raised when an agent cannot be exported as a document import would accept."""
class _SafeNoAliasLoader(yaml.SafeLoader):
"""SafeLoader that also rejects anchors/aliases (billion-laughs guard)."""
def compose_node(self, parent, index):
if self.check_event(yaml.events.AliasEvent):
raise AgentImportError("YAML anchors/aliases are not allowed")
return super().compose_node(parent, index)
# ---------------------------------------------------------------------------
# Slug helpers
# ---------------------------------------------------------------------------
def slugify(name: Optional[str]) -> str:
"""Turn an agent name into a url-safe slug; never empty."""
s = (name or "").strip().lower()
s = re.sub(r"[^a-z0-9]+", "-", s)
s = re.sub(r"-{2,}", "-", s).strip("-")
return s or "agent"
def _unique_slug(
repo: AgentsRepository,
user: str,
base: str,
*,
exclude_id: Optional[str] = None,
) -> str:
"""Return ``base`` (slugified) or the first ``-N`` variant free for ``user``."""
base = slugify(base)
candidate = base
n = 2
while True:
found = repo.find_by_slug(user, candidate)
if found is None or (exclude_id and str(found["id"]) == str(exclude_id)):
return candidate
candidate = f"{base}-{n}"
n += 1
def ensure_agent_slug(conn, agent: dict, user: str) -> str:
"""Return the agent's slug, lazily assigning and persisting one if absent."""
existing = agent.get("slug")
if existing:
return str(existing)
repo = AgentsRepository(conn)
slug = _unique_slug(repo, user, agent.get("name"), exclude_id=str(agent["id"]))
repo.update(str(agent["id"]), user, {"slug": slug})
return slug
# ---------------------------------------------------------------------------
# Tool-manager access (lazy — avoids importing the tool registry at module load)
# ---------------------------------------------------------------------------
def _tool_manager():
from docsgpt.api.user.tools.routes import tool_manager
return tool_manager
def _tool_instance(tool_type: str):
return _tool_manager().tools.get(tool_type)
def _secret_field_names(config_requirements: dict, *, required_only: bool = False) -> list:
out = []
for key, spec in (config_requirements or {}).items():
if not isinstance(spec, dict) or not spec.get("secret"):
continue
if required_only and not spec.get("required"):
continue
out.append(key)
return out
def _live_requires_secrets(tool_type: str) -> list:
inst = _tool_instance(tool_type)
if inst is None:
return []
return _secret_field_names(inst.get_config_requirements() or {})
def _safe_export_config(stored_config: dict, config_requirements: dict) -> dict:
"""Return only config values provably non-secret (allowlist).
Export must never leak a credential. A key is emitted ONLY if
``config_requirements`` declares it and does not mark it secret.
Anything not described by requirements — e.g. MCP free-form ``headers``,
``url`` query strings, ``server_url`` that routinely embed tokens, or
tools with empty ``config_requirements`` — is dropped. Reuse on import
matches the existing tool; create re-collects config from the user. This
blocklist-free approach can't be defeated by a stale or empty
``config_requirements``. (``api_tool`` is the one exception: its
user-authored ``config["actions"]`` travel via
``_export_api_tool_actions``, redacted under the same posture.)
"""
requirements = config_requirements or {}
safe: dict[str, Any] = {}
for key, value in (stored_config or {}).items():
spec = requirements.get(key)
if isinstance(spec, dict) and not spec.get("secret"):
safe[key] = value
return safe
# Param-type sections an action may carry. Fixed ``value``s are exported only
# for ``parameters`` (declared function args, treated as non-secret throughout:
# returned to the client, shown in the actions editor). The free-form HTTP
# types routinely hold credentials (an ``Authorization`` header fixed as a
# value), so for those only the ``filled_by_llm`` flag travels — same posture
# as ``_safe_export_config`` dropping free-form config.
#
# Scope note: ``actions``-column rows come from ``transform_actions`` over live
# tool metadata, which today only ever emits ``parameters`` sections — the
# HTTP types appear only in ``api_tool``'s ``config["actions"]``, which travel
# through ``_export_api_tool_actions`` / ``_sanitize_api_tool_actions`` below.
# The full tuple is kept here so a hand-edited file naming those sections is
# handled (flags applied, values refused) rather than silently ignored.
_ACTION_PARAM_TYPES = ("parameters", "query_params", "headers", "body")
_ACTION_VALUE_SAFE_TYPES = frozenset({"parameters"})
# Action keys copied verbatim by ``_export_api_tool_actions`` — everything
# except ``url`` (query-string-stripped) and the param sections (redacted).
_API_ACTION_SAFE_KEYS = (
"name",
"description",
"method",
"active",
"require_approval",
"body_content_type",
"body_encoding_rules",
)
def _export_api_tool_actions(stored_config: dict) -> dict:
"""Redacted copy of an ``api_tool``'s ``config["actions"]``.
``api_tool`` is the one tool type whose actions are user-authored data in
``config["actions"]`` (a dict keyed by action name) rather than the
``actions`` column (the executor reads exactly that — see ToolExecutor).
The action *shape* — endpoint, method, param schemas, ``filled_by_llm``
flags, fixed ``parameters`` values — is what makes the tool work, so it
travels. Fixed values in the free-form HTTP sections and URL query
strings routinely hold credentials, so those are blanked — same posture
as ``_safe_export_config``.
"""
actions = stored_config.get("actions")
if not isinstance(actions, dict):
return {}
out: dict[str, Any] = {}
for name, action in actions.items():
if not isinstance(action, dict):
continue
entry: dict[str, Any] = {
k: copy.deepcopy(action[k]) for k in _API_ACTION_SAFE_KEYS if k in action
}
entry["url"] = str(action.get("url") or "").split("?", 1)[0]
for param_type in _ACTION_PARAM_TYPES:
section = action.get(param_type)
if not isinstance(section, dict):
continue
section = copy.deepcopy(section)
props = section.get("properties")
if param_type not in _ACTION_VALUE_SAFE_TYPES and isinstance(props, dict):
for details in props.values():
if isinstance(details, dict) and details.get("value"):
details["value"] = ""
entry[param_type] = section
out[str(name)] = entry
return out
def _sanitize_api_tool_actions(raw, tool_label: str, warnings: list) -> dict:
"""Vet a hand-editable ``config["actions"]`` block before storing it.
Unlike column actions — whose schema comes from live tool code, so the
YAML may contribute flags and scalar values only — ``api_tool`` actions
ARE user-authored schema, and the whole definition is accepted. Each
action's URL goes through the same SSRF gate as the tool-config routes
(``_validate_tool_urls``); an action with a missing or unsafe URL is
dropped with a warning rather than failing the import.
"""
if not isinstance(raw, dict):
return {}
if len(raw) > MAX_LIST_ITEMS:
raise AgentImportError(f"Too many api_tool actions (max {MAX_LIST_ITEMS})")
out: dict[str, Any] = {}
for name, action in raw.items():
if not isinstance(action, dict):
continue
url = str(action.get("url") or "").strip()
if not url:
warnings.append(
f"Tool '{tool_label}' action '{name}' has no URL; dropped"
)
continue
try:
validate_url(url)
except SSRFError:
warnings.append(
f"Tool '{tool_label}' action '{name}' has an unsafe URL; dropped"
)
continue
# Export blanks fixed values in the free-form HTTP sections (they
# routinely hold credentials); a param left pinned with an empty
# value is silently omitted from the request at run time, so
# surface it now instead.
pinned_empty = [
f"{param_type}.{pname}"
for param_type in _ACTION_PARAM_TYPES
if param_type not in _ACTION_VALUE_SAFE_TYPES
and isinstance(action.get(param_type), dict)
for pname, details in (action[param_type].get("properties") or {}).items()
if isinstance(details, dict)
and details.get("filled_by_llm") is False
and not details.get("value")
]
if pinned_empty:
warnings.append(
f"Tool '{tool_label}' action '{name}': fixed values for "
f"{', '.join(pinned_empty)} don't travel in the file; set them "
"in the tool's settings or the request will omit them"
)
out[str(name)] = action
return out
def _export_action_overrides(actions) -> list:
"""Compact per-action customizations that deviate from factory defaults.
``transform_actions`` initializes every action ``active`` with every param
``filled_by_llm: true, value: ""`` — anything else is a user customization
(a fixed ``chat_id``, a disabled action) that must survive export, or the
imported tool silently reverts to "let the LLM fill everything".
"""
out = []
for action in actions or []:
if not isinstance(action, dict) or not action.get("name"):
continue
override: dict[str, Any] = {"name": action["name"]}
customized = False
if action.get("active", True) is False:
override["active"] = False
customized = True
for param_type in _ACTION_PARAM_TYPES:
section = action.get(param_type)
props = section.get("properties") if isinstance(section, dict) else None
if not isinstance(props, dict):
continue
for pname, details in props.items():
if not isinstance(details, dict):
continue
filled = details.get("filled_by_llm", True)
value = details.get("value", "")
if filled and not value:
continue
entry: dict[str, Any] = {"filled_by_llm": bool(filled)}
if param_type in _ACTION_VALUE_SAFE_TYPES:
entry["value"] = value
override.setdefault(param_type, {})[pname] = entry
customized = True
if customized:
out.append(override)
return out
def _apply_action_overrides(base_actions: list, overrides, tool_type: str, warnings: list) -> list:
"""Overlay exported customizations onto freshly transformed live metadata.
Matched strictly by (action name, param-type, param name): the YAML
contributes flags and scalar values only, never schema — a hand-edited
file can't inject a doctored action definition — and the merge survives
tool-version drift. Unknown names are skipped with a warning. Values are
accepted only for the ``parameters`` type, mirroring export.
"""
if not isinstance(overrides, list) or not overrides:
return base_actions
by_name = {a.get("name"): a for a in base_actions if isinstance(a, dict)}
for override in overrides:
if not isinstance(override, dict):
continue
name = override.get("name") or ""
action = by_name.get(name)
if action is None:
warnings.append(
f"Tool '{tool_type}' action '{name}' no longer exists; customization skipped"
)
continue
if override.get("active") is False:
action["active"] = False
for param_type in _ACTION_PARAM_TYPES:
override_props = override.get(param_type)
if not isinstance(override_props, dict):
continue
section = action.get(param_type)
props = section.get("properties") if isinstance(section, dict) else {}
for pname, entry in override_props.items():
details = props.get(pname) if isinstance(props, dict) else None
if not isinstance(details, dict) or not isinstance(entry, dict):
warnings.append(
f"Tool '{tool_type}' action '{name}' has no parameter "
f"'{pname}'; customization skipped"
)
continue
details["filled_by_llm"] = bool(entry.get("filled_by_llm", True))
if param_type in _ACTION_VALUE_SAFE_TYPES:
value = entry.get("value", "")
if isinstance(value, (str, int, float, bool)) or value is None:
details["value"] = "" if value is None else value
elif not details["filled_by_llm"] and not details.get("value"):
# Pinned but valueless: the executor omits the param from
# the LLM schema AND skips the empty fixed value, so the
# request silently goes out without it. Values for the
# free-form HTTP types never travel in the file, so this
# can only be repaired in the tool's settings.
warnings.append(
f"Tool '{tool_type}' action '{name}': fixed value for "
f"{param_type} '{pname}' doesn't travel in the file; set "
"it in the tool's settings or the request will omit it"
)
return base_actions
# ---------------------------------------------------------------------------
# Export
# ---------------------------------------------------------------------------
def _serialize_prompt(conn, agent: dict, user: str):
prompt_id = agent.get("prompt_id")
if not prompt_id:
return "default"
row = PromptsRepository(conn).get(str(prompt_id), user)
if not row:
return "default"
return {"name": row.get("name") or "", "content": row.get("content") or ""}
def _serialize_sources(conn, agent: dict, user: str) -> list:
ids: list[str] = []
if agent.get("source_id"):
ids.append(str(agent["source_id"]))
for sid in agent.get("extra_source_ids") or []:
s = str(sid)
if s and s not in ids:
ids.append(s)
repo = SourcesRepository(conn)
out = []
for sid in ids:
row = repo.get(sid, user)
if not row:
continue
out.append(
{
"name": row.get("name") or "",
"type": row.get("type") or "",
"ref": sid,
}
)
return out
def _serialize_tools(conn, agent: dict, user: str) -> list:
repo = UserToolsRepository(conn)
out = []
for tid in agent.get("tools") or []:
tid_str = str(tid)
if is_synthesized_tool_id(tid_str):
out.append(
{
"type": synthesized_tool_name_for_id(tid_str),
"builtin": True,
"ref": tid_str,
}
)
continue
row = repo.get_any(tid_str, user)
if not row:
continue
requirements = row.get("config_requirements") or {}
requires_secrets = _secret_field_names(requirements)
config = _safe_export_config(row.get("config") or {}, requirements)
if (row.get("name") or "") == "api_tool":
# api_tool's actions live in config, not the actions column — an
# export without them imports as a dead tool.
exported_actions = _export_api_tool_actions(row.get("config") or {})
if exported_actions:
config["actions"] = exported_actions
entry: dict[str, Any] = {
"type": row.get("name") or "",
# Raw custom_name (may be "") so import matching is symmetric and
# idempotent — see _find_user_tool.
"name": row.get("custom_name") or "",
"display_name": row.get("display_name") or "",
"description": row.get("description") or "",
"config": config,
"ref": tid_str,
}
action_overrides = _export_action_overrides(row.get("actions"))
if action_overrides:
entry["actions"] = action_overrides
if requires_secrets:
entry["requires_secrets"] = requires_secrets
out.append(entry)
return out
def _serialize_models(conn, agent: dict, user: str) -> dict:
repo = UserCustomModelsRepository(conn)
def describe(model_id: str):
mid = str(model_id)
if looks_like_uuid(mid):
row = repo.get(mid, user)
if not row:
return None
return {
"type": "custom",
"display_name": row.get("display_name") or "",
"upstream_model_id": row.get("upstream_model_id") or "",
"base_url": row.get("base_url") or "",
"capabilities": row.get("capabilities") or {},
"requires_secrets": ["api_key"],
"ref": mid,
}
return mid
available = []
for m in agent.get("models") or []:
described = describe(m)
if described is not None:
available.append(described)
default_value = ""
raw_default = agent.get("default_model_id") or ""
if raw_default:
described = describe(raw_default)
if isinstance(described, dict):
default_value = described["display_name"]
elif isinstance(described, str):
default_value = described
return {"default": default_value, "available": available}
def _effective_node_config(config: dict) -> dict:
"""Return the dict a node's settings actually live in.
The builder writes some node configs nested under ``config`` while other
payloads keep them flat; the engine reads
``node.config.get("config", node.config)``. Mirror that here so reference
rewriting touches the dict the engine will read.
"""
nested = config.get("config")
return nested if isinstance(nested, dict) else config
def _strip_shadowed_ref_keys(config: dict) -> None:
"""Drop flat ``tools``/``sources``/``model_id`` shadowed by a nested config.
When a node stores its settings under ``config["config"]``, the engine
reads only the nested dict — flat leftovers of those keys are inert, but
on export they would leak the exporter's raw ids into the file, and on
import they would store a hand-edited file's raw ids verbatim.
"""
if isinstance(config.get("config"), dict):
for key in ("tools", "sources", "model_id"):
config.pop(key, None)
def _serialize_workflow(conn, agent: dict, user: str):
"""Build the portable ``spec.workflow`` block plus its reference sections.
Agent nodes hold raw tool/source/model ids; those are collected (deduped,
in first-appearance order) into the same top-level ``tools`` / ``sources``
/ ``model`` sections a classic export uses, and the node configs are
rewritten to portable keys: ``tool-N``, the source's name, a custom
model's display name. Returns ``(workflow_block, tools, sources, model)``;
the block is None when the agent has no resolvable workflow.
"""
empty_model = {"default": "", "available": []}
workflow_id = agent.get("workflow_id")
if not workflow_id:
return None, [], [], empty_model
row = WorkflowsRepository(conn).get(str(workflow_id), user)
if row is None:
return None, [], [], empty_model
# Lazy import, same as _apply_workflow — keeps route modules decoupled at load.
from docsgpt.api.user.workflows.routes import get_workflow_graph_version
version = get_workflow_graph_version(row)
node_rows = WorkflowNodesRepository(conn).find_by_version(str(row["id"]), version)
edge_rows = WorkflowEdgesRepository(conn).find_by_version(str(row["id"]), version)
# Pass 1: collect referenced ids across agent nodes, deduped and ordered.
tool_ids: list[str] = []
source_ids: list[str] = []
model_ids: list[str] = []
for n in node_rows:
if n.get("node_type") != "agent":
continue
cfg = _effective_node_config(n.get("config") or {})
for tid in cfg.get("tools") or []:
if str(tid) and str(tid) not in tool_ids:
tool_ids.append(str(tid))
for sid in cfg.get("sources") or []:
if str(sid) and str(sid) not in source_ids:
source_ids.append(str(sid))
mid = cfg.get("model_id")
if mid and str(mid) not in model_ids:
model_ids.append(str(mid))
# Serialize the references through the classic per-resource serializers so
# secret redaction and identity rules stay in one place. Unresolvable ids
# (dangling rows) are dropped from both the section and the node config;
# every emitted entry carries its ``ref``, which keys the maps below.
tools: list = []
tool_key_by_id: dict[str, str] = {}
for entry in _serialize_tools(conn, {"tools": tool_ids}, user):
tool_key_by_id[str(entry["ref"])] = _tool_key(len(tools))
tools.append(entry)
sources: list = []
source_name_by_id: dict[str, str] = {}
for entry in _serialize_sources(conn, {"extra_source_ids": source_ids}, user):
source_name_by_id[str(entry["ref"])] = entry["name"]
sources.append(entry)
model = _serialize_models(conn, {"models": model_ids, "default_model_id": ""}, user)
model_name_by_id: dict[str, str] = {}
for entry in model["available"]:
if isinstance(entry, dict):
model_name_by_id[str(entry.get("ref") or "")] = entry.get("display_name") or ""
elif isinstance(entry, str):
model_name_by_id[entry] = entry
# ``parse_agent_yaml`` bounds these sections at import (DoS cap), so a
# graph referencing more would export a file that import then rejects —
# refuse here with a clear message instead.
if (
len(tools) > MAX_LIST_ITEMS
or len(sources) > MAX_LIST_ITEMS
or len(model["available"]) > MAX_LIST_ITEMS
):
raise AgentExportError(
"This workflow references more than "
f"{MAX_LIST_ITEMS} tools, sources, or models and cannot be exported"
)
# Pass 2: emit nodes/edges with agent node configs rewritten to the
# portable keys. (Node settings sit under ``config`` here; the workflow
# API's wire shape calls that field ``data`` — import converts.)
nodes_out = []
for n in node_rows:
config = copy.deepcopy(n.get("config") or {})
if n.get("node_type") == "agent":
cfg = _effective_node_config(config)
cfg["tools"] = [
tool_key_by_id[str(t)]
for t in cfg.get("tools") or []
if str(t) in tool_key_by_id
]
cfg["sources"] = [
source_name_by_id[str(s)]
for s in cfg.get("sources") or []
if str(s) in source_name_by_id
]
if cfg.get("model_id"):
cfg["model_id"] = model_name_by_id.get(str(cfg["model_id"]))
_strip_shadowed_ref_keys(config)
nodes_out.append(
{
"id": n["node_id"],
"type": n["node_type"],
"title": n.get("title") or "",
"description": n.get("description") or "",
"position": n.get("position") or {"x": 0, "y": 0},
"config": config,
}
)
edges_out = [
{
"id": e["edge_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
]
block = {
"name": row.get("name") or "",
"description": row.get("description") or "",
"nodes": nodes_out,
"edges": edges_out,
}
return block, tools, sources, model
def _import_config(spec: dict) -> dict:
"""Validate a spec's ``config`` through the same gate as the API.
A YAML is hand-editable, so it must not be a way to install a control the
write path would have rejected. An invalid block is dropped rather than
failing the whole import, and the caller surfaces it as a warning.
"""
from docsgpt.api.user.agents.routes import normalize_agent_config
try:
return normalize_agent_config(spec.get("config")) or {}
except ValueError:
current_app.logger.warning(
"Dropping invalid guardrails config during agent import"
)
return {}
def serialize_agent(conn, agent: dict, user: str) -> dict:
"""Build the portable export document for an agent row."""
is_workflow = (agent.get("agent_type") or "") == "workflow"
if is_workflow:
# A workflow agent's behavior lives in its graph: the agent-level
# prompt/source/tool/model slots are unused, so the top-level sections
# instead carry the references collected from the graph's agent nodes.
workflow_block, tools, sources, model = _serialize_workflow(conn, agent, user)
prompt = "default"
else:
tools = _serialize_tools(conn, agent, user)
sources = _serialize_sources(conn, agent, user)
model = _serialize_models(conn, agent, user)
prompt = _serialize_prompt(conn, agent, user)
spec = {
"name": agent.get("name") or "",
"description": agent.get("description") or "",
"agent_type": agent.get("agent_type") or "classic",
"retriever": agent.get("retriever") or "classic",
"chunks": int(agent["chunks"]) if agent.get("chunks") is not None else None,
"prompt": prompt,
"model": model,
"sources": sources,
"tools": tools,
"limits": {
"limited_token_mode": bool(agent.get("limited_token_mode", False)),
"token_limit": agent.get("token_limit"),
"limited_request_mode": bool(agent.get("limited_request_mode", False)),
"request_limit": agent.get("request_limit"),
},
"json_schema": agent.get("json_schema"),
"allow_system_prompt_override": bool(agent.get("allow_system_prompt_override", False)),
"config": agent.get("config") or {},
}
if is_workflow:
spec["workflow"] = workflow_block
return {
"apiVersion": API_VERSION,
"kind": KIND,
"metadata": {"id": str(agent["id"]), "slug": agent.get("slug") or ""},
"spec": spec,
}
class _AgentYamlDumper(yaml.SafeDumper):
"""SafeDumper that renders multi-line strings as block scalars."""
def _str_representer(dumper, data):
style = "|" if "\n" in data else None
return dumper.represent_scalar("tag:yaml.org,2002:str", data, style=style)
_AgentYamlDumper.add_representer(str, _str_representer)
def agent_to_yaml(export: dict) -> str:
"""Serialize an export document to YAML, preserving key order."""
return yaml.dump(
export,
Dumper=_AgentYamlDumper,
sort_keys=False,
default_flow_style=False,
allow_unicode=True,
)
# ---------------------------------------------------------------------------
# Parse
# ---------------------------------------------------------------------------
def _validate_workflow_shape(spec: dict) -> None:
"""Shallow-validate ``spec.workflow`` (deep graph validation runs at apply).
A missing key or explicit null is allowed — a draft workflow agent may have
no graph yet; only a present block must be well-formed and bounded.
"""
workflow = spec.get("workflow")
if workflow is None:
return
if not isinstance(workflow, dict):
raise AgentImportError("spec.workflow must be a mapping")
nodes = workflow.get("nodes")
edges = workflow.get("edges")
if not isinstance(nodes, list) or not isinstance(edges, list):
raise AgentImportError("spec.workflow.nodes and spec.workflow.edges must be lists")
if len(nodes) > MAX_WORKFLOW_NODES:
raise AgentImportError(f"Too many workflow nodes (max {MAX_WORKFLOW_NODES})")
if len(edges) > MAX_WORKFLOW_EDGES:
raise AgentImportError(f"Too many workflow edges (max {MAX_WORKFLOW_EDGES})")
for node in nodes:
if not isinstance(node, dict) or not node.get("id") or not node.get("type"):
raise AgentImportError("Every workflow node needs an 'id' and a 'type'")
for edge in edges:
if not isinstance(edge, dict) or not edge.get("id"):
raise AgentImportError("Every workflow edge needs an 'id'")
def parse_agent_yaml(text: str) -> dict:
"""Parse and shallow-validate an agent YAML document."""
if text and len(text) > MAX_IMPORT_BYTES:
raise AgentImportError("Import document too large")
try:
doc = yaml.load(text, Loader=_SafeNoAliasLoader)
except yaml.YAMLError as exc:
raise AgentImportError(f"Invalid YAML: {exc}") from exc
if not isinstance(doc, dict):
raise AgentImportError("Top-level YAML must be a mapping")
if doc.get("kind") != KIND:
raise AgentImportError(f"Unsupported kind {doc.get('kind')!r}; expected {KIND!r}")
if not str(doc.get("apiVersion") or "").startswith("docsgpt."):
raise AgentImportError(f"Unsupported apiVersion {doc.get('apiVersion')!r}")
spec = doc.get("spec")
if not isinstance(spec, dict):
raise AgentImportError("Missing or invalid 'spec'")
if not spec.get("name"):
raise AgentImportError("spec.name is required")
if spec.get("agent_type") == "workflow":
_validate_workflow_shape(spec)
for key in ("sources", "tools"):
value = spec.get(key)
if isinstance(value, list) and len(value) > MAX_LIST_ITEMS:
raise AgentImportError(f"Too many {key} (max {MAX_LIST_ITEMS})")
available = (spec.get("model") or {}).get("available")
if isinstance(available, list) and len(available) > MAX_LIST_ITEMS:
raise AgentImportError(f"Too many models (max {MAX_LIST_ITEMS})")
return doc
# ---------------------------------------------------------------------------
# Reference resolution (shared by plan + apply)
# ---------------------------------------------------------------------------
def _resolve_target(conn, user: str, metadata: dict) -> dict:
repo = AgentsRepository(conn)
metadata = metadata or {}
agent_id = metadata.get("id")
if agent_id and looks_like_uuid(str(agent_id)):
row = repo.get(str(agent_id), user)
if row:
return {
"action": "update",
"agent_id": str(row["id"]),
"matched_by": "id",
"status": row.get("status"),
}
slug = metadata.get("slug")
if slug:
row = repo.find_by_slug(user, str(slug))
if row:
return {
"action": "update",
"agent_id": str(row["id"]),
"matched_by": "slug",
"status": row.get("status"),
}
return {"action": "create", "agent_id": None, "matched_by": None, "status": None}
def _find_user_tool(user_tools: list, tool_type: str, custom_name: str):
"""Match a stored tool by ``(type, custom_name)`` for idempotent reuse.
Compares against the stored ``custom_name`` directly (export emits the
raw value), so a named tool matches exactly and an unnamed one (``""``)
reuses the first same-type tool rather than spawning a duplicate.
"""
target = custom_name or ""
for row in user_tools:
if (row.get("name") or "") != tool_type:
continue
if (row.get("custom_name") or "") == target:
return row
return None
def _find_custom_model(customs: list, spec_model: dict):
for row in customs:
if (
(row.get("display_name") or "") == (spec_model.get("display_name") or "")
and (row.get("upstream_model_id") or "") == (spec_model.get("upstream_model_id") or "")
and (row.get("base_url") or "") == (spec_model.get("base_url") or "")
):
return row
return None
def _tool_key(index: int) -> str:
return f"tool-{index}"
# ---------------------------------------------------------------------------
# Plan (dry run)
# ---------------------------------------------------------------------------
def _plan_workflow_removal(conn, user: str, target: dict) -> Optional[dict]:
"""Describe what an explicit ``workflow: null`` would destroy, if anything.
Returns a ``delete`` block sized with the current graph, or None when the
workflow survives — the target is new, published (``_apply_workflow``
keeps a published agent's graph), has no workflow, or another agent still
references it.
"""
if target.get("action") != "update" or not target.get("agent_id"):
return None
if (target.get("status") or "") == "published":
return None
agent_row = AgentsRepository(conn).get(str(target["agent_id"]), user) or {}
existing_id = agent_row.get("workflow_id")
if not existing_id:
return None
row = WorkflowsRepository(conn).get(str(existing_id), user)
if row is None:
return None
if AgentsRepository(conn).count_by_workflow(str(existing_id), user) > 1:
return None
# Lazy import, same as _apply_workflow — keeps route modules decoupled at load.
from docsgpt.api.user.workflows.routes import get_workflow_graph_version
version = get_workflow_graph_version(row)
return {
"action": "delete",
"nodes": len(WorkflowNodesRepository(conn).find_by_version(str(row["id"]), version)),
"edges": len(WorkflowEdgesRepository(conn).find_by_version(str(row["id"]), version)),
}
def plan_import(conn, user: str, doc: dict) -> dict:
"""Resolve every reference without writing; returns a resolution report."""
spec = doc["spec"]
target = _resolve_target(conn, user, doc.get("metadata") or {})
sources_repo = SourcesRepository(conn)
sources = []
for src in spec.get("sources") or []:
name = src.get("name") or ""
match = sources_repo.find_by_name(user, name)
sources.append(
{
"name": name,
"type": src.get("type") or "",
"status": "matched" if match else "missing",
"target_id": str(match["id"]) if match else None,
}
)
user_tools = UserToolsRepository(conn).list_for_user(user)
tools = []
for index, tool in enumerate(spec.get("tools") or []):
tool_type = tool.get("type") or ""
if tool.get("builtin"):
available = synthesize_tool_by_name(tool_type) is not None
tools.append(
{
"key": _tool_key(index),
"type": tool_type,
"builtin": True,
"status": "builtin" if available else "unavailable",
"target_id": default_tool_id(tool_type) if available else None,
}
)
continue
custom_name = tool.get("name") or ""
match = _find_user_tool(user_tools, tool_type, custom_name)
if match:
tools.append(
{
"key": _tool_key(index),
"type": tool_type,
"name": custom_name,
"status": "reuse",
"target_id": str(match["id"]),
}
)
continue
available = _tool_instance(tool_type) is not None
tools.append(
{
"key": _tool_key(index),
"type": tool_type,
"name": custom_name,
"status": "create" if available else "unavailable",
"requires_secrets": tool.get("requires_secrets") or _live_requires_secrets(tool_type),
}
)
prompt_spec = spec.get("prompt")
if isinstance(prompt_spec, dict):
existing = PromptsRepository(conn).find(
user, prompt_spec.get("name") or "", prompt_spec.get("content") or ""
)
prompt = {"status": "reuse" if existing else "create", "name": prompt_spec.get("name") or ""}
else:
prompt = {"status": "default"}
customs = UserCustomModelsRepository(conn).list_for_user(user)
models = []
for entry in (spec.get("model") or {}).get("available") or []:
if isinstance(entry, str):
models.append(
{
"id": entry,
"status": "matched" if validate_model_id(entry, user) else "unavailable",
}
)
elif isinstance(entry, dict):
match = _find_custom_model(customs, entry)
models.append(
{
"display_name": entry.get("display_name") or "",
"status": "reuse" if match else "create",
"requires_secrets": ["api_key"] if not match else [],
}
)
workflow_plan = None
wf_spec = spec.get("workflow")
if spec.get("agent_type") == "workflow":
if isinstance(wf_spec, dict):
action = "create"
if target.get("action") == "update" and target.get("agent_id"):
agent_row = AgentsRepository(conn).get(str(target["agent_id"]), user)
existing_wf = (agent_row or {}).get("workflow_id")
if existing_wf and WorkflowsRepository(conn).get(str(existing_wf), user):
action = "update"
workflow_plan = {
"nodes": len(wf_spec.get("nodes") or []),
"edges": len(wf_spec.get("edges") or []),
"action": action,
}
elif "workflow" in spec:
# Explicit ``workflow: null`` (vs. an omitted key, which changes
# nothing). Apply clears the link and reaps the row when no other
# agent references it, destroying the graph, its run history and
# its artifacts — the dry-run has to say so before the user commits.
workflow_plan = _plan_workflow_removal(conn, user, target)
return {
"target": target,
"sources": sources,
"tools": tools,
"prompt": prompt,
"models": models,
"workflow": workflow_plan,
}
# ---------------------------------------------------------------------------
# Apply
# ---------------------------------------------------------------------------
def _apply_prompt(conn, user: str, spec: dict) -> Optional[str]:
prompt_spec = spec.get("prompt")
if not isinstance(prompt_spec, dict):
return None
row = PromptsRepository(conn).find_or_create(
user, prompt_spec.get("name") or "Imported prompt", prompt_spec.get("content") or ""
)
return str(row["id"])
def _apply_sources(conn, user: str, spec: dict, resolution: dict, warnings: list):
"""Resolve spec sources; returns ``(ids, id_by_name)``.
``id_by_name`` maps each spec source name to its resolved id (or None) so
a workflow graph's node configs — which reference sources by name — can
be rewritten from the same resolution.
"""
repo = SourcesRepository(conn)
mapping = resolution.get("sources")
if not isinstance(mapping, dict):
mapping = {}
resolved: list[str] = []
id_by_name: dict[str, Optional[str]] = {}
entries = [src for src in spec.get("sources") or [] if isinstance(src, dict)]
# ``id_by_name`` (and a workflow node's source reference) keys on the name
# alone, but names aren't unique per user — ``find_by_name`` returns the
# oldest match, so duplicates collapse onto one source. Same hazard the
# custom-model display names carry below; surface it rather than silently
# rewiring a node to a different source.
names = [src.get("name") or "" for src in entries]
for dup in sorted({n for n in names if n and names.count(n) > 1}):
warnings.append(
f"Multiple sources share the name '{dup}'; references to it all "
"resolve to the oldest one"
)
for src in entries:
name = src.get("name") or ""
id_by_name.setdefault(name, None)
mapped = mapping.get(name)
if mapped:
# Ownership-check the client-supplied id before linking (IDOR guard).
owned = repo.get_any(str(mapped), user)
if owned:
resolved.append(str(owned["id"]))
id_by_name[name] = str(owned["id"])
else:
warnings.append(f"Source mapping for '{name}' is not yours; ignored")
continue
match = repo.find_by_name(user, name)
if match:
resolved.append(str(match["id"]))
id_by_name[name] = str(match["id"])
else:
warnings.append(f"Source '{name}' not found; left unattached")
seen: set = set()
out = []
for sid in resolved:
if sid not in seen:
seen.add(sid)
out.append(sid)
return out, id_by_name
def _validate_tool_urls(tool_type: str, config: dict) -> Optional[str]:
"""SSRF-guard imported tool config to parity with the create routes.
Returns an error message if a URL is unsafe, else None.
"""
try:
if tool_type == "mcp_tool":
server_url = (config.get("server_url") or "").strip()
if server_url:
validate_url(server_url)
elif tool_type == "api_tool":
url = (config.get("url") or "").strip()
if url:
validate_url(url)
except SSRFError:
return f"Tool '{tool_type}' has an unsafe URL; not created"
return None
def _create_tool_from_spec(conn, user: str, tool: dict, secrets: dict, warnings: list) -> Optional[str]:
from docsgpt.api.user.tools.routes import _encrypt_secret_fields, transform_actions
tool_type = tool.get("type") or ""
inst = _tool_instance(tool_type)
if inst is None:
warnings.append(f"Tool type '{tool_type}' not available on this instance; skipped")
return None
config_requirements = inst.get_config_requirements() or {}
config = dict(tool.get("config") or {})
config.update(secrets or {})
if tool_type == "api_tool":
label = tool.get("display_name") or tool.get("name") or tool_type
config["actions"] = _sanitize_api_tool_actions(
config.get("actions"), label, warnings
)
if not config["actions"]:
# Zero actions means a dead tool — the executor resolves api_tool
# calls exclusively through config["actions"]. Reachable via older
# exports (which couldn't carry actions) or a fully-dropped block.
warnings.append(
f"API tool '{label}' was created without any actions; "
"configure its actions in the tool's settings before use"
)
missing = [k for k in _secret_field_names(config_requirements, required_only=True) if not config.get(k)]
if missing:
warnings.append(f"Tool '{tool_type}' needs secret(s) {missing}; not created")
return None
url_error = _validate_tool_urls(tool_type, config)
if url_error:
warnings.append(url_error)
return None
provided_secrets = [k for k in _secret_field_names(config_requirements) if config.get(k)]
storage_config = _encrypt_secret_fields(config, config_requirements, user)
if provided_secrets and not storage_config.get("encrypted_credentials"):
warnings.append(f"Tool '{tool_type}' secret encryption failed; not created")
return None
actions = transform_actions(inst.get_actions_metadata() or [])
actions = _apply_action_overrides(actions, tool.get("actions"), tool_type, warnings)
created = UserToolsRepository(conn).create(
user,
tool_type,
config=storage_config,
custom_name=tool.get("name") or "",
display_name=tool.get("display_name") or tool_type,
description=tool.get("description") or "",
config_requirements=config_requirements,
actions=actions,
status=True,
)
return str(created["id"])
def _apply_tools(conn, user: str, spec: dict, resolution: dict, warnings: list):
"""Resolve spec tools; returns ``(ids, id_by_key)``.
``id_by_key`` maps each spec entry's ``tool-N`` key to its resolved id
(or None when skipped/unavailable) so a workflow graph's node configs —
which reference tools by those keys — can be rewritten from the same
resolution.
"""
repo = UserToolsRepository(conn)
user_tools = repo.list_for_user(user)
decisions = resolution.get("tools")
if not isinstance(decisions, dict):
decisions = {}
out: list[str] = []
id_by_key: dict[str, Optional[str]] = {}
def _resolve(key: str, tid: Optional[str]) -> None:
id_by_key[key] = tid
if tid:
out.append(tid)
for index, tool in enumerate(spec.get("tools") or []):
key = _tool_key(index)
tool_type = tool.get("type") or ""
decision = decisions.get(key)
if not isinstance(decision, dict):
decision = {}
choice = decision.get("decision")
if tool.get("builtin"):
if synthesize_tool_by_name(tool_type) is None:
warnings.append(f"Built-in tool '{tool_type}' unavailable; skipped")
_resolve(key, None)
continue
_resolve(key, default_tool_id(tool_type))
continue
if choice == "skip":
_resolve(key, None)
continue
if choice == "reuse" and decision.get("tool_id"):
# Ownership-check the client-supplied id before linking (IDOR guard).
tid = str(decision["tool_id"])
if is_synthesized_tool_id(tid) or repo.get_any(tid, user):
_resolve(key, tid)
else:
warnings.append(f"Tool reuse id '{tid}' is not yours; ignored")
_resolve(key, None)
continue
custom_name = tool.get("name") or ""
match = _find_user_tool(user_tools, tool_type, custom_name)
if match and choice != "create":
_resolve(key, str(match["id"]))
continue
created_id = _create_tool_from_spec(conn, user, tool, decision.get("secrets") or {}, warnings)
_resolve(key, created_id)
if created_id:
# Record so a later identical spec entry reuses rather than re-creates.
user_tools.append({"id": created_id, "name": tool_type, "custom_name": custom_name})
return out, id_by_key
def _apply_models(conn, user: str, spec: dict, resolution: dict, warnings: list):
model_spec = spec.get("model") or {}
repo = UserCustomModelsRepository(conn)
customs = repo.list_for_user(user)
decisions = resolution.get("models")
if not isinstance(decisions, dict):
decisions = {}
name_to_id: dict[str, str] = {}
models: list[str] = []
available = model_spec.get("available") or []
# ``name_to_id`` (and a workflow node's ``model_id`` reference) keys on
# display_name alone, while custom-model identity is the full
# (display_name, upstream_model_id, base_url) triple — duplicates
# collapse onto the last entry, so surface that instead of silently
# rewiring nodes.
custom_names = [
e.get("display_name") or "" for e in available if isinstance(e, dict)
]
for dup in sorted({n for n in custom_names if n and custom_names.count(n) > 1}):
warnings.append(
f"Multiple custom models share the display name '{dup}'; "
"references to it resolve to the last one in the file"
)
for entry in available:
if isinstance(entry, str):
if validate_model_id(entry, user):
models.append(entry)
name_to_id[entry] = entry
else:
warnings.append(f"Model '{entry}' not available; skipped")
continue
if not isinstance(entry, dict):
continue
display_name = entry.get("display_name") or ""
match = _find_custom_model(customs, entry)
if match:
models.append(str(match["id"]))
name_to_id[display_name] = str(match["id"])
continue
decision = decisions.get(display_name)
api_key = decision.get("api_key") if isinstance(decision, dict) else None
if not api_key:
warnings.append(f"Custom model '{display_name}' needs an API key; skipped")
continue
base_url = entry.get("base_url") or ""
try:
validate_user_base_url(base_url)
except UnsafeUserUrlError:
warnings.append(f"Custom model '{display_name}' has an unsafe base URL; skipped")
continue
created = repo.create(
user,
entry.get("upstream_model_id") or "",
display_name,
base_url,
api_key,
capabilities=entry.get("capabilities") or {},
)
if not created.get("api_key_encrypted"):
warnings.append(f"Custom model '{display_name}' key encryption failed; skipped")
continue
models.append(str(created["id"]))
name_to_id[display_name] = str(created["id"])
default_id = ""
raw_default = model_spec.get("default") or ""
if raw_default:
if raw_default in name_to_id:
default_id = name_to_id[raw_default]
elif validate_model_id(raw_default, user):
default_id = raw_default
else:
# Custom-model default carried as its display_name — resolve against
# the user's own models even if it wasn't in `available`.
owned = next(
(c for c in customs if (c.get("display_name") or "") == raw_default), None
)
if owned:
default_id = str(owned["id"])
else:
warnings.append(f"Default model '{raw_default}' not resolved")
return (models or None), (default_id or None), name_to_id
def _limit(spec: dict, key: str):
return (spec.get("limits") or {}).get(key)
# Sentinel: the file omitted ``spec.workflow`` entirely, so the agent's
# existing workflow link is left untouched (vs. explicit null = clear).
_WORKFLOW_UNCHANGED = object()
def _rewrite_node_refs(
conn,
user: str,
cfg: dict,
tool_ids_by_key: dict,
source_ids_by_name: dict,
model_ids_by_name: dict,
warnings: list,
) -> None:
"""Resolve an agent node's portable references back to the importer's ids.
Export writes ``tool-N`` keys / source names / model display names, and a
hand-edited file may also carry raw ids — those are accepted only when
they are a builtin synthetic id or ownership-check against the importing
user (IDOR guard). Anything unresolvable is dropped with a warning so the
graph imports degraded rather than pointing at foreign resources.
"""
tools_repo = UserToolsRepository(conn)
sources_repo = SourcesRepository(conn)
resolved_tools = []
for ref in cfg.get("tools") or []:
ref = str(ref)
if ref in tool_ids_by_key:
tid = tool_ids_by_key[ref]
if tid:
resolved_tools.append(tid)
else:
warnings.append(f"Workflow node tool '{ref}' not resolved; removed")
elif is_synthesized_tool_id(ref):
resolved_tools.append(ref)
elif looks_like_uuid(ref) and tools_repo.get_any(ref, user):
resolved_tools.append(ref)
else:
warnings.append(f"Workflow node tool '{ref}' not resolved; removed")
cfg["tools"] = resolved_tools
resolved_sources = []
for ref in cfg.get("sources") or []:
ref = str(ref)
if ref in source_ids_by_name:
sid = source_ids_by_name[ref]
if sid:
resolved_sources.append(sid)
else:
warnings.append(f"Workflow node source '{ref}' not resolved; removed")
elif looks_like_uuid(ref) and sources_repo.get_any(ref, user):
resolved_sources.append(ref)
else:
warnings.append(f"Workflow node source '{ref}' not resolved; removed")
# Distinct names can collapse onto one source (see ``_apply_sources``), so
# dedupe rather than storing the same id twice on the node.
cfg["sources"] = list(dict.fromkeys(resolved_sources))
raw_model = cfg.get("model_id")
if raw_model:
raw_model = str(raw_model)
if raw_model in model_ids_by_name:
cfg["model_id"] = model_ids_by_name[raw_model]
elif validate_model_id(raw_model, user):
cfg["model_id"] = raw_model
else:
warnings.append(f"Workflow node model '{raw_model}' not resolved; removed")
cfg["model_id"] = None
def _apply_workflow(
conn,
user: str,
spec: dict,
target: dict,
tool_ids_by_key: dict,
source_ids_by_name: dict,
model_ids_by_name: dict,
warnings: list,
):
"""Write the imported graph; returns the workflow id, None, or the sentinel.
The graph goes through the same validation gate as the workflow API
(``validate_workflow_structure``) so a hand-edited YAML can't install a
graph the write path would have rejected. When the matched agent already
owns a workflow the graph lands as a new version of it (mirroring
``PUT /api/workflows/<id>``); otherwise a new workflow is created.
"""
if "workflow" not in spec:
return _WORKFLOW_UNCHANGED
wf_spec = spec.get("workflow")
if not isinstance(wf_spec, dict):
return None
from docsgpt.api.user.workflows.routes import (
_write_graph,
get_workflow_graph_version,
normalize_agent_node_json_schemas,
validate_workflow_structure,
)
nodes_data = []
for node in wf_spec.get("nodes") or []:
config = copy.deepcopy(node.get("config") or {})
if not isinstance(config, dict):
config = {}
if node.get("type") == "agent":
_rewrite_node_refs(
conn,
user,
_effective_node_config(config),
tool_ids_by_key,
source_ids_by_name,
model_ids_by_name,
warnings,
)
_strip_shadowed_ref_keys(config)
position = node.get("position")
nodes_data.append(
{
"id": str(node.get("id")),
"type": str(node.get("type")),
"title": node.get("title") or "",
"description": node.get("description") or "",
"position": position if isinstance(position, dict) else {"x": 0, "y": 0},
"data": config,
}
)
edges_data = [
{
"id": str(edge.get("id")),
"source": edge.get("source"),
"target": edge.get("target"),
"sourceHandle": edge.get("sourceHandle"),
"targetHandle": edge.get("targetHandle"),
}
for edge in wf_spec.get("edges") or []
]
errors = validate_workflow_structure(nodes_data, edges_data, user_id=user)
if errors:
raise AgentImportError("Workflow validation failed: " + "; ".join(errors[:5]))
nodes_data = normalize_agent_node_json_schemas(nodes_data)
wf_repo = WorkflowsRepository(conn)
name = wf_spec.get("name") or spec.get("name") or "Imported workflow"
description = wf_spec.get("description") or ""
existing = None
if target.get("action") == "update" and target.get("agent_id"):
agent_row = AgentsRepository(conn).get(str(target["agent_id"]), user)
existing_id = (agent_row or {}).get("workflow_id")
if existing_id:
existing = wf_repo.get(str(existing_id), user)
if existing is not None:
pg_workflow_id = str(existing["id"])
next_version = get_workflow_graph_version(existing) + 1
_write_graph(conn, pg_workflow_id, next_version, nodes_data, edges_data)
wf_repo.update(
pg_workflow_id,
user,
{
"name": name,
"description": description,
"current_graph_version": next_version,
},
)
WorkflowNodesRepository(conn).delete_other_versions(pg_workflow_id, next_version)
WorkflowEdgesRepository(conn).delete_other_versions(pg_workflow_id, next_version)
return pg_workflow_id
created = wf_repo.create(user, name, description=description)
_write_graph(conn, str(created["id"]), 1, nodes_data, edges_data)
return str(created["id"])
def apply_import(conn, user: str, doc: dict, resolution: Optional[dict] = None) -> dict:
"""Create or update an agent from a parsed YAML doc.
A newly created agent lands as a draft. An update to an existing agent
(matched by id or slug) preserves that agent's current status, so
re-importing over a published agent keeps it live — its API key and any
active users are unaffected — rather than silently reverting it to draft.
On update the YAML is authoritative: fields it specifies are written even
when that clears a value (e.g. removing all models, dropping ``json_schema``,
or switching ``prompt`` back to default), so re-importing an edited file is
a true sync rather than an additive merge.
"""
if not isinstance(resolution, dict):
resolution = {}
spec = doc["spec"]
metadata = doc.get("metadata") or {}
warnings: list[str] = []
target = _resolve_target(conn, user, metadata)
prompt_id = _apply_prompt(conn, user, spec)
source_ids, source_ids_by_name = _apply_sources(conn, user, spec, resolution, warnings)
tool_ids, tool_ids_by_key = _apply_tools(conn, user, spec, resolution, warnings)
models, default_model_id, model_ids_by_name = _apply_models(
conn, user, spec, resolution, warnings
)
is_workflow = (spec.get("agent_type") or "classic") == "workflow"
agents_repo = AgentsRepository(conn)
is_update = bool(target.get("action") == "update" and target.get("agent_id"))
exclude_id = str(target["agent_id"]) if is_update else None
# Slug is recomputed from the file, not preserved from the matched row:
# a round-tripped export carries metadata.slug so it stays stable, but an
# update-by-id whose file omits the slug and renames the agent rewrites the
# slug to follow the new name (excluding the row's own slug from collision
# checks). Intended — the slug tracks the file.
slug = _unique_slug(agents_repo, user, metadata.get("slug") or spec.get("name"), exclude_id=exclude_id)
try:
chunks_value = int(spec["chunks"]) if spec.get("chunks") is not None else 6
except (TypeError, ValueError):
chunks_value = 6
# YAML-authoritative fields — written even when the resolved value is None,
# so a re-import can CLEAR models / json_schema / prompt-to-default. On
# create, AgentsRepository.create() skips None kwargs, so Nones become
# column defaults.
authoritative: dict[str, Any] = {
"description": spec.get("description") or "",
"agent_type": spec.get("agent_type") or "classic",
"chunks": chunks_value,
"prompt_id": prompt_id,
"tools": tool_ids,
"json_schema": spec.get("json_schema"),
"models": models,
"default_model_id": default_model_id,
"extra_source_ids": source_ids,
"limited_token_mode": bool(_limit(spec, "limited_token_mode")),
"limited_request_mode": bool(_limit(spec, "limited_request_mode")),
"allow_system_prompt_override": bool(spec.get("allow_system_prompt_override")),
"config": _import_config(spec),
"slug": slug,
}
orphaned_workflow_id: Optional[str] = None
if is_workflow:
# A workflow agent's behavior lives in its graph — the agent-level
# prompt/tool/model/source slots in the file exist only to carry the
# graph's references through resolution, so they never land on the row.
authoritative.update(
{
"prompt_id": None,
"tools": [],
"models": None,
"default_model_id": None,
"extra_source_ids": [],
}
)
prior_workflow_id = None
if is_update:
prior_workflow_id = (
agents_repo.get(str(target["agent_id"]), user) or {}
).get("workflow_id")
published = (target.get("status") or "") == "published"
workflow_result = _apply_workflow(
conn,
user,
spec,
target,
tool_ids_by_key,
source_ids_by_name,
model_ids_by_name,
warnings,
)
# A published workflow agent with no graph can't run — the create and
# update routes both reject that state, so import must not install it
# through the back door (e.g. a file that flips ``agent_type`` to
# workflow over a published classic agent).
if workflow_result is _WORKFLOW_UNCHANGED:
if published and not prior_workflow_id:
raise AgentImportError(
"A published workflow agent needs a workflow; this file has none"
)
elif workflow_result is None and published:
if not prior_workflow_id:
raise AgentImportError(
"A published workflow agent needs a workflow; this file has none"
)
warnings.append(
"File has no workflow; the published agent kept its existing one"
)
else:
authoritative["workflow_id"] = workflow_result
if workflow_result is None and prior_workflow_id:
# Explicit ``workflow: null`` on a draft: remember the old
# link so the row can be reaped after the update — there is
# no workflow-list API independent of agents, so an
# unlinked workflow would be unreachable forever.
orphaned_workflow_id = str(prior_workflow_id)
# Optional fields — applied only when present so a partial file doesn't wipe them.
optional = {
"retriever": spec.get("retriever") or ("classic" if not source_ids else None),
"token_limit": _limit(spec, "token_limit"),
"request_limit": _limit(spec, "request_limit"),
}
optional = {k: v for k, v in optional.items() if v is not None}
if is_update:
# Status is intentionally omitted so an update never changes it:
# re-importing over a published agent keeps it live (its API key and
# any active users are unaffected) instead of reverting it to draft.
# Only a brand-new agent (the create path below) starts as a draft.
fields = {**authoritative, **optional, "name": spec.get("name")}
if agents_repo.update(str(target["agent_id"]), user, fields):
if orphaned_workflow_id and not agents_repo.count_by_workflow(
orphaned_workflow_id, user
):
# Nothing references it anymore; delete (cascades nodes/edges
# and reaps run artifacts) rather than stranding the row.
WorkflowsRepository(conn).delete(orphaned_workflow_id, user)
# Destructive and irreversible — the graph, its run history and
# its artifacts (rows and stored bytes) all go. Say so.
warnings.append(
"File has no workflow; the agent's workflow was deleted "
"along with its run history and artifacts"
)
return {
"agent_id": str(target["agent_id"]),
"action": "updated",
"status": target.get("status") or "draft",
"agent_type": authoritative["agent_type"],
"slug": slug,
"warnings": warnings,
}
# Row vanished between resolve and write — fall through to create.
row = agents_repo.create(user, spec.get("name"), "draft", **{**authoritative, **optional})
return {
"agent_id": str(row["id"]),
"action": "created",
"status": "draft",
"agent_type": authoritative["agent_type"],
"slug": slug,
"warnings": warnings,
}
# ---------------------------------------------------------------------------
# HTTP surface
# ---------------------------------------------------------------------------
agents_portability_ns = Namespace(
"agents", description="Agent import/export operations", path="/api"
)
def _read_import_payload(req):
"""Pull (yaml_text, resolution) from a JSON body or uploaded file.
Enforces a body-size cap before reading so a giant document can't be
buffered into memory.
"""
if req.content_length and req.content_length > MAX_IMPORT_BYTES:
raise AgentImportError("Import document too large")
content_type = req.content_type or ""
if "application/json" in content_type:
data = req.get_json(silent=True) or {}
resolution = data.get("resolution")
return (
data.get("yaml") or data.get("content") or "",
resolution if isinstance(resolution, dict) else {},
)
if "file" in req.files:
raw = req.files["file"].read(MAX_IMPORT_BYTES + 1)
if len(raw) > MAX_IMPORT_BYTES:
raise AgentImportError("Import document too large")
return raw.decode("utf-8", "replace"), {}
# Raw body: read from the stream with a hard cap so a chunked request
# (no Content-Length) can't be buffered unbounded into memory.
raw = req.stream.read(MAX_IMPORT_BYTES + 1)
if len(raw) > MAX_IMPORT_BYTES:
raise AgentImportError("Import document too large")
return raw.decode("utf-8", "replace"), {}
def _restricted_token_denial(conn, user: str, doc: dict) -> Optional[str]:
"""Why a resource-restricted personal access token may not import ``doc``, if it may not.
An import resolves sources, prompts and tools by name and may create them,
so a token restricted on any of those families cannot be held to its
allowlist here. A token restricted to specific agents may update exactly
those; it can never create one.
"""
for family in ("sources", "prompts", "tools", "workflows"):
if allowed_ids(request, family) is not None:
return f"A token restricted to specific {family} cannot import agents"
allowed_agents = allowed_ids(request, "agents")
if allowed_agents is None:
return None
target = _resolve_target(conn, user, doc.get("metadata") or {})
if target["action"] != "update" or target["agent_id"] not in allowed_agents:
return "This token is restricted to specific agents and may only update those"
return None
def _token_denied_response(reason: str):
return make_response(
jsonify({"success": False, "error": "resource_not_allowed", "message": reason}), 403
)
@agents_portability_ns.route("/export_agent")
class ExportAgent(Resource):
@api.doc(params={"id": "Agent ID"}, description="Export an agent as YAML")
def get(self):
if not (decoded_token := request.decoded_token):
return make_response(jsonify({"success": False}), 401)
user = decoded_token.get("sub")
agent_id = request.args.get("id")
if not agent_id:
return make_response(jsonify({"success": False, "message": "id is required"}), 400)
with db_session() as conn:
repo = AgentsRepository(conn)
try:
ra = require(conn, "agent", agent_id, user, "export")
except AccessDenied as denied:
return make_response(
jsonify({"success": False, "message": denied.message}), denied.status
)
agent = repo.get_by_id(ra.resource_id)
if not agent:
return make_response(
jsonify({"success": False, "message": "Agent not found"}), 404
)
# Serialized as the owner: the agent's prompt, sources, tools and
# workflow are the owner's (secrets are never exported).
owner_id = ra.owner_id
agent["slug"] = ensure_agent_slug(conn, agent, owner_id)
try:
export = serialize_agent(conn, agent, owner_id)
except AgentExportError as exc:
return make_response(
jsonify({"success": False, "message": str(exc)}), 400
)
body = agent_to_yaml(export)
filename = f"{export['metadata'].get('slug') or 'agent'}.agent.yaml"
response = make_response(body, 200)
response.headers["Content-Type"] = "application/x-yaml; charset=utf-8"
response.headers["Content-Disposition"] = f'attachment; filename="{filename}"'
return response
@agents_portability_ns.route("/import_agent/plan")
class ImportAgentPlan(Resource):
@api.doc(description="Dry-run an agent YAML import and return the resolution plan")
def post(self):
if not (decoded_token := request.decoded_token):
return make_response(jsonify({"success": False}), 401)
user = decoded_token.get("sub")
try:
yaml_text, _ = _read_import_payload(request)
doc = parse_agent_yaml(yaml_text)
except AgentImportError as exc:
return make_response(jsonify({"success": False, "message": str(exc)}), 400)
try:
with db_readonly() as conn:
if reason := _restricted_token_denial(conn, user, doc):
return _token_denied_response(reason)
plan = plan_import(conn, user, doc)
except Exception:
current_app.logger.error("Agent import plan failed", exc_info=True)
return make_response(
jsonify({"success": False, "message": "Could not analyze the agent file"}), 500
)
return make_response(jsonify({"success": True, "plan": plan}), 200)
@agents_portability_ns.route("/import_agent")
class ImportAgent(Resource):
@api.doc(description="Import an agent from YAML (created as a draft)")
def post(self):
if not (decoded_token := request.decoded_token):
return make_response(jsonify({"success": False}), 401)
user = decoded_token.get("sub")
try:
yaml_text, resolution = _read_import_payload(request)
doc = parse_agent_yaml(yaml_text)
except AgentImportError as exc:
return make_response(jsonify({"success": False, "message": str(exc)}), 400)
try:
with db_session() as conn:
if reason := _restricted_token_denial(conn, user, doc):
return _token_denied_response(reason)
result = apply_import(conn, user, doc, resolution)
except AgentImportError as exc:
# Apply-time rejection of the user's document (e.g. the workflow
# graph failing the API's validation gate) — a client error, with
# the detail surfaced, matching how PUT /workflows reports it.
return make_response(jsonify({"success": False, "message": str(exc)}), 400)
except Exception:
current_app.logger.error("Agent import failed", exc_info=True)
return make_response(
jsonify({"success": False, "message": "Import failed"}), 500
)
return make_response(jsonify({"success": True, **result}), 200)