From 381fb3335dd0cdaabfac0f2fa75720d0d8aae88c Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 24 Jun 2026 12:30:18 +0100 Subject: [PATCH] Add document generation tool with editable spec-backed artifacts Add an artifact generator tool with create, edit, and rewrite actions that render presentations, documents, spreadsheets, and PDFs from a validated JSON spec by running a fixed program in the sandbox (the spec travels as data, so its contents can never execute). The spec is stored as the source of truth on each artifact version; edits apply a JSON merge-patch and re-render, appending a new version while earlier versions stay intact. --- .../agents/tools/artifact_generator.py | 545 ++++++++++++++++++ application/agents/tools/tool_manager.py | 22 +- .../tools/test_artifact_generator_unit.py | 160 +++++ .../test_artifact_generator_e2e.py | 274 +++++++++ tests/requirements.txt | 7 + 5 files changed, 1006 insertions(+), 2 deletions(-) create mode 100644 application/agents/tools/artifact_generator.py create mode 100644 tests/agents/tools/test_artifact_generator_unit.py create mode 100644 tests/integration/test_artifact_generator_e2e.py diff --git a/application/agents/tools/artifact_generator.py b/application/agents/tools/artifact_generator.py new file mode 100644 index 00000000..2ed15620 --- /dev/null +++ b/application/agents/tools/artifact_generator.py @@ -0,0 +1,545 @@ +"""Artifact Generator tool: render editable documents from a JSON spec and version them append-only. + +The ``artifact_versions.spec`` JSONB is the source of truth; the rendered +``.pptx``/``.docx``/``.xlsx``/``.pdf`` is derived. ``create_artifact`` stores +v1, ``edit_artifact`` applies an RFC 7386 merge-patch to the current spec and +appends a version, ``rewrite_artifact`` replaces the spec wholesale and appends +a version. Rendering runs a FIXED program in the sandbox that reads the spec as +DATA (``json.loads``) — spec values are never interpolated into the program, so +a spec string containing code/quotes is rendered as literal text, not executed. +""" + +from __future__ import annotations + +import copy +import json +import logging +import uuid +from typing import Any, Dict, List, Optional + +from application.agents.tools.base import Tool +from application.core.settings import settings +from application.sandbox.artifacts_capture import append_artifact_version, persist_new_artifact +from application.sandbox.sandbox_creator import SandboxCreator +from application.storage.db.repositories.artifacts import ArtifactsRepository +from application.storage.db.session import db_readonly + +logger = logging.getLogger(__name__) + +try: + import jsonschema +except Exception: # pragma: no cover - jsonschema is a declared dependency + jsonschema = None # type: ignore[assignment] + +# Per-kind output metadata: artifact ``kind`` + produced file extension + mime. +_KIND_INFO: Dict[str, Dict[str, str]] = { + "presentation": { + "kind": "presentation", + "ext": "pptx", + "mime": "application/vnd.openxmlformats-officedocument.presentationml.presentation", + }, + "document": { + "kind": "document", + "ext": "docx", + "mime": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + }, + "spreadsheet": { + "kind": "spreadsheet", + "ext": "xlsx", + "mime": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + }, + "pdf": { + "kind": "document", + "ext": "pdf", + "mime": "application/pdf", + }, +} + +# Tight per-kind JSON schemas. ``additionalProperties: false`` keeps specs minimal +# and rejects stray keys before any rendering happens. +_SCHEMAS: Dict[str, Dict[str, Any]] = { + "presentation": { + "type": "object", + "additionalProperties": False, + "properties": { + "title": {"type": "string"}, + "slides": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "properties": { + "title": {"type": "string"}, + "bullets": {"type": "array", "items": {"type": "string"}}, + "notes": {"type": "string"}, + }, + }, + }, + }, + "required": ["slides"], + }, + "document": { + "type": "object", + "additionalProperties": False, + "properties": { + "title": {"type": "string"}, + "sections": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "properties": { + "heading": {"type": "string"}, + "paragraphs": {"type": "array", "items": {"type": "string"}}, + }, + }, + }, + }, + "required": ["sections"], + }, + "spreadsheet": { + "type": "object", + "additionalProperties": False, + "properties": { + "sheets": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "properties": { + "name": {"type": "string"}, + "rows": { + "type": "array", + "items": {"type": "array", "items": {}}, + }, + }, + "required": ["rows"], + }, + }, + }, + "required": ["sheets"], + }, + "pdf": { + "type": "object", + "additionalProperties": False, + "properties": { + "title": {"type": "string"}, + "blocks": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "properties": { + "type": {"type": "string", "enum": ["heading", "paragraph"]}, + "text": {"type": "string"}, + }, + "required": ["type", "text"], + }, + }, + }, + "required": ["blocks"], + }, +} + +# FIXED renderer programs. Each reads ``spec.json`` from the workspace as DATA +# and writes ``out.``. The spec is NEVER string-interpolated into the +# program; ``{spec_path}``/``{out_path}`` are server-controlled path literals. +_RENDERERS: Dict[str, str] = { + "presentation": ( + "import json\n" + "from pptx import Presentation\n" + "from pptx.util import Pt\n" + "spec = json.load(open({spec_path!r}))\n" + "prs = Presentation()\n" + "blank = prs.slide_layouts[6]\n" + "title_only = prs.slide_layouts[5]\n" + "for s in spec.get('slides', []):\n" + " slide = prs.slides.add_slide(title_only)\n" + " slide.shapes.title.text = str(s.get('title', '') or '')\n" + " bullets = s.get('bullets') or []\n" + " if bullets:\n" + " left = top = Pt(72)\n" + " width = prs.slide_width - Pt(144)\n" + " height = prs.slide_height - Pt(216)\n" + " box = slide.shapes.add_textbox(left, Pt(150), width, height)\n" + " tf = box.text_frame\n" + " tf.word_wrap = True\n" + " for i, b in enumerate(bullets):\n" + " para = tf.paragraphs[0] if i == 0 else tf.add_paragraph()\n" + " para.text = str(b)\n" + " notes = s.get('notes')\n" + " if notes:\n" + " slide.notes_slide.notes_text_frame.text = str(notes)\n" + "prs.save({out_path!r})\n" + ), + "document": ( + "import json\n" + "from docx import Document\n" + "spec = json.load(open({spec_path!r}))\n" + "doc = Document()\n" + "title = spec.get('title')\n" + "if title:\n" + " doc.add_heading(str(title), level=0)\n" + "for sec in spec.get('sections', []):\n" + " heading = sec.get('heading')\n" + " if heading:\n" + " doc.add_heading(str(heading), level=1)\n" + " for p in (sec.get('paragraphs') or []):\n" + " doc.add_paragraph(str(p))\n" + "doc.save({out_path!r})\n" + ), + "spreadsheet": ( + "import json\n" + "from openpyxl import Workbook\n" + "spec = json.load(open({spec_path!r}))\n" + "wb = Workbook()\n" + "wb.remove(wb.active)\n" + "for idx, sheet in enumerate(spec.get('sheets', [])):\n" + " name = str(sheet.get('name') or ('Sheet%d' % (idx + 1)))[:31]\n" + " ws = wb.create_sheet(title=name)\n" + " for row in (sheet.get('rows') or []):\n" + " ws.append([('' if c is None else c) for c in row])\n" + "if not wb.sheetnames:\n" + " wb.create_sheet(title='Sheet1')\n" + "wb.save({out_path!r})\n" + ), + "pdf": ( + "import json\n" + "from reportlab.lib.pagesizes import letter\n" + "from reportlab.lib.styles import getSampleStyleSheet\n" + "from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer\n" + "from xml.sax.saxutils import escape\n" + "spec = json.load(open({spec_path!r}))\n" + "styles = getSampleStyleSheet()\n" + "story = []\n" + "title = spec.get('title')\n" + "if title:\n" + " story.append(Paragraph(escape(str(title)), styles['Title']))\n" + " story.append(Spacer(1, 12))\n" + "for block in spec.get('blocks', []):\n" + " style = styles['Heading1'] if block.get('type') == 'heading' else styles['BodyText']\n" + " story.append(Paragraph(escape(str(block.get('text', ''))), style))\n" + " story.append(Spacer(1, 6))\n" + "SimpleDocTemplate({out_path!r}, pagesize=letter).build(story)\n" + ), +} + + +def merge_patch(target: Any, patch: Any) -> Any: + """Apply an RFC 7386 JSON Merge Patch to ``target`` and return the result.""" + if not isinstance(patch, dict): + return copy.deepcopy(patch) + if not isinstance(target, dict): + target = {} + result = copy.deepcopy(target) + for key, value in patch.items(): + if value is None: + result.pop(key, None) + else: + result[key] = merge_patch(result.get(key), value) + return result + + +class ArtifactGeneratorTool(Tool): + """Artifact Generator + Render editable documents (presentation/document/spreadsheet/pdf) from a JSON spec and version them. + """ + + def __init__(self, tool_config: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None) -> None: + """Bind the tool to the invoker and its conversation/run-scoped sandbox session.""" + self.config: Dict[str, Any] = tool_config or {} + self.user_id: Optional[str] = user_id + self.tool_id: Optional[str] = self.config.get("tool_id") + self.conversation_id: Optional[str] = self.config.get("conversation_id") + self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id") + self._last_artifact_id: Optional[str] = None + + # ------------------------------------------------------------------ + # Tool ABC + # ------------------------------------------------------------------ + def get_actions_metadata(self) -> List[Dict[str, Any]]: + """Return JSON metadata describing the create/edit/rewrite actions for tool schemas.""" + kinds = sorted(_KIND_INFO.keys()) + return [ + { + "name": "create_artifact", + "description": ( + "Render a new editable document from a JSON spec and store it as version 1. " + "The spec is the source of truth; the rendered file is derived." + ), + "active": True, + "parameters": { + "type": "object", + "properties": { + "kind": {"type": "string", "enum": kinds, "description": "Document kind to render."}, + "title": {"type": "string", "description": "Optional artifact title."}, + "spec": {"type": "object", "description": "Document spec matching the kind's schema."}, + }, + "required": ["kind", "spec"], + }, + }, + { + "name": "edit_artifact", + "description": ( + "Apply a JSON merge-patch (RFC 7386) to the current spec, re-render, and append a " + "new version. Preferred for small, targeted changes." + ), + "active": True, + "parameters": { + "type": "object", + "properties": { + "id": {"type": "string", "description": "Artifact id to edit."}, + "spec_patch": { + "type": "object", + "description": "RFC 7386 merge-patch; null values delete keys.", + }, + }, + "required": ["id", "spec_patch"], + }, + }, + { + "name": "rewrite_artifact", + "description": "Replace the spec wholesale, re-render, and append a new version.", + "active": True, + "parameters": { + "type": "object", + "properties": { + "id": {"type": "string", "description": "Artifact id to rewrite."}, + "spec": {"type": "object", "description": "Replacement spec matching the kind's schema."}, + }, + "required": ["id", "spec"], + }, + }, + ] + + def get_config_requirements(self) -> Dict[str, Any]: + """Return configuration requirements (none beyond the deployment sandbox backend).""" + return {} + + def get_artifact_id(self, action_name: str, **kwargs: Any) -> Optional[str]: + """Return the produced artifact id so the UI artifact rail lights up.""" + return self._last_artifact_id + + # ------------------------------------------------------------------ + # Dispatch + # ------------------------------------------------------------------ + def execute_action(self, action_name: str, **kwargs: Any) -> Dict[str, Any]: + """Dispatch a create/edit/rewrite action.""" + self._last_artifact_id = None + if not self.user_id: + return {"status": "error", "error": "artifact_generator requires a valid user_id."} + if self.conversation_id is None and self.workflow_run_id is None: + return {"status": "error", "error": "artifact_generator requires a conversation_id or workflow_run_id."} + if jsonschema is None: + return {"status": "error", "error": "jsonschema is required for spec validation."} + if action_name == "create_artifact": + return self._create(**kwargs) + if action_name == "edit_artifact": + return self._edit(**kwargs) + if action_name == "rewrite_artifact": + return self._rewrite(**kwargs) + return {"status": "error", "error": f"unknown action: {action_name}"} + + # ------------------------------------------------------------------ + # Actions + # ------------------------------------------------------------------ + def _create(self, **kwargs: Any) -> Dict[str, Any]: + """Validate, render, and persist a new artifact at version 1.""" + kind = kwargs.get("kind") + spec = kwargs.get("spec") + title = kwargs.get("title") + if kind not in _KIND_INFO: + return {"status": "error", "error": f"unsupported kind: {kind!r}; expected one of {sorted(_KIND_INFO)}."} + valid = self._validate(kind, spec) + if valid is not None: + return valid + + rendered = self._render(kind, spec) + if rendered.get("error"): + return {"status": "error", "error": rendered["error"]} + + info = _KIND_INFO[kind] + filename = self._filename(title, info["ext"]) + ref = persist_new_artifact( + user_id=self.user_id, + kind=info["kind"], + data=rendered["data"], + filename=filename, + mime_type=info["mime"], + title=title, + conversation_id=self.conversation_id, + workflow_run_id=self.workflow_run_id, + spec=spec, + produced_by=self._produced_by("create_artifact", kind), + ) + if ref is None: + return {"status": "error", "error": "failed to persist artifact."} + self._last_artifact_id = ref["artifact_id"] + return {"status": "ok", **ref} + + def _edit(self, **kwargs: Any) -> Dict[str, Any]: + """Merge-patch the current spec, re-render, and append a version.""" + artifact_id = kwargs.get("id") + spec_patch = kwargs.get("spec_patch") + if not isinstance(spec_patch, dict): + return {"status": "error", "error": "spec_patch must be a JSON object (merge-patch)."} + loaded = self._load_current(artifact_id) + if loaded.get("error"): + return {"status": "error", "error": loaded["error"]} + kind = loaded["kind"] + new_spec = merge_patch(loaded["spec"], spec_patch) + return self._reversion(artifact_id, kind, new_spec, "edit_artifact") + + def _rewrite(self, **kwargs: Any) -> Dict[str, Any]: + """Replace the spec wholesale, re-render, and append a version.""" + artifact_id = kwargs.get("id") + spec = kwargs.get("spec") + loaded = self._load_current(artifact_id) + if loaded.get("error"): + return {"status": "error", "error": loaded["error"]} + return self._reversion(artifact_id, loaded["kind"], spec, "rewrite_artifact") + + def _reversion(self, artifact_id: str, kind: str, spec: Any, action: str) -> Dict[str, Any]: + """Validate the new spec, re-render, and append the next version of an existing artifact.""" + valid = self._validate(kind, spec) + if valid is not None: + return valid + rendered = self._render(kind, spec) + if rendered.get("error"): + return {"status": "error", "error": rendered["error"]} + info = _KIND_INFO[kind] + filename = self._filename(None, info["ext"]) + ref = append_artifact_version( + user_id=self.user_id, + artifact_id=artifact_id, + data=rendered["data"], + filename=filename, + mime_type=info["mime"], + spec=spec, + produced_by=self._produced_by(action, kind), + ) + if ref is None: + return {"status": "error", "error": "failed to persist artifact version."} + self._last_artifact_id = ref["artifact_id"] + return {"status": "ok", **ref} + + # ------------------------------------------------------------------ + # Spec / render helpers + # ------------------------------------------------------------------ + def _validate(self, kind: str, spec: Any) -> Optional[Dict[str, Any]]: + """Return an error payload when ``spec`` is invalid for ``kind``, else None.""" + if not isinstance(spec, dict): + return {"status": "error", "error": "spec must be a JSON object."} + try: + jsonschema.validate(spec, _SCHEMAS[kind]) + except jsonschema.ValidationError as exc: + return {"status": "error", "error": f"invalid {kind} spec: {exc.message}"} + return None + + def _load_current(self, artifact_id: Any) -> Dict[str, Any]: + """Fetch the parent-scoped artifact and its current-version spec for edit/rewrite.""" + if not isinstance(artifact_id, str) or not artifact_id.strip(): + return {"error": "id is required."} + try: + with db_readonly() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact_in_parent( + artifact_id, + conversation_id=self.conversation_id, + workflow_run_id=self.workflow_run_id, + ) + if artifact is None: + return {"error": f"artifact {artifact_id} not found in this conversation/run."} + version = repo.get_version(artifact_id, artifact["current_version"]) + except Exception: + logger.exception("artifact_generator: failed to load artifact") + return {"error": f"failed to load artifact {artifact_id}."} + if not version or version.get("spec") is None: + return {"error": f"artifact {artifact_id} has no editable spec."} + kind = self._kind_for(artifact, version) + if kind is None: + return {"error": f"artifact {artifact_id} is not a spec-rendered document."} + return {"kind": kind, "spec": version["spec"]} + + @staticmethod + def _kind_for(artifact: Dict[str, Any], version: Dict[str, Any]) -> Optional[str]: + """Resolve the spec kind from ``produced_by`` (preferred) or the version mime type.""" + produced = version.get("produced_by") + if isinstance(produced, dict): + spec_kind = produced.get("spec_kind") + if spec_kind in _KIND_INFO: + return spec_kind + mime = version.get("mime_type") or "" + for spec_kind, info in _KIND_INFO.items(): + if info["mime"] == mime: + return spec_kind + return None + + def _render(self, kind: str, spec: Any) -> Dict[str, Any]: + """Run the fixed renderer in the sandbox and return the produced file bytes.""" + session_id = self._resolve_session_id() + if session_id is None: + return {"error": "artifact_generator requires a conversation_id or workflow_run_id."} + + token = uuid.uuid4().hex + spec_path = f"artifacts/{token}/spec.json" + out_path = f"artifacts/{token}/out.{_KIND_INFO[kind]['ext']}" + program = _RENDERERS[kind].format(spec_path=spec_path, out_path=out_path) + timeout = float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60)) + + manager = SandboxCreator.get_manager() + try: + manager.open(session_id, ttl=timeout) + except Exception as exc: + logger.exception("artifact_generator: failed to open sandbox session") + return {"error": f"sandbox unavailable: {type(exc).__name__}: {exc}"} + try: + # The spec rides in as a JSON file the program ``json.load``s; it is + # never interpolated into the program, so its contents stay data. + manager.put_file(session_id, spec_path, json.dumps(spec).encode("utf-8")) + result = manager.exec(session_id, program, timeout=timeout) + if not result.ok: + detail = ( + f"{result.error_name}: {result.error_value}" + if result.error_name + else (result.error_value or "render failed") + ) + return {"error": f"render failed: {detail}"} + data = manager.get_file(session_id, out_path) + except Exception as exc: + logger.exception("artifact_generator: render failed") + return {"error": f"render failed: {type(exc).__name__}: {exc}"} + finally: + try: + manager.close(session_id) + except Exception: + logger.exception("artifact_generator: session close failed") + if not data: + return {"error": "renderer produced an empty file."} + return {"data": data} + + # ------------------------------------------------------------------ + # Misc helpers + # ------------------------------------------------------------------ + def _produced_by(self, action: str, kind: str) -> Dict[str, Any]: + """Build the ``produced_by`` provenance record, carrying the spec kind for re-editing.""" + return { + "tool": "artifact_generator", + "action": action, + "spec_kind": kind, + "tool_id": self.tool_id, + } + + @staticmethod + def _filename(title: Optional[str], ext: str) -> str: + """Derive a download filename from a title (or a generic stem) plus the kind extension.""" + stem = (title or "artifact").strip() or "artifact" + return f"{stem}.{ext}" + + def _resolve_session_id(self) -> Optional[str]: + """Derive the sandbox session id from the bound conversation/run; sanitize to the gateway charset.""" + raw = self.conversation_id or self.workflow_run_id + if not raw: + return None + sanitized = "".join(c if c.isalnum() or c in "-_" else "-" for c in str(raw)) + return sanitized or None diff --git a/application/agents/tools/tool_manager.py b/application/agents/tools/tool_manager.py index a0dc35e6..579d5353 100644 --- a/application/agents/tools/tool_manager.py +++ b/application/agents/tools/tool_manager.py @@ -30,7 +30,16 @@ class ToolManager: if issubclass(obj, Tool) and obj is not Tool: if ( tool_name - in {"mcp_tool", "notes", "memory", "todo_list", "scheduler", "remote_device", "code_executor"} + in { + "mcp_tool", + "notes", + "memory", + "todo_list", + "scheduler", + "remote_device", + "code_executor", + "artifact_generator", + } and user_id ): return obj(tool_config, user_id) @@ -42,7 +51,16 @@ class ToolManager: raise ValueError(f"Tool '{tool_name}' not loaded") if ( tool_name - in {"mcp_tool", "memory", "todo_list", "notes", "scheduler", "remote_device", "code_executor"} + in { + "mcp_tool", + "memory", + "todo_list", + "notes", + "scheduler", + "remote_device", + "code_executor", + "artifact_generator", + } and user_id ): tool_config = self.config.get(tool_name, {}) diff --git a/tests/agents/tools/test_artifact_generator_unit.py b/tests/agents/tools/test_artifact_generator_unit.py new file mode 100644 index 00000000..5bcd6c00 --- /dev/null +++ b/tests/agents/tools/test_artifact_generator_unit.py @@ -0,0 +1,160 @@ +"""Unit tests for ArtifactGeneratorTool: spec validation, merge-patch, and renderer injection safety. + +No DB or sandbox: these exercise the pure logic (schema gate + RFC 7386 merge) +and the FIXED renderer programs directly (executed in-process against a temp dir) +to prove a spec value containing Python/quotes is rendered as literal text. +""" + +from __future__ import annotations + +import json +import os +import tempfile + +import pytest + +from application.agents.tools.artifact_generator import ( + _KIND_INFO, + _RENDERERS, + ArtifactGeneratorTool, + merge_patch, +) + +pytest.importorskip("pptx") +pytest.importorskip("docx") +pytest.importorskip("openpyxl") +pytest.importorskip("reportlab") + + +def _tool(): + return ArtifactGeneratorTool( + tool_config={"conversation_id": "conv-1", "tool_id": "t-1"}, user_id="u-1" + ) + + +# --------------------------------------------------------------------------- +# Spec validation +# --------------------------------------------------------------------------- + + +def test_validate_accepts_minimal_presentation(): + assert _tool()._validate("presentation", {"slides": [{"title": "x"}]}) is None + + +def test_validate_rejects_missing_required_key(): + err = _tool()._validate("presentation", {"title": "no slides"}) + assert err["status"] == "error" + assert "invalid presentation spec" in err["error"] + + +def test_validate_rejects_unknown_key(): + err = _tool()._validate("document", {"sections": [], "bogus": 1}) + assert err["status"] == "error" + + +def test_validate_rejects_wrong_type(): + err = _tool()._validate("spreadsheet", {"sheets": "not-a-list"}) + assert err["status"] == "error" + + +def test_validate_rejects_non_object_spec(): + err = _tool()._validate("pdf", "just a string") + assert err["status"] == "error" + assert "spec must be a JSON object" in err["error"] + + +def test_create_rejects_unknown_kind(): + out = _tool()._create(kind="hologram", spec={}) + assert out["status"] == "error" + assert "unsupported kind" in out["error"] + + +# --------------------------------------------------------------------------- +# RFC 7386 JSON merge-patch +# --------------------------------------------------------------------------- + + +def test_merge_patch_adds_and_overwrites(): + assert merge_patch({"a": 1, "b": 2}, {"b": 3, "c": 4}) == {"a": 1, "b": 3, "c": 4} + + +def test_merge_patch_null_deletes_key(): + assert merge_patch({"a": 1, "b": 2}, {"b": None}) == {"a": 1} + + +def test_merge_patch_recurses_into_objects(): + assert merge_patch({"x": {"a": 1, "b": 2}}, {"x": {"b": None, "c": 3}}) == {"x": {"a": 1, "c": 3}} + + +def test_merge_patch_replaces_array_wholesale(): + # RFC 7386: arrays are replaced, not merged element-wise. + assert merge_patch({"l": [1, 2, 3]}, {"l": [9]}) == {"l": [9]} + + +def test_merge_patch_non_object_patch_replaces_target(): + assert merge_patch({"a": 1}, "scalar") == "scalar" + + +def test_merge_patch_does_not_mutate_inputs(): + target = {"a": {"b": 1}} + patch = {"a": {"c": 2}} + merge_patch(target, patch) + assert target == {"a": {"b": 1}} + assert patch == {"a": {"c": 2}} + + +# --------------------------------------------------------------------------- +# Renderer injection safety +# --------------------------------------------------------------------------- + + +def _render_in_process(kind: str, spec: dict) -> str: + """Execute the FIXED renderer program against a temp dir; return the output path.""" + workdir = tempfile.mkdtemp() + spec_path = os.path.join(workdir, "spec.json") + out_path = os.path.join(workdir, f"out.{_KIND_INFO[kind]['ext']}") + with open(spec_path, "w") as handle: + json.dump(spec, handle) + program = _RENDERERS[kind].format(spec_path=spec_path, out_path=out_path) + namespace: dict = {} + exec(compile(program, "", "exec"), namespace, namespace) # noqa: S102 + return out_path + + +def test_renderer_does_not_execute_spec_code(capfd, tmp_path): + # A spec whose values are Python source / shell payloads must be treated as + # literal text. If the renderer string-interpolated the spec it would run + # this; instead it json.loads the spec as data, so nothing executes. + sentinel = tmp_path / "pwned.txt" + payload = ( + f"'''__import__('os').system('echo PWNED > {sentinel}')'''; " + "print('SHOULD_NOT_PRINT')" + ) + spec = { + "title": payload, + "slides": [{"title": payload, "bullets": [payload], "notes": payload}], + } + out_path = _render_in_process("presentation", spec) + + captured = capfd.readouterr() + assert "SHOULD_NOT_PRINT" not in captured.out + assert "PWNED" not in captured.out + assert not sentinel.exists() + assert os.path.getsize(out_path) > 0 + + +def test_renderer_keeps_injection_text_as_literal_content(): + from pptx import Presentation + + payload = "'''; import os; os.system('echo HACK'); x = '''" + spec = {"slides": [{"title": payload, "bullets": [payload]}]} + prs = Presentation(_render_in_process("presentation", spec)) + assert len(prs.slides) == 1 + assert prs.slides[0].shapes.title.text == payload + + +def test_pdf_renderer_escapes_markup_and_does_not_execute(): + payload = "not bold & \"\" '''os.system('x')'''" + spec = {"title": payload, "blocks": [{"type": "paragraph", "text": payload}]} + out_path = _render_in_process("pdf", spec) + assert os.path.getsize(out_path) > 0 diff --git a/tests/integration/test_artifact_generator_e2e.py b/tests/integration/test_artifact_generator_e2e.py new file mode 100644 index 00000000..20f8b54c --- /dev/null +++ b/tests/integration/test_artifact_generator_e2e.py @@ -0,0 +1,274 @@ +"""End-to-end ArtifactGeneratorTool: live Jupyter gateway + ephemeral Postgres + local storage. + +Launches a real ``jupyter kernelgateway`` (no Docker — the credential helper +hangs on dev machines), wires the tool to the ephemeral pytest-postgresql DB and +a temp-dir ``LocalStorage``, and drives create/edit/rewrite through the full +path: a spec is rendered in the kernel by the FIXED renderer, the produced file +is persisted as an artifact version (spec = source of truth), and the rendered +bytes are re-opened with the real library to assert structure. + +Skips gracefully when the gateway binary, websocket-client, or a renderer +library is unavailable. +""" + +from __future__ import annotations + +import io +import shutil +import socket +import subprocess +import time +import uuid + +import pytest + +requests = pytest.importorskip("requests") +pytest.importorskip("websocket") # websocket-client +pytest.importorskip("pptx") +pytest.importorskip("docx") +pytest.importorskip("openpyxl") + +from application.agents.tools.artifact_generator import ArtifactGeneratorTool # noqa: E402 +from application.sandbox.jupyter_gateway import JupyterKernelGatewaySandbox # noqa: E402 +from application.sandbox.manager import SandboxManager # noqa: E402 +from application.sandbox.sandbox_creator import SandboxCreator # noqa: E402 +from application.storage.db.repositories.artifacts import ArtifactsRepository # noqa: E402 +from application.storage.local import LocalStorage # noqa: E402 +from application.storage.storage_creator import StorageCreator # noqa: E402 + +_GATEWAY_BIN = shutil.which("jupyter-kernelgateway") or shutil.which("jupyter") + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + _GATEWAY_BIN is None, + reason="jupyter kernel gateway not installed (pip install jupyter-kernel-gateway)", + ), +] + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _gateway_cmd(port: int) -> list: + if _GATEWAY_BIN.endswith("jupyter-kernelgateway"): + base = [_GATEWAY_BIN] + else: + base = [_GATEWAY_BIN, "kernelgateway"] + return base + [ + "--KernelGatewayApp.ip=127.0.0.1", + f"--KernelGatewayApp.port={port}", + "--ZMQChannelsWebsocketConnection.limit_rate=False", + ] + + +@pytest.fixture(scope="module") +def gateway_url(): + port = _free_port() + proc = subprocess.Popen(_gateway_cmd(port), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + url = f"http://127.0.0.1:{port}" + deadline = time.time() + 30 + ready = False + try: + while time.time() < deadline: + if proc.poll() is not None: + pytest.skip("jupyter kernelgateway process exited during startup") + try: + if requests.get(f"{url}/api", timeout=1).status_code == 200: + ready = True + break + except requests.RequestException: + time.sleep(0.3) + if not ready: + pytest.skip("jupyter kernelgateway did not become ready in time") + yield url + finally: + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + +@pytest.fixture() +def wired_tool(gateway_url, pg_engine, tmp_path, monkeypatch): + """An ArtifactGeneratorTool wired to the live gateway, ephemeral PG, and a temp local store.""" + backend = JupyterKernelGatewaySandbox(gateway_url=gateway_url, default_timeout=60.0) + SandboxCreator._instance = SandboxManager(backend=backend, max_ttl=1200.0) + + storage = LocalStorage(base_dir=str(tmp_path)) + monkeypatch.setattr(StorageCreator, "_instance", storage, raising=False) + + monkeypatch.setattr("application.storage.db.session.get_engine", lambda: pg_engine) + + conversation_id = str(uuid.uuid4()) + tool = ArtifactGeneratorTool( + tool_config={"conversation_id": conversation_id, "tool_id": str(uuid.uuid4())}, + user_id="user-e2e", + ) + try: + yield tool, conversation_id, pg_engine, storage + finally: + SandboxCreator.reset() + + +def test_create_edit_rewrite_presentation_versions(wired_tool): + tool, conversation_id, pg_engine, storage = wired_tool + from pptx import Presentation + + # create -> v1 + spec = {"title": "Q3", "slides": [{"title": "Intro", "bullets": ["one", "two"]}]} + created = tool.execute_action("create_artifact", kind="presentation", title="Deck", spec=spec) + assert created["status"] == "ok", created + assert created["version"] == 1 + assert created["mime_type"].endswith("presentationml.presentation") + assert created["filename"] == "Deck.pptx" + artifact_id = created["artifact_id"] + assert tool.get_artifact_id("create_artifact") == artifact_id + + # v1 spec is the source of truth; the rendered .pptx re-opens with one slide. + with pg_engine.connect() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact_in_parent(artifact_id, conversation_id=conversation_id) + v1 = repo.get_version(artifact_id, 1) + assert artifact["current_version"] == 1 + assert v1["spec"] == spec + assert v1["produced_by"]["spec_kind"] == "presentation" + prs = Presentation(storage.get_file(v1["storage_path"])) + assert len(prs.slides) == 1 + assert prs.slides[0].shapes.title.text == "Intro" + + # edit (merge-patch) -> v2: add a second slide, keep title. + patch = {"slides": [{"title": "Intro", "bullets": ["one", "two"]}, {"title": "Outro"}]} + edited = tool.execute_action("edit_artifact", id=artifact_id, spec_patch=patch) + assert edited["status"] == "ok", edited + assert edited["version"] == 2 + assert edited["artifact_id"] == artifact_id + + with pg_engine.connect() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact_in_parent(artifact_id, conversation_id=conversation_id) + v1_again = repo.get_version(artifact_id, 1) + v2 = repo.get_version(artifact_id, 2) + assert artifact["current_version"] == 2 + # Append-only: v1 unchanged and still present. + assert v1_again["spec"] == spec + assert v2["spec"]["title"] == "Q3" + assert len(v2["spec"]["slides"]) == 2 + prs2 = Presentation(storage.get_file(v2["storage_path"])) + assert len(prs2.slides) == 2 + + # rewrite -> v3: wholesale new spec. + new_spec = {"title": "Fresh", "slides": [{"title": "A"}, {"title": "B"}, {"title": "C"}]} + rewritten = tool.execute_action("rewrite_artifact", id=artifact_id, spec=new_spec) + assert rewritten["status"] == "ok", rewritten + assert rewritten["version"] == 3 + + with pg_engine.connect() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact_in_parent(artifact_id, conversation_id=conversation_id) + versions = repo.list_versions(artifact_id) + v3 = repo.get_version(artifact_id, 3) + assert artifact["current_version"] == 3 + assert [v["version"] for v in versions] == [1, 2, 3] + assert v3["spec"] == new_spec + assert len(Presentation(storage.get_file(v3["storage_path"])).slides) == 3 + + +def test_create_document_docx(wired_tool): + tool, conversation_id, pg_engine, storage = wired_tool + from docx import Document + + spec = {"title": "Report", "sections": [{"heading": "Intro", "paragraphs": ["Hello.", "World."]}]} + created = tool.execute_action("create_artifact", kind="document", spec=spec) + assert created["status"] == "ok", created + assert created["mime_type"].endswith("wordprocessingml.document") + + with pg_engine.connect() as conn: + version = ArtifactsRepository(conn).get_version(created["artifact_id"], 1) + doc = Document(storage.get_file(version["storage_path"])) + texts = [p.text for p in doc.paragraphs] + assert "Hello." in texts and "World." in texts + + +def test_create_spreadsheet_xlsx(wired_tool): + tool, conversation_id, pg_engine, storage = wired_tool + from openpyxl import load_workbook + + spec = {"sheets": [{"name": "Data", "rows": [["a", "b"], [1, 2]]}]} + created = tool.execute_action("create_artifact", kind="spreadsheet", spec=spec) + assert created["status"] == "ok", created + assert created["mime_type"].endswith("spreadsheetml.sheet") + + with pg_engine.connect() as conn: + version = ArtifactsRepository(conn).get_version(created["artifact_id"], 1) + wb = load_workbook(storage.get_file(version["storage_path"])) + ws = wb["Data"] + assert ws.cell(row=1, column=1).value == "a" + assert ws.cell(row=2, column=2).value == 2 + + +def test_invalid_spec_creates_no_artifact(wired_tool): + tool, conversation_id, pg_engine, _storage = wired_tool + + out = tool.execute_action("create_artifact", kind="presentation", spec={"title": "no slides"}) + assert out["status"] == "error" + assert "invalid presentation spec" in out["error"] + + with pg_engine.connect() as conn: + rows = ArtifactsRepository(conn).list_artifacts(conversation_id=conversation_id) + assert rows == [] + + +def test_edit_cross_tenant_denied(wired_tool): + tool, conversation_id, pg_engine, storage = wired_tool + + # An artifact owned by a DIFFERENT conversation must not be editable here. + other_conversation = str(uuid.uuid4()) + foreign_id = _seed_presentation(pg_engine, storage, other_conversation) + + out = tool.execute_action("edit_artifact", id=foreign_id, spec_patch={"title": "x"}) + assert out["status"] == "error" + assert "not found in this conversation/run" in out["error"] + + # The foreign artifact stayed at v1 (no version appended). + with pg_engine.connect() as conn: + artifact = ArtifactsRepository(conn).get_artifact(foreign_id) + assert artifact["current_version"] == 1 + + +def _seed_presentation(pg_engine, storage, conversation_id) -> str: + """Create a presentation artifact (row + version + bytes) in a given conversation; return its id.""" + spec = {"slides": [{"title": "seed"}]} + data = b"%PDF-stub-bytes" # content is irrelevant to the authz check + with pg_engine.begin() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.create_artifact( + "user-e2e", + "presentation", + conversation_id=conversation_id, + title="seed", + mime_type="application/vnd.openxmlformats-officedocument.presentationml.presentation", + filename="seed.pptx", + storage_path=None, + size=len(data), + sha256="0" * 64, + spec=spec, + produced_by={"tool": "artifact_generator", "action": "create_artifact", "spec_kind": "presentation"}, + ) + artifact_id = str(artifact["id"]) + from sqlalchemy import text + + storage_path = f"inputs/user-e2e/artifacts/{artifact_id}/v1/seed.pptx" + storage.save_file(io.BytesIO(data), storage_path) + conn.execute( + text( + "UPDATE artifact_versions SET storage_path = :p " + "WHERE artifact_id = CAST(:aid AS uuid) AND version = 1" + ), + {"p": storage_path, "aid": artifact_id}, + ) + return artifact_id diff --git a/tests/requirements.txt b/tests/requirements.txt index f34900ea..04d6d036 100644 --- a/tests/requirements.txt +++ b/tests/requirements.txt @@ -11,3 +11,10 @@ cryptography>=46.0.0 jupyter-kernel-gateway>=2.5.0 ipykernel>=6.0.0 websocket-client>=1.6.0 +# Document renderers the artifact_generator runs inside the sandbox kernel. +# In production these are pre-baked into the sandbox image; locally the kernel +# uses the app .venv, so the integration test needs them installed. python-pptx +# and openpyxl already ship in application/requirements.txt; only the rest are +# test-only here (kept out of application/requirements.txt — sandbox-image libs). +python-docx>=1.1 +reportlab>=4.0,<5