From 12dfefdd055f8251d449ec651687b8718c10e560 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 8 Jul 2026 20:06:31 +0100 Subject: [PATCH] fix: clean up and better artefact consistency --- .../agents/tools/artifact_generator.py | 23 +++++-- application/agents/workflow_agent.py | 3 +- application/agents/workflows/schemas.py | 2 - application/api/answer/routes/base.py | 17 ++--- application/core/settings.py | 6 +- application/sandbox/jupyter_gateway.py | 7 +- deployment/sandbox/gateway-launch.sh | 11 +++- frontend/src/agents/types/workflow.ts | 1 - .../agents/workflow/codeNodeConfig.test.ts | 10 --- .../src/agents/workflow/codeNodeConfig.ts | 8 --- tests/agents/test_workflow_input_documents.py | 8 ++- .../tools/test_artifact_generator_unit.py | 43 ++++++++++++ tests/api/answer/routes/test_base.py | 65 +++++++++++++++++++ 13 files changed, 160 insertions(+), 44 deletions(-) diff --git a/application/agents/tools/artifact_generator.py b/application/agents/tools/artifact_generator.py index e1457e53..a96425f7 100644 --- a/application/agents/tools/artifact_generator.py +++ b/application/agents/tools/artifact_generator.py @@ -551,7 +551,9 @@ class ArtifactGeneratorTool(Tool): if loaded.get("error"): return {"status": "error", "error": loaded["error"]} new_spec = merge_patch(loaded["spec"], spec_patch) - return self._reversion(loaded["artifact_id"], loaded["kind"], new_spec, "edit_artifact") + return self._reversion( + loaded["artifact_id"], loaded["kind"], new_spec, "edit_artifact", loaded.get("title") + ) def _rewrite(self, **kwargs: Any) -> Dict[str, Any]: """Replace the spec wholesale, re-render, and append a version.""" @@ -559,9 +561,13 @@ class ArtifactGeneratorTool(Tool): loaded = self._load_current(kwargs.get("id")) if loaded.get("error"): return {"status": "error", "error": loaded["error"]} - return self._reversion(loaded["artifact_id"], loaded["kind"], spec, "rewrite_artifact") + return self._reversion( + loaded["artifact_id"], loaded["kind"], spec, "rewrite_artifact", loaded.get("title") + ) - def _reversion(self, artifact_id: str, kind: str, spec: Any, action: str) -> Dict[str, Any]: + def _reversion( + self, artifact_id: str, kind: str, spec: Any, action: str, title: Optional[str] = None + ) -> 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: @@ -570,7 +576,9 @@ class ArtifactGeneratorTool(Tool): if rendered.get("error"): return {"status": "error", "error": rendered["error"]} info = _KIND_INFO[kind] - filename = self._filename(None, kind) + # Keep the original artifact's download name across versions (v2 of "Q3 Deck" + # must stay "Q3 Deck.pptx", not a generic "artifact.pptx"). + filename = self._filename(title, kind) try: ref = append_artifact_version( user_id=self.user_id, @@ -637,7 +645,12 @@ class ArtifactGeneratorTool(Tool): kind = self._kind_for(artifact, version) if kind is None: return {"error": f"artifact {raw_id} is not a spec-rendered document."} - return {"artifact_id": artifact_id, "kind": kind, "spec": version["spec"]} + return { + "artifact_id": artifact_id, + "kind": kind, + "spec": version["spec"], + "title": artifact.get("title"), + } @staticmethod def _kind_for(artifact: Dict[str, Any], version: Dict[str, Any]) -> Optional[str]: diff --git a/application/agents/workflow_agent.py b/application/agents/workflow_agent.py index b75049a0..331cd263 100644 --- a/application/agents/workflow_agent.py +++ b/application/agents/workflow_agent.py @@ -81,6 +81,7 @@ class WorkflowAgent(BaseAgent): self._bridge_error = str(exc) yield { "type": "error", + "user_facing": True, "error": ( "This run's input documents exceed your artifact storage quota. " "Delete some artifacts and try again." @@ -92,7 +93,7 @@ class WorkflowAgent(BaseAgent): # Non-fatal: some attachments were dropped (oversize / unreadable). Tell the # user which, then still run with the documents that did bridge. if dropped: - yield {"type": "error", "error": " ".join(dropped)} + yield {"type": "error", "user_facing": True, "error": " ".join(dropped)} self._engine.run_persisted = self._run_persisted interrupted = True diff --git a/application/agents/workflows/schemas.py b/application/agents/workflows/schemas.py index 4ccd1762..f5a806c4 100644 --- a/application/agents/workflows/schemas.py +++ b/application/agents/workflows/schemas.py @@ -63,9 +63,7 @@ class AgentNodeConfig(BaseModel): class CodeNodeConfig(BaseModel): model_config = ConfigDict(extra="allow") code: str = "" - language: str = "python" inputs: List[str] = Field(default_factory=list) - libraries: List[str] = Field(default_factory=list) output_variable: Optional[str] = None timeout: Optional[int] = None json_schema: Optional[Dict[str, Any]] = None diff --git a/application/api/answer/routes/base.py b/application/api/answer/routes/base.py index c5654afa..6fa3574b 100644 --- a/application/api/answer/routes/base.py +++ b/application/api/answer/routes/base.py @@ -543,14 +543,15 @@ class BaseAnswerResource: paused = True yield _emit(line) elif line.get("type") == "error": - yield _emit( - { - "type": "error", - "error": sanitize_api_error( - line.get("error", "An error occurred") - ), - } - ) + # An event flagged ``user_facing`` already carries a curated, + # actionable message (e.g. an artifact-quota notice). Passing it + # through sanitize_api_error would substring-match words like + # "quota" and rewrite it into a misleading rate-limit message, so + # emit it verbatim; sanitize only raw/technical errors. + error_text = line.get("error", "An error occurred") + if not line.get("user_facing"): + error_text = sanitize_api_error(error_text) + yield _emit({"type": "error", "error": error_text}) else: yield _emit(line) if is_structured and structured_chunks: diff --git a/application/core/settings.py b/application/core/settings.py index e02c83ed..dcfd3946 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -334,7 +334,11 @@ class Settings(BaseSettings): # URL of the Jupyter Kernel Gateway runner (the docsgpt-sandbox service). SANDBOX_GATEWAY_URL: str = "http://localhost:8888" SANDBOX_GATEWAY_AUTH_TOKEN: Optional[str] = None # gateway auth token, if set - SANDBOX_KERNEL_NAME: str = "python3" # kernelspec name to launch per session + # Kernelspec launched per session. Defaults to the env-scrubbing "docsgpt-python" + # spec (shipped by the docsgpt-sandbox runner) so kernel code cannot read the + # gateway auth token or operator secrets from os.environ. The stock "python3" + # spec inherits the gateway env verbatim and must not be used with untrusted code. + SANDBOX_KERNEL_NAME: str = "docsgpt-python" SANDBOX_MAX_TTL: int = 1200 # hard cap (s) on agent-selectable keep-alive TTL # Per-process/worker cap on concurrent live sandbox sessions. Backend-agnostic # (complements DAYTONA_MAX_SANDBOXES); when reached, an LRU-idle session is diff --git a/application/sandbox/jupyter_gateway.py b/application/sandbox/jupyter_gateway.py index 0155021f..fe917e6b 100644 --- a/application/sandbox/jupyter_gateway.py +++ b/application/sandbox/jupyter_gateway.py @@ -9,7 +9,7 @@ import threading import time import uuid from typing import Dict, List, Optional -from urllib.parse import urlencode, urlparse, urlunparse +from urllib.parse import urlparse, urlunparse import requests import websocket @@ -102,11 +102,12 @@ class JupyterKernelGatewaySandbox(CodeSandbox): return headers def _ws_url(self, kernel_id: str) -> str: + # The token is sent only via the Authorization header (_ws_headers); keeping it + # out of the URL query avoids leaking it into gateway/proxy access logs. parsed = urlparse(self._base_url) scheme = "wss" if parsed.scheme == "https" else "ws" path = f"/api/kernels/{kernel_id}/channels" - query = urlencode({"token": self._auth_token}) if self._auth_token else "" - return urlunparse((scheme, parsed.netloc, path, "", query, "")) + return urlunparse((scheme, parsed.netloc, path, "", "", "")) def _get_kernel(self, session_id: str) -> _Kernel: with self._lock: diff --git a/deployment/sandbox/gateway-launch.sh b/deployment/sandbox/gateway-launch.sh index 0e1ab601..c49d811f 100755 --- a/deployment/sandbox/gateway-launch.sh +++ b/deployment/sandbox/gateway-launch.sh @@ -9,9 +9,14 @@ # fail closed if it is missing -- an unauthenticated gateway must never start. # # The token is shared with the app via SANDBOX_GATEWAY_AUTH_TOKEN (the app sends -# it as `Authorization: token <...>`). Kernel code cannot read it: kernel-launch.sh -# re-execs ipykernel under a scrubbed `env -i` allowlist that excludes the token, -# so it is present for the gateway process only. +# it as `Authorization: token <...>`). kernel-launch.sh re-execs ipykernel under a +# scrubbed `env -i` allowlist that excludes the token, so it is NOT in kernel code's +# os.environ. NOTE: this is not a hard secret boundary -- the gateway and kernels +# share one container and uid, so kernel code can still read the gateway's env via +# /proc//environ. The token limits the blast radius (an outside caller +# cannot drive the gateway), but the Jupyter runner is a single trust domain: for +# real isolation against untrusted code use the Daytona backend or run under gVisor +# (see deployment/sandbox/README.md). set -eu TOKEN="${SANDBOX_GATEWAY_AUTH_TOKEN:-}" diff --git a/frontend/src/agents/types/workflow.ts b/frontend/src/agents/types/workflow.ts index d2f5c27b..370dfb25 100644 --- a/frontend/src/agents/types/workflow.ts +++ b/frontend/src/agents/types/workflow.ts @@ -9,7 +9,6 @@ export type NodeType = export interface CodeNodeConfig { code: string; - language: string; inputs: string[]; output_variable?: string; timeout?: number; diff --git a/frontend/src/agents/workflow/codeNodeConfig.test.ts b/frontend/src/agents/workflow/codeNodeConfig.test.ts index 2c4490f5..18373f34 100644 --- a/frontend/src/agents/workflow/codeNodeConfig.test.ts +++ b/frontend/src/agents/workflow/codeNodeConfig.test.ts @@ -81,7 +81,6 @@ describe('createDefaultCodeConfig', () => { it('returns python defaults with empty code and inputs', () => { expect(createDefaultCodeConfig()).toEqual({ code: '', - language: 'python', inputs: [], }); }); @@ -91,13 +90,11 @@ describe('serializeCodeConfig', () => { it('omits optional fields when empty', () => { const serialized = serializeCodeConfig({ code: 'print(1)', - language: 'python', inputs: ['', ' '], output_variable: ' ', }); expect(serialized).toEqual({ code: 'print(1)', - language: 'python', inputs: [], }); expect('output_variable' in serialized).toBe(false); @@ -108,7 +105,6 @@ describe('serializeCodeConfig', () => { it('keeps all CodeNodeConfig fields when present', () => { const serialized = serializeCodeConfig({ code: 'print(1)', - language: 'python', inputs: ['A1', 'state_var'], output_variable: 'result', timeout: 30, @@ -116,7 +112,6 @@ describe('serializeCodeConfig', () => { }); expect(serialized).toEqual({ code: 'print(1)', - language: 'python', inputs: ['A1', 'state_var'], output_variable: 'result', timeout: 30, @@ -127,7 +122,6 @@ describe('serializeCodeConfig', () => { it('falls back to python defaults for an undefined config', () => { expect(serializeCodeConfig(undefined)).toEqual({ code: '', - language: 'python', inputs: [], }); }); @@ -137,7 +131,6 @@ describe('normalizeCodeConfig', () => { it('tolerates a missing/empty payload', () => { expect(normalizeCodeConfig(undefined)).toEqual({ code: '', - language: 'python', inputs: [], }); }); @@ -145,7 +138,6 @@ describe('normalizeCodeConfig', () => { it('coerces bad field types without throwing', () => { const config = normalizeCodeConfig({ code: 42, - language: '', inputs: ['A1', 5, 'B2'], output_variable: '', timeout: 'soon', @@ -153,7 +145,6 @@ describe('normalizeCodeConfig', () => { } as unknown as Record); expect(config).toEqual({ code: '', - language: 'python', inputs: ['A1', 'B2'], }); }); @@ -161,7 +152,6 @@ describe('normalizeCodeConfig', () => { it('preserves a fully populated saved config (round-trip with serialize)', () => { const saved = { code: 'print(1)', - language: 'python', inputs: ['A1'], output_variable: 'result', timeout: 30, diff --git a/frontend/src/agents/workflow/codeNodeConfig.ts b/frontend/src/agents/workflow/codeNodeConfig.ts index df370c0a..ff03dbfd 100644 --- a/frontend/src/agents/workflow/codeNodeConfig.ts +++ b/frontend/src/agents/workflow/codeNodeConfig.ts @@ -1,7 +1,5 @@ import type { CodeNodeConfig } from '../types/workflow'; -const DEFAULT_LANGUAGE = 'python'; - /** * Validate a parsed JSON schema object the same way the agent node does: * it must be a non-array object that declares a "type" or "schema" field. @@ -60,7 +58,6 @@ export function stringifyCodeInputs(inputs: string[] | undefined): string { export function createDefaultCodeConfig(): CodeNodeConfig { return { code: '', - language: DEFAULT_LANGUAGE, inputs: [], }; } @@ -75,10 +72,6 @@ export function normalizeCodeConfig( const source = raw || {}; const config: CodeNodeConfig = { code: typeof source.code === 'string' ? source.code : '', - language: - typeof source.language === 'string' && source.language.trim() !== '' - ? source.language - : DEFAULT_LANGUAGE, inputs: Array.isArray(source.inputs) ? source.inputs.filter( (entry): entry is string => typeof entry === 'string', @@ -111,7 +104,6 @@ export function serializeCodeConfig( const source = config || createDefaultCodeConfig(); const serialized: CodeNodeConfig = { code: source.code || '', - language: source.language || DEFAULT_LANGUAGE, inputs: (source.inputs || []).filter( (entry) => typeof entry === 'string' && entry.trim() !== '', ), diff --git a/tests/agents/test_workflow_input_documents.py b/tests/agents/test_workflow_input_documents.py index 4414708e..f7e53730 100644 --- a/tests/agents/test_workflow_input_documents.py +++ b/tests/agents/test_workflow_input_documents.py @@ -378,9 +378,11 @@ def test_quota_exceeded_fails_run_and_does_not_execute(pg_engine, tmp_path, monk events = list(agent._gen_inner("summarize", log_context=None)) engine = _RecordingEngine.instances[-1] - # A fatal error was surfaced on the stream. + # A fatal error was surfaced on the stream, flagged user_facing so the route's + # sanitize_api_error does not rewrite the "quota" wording into a rate-limit message. errors = [e for e in events if e.get("type") == "error"] assert errors and "quota" in errors[0]["error"].lower() + assert errors[0].get("user_facing") is True # The engine never executed (execute was not reached -> no captured inputs). assert engine.captured_inputs is None # The pre-created RUNNING row was finalized FAILED (not left dangling), and no @@ -413,9 +415,11 @@ def test_oversize_declared_attachment_skipped_with_notice(pg_engine, tmp_path, m # The oversize doc was dropped -> no input documents bridged, but the run ran. assert engine.captured_inputs is not None assert engine.captured_inputs["input_documents"] == [] - # A non-fatal notice naming the dropped document was surfaced. + # A non-fatal notice naming the dropped document was surfaced, flagged user_facing + # so the route emits the filename verbatim instead of a generic sanitized message. errors = [e for e in events if e.get("type") == "error"] assert errors and "big.txt" in errors[0]["error"] + assert errors[0].get("user_facing") is True # Nothing was persisted for the oversize doc. with pg_engine.connect() as conn: n = conn.execute( diff --git a/tests/agents/tools/test_artifact_generator_unit.py b/tests/agents/tools/test_artifact_generator_unit.py index 6769d893..1b04e35f 100644 --- a/tests/agents/tools/test_artifact_generator_unit.py +++ b/tests/agents/tools/test_artifact_generator_unit.py @@ -53,6 +53,49 @@ def test_validate_rejects_unknown_key(): assert err["status"] == "error" +# --------------------------------------------------------------------------- +# Reversion download filename +# --------------------------------------------------------------------------- + + +def _patch_reversion(monkeypatch, tool): + """Stub render + persist so _reversion runs without a sandbox or DB; capture kwargs.""" + import application.agents.tools.artifact_generator as ag + + captured: dict = {} + + def _fake_append(**kwargs): + captured.update(kwargs) + return {"artifact_id": kwargs["artifact_id"], "version": 2, "id": "v2"} + + monkeypatch.setattr(ag, "append_artifact_version", _fake_append) + monkeypatch.setattr(tool, "_render", lambda kind, spec: {"data": b"x"}) + return captured + + +def test_reversion_preserves_title_in_filename(monkeypatch): + """An edited/rewritten version keeps the original artifact's download name.""" + tool = _tool() + captured = _patch_reversion(monkeypatch, tool) + + result = tool._reversion( + "art-1", "presentation", {"slides": [{"title": "x"}]}, "edit_artifact", "Q3 Deck" + ) + + assert result["status"] == "ok" + assert captured["filename"] == "Q3 Deck.pptx" + + +def test_reversion_without_title_falls_back_to_generic(monkeypatch): + """A missing title still yields a valid generic filename (never crashes).""" + tool = _tool() + captured = _patch_reversion(monkeypatch, tool) + + tool._reversion("art-1", "presentation", {"slides": [{"title": "x"}]}, "rewrite_artifact", None) + + assert captured["filename"] == "artifact.pptx" + + def test_validate_rejects_wrong_type(): err = _tool()._validate("spreadsheet", {"sheets": "not-a-list"}) assert err["status"] == "error" diff --git a/tests/api/answer/routes/test_base.py b/tests/api/answer/routes/test_base.py index 63745692..65da96d0 100644 --- a/tests/api/answer/routes/test_base.py +++ b/tests/api/answer/routes/test_base.py @@ -203,6 +203,71 @@ class TestCompleteStreamMethod: assert any('"type": "error"' in s for s in stream) + def test_user_facing_error_is_not_sanitized(self, mock_mongo_db, flask_app): + """A user_facing error (e.g. an artifact-quota notice) streams verbatim. + + Without the flag, sanitize_api_error substring-matches "quota" and rewrites the + message into a misleading rate-limit notice. + """ + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + + mock_agent = MagicMock() + mock_agent.gen.return_value = iter( + [ + { + "type": "error", + "user_facing": True, + "error": "This run's input documents exceed your artifact storage quota.", + } + ] + ) + + stream = list( + resource.complete_stream( + question="Test?", + agent=mock_agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "user123"}, + should_persist=False, + ) + ) + + error_chunks = [s for s in stream if '"type": "error"' in s] + assert error_chunks + assert "artifact storage quota" in error_chunks[0] + assert "Rate limit exceeded" not in error_chunks[0] + + def test_non_user_facing_error_is_sanitized(self, mock_mongo_db, flask_app): + """A raw error without the flag is still routed through sanitize_api_error.""" + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + + mock_agent = MagicMock() + mock_agent.gen.return_value = iter( + [{"type": "error", "error": "OpenAI 429: quota exceeded for this key"}] + ) + + stream = list( + resource.complete_stream( + question="Test?", + agent=mock_agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "user123"}, + should_persist=False, + ) + ) + + error_chunks = [s for s in stream if '"type": "error"' in s] + assert error_chunks + assert "Rate limit exceeded" in error_chunks[0] + def test_saves_conversation_when_enabled(self, mock_mongo_db, flask_app): from application.api.answer.routes.base import BaseAnswerResource