diff --git a/application/alembic/versions/0019_agent_slug.py b/application/alembic/versions/0019_agent_slug.py new file mode 100644 index 00000000..8d4df0b2 --- /dev/null +++ b/application/alembic/versions/0019_agent_slug.py @@ -0,0 +1,36 @@ +"""0019 agent slug — stable per-user identifier for YAML export/import. + +Adds ``agents.slug`` (CITEXT) plus a partial unique index on +``(user_id, slug)`` where slug is not null, so an exported agent can be +re-imported idempotently (and mapped from a repo file for GitOps) while +still allowing multiple agents to carry NULL. Idempotent both ways. + +Revision ID: 0019_agent_slug +Revises: 0018_tool_attempts_attribution +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0019_agent_slug" +down_revision: Union[str, None] = "0018_tool_attempts_attribution" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute("ALTER TABLE agents ADD COLUMN IF NOT EXISTS slug CITEXT;") + op.execute( + """ + CREATE UNIQUE INDEX IF NOT EXISTS ix_agents_user_slug + ON agents (user_id, slug) + WHERE slug IS NOT NULL; + """ + ) + + +def downgrade() -> None: + op.execute("DROP INDEX IF EXISTS ix_agents_user_slug;") + op.execute("ALTER TABLE agents DROP COLUMN IF EXISTS slug;") diff --git a/application/api/user/agents/__init__.py b/application/api/user/agents/__init__.py index 30199c85..4892c0e7 100644 --- a/application/api/user/agents/__init__.py +++ b/application/api/user/agents/__init__.py @@ -4,5 +4,12 @@ from .routes import agents_ns from .sharing import agents_sharing_ns from .webhooks import agents_webhooks_ns from .folders import agents_folders_ns +from .portability import agents_portability_ns -__all__ = ["agents_ns", "agents_sharing_ns", "agents_webhooks_ns", "agents_folders_ns"] +__all__ = [ + "agents_ns", + "agents_sharing_ns", + "agents_webhooks_ns", + "agents_folders_ns", + "agents_portability_ns", +] diff --git a/application/api/user/agents/portability.py b/application/api/user/agents/portability.py new file mode 100644 index 00000000..498f6cd2 --- /dev/null +++ b/application/api/user/agents/portability.py @@ -0,0 +1,949 @@ +"""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. +""" + +from __future__ import annotations + +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 application.agents.default_tools import ( + default_tool_id, + is_synthesized_tool_id, + synthesize_tool_by_name, + synthesized_tool_name_for_id, +) +from application.api import api +from application.core.model_utils import validate_model_id +from application.core.url_validation import SSRFError, validate_url +from application.security.safe_url import UnsafeUserUrlError, validate_user_base_url +from application.storage.db.base_repository import looks_like_uuid +from application.storage.db.repositories.agents import AgentsRepository +from application.storage.db.repositories.prompts import PromptsRepository +from application.storage.db.repositories.sources import SourcesRepository +from application.storage.db.repositories.user_custom_models import ( + UserCustomModelsRepository, +) +from application.storage.db.repositories.user_tools import UserToolsRepository +from application.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 + + +class AgentImportError(Exception): + """Raised when an agent YAML document is malformed or unsupported.""" + + +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 application.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. ``api_tool``/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``. + """ + 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 + + +# --------------------------------------------------------------------------- +# 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}) + continue + row = repo.get_any(tid_str, user) + if not row: + continue + requirements = row.get("config_requirements") or {} + requires_secrets = _secret_field_names(requirements) + 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": _safe_export_config(row.get("config") or {}, requirements), + "ref": tid_str, + } + 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 serialize_agent(conn, agent: dict, user: str) -> dict: + """Build the portable export document for an agent row.""" + 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": _serialize_prompt(conn, agent, user), + "model": _serialize_models(conn, agent, user), + "sources": _serialize_sources(conn, agent, user), + "tools": _serialize_tools(conn, agent, user), + "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)), + } + 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 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": + raise AgentImportError("Workflow agents are not supported for import yet") + 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_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 [], + } + ) + + return {"target": target, "sources": sources, "tools": tools, "prompt": prompt, "models": models} + + +# --------------------------------------------------------------------------- +# 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) -> list: + repo = SourcesRepository(conn) + mapping = resolution.get("sources") + if not isinstance(mapping, dict): + mapping = {} + resolved: list[str] = [] + for src in spec.get("sources") or []: + name = src.get("name") or "" + 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"])) + 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"])) + 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 + + +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 application.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 {}) + 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 []) + 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) -> list: + repo = UserToolsRepository(conn) + user_tools = repo.list_for_user(user) + decisions = resolution.get("tools") + if not isinstance(decisions, dict): + decisions = {} + out: list[str] = [] + for index, tool in enumerate(spec.get("tools") or []): + tool_type = tool.get("type") or "" + decision = decisions.get(_tool_key(index)) + 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") + continue + out.append(default_tool_id(tool_type)) + continue + + if choice == "skip": + 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): + out.append(tid) + else: + warnings.append(f"Tool reuse id '{tid}' is not yours; ignored") + continue + + custom_name = tool.get("name") or "" + match = _find_user_tool(user_tools, tool_type, custom_name) + if match and choice != "create": + out.append(str(match["id"])) + continue + + created_id = _create_tool_from_spec(conn, user, tool, decision.get("secrets") or {}, warnings) + if created_id: + out.append(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 + + +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] = [] + + for entry in model_spec.get("available") or []: + 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) + + +def _limit(spec: dict, key: str): + return (spec.get("limits") or {}).get(key) + + +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 = _apply_sources(conn, user, spec, resolution, warnings) + tool_ids = _apply_tools(conn, user, spec, resolution, warnings) + models, default_model_id = _apply_models(conn, user, spec, resolution, warnings) + + 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 2 + except (TypeError, ValueError): + chunks_value = 2 + + # 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")), + "slug": slug, + } + # 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): + return { + "agent_id": str(target["agent_id"]), + "action": "updated", + "status": target.get("status") or "draft", + "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", + "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"), {} + + +@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) + agent = repo.get_any(agent_id, user) + if not agent: + return make_response( + jsonify({"success": False, "message": "Agent not found"}), 404 + ) + if (agent.get("agent_type") or "") == "workflow": + return make_response( + jsonify( + {"success": False, "message": "Workflow agents can't be exported yet"} + ), + 400, + ) + agent["slug"] = ensure_agent_slug(conn, agent, user) + export = serialize_agent(conn, agent, user) + 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: + 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: + result = apply_import(conn, user, doc, resolution) + 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) diff --git a/application/api/user/agents/routes.py b/application/api/user/agents/routes.py index ee2f5dd0..fa097e51 100644 --- a/application/api/user/agents/routes.py +++ b/application/api/user/agents/routes.py @@ -181,6 +181,7 @@ def _format_agent_output(agent: dict, *, pinned: bool = False, include_key_maske out = { "id": str(agent["id"]), "name": agent.get("name", ""), + "slug": agent.get("slug", "") or "", "description": agent.get("description", "") or "", "image": ( generate_image_url(agent["image"]) if agent.get("image") else "" @@ -1297,6 +1298,10 @@ class AdoptAgent(Resource): now = datetime.datetime.now(datetime.timezone.utc) new_key = str(uuid.uuid4()) + # Copy content columns only. Identity / idempotency columns + # (key, slug, shared_token, incoming_webhook_token) are + # deliberately NOT copied — the adopted agent is a fresh copy + # and must get its own values (slug stays NULL until exported). create_kwargs: dict = {} for col in ( "description", "agent_type", "image", "retriever", diff --git a/application/api/user/routes.py b/application/api/user/routes.py index 675f0d01..53f0daad 100644 --- a/application/api/user/routes.py +++ b/application/api/user/routes.py @@ -5,7 +5,13 @@ Main user API routes - registers all namespace modules. from flask import Blueprint from application.api import api -from .agents import agents_ns, agents_sharing_ns, agents_webhooks_ns, agents_folders_ns +from .agents import ( + agents_folders_ns, + agents_ns, + agents_portability_ns, + agents_sharing_ns, + agents_webhooks_ns, +) from .analytics import analytics_ns from .attachments import attachments_ns from .conversations import conversations_ns @@ -32,11 +38,12 @@ api.add_namespace(conversations_ns) # Models api.add_namespace(models_ns) -# Agents (main, sharing, webhooks, folders) +# Agents (main, sharing, webhooks, folders, import/export) api.add_namespace(agents_ns) api.add_namespace(agents_sharing_ns) api.add_namespace(agents_webhooks_ns) api.add_namespace(agents_folders_ns) +api.add_namespace(agents_portability_ns) # Prompts api.add_namespace(prompts_ns) diff --git a/application/storage/db/models.py b/application/storage/db/models.py index 2099c2ec..5061c2d8 100644 --- a/application/storage/db/models.py +++ b/application/storage/db/models.py @@ -23,6 +23,7 @@ from sqlalchemy import ( DateTime, ForeignKey, ForeignKeyConstraint, + Index, Integer, MetaData, UniqueConstraint, @@ -204,6 +205,11 @@ agents_table = Table( Column("agent_type", Text), Column("status", Text, nullable=False), Column("key", CITEXT, unique=True), + # Stable per-user human identifier used to match an agent across + # YAML export/import (idempotent re-import / GitOps). Uniqueness is + # enforced by a partial unique index in migration 0019, not here, so + # multiple agents may carry NULL. + Column("slug", CITEXT), Column("image", Text), Column("source_id", UUID(as_uuid=True), ForeignKey("sources.id", ondelete="SET NULL")), Column("extra_source_ids", ARRAY(UUID(as_uuid=True)), nullable=False, server_default="{}"), @@ -231,6 +237,17 @@ agents_table = Table( Column("legacy_mongo_id", Text), ) +# Per-user uniqueness of the export/import slug. Mirrors the partial unique +# index created in migration 0019 so a schema built from this metadata +# (e.g. create_all) matches an Alembic-built one. +Index( + "ix_agents_user_slug", + agents_table.c.user_id, + agents_table.c.slug, + unique=True, + postgresql_where=agents_table.c.slug.isnot(None), +) + user_custom_models_table = Table( "user_custom_models", metadata, diff --git a/application/storage/db/repositories/agents.py b/application/storage/db/repositories/agents.py index 3e27bb83..930c0582 100644 --- a/application/storage/db/repositories/agents.py +++ b/application/storage/db/repositories/agents.py @@ -27,7 +27,7 @@ class AgentsRepository: @staticmethod def _normalize_unique_text(col: str, val): """Coerce blank strings for nullable unique text columns to NULL.""" - if col == "key" and val == "": + if col in ("key", "slug") and val == "": return None return val @@ -35,7 +35,7 @@ class AgentsRepository: values: dict = {"user_id": user_id, "name": name, "status": status} _ALLOWED = { - "description", "agent_type", "key", "retriever", + "description", "agent_type", "key", "slug", "retriever", "default_model_id", "incoming_webhook_token", "source_id", "prompt_id", "folder_id", "workflow_id", "extra_source_ids", "image", @@ -142,6 +142,17 @@ class AgentsRepository: row = result.fetchone() return row_to_dict(row) if row is not None else None + def find_by_slug(self, user_id: str, slug: str) -> Optional[dict]: + """Resolve one of a user's agents by its stable export/import slug.""" + if not slug: + return None + result = self._conn.execute( + text("SELECT * FROM agents WHERE user_id = :user_id AND slug = :slug"), + {"user_id": user_id, "slug": slug}, + ) + row = result.fetchone() + return row_to_dict(row) if row is not None else None + def list_for_user(self, user_id: str) -> list[dict]: result = self._conn.execute( text("SELECT * FROM agents WHERE user_id = :user_id ORDER BY created_at DESC"), @@ -150,14 +161,20 @@ class AgentsRepository: return [row_to_dict(r) for r in result.fetchall()] def list_templates(self) -> list[dict]: + # Template/system agents are seeded under ``__system__`` (migration + # 0001 + the seeder); the legacy ``'system'`` value is kept in the + # match so older rows still surface. result = self._conn.execute( - text("SELECT * FROM agents WHERE user_id = 'system' ORDER BY name"), + text( + "SELECT * FROM agents WHERE user_id IN ('system', '__system__') " + "ORDER BY name" + ), ) return [row_to_dict(r) for r in result.fetchall()] def update(self, agent_id: str, user_id: str, fields: dict) -> bool: allowed = { - "name", "description", "agent_type", "status", "key", "source_id", + "name", "description", "agent_type", "status", "key", "slug", "source_id", "chunks", "retriever", "prompt_id", "tools", "json_schema", "models", "default_model_id", "folder_id", "workflow_id", "extra_source_ids", "image", diff --git a/application/storage/db/repositories/prompts.py b/application/storage/db/repositories/prompts.py index 609c13c5..02babf5c 100644 --- a/application/storage/db/repositories/prompts.py +++ b/application/storage/db/repositories/prompts.py @@ -161,11 +161,8 @@ class PromptsRepository: ) return result.rowcount > 0 - def find_or_create(self, user_id: str, name: str, content: str) -> dict: - """Return existing prompt matching (user, name, content), or create one. - - Used by the seeder to avoid duplicating template prompts. - """ + def find(self, user_id: str, name: str, content: str) -> Optional[dict]: + """Return a prompt exactly matching (user, name, content), or None.""" result = self._conn.execute( text( "SELECT * FROM prompts WHERE user_id = :user_id AND name = :name AND content = :content" @@ -173,6 +170,14 @@ class PromptsRepository: {"user_id": user_id, "name": name, "content": content}, ) row = result.fetchone() - if row is not None: - return row_to_dict(row) + return row_to_dict(row) if row is not None else None + + def find_or_create(self, user_id: str, name: str, content: str) -> dict: + """Return existing prompt matching (user, name, content), or create one. + + Used by the seeder and agent YAML import to avoid duplicating prompts. + """ + existing = self.find(user_id, name, content) + if existing is not None: + return existing return self.create(user_id, name, content) diff --git a/application/storage/db/repositories/sources.py b/application/storage/db/repositories/sources.py index c9132565..31f54926 100644 --- a/application/storage/db/repositories/sources.py +++ b/application/storage/db/repositories/sources.py @@ -174,6 +174,25 @@ class SourcesRepository: return row return self.get_by_legacy_id(source_id, user_id) + def find_by_name(self, user_id: str, name: str) -> Optional[dict]: + """Return a user's source whose name matches ``name`` (case-insensitive). + + Used by agent YAML import to map a portable source name back to a + concrete source id. Returns the oldest match when names collide. + """ + if not name: + return None + result = self._conn.execute( + text( + "SELECT * FROM sources " + "WHERE user_id = :user_id AND lower(name) = lower(:name) " + "ORDER BY created_at LIMIT 1" + ), + {"user_id": user_id, "name": name}, + ) + row = result.fetchone() + return row_to_dict(row) if row is not None else None + def list_for_user( self, user_id: str, diff --git a/frontend/src/agents/AgentCard.tsx b/frontend/src/agents/AgentCard.tsx index 42968fae..f4091603 100644 --- a/frontend/src/agents/AgentCard.tsx +++ b/frontend/src/agents/AgentCard.tsx @@ -4,6 +4,7 @@ import { useNavigate } from 'react-router-dom'; import { useTranslation } from 'react-i18next'; import userService from '../api/services/userService'; +import Download from '../assets/download.svg'; import Duplicate from '../assets/duplicate.svg'; import Edit from '../assets/edit.svg'; import FolderIcon from '../assets/folder.svg'; @@ -104,6 +105,17 @@ export default function AgentCard({ iconWidth: 14, iconHeight: 14, }, + { + icon: Download, + label: t('agents.exportAgent'), + onClick: (e: SyntheticEvent) => { + e.stopPropagation(); + handleExport(); + }, + variant: 'default', + iconWidth: 14, + iconHeight: 14, + }, ...(agent.status === 'published' ? [ { @@ -223,6 +235,25 @@ export default function AgentCard({ } }; + const handleExport = async () => { + try { + const response = await userService.exportAgent(agent.id ?? '', token); + if (!response.ok) throw new Error('Failed to export agent'); + const yamlText = await response.text(); + const blob = new Blob([yamlText], { type: 'application/x-yaml' }); + const url = URL.createObjectURL(blob); + const link = document.createElement('a'); + link.href = url; + link.download = `${agent.slug || agent.name || 'agent'}.agent.yaml`; + document.body.appendChild(link); + link.click(); + document.body.removeChild(link); + URL.revokeObjectURL(url); + } catch (error) { + console.error('Error:', error); + } + }; + const handleDelete = async () => { try { const response = await userService.deleteAgent(agent.id ?? '', token); diff --git a/frontend/src/agents/AgentsList.tsx b/frontend/src/agents/AgentsList.tsx index a1902a48..6d9438c9 100644 --- a/frontend/src/agents/AgentsList.tsx +++ b/frontend/src/agents/AgentsList.tsx @@ -12,6 +12,7 @@ import { setConversation, updateConversationId, } from '../conversation/conversationSlice'; +import ImportAgentModal from '../modals/ImportAgentModal'; import { selectAgentFolders, selectSelectedAgent, @@ -284,6 +285,7 @@ function AgentSection({ const allAgents = useSelector(config.selectData); const [isCreatingFolder, setIsCreatingFolder] = useState(false); const [newFolderName, setNewFolderName] = useState(''); + const [showImportModal, setShowImportModal] = useState(false); const newFolderInputRef = useRef(null); const currentFolderId = @@ -540,6 +542,16 @@ function AgentSection({ {t('agents.folders.newFolder')} ))} + {config.id === 'user' && ( + + )} {config.showNewAgentButton && ( + ); + } + return ( + <> + + {!plan ? ( + + ) : ( + + )} + + ); + }; + + return ( + !o && handleClose()} + title={t('modals.importAgent.title')} + size="lg" + contentClassName="max-h-[70vh]" + footer={renderFooter()} + > +
+ {warnings ? ( +
+

+ {importedStatus === 'published' + ? t('modals.importAgent.warningsTitlePublished') + : t('modals.importAgent.warningsTitle')} +

+
    + {warnings.map((w, i) => ( +
  • {w}
    • + ))} +
+
+ ) : !plan ? ( +
+

+ {t('modals.importAgent.description')} +

+
fileInputRef.current?.click()} + className="border-border hover:border-primary flex cursor-pointer flex-col items-center justify-center rounded-xl border-2 border-dashed p-8 transition-colors" + > + +

+ {fileName || t('modals.importAgent.dropzoneText')} +

+
+ + {error &&

{error}

} +
+ ) : ( +
+
+ {plan.target.action === 'update' + ? t('modals.importAgent.willUpdate', { + matchedBy: plan.target.matched_by, + }) + : t('modals.importAgent.willCreate')} +
+ + {plan.sources.length > 0 && ( +
+ {plan.sources.map((s) => + s.status === 'matched' ? ( +

+ + {t('modals.importAgent.sourceMatched', { name: s.name })} +

+ ) : ( +
+

+ + {t('modals.importAgent.sourceMissing', { + name: s.name, + })} +

+ +
+ ), + )} +
+ )} + + {plan.tools.length > 0 && ( +
+ {plan.tools.map((tool) => ( +
+ {tool.status === 'builtin' && ( +

+ + {t('modals.importAgent.toolBuiltin', { + type: tool.type, + })} +

+ )} + {tool.status === 'reuse' && ( +

+ + {t('modals.importAgent.toolReuse', { + name: tool.name || tool.type, + })} +

+ )} + {tool.status === 'unavailable' && ( +

+ + {t('modals.importAgent.toolUnavailable', { + type: tool.type, + })} +

+ )} + {tool.status === 'create' && ( +
+

+ {t('modals.importAgent.toolCreate', { + name: tool.name || tool.type, + })} +

+ {(tool.requires_secrets || []).map((field) => ( + + setToolSecret(tool.key, field, e.target.value) + } + className="bg-card h-auto rounded-lg px-3 py-2 text-sm md:text-sm" + /> + ))} +
+ )} +
+ ))} +
+ )} + + {plan.models.some((m) => m.status === 'create') && ( +
+ {plan.models + .filter((m) => m.status === 'create') + .map((m) => ( +
+

+ {t('modals.importAgent.modelCustom', { + name: m.display_name, + })} +

+ + setModelKeys((prev) => ({ + ...prev, + [m.display_name || '']: e.target.value, + })) + } + className="bg-card h-auto rounded-lg px-3 py-2 text-sm md:text-sm" + /> +
+ ))} +
+ )} + + {error &&

{error}

} +
+ )} +
+ + ); +} + +function Section({ + title, + children, +}: { + title: string; + children: React.ReactNode; +}) { + return ( +
+

{title}

+
{children}
+
+ ); +} diff --git a/tests/api/test_agent_portability.py b/tests/api/test_agent_portability.py new file mode 100644 index 00000000..8ef90167 --- /dev/null +++ b/tests/api/test_agent_portability.py @@ -0,0 +1,480 @@ +"""Tests for agent YAML export/import (application.api.user.agents.portability). + +These exercise real SQL against the ephemeral ``pg_conn`` fixture, calling +the dependency-injected serialize/plan/apply functions directly so the +logic is covered without a running Flask app. +""" + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest + +from application.agents.default_tools import default_tool_id +from application.api.user.agents.portability import ( + API_VERSION, + AgentImportError, + agent_to_yaml, + apply_import, + ensure_agent_slug, + parse_agent_yaml, + plan_import, + serialize_agent, +) +from application.storage.db.repositories.agents import AgentsRepository +from application.storage.db.repositories.prompts import PromptsRepository +from application.storage.db.repositories.sources import SourcesRepository +from application.storage.db.repositories.user_tools import UserToolsRepository + + +pytestmark = pytest.mark.integration + + +def _make_agent(conn, user, **kwargs): + defaults = dict(description="d", chunks=2, retriever="classic") + defaults.update(kwargs) + return AgentsRepository(conn).create(user, kwargs.pop("name", "Bot"), "published", **defaults) + + +def _doc(**spec): + spec.setdefault("name", "Imported") + return { + "apiVersion": API_VERSION, + "kind": "Agent", + "metadata": {"slug": spec.pop("_slug", "imported")}, + "spec": spec, + } + + +def test_serialize_redacts_tool_secret(pg_conn): + user = "u_redact" + src = SourcesRepository(pg_conn).create("Docs", user_id=user, type="file") + prompt = PromptsRepository(pg_conn).create(user, "P", "You are helpful") + tool = UserToolsRepository(pg_conn).create( + user, + "brave", + config={"encrypted_credentials": "SECRETBLOB", "region": "US"}, + custom_name="My Brave", + display_name="Brave Search", + description="search", + config_requirements={"token": {"secret": True, "required": True}, "region": {}}, + actions=[], + ) + agent = AgentsRepository(pg_conn).create( + user, + "Support Bot", + "published", + description="help desk", + chunks=3, + retriever="classic", + prompt_id=str(prompt["id"]), + source_id=str(src["id"]), + tools=[str(tool["id"])], + ) + + export = serialize_agent(pg_conn, agent, user) + text = agent_to_yaml(export) + + assert "SECRETBLOB" not in text + assert "encrypted_credentials" not in text + tool_entry = export["spec"]["tools"][0] + assert tool_entry["type"] == "brave" + assert tool_entry["requires_secrets"] == ["token"] + assert "encrypted_credentials" not in tool_entry["config"] + assert tool_entry["config"] == {"region": "US"} + assert export["spec"]["prompt"] == {"name": "P", "content": "You are helpful"} + assert export["spec"]["sources"][0]["name"] == "Docs" + + +def test_export_assigns_unique_slug(pg_conn): + user = "u_slug" + a1 = AgentsRepository(pg_conn).create(user, "My Bot", "draft") + a2 = AgentsRepository(pg_conn).create(user, "My Bot", "draft") + + s1 = ensure_agent_slug(pg_conn, a1, user) + a2_reload = AgentsRepository(pg_conn).get(str(a2["id"]), user) + s2 = ensure_agent_slug(pg_conn, a2_reload, user) + + assert s1 == "my-bot" + assert s2 == "my-bot-2" + + +def test_round_trip_same_user_idempotent(pg_conn): + user = "u_round" + src = SourcesRepository(pg_conn).create("KB", user_id=user, type="file") + prompt = PromptsRepository(pg_conn).create(user, "Sys", "Be nice") + scheduler_id = default_tool_id("scheduler") + agent = AgentsRepository(pg_conn).create( + user, + "Bot", + "published", + description="d", + chunks=2, + retriever="classic", + prompt_id=str(prompt["id"]), + source_id=str(src["id"]), + tools=[scheduler_id], + ) + ensure_agent_slug(pg_conn, agent, user) + agent = AgentsRepository(pg_conn).get(str(agent["id"]), user) + + doc = parse_agent_yaml(agent_to_yaml(serialize_agent(pg_conn, agent, user))) + r1 = apply_import(pg_conn, user, doc) + r2 = apply_import(pg_conn, user, doc) + + assert r1["agent_id"] == r2["agent_id"] == str(agent["id"]) # matched by id + assert len(PromptsRepository(pg_conn).list_for_user(user)) == 1 # no dup prompt + imported = AgentsRepository(pg_conn).get(r1["agent_id"], user) + assert imported["status"] == "published" # update preserves the live status + assert [str(s) for s in imported["extra_source_ids"]] == [str(src["id"])] + assert imported["tools"] == [scheduler_id] # builtin passthrough + + +def test_update_preserves_published_status(pg_conn): + """Re-importing over a published agent must not revert it to draft. + + The agent's API key and any active users keep working; only the content + is synced. A new agent (no match) still lands as a draft. + """ + user = "u_pub" + agent = AgentsRepository(pg_conn).create( + user, + "Live Bot", + "published", + description="old", + chunks=2, + retriever="classic", + ) + slug = ensure_agent_slug(pg_conn, agent, user) + doc = { + "apiVersion": API_VERSION, + "kind": "Agent", + "metadata": {"id": str(agent["id"]), "slug": slug}, + "spec": {"name": "Live Bot", "description": "new", "retriever": "classic"}, + } + + result = apply_import(pg_conn, user, doc) + + assert result["action"] == "updated" + assert result["status"] == "published" # response reports the preserved status + updated = AgentsRepository(pg_conn).get(str(agent["id"]), user) + assert updated["status"] == "published" # stayed live + assert updated["description"] == "new" # content still synced + + +def test_import_by_slug_idempotent(pg_conn): + user = "u_slug_import" + doc = _doc(name="Slug Bot", _slug="slug-bot", retriever="classic") + + r1 = apply_import(pg_conn, user, doc) + r2 = apply_import(pg_conn, user, doc) + + assert r1["action"] == "created" + assert r2["action"] == "updated" + assert r1["agent_id"] == r2["agent_id"] + assert len(AgentsRepository(pg_conn).list_for_user(user)) == 1 + + +def test_import_missing_source_drafts_and_warns(pg_conn): + user = "u_missing" + doc = _doc( + name="No Source Bot", + _slug="no-source", + sources=[{"name": "Nonexistent KB", "type": "file"}], + ) + + result = apply_import(pg_conn, user, doc) + + assert result["status"] == "draft" # new agents are created as drafts + assert any("Nonexistent KB" in w for w in result["warnings"]) + agent = AgentsRepository(pg_conn).get(result["agent_id"], user) + assert agent["status"] == "draft" + assert list(agent["extra_source_ids"]) == [] + + +def test_import_matches_existing_source_by_name(pg_conn): + user = "u_match" + src = SourcesRepository(pg_conn).create("Product Docs", user_id=user, type="file") + doc = _doc( + name="Matcher", + _slug="matcher", + sources=[{"name": "product docs", "type": "file"}], # case-insensitive + ) + + result = apply_import(pg_conn, user, doc) + + agent = AgentsRepository(pg_conn).get(result["agent_id"], user) + assert [str(s) for s in agent["extra_source_ids"]] == [str(src["id"])] + assert result["warnings"] == [] + + +def test_plan_classifies_references(pg_conn): + user = "u_plan" + SourcesRepository(pg_conn).create("Known", user_id=user, type="file") + tool = UserToolsRepository(pg_conn).create( + user, "brave", custom_name="My Brave", display_name="Brave", config={}, + ) + doc = _doc( + name="Plan Bot", + sources=[{"name": "Known"}, {"name": "Unknown"}], + tools=[ + {"type": "scheduler", "builtin": True}, + {"type": "brave", "name": "My Brave"}, + ], + ) + + plan = plan_import(pg_conn, user, doc) + + statuses = {s["name"]: s["status"] for s in plan["sources"]} + assert statuses == {"Known": "matched", "Unknown": "missing"} + tool_statuses = [t["status"] for t in plan["tools"]] + assert tool_statuses[0] == "builtin" + assert tool_statuses[1] == "reuse" + assert plan["tools"][1]["target_id"] == str(tool["id"]) + assert plan["target"]["action"] == "create" + + +def test_apply_creates_tool_with_supplied_secret(pg_conn, monkeypatch): + user = "u_tool" + fake_tool = Mock() + fake_tool.get_config_requirements.return_value = { + "api_key": {"secret": True, "required": True, "label": "API Key"} + } + fake_tool.get_actions_metadata.return_value = [] + monkeypatch.setattr( + "application.api.user.agents.portability._tool_instance", + lambda tool_type: fake_tool, + ) + doc = _doc( + name="Tool Bot", + _slug="tool-bot", + tools=[ + { + "type": "brave", + "name": "My API", + "config": {}, + "requires_secrets": ["api_key"], + } + ], + ) + + # No secret supplied -> tool skipped, agent still created. + res = apply_import(pg_conn, user, doc) + agent = AgentsRepository(pg_conn).get(res["agent_id"], user) + assert agent["tools"] == [] + assert any("api_key" in w for w in res["warnings"]) + + # Secret supplied via resolution -> tool created and linked, secret encrypted. + res2 = apply_import( + pg_conn, + user, + doc, + resolution={"tools": {"tool-0": {"secrets": {"api_key": "PLAINTEXT_KEY"}}}}, + ) + agent2 = AgentsRepository(pg_conn).get(res2["agent_id"], user) + assert len(agent2["tools"]) == 1 + tool_row = UserToolsRepository(pg_conn).get_any(agent2["tools"][0], user) + assert "PLAINTEXT_KEY" not in str(tool_row["config"]) + assert "encrypted_credentials" in tool_row["config"] + + +def test_parse_rejects_bad_documents(): + with pytest.raises(AgentImportError): + parse_agent_yaml("kind: NotAgent\napiVersion: docsgpt.arc53.com/v1\nspec:\n name: x\n") + with pytest.raises(AgentImportError): + parse_agent_yaml("kind: Agent\napiVersion: other/v1\nspec:\n name: x\n") + with pytest.raises(AgentImportError): + parse_agent_yaml("kind: Agent\napiVersion: docsgpt.arc53.com/v1\nspec: {}\n") + with pytest.raises(AgentImportError): + parse_agent_yaml( + "kind: Agent\napiVersion: docsgpt.arc53.com/v1\n" + "spec:\n name: x\n agent_type: workflow\n" + ) + + +def test_parse_rejects_yaml_aliases(): + bomb = ( + "kind: Agent\napiVersion: docsgpt.arc53.com/v1\n" + "spec:\n name: &a x\n description: *a\n" + ) + with pytest.raises(AgentImportError): + parse_agent_yaml(bomb) + + +def test_update_clears_removed_json_schema(pg_conn): + """On update the YAML is authoritative — removing json_schema clears it.""" + user = "u_clear" + agent = AgentsRepository(pg_conn).create( + user, + "Clear Bot", + "published", + description="d", + chunks=2, + retriever="classic", + json_schema={"type": "object"}, + ) + slug = ensure_agent_slug(pg_conn, agent, user) + doc = { + "apiVersion": API_VERSION, + "kind": "Agent", + "metadata": {"slug": slug}, + "spec": {"name": "Clear Bot", "retriever": "classic", "json_schema": None}, + } + + result = apply_import(pg_conn, user, doc) + + assert result["action"] == "updated" + updated = AgentsRepository(pg_conn).get(str(agent["id"]), user) + assert updated["json_schema"] is None + + +def test_custom_tool_dedup_across_reimport(pg_conn): + user = "u_dedup" + tool = UserToolsRepository(pg_conn).create( + user, "brave", custom_name="My Brave", display_name="Brave", config={}, + ) + agent = AgentsRepository(pg_conn).create( + user, + "Dedup Bot", + "published", + description="d", + chunks=2, + retriever="classic", + tools=[str(tool["id"])], + ) + ensure_agent_slug(pg_conn, agent, user) + agent = AgentsRepository(pg_conn).get(str(agent["id"]), user) + doc = parse_agent_yaml(agent_to_yaml(serialize_agent(pg_conn, agent, user))) + + apply_import(pg_conn, user, doc) + apply_import(pg_conn, user, doc) + + assert len(UserToolsRepository(pg_conn).list_for_user(user)) == 1 # reused, never duplicated + final = AgentsRepository(pg_conn).get(str(agent["id"]), user) + assert final["tools"] == [str(tool["id"])] + + +def test_import_rejects_ssrf_tool_url(pg_conn, monkeypatch): + user = "u_ssrf" + fake_tool = Mock() + fake_tool.get_config_requirements.return_value = {} + fake_tool.get_actions_metadata.return_value = [] + monkeypatch.setattr( + "application.api.user.agents.portability._tool_instance", + lambda tool_type: fake_tool, + ) + doc = _doc( + name="SSRF Bot", + _slug="ssrf", + tools=[ + { + "type": "mcp_tool", + "name": "evil", + "config": {"server_url": "http://169.254.169.254/latest/meta-data"}, + } + ], + ) + + result = apply_import(pg_conn, user, doc) + + agent = AgentsRepository(pg_conn).get(result["agent_id"], user) + assert agent["tools"] == [] + assert any("unsafe" in w.lower() for w in result["warnings"]) + + +def test_import_rejects_unowned_source_mapping(pg_conn): + owner, attacker = "u_owner", "u_attacker" + src = SourcesRepository(pg_conn).create("Secret KB", user_id=owner, type="file") + doc = _doc(name="IDOR Bot", _slug="idor", sources=[{"name": "Their Source"}]) + + # Attacker maps the missing source to the owner's source id. + result = apply_import( + pg_conn, attacker, doc, resolution={"sources": {"Their Source": str(src["id"])}} + ) + + agent = AgentsRepository(pg_conn).get(result["agent_id"], attacker) + assert list(agent["extra_source_ids"]) == [] # not linked + assert any("not yours" in w for w in result["warnings"]) + + +def test_apply_tolerates_plan_shaped_resolution(pg_conn): + """Echoing the plan's list-shaped resolution back must not crash apply.""" + user = "u_shape" + doc = _doc(name="Shape Bot", _slug="shape", sources=[{"name": "X"}]) + + result = apply_import( + pg_conn, user, doc, resolution={"sources": [], "tools": [], "models": []} + ) + + assert result["action"] == "created" + + +def test_update_clears_models_and_resets_prompt(pg_conn): + """A re-imported file that drops models and prompt clears them on update.""" + user = "u_clear2" + prompt = PromptsRepository(pg_conn).create(user, "Sys", "Be terse") + agent = AgentsRepository(pg_conn).create( + user, + "Model Bot", + "published", + description="d", + chunks=2, + retriever="classic", + prompt_id=str(prompt["id"]), + models=["gpt-4o"], + default_model_id="gpt-4o", + ) + slug = ensure_agent_slug(pg_conn, agent, user) + doc = { + "apiVersion": API_VERSION, + "kind": "Agent", + "metadata": {"slug": slug}, + "spec": { + "name": "Model Bot", + "retriever": "classic", + "prompt": "default", + "model": {"default": "", "available": []}, + }, + } + + apply_import(pg_conn, user, doc) + + updated = AgentsRepository(pg_conn).get(str(agent["id"]), user) + assert not updated["models"] + assert not updated["default_model_id"] + assert updated["prompt_id"] is None # reset to default + + +def test_tool_skipped_when_encryption_fails(pg_conn, monkeypatch): + user = "u_encfail" + fake_tool = Mock() + fake_tool.get_config_requirements.return_value = { + "api_key": {"secret": True, "required": True} + } + fake_tool.get_actions_metadata.return_value = [] + monkeypatch.setattr( + "application.api.user.agents.portability._tool_instance", + lambda tool_type: fake_tool, + ) + # Force credential encryption to fail (returns ""). + monkeypatch.setattr( + "application.api.user.tools.routes.encrypt_credentials", + lambda creds, user_id: "", + ) + doc = _doc( + name="EncFail Bot", + _slug="encfail", + tools=[{"type": "brave", "name": "X", "config": {}, "requires_secrets": ["api_key"]}], + ) + + result = apply_import( + pg_conn, + user, + doc, + resolution={"tools": {"tool-0": {"secrets": {"api_key": "K"}}}}, + ) + + agent = AgentsRepository(pg_conn).get(result["agent_id"], user) + assert agent["tools"] == [] + assert any("encryption failed" in w for w in result["warnings"])