diff --git a/.gitignore b/.gitignore index ffd7a66c4..da134e0bc 100644 --- a/.gitignore +++ b/.gitignore @@ -71,6 +71,12 @@ tmp/ temp/ *.tmp +# Local sandbox artifact dumps (regeneratable, never committed) +artifacts_out*/ + +# OpenShell policy files generated by scripts/setup_openshell.sh +configs/openshell/generated/ + # ============================================================================ # Evaluation Dataset, Results and Outputs (Added for evaluation suite) diff --git a/configs/config_openshell.yml b/configs/config_openshell.yml new file mode 100644 index 000000000..1d6537d18 --- /dev/null +++ b/configs/config_openshell.yml @@ -0,0 +1,176 @@ +# EXPERIMENTAL: AI-Q deep research over a local, single-operator OpenShell sandbox. +# Run ./scripts/setup_openshell.sh first to provision the gateway and named sandbox. +# Jobs attach to that shared sandbox; per-job directories prevent filename collisions +# but are not a security boundary. Do not run mutually untrusted jobs concurrently. + +general: + use_uvloop: true + telemetry: + logging: + console: + _type: console + level: INFO + + front_end: + _type: aiq_api + runner_class: aiq_api.plugin.AIQAPIWorker + db_url: ${NAT_JOB_STORE_DB_URL:-sqlite+aiosqlite:///./jobs.db} + expiry_seconds: 86400 + cors: + allow_origin_regex: 'http://localhost(:\d+)?|http://127.0.0.1(:\d+)?' + allow_methods: + - GET + - POST + - DELETE + - OPTIONS + allow_headers: + - "*" + allow_credentials: true + expose_headers: + - "*" + +llms: + nemotron_llm_intent: + _type: nim + model_name: nvidia/nemotron-3-super-120b-a12b + base_url: "https://integrate.api.nvidia.com/v1" + temperature: 0.5 + top_p: 0.9 + max_tokens: 4096 + num_retries: 5 + chat_template_kwargs: + enable_thinking: true + + nemotron_super_llm: + _type: nim + model_name: nvidia/nemotron-3-super-120b-a12b + base_url: "https://integrate.api.nvidia.com/v1" + temperature: 0.7 + top_p: 0.7 + max_tokens: 65536 + num_retries: 5 + chat_template_kwargs: + enable_thinking: true + + gpt_oss_llm: + _type: nim + model_name: openai/gpt-oss-120b + base_url: https://integrate.api.nvidia.com/v1 + temperature: 1.0 + top_p: 1.0 + max_tokens: 256000 + api_key: ${NVIDIA_API_KEY} + max_retries: 10 + + summary_llm: + _type: nim + model_name: nvidia/nemotron-mini-4b-instruct + base_url: "https://integrate.api.nvidia.com/v1" + api_key: ${NVIDIA_API_KEY} + temperature: 0.3 + max_tokens: 100 + +functions: + data_sources: + _type: data_source_registry + sources: + - id: web_search + name: "Web Search" + description: "Search the web for real-time information." + tools: + - web_search_tool + - advanced_web_search_tool + - id: knowledge_layer + name: "Knowledge Base" + description: "Search uploaded documents and files." + tools: + - knowledge_search + + web_search_tool: + _type: tavily_web_search + max_results: 5 + max_content_length: 1000 + + advanced_web_search_tool: + _type: tavily_web_search + max_results: 2 + advanced_search: true + + knowledge_search: + _type: knowledge_retrieval + backend: llamaindex + collection_name: ${COLLECTION_NAME:-test_collection} + generate_summary: true + summary_model: summary_llm + summary_db: ${AIQ_SUMMARY_DB:-sqlite+aiosqlite:///./summaries.db} + top_k: 5 + chroma_dir: ${AIQ_CHROMA_DIR:-/tmp/chroma_data} + + intent_classifier: + _type: intent_classifier + llm: nemotron_llm_intent + verbose: true + + clarifier_agent: + _type: clarifier_agent + llm: nemotron_super_llm + planner_llm: nemotron_super_llm + max_turns: 3 + enable_plan_approval: true + log_response_max_chars: 2000 + verbose: true + + shallow_research_agent: + _type: shallow_research_agent + llm: nemotron_super_llm + exclude_tools: + - advanced_web_search_tool + verbose: true + max_llm_turns: 10 + max_tool_iterations: 5 + + # OpenShell configuration: skills + sandbox + deep_research_skills: + _type: deep_research_skills + agents: + researcher-agent: [research] + writer-agent: [synthesis, research] + require_sandbox: + - research + + deep_research_sandbox: + _type: deep_research_sandbox + provider: openshell + sandbox_name: ${AIQ_OPENSHELL_SANDBOX_NAME:-aiq-openshell-demo} + policy: ${AIQ_OPENSHELL_POLICY_FILE:-configs/openshell/generated/aiq-openshell-policy.yaml} + workdir: /sandbox + network: blocked + timeout: 1200 + idle_timeout: 1800 + delete_on_exit: false + artifact_capture: + enabled: true + max_file_bytes: 50000000 + allow_extensions: [.png, .jpg, .jpeg, .webp, .csv, .json, .md, .ipynb, .pdf] + + deep_research_agent: + _type: deep_research_agent + enable_citation_verification: true + orchestrator_llm: gpt_oss_llm + source_router_llm: nemotron_super_llm + researcher_llm: nemotron_super_llm + planner_llm: gpt_oss_llm + writer_llm: gpt_oss_llm + exclude_tools: + - web_search_tool + verbose: true + skills: deep_research_skills + sandbox: deep_research_sandbox + +workflow: + _type: chat_deepresearcher_agent + verbose: true + enable_escalation: true + enable_clarifier: true + use_async_deep_research: true + checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} diff --git a/configs/openshell/aiq-research-policy.yaml b/configs/openshell/aiq-research-policy.yaml new file mode 100644 index 000000000..8bb9775fa --- /dev/null +++ b/configs/openshell/aiq-research-policy.yaml @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Generated by scripts/setup_openshell.sh. + +version: 1 + +filesystem_policy: + include_workdir: true + read_only: + - /usr + - /lib + - /etc + - /app + - /var/log + - /proc/self + - /dev/urandom + read_write: + - /sandbox + - /workspace + - /tmp + - /dev/null + +# best_effort lets the sandbox start on hosts without Landlock (e.g. Docker Desktop on +# macOS), but there filesystem confinement is silently dropped. Acceptable only for the +# local single-operator demo; production must use `hard_requirement` (fail closed). +landlock: + compatibility: best_effort + +process: + run_as_user: sandbox + run_as_group: sandbox + +network_policies: + github: + name: github-readonly + endpoints: + - host: api.github.com + port: 443 + protocol: rest + enforcement: enforce + access: read-only + - host: github.com + port: 443 + protocol: rest + enforcement: enforce + access: read-only + binaries: + - { path: /usr/bin/curl } + - { path: /usr/local/bin/python3 } + - { path: /usr/local/bin/python } + - { path: /usr/local/bin/pip } + - { path: /usr/local/bin/pip3 } + nvidia: + name: nvidia-api-readonly + endpoints: + - host: integrate.api.nvidia.com + port: 443 + protocol: rest + enforcement: enforce + access: read-only + binaries: + - { path: /usr/bin/curl } + - { path: /usr/local/bin/python3 } + - { path: /usr/local/bin/python } + - { path: /usr/local/bin/pip } + - { path: /usr/local/bin/pip3 } + tavily: + name: tavily-api-readonly + endpoints: + - host: api.tavily.com + port: 443 + protocol: rest + enforcement: enforce + access: read-only + binaries: + - { path: /usr/bin/curl } + - { path: /usr/local/bin/python3 } + - { path: /usr/local/bin/python } + - { path: /usr/local/bin/pip } + - { path: /usr/local/bin/pip3 } + serper: + name: serper-api-readonly + endpoints: + - host: google.serper.dev + port: 443 + protocol: rest + enforcement: enforce + access: read-only + binaries: + - { path: /usr/bin/curl } + - { path: /usr/local/bin/python3 } + - { path: /usr/local/bin/python } + - { path: /usr/local/bin/pip } + - { path: /usr/local/bin/pip3 } diff --git a/deploy/openshell/Dockerfile.aiq-demo b/deploy/openshell/Dockerfile.aiq-demo new file mode 100644 index 000000000..06ff906ce --- /dev/null +++ b/deploy/openshell/Dockerfile.aiq-demo @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +FROM python:3.13-slim + +ENV DEBIAN_FRONTEND=noninteractive + +# Optional OpenShell sandbox log verbosity. Defaults to `warn` (OpenShell's stock +# sandbox level, so default behavior is unchanged). Rebuild with +# `--build-arg OPENSHELL_SANDBOX_LOG_LEVEL=debug` (or `setup_openshell.sh +# --sandbox-log-level debug`) to surface process/relay detail (`OCSF PROC:` etc.) +# in the container logs (`/var/log/openshell.*.log`, `openshell logs `). +ARG OPENSHELL_SANDBOX_LOG_LEVEL=warn +ENV RUST_LOG=${OPENSHELL_SANDBOX_LOG_LEVEL} + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + bash \ + ca-certificates \ + curl \ + iproute2 \ + iptables \ + procps \ + && python -m pip install --no-cache-dir numpy pandas matplotlib pillow tabulate requests \ + && rm -rf /var/lib/apt/lists/* \ + && groupadd -r sandbox \ + && useradd -r -g sandbox -d /sandbox -s /bin/bash sandbox \ + && mkdir -p /sandbox /workspace /tmp \ + && chown -R sandbox:sandbox /sandbox /workspace + +WORKDIR /workspace +USER sandbox +ENTRYPOINT ["/bin/bash"] diff --git a/docs/source/architecture/agents/sandbox.md b/docs/source/architecture/agents/sandbox.md index f33988bf4..5b92b8c88 100644 --- a/docs/source/architecture/agents/sandbox.md +++ b/docs/source/architecture/agents/sandbox.md @@ -5,43 +5,52 @@ SPDX-License-Identifier: Apache-2.0 # Deep Research Sandbox Notes -Deep research can optionally run DeepAgents `execute` calls in a Modal sandbox. -In this release, sandboxes are scoped to a single async job: the Modal sandbox -name is the resolved job ID. This prevents unrelated jobs from sharing sandbox -filesystem state when each request receives a unique job ID. +Deep research can optionally run DeepAgents `execute` calls through a sandbox provider +(Modal, OpenShell, or any registered provider). Modal creates a sandbox per job. The +experimental OpenShell path attaches to a pre-created named sandbox shared by its jobs; +job-scoped directories prevent filename collisions but are not a security boundary. The sandbox is an internal execution detail. There are no sandbox-specific API endpoints, and job-level auth remains responsible for submit, stream, status, -cancel, state, and report access. +cancel, state, and report access. The one user-visible surface is the artifact +runtime (`.../job/{job_id}/artifacts`), which is also auth-scoped to the job. + +> **Developer reference:** the full architecture, provider contract, config schema, +> artifact pipeline, and troubleshooting live next to the code in +> [`src/aiq_agent/agents/deep_researcher/sandbox/README.md`](../../../../src/aiq_agent/agents/deep_researcher/sandbox/README.md). ## Current Behavior -- One sandbox name is used per deep research job when sandboxing is enabled. -- The sandbox name is the resolved job ID. -- Different jobs produce different sandbox names. +- Modal uses one sandbox per deep research job. OpenShell currently attaches jobs to + the configured shared sandbox name and is intended for local, single-operator testing. - Synchronous sandbox-enabled runs use an internal per-agent runtime ID. -- Job IDs must be valid Modal object names: 64 characters or fewer, using only - alphanumeric characters, dashes, periods, and underscores. -- Modal `timeout` and `idle_timeout` control sandbox lifetime. -- Files written inside the Modal workdir are temporary scratch state. -- Durable results should be returned by the agent or written through DeepAgents - virtual filesystem paths such as `/shared/`. +- Providers are selected by config (`sandbox.provider` + `providers.`); the + provider is validated against the registry and gated by its declared capabilities. + OpenShell policy is provisioned externally and is not verified when AI-Q attaches. +- Job IDs must satisfy each provider's object-name rules (Modal: 64 chars or fewer, + alphanumeric plus dash/period/underscore). +- `timeout` bounds individual execution. Other lifecycle controls are provider-dependent. +- Files written inside the workdir are temporary scratch state. Durable text should + be written through DeepAgents virtual paths such as `/shared/`; durable binaries + (charts, CSVs) are captured by the artifact runtime. ## Operational Notes -- High concurrency creates one Modal sandbox per concurrent sandbox-enabled job. -- If clients provide custom job IDs, they must not reuse a job ID for a new job. - Reuse can attach the job to an existing Modal sandbox until Modal terminates it. -- Cancelled or failed jobs may leave sandbox scratch files until Modal terminates - the sandbox according to timeout settings. -- If Modal removes a container mid-job, the job may fail and should be retried. +- High-concurrency Modal runs create one sandbox per job. OpenShell runs share the named + sandbox and must not be used concurrently for mutually untrusted jobs. Optional submit-path + caps (`AIQ_MAX_SANDBOXES_PER_PRINCIPAL` / `AIQ_MAX_SANDBOXES_GLOBAL`, default-off) bound + concurrency/cost but do not provide filesystem isolation. +- Custom client-supplied job IDs must not be reused for a new job. +- The runtime closes provider sessions on success, failure, cancellation, and timeout. + A named OpenShell sandbox persists when `delete_on_exit` is disabled. -## Deferred Hardening +## Current Safeguards -Planned follow-up work for production deployments: +The following safeguards are in place: -- Explicit sandbox cleanup on job success, failure, cancellation, and timeout. -- Retry-on-stale-container handling for Modal `NotFoundError`. -- Artifact capture rules for generated charts and binary outputs before cleanup. -- Sandbox quota and concurrency controls. -- Metrics and structured logs for sandbox create, reuse, failure, and cleanup. +- Explicit sandbox cleanup on success, failure, cancellation, and timeout. +- Idempotency-gated retry-on-stale-container handling. +- Artifact capture for generated charts/binaries (validate -> store -> serve/embed), + with MIME-from-bytes spoof rejection, SVG sanitization, and an inline-render allowlist. +- Sandbox quota and concurrency controls, and artifact retention via job-expiry cleanup. +- Structured lifecycle logging for sandbox create, reuse, failure, and cleanup. diff --git a/frontends/aiq_api/src/aiq_api/jobs/access.py b/frontends/aiq_api/src/aiq_api/jobs/access.py index 411236771..6842ae86a 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/access.py +++ b/frontends/aiq_api/src/aiq_api/jobs/access.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio +import logging import os from collections.abc import Mapping from typing import Any @@ -29,8 +30,14 @@ from aiq_agent.auth import Principal from aiq_agent.auth import get_current_principal +logger = logging.getLogger(__name__) + _job_access_schema_initialized: set[str] = set() +# Statuses that mean a job no longer holds a live sandbox. Anything else (running, +# pending, submitted, etc.) counts as active for the concurrency guard. +_TERMINAL_STATUS_SQL = "('success','failure','failed','interrupted','cancelled','completed','error')" + _JOB_ACCESS_INDEX_SQL = "CREATE INDEX IF NOT EXISTS idx_job_access_owner ON job_access(owner_auth_type, owner_subject)" _JOB_ACCESS_SELECT_SQL = text( "SELECT job_id, owner_auth_type, owner_subject, owner_email, created_at FROM job_access WHERE job_id = :job_id" @@ -44,6 +51,7 @@ def _is_postgres(db_url: str) -> bool: + """Return whether the database URL targets PostgreSQL.""" return db_url.startswith("postgres") @@ -110,6 +118,50 @@ def rollback_job_submission(job_id: str, db_url: str) -> None: conn.commit() +def count_active_jobs_for_owner(db_url: str, principal: Principal) -> int | None: + """Count an owner's non-terminal, non-expired jobs. + + Returns ``None`` if the count cannot be computed (e.g. the NAT ``job_info`` + schema differs); callers should fail open so a query mismatch never blocks + legitimate submissions. Used by the submit-path sandbox concurrency guard. + """ + try: + with _job_access_connection(db_url) as conn: + _ensure_job_access_schema(conn, db_url) + row = conn.execute( + text( + "SELECT COUNT(*) FROM job_access ja JOIN job_info ji ON ja.job_id = ji.job_id " + "WHERE ja.owner_auth_type = :t AND ja.owner_subject = :s " + "AND (ji.is_expired IS NOT TRUE) " + f"AND lower(ji.status) NOT IN {_TERMINAL_STATUS_SQL}" + ), + {"t": principal.type, "s": principal.sub}, + ).scalar() + return int(row or 0) + except Exception as exc: # noqa: BLE001 - guard must fail open, never block submits + logger.warning("Could not count active jobs for owner; allowing submit: %s", exc) + return None + + +def count_active_jobs_global(db_url: str) -> int | None: + """Count all non-terminal, non-expired jobs (global capacity guard). + + Returns ``None`` on query failure so callers fail open. + """ + try: + with _job_access_connection(db_url) as conn: + row = conn.execute( + text( + "SELECT COUNT(*) FROM job_info " + f"WHERE (is_expired IS NOT TRUE) AND lower(status) NOT IN {_TERMINAL_STATUS_SQL}" + ) + ).scalar() + return int(row or 0) + except Exception as exc: # noqa: BLE001 - guard must fail open, never block submits + logger.warning("Could not count active jobs globally; allowing submit: %s", exc) + return None + + def _make_no_auth_principal(owner: str | None = None) -> Principal: """Synthesize a principal for deployments with auth disabled (REQUIRE_AUTH=false). @@ -173,10 +225,12 @@ async def authorize_job_access(job_store: Any, db_url: str, job_id: str, princip def _principal_matches_access(principal: Principal, access: Mapping[str, Any]) -> bool: + """Return whether a principal matches a job-access row's owner identity.""" return principal.type == access.get("owner_auth_type") and principal.sub == access.get("owner_subject") def _job_access_connection(db_url: str): + """Open a sync connection on the shared event-store engine for the URL.""" from .event_store import EventStore engine = EventStore._get_or_create_sync_engine(db_url) @@ -184,6 +238,7 @@ def _job_access_connection(db_url: str): def _ensure_job_access_schema(conn: Connection, db_url: str) -> None: + """Create the ``job_access`` table and index once per database URL.""" if db_url in _job_access_schema_initialized: return conn.execute(text(_job_access_table_sql(db_url))) @@ -192,6 +247,7 @@ def _ensure_job_access_schema(conn: Connection, db_url: str) -> None: def _job_access_table_sql(db_url: str) -> str: + """Return the ``CREATE TABLE`` SQL for ``job_access``, dialect-aware for the URL.""" created_at_type = ( "TIMESTAMP WITH TIME ZONE DEFAULT NOW()" if _is_postgres(db_url) else "DATETIME DEFAULT CURRENT_TIMESTAMP" ) @@ -207,6 +263,7 @@ def _job_access_table_sql(db_url: str) -> str: def _job_access_upsert_sql(db_url: str): + """Return the dialect-appropriate upsert statement for ``job_access``.""" postgres_upsert = ( "INSERT INTO job_access (job_id, owner_auth_type, owner_subject, owner_email) " "VALUES (:job_id, :owner_auth_type, :owner_subject, :owner_email) " @@ -223,6 +280,7 @@ def _job_access_upsert_sql(db_url: str): def _principal_params(job_id: str, principal: Principal) -> dict[str, str | None]: + """Return the SQL bind params for a job's owner identity.""" return { "job_id": job_id, "owner_auth_type": principal.type, diff --git a/frontends/aiq_api/src/aiq_api/jobs/callbacks.py b/frontends/aiq_api/src/aiq_api/jobs/callbacks.py index ec2c7b91f..8c754ab4a 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/callbacks.py +++ b/frontends/aiq_api/src/aiq_api/jobs/callbacks.py @@ -136,6 +136,7 @@ class ToolArtifactMapping: """ def __init__(self): + """Initialize an empty tool-to-artifact mapping and register the defaults.""" self._mappings: dict[str, dict] = {} self._register_defaults() @@ -208,6 +209,9 @@ class AgentEventCallback(BaseCallbackHandler): URL_PATTERN = re.compile(r'https?://[^\s<>"\')\]}>]+', re.IGNORECASE) SEARCH_TOOL_PATTERNS = {"search", "tavily", "web_search", "google", "bing"} TOOL_CALL_PATTERN = re.compile(r'\b[a-z][a-z0-9_]*\s*\(\s*(?:["\'{]|[a-z_]+\s*=)', re.IGNORECASE) + SANDBOX_EXEC_TOOLS = frozenset({"execute"}) + SANDBOX_FILE_TOOLS = frozenset({"write_file", "read_file", "edit_file", "ls"}) + SHARED_FS_PREFIX = "/shared" AGENT_PATTERNS = {"agent"} AGENT_EXCLUDE_PATTERNS = {"middleware", "handler", "callback"} @@ -221,6 +225,12 @@ def __init__( event_store: EventStore | None = None, tool_artifact_mapping: ToolArtifactMapping | None = None, ): + """Wire the event store and tool/artifact mapping and init per-job URL caches. + + Args: + event_store: Sink for SSE events; None disables emission. + tool_artifact_mapping: Mapping of tools to artifact types; a default is used when omitted. + """ super().__init__() self._event_store = event_store self._tool_mapping = tool_artifact_mapping or ToolArtifactMapping() @@ -228,6 +238,7 @@ def __init__( self._run_id_to_name: dict[str, str] = {} self._run_id_to_parent: dict[str, str] = {} self._agent_run_ids: dict[str, str] = {} # {run_id: name} + self._run_id_to_sandbox_tool: dict[str, bool] = {} self._job_id = event_store.job_id if event_store else None self._instance_discovered_urls: set[str] = set() @@ -299,6 +310,7 @@ def _build_metadata_for_run(self, run_id: str, **extra: Any) -> dict[str, Any] | return metadata if metadata else None def _emit(self, event: IntermediateStepEvent): + """Store an event as an SSE dict when an event store is configured.""" if self._event_store: self._event_store.store(event.to_sse_dict()) @@ -341,12 +353,23 @@ def _emit_artifact( ) def _get_chain_name(self, serialized: dict | None, **kwargs) -> str: + """Resolve a human-readable chain name from serialized data or kwargs.""" if serialized: name = serialized.get("name") or serialized.get("id", [""])[-1] if name: return name return kwargs.get("name", "unknown") + def _is_sandbox_tool(self, tool_name: str, parsed_input: Any) -> bool: + """Return whether a tool call runs against the provisioned sandbox.""" + if tool_name in self.SANDBOX_EXEC_TOOLS: + return True + if tool_name not in self.SANDBOX_FILE_TOOLS or not isinstance(parsed_input, dict): + return False + path = str(parsed_input.get("file_path") or parsed_input.get("path") or parsed_input.get("filename") or "") + in_shared = path == self.SHARED_FS_PREFIX or path.startswith(self.SHARED_FS_PREFIX + "/") + return bool(path and not in_shared) + def _get_source_registry(self): """Return the session-scoped SourceRegistry if set, otherwise None.""" return get_session_registry() @@ -521,6 +544,7 @@ def _emit_tool_artifact(self, tool_name: str, tool_input: Any, run_id: str = "") self._emit_artifact(artifact_type, content, name=name, **extra_data) def on_chain_start(self, serialized: dict | None, inputs: dict, **kwargs) -> None: + """Track the run/parent lineage and emit an agent.start event for agent-like chains.""" name = self._get_chain_name(serialized, **kwargs) run_id = str(kwargs.get("run_id", "")) parent_run_id = str(kwargs.get("parent_run_id", "")) if kwargs.get("parent_run_id") else "" @@ -544,6 +568,7 @@ def on_chain_start(self, serialized: dict | None, inputs: dict, **kwargs) -> Non ) def on_chain_end(self, outputs: dict, **kwargs) -> None: + """Emit an agent.end event for agent-like chains and clear run bookkeeping.""" run_id = str(kwargs.get("run_id", "")) name = self._run_id_to_name.pop(run_id, kwargs.get("name", "")) @@ -592,6 +617,7 @@ def _trim_tool_input(self, parsed_input: Any) -> Any: return serialized[: self.TOOL_INPUT_TRIM_LIMIT] + "..." def on_tool_start(self, serialized: dict | None, input_str: str, **kwargs) -> None: + """Emit a tool.start event and record sandbox/lineage state for the tool run.""" tool_name = serialized.get("name", "unknown") if serialized else "unknown" run_id = str(kwargs.get("run_id", "")) parent_run_id = str(kwargs.get("parent_run_id", "")) if kwargs.get("parent_run_id") else "" @@ -604,6 +630,9 @@ def on_tool_start(self, serialized: dict | None, input_str: str, **kwargs) -> No parsed_input = self._parse_tool_input(input_str) emit_input = self._trim_tool_input(parsed_input) + is_sandbox_tool = self._is_sandbox_tool(tool_name, parsed_input) + if run_id: + self._run_id_to_sandbox_tool[run_id] = is_sandbox_tool self._emit( IntermediateStepEvent( @@ -611,15 +640,17 @@ def on_tool_start(self, serialized: dict | None, input_str: str, **kwargs) -> No state=EventState.START, name=tool_name, data=EventData(input=emit_input) if emit_input else None, - metadata=self._build_metadata_for_run(run_id), + metadata=self._build_metadata_for_run(run_id, sandbox=True if is_sandbox_tool else None), ) ) self._emit_tool_artifact(tool_name, parsed_input, run_id=run_id) def on_tool_end(self, output: str, **kwargs) -> None: + """Emit a tool.end event and clear the tool run's lineage/sandbox state.""" run_id = str(kwargs.get("run_id", "")) tool_name = self._run_id_to_name.pop(run_id, kwargs.get("name", "unknown")) + is_sandbox_tool = self._run_id_to_sandbox_tool.pop(run_id, False) agent_info = self._find_agent_for_run(run_id) @@ -629,7 +660,7 @@ def on_tool_end(self, output: str, **kwargs) -> None: state=EventState.END, name=tool_name, data=None, - metadata=self._build_metadata_for_run(run_id), + metadata=self._build_metadata_for_run(run_id, sandbox=True if is_sandbox_tool else None), ) ) @@ -652,6 +683,7 @@ def on_tool_end(self, output: str, **kwargs) -> None: self._run_id_to_parent.pop(run_id, None) def on_llm_start(self, serialized: dict, prompts: list, **kwargs) -> None: + """Emit an llm.start event for a completion-style model call.""" model_name = "unknown" if serialized: model_name = serialized.get("name") or serialized.get("id", ["unknown"])[-1] @@ -677,6 +709,7 @@ def on_llm_start(self, serialized: dict, prompts: list, **kwargs) -> None: ) def on_llm_new_token(self, token: str, **kwargs) -> None: + """Emit a streaming llm.chunk event for each non-empty token.""" if token: self._emit( IntermediateStepEvent( @@ -690,6 +723,7 @@ def on_llm_new_token(self, token: str, **kwargs) -> None: THINKING_TRIM_SUFFIX = " [Trimmed - check traces for full logs]" def on_llm_end(self, response, **kwargs) -> None: + """Emit an llm.end event with content, thinking, and token usage.""" run_id = str(kwargs.get("run_id", "")) model_name = self._run_id_to_name.pop(run_id, "unknown") @@ -734,6 +768,7 @@ def on_llm_end(self, response, **kwargs) -> None: self._run_id_to_parent.pop(run_id, None) def on_chat_model_start(self, serialized: dict, messages: list, **kwargs) -> None: + """Emit an llm.start event for a chat-model call.""" model_name = "unknown" if serialized: model_name = serialized.get("name") or serialized.get("kwargs", {}).get("model", "unknown") diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index 9c068e0de..717b3fd5b 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -72,6 +72,14 @@ def __init__( job_id: str, poll_interval: float = 1.0, ): + """Configure the job-status poller used to detect interruption. + + Args: + scheduler_address: Dask scheduler address (unused by polling, kept for context). + db_url: Database URL of the job store to poll. + job_id: Job whose status is monitored. + poll_interval: Seconds between status polls. + """ self.scheduler_address = scheduler_address self.db_url = db_url self.job_id = job_id @@ -81,6 +89,7 @@ def __init__( @property def is_cancelled(self) -> bool: + """Return whether the monitored job has been interrupted.""" return self._cancelled.is_set() async def _poll_job_status(self) -> None: @@ -315,6 +324,9 @@ async def run_agent_job( job_store: JobStore | None = None cancellation_monitor: CancellationMonitor | None = None event_store: EventStore | BatchingEventStore | None = None + # Sandbox runtime is released on the terminal path; interrupted forces terminate() over close(). + sandbox_runtime: Any | None = None + interrupted = False logger.info( "Dask worker received: agent=%s, config=%s, job_id=%s", agent_class_path, @@ -481,8 +493,17 @@ async def run_agent_job( verbose=verbose, callbacks=callbacks, job_id=job_id, + # Artifact harvesting rides 284's job store + event stream: the same db_url + # backs the SqlArtifactStore, and event_store.store carries artifact SSE + # events. Inert unless sandbox.artifact_capture is enabled in config. + artifact_db_url=db_url, + artifact_emit=event_store.store, ) + # Capture the runtime so the terminal path can release the sandbox. None for + # agents without a sandbox runtime; close()/terminate() are then no-ops. + sandbox_runtime = getattr(agent, "deepagents_runtime", None) + # Run agent - LLM/tool events will be nested under workflow span result = await _run_agent( agent=agent, @@ -529,6 +550,7 @@ async def run_agent_job( except asyncio.CancelledError: logger.info("Job %s cancelled", job_id) + interrupted = True if job_store: try: job = await job_store.get_job(job_id) @@ -575,6 +597,10 @@ async def run_agent_job( event_store.flush() if cancellation_monitor: cancellation_monitor.stop() + # Release the sandbox off the event loop so the SDK session close never blocks the Dask + # worker. The single artifact harvest already ran in agent.run() before this point, so + # teardown only closes/terminates; interrupted jobs terminate() to preempt a live execute. + await asyncio.to_thread(_teardown_sandbox, sandbox_runtime, job_id=job_id, interrupted=interrupted) # Clean up job-scoped auth token if _auth_token_reset is not None: from ._auth_context import job_auth_token @@ -582,6 +608,26 @@ async def run_agent_job( job_auth_token.reset(_auth_token_reset) +def _teardown_sandbox(sandbox_runtime: Any | None, *, job_id: str, interrupted: bool) -> None: + """Release sandbox resources on a terminal path (best-effort, never raises). + + Interrupted jobs (cancel/timeout) call ``terminate()`` so a still-running ``execute`` is + forcibly preempted; normal paths call ``close()`` gracefully. Both are idempotent. This runs + off the event loop (``asyncio.to_thread``) so the SDK session close cannot block the worker. + """ + if sandbox_runtime is None: + return + teardown = getattr(sandbox_runtime, "terminate", None) if interrupted else None + if teardown is None: + teardown = getattr(sandbox_runtime, "close", None) + if teardown is None: + return + try: + teardown() + except Exception: # noqa: BLE001 - cleanup must never raise on the terminal path + logger.warning("Sandbox cleanup failed for job %s", job_id, exc_info=True) + + def _create_agent_instance( agent_cls: type, llm_provider, @@ -591,6 +637,8 @@ def _create_agent_instance( verbose: bool, callbacks: list, job_id: str | None = None, + artifact_db_url: str | None = None, + artifact_emit=None, ): """ Create an agent instance, supporting different constructor patterns. @@ -613,6 +661,8 @@ def _create_agent_instance( skills=fn_config.skills, sandbox=fn_config.sandbox, job_id=job_id, + artifact_db_url=artifact_db_url, + artifact_emit=artifact_emit, max_research_concurrency=fn_config.max_research_concurrency, max_concurrent_source_tool_calls=fn_config.max_concurrent_source_tool_calls, max_source_tool_batch_size=fn_config.max_source_tool_batch_size, diff --git a/frontends/aiq_api/src/aiq_api/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index b26f9dedc..61003b124 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/frontends/aiq_api/src/aiq_api/routes/jobs.py @@ -32,9 +32,11 @@ import asyncio import json import logging +import os import time from typing import TYPE_CHECKING from typing import Annotated +from typing import Any from fastapi import Body from fastapi import FastAPI @@ -60,6 +62,70 @@ logger = logging.getLogger(__name__) +def _int_env(name: str, default: int) -> int: + """Read a non-negative integer ops knob from the environment. + + A missing, non-integer, or negative value falls back to ``default`` so a + misconfigured cap can never silently invert into "block all submissions". + """ + try: + value = int(os.environ[name]) + except (KeyError, ValueError): + return default + if value < 0: + logger.warning("%s=%d is negative; using default %d", name, value, default) + return default + return value + + +def _sandbox_caps_configured() -> bool: + """Whether an operator has opted into sandbox concurrency caps via env. + + Default-off so the guard never adds a function-config lookup (or behavior change) + to submits unless caps are explicitly configured. + """ + return "AIQ_MAX_SANDBOXES_PER_PRINCIPAL" in os.environ or "AIQ_MAX_SANDBOXES_GLOBAL" in os.environ + + +def _agent_uses_sandbox(builder: Any, config_name: str) -> bool: + """Return whether the agent's function config enables a sandbox.""" + try: + fn_config = builder.get_function_config(config_name) + except Exception: # noqa: BLE001 - missing/odd config means "no sandbox guard" + return False + sandbox = getattr(fn_config, "sandbox", None) + if sandbox is None: + return False + return bool(getattr(sandbox, "enabled", True)) + + +async def _enforce_sandbox_concurrency(db_url: str, principal: Any) -> None: + """Reject submission when per-principal or global sandbox limits are reached. + + Option A: enforced at the API submit path so cost is stopped before a Dask worker + spins up a sandbox. Counts fail open (None) so a query mismatch never blocks submits. + Configurable via AIQ_MAX_SANDBOXES_PER_PRINCIPAL / AIQ_MAX_SANDBOXES_GLOBAL. + """ + from ..jobs.access import count_active_jobs_for_owner + from ..jobs.access import count_active_jobs_global + + per_principal = _int_env("AIQ_MAX_SANDBOXES_PER_PRINCIPAL", 5) + global_cap = _int_env("AIQ_MAX_SANDBOXES_GLOBAL", 50) + loop = asyncio.get_running_loop() + + owner_count = await loop.run_in_executor(None, count_active_jobs_for_owner, db_url, principal) + if owner_count is not None and owner_count >= per_principal: + raise HTTPException( + 429, + f"Active job limit reached for this principal ({per_principal}). " + "Wait for running jobs to finish before submitting more.", + ) + + global_count = await loop.run_in_executor(None, count_active_jobs_global, db_url) + if global_count is not None and global_count >= global_cap: + raise HTTPException(503, "Server is at sandbox capacity; please retry shortly.") + + class JobSubmitRequest(BaseModel): """Request to submit an async job.""" @@ -458,6 +524,13 @@ async def submit_job( len(req.data_sources) if req.data_sources is not None else "none", ) + # Sandbox concurrency / cost guard (Option A): cap concurrent sandbox-enabled + # jobs per principal and globally, enforced at submit so cost is stopped before + # a worker spins up. Opt-in (default-off) via AIQ_MAX_SANDBOXES_* env vars so the + # default submit path stays lazy; fail-open if the active-job count is unknown. + if _sandbox_caps_configured() and _agent_uses_sandbox(builder, agent_config.config_name): + await _enforce_sandbox_concurrency(db_url, principal) + # Propagate auth token to Dask worker for requires_auth data sources from aiq_agent.auth import get_auth_token @@ -607,6 +680,75 @@ async def get_job_state(job_id: str) -> JobStateResponse: artifacts=artifacts, ) + @app.get( + "/v1/jobs/async/job/{job_id}/artifacts", + tags=["async jobs"], + summary="List durable artifacts", + description="List generated artifacts (charts, CSVs, notebooks) harvested from the sandbox.", + responses={404: {"description": "Job not found"}}, + ) + async def list_job_artifacts(job_id: str) -> dict: + """List durable artifact metadata for a job (no bytes).""" + from aiq_agent.agents.deep_researcher.sandbox.artifacts import SqlArtifactStore + + principal = require_verified_principal() + await authorize_job_access(job_store, db_url, job_id, principal) + + store = SqlArtifactStore(db_url) + artifacts = await asyncio.to_thread(store.list, job_id) + # Exclude storage internals (storage_uri embeds the db_url, which may carry + # credentials/hostnames; sandbox_path is an internal layout detail) from the + # client-facing payload. Clients use the content endpoint, not these fields. + return { + "job_id": job_id, + "artifacts": [a.model_dump(mode="json", exclude={"storage_uri", "sandbox_path"}) for a in artifacts], + } + + @app.get( + "/v1/jobs/async/job/{job_id}/artifacts/{artifact_id}/content", + tags=["async jobs"], + summary="Download artifact content", + description="Stream the bytes of a single artifact. Job-ownership checks apply.", + responses={404: {"description": "Job or artifact not found"}}, + ) + async def get_job_artifact_content(job_id: str, artifact_id: str) -> StreamingResponse: + """Stream an artifact's bytes (auth-scoped to the owning job).""" + from aiq_agent.agents.deep_researcher.sandbox.artifacts import SqlArtifactStore + + principal = require_verified_principal() + await authorize_job_access(job_store, db_url, job_id, principal) + + store = SqlArtifactStore(db_url) + artifact = await asyncio.to_thread(store.get, job_id, artifact_id) + if artifact is None: + raise HTTPException(404, f"Artifact not found: {artifact_id}") + + # The filename is sandbox-controlled; strip control chars and quotes so it cannot + # break out of the header value (response-splitting / header injection). + safe_filename = "".join(c for c in artifact.filename if c.isprintable() and c not in '"\\') or "artifact" + # Starlette encodes header values as Latin-1, so a non-Latin-1 filename (emoji, CJK) + # would raise UnicodeEncodeError. Provide an ASCII-only fallback plus an RFC 5987 + # filename* with the UTF-8 percent-encoded original for clients that support it. + from urllib.parse import quote + + ascii_filename = safe_filename.encode("ascii", "ignore").decode() or "artifact" + encoded_filename = quote(safe_filename, safe="") + # Only magic-verified raster images may render inline; everything else (SVG, HTML, + # notebooks, PDFs) is forced to download with nosniff to prevent stored-XSS if a + # user opens the content URL directly in a browser. + inline_safe = artifact.mime_type in {"image/png", "image/jpeg", "image/webp"} + disposition = "inline" if inline_safe else "attachment" + return StreamingResponse( + store.open_bytes(job_id, artifact_id), + media_type=artifact.mime_type, + headers={ + "Content-Disposition": ( + f"{disposition}; filename=\"{ascii_filename}\"; filename*=UTF-8''{encoded_filename}" + ), + "X-Content-Type-Options": "nosniff", + }, + ) + @app.get( "/v1/jobs/async/job/{job_id}/report", response_model=JobReportResponse, @@ -868,6 +1010,7 @@ async def _run_event_cleanup(db_url: str, retention_seconds: int, is_postgres: b loop = asyncio.get_running_loop() def _do_cleanup() -> tuple[int, int, int]: + """Delete expired events/jobs/access rows synchronously; return removal counts.""" from sqlalchemy import text engine = EventStore._get_or_create_sync_engine(db_url) @@ -911,6 +1054,19 @@ def _do_cleanup() -> tuple[int, int, int]: time_deleted, expired_deleted, access_deleted = await loop.run_in_executor(None, _do_cleanup) + # Artifact retention shares the job expiry boundary (best-effort; the artifacts + # table only exists when artifact capture has been used). + try: + from aiq_agent.agents.deep_researcher.sandbox.artifacts import SqlArtifactStore + + artifacts_deleted = await loop.run_in_executor( + None, lambda: SqlArtifactStore(db_url).cleanup_old_artifacts(retention_seconds) + ) + if artifacts_deleted: + logger.info("Artifact cleanup: %d old artifacts removed", artifacts_deleted) + except Exception as e: # noqa: BLE001 - retention is best-effort + logger.debug("Artifact cleanup skipped: %s", e) + if time_deleted > 0 or expired_deleted > 0 or access_deleted > 0: logger.info( "Event cleanup: %d old events removed, %d events for expired jobs removed, %d access rows removed", @@ -971,6 +1127,7 @@ def _process_tool_start(event: dict, data: dict, metadata: dict, tool_call_map: "output": None, "status": "running", "workflow": metadata.get("workflow"), + "is_sandbox": bool(metadata.get("sandbox")), "timestamp": event.get("timestamp"), } @@ -984,6 +1141,7 @@ def _process_tool_end(event: dict, data: dict, metadata: dict, tool_call_map: di if tool_id in tool_call_map: tool_call_map[tool_id]["output"] = tool_output tool_call_map[tool_id]["status"] = "completed" + tool_call_map[tool_id]["is_sandbox"] = tool_call_map[tool_id].get("is_sandbox") or bool(metadata.get("sandbox")) else: tool_call_map[tool_id] = { "id": tool_id, @@ -992,6 +1150,7 @@ def _process_tool_end(event: dict, data: dict, metadata: dict, tool_call_map: di "output": tool_output, "status": "completed", "workflow": metadata.get("workflow"), + "is_sandbox": bool(metadata.get("sandbox")), "timestamp": event.get("timestamp"), } @@ -1162,6 +1321,7 @@ async def _sse_generator_postgres(job_store, job_id: str, db_url: str, start_eve is_reconnect = start_event_id > 0 def format_sse(event_type: str, data: dict, event_id: int | None = None) -> str: + """Format an SSE frame and advance (or set) the monotonic event sequence id.""" nonlocal sequence_id if event_id is not None: sequence_id = event_id @@ -1183,6 +1343,7 @@ def format_sse(event_type: str, data: dict, event_id: int | None = None) -> str: notification_queue: asyncio.Queue = asyncio.Queue() def notification_handler(connection, pid, channel_name, payload): + """Enqueue a Postgres LISTEN/NOTIFY payload, dropping it if the queue is full.""" try: notification_queue.put_nowait(payload) except asyncio.QueueFull: @@ -1348,6 +1509,7 @@ async def _sse_generator_polling(job_store, job_id: str, db_url: str, start_even replay_mode_announced = False def format_sse(event_type: str, data: dict, event_id: int | None = None) -> str: + """Format an SSE frame and advance (or set) the monotonic event sequence id.""" nonlocal sequence_id if event_id is not None: sequence_id = event_id diff --git a/frontends/aiq_api/tests/test_sandbox_concurrency.py b/frontends/aiq_api/tests/test_sandbox_concurrency.py new file mode 100644 index 000000000..5af473990 --- /dev/null +++ b/frontends/aiq_api/tests/test_sandbox_concurrency.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the submit-path sandbox concurrency guard (Option A).""" + +from __future__ import annotations + +import asyncio +import os +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from fastapi import HTTPException + +from aiq_api.routes import jobs as jobs_module + +_PRINCIPAL = SimpleNamespace(type="user", sub="subject-1") + +# Pin the caps so the assertions are independent of any AIQ_MAX_SANDBOXES_* set in CI/dev. +_PINNED_CAPS = {"AIQ_MAX_SANDBOXES_PER_PRINCIPAL": "5", "AIQ_MAX_SANDBOXES_GLOBAL": "50"} + + +def _run(coro): + return asyncio.run(coro) + + +class TestAgentUsesSandbox: + def test_true_when_enabled(self) -> None: + builder = SimpleNamespace(get_function_config=lambda _n: SimpleNamespace(sandbox=SimpleNamespace(enabled=True))) + assert jobs_module._agent_uses_sandbox(builder, "cfg") is True + + def test_false_when_disabled(self) -> None: + cfg = SimpleNamespace(sandbox=SimpleNamespace(enabled=False)) + builder = SimpleNamespace(get_function_config=lambda _n: cfg) + assert jobs_module._agent_uses_sandbox(builder, "cfg") is False + + def test_false_when_no_sandbox(self) -> None: + builder = SimpleNamespace(get_function_config=lambda _n: SimpleNamespace(sandbox=None)) + assert jobs_module._agent_uses_sandbox(builder, "cfg") is False + + def test_false_when_config_lookup_errors(self) -> None: + def _boom(_name: str): + raise RuntimeError("no config") + + builder = SimpleNamespace(get_function_config=_boom) + assert jobs_module._agent_uses_sandbox(builder, "cfg") is False + + +class TestEnforceSandboxConcurrency: + def test_rejects_when_owner_over_limit(self) -> None: + with ( + patch.dict(os.environ, _PINNED_CAPS, clear=False), + patch("aiq_api.jobs.access.count_active_jobs_for_owner", return_value=5), + patch("aiq_api.jobs.access.count_active_jobs_global", return_value=0), + pytest.raises(HTTPException) as exc, + ): + _run(jobs_module._enforce_sandbox_concurrency("sqlite:///x", _PRINCIPAL)) + assert exc.value.status_code == 429 + + def test_rejects_when_global_over_limit(self) -> None: + with ( + patch.dict(os.environ, _PINNED_CAPS, clear=False), + patch("aiq_api.jobs.access.count_active_jobs_for_owner", return_value=0), + patch("aiq_api.jobs.access.count_active_jobs_global", return_value=50), + pytest.raises(HTTPException) as exc, + ): + _run(jobs_module._enforce_sandbox_concurrency("sqlite:///x", _PRINCIPAL)) + assert exc.value.status_code == 503 + + def test_allows_under_limit(self) -> None: + with ( + patch.dict(os.environ, _PINNED_CAPS, clear=False), + patch("aiq_api.jobs.access.count_active_jobs_for_owner", return_value=1), + patch("aiq_api.jobs.access.count_active_jobs_global", return_value=1), + ): + _run(jobs_module._enforce_sandbox_concurrency("sqlite:///x", _PRINCIPAL)) + + def test_fails_open_when_counts_unknown(self) -> None: + with ( + patch.dict(os.environ, _PINNED_CAPS, clear=False), + patch("aiq_api.jobs.access.count_active_jobs_for_owner", return_value=None), + patch("aiq_api.jobs.access.count_active_jobs_global", return_value=None), + ): + _run(jobs_module._enforce_sandbox_concurrency("sqlite:///x", _PRINCIPAL)) diff --git a/frontends/ui/src/adapters/api/deep-research-client.ts b/frontends/ui/src/adapters/api/deep-research-client.ts index 75c18fb54..5377e3d7c 100644 --- a/frontends/ui/src/adapters/api/deep-research-client.ts +++ b/frontends/ui/src/adapters/api/deep-research-client.ts @@ -203,8 +203,8 @@ export interface DeepResearchCallbacks { usage?: { input_tokens: number; output_tokens: number } ) => void /** Called on tool events */ - onToolStart?: (name: string, input?: Record, workflow?: string, eventId?: string, agentId?: string) => void - onToolEnd?: (name: string, output?: string, eventId?: string, agentId?: string) => void + onToolStart?: (name: string, input?: Record, workflow?: string, eventId?: string, agentId?: string, isSandbox?: boolean) => void + onToolEnd?: (name: string, output?: string, eventId?: string, agentId?: string, isSandbox?: boolean) => void /** Called on artifact updates */ onTodoUpdate?: (todos: TodoItem[], workflow?: string) => void onCitationUpdate?: (url: string, content: string, isCited?: boolean) => void @@ -463,22 +463,22 @@ export const createDeepResearchClient = (options: DeepResearchStreamOptions): De } case 'tool.start': { - // tool events have nested structure: { id, name, timestamp, data: { input }, metadata: { workflow, agent_id } } + // tool events have nested structure: { id, name, timestamp, data: { input }, metadata: { workflow, agent_id, sandbox } } // Note: data.input may be a Python repr string when the backend trims large inputs via str() const toolData = rawData as { id?: string name: string data?: { input?: unknown } - metadata?: { workflow?: string; agent_id?: string } + metadata?: { workflow?: string; agent_id?: string; sandbox?: boolean } } const normalizedInput = normalizeToolInput(toolData.data?.input) - callbacks.onToolStart?.(toolData.name, normalizedInput, toolData.metadata?.workflow, toolData.id, toolData.metadata?.agent_id) + callbacks.onToolStart?.(toolData.name, normalizedInput, toolData.metadata?.workflow, toolData.id, toolData.metadata?.agent_id, Boolean(toolData.metadata?.sandbox)) break } case 'tool.end': { - const toolData = rawData as { id?: string; name: string; data?: { output?: string }; metadata?: { agent_id?: string } } - callbacks.onToolEnd?.(toolData.name, toolData.data?.output, toolData.id, toolData.metadata?.agent_id) + const toolData = rawData as { id?: string; name: string; data?: { output?: string }; metadata?: { agent_id?: string; sandbox?: boolean } } + callbacks.onToolEnd?.(toolData.name, toolData.data?.output, toolData.id, toolData.metadata?.agent_id, Boolean(toolData.metadata?.sandbox)) break } diff --git a/frontends/ui/src/app/api/jobs/async/[...path]/route.ts b/frontends/ui/src/app/api/jobs/async/[...path]/route.ts index 5d4f450fe..50c4eecec 100644 --- a/frontends/ui/src/app/api/jobs/async/[...path]/route.ts +++ b/frontends/ui/src/app/api/jobs/async/[...path]/route.ts @@ -89,6 +89,8 @@ export async function GET( const { path } = await params const backendUrl = buildBackendUrl(path) const isStreamRequest = path.includes('stream') + // Artifact bytes (.../artifacts/{id}/content) are binary — never JSON-parse them. + const isArtifactContent = path.includes('artifacts') && path[path.length - 1] === 'content' console.log('[Deep Research API] GET:', backendUrl, isStreamRequest ? '(SSE)' : '') @@ -97,11 +99,16 @@ export async function GET( console.log('[Deep Research API] idToken cookie present:', !!authHeaders.Cookie) // Forward the request to the backend + const acceptHeader = isStreamRequest + ? 'text/event-stream' + : isArtifactContent + ? '*/*' + : 'application/json' const response = await fetch(backendUrl, { method: 'GET', headers: { ...authHeaders, - Accept: isStreamRequest ? 'text/event-stream' : 'application/json', + Accept: acceptHeader, }, ...(isStreamRequest ? { signal: req.signal } : {}), }) @@ -154,6 +161,26 @@ export async function GET( }) } + // For artifact content, stream the raw bytes through with the upstream content type + // (JSON-parsing here would corrupt binary payloads like PNGs). + if (isArtifactContent) { + if (!response.body) { + return new NextResponse( + JSON.stringify({ + error: { code: 'NO_RESPONSE_BODY', message: 'Backend returned no artifact content' }, + }), + { status: 502, headers: { 'Content-Type': 'application/json' } } + ) + } + const passthroughHeaders: Record = { + 'Content-Type': response.headers.get('Content-Type') ?? 'application/octet-stream', + 'Cache-Control': 'private, max-age=3600', + } + const disposition = response.headers.get('Content-Disposition') + if (disposition) passthroughHeaders['Content-Disposition'] = disposition + return new NextResponse(response.body, { status: response.status, headers: passthroughHeaders }) + } + // For regular JSON responses const data = await response.json() return NextResponse.json(data) diff --git a/frontends/ui/src/features/chat/hooks/use-deep-research.ts b/frontends/ui/src/features/chat/hooks/use-deep-research.ts index 124434a14..4cbcae37c 100644 --- a/frontends/ui/src/features/chat/hooks/use-deep-research.ts +++ b/frontends/ui/src/features/chat/hooks/use-deep-research.ts @@ -209,7 +209,7 @@ export const useDeepResearch = (): UseDeepResearchReturn => { activeToolStacks: new Map(), agents: new Map(), llmSteps: new Map(), - toolCalls: new Map; output?: string; workflow?: string; agentId?: string }>(), + toolCalls: new Map; output?: string; workflow?: string; agentId?: string; isSandbox?: boolean }>(), todos: null as TodoItem[] | null, citations: [] as Array<{ url: string; content: string; isCited: boolean }>, files: new Map(), @@ -236,7 +236,7 @@ export const useDeepResearch = (): UseDeepResearchReturn => { const now = new Date() const agents = Array.from(buf.agents.entries()).map(([id, a]) => ({ id, name: a.name, input: a.input, output: a.output, status: 'complete' as const, startedAt: now, completedAt: now })) const llmSteps = Array.from(buf.llmSteps.entries()).map(([id, s]) => ({ id, name: s.name, workflow: s.workflow, content: s.content, thinking: s.thinking, usage: s.usage, isComplete: true, timestamp: now })) - const toolCalls = Array.from(buf.toolCalls.entries()).map(([id, t]) => ({ id, name: t.name, input: t.input, output: t.output, workflow: t.workflow, agentId: t.agentId, status: 'complete' as const, timestamp: now })) + const toolCalls = Array.from(buf.toolCalls.entries()).map(([id, t]) => ({ id, name: t.name, input: t.input, output: t.output, workflow: t.workflow, agentId: t.agentId, isSandbox: t.isSandbox, status: 'complete' as const, timestamp: now })) const citations = buf.citations.map((c, i) => ({ id: `citation-${i}`, url: c.url, content: c.content, isCited: c.isCited, timestamp: now })) const files = Array.from(buf.files.entries()).map(([filename, content], i) => ({ id: `file-${i}`, filename, content, timestamp: now })) const todos = buf.todos ? normalizeDeepResearchTodos(buf.todos) : undefined @@ -439,10 +439,10 @@ export const useDeepResearch = (): UseDeepResearchReturn => { if (llmStepKeys.length > 0) { const [key, llmStepId] = llmStepKeys[llmStepKeys.length - 1]; completeDeepResearchLLMStep(llmStepId, thinking, usage); activeStepIdsRef.current.delete(key) } }, - onToolStart: (name, input, workflow, _eventId, agentId) => { + onToolStart: (name, input, workflow, _eventId, agentId, isSandbox) => { if (name === 'task') return if (buf.active) { - const id = `tool-${buf.idCounter++}`; buf.toolCalls.set(id, { name, input, workflow, agentId }) + const id = `tool-${buf.idCounter++}`; buf.toolCalls.set(id, { name, input, workflow, agentId, isSandbox }) let stack = buf.activeToolStacks.get(name); if (!stack) { stack = []; buf.activeToolStacks.set(name, stack) }; stack.push(id); return } if (!isActiveJob()) return @@ -453,7 +453,7 @@ export const useDeepResearch = (): UseDeepResearchReturn => { const stepId = addThinkingStep({ category: 'tools', functionName: name, displayName: name, content: inputText ? `Input: ${inputText}\n` : 'Executing...\n', isComplete: false, isDeepResearch: true }) activeStepIdsRef.current.set(`tool:${name}`, stepId) } - const toolCallId = addDeepResearchToolCall({ name, input, workflow, agentId }) + const toolCallId = addDeepResearchToolCall({ name, input, workflow, agentId, isSandbox }) activeStepIdsRef.current.set(`toolCall:${name}`, toolCallId) }, diff --git a/frontends/ui/src/features/chat/hooks/use-load-job-data.ts b/frontends/ui/src/features/chat/hooks/use-load-job-data.ts index 87195561a..b39b76543 100644 --- a/frontends/ui/src/features/chat/hooks/use-load-job-data.ts +++ b/frontends/ui/src/features/chat/hooks/use-load-job-data.ts @@ -281,11 +281,12 @@ export const useLoadJobData = (): UseLoadJobDataReturn => { const { tools, outputs } = stateResponse.artifacts tools?.forEach( - (tool: { name: string; input?: Record; output?: string }) => { + (tool: { name: string; input?: Record; output?: string; is_sandbox?: boolean }) => { const toolCallId = addDeepResearchToolCall({ name: tool.name, input: tool.input, workflow: undefined, + isSandbox: tool.is_sandbox, }) if (tool.output) { completeDeepResearchToolCall(toolCallId, tool.output) @@ -365,6 +366,7 @@ export const useLoadJobData = (): UseLoadJobDataReturn => { output?: string workflow?: string agentId?: string + isSandbox?: boolean } >(), todos: null as TodoItem[] | null, @@ -412,6 +414,7 @@ export const useLoadJobData = (): UseLoadJobDataReturn => { output: t.output, workflow: t.workflow, agentId: t.agentId, + isSandbox: t.isSandbox, status: 'complete' as const, timestamp: now, })) @@ -538,10 +541,10 @@ export const useLoadJobData = (): UseLoadJobDataReturn => { } }, - onToolStart: (name, input, workflow, _eventId, agentId) => { + onToolStart: (name, input, workflow, _eventId, agentId, isSandbox) => { if (name === 'task') return const uniqueId = `tool-${idCounter++}` - buffer.toolCalls.set(uniqueId, { name, input, workflow, agentId }) + buffer.toolCalls.set(uniqueId, { name, input, workflow, agentId, isSandbox }) let stack = activeToolStacks.get(name) if (!stack) { stack = [] diff --git a/frontends/ui/src/features/chat/store.ts b/frontends/ui/src/features/chat/store.ts index 8db300d37..2e8ef16c4 100644 --- a/frontends/ui/src/features/chat/store.ts +++ b/frontends/ui/src/features/chat/store.ts @@ -3198,6 +3198,19 @@ export const selectHasConnectionError = (state: ChatStore): boolean => (m) => m.messageType === 'error' && m.errorData?.errorCode?.startsWith('connection.') ) ?? false +/** + * Resolve the deep-research job id for artifact resolution (report images, PDF/markdown + * export). Prefers the active streaming job, then falls back to the latest deep-research + * message in the conversation — the active id is cleared once a job stops streaming, but a + * finished report still needs the id to fetch its artifact content. + */ +export const selectResolvedDeepResearchJobId = (state: ChatStore): string | undefined => { + if (state.deepResearchJobId) return state.deepResearchJobId + const conversation = state.currentConversation + if (!conversation) return undefined + return getLatestDeepResearchMessage(conversation)?.deepResearchJobId ?? undefined +} + // ============================================================ // Storage Event Monitoring (for debugging session clearing) // ============================================================ diff --git a/frontends/ui/src/features/chat/types.ts b/frontends/ui/src/features/chat/types.ts index 6f48192aa..a62785299 100644 --- a/frontends/ui/src/features/chat/types.ts +++ b/frontends/ui/src/features/chat/types.ts @@ -346,6 +346,8 @@ export interface DeepResearchToolCall { workflow?: string /** Parent agent ID that invoked the tool (for grouping under agents) */ agentId?: string + /** Whether this tool call executed against the sandbox runtime */ + isSandbox?: boolean /** Current execution status */ status: 'running' | 'complete' | 'error' /** When tool was called */ diff --git a/frontends/ui/src/features/layout/components/AgentCard.tsx b/frontends/ui/src/features/layout/components/AgentCard.tsx index cf5a217c9..c6fa73a5f 100644 --- a/frontends/ui/src/features/layout/components/AgentCard.tsx +++ b/frontends/ui/src/features/layout/components/AgentCard.tsx @@ -298,6 +298,11 @@ export const AgentCard: FC = ({ agent, defaultExpanded = true }) ) : ( <> {getToolDisplayName(toolCall.name)} + {toolCall.isSandbox && ( + + Sandbox + + )} {description !== toolCall.name && ( : {description} )} diff --git a/frontends/ui/src/features/layout/components/ExportFooter.spec.tsx b/frontends/ui/src/features/layout/components/ExportFooter.spec.tsx index ff3579c53..a94e7c211 100644 --- a/frontends/ui/src/features/layout/components/ExportFooter.spec.tsx +++ b/frontends/ui/src/features/layout/components/ExportFooter.spec.tsx @@ -15,6 +15,7 @@ let mockChatState: Record = { isDeepResearchStreaming: false, deepResearchStatus: null as 'submitted' | 'running' | 'success' | 'failure' | 'interrupted' | null, currentConversation: { title: 'AI Market Trends' }, + deepResearchJobId: 'job-123', } vi.mock('@/features/chat', () => ({ @@ -25,6 +26,7 @@ vi.mock('@/features/chat', () => ({ return mockChatState }, useIsCurrentSessionBusy: () => mockIsBusy, + selectResolvedDeepResearchJobId: (state: any) => state?.deepResearchJobId, })) // Mock the download utilities @@ -75,7 +77,7 @@ describe('ExportFooter', () => { await user.click(screen.getByRole('button', { name: /pdf/i })) - expect(mockDownloadPdf).toHaveBeenCalledWith('Some report content', 'AI Market Trends') + expect(mockDownloadPdf).toHaveBeenCalledWith('Some report content', 'AI Market Trends', 'job-123') }) test('disables buttons when disabled prop is true', () => { diff --git a/frontends/ui/src/features/layout/components/ExportFooter.tsx b/frontends/ui/src/features/layout/components/ExportFooter.tsx index 6c9cddceb..4016a5e58 100644 --- a/frontends/ui/src/features/layout/components/ExportFooter.tsx +++ b/frontends/ui/src/features/layout/components/ExportFooter.tsx @@ -12,9 +12,10 @@ import { type FC, useCallback, useState } from 'react' import { Banner, Flex, Button } from '@/adapters/ui' -import { useChatStore, useIsCurrentSessionBusy } from '@/features/chat' +import { useChatStore, useIsCurrentSessionBusy, selectResolvedDeepResearchJobId } from '@/features/chat' import { downloadAsMarkdown } from '@/utils/download-as-markdown' import { useDownloadPdfRoute } from '@/hooks/use-download-pdf' +import { rewriteArtifactRefs } from '@/shared/components/MarkdownRenderer' import { Download } from '@/adapters/ui/icons' interface ExportFooterProps { @@ -29,6 +30,8 @@ interface ExportFooterProps { export const ExportFooter: FC = ({ disabled }) => { const reportContent = useChatStore((state) => state.reportContent) const conversationTitle = useChatStore((state) => state.currentConversation?.title) + // Active job id, or the latest finished deep-research job, so export can resolve artifacts. + const deepResearchJobId = useChatStore(selectResolvedDeepResearchJobId) const { downloadPdf, isLoading: isPdfLoading, error: pdfError, clearError: clearPdfError } = useDownloadPdfRoute() const [mdError, setMdError] = useState(null) @@ -52,16 +55,22 @@ export const ExportFooter: FC = ({ disabled }) => { const handleExportMarkdown = useCallback(() => { if (isExportDisabled) return setMdError(null) - const result = downloadAsMarkdown(reportContentStr, conversationTitle ?? undefined) + // Rewrite artifact:// refs to absolute content URLs so the downloaded .md renders + // images while the backend is reachable. + const origin = typeof window !== 'undefined' ? window.location.origin : '' + const markdown = deepResearchJobId + ? rewriteArtifactRefs(reportContentStr, deepResearchJobId, origin) + : reportContentStr + const result = downloadAsMarkdown(markdown, conversationTitle ?? undefined) if (!result.success && result.error) { setMdError(result.error) } - }, [isExportDisabled, reportContentStr, conversationTitle]) + }, [isExportDisabled, reportContentStr, conversationTitle, deepResearchJobId]) const handleExportPDF = useCallback(() => { if (isExportDisabled || isPdfLoading) return - downloadPdf(reportContentStr, conversationTitle ?? undefined) - }, [isExportDisabled, isPdfLoading, reportContentStr, downloadPdf, conversationTitle]) + downloadPdf(reportContentStr, conversationTitle ?? undefined, deepResearchJobId ?? undefined) + }, [isExportDisabled, isPdfLoading, reportContentStr, downloadPdf, conversationTitle, deepResearchJobId]) const exportError = mdError || pdfError const clearExportError = useCallback(() => { diff --git a/frontends/ui/src/features/layout/components/ReportTab.spec.tsx b/frontends/ui/src/features/layout/components/ReportTab.spec.tsx index 58502179a..ee7b02981 100644 --- a/frontends/ui/src/features/layout/components/ReportTab.spec.tsx +++ b/frontends/ui/src/features/layout/components/ReportTab.spec.tsx @@ -15,6 +15,7 @@ vi.mock('@/features/chat', () => ({ } return selector ? selector(state) : state }), + selectResolvedDeepResearchJobId: (state: any) => state?.deepResearchJobId, })) // Mock MarkdownRenderer diff --git a/frontends/ui/src/features/layout/components/ReportTab.tsx b/frontends/ui/src/features/layout/components/ReportTab.tsx index aa9ed5526..c7898811a 100644 --- a/frontends/ui/src/features/layout/components/ReportTab.tsx +++ b/frontends/ui/src/features/layout/components/ReportTab.tsx @@ -19,7 +19,7 @@ import { Flex, Text } from '@/adapters/ui' import { useShallow } from 'zustand/react/shallow' import { Document } from '@/adapters/ui/icons' import { MarkdownRenderer } from '@/shared/components/MarkdownRenderer' -import { useChatStore } from '@/features/chat' +import { useChatStore, selectResolvedDeepResearchJobId } from '@/features/chat' import { ExportFooter } from './ExportFooter' interface ReportTabProps { @@ -40,6 +40,8 @@ export const ReportTab: FC = ({ children }) => { isStreaming: s.isStreaming, currentStatus: s.currentStatus, }))) + // Resolve the owning job id (active or latest finished) so artifact:// images render. + const deepResearchJobId = useChatStore(selectResolvedDeepResearchJobId) const reportContentStr = typeof reportContent === 'string' ? reportContent : '' const isEmpty = !reportContentStr.trim() @@ -77,6 +79,7 @@ export const ReportTab: FC = ({ children }) => { content={reportContentStr} isStreaming={false} className="max-w-none" + artifactJobId={deepResearchJobId ?? undefined} /> @@ -87,6 +90,7 @@ export const ReportTab: FC = ({ children }) => { content={reportContentStr} isStreaming={isGeneratingReport} className="max-w-none" + artifactJobId={deepResearchJobId ?? undefined} /> )} diff --git a/frontends/ui/src/features/layout/components/ToolCallCard.tsx b/frontends/ui/src/features/layout/components/ToolCallCard.tsx index 04a52b2cc..749077f8e 100644 --- a/frontends/ui/src/features/layout/components/ToolCallCard.tsx +++ b/frontends/ui/src/features/layout/components/ToolCallCard.tsx @@ -33,6 +33,8 @@ export interface ToolCallInfo { timestamp?: Date | string /** Parent agent that invoked the tool */ workflow?: string + /** Whether this tool call executed against the sandbox runtime */ + isSandbox?: boolean /** Error message if status is error */ error?: string } @@ -139,9 +141,16 @@ export const ToolCallCard: FC = ({ toolCall }) => { {/* Tool Info */} - - {toolCall.name} - + + + {toolCall.name} + + {toolCall.isSandbox && ( + + Sandbox + + )} + {toolCall.workflow && ( via {toolCall.workflow} diff --git a/frontends/ui/src/hooks/use-download-pdf.ts b/frontends/ui/src/hooks/use-download-pdf.ts index 19b92646b..56e384536 100644 --- a/frontends/ui/src/hooks/use-download-pdf.ts +++ b/frontends/ui/src/hooks/use-download-pdf.ts @@ -12,7 +12,7 @@ export const useDownloadPdfRoute = () => { setError(null) }, []) - const downloadPdf = async (markdown: string, filename?: string) => { + const downloadPdf = async (markdown: string, filename?: string, jobId?: string) => { setIsLoading(true) setError(null) @@ -20,7 +20,7 @@ export const useDownloadPdfRoute = () => { const response = await fetch('/api/generate-pdf', { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ markdown }), + body: JSON.stringify({ markdown, jobId }), }) if (!response.ok) { diff --git a/frontends/ui/src/lib/pdf/ReactPdfDocument.tsx b/frontends/ui/src/lib/pdf/ReactPdfDocument.tsx index d316bd95b..f0e23bec8 100644 --- a/frontends/ui/src/lib/pdf/ReactPdfDocument.tsx +++ b/frontends/ui/src/lib/pdf/ReactPdfDocument.tsx @@ -2,10 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 import React from 'react' -import { Document, Page, Text, View, StyleSheet, Font, Link } from '@react-pdf/renderer' +import { Document, Page, Text, View, StyleSheet, Font, Link, Image } from '@react-pdf/renderer' import { marked } from 'marked' type Token = ReturnType[number] +type ImageToken = { type: 'image'; href: string; text?: string; title?: string } type HeadingToken = Extract type ParagraphToken = Extract type ListToken = Extract @@ -140,6 +141,25 @@ const styles = StyleSheet.create({ color: '#0066cc', textDecoration: 'underline', }, + figure: { + marginTop: 10, + marginBottom: 10, + alignItems: 'center', + }, + image: { + // Explicit width is required: react-pdf renders an Image at intrinsic pixel size when + // width is unset, so a large chart overflows the page. '100%' fits the parent and scales + // height by aspect ratio. + width: '100%', + objectFit: 'contain', + }, + imageCaption: { + fontSize: 9, + color: '#555555', + marginTop: 4, + textAlign: 'center', + fontStyle: 'italic', + }, }) interface MarkdownPDFProps { @@ -197,11 +217,78 @@ function renderHeading(token: HeadingToken, index: number): React.ReactNode { ) } +// Matches markdown image syntax so it can be stripped from inline text (the image is rendered +// as a block figure instead of leaking through the inline link parser as a stray link). +const IMAGE_MD_RE = /!\[[^\]]*\]\([^)]*\)/g + +// Only embed raster image data URIs within a bounded decoded size — the markdown may carry an +// arbitrary `data:` URI (e.g. one the report wrote directly), which we must not feed unchecked +// into the PDF renderer. +const MAX_PDF_EMBED_BYTES = 8 * 1024 * 1024 +// @react-pdf/renderer supports PNG and JPEG only; a WebP data URI would pass the size +// guard and then fail silently during PDF rendering, so it is excluded here. +const DATA_IMAGE_RE = /^data:image\/(?:png|jpe?g);base64,([A-Za-z0-9+/=]+)$/ + +function isEmbeddableDataImage(href: string): boolean { + const match = DATA_IMAGE_RE.exec(href) + if (!match) return false + const b64 = match[1] + const padding = b64.endsWith('==') ? 2 : b64.endsWith('=') ? 1 : 0 + const decodedBytes = Math.floor((b64.length * 3) / 4) - padding + return decodedBytes <= MAX_PDF_EMBED_BYTES +} + +/** + * Recursively collect embeddable image tokens. Only bounded raster `data:` URIs are embeddable + * (artifact refs are pre-resolved to data URIs server-side); remote/unresolved/oversized images + * are skipped so the PDF never shows a broken figure or blows up memory. Walks nested `tokens` + * so images inside list items, blockquotes, etc. are found, not just top-level paragraphs. + */ +function collectEmbeddableImages(tokens: unknown): ImageToken[] { + if (!Array.isArray(tokens)) return [] + const found: ImageToken[] = [] + for (const token of tokens) { + const t = token as { type?: string; href?: string; tokens?: unknown } + if (t.type === 'image' && typeof t.href === 'string' && isEmbeddableDataImage(t.href)) { + found.push(t as ImageToken) + } else if (Array.isArray(t.tokens)) { + found.push(...collectEmbeddableImages(t.tokens)) + } + } + return found +} + +/** Render embeddable images as centered block figures with optional captions. */ +function renderFigures(images: ImageToken[], keyPrefix: string): React.ReactNode[] { + return images.map((img, imgIndex) => ( + + + {(img.text || img.title) && {img.text || img.title}} + + )) +} + function renderParagraph(token: ParagraphToken, index: number): React.ReactNode { + const images = collectEmbeddableImages((token as ParagraphToken & { tokens?: Token[] }).tokens) + // Strip image markdown so a standalone `![alt](data:...)` does not also render as a link. + const textWithoutImages = token.text.replace(IMAGE_MD_RE, '').trim() + + if (images.length === 0) { + if (!textWithoutImages) return null + return ( + + {parseInlineFormatting(textWithoutImages)} + + ) + } + return ( - - {parseInlineFormatting(token.text)} - + + {textWithoutImages && ( + {parseInlineFormatting(textWithoutImages)} + )} + {renderFigures(images, String(index))} + ) } @@ -221,7 +308,7 @@ function renderList(token: ListToken, index: number, _nested: boolean = false): }) } - const mainText = textTokens.length + const rawText = textTokens.length ? textTokens .map((t: any) => { if ('text' in t) { @@ -235,7 +322,12 @@ function renderList(token: ListToken, index: number, _nested: boolean = false): .join(' ') : stripHtml(preserveHtmlLinks(item.text)) - const hasNestedContent = nestedLists.length > 0 + // A bullet may carry an embedded figure (e.g. "- The chart: ![alt](data:...)"). Collect + // only from this item's own text tokens (not nestedLists, which render their own images) + // so images in nested list items aren't rendered twice. + const images = collectEmbeddableImages(textTokens) + const mainText = rawText.replace(IMAGE_MD_RE, '').trim() + const hasNestedContent = nestedLists.length > 0 || images.length > 0 return ( {bullet} - {parseInlineFormatting(mainText)} + {mainText && {parseInlineFormatting(mainText)}} + {renderFigures(images, `${index}-${itemIndex}`)} {nestedLists.map((nestedList: any, nlIndex: number) => renderList(nestedList as ListToken, nlIndex, true) )} diff --git a/frontends/ui/src/pages/api/generate-pdf.ts b/frontends/ui/src/pages/api/generate-pdf.ts index 159a62715..b18455afb 100644 --- a/frontends/ui/src/pages/api/generate-pdf.ts +++ b/frontends/ui/src/pages/api/generate-pdf.ts @@ -5,10 +5,86 @@ import type { NextApiRequest, NextApiResponse } from 'next' import React from 'react' import { renderToStream } from '@react-pdf/renderer' import { MarkdownPDF } from '../../lib/pdf/ReactPdfDocument' +import { extractArtifactIds, replaceArtifactImages } from '../../shared/components/MarkdownRenderer/artifact-url' +import { isAuthRequired } from '@/adapters/auth/config' + +// Cap the bytes we embed per image. Charts are tiny; this guards against base64-inflating a +// large artifact (up to the 50 MB harvest cap) into a pathological PDF. +const MAX_PDF_IMAGE_BYTES = 8 * 1024 * 1024 +// Bound the number of artifact fetches per PDF so a report with many refs can't fan out into +// excessive backend load / memory. +const MAX_PDF_ARTIFACT_REFS = 25 + +const getBackendUrl = (): string => { + const url = process.env.BACKEND_URL || process.env.NEXT_PUBLIC_BACKEND_URL || 'http://localhost:8000' + return url.replace(/\/$/, '') +} + +// This runs server-side and calls the backend directly, so it targets the backend's `/v1` +// route. (The `/api/...` path that the browser uses only exists on the Next.js proxy.) +const backendArtifactContentPath = (jobId: string, artifactId: string): string => + `/v1/jobs/async/job/${encodeURIComponent(jobId)}/artifacts/${encodeURIComponent(artifactId)}/content` + +/** + * Replace every `![alt](artifact://)` with a self-contained `data:` URI by fetching the + * artifact bytes from the backend. This keeps the PDF fully embedded (no runtime network or + * auth needed during react-pdf rendering). Refs that fail to resolve are left untouched and + * the PDF renderer skips them. + */ +const inlineArtifactImages = async ( + markdown: string, + jobId: string | undefined, + authHeaders: Record +): Promise => { + if (!jobId) return markdown + + const ids = extractArtifactIds(markdown).slice(0, MAX_PDF_ARTIFACT_REFS) + console.log(`[PDF] inline: jobId=${jobId} artifactRefs=${ids.length}`) + if (ids.length === 0) return markdown + + const backend = getBackendUrl() + const dataUris = new Map() + await Promise.all( + ids.map(async (id) => { + try { + const resp = await fetch(`${backend}${backendArtifactContentPath(jobId, id)}`, { + headers: { ...authHeaders, Accept: '*/*' }, + // Bound the call so a stalled backend can't hang PDF generation indefinitely. + signal: AbortSignal.timeout(15_000), + }) + const contentType = resp.headers.get('Content-Type') ?? 'application/octet-stream' + console.log(`[PDF] fetch ${id}: status=${resp.status} type=${contentType}`) + if (!resp.ok) return + // Only raster images are embeddable in the PDF. + if (!contentType.startsWith('image/')) return + // Reject oversized artifacts by declared length before buffering the whole body. + const declaredLen = Number(resp.headers.get('Content-Length')) + if (Number.isFinite(declaredLen) && declaredLen > MAX_PDF_IMAGE_BYTES) { + console.warn(`[PDF] Skipping artifact ${id}: declared ${declaredLen} bytes exceeds embed cap`) + return + } + const buffer = Buffer.from(await resp.arrayBuffer()) + if (buffer.byteLength > MAX_PDF_IMAGE_BYTES) { + console.warn(`[PDF] Skipping artifact ${id}: ${buffer.byteLength} bytes exceeds embed cap`) + return + } + dataUris.set(id, `data:${contentType};base64,${buffer.toString('base64')}`) + } catch (err) { + console.error('[PDF] Failed to fetch artifact', id, err) + } + }) + ) + + console.log(`[PDF] inline: embedded ${dataUris.size}/${ids.length} image(s) as data URIs`) + return replaceArtifactImages(markdown, (alt, id) => { + const uri = dataUris.get(id) + return uri ? `![${alt}](${uri})` : null + }) +} /** * POST /api/generate-pdf - * Receives: { markdown: string } in JSON body + * Receives: { markdown: string, jobId?: string } in JSON body * Returns: PDF file generated from markdown as application/pdf */ export default async function handler(req: NextApiRequest, res: NextApiResponse) { @@ -17,14 +93,33 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse) } try { - const { markdown } = req.body + const { markdown, jobId } = req.body if (!markdown || typeof markdown !== 'string') { return res.status(400).json({ error: 'Invalid or missing markdown content' }) } + // Forward only the auth the artifact endpoint needs — the Authorization header and the + // idToken cookie — rather than the caller's entire cookie jar. Gated on isAuthRequired() + // so anonymous mode (REQUIRE_AUTH=false) sends no identity headers, matching the jobs proxy. + const authHeaders: Record = {} + if (isAuthRequired()) { + if (req.headers.authorization) authHeaders.Authorization = req.headers.authorization + const idTokenCookie = req.headers.cookie + ?.split(';') + .map((c) => c.trim()) + .find((c) => c.startsWith('idToken=')) + if (idTokenCookie) authHeaders.Cookie = idTokenCookie + } + + const resolvedMarkdown = await inlineArtifactImages( + markdown, + typeof jobId === 'string' ? jobId : undefined, + authHeaders + ) + const stream = await renderToStream( - React.createElement(MarkdownPDF, { markdown }) as React.ReactElement + React.createElement(MarkdownPDF, { markdown: resolvedMarkdown }) as React.ReactElement ) const chunks: Buffer[] = [] diff --git a/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.spec.tsx b/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.spec.tsx index 782c364c3..79af7ebe6 100644 --- a/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.spec.tsx +++ b/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.spec.tsx @@ -174,6 +174,36 @@ Paragraph 2.`} />) }) }) + describe('images', () => { + test('resolves artifact:// refs to the content endpoint and renders a caption', () => { + render( + + ) + + const img = screen.getByRole('img', { name: 'Population chart' }) + expect(img).toHaveAttribute('src', '/api/jobs/async/job/job-9/artifacts/art_abc123/content') + // Caption renders in a phrasing-safe (not
) to keep valid + // nesting inside the markdown

wrapper. + expect(screen.getByText('Population chart').tagName).toBe('SPAN') + }) + + test('skips an artifact:// image when no job id is available', () => { + render() + + expect(screen.queryByRole('img')).not.toBeInTheDocument() + }) + + test('passes through a normal image url', () => { + render() + + const img = screen.getByRole('img', { name: 'Logo' }) + expect(img).toHaveAttribute('src', 'https://example.com/logo.png') + }) + }) + describe('code blocks', () => { test('renders code block with language', () => { const code = '```javascript\nconst x = 1;\n```' diff --git a/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.tsx b/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.tsx index eba9bff0c..efb852fd3 100644 --- a/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.tsx +++ b/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.tsx @@ -4,11 +4,18 @@ 'use client' import { type FC, type ReactNode, memo, useMemo } from 'react' -import ReactMarkdown, { type Components, type ExtraProps } from 'react-markdown' +import ReactMarkdown, { type Components, type ExtraProps, defaultUrlTransform } from 'react-markdown' import remarkGfm from 'remark-gfm' import { Text, CodeSnippet, Anchor } from '@/adapters/ui' import type { MarkdownRendererProps } from './types' import { getLanguageFromClassName } from './utils' +import { ARTIFACT_SCHEME, isArtifactRef, resolveArtifactUrl } from './artifact-url' + +// react-markdown's default sanitizer strips non-standard URL schemes, which would blank the +// src of `artifact://` images before the `img` renderer can resolve them. Preserve that +// scheme and defer to the default transform for everything else (keeps XSS protection). +const urlTransform = (url: string): string => + url.startsWith(ARTIFACT_SCHEME) ? url : defaultUrlTransform(url) function getTextFromChildren(node: ReactNode): string { if (typeof node === 'string') return node @@ -38,7 +45,7 @@ function slugify(text: string): string { * @param compact - Use smaller text sizes for chat bubbles */ export const MarkdownRenderer: FC = memo( - ({ content, className = '', compact = false }) => { + ({ content, className = '', compact = false, artifactJobId }) => { // Custom component mappings const components: Components = useMemo( () => ({ @@ -162,6 +169,36 @@ export const MarkdownRenderer: FC = memo( ) }, + // Images — resolve durable artifact:// refs to the content endpoint and render + // as a captioned figure; pass other images through with responsive styling. + img: ({ src, alt }) => { + const rawSrc = typeof src === 'string' ? src : '' + const resolved = isArtifactRef(rawSrc) + ? resolveArtifactUrl(rawSrc, artifactJobId) + : rawSrc + // An unresolved artifact ref (no job id) would be a broken image — skip it. + if (!resolved || isArtifactRef(resolved)) return null + const caption = alt ?? '' + return ( + // react-markdown renders images inside a

, so use phrasing-content spans + // (not

/
, which are invalid inside

). + + {/* eslint-disable-next-line @next/next/no-img-element */} + {caption} + {caption && ( + + {caption} + + )} + + ) + }, + // Emphasis strong: ({ children }) => ( {children} @@ -198,12 +235,12 @@ export const MarkdownRenderer: FC = memo( ), }), - [compact] + [compact, artifactJobId] ) return (

*:last-child]:mb-0 ${className}`}> - + {content}
diff --git a/frontends/ui/src/shared/components/MarkdownRenderer/artifact-url.spec.ts b/frontends/ui/src/shared/components/MarkdownRenderer/artifact-url.spec.ts new file mode 100644 index 000000000..b01374fc8 --- /dev/null +++ b/frontends/ui/src/shared/components/MarkdownRenderer/artifact-url.spec.ts @@ -0,0 +1,61 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, test, expect } from 'vitest' +import { isArtifactRef, artifactIdFromRef, resolveArtifactUrl, rewriteArtifactRefs } from './artifact-url' + +describe('artifact-url', () => { + describe('isArtifactRef', () => { + test('detects the artifact scheme', () => { + expect(isArtifactRef('artifact://art_1')).toBe(true) + expect(isArtifactRef('https://example.com/x.png')).toBe(false) + expect(isArtifactRef(undefined)).toBe(false) + }) + }) + + describe('artifactIdFromRef', () => { + test('strips the scheme and trims', () => { + expect(artifactIdFromRef('artifact:// art_42 ')).toBe('art_42') + }) + }) + + describe('resolveArtifactUrl', () => { + test('builds the same-origin content URL', () => { + expect(resolveArtifactUrl('artifact://art_42', 'job-1')).toBe( + '/api/jobs/async/job/job-1/artifacts/art_42/content' + ) + }) + + test('returns the original src for non-artifact refs', () => { + expect(resolveArtifactUrl('https://example.com/x.png', 'job-1')).toBe( + 'https://example.com/x.png' + ) + }) + + test('returns the original src when job id is missing', () => { + expect(resolveArtifactUrl('artifact://art_42', undefined)).toBe('artifact://art_42') + }) + }) + + describe('rewriteArtifactRefs', () => { + test('rewrites every artifact image ref to a relative content URL', () => { + const md = 'See ![Chart](artifact://art_1) and ![Other](artifact://art_2).' + expect(rewriteArtifactRefs(md, 'job-1')).toBe( + 'See ![Chart](/api/jobs/async/job/job-1/artifacts/art_1/content) and ' + + '![Other](/api/jobs/async/job/job-1/artifacts/art_2/content).' + ) + }) + + test('supports an absolute origin for downloaded markdown', () => { + const md = '![Chart](artifact://art_1)' + expect(rewriteArtifactRefs(md, 'job-1', 'https://host')).toBe( + '![Chart](https://host/api/jobs/async/job/job-1/artifacts/art_1/content)' + ) + }) + + test('leaves markdown untouched without a job id', () => { + const md = '![Chart](artifact://art_1)' + expect(rewriteArtifactRefs(md, undefined)).toBe(md) + }) + }) +}) diff --git a/frontends/ui/src/shared/components/MarkdownRenderer/artifact-url.ts b/frontends/ui/src/shared/components/MarkdownRenderer/artifact-url.ts new file mode 100644 index 000000000..bb406dd91 --- /dev/null +++ b/frontends/ui/src/shared/components/MarkdownRenderer/artifact-url.ts @@ -0,0 +1,74 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +/** + * Resolution of the backend `artifact://` scheme to loadable URLs. + * + * The deep-research report stores image references as `![caption](artifact://)`. + * Each render surface rewrites that scheme to a real source at its own edge; this module is + * the single source of truth for building the same-origin content URL used by the UI. + */ + +export const ARTIFACT_SCHEME = 'artifact://' + +/** True when a markdown image `src` uses the durable `artifact://` scheme. */ +export const isArtifactRef = (src: string | undefined): src is string => + typeof src === 'string' && src.startsWith(ARTIFACT_SCHEME) + +/** Extract the artifact id from an `artifact://` reference (trims any path/query noise). */ +export const artifactIdFromRef = (src: string): string => + src.slice(ARTIFACT_SCHEME.length).trim() + +/** + * Build the same-origin proxy URL that streams an artifact's bytes. + * + * Returns the original `src` unchanged when it is not an `artifact://` ref or when the owning + * `jobId` is unknown (so a bad ref degrades to a broken image rather than a wrong fetch). + */ +export const resolveArtifactUrl = (src: string | undefined, jobId?: string): string | undefined => { + if (!isArtifactRef(src) || !jobId) return src + const id = artifactIdFromRef(src) + if (!id) return src + return artifactContentPath(jobId, id) +} + +// Single source of truth for the markdown image -> artifact ref pattern. `matchAll` +// clones the regex internally and `replace` resets lastIndex, so sharing this module-level +// instance across the helpers below is safe. +const ARTIFACT_IMG_RE = /!\[([^\]]*)\]\(artifact:\/\/([^)]+)\)/g + +/** Build the relative same-origin content path for an artifact (no leading origin). */ +export const artifactContentPath = (jobId: string, artifactId: string): string => + `/api/jobs/async/job/${encodeURIComponent(jobId)}/artifacts/${encodeURIComponent(artifactId)}/content` + +/** + * Replace every `![alt](artifact://)` image using `replacer`. Returning `null` from the + * replacer leaves the original token untouched. Shared by the UI/download rewrite and the + * server-side PDF inliner so the matching rule lives in exactly one place. + */ +export const replaceArtifactImages = ( + markdown: string, + replacer: (alt: string, artifactId: string) => string | null +): string => + markdown.replace(ARTIFACT_IMG_RE, (full, alt: string, id: string) => replacer(alt, id.trim()) ?? full) + +/** Collect the unique artifact ids referenced as images in the markdown. */ +export const extractArtifactIds = (markdown: string): string[] => { + const ids = new Set() + for (const match of markdown.matchAll(ARTIFACT_IMG_RE)) ids.add(match[2].trim()) + return Array.from(ids) +} + +/** + * Rewrite every `![alt](artifact://)` in a markdown string to a content URL. + * + * @param markdown report markdown that may contain `artifact://` image refs + * @param jobId owning job id used to build the content URL + * @param origin optional absolute origin (e.g. `https://host`); when provided the URL is + * absolute so the markdown renders outside the app (downloaded `.md`). Defaults to a + * same-origin relative URL. + */ +export const rewriteArtifactRefs = (markdown: string, jobId?: string, origin = ''): string => { + if (!jobId) return markdown + return replaceArtifactImages(markdown, (alt, id) => `![${alt}](${origin}${artifactContentPath(jobId, id)})`) +} diff --git a/frontends/ui/src/shared/components/MarkdownRenderer/index.ts b/frontends/ui/src/shared/components/MarkdownRenderer/index.ts index 0c83396ce..cab48f1a5 100644 --- a/frontends/ui/src/shared/components/MarkdownRenderer/index.ts +++ b/frontends/ui/src/shared/components/MarkdownRenderer/index.ts @@ -3,3 +3,13 @@ export { MarkdownRenderer } from './MarkdownRenderer' export type { MarkdownRendererProps, SupportedLanguage } from './types' +export { + ARTIFACT_SCHEME, + isArtifactRef, + artifactIdFromRef, + artifactContentPath, + resolveArtifactUrl, + replaceArtifactImages, + extractArtifactIds, + rewriteArtifactRefs, +} from './artifact-url' diff --git a/frontends/ui/src/shared/components/MarkdownRenderer/types.ts b/frontends/ui/src/shared/components/MarkdownRenderer/types.ts index 91b3deb0b..68f00074d 100644 --- a/frontends/ui/src/shared/components/MarkdownRenderer/types.ts +++ b/frontends/ui/src/shared/components/MarkdownRenderer/types.ts @@ -10,6 +10,11 @@ export interface MarkdownRendererProps { className?: string /** Use compact text sizes (for chat bubbles vs full reports) */ compact?: boolean + /** + * Owning job id used to resolve `artifact://` image refs to the content endpoint. + * When omitted, artifact images are not resolved (rendered as-is). + */ + artifactJobId?: string } /** Supported languages for syntax highlighting */ diff --git a/scripts/README.md b/scripts/README.md index a19c6a9df..12a7dbbef 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -49,6 +49,63 @@ Starts the agent in CLI mode with browser-based authentication. | `--verbose` or `-v` | Enable verbose logging | | `--config_file ` | Use a custom configuration file | +### `setup_openshell.sh` - OpenShell Sandbox Setup + +Sets up the experimental, local single-operator NVIDIA OpenShell path for AI-Q. Run this once before using +`configs/config_openshell.yml` with `start_cli.sh` or `start_e2e.sh`. It installs +the `openshell` SDK and the `langchain-nvidia-openshell` adapter, starts/verifies +the local OpenShell gateway, builds the sandbox image, generates a network policy, +and creates the named sandbox `aiq-openshell-demo`. Inference is unaffected (it +stays host-side, routed to NVIDIA Build); only generated code runs in the +network-blocked sandbox. + +The generated configuration attaches all jobs to one named sandbox. Per-job directories +avoid filename collisions but do not isolate mutually untrusted jobs, and AI-Q does not +verify the provisioned policy when attaching. Do not treat this setup as a multi-tenant +security boundary. + +```bash +./scripts/setup_openshell.sh --policy offline +./scripts/start_e2e.sh --config_file configs/config_openshell.yml +# or direct serve: +dotenv -f deploy/.env run .venv/bin/nat serve --config_file configs/config_openshell.yml --host 0.0.0.0 --port 8000 +``` + +Useful version examples: + +```bash +./scripts/setup_openshell.sh --openshell-version 0.0.72 +./scripts/setup_openshell.sh --openshell-version latest +./scripts/setup_openshell.sh --list-openshell-versions +``` + +In the interactive version prompt, pressing Enter selects `0.0.72`. + +The setup installs the `openshell` SDK plus the official `langchain-nvidia-openshell` +adapter (`OpenShellSandbox`), published on PyPI. The script installs it from PyPI by +default; set `LANGCHAIN_NVIDIA_REPO` or pass `--langchain-nvidia` to use another +`uv pip install` spec or a local checkout. + +Useful policy examples: + +```bash +./scripts/setup_openshell.sh --policy offline +./scripts/setup_openshell.sh --policy research +./scripts/setup_openshell.sh --policy python-packages +./scripts/setup_openshell.sh --policy custom --allow github,pypi,nvidia,tavily +``` + +Verify and clean up: + +```bash +.venv/bin/openshell status +.venv/bin/openshell sandbox list # expect: aiq-openshell-demo ... Ready +.venv/bin/openshell sandbox delete aiq-openshell-demo +# Inspect, then stop only the gateway you started (avoid killing other sessions): +pgrep -fl openshell-gateway # find the PID(s) +kill # stop the specific process +``` + ### `start_server_in_debug_mode.sh` - Server Mode @@ -112,6 +169,7 @@ Starts both backend and frontend for full WebSocket support and HITL workflows. ```bash ./scripts/start_e2e.sh +./scripts/start_e2e.sh --config_file configs/config_openshell.yml ``` **Services:** @@ -128,6 +186,8 @@ Starts both backend and frontend for full WebSocket support and HITL workflows. | `configs/config_cli_default.yml` | CLI mode with web search (default) | | `configs/config_web_frag.yml` | Server/E2E mode with Foundational RAG | | `configs/config_web_default_llamaindex.yml` | Server/E2E mode with LlamaIndex | +| `configs/config_skills.yml` | Deep research with DeepAgents skills + Modal sandbox | +| `configs/config_openshell.yml` | Experimental local single-operator OpenShell sandbox + artifact capture (run `setup_openshell.sh` first) | ## Development Workflow diff --git a/scripts/setup_openshell.sh b/scripts/setup_openshell.sh new file mode 100755 index 000000000..9a5a62493 --- /dev/null +++ b/scripts/setup_openshell.sh @@ -0,0 +1,1087 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Set up NVIDIA OpenShell for AI-Q with a named, policy-backed sandbox. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(dirname "$SCRIPT_DIR")" +VENV_DIR="$REPO_ROOT/.venv" + +# Floor aligned with the published langchain-nvidia-openshell 0.1.0 adapter, which is +# tested against OpenShell 0.0.72+. Anything below this is upgraded by the adapter during +# its install, so do not pin under it. +MIN_OPENSHELL_VERSION="0.0.72" +DEFAULT_OPENSHELL_VERSION="0.0.72" +OPENSHELL_VERSION="${AIQ_OPENSHELL_VERSION:-}" +OPENSHELL_VERSION_USER_SUPPLIED=false +if [[ -n "$OPENSHELL_VERSION" ]]; then + OPENSHELL_VERSION_USER_SUPPLIED=true +fi +OPENSHELL_LATEST_VERSION="" +OPENSHELL_AVAILABLE_VERSIONS="" +PYTHON_BIN="" +# Official OpenShell deepagents adapter: the `langchain-nvidia-openshell` partner +# package, now published on PyPI. Override with LANGCHAIN_NVIDIA_REPO to use a git +# spec or a local checkout (e.g. to test an unreleased adapter build). +DEFAULT_LANGCHAIN_NVIDIA_INSTALL_SPEC="langchain-nvidia-openshell==0.1.0" +LANGCHAIN_NVIDIA_REPO="${LANGCHAIN_NVIDIA_REPO:-$DEFAULT_LANGCHAIN_NVIDIA_INSTALL_SPEC}" +SANDBOX_NAME="${AIQ_OPENSHELL_SANDBOX_NAME:-aiq-openshell-demo}" +IMAGE_NAME="${AIQ_OPENSHELL_IMAGE:-aiq-openshell-demo:latest}" +# Sandbox log verbosity baked into the image (RUST_LOG). Default `warn` is OpenShell's +# stock sandbox level; set to `debug` to surface in-container process/relay detail. +SANDBOX_LOG_LEVEL="${AIQ_OPENSHELL_SANDBOX_LOG_LEVEL:-warn}" +POLICY_PRESET="${AIQ_OPENSHELL_POLICY:-}" +POLICY_ALLOWLIST="${AIQ_OPENSHELL_POLICY_ALLOWLIST:-${AIQ_OPENSHELL_POLICY_SERVICES:-}}" +POLICY_FILE="${AIQ_OPENSHELL_POLICY_FILE:-$REPO_ROOT/configs/openshell/generated/aiq-openshell-policy.yaml}" +GATEWAY_NAME="${AIQ_OPENSHELL_GATEWAY_NAME:-aiq-local}" +GATEWAY_PORT="${AIQ_OPENSHELL_GATEWAY_PORT:-8080}" +DOCKER_BIN="${DOCKER_BIN:-}" +OPENSHELL_BIN="${OPENSHELL_BIN:-}" +OPENSHELL_GATEWAY_LAUNCH_BIN="${OPENSHELL_GATEWAY_LAUNCH_BIN:-}" +GATEWAY_ENDPOINT="" +GATEWAY_DISABLE_TLS=false +RESTART_GATEWAY=true +BUILD_IMAGE=true +CREATE_SANDBOX=true +LIST_OPENSHELL_VERSIONS=false + +SUPPORTED_SERVICES="github,pypi,nvidia,tavily,serper,huggingface,arxiv,semantic-scholar,npm" +SUPPORTED_POLICIES="offline,research,python-packages,ai-dev,custom" + +usage() { + cat <= 0.0.72. + -h, --help Show this help. + +Examples: + scripts/setup_openshell.sh + scripts/setup_openshell.sh --policy python-packages + scripts/setup_openshell.sh --policy custom --allow github,pypi,nvidia,tavily + scripts/setup_openshell.sh --openshell-version latest --policy offline +EOF +} + +log() { + echo "" + echo "==> $*" +} + +fail() { + echo "ERROR: $*" >&2 + exit 1 +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --openshell-version) + OPENSHELL_VERSION="$2" + OPENSHELL_VERSION_USER_SUPPLIED=true + shift 2 + ;; + --policy) + POLICY_PRESET="$2" + shift 2 + ;; + --allow) + POLICY_ALLOWLIST="$2" + shift 2 + ;; + --policy-services) + POLICY_ALLOWLIST="$2" + shift 2 + ;; + --policy-file) + POLICY_FILE="$2" + shift 2 + ;; + --sandbox-name) + SANDBOX_NAME="$2" + shift 2 + ;; + --image-name) + IMAGE_NAME="$2" + shift 2 + ;; + --sandbox-log-level) + SANDBOX_LOG_LEVEL="$2" + shift 2 + ;; + --langchain-nvidia) + LANGCHAIN_NVIDIA_REPO="$2" + shift 2 + ;; + --gateway-name) + GATEWAY_NAME="$2" + shift 2 + ;; + --gateway-port) + GATEWAY_PORT="$2" + shift 2 + ;; + --docker-bin) + DOCKER_BIN="$2" + shift 2 + ;; + --gateway-bin) + OPENSHELL_GATEWAY_LAUNCH_BIN="$2" + shift 2 + ;; + --no-restart-gateway) + RESTART_GATEWAY=false + shift + ;; + --skip-build) + BUILD_IMAGE=false + shift + ;; + --skip-sandbox) + CREATE_SANDBOX=false + shift + ;; + --list-policies) + echo "$SUPPORTED_POLICIES" | tr ',' '\n' + exit 0 + ;; + --list-services) + echo "$SUPPORTED_SERVICES" | tr ',' '\n' + exit 0 + ;; + --list-openshell-versions) + LIST_OPENSHELL_VERSIONS=true + shift + ;; + -h|--help) + usage + exit 0 + ;; + *) + usage + fail "Unknown option: $1" + ;; + esac +done + +detect_os() { + case "$(uname -s)" in + Darwin) + OS_NAME=macos + ;; + Linux) + OS_NAME=linux + ;; + *) + fail "Unsupported operating system: $(uname -s). This script supports macOS and Linux." + ;; + esac + echo "Detected OS: $OS_NAME" +} + +require_uv() { + if ! command -v uv >/dev/null 2>&1; then + cat <<'EOF' +uv was not found on PATH. + +Install uv, then rerun: + + curl -LsSf https://astral.sh/uv/install.sh | sh + +If uv is already installed, open a new shell or add it to PATH before rerunning. +EOF + exit 1 + fi +} + +resolve_python() { + if [[ -n "$PYTHON_BIN" && -x "$PYTHON_BIN" ]]; then + return + fi + PYTHON_BIN="$(uv python find 2>/dev/null || true)" + if [[ -z "$PYTHON_BIN" || ! -x "$PYTHON_BIN" ]]; then + PYTHON_BIN="$(command -v python3 || command -v python || true)" + fi + if [[ -z "$PYTHON_BIN" || ! -x "$PYTHON_BIN" ]]; then + fail "Python was not found. Install Python 3.11+ or run ./scripts/setup.sh first." + fi +} + +fetch_openshell_versions() { + resolve_python + log "Checking available OpenShell versions" + local output + if ! output="$("$PYTHON_BIN" - "$MIN_OPENSHELL_VERSION" <<'PY' +import json +import re +import sys +import urllib.request + +min_version = sys.argv[1] +version_re = re.compile(r"^\d+\.\d+\.\d+$") + +def parse(version: str) -> tuple[int, int, int]: + return tuple(int(part) for part in version.split(".")) + +try: + with urllib.request.urlopen("https://pypi.org/pypi/openshell/json", timeout=15) as response: + payload = json.load(response) +except Exception as exc: + raise SystemExit(f"failed to fetch OpenShell versions from PyPI: {exc}") + +minimum = parse(min_version) +versions = [] +for version, files in payload.get("releases", {}).items(): + if not version_re.match(version): + continue + if not files: + continue + parsed = parse(version) + if parsed >= minimum: + versions.append((parsed, version)) + +if not versions: + raise SystemExit(f"no OpenShell releases found at or above {min_version}") + +versions.sort() +print(versions[-1][1]) +print(",".join(version for _, version in versions)) +PY +)"; then + fail "$output" + fi + + OPENSHELL_LATEST_VERSION="$(printf '%s\n' "$output" | sed -n '1p')" + OPENSHELL_AVAILABLE_VERSIONS="$(printf '%s\n' "$output" | sed -n '2p')" + + if [[ -z "$OPENSHELL_LATEST_VERSION" || -z "$OPENSHELL_AVAILABLE_VERSIONS" ]]; then + fail "Could not determine available OpenShell versions from PyPI." + fi + echo "OpenShell version range: $MIN_OPENSHELL_VERSION through $OPENSHELL_LATEST_VERSION" +} + +version_is_available() { + case ",$OPENSHELL_AVAILABLE_VERSIONS," in + *",$1,"*) + return 0 + ;; + *) + return 1 + ;; + esac +} + +version_prompt_text() { + echo "OpenShell version [press Enter for ${DEFAULT_OPENSHELL_VERSION}; type latest for ${OPENSHELL_LATEST_VERSION}]: " +} + +choose_openshell_version_interactively() { + local candidate + while true; do + read -r -p "$(version_prompt_text)" candidate + candidate="${candidate:-$DEFAULT_OPENSHELL_VERSION}" + if [[ "$candidate" == "latest" ]]; then + candidate="$OPENSHELL_LATEST_VERSION" + fi + if version_is_available "$candidate"; then + OPENSHELL_VERSION="$candidate" + return + fi + echo "OpenShell version '$candidate' was not found between $MIN_OPENSHELL_VERSION and $OPENSHELL_LATEST_VERSION." + echo "Try an exact released version, '$DEFAULT_OPENSHELL_VERSION', or 'latest'." + done +} + +resolve_openshell_version() { + fetch_openshell_versions + + if [[ -n "$OPENSHELL_VERSION" ]]; then + if [[ "$OPENSHELL_VERSION" == "latest" ]]; then + OPENSHELL_VERSION="$OPENSHELL_LATEST_VERSION" + fi + if version_is_available "$OPENSHELL_VERSION"; then + echo "OpenShell version selected: $OPENSHELL_VERSION" + return + fi + + if [[ -t 0 && "$OPENSHELL_VERSION_USER_SUPPLIED" == "true" ]]; then + echo "OpenShell version '$OPENSHELL_VERSION' was not found between $MIN_OPENSHELL_VERSION and $OPENSHELL_LATEST_VERSION." + choose_openshell_version_interactively + echo "OpenShell version selected: $OPENSHELL_VERSION" + return + fi + fail "OpenShell version '$OPENSHELL_VERSION' was not found between $MIN_OPENSHELL_VERSION and $OPENSHELL_LATEST_VERSION." + fi + + if [[ -t 0 ]]; then + choose_openshell_version_interactively + else + OPENSHELL_VERSION="$DEFAULT_OPENSHELL_VERSION" + if ! version_is_available "$OPENSHELL_VERSION"; then + fail "Default OpenShell version '$OPENSHELL_VERSION' was not found on PyPI." + fi + fi + echo "OpenShell version selected: $OPENSHELL_VERSION" +} + +ensure_aiq_env() { + cd "$REPO_ROOT" + if [[ ! -d "$VENV_DIR" ]]; then + log "Creating AI-Q virtual environment" + ./scripts/setup.sh + else + log "Using existing AI-Q virtual environment" + fi +} + +install_openshell_python() { + log "Installing OpenShell Python package exactly: openshell==$OPENSHELL_VERSION" + uv pip install "openshell==$OPENSHELL_VERSION" + + local adapter_install_spec="$LANGCHAIN_NVIDIA_REPO" + local editable_args=() + if [[ -d "$LANGCHAIN_NVIDIA_REPO" ]]; then + if [[ -f "$LANGCHAIN_NVIDIA_REPO/libs/openshell/pyproject.toml" ]]; then + adapter_install_spec="$LANGCHAIN_NVIDIA_REPO/libs/openshell" + elif [[ -f "$LANGCHAIN_NVIDIA_REPO/pyproject.toml" ]]; then + adapter_install_spec="$LANGCHAIN_NVIDIA_REPO" + else + fail "langchain-nvidia-openshell package not found in local checkout: $LANGCHAIN_NVIDIA_REPO" + fi + editable_args=(-e) + fi + + log "Installing langchain-nvidia-openshell adapter: $adapter_install_spec" + # NOTE: expand editable_args only when non-empty; macOS bash 3.2 errors on + # "${arr[@]}" for an empty array under `set -u`. + local adapter_install_args=() + if [[ ${#editable_args[@]} -eq 0 ]]; then + adapter_install_args=(--reinstall-package langchain-nvidia-openshell) + else + adapter_install_args=("${editable_args[@]}") + fi + if ! uv pip install "${adapter_install_args[@]}" "$adapter_install_spec"; then + cat <=0.6.5). The adapter's + # code only uses the stable deepagents BaseSandbox/protocol surface (the same imports + # AI-Q's own sandbox package uses on 0.6.x), so reasserting the floor AI-Q needs is + # safe. This is the OpenShell setup script, so keeping AI-Q runnable is the goal. + log "Reasserting deepagents>=0.6.5 (AI-Q runtime floor) after adapter install" + uv pip install "deepagents>=0.6.5" + + local installed + installed="$("$VENV_DIR/bin/python" - <<'PY' +import openshell +print(getattr(openshell, "__version__", "unknown")) +PY +)" + # The adapter pins openshell>=0.0.68 and may upgrade the package above the requested + # version during its own install; only a version BELOW the requested floor is an error + # (an exact-match check would spuriously fail on that allowed adapter-driven upgrade). + if ! "$VENV_DIR/bin/python" - "$OPENSHELL_VERSION" "$installed" <<'PY' +import sys + + +def parts(v): + return tuple(int(p) for p in v.split(".")[:3] if p.isdigit()) + + +sys.exit(0 if parts(sys.argv[2]) >= parts(sys.argv[1]) else 1) +PY + then + fail "Installed openshell $installed is older than the requested floor $OPENSHELL_VERSION" + fi + "$VENV_DIR/bin/python" - <<'PY' +import langchain_nvidia_openshell # noqa: F401 +PY + echo "OpenShell Python package verified: $installed" + echo "langchain-nvidia-openshell adapter verified" +} + +resolve_openshell_cli() { + local candidates=( + "$OPENSHELL_BIN" + "$VENV_DIR/bin/openshell" + "$(command -v openshell || true)" + "$HOME/.local/bin/openshell" + ) + local candidate + for candidate in "${candidates[@]}"; do + if [[ -n "$candidate" && -x "$candidate" ]]; then + OPENSHELL_BIN="$candidate" + echo "OpenShell CLI: $OPENSHELL_BIN ($("$OPENSHELL_BIN" --version 2>/dev/null || true))" + return + fi + done + fail "OpenShell CLI was not found after installing openshell==$OPENSHELL_VERSION" +} + +resolve_docker() { + log "Resolving Docker CLI" + local candidates=( + "$DOCKER_BIN" + "$(command -v docker || true)" + "/opt/homebrew/bin/docker" + "/usr/local/bin/docker" + "/usr/bin/docker" + ) + local candidate + for candidate in "${candidates[@]}"; do + if [[ -n "$candidate" && -x "$candidate" ]]; then + DOCKER_BIN="$candidate" + echo "Docker CLI: $DOCKER_BIN" + return + fi + done + + if [[ "$OS_NAME" == "macos" ]]; then + candidate="$(find /opt/homebrew/Cellar/docker /usr/local/Cellar/docker -path '*/bin/docker' -type f 2>/dev/null | sort | tail -1 || true)" + if [[ -n "$candidate" && -x "$candidate" ]]; then + DOCKER_BIN="$candidate" + echo "Docker CLI: $DOCKER_BIN" + return + fi + cat <<'EOF' +Docker CLI was not found. + +Install or link Docker CLI, then rerun: + + brew install docker + +If Docker is installed but not on PATH, rerun with: + + scripts/setup_openshell.sh --docker-bin /path/to/docker +EOF + else + cat <<'EOF' +Docker CLI was not found. + +Install Docker for your Linux distribution, then rerun. For example: + + sudo apt-get update + sudo apt-get install -y docker.io + +If Docker is installed but not on PATH, rerun with: + + scripts/setup_openshell.sh --docker-bin /path/to/docker +EOF + fi + exit 1 +} + +configure_docker_host() { + if [[ -n "${DOCKER_HOST:-}" ]]; then + echo "DOCKER_HOST already set: $DOCKER_HOST" + return + fi + local colima_openshell="$HOME/.colima/openshell/docker.sock" + local colima_default="$HOME/.colima/default/docker.sock" + if [[ -S "$colima_openshell" ]]; then + export DOCKER_HOST="unix://$colima_openshell" + echo "DOCKER_HOST: $DOCKER_HOST" + elif [[ -S "$colima_default" ]]; then + export DOCKER_HOST="unix://$colima_default" + echo "DOCKER_HOST: $DOCKER_HOST" + fi +} + +verify_docker_runtime() { + log "Verifying Docker runtime" + if "$DOCKER_BIN" info >/dev/null 2>&1; then + echo "Docker daemon is reachable" + return + fi + + if [[ -n "${DOCKER_HOST:-}" ]]; then + cat </dev/null 2>&1 || true + sleep 2 + rm -f /tmp/aiq-openshell-gateway.log + if [[ "$GATEWAY_DISABLE_TLS" == "true" ]]; then + nohup env OPENSHELL_SERVER_PORT="$GATEWAY_PORT" \ + OPENSHELL_DRIVERS=docker \ + DOCKER_HOST="${DOCKER_HOST:-}" \ + "$OPENSHELL_GATEWAY_LAUNCH_BIN" --disable-tls >/tmp/aiq-openshell-gateway.log 2>&1 & + else + nohup env OPENSHELL_SERVER_PORT="$GATEWAY_PORT" \ + OPENSHELL_DRIVERS=docker \ + DOCKER_HOST="${DOCKER_HOST:-}" \ + "$OPENSHELL_GATEWAY_LAUNCH_BIN" >/tmp/aiq-openshell-gateway.log 2>&1 & + fi + "$OPENSHELL_BIN" gateway remove "$GATEWAY_NAME" >/dev/null 2>&1 || true + if [[ "$GATEWAY_DISABLE_TLS" == "true" ]]; then + "$OPENSHELL_BIN" gateway add --name "$GATEWAY_NAME" "$GATEWAY_ENDPOINT" --local + else + "$OPENSHELL_BIN" gateway add --name "$GATEWAY_NAME" "$GATEWAY_ENDPOINT" --local --gateway-insecure + fi + "$OPENSHELL_BIN" gateway select "$GATEWAY_NAME" + fi + + local attempt + for attempt in $(seq 1 60); do + if "$OPENSHELL_BIN" status; then + return + fi + sleep 1 + done + + echo "OpenShell gateway log:" + LC_ALL=C tr -d '\000' /dev/null | sed -n '1,220p' || true + fail "OpenShell gateway did not become reachable" +} + +print_policy_menu() { + cat <<'EOF' + +Choose an OpenShell sandbox network policy: + + 1. Offline (recommended) + No sandbox network access. AI-Q tools gather data; sandbox code computes on inputs. + + 2. Research APIs + Allow GitHub, NVIDIA API, Tavily, and Serper. + + 3. Python packages + Allow GitHub and PyPI for package metadata/download checks. + + 4. AI development + Allow GitHub, PyPI, NVIDIA API, Tavily, Serper, Hugging Face, arXiv, + Semantic Scholar, and npm. + + 5. Custom + Type a comma-separated allowlist. + +EOF +} + +choose_policy_interactively() { + if [[ -n "$POLICY_PRESET" || -n "$POLICY_ALLOWLIST" ]]; then + return + fi + if [[ ! -t 0 ]]; then + POLICY_PRESET=offline + return + fi + + print_policy_menu + local choice + while true; do + read -r -p "Policy choice [1]: " choice + choice="${choice:-1}" + case "$choice" in + 1|offline) + POLICY_PRESET=offline + return + ;; + 2|research) + POLICY_PRESET=research + return + ;; + 3|python|python-packages) + POLICY_PRESET=python-packages + return + ;; + 4|ai|ai-dev) + POLICY_PRESET=ai-dev + return + ;; + 5|custom) + POLICY_PRESET=custom + read -r -p "Allow services ($SUPPORTED_SERVICES): " POLICY_ALLOWLIST + return + ;; + *) + echo "Choose 1, 2, 3, 4, or 5." + ;; + esac + done +} + +resolve_policy() { + choose_policy_interactively + POLICY_PRESET="$(echo "${POLICY_PRESET:-}" | tr '[:upper:]' '[:lower:]' | tr -d ' ')" + POLICY_ALLOWLIST="$(echo "${POLICY_ALLOWLIST:-}" | tr '[:upper:]' '[:lower:]' | tr -d ' ')" + + if [[ -z "$POLICY_PRESET" && -n "$POLICY_ALLOWLIST" ]]; then + POLICY_PRESET=custom + fi + if [[ -z "$POLICY_PRESET" ]]; then + POLICY_PRESET=offline + fi + if [[ "$POLICY_PRESET" == "none" ]]; then + POLICY_PRESET=offline + fi + + case "$POLICY_PRESET" in + offline) + POLICY_ALLOWLIST=offline + ;; + research) + POLICY_ALLOWLIST=github,nvidia,tavily,serper + ;; + python-packages) + POLICY_ALLOWLIST=github,pypi + ;; + ai-dev) + POLICY_ALLOWLIST=github,pypi,nvidia,tavily,serper,huggingface,arxiv,semantic-scholar,npm + ;; + custom) + if [[ -z "$POLICY_ALLOWLIST" ]]; then + fail "--policy custom requires --allow, for example: --policy custom --allow github,pypi" + fi + ;; + *) + fail "Unsupported policy '$POLICY_PRESET'. Choices: $SUPPORTED_POLICIES" + ;; + esac + + if [[ "$POLICY_ALLOWLIST" == *offline* && "$POLICY_ALLOWLIST" != "offline" ]]; then + fail "Use either offline or a service allowlist, not both: $POLICY_ALLOWLIST" + fi +} + +validate_service() { + case "$1" in + offline|github|pypi|nvidia|tavily|serper|huggingface|arxiv|semantic-scholar|npm) + ;; + *) + fail "Unsupported policy service '$1'. Supported: $SUPPORTED_SERVICES" + ;; + esac +} + +emit_policy_header() { + cat >"$POLICY_FILE" <<'EOF' +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Generated by scripts/setup_openshell.sh. + +version: 1 + +filesystem_policy: + include_workdir: true + read_only: + - /usr + - /lib + - /etc + - /app + - /var/log + - /proc/self + - /dev/urandom + read_write: + - /sandbox + - /workspace + - /tmp + - /dev/null + +# best_effort lets the sandbox start on hosts without Landlock (e.g. Docker Desktop on +# macOS), but on those hosts filesystem confinement is silently dropped. This is acceptable +# only for the local single-operator demo; production must use `hard_requirement` so a +# missing LSM fails closed instead of running unconfined. +landlock: + compatibility: best_effort + +process: + run_as_user: sandbox + run_as_group: sandbox + +EOF +} + +emit_policy_entry() { + local key="$1" + local name="$2" + shift 2 + { + echo " $key:" + echo " name: $name" + echo " endpoints:" + local host + for host in "$@"; do + echo " - host: $host" + echo " port: 443" + echo " protocol: rest" + echo " enforcement: enforce" + echo " access: read-only" + done + echo " binaries:" + echo " - { path: /usr/bin/curl }" + echo " - { path: /usr/local/bin/python3 }" + echo " - { path: /usr/local/bin/python }" + echo " - { path: /usr/local/bin/pip }" + echo " - { path: /usr/local/bin/pip3 }" + } >>"$POLICY_FILE" +} + +emit_policy_service() { + case "$1" in + github) + emit_policy_entry github github-readonly api.github.com github.com + ;; + pypi) + emit_policy_entry pypi pypi-readonly pypi.org files.pythonhosted.org + ;; + nvidia) + emit_policy_entry nvidia nvidia-api-readonly integrate.api.nvidia.com + ;; + tavily) + emit_policy_entry tavily tavily-api-readonly api.tavily.com + ;; + serper) + emit_policy_entry serper serper-api-readonly google.serper.dev + ;; + huggingface) + emit_policy_entry huggingface huggingface-readonly huggingface.co cdn-lfs.huggingface.co + ;; + arxiv) + emit_policy_entry arxiv arxiv-readonly export.arxiv.org + ;; + semantic-scholar) + emit_policy_entry semantic_scholar semantic-scholar-readonly api.semanticscholar.org + ;; + npm) + emit_policy_entry npm npm-readonly registry.npmjs.org + ;; + esac +} + +generate_policy() { + log "Generating OpenShell policy" + resolve_policy + mkdir -p "$(dirname "$POLICY_FILE")" + emit_policy_header + if [[ "$POLICY_ALLOWLIST" == "offline" ]]; then + echo "network_policies: {}" >>"$POLICY_FILE" + else + echo "network_policies:" >>"$POLICY_FILE" + IFS=',' read -r -a services <<<"$POLICY_ALLOWLIST" + local service + for service in "${services[@]}"; do + validate_service "$service" + emit_policy_service "$service" + done + fi + echo "Policy file: $POLICY_FILE" + echo "Policy: $POLICY_PRESET" + echo "Allowed services: $POLICY_ALLOWLIST" +} + +build_image() { + if [[ "$BUILD_IMAGE" != "true" ]]; then + return + fi + log "Building sandbox image: $IMAGE_NAME (sandbox log level: $SANDBOX_LOG_LEVEL)" + "$DOCKER_BIN" build -t "$IMAGE_NAME" \ + --build-arg OPENSHELL_SANDBOX_LOG_LEVEL="$SANDBOX_LOG_LEVEL" \ + -f "$REPO_ROOT/deploy/openshell/Dockerfile.aiq-demo" "$REPO_ROOT/deploy/openshell" +} + +create_sandbox() { + if [[ "$CREATE_SANDBOX" != "true" ]]; then + return + fi + log "Creating named OpenShell sandbox: $SANDBOX_NAME" + "$OPENSHELL_BIN" sandbox delete "$SANDBOX_NAME" >/dev/null 2>&1 || true + local create_log="/tmp/${SANDBOX_NAME}-openshell-create.log" + local policy_label="${POLICY_PRESET//,/_}" + rm -f "$create_log" + "$OPENSHELL_BIN" sandbox create \ + --name "$SANDBOX_NAME" \ + --from "$IMAGE_NAME" \ + --policy "$POLICY_FILE" \ + --label aiq=openshell \ + --label aiq-policy="$policy_label" \ + --no-tty \ + -- sleep infinity >"$create_log" 2>&1 & + + local attempt + for attempt in $(seq 1 120); do + if "$OPENSHELL_BIN" sandbox list | grep -F "$SANDBOX_NAME" | grep -F "Ready" >/dev/null 2>&1; then + "$OPENSHELL_BIN" sandbox list | grep -F "$SANDBOX_NAME" || true + return + fi + sleep 1 + done + + echo "Sandbox create log:" + sed -n '1,220p' "$create_log" || true + fail "Timed out waiting for sandbox '$SANDBOX_NAME' to become Ready" +} + +print_next_steps() { + cat < Use `status` to inspect job status and saved artifacts. Use `report` when the job has already finished and you only need the final output. Use `research_poll` to keep waiting for completion. +The final report may reference generated artifacts (charts, CSVs) as `artifact://` links. To materialize them as local +files, run `python3 $SKILL_DIR/scripts/aiq.py artifacts --download-dir ./aiq-artifacts`; it downloads each artifact +and prints the local path. Do not expect base64 image data in the report itself. + +For a self-contained, shareable report, run `python3 $SKILL_DIR/scripts/aiq.py report --out-dir ./my-report`. It writes `report.md` plus an +`artifacts/` folder and rewrites each `artifact://` link to the matching local file, so the report renders (charts and +all) in any markdown viewer without a running backend. + ### Step 5 - Present the report When `research_poll` completes successfully, fetch and present the full report. Keep citations and source URLs intact. @@ -237,7 +245,8 @@ If your Blueprint version is not compatible: | `scripts/aiq.py research_poll` | Resume polling an existing async job | `` | | `scripts/aiq.py status` | Fetch job status plus `/state` artifacts | `` | | `scripts/aiq.py state` | Fetch event-store artifacts only | `` | -| `scripts/aiq.py report` | Fetch the final report for a completed job | `` | +| `scripts/aiq.py report` | Fetch the final report; with `--out-dir DIR`, export a portable `report.md` + `artifacts/` folder with links rewritten to local files | ` [--out-dir DIR]` | +| `scripts/aiq.py artifacts` | List durable artifacts; with `--download-dir DIR`, download them and print local paths | ` [--download-dir DIR]` | | `scripts/aiq.py stream` | Stream SSE events from a job | `` | | `scripts/aiq.py cancel` | Cancel a running job | `` | diff --git a/skills/aiq-research/scripts/aiq.py b/skills/aiq-research/scripts/aiq.py index 61e1ab887..8a7d75cb2 100644 --- a/skills/aiq-research/scripts/aiq.py +++ b/skills/aiq-research/scripts/aiq.py @@ -233,6 +233,42 @@ def get_report(job_id: str) -> dict[str, Any]: return _api_request("GET", f"/v1/jobs/async/job/{_validate_job_id(job_id)}/report") +_ARTIFACT_ID_RE = re.compile(r"^art_[0-9a-f]{32}$") + + +def list_artifacts(job_id: str) -> dict[str, Any]: + """Fetch durable artifact metadata (no bytes) for an async AI-Q job.""" + return _api_request("GET", f"/v1/jobs/async/job/{_validate_job_id(job_id)}/artifacts") + + +def _binary_request(path: str, *, timeout: int = DEFAULT_API_TIMEOUT_SECONDS) -> bytes: + """Fetch raw bytes from an AI-Q endpoint (artifact content).""" + _validate_api_path(path) + url = f"{_validate_base_url(AIQ_SERVER_URL)}{path}" + req = urllib.request.Request(url, method="GET") + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + return resp.read() + except urllib.error.HTTPError as exc: + error_body = exc.read().decode("utf-8", errors="replace") + print(f"HTTP {exc.code}: {error_body[:ERROR_BODY_PREVIEW_CHARS]}", file=sys.stderr) + raise RuntimeError(f"HTTP {exc.code}") from exc + except urllib.error.URLError as exc: + print(f"Connection failed for {url}: {exc.reason}", file=sys.stderr) + raise RuntimeError(f"Connection failed: {exc.reason}") from exc + + +def download_artifact(job_id: str, artifact_id: str) -> bytes: + """Download a single artifact's bytes by id.""" + valid_job = _validate_job_id(job_id) + if not _ARTIFACT_ID_RE.fullmatch(artifact_id): + raise RuntimeError("artifact_id must look like art_") + return _binary_request( + f"/v1/jobs/async/job/{valid_job}/artifacts/{artifact_id}/content", + timeout=DEFAULT_LONG_HTTP_TIMEOUT_SECONDS, + ) + + def cancel_job(job_id: str) -> dict[str, Any]: """Request cancellation for a running async AI-Q job.""" return _api_request("POST", f"/v1/jobs/async/job/{_validate_job_id(job_id)}/cancel") @@ -321,7 +357,8 @@ def _print_usage() -> None: print(" status Job status plus /state artifacts") print(" state Event-store artifacts for one async job") print(" stream Stream SSE events from an async job") - print(" report Get final report from an async job") + print(" report [--out-dir DIR] Get final report (or export a portable report.md + artifacts/ folder)") + print(" artifacts [--download-dir DIR] List or download durable artifacts") print(" research [agent_type] Submit async job, poll, and return report") print(" research_poll Resume polling an existing async job") print(" cancel Cancel a running async job") @@ -391,11 +428,112 @@ def _command_stream(args: list[str]) -> None: stream_job(job_id) +OUT_DIR_FLAG = "--out-dir" +# Matches a markdown image whose target is an artifact:// reference, capturing the +# "![alt](" prefix, the reference body, and the ")" suffix so only the URL is rewritten. +_REPORT_ARTIFACT_RE = re.compile(r"(!\[[^\]]*\]\()artifact://([^)]+)(\))") + + +def _unique_filename(directory: str, filename: str, used: set[str]) -> str: + """Return a filename that doesn't collide within ``directory``/``used`` (suffix -1, -2, ...).""" + if filename not in used and not os.path.exists(os.path.join(directory, filename)): + used.add(filename) + return filename + stem, ext = os.path.splitext(filename) + i = 1 + while True: + candidate = f"{stem}-{i}{ext}" + if candidate not in used and not os.path.exists(os.path.join(directory, candidate)): + used.add(candidate) + return candidate + i += 1 + + +def _export_report_bundle(job_id: str, out_dir: str) -> None: + """Write a portable report folder: ``report.md`` plus an ``artifacts/`` directory. + + Downloads every durable artifact for the job and rewrites each + ``![caption](artifact://)`` reference to the matching local file so the report + renders in any markdown viewer without a running backend. + """ + report = get_report(job_id).get("report") + if not report: + print(f"No report available for job {job_id}", file=sys.stderr) + sys.exit(EXIT_FAILURE) + + artifacts_dir = os.path.join(out_dir, "artifacts") + os.makedirs(artifacts_dir, exist_ok=True) + + id_to_relpath: dict[str, str] = {} + saved: list[dict[str, Any]] = [] + used: set[str] = set() + for artifact in list_artifacts(job_id).get("artifacts", []): + artifact_id = artifact.get("artifact_id", "") + if not artifact_id: + continue + filename = _unique_filename(artifacts_dir, os.path.basename(artifact.get("filename") or artifact_id), used) + data = download_artifact(job_id, artifact_id) + dest = os.path.join(artifacts_dir, filename) + with open(dest, "wb") as handle: + handle.write(data) + id_to_relpath[artifact_id] = f"artifacts/{filename}" + saved.append({"artifact_id": artifact_id, "path": dest, "size_bytes": len(data)}) + + def _rewrite(match: re.Match[str]) -> str: + relpath = id_to_relpath.get(match.group(2).strip()) + return f"{match.group(1)}{relpath}{match.group(3)}" if relpath else match.group(0) + + report_path = os.path.join(out_dir, "report.md") + with open(report_path, "w", encoding="utf-8") as handle: + handle.write(_REPORT_ARTIFACT_RE.sub(_rewrite, report)) + print(json.dumps({"job_id": job_id, "report": report_path, "artifacts": saved}, indent=JSON_INDENT_SPACES)) + + def _command_report(args: list[str]) -> None: - job_id = _require_arg(args, "Usage: aiq.py report ") + job_id = _require_arg(args, f"Usage: aiq.py report [{OUT_DIR_FLAG} DIR]") + if OUT_DIR_FLAG in args: + flag_index = args.index(OUT_DIR_FLAG) + if flag_index + 1 >= len(args): + print(f"Usage: aiq.py report {OUT_DIR_FLAG} DIR", file=sys.stderr) + sys.exit(EXIT_FAILURE) + _export_report_bundle(job_id, args[flag_index + 1]) + return print(json.dumps(get_report(job_id), indent=JSON_INDENT_SPACES)) +DOWNLOAD_DIR_FLAG = "--download-dir" + + +def _command_artifacts(args: list[str]) -> None: + job_id = _require_arg(args, f"Usage: aiq.py artifacts [{DOWNLOAD_DIR_FLAG} DIR]") + listing = list_artifacts(job_id) + + if DOWNLOAD_DIR_FLAG not in args: + print(json.dumps(listing, indent=JSON_INDENT_SPACES)) + return + + flag_index = args.index(DOWNLOAD_DIR_FLAG) + if flag_index + 1 >= len(args): + print(f"Usage: aiq.py artifacts {DOWNLOAD_DIR_FLAG} DIR", file=sys.stderr) + sys.exit(EXIT_FAILURE) + download_dir = args[flag_index + 1] + os.makedirs(download_dir, exist_ok=True) + + saved: list[dict[str, Any]] = [] + used: set[str] = set() + for artifact in listing.get("artifacts", []): + artifact_id = artifact.get("artifact_id", "") + if not artifact_id: + continue + filename = _unique_filename(download_dir, os.path.basename(artifact.get("filename") or artifact_id), used) + data = download_artifact(job_id, artifact_id) + dest = os.path.join(download_dir, filename) + with open(dest, "wb") as handle: + handle.write(data) + saved.append({"artifact_id": artifact_id, "path": dest, "size_bytes": len(data)}) + print(json.dumps({"job_id": job_id, "downloaded": saved}, indent=JSON_INDENT_SPACES)) + + def _command_research(args: list[str]) -> None: query = _require_arg(args, "Usage: aiq.py research [agent_type]") agent_type = args[OPTIONAL_AGENT_TYPE_POSITION] if len(args) > OPTIONAL_AGENT_TYPE_POSITION else DEFAULT_AGENT_TYPE @@ -464,6 +602,7 @@ def main() -> None: "state": _command_state, "stream": _command_stream, "report": _command_report, + "artifacts": _command_artifacts, "research": _command_research, "research_poll": _command_research_poll, "cancel": _command_cancel, diff --git a/src/aiq_agent/agents/deep_researcher/agent.py b/src/aiq_agent/agents/deep_researcher/agent.py index 94e035713..4384e2565 100644 --- a/src/aiq_agent/agents/deep_researcher/agent.py +++ b/src/aiq_agent/agents/deep_researcher/agent.py @@ -17,8 +17,10 @@ from __future__ import annotations +import asyncio import logging import re +from collections.abc import Callable from collections.abc import Sequence from pathlib import Path from typing import Any @@ -77,6 +79,8 @@ def __init__( skills: DeepResearchSkillsConfig | None = None, sandbox: DeepResearchSandboxConfig | None = None, job_id: str | None = None, + artifact_db_url: str | None = None, + artifact_emit: Callable[[dict[str, Any]], None] | None = None, max_research_concurrency: int = DEFAULT_MAX_RESEARCH_CONCURRENCY, max_concurrent_source_tool_calls: int = DEFAULT_MAX_CONCURRENT_SOURCE_TOOL_CALLS, max_source_tool_batch_size: int = DEFAULT_MAX_SOURCE_TOOL_BATCH_SIZE, @@ -112,7 +116,13 @@ def __init__( self.enable_citation_verification = enable_citation_verification self.job_id = str(job_id) if job_id is not None else str(uuid4()) - self.deepagents_runtime = DeepAgentsRuntime(skills=skills, sandbox=sandbox, job_id=self.job_id) + self.deepagents_runtime = DeepAgentsRuntime( + skills=skills, + sandbox=sandbox, + job_id=self.job_id, + artifact_db_url=artifact_db_url, + artifact_emit=artifact_emit, + ) self._prompts = self._load_prompts() source_tool_names = {tool.name for tool in self.tools} @@ -284,6 +294,27 @@ async def run(self, state: DeepResearchAgentState) -> DeepResearchAgentState: sanitization = sanitize_report(final_message) final_message = sanitization.sanitized_report + # Post-process: harvest sandbox artifacts and resolve artifact:// references so + # generated charts/files render in the report. Inert (manager is None) unless a + # sandbox + artifact_capture + db_url are configured. Blocking I/O off the loop. + manager = self.deepagents_runtime.artifact_manager + if manager is not None: + try: + await asyncio.to_thread(manager.final_harvest) + produced = await asyncio.to_thread(manager.store.list, manager.job_id) + final_message = await asyncio.to_thread(manager.resolve_report_references, final_message, produced) + final_message = await asyncio.to_thread( + manager.ensure_inline_artifacts_embedded, final_message, produced + ) + final_message = await asyncio.to_thread(manager.append_artifact_index, final_message, produced) + except Exception: + # Best-effort: never discard an already verified/sanitized report because + # artifact harvest or embedding failed. final_message stays as-is. + logger.warning( + "Artifact post-processing failed; returning report without embedded artifacts", + exc_info=True, + ) + # Re-emit the verified/sanitized report so the frontend overwrites # the raw version that on_llm_end auto-emitted during ainvoke(). for cb in self.callbacks: diff --git a/src/aiq_agent/agents/deep_researcher/custom_middleware.py b/src/aiq_agent/agents/deep_researcher/custom_middleware.py index 7ae913a1e..ad325ab3f 100644 --- a/src/aiq_agent/agents/deep_researcher/custom_middleware.py +++ b/src/aiq_agent/agents/deep_researcher/custom_middleware.py @@ -96,6 +96,7 @@ class ToolNameSanitizationMiddleware(AgentMiddleware): """ def __init__(self, valid_tool_names: list[str]): + """Store the set of valid tool names used to correct malformed tool calls.""" self.valid_tool_names = set(valid_tool_names) def _sanitize_tool_name(self, name: str) -> str: @@ -184,9 +185,11 @@ class ToolVisibilityMiddleware(AgentMiddleware): """Hide selected tools from model requests without removing scaffolding middleware.""" def __init__(self, hidden_tool_names: set[str]) -> None: + """Store the tool names to hide from model requests.""" self.hidden_tool_names = hidden_tool_names def _filter_tools(self, tools: list[object]) -> list[object]: + """Return the tool list with hidden tools removed.""" if not self.hidden_tool_names: return tools return [tool for tool in tools if _request_tool_name(tool) not in self.hidden_tool_names] @@ -255,6 +258,7 @@ def __init__( backoff_factor: float = 2.0, initial_delay: float = 1.0, ): + """Configure retry count and exponential backoff for failed tool calls.""" self.max_retries = max_retries self.backoff_factor = backoff_factor self.initial_delay = initial_delay @@ -304,6 +308,7 @@ class SourceRegistryMiddleware(AgentMiddleware): """ def __init__(self, source_tool_names: set[str] | None = None) -> None: + """Create a source registry scoped to the given source-producing tool names.""" self.registry = SourceRegistry() self._source_tool_names = source_tool_names or set() self._compact_source_keys: set[str] = set() @@ -519,6 +524,7 @@ class ToolResultPruningMiddleware(AgentMiddleware): """ def __init__(self, keep_last_n: int = 3, max_chars: int = 500): + """Configure how many recent tool results to keep intact and the truncation cap.""" self.keep_last_n = keep_last_n self.max_chars = max_chars diff --git a/src/aiq_agent/agents/deep_researcher/deepagents_runtime.py b/src/aiq_agent/agents/deep_researcher/deepagents_runtime.py index e8d17c229..6fe83f45f 100644 --- a/src/aiq_agent/agents/deep_researcher/deepagents_runtime.py +++ b/src/aiq_agent/agents/deep_researcher/deepagents_runtime.py @@ -19,9 +19,7 @@ import importlib.util import logging -import re -import shlex -import threading +from collections.abc import Callable from pathlib import Path from typing import Any from typing import Literal @@ -30,22 +28,21 @@ from deepagents.backends import CompositeBackend from deepagents.backends import FilesystemBackend from deepagents.backends import StateBackend -from deepagents.backends.protocol import ExecuteResponse -from deepagents.backends.protocol import FileDownloadResponse -from deepagents.backends.protocol import FileUploadResponse -from deepagents.backends.sandbox import BaseSandbox from pydantic import ConfigDict from pydantic import Field from pydantic import field_validator from nat.data_models.function import FunctionBaseConfig +from .sandbox.config import ArtifactCaptureConfig + logger = logging.getLogger(__name__) BUILTIN_SKILLS_DIR = Path(__file__).with_name("skills") BUILTIN_SKILL_SOURCE = "/skills/" SHARED_ROUTE = "/shared/" SKILL_AGENT_NAMES = frozenset({"researcher-agent", "writer-agent"}) +DEFAULT_WORKDIR = "/workspace" class DeepResearchSkillsConfig(FunctionBaseConfig, name="deep_research_skills"): @@ -65,6 +62,7 @@ class DeepResearchSkillsConfig(FunctionBaseConfig, name="deep_research_skills"): @field_validator("agents") @classmethod def _validate_agent_names(cls, value: dict[str, tuple[str, ...]]) -> dict[str, tuple[str, ...]]: + """Reject skill assignments to agent names that are not skill-bearing agents.""" unknown = sorted(set(value) - SKILL_AGENT_NAMES) if unknown: raise ValueError( @@ -78,19 +76,42 @@ class DeepResearchSandboxConfig(FunctionBaseConfig, name="deep_research_sandbox" model_config = ConfigDict(extra="forbid") - provider: Literal["modal"] = Field(default="modal", description="Sandbox backend provider.") + provider: str = Field(default="openshell", description="Sandbox backend provider (resolved by registry).") + # Modal-specific (used when provider == "modal"). app_name: str = Field(default="aiq-deep-research", description="Modal app name for deep research sandboxes") image: str = Field(default="python:3.13-slim", description="Container image for Modal sandboxes") packages: tuple[str, ...] = Field( default=(), description="Python packages to install into the Modal sandbox image.", ) - workdir: str = Field(default="/workspace", description="Working directory inside Modal sandboxes") - timeout: int = Field(default=1200, description="Maximum Modal sandbox lifetime in seconds") - idle_timeout: int = Field(default=1800, description="Modal sandbox idle timeout in seconds") + # OpenShell-specific (used when provider == "openshell"). The named sandbox is created + # out-of-band by scripts/setup_openshell.sh and attached to by name. + sandbox_name: str | None = Field(default=None, description="Existing named OpenShell sandbox to attach to.") + gateway: str | None = Field( + default=None, + description="OpenShell gateway endpoint/name (null uses the locally selected gateway).", + ) + policy: str | None = Field(default=None, description="OpenShell policy file path (requires a named sandbox).") + ready_timeout_seconds: float = Field( + default=300.0, + description="Seconds to wait for the OpenShell sandbox to become ready.", + ) + delete_on_exit: bool = Field(default=False, description="Delete the OpenShell sandbox when the session closes.") + shell: tuple[str, ...] = Field( + default=("bash", "-c"), + description="Shell argv prefix passed to the langchain-nvidia-openshell adapter.", + ) + # Shared across providers. + workdir: str | None = Field(default=None, description="Working directory inside the sandbox") + timeout: int = Field(default=1200, description="Maximum sandbox lifetime in seconds") + idle_timeout: int = Field(default=1800, description="Sandbox idle timeout in seconds") network: Literal["blocked", "enabled"] = Field( default="blocked", - description="Outbound network policy for Modal sandboxes.", + description="Outbound network policy for the sandbox.", + ) + artifact_capture: ArtifactCaptureConfig = Field( + default_factory=ArtifactCaptureConfig, + description="Durable harvesting of generated artifacts (charts/CSVs). Disabled by default.", ) @property @@ -108,11 +129,24 @@ def __init__( skills: DeepResearchSkillsConfig | None = None, sandbox: DeepResearchSandboxConfig | None = None, job_id: str | None = None, + artifact_db_url: str | None = None, + artifact_emit: Callable[[dict[str, Any]], None] | None = None, ) -> None: + """Resolve skill sources and eagerly build the sandbox provider/artifact manager. + + Args: + skills: Skill collections assigned per agent, or None for no skills. + sandbox: Sandbox configuration, or None to run without a sandbox. + job_id: Owning job id; a random id is generated when omitted. + artifact_db_url: Database URL for durable artifact storage. + artifact_emit: Optional SSE emitter for artifact events. + """ self._skills = skills self._sandbox = sandbox self._job_id = str(job_id) if job_id is not None else str(uuid4()) self._backend: Any | None = None + self._sandbox_provider: Any | None = None + self.artifact_manager: Any | None = None self._skill_sources_by_agent = _resolve_agent_skill_sources(skills) self._skill_sources = tuple( dict.fromkeys(source for sources in self._skill_sources_by_agent.values() for source in sources) @@ -121,6 +155,17 @@ def __init__( skills=skills, sandbox=sandbox, ) + # Build the provider-neutral sandbox provider eagerly (lazy SDK session) so the + # runtime can expose its job-scoped workdir/artifact_dir and own its lifecycle. + if sandbox is not None: + self._sandbox_provider = _create_sandbox_backend(sandbox, self._job_id) + self.artifact_manager = _maybe_build_artifact_manager( + provider=self._sandbox_provider, + job_id=self._job_id, + artifact_dir=self.artifact_dir, + artifact_db_url=artifact_db_url, + artifact_emit=artifact_emit, + ) @property def execution_enabled(self) -> bool: @@ -137,36 +182,104 @@ def skill_sources_for(self, agent_name: str) -> list[str] | None: sources = self._skill_sources_by_agent.get(agent_name) return list(sources) if sources else None + @property + def workdir(self) -> str: + """Job-scoped sandbox working directory (or the default when no sandbox).""" + if self._sandbox_provider is None: + return DEFAULT_WORKDIR + return self._sandbox_provider.workdir + + @property + def artifact_dir(self) -> str: + """Job-scoped sandbox artifact directory (harvest root) or the default.""" + if self._sandbox_provider is None: + return f"{DEFAULT_WORKDIR}/aiq-artifacts" + return self._sandbox_provider.artifact_dir + @property def backend(self) -> Any: """Return the concrete backend instance passed to DeepAgents.""" if self._backend is None: self._backend = _build_backend( - sandbox=self._sandbox, - job_id=self._job_id, + provider=self._sandbox_provider, skills_enabled=self.skills_enabled, ) return self._backend + def final_harvest(self) -> None: + """Best-effort final artifact harvest before cleanup (terminal job path).""" + manager = self.artifact_manager + if manager is None: + return + try: + manager.final_harvest() + except Exception: # noqa: BLE001 - harvest is best-effort on the terminal path + logger.warning("Final artifact harvest failed for job %s", self._job_id, exc_info=True) + + def close(self) -> None: + """Release the sandbox provider on a normal terminal job path (idempotent).""" + provider = self._sandbox_provider + if provider is not None and hasattr(provider, "close"): + provider.close() + + def terminate(self) -> None: + """Forcibly stop the sandbox on an interrupted job (cancel/timeout), idempotent.""" + provider = self._sandbox_provider + if provider is not None and hasattr(provider, "terminate"): + provider.terminate() + def _build_backend( *, - sandbox: DeepResearchSandboxConfig | None, - job_id: str, + provider: Any | None, skills_enabled: bool, ) -> Any: - """Build the smallest stock DeepAgents backend needed for this run.""" - default = _create_sandbox_backend(sandbox, job_id) if sandbox is not None else StateBackend() + """Build the smallest stock DeepAgents backend needed for this run. + + The sandbox provider (a ``BaseSandbox``) is created once by the runtime so it can + own the artifact manager and lifecycle; here it is simply used as the default backend. + """ + default = provider if provider is not None else StateBackend() routes: dict[str, Any] = {} if skills_enabled: routes[BUILTIN_SKILL_SOURCE] = _skills_backend() - if sandbox is not None: + if provider is not None: routes[SHARED_ROUTE] = StateBackend() if not routes: return default return CompositeBackend(default=default, routes=routes) +def _maybe_build_artifact_manager( + *, + provider: Any | None, + job_id: str, + artifact_dir: str, + artifact_db_url: str | None, + artifact_emit: Callable[[dict[str, Any]], None] | None, +) -> Any | None: + """Build an ArtifactManager only when capture is enabled and a store URL is provided. + + Defaults to ``None`` (no harvesting) so adding the sandbox alone never requires a DB. + """ + if provider is None or artifact_db_url is None: + return None + capture = getattr(getattr(provider, "config", None), "artifact_capture", None) + if capture is None or not getattr(capture, "enabled", False): + return None + from .sandbox.artifacts import ArtifactManager + from .sandbox.artifacts import SqlArtifactStore + + return ArtifactManager( + job_id=job_id, + backend=provider, + store=SqlArtifactStore(artifact_db_url), + config=capture, + artifact_dir=artifact_dir, + emit=artifact_emit, + ) + + def _skills_backend() -> FilesystemBackend: """Return the filesystem-backed built-in skills route.""" return FilesystemBackend(root_dir=BUILTIN_SKILLS_DIR.resolve(), virtual_mode=True) @@ -197,6 +310,7 @@ def resolve_skill_collections(collection_names: tuple[str, ...]) -> tuple[str, . def _resolve_agent_skill_sources(skills: DeepResearchSkillsConfig | None) -> dict[str, tuple[str, ...]]: + """Map each agent to its resolved skill source paths, skipping empty assignments.""" if skills is None: return {} return { @@ -211,6 +325,7 @@ def _validate_sandbox_requirements( skills: DeepResearchSkillsConfig | None, sandbox: DeepResearchSandboxConfig | None, ) -> None: + """Fail fast when a skill collection that requires a sandbox is assigned without one.""" if skills is None or not skills.require_sandbox: return @@ -231,27 +346,58 @@ def _validate_sandbox_requirements( ) -def _validate_modal_sandbox_name(job_id: str) -> str: - if len(job_id) > 64 or re.match(r"^[a-zA-Z0-9-_.]+$", job_id) is None or re.match(r"^ap-[a-zA-Z0-9]{22}$", job_id): - raise ValueError( - "Deep research job_id must be a valid Modal sandbox name: " - "64 characters or fewer, using only alphanumeric characters, dashes, periods, and underscores." - ) - return job_id - - def _create_sandbox_backend(config: DeepResearchSandboxConfig, job_id: str) -> Any: - if config.provider == "modal": - return _create_modal_backend(config, job_id) - raise ValueError(f"Unsupported sandbox provider: {config.provider}. Supported providers: modal") - - -def _create_modal_backend(config: DeepResearchSandboxConfig, job_id: str) -> Any: - _ensure_modal_dependencies() - return _LazyModalSandboxBackend(config, job_id) + """Resolve the AI-Q sandbox config to a provider-neutral sandbox backend. + + Keeps the Modal dependency pre-check (clear early error when Modal is configured but + not installed), then maps the config to the provider-neutral ``SandboxConfig`` and + dispatches through the sandbox provider registry. + """ + from .sandbox import create_sandbox_backend as registry_create + from .sandbox.config import SandboxConfig as ProviderSandboxConfig + + provider = config.provider.lower() + workdir = config.workdir or ("/workspace" if provider == "modal" else "/sandbox") + + if provider == "modal": + _ensure_modal_dependencies() + providers = { + "modal": { + "app_name": config.app_name, + "image": config.image, + "python_packages": config.packages, + } + } + elif provider == "openshell": + providers = { + "openshell": { + "gateway": config.gateway, + "sandbox_name": config.sandbox_name, + "policy": config.policy, + "ready_timeout_seconds": config.ready_timeout_seconds, + "delete_on_exit": config.delete_on_exit, + "shell": list(config.shell), + } + } + else: + providers = {} + + provider_config = ProviderSandboxConfig.model_validate( + { + "provider": provider, + "workdir": workdir, + "timeout": config.timeout, + "idle_timeout": config.idle_timeout, + "network": {"mode": "blocked" if config.block_network else "open"}, + "artifact_capture": config.artifact_capture.model_dump(), + "providers": providers, + } + ) + return registry_create(provider_config, job_id) def _ensure_modal_dependencies() -> None: + """Raise ImportError listing any missing Modal packages when Modal is configured.""" missing = [ package for module_name, package in (("modal", "modal"), ("langchain_modal", "langchain-modal")) @@ -263,152 +409,3 @@ def _ensure_modal_dependencies() -> None: "Modal sandbox is configured, but required package(s) are missing: " f"{packages}. Install the Modal sandbox dependencies or remove the sandbox config." ) - - -class _LazyModalSandboxBackend(BaseSandbox): - """Job-scoped Modal backend that creates and recreates the sandbox on demand.""" - - def __init__(self, config: DeepResearchSandboxConfig, job_id: str) -> None: - self.config = config - self.sandbox_name = _validate_modal_sandbox_name(job_id) - self._backend: Any | None = None - self._lock = threading.Lock() - - @property - def id(self) -> str: - backend = self._backend - if backend is None: - return self.sandbox_name - return backend.id - - def execute(self, command: str, *, timeout: int | None = None) -> ExecuteResponse: - for attempt in range(2): - try: - return self._get_backend().execute(command, timeout=timeout) - except Exception as exc: - if attempt == 0 and _is_modal_not_found_error(exc): - logger.warning( - "Modal sandbox %s disappeared during execute; recreating and retrying once", - self.sandbox_name, - ) - self._reset_backend() - continue - raise - raise RuntimeError("unreachable") - - def upload_files(self, files: list[tuple[str, bytes]]) -> list[FileUploadResponse]: - for attempt in range(2): - try: - return self._get_backend().upload_files(files) - except Exception as exc: - if attempt == 0 and _is_modal_not_found_error(exc): - logger.warning( - "Modal sandbox %s disappeared during file upload; recreating and retrying once", - self.sandbox_name, - ) - self._reset_backend() - continue - raise - raise RuntimeError("unreachable") - - def download_files(self, paths: list[str]) -> list[FileDownloadResponse]: - for attempt in range(2): - try: - return self._get_backend().download_files(paths) - except Exception as exc: - if attempt == 0 and _is_modal_not_found_error(exc): - logger.warning( - "Modal sandbox %s disappeared during file download; recreating and retrying once", - self.sandbox_name, - ) - self._reset_backend() - continue - raise - raise RuntimeError("unreachable") - - def _get_backend(self) -> Any: - backend = self._backend - if backend is not None: - return backend - - with self._lock: - if self._backend is None: - logger.info( - "Modal sandbox backend init: sandbox_name=%s app=%s", - self.sandbox_name, - self.config.app_name, - ) - self._backend = _create_modal_backend_now(self.config, self.sandbox_name) - return self._backend - - def _reset_backend(self) -> None: - with self._lock: - logger.warning( - "Modal sandbox backend RESET: sandbox_name=%s app=%s " - "(any uploaded files in the previous sandbox are now lost)", - self.sandbox_name, - self.config.app_name, - ) - self._backend = _create_modal_backend_now(self.config, self.sandbox_name, force_new=True) - - -def _create_modal_backend_now( - config: DeepResearchSandboxConfig, - sandbox_name: str, - *, - force_new: bool = False, -) -> Any: - try: - import modal - from langchain_modal import ModalSandbox - except ImportError as exc: - raise ImportError( - "The Modal sandbox backend requires the `langchain-modal` and `modal` packages. " - "Install the updated AIQ dependencies and run `modal setup` before enabling a Modal sandbox." - ) from exc - - app = modal.App.lookup(name=config.app_name, create_if_missing=True) - if not force_new: - try: - sandbox = modal.Sandbox.from_name(config.app_name, sandbox_name) - logger.info("Modal sandbox attached to existing instance: name=%s", sandbox_name) - return ModalSandbox(sandbox=sandbox) - except modal.exception.NotFoundError: - logger.info("Modal sandbox not found, creating fresh instance: name=%s", sandbox_name) - - image = modal.Image.from_registry(config.image) - if config.packages: - image = image.pip_install(*config.packages) - if config.workdir: - image = image.run_commands(f"mkdir -p {shlex.quote(config.workdir)}") - - try: - sandbox = modal.Sandbox.create( - app=app, - image=image, - workdir=config.workdir, - name=sandbox_name, - timeout=config.timeout, - idle_timeout=config.idle_timeout, - block_network=config.block_network, - ) - logger.info( - "Modal sandbox CREATED: name=%s image=%s workdir=%s timeout=%ds", - sandbox_name, - config.image, - config.workdir, - config.timeout, - ) - except modal.exception.AlreadyExistsError: - sandbox = modal.Sandbox.from_name(config.app_name, sandbox_name) - logger.info("Modal sandbox attached after AlreadyExistsError: name=%s", sandbox_name) - return ModalSandbox(sandbox=sandbox) - - -def _is_modal_not_found_error(exc: Exception) -> bool: - try: - import modal - - return isinstance(exc, modal.exception.NotFoundError) - except ImportError: - return exc.__class__.__name__ == "NotFoundError" and exc.__class__.__module__.startswith("modal") diff --git a/src/aiq_agent/agents/deep_researcher/factory.py b/src/aiq_agent/agents/deep_researcher/factory.py index f7296417a..2167ee0a9 100644 --- a/src/aiq_agent/agents/deep_researcher/factory.py +++ b/src/aiq_agent/agents/deep_researcher/factory.py @@ -125,24 +125,36 @@ class DeepResearchGraphContext: @property def available_documents(self) -> list[dict[str, Any]]: + """Return the user-uploaded documents for this run as serialized dicts.""" return [doc.model_dump() for doc in (self.state.available_documents or [])] def render_prompt(self, prompt_name: str, **values: Any) -> str: + """Render a named prompt template with shared context plus any overrides.""" + prompt_values = { + "current_datetime": self.current_datetime, + "user_info": self.state.user_info, + "available_documents": self.available_documents, + "execution_enabled": self.runtime.execution_enabled, + "skills_enabled": self.runtime.skills_enabled, + "sandbox_workdir": self.runtime.workdir, + "sandbox_artifact_dir": self.runtime.artifact_dir, + **values, + } return render_prompt_template( self.prompts[prompt_name], - current_datetime=self.current_datetime, - user_info=self.state.user_info, - available_documents=self.available_documents, - **values, + **prompt_values, ) def middleware(self, base: Sequence[Any]) -> list[Any]: + """Return the base middleware stack extended with tool-visibility middleware.""" return [*base, *self.visibility_middleware] def permissions(self, agent_name: str) -> list[FilesystemPermission]: + """Return the skill-derived filesystem permissions for an agent.""" return runtime_skill_filesystem_permissions(self.runtime, agent_name) def skill_sources(self, agent_name: str) -> list[str] | None: + """Return the resolved skill source paths for an agent, or None.""" return self.runtime.skill_sources_for(agent_name) @@ -211,6 +223,7 @@ def build_deep_research_middleware_set( """Build researcher, writer, and orchestrator middleware stacks.""" def common(extra_valid_tool_names: Sequence[str] = ()) -> list[Any]: + """Build the shared middleware stack, allowing extra valid tool names.""" return build_common_middleware( tool_set=tool_set, source_registry_middleware=source_registry_middleware, @@ -321,6 +334,7 @@ def _subagent_spec( response_format: Any = None, skills: list[str] | None = None, ) -> dict[str, Any]: + """Assemble a deepagents subagent spec (prompt, model, tools, permissions, middleware).""" spec: dict[str, Any] = { "name": name, "description": description, diff --git a/src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2 b/src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2 index 24667162b..2eea59a96 100644 --- a/src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2 +++ b/src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2 @@ -14,7 +14,9 @@ Research is handled through the `run_research_batch` tool. ## Filesystem Notes - `/shared/` is a virtual path managed by the agent runtime. Inspect it only via `read_file`, `write_file`, `ls`, or `glob`. -{% if execution_enabled %}- Use `/workspace/` for sandbox-local files and copy durable text artifacts back through DeepAgents filesystem tools. Shell commands cannot see `/shared/`. +{% if execution_enabled %}- `{{ sandbox_workdir }}/` is the sandbox real filesystem for scripts and intermediate files; reachable via `execute`/shell. +- `{{ sandbox_artifact_dir }}/` is the sandbox artifact directory. Charts, images, CSVs, and manifest files written here are harvested automatically and can be embedded in the report with `artifact://`. +- Shell commands cannot see `/shared/`. Do not write `/shared/` paths from inside sandbox scripts. {% endif %} - `write_file` cannot overwrite an existing file. To revise, use `edit_file` or write to a new path. @@ -113,6 +115,8 @@ Read /shared/plan.json, all research note files under /shared/, and the verified Before writing, inspect Available Skills. If an applicable writer skill exists, read its SKILL.md first and use it as the controlling synthesis protocol. Do not draft the final answer until that skill has been read. Synthesize a final response from the detailed research context. For broad reports, produce a cross-synthesized narrative with developed paragraphs; do not compress the answer into a checklist of short component summaries unless the user asked for that format. +{% if execution_enabled %}If the report includes generated figures, embed each earned chart once with `![](artifact://)`; never paste sandbox paths or base64 data.{% endif %} +{% if execution_enabled %}Only generate or embed a chart when its data is source-anchored and reasonably complete; otherwise present the table with explicit gaps and state the limitation.{% endif %} Write the final Markdown answer to /shared/output.md. Return only the short completion marker `Wrote /shared/output.md` after /shared/output.md is written. Do not return JSON and do not echo the full Markdown. ``` diff --git a/src/aiq_agent/agents/deep_researcher/prompts/researcher.j2 b/src/aiq_agent/agents/deep_researcher/prompts/researcher.j2 index 2fdcaf099..901c50fe1 100644 --- a/src/aiq_agent/agents/deep_researcher/prompts/researcher.j2 +++ b/src/aiq_agent/agents/deep_researcher/prompts/researcher.j2 @@ -11,6 +11,7 @@ Your final response is validated against the ResearchNotes schema. The batch too 6. **Mandatory skill use**: If the Skills System lists an available skill whose metadata, triggers, or description matches the ResearchQuery, target components, plan context, or any analysis/transformation you are about to perform, you MUST read that skill's `SKILL.md` using `read_file` before doing the matching work. Follow the loaded skill as the controlling procedure, including any tools it instructs you to use. Do not substitute generic reasoning, ad hoc code, source-tool summarization, or memory for a matching available skill. If multiple listed skills match, read and follow the most specific one first. If no listed skill applies, continue with this base researcher prompt. 7. **Stop when you can answer confidently** - Don't keep searching for perfection 8. **Use subqueries selectively** - The user message contains a ResearchQuery JSON object. Use the first `preferred_tools` item as the primary source tool, then use any additional `preferred_tools` in order. Run the main `query` first. If `subqueries` are present, use them as optional focused angles for distinct gaps or facets not covered by the main query. Do not run a subquery when the main query already produced enough evidence for that angle. +9. **Visualization is earned** - For quantitative deliverables, only chart data that is source-anchored and reasonably complete. If a series is mostly undisclosed or mixes metric definitions, produce the table with explicit gaps and note the limitation rather than a misleading chart. ## Guidelines - If possible, cross-reference multiple sources for accuracy @@ -47,7 +48,9 @@ Your output will be used to write a cited final answer. Produce **in-depth, deta ## Tool Use - Use the fewest high-signal source-tool calls needed for the assigned query. {% if execution_enabled %}- `execute` runs shell commands in the sandbox only. +- Each `execute` runs in a fresh shell; `cd` does not persist between calls. Never use a standalone `cd` — use absolute paths in every command (or `cd && ` in one line). - Never pass agent source tool names to `execute`. +- When a skill has you run code, write scripts and intermediate files under `{{ sandbox_workdir }}` and durable artifacts (charts, CSVs, plus their `manifest.json`) under `{{ sandbox_artifact_dir }}` exactly as the skill describes. Both are per-job paths, so use them verbatim and do not write artifacts to any other location (the runtime only harvests `{{ sandbox_artifact_dir }}`). {% endif %} - Default source budget per ResearchQuery: one primary source-tool call, plus at most one fallback or corroboration call. - For broad survey queries, make at most one extra targeted follow-up after the first results, only when needed for target_components. diff --git a/src/aiq_agent/agents/deep_researcher/prompts/writer.j2 b/src/aiq_agent/agents/deep_researcher/prompts/writer.j2 index fea3e169a..bd6cb7495 100644 --- a/src/aiq_agent/agents/deep_researcher/prompts/writer.j2 +++ b/src/aiq_agent/agents/deep_researcher/prompts/writer.j2 @@ -63,7 +63,17 @@ When synthesizing: - Do not expose internal `evidence_judgment` scores or rationale in the final answer unless the user explicitly asks for methodology. - Err on the side of more useful information rather than less, while staying focused on the requested answer shape. -## Citations +{% if execution_enabled %}## Figures + +When `answer_strategy.required_components` calls for a chart, GENERATE it yourself with the chart skill and embed it once, where it is discussed, with `![](artifact://)`. Read the relevant chart `SKILL.md` first, then write the plotting script under `{{ sandbox_workdir }}` and `execute` it so the script writes the image plus its `manifest.json` under `{{ sandbox_artifact_dir }}` (the only directory the runtime harvests). Use those per-job paths verbatim. + +Each `execute` runs in a fresh shell, so `cd` does NOT persist between calls and a standalone `cd` accomplishes nothing — never navigate with `cd`; put absolute paths in every command (or chain in one line as `cd && `). + +NEVER paste the plotting code, a sandbox file path (e.g. `{{ sandbox_artifact_dir }}/chart.png`), or base64 image data into the report as prose - the `artifact://` token is the ONLY way a figure renders. A "plotting code" or "chart plan" section is not a substitute for the actual figure. Precede each embed with a one-sentence description of what it shows. + +A figure is earned: generate or embed one only when its data is source-anchored and reasonably complete. If a quantitative series is mostly undisclosed or mixes metric definitions, present the table (which shows the gaps) and state the limitation in one sentence instead of a misleading chart. + +{% endif %}## Citations - Citations are mandatory for any sourced final answer. Every material source-derived factual claim, number, date, quote, source-specific caveat, or source claim must have an inline citation like `[1]`. - For arithmetic, rankings, ranges, means, medians, or summary statistics computed solely from already cited values in the report, cite the underlying source data where those values appear. Then state that the derived statistics are computed from that cited table or cited values; do not attach source numbers to every computed-statistic line. - In a summary-statistics block, cite only source-derived inputs, source-specific caveats, or any newly introduced external value. Lines such as count, mean, median, range, min, or max do not need their own citation when they are computed from the immediately preceding cited table or cited value set. diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/README.md b/src/aiq_agent/agents/deep_researcher/sandbox/README.md new file mode 100644 index 000000000..3c0ed2b1d --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/README.md @@ -0,0 +1,301 @@ + + +# Deep Research Sandbox + Artifact Runtime + +Provider-neutral sandbox execution for agent-generated code, plus a durable artifact +runtime that harvests generated files (charts, CSVs, notebooks) so they survive the +sandbox and can be served to the UI, embedded in reports, or downloaded via the skill CLI. + +Design pattern: **sandbox-as-tool**. AI-Q keeps the orchestrator, auth, tools, event +store, and report state in-process; only generated code runs remotely. Secrets and +data-source credentials never enter the sandbox. + +## Architecture + +```text +config YAML (sandbox.provider + providers.) + -> SandboxConfig (config.py) + -> registry.create_sandbox_backend --(fail-closed capability gate)--> SandboxProvider + | +DeepAgentsRuntime (deepagents_runtime.py) holds the provider and composes: | + CompositeBackend(default = provider, routes = {/shared/, /skills/ -> StateBackend}) + - workdir (default route): real sandbox FS, reached via execute. The EFFECTIVE + workdir is per-job: / (e.g. /sandbox/), with + artifacts nested at /aiq-artifacts. See "Workspace isolation" below. + - /shared/, /skills/: in-process virtual FS (durable text, never the sandbox) + ArtifactManager (artifacts/manager.py): download_files -> validate -> ArtifactStore -> SSE +``` + +## Workspace organization and isolation limits + +The effective working directory is scoped per job to `/`, and +the artifact directory is nested under it at `/aiq-artifacts`. The provider base +creates these on session start (`_prepare_workspace`, an idempotent `mkdir -p`) and the +runtime injects them into prompts/skills as `sandbox_workdir`/`sandbox_artifact_dir`. This +prevents accidental filename collisions and keeps harvesting scoped to the current job. + +Modal creates a fresh sandbox for each job. The experimental OpenShell configuration instead +attaches jobs to one pre-created named sandbox because the SDK cannot apply the configured +policy when creating an anonymous sandbox. Per-job directories inside that sandbox are not +an access-control boundary: executed code can access sibling job directories allowed by the +shared policy. Use OpenShell only for local, single-operator testing, and do not run mutually +untrusted jobs concurrently. Physical per-job OpenShell isolation and attach-time policy +verification are follow-up work. The default policy also sets `landlock.compatibility: +best_effort`, so on hosts without Landlock (e.g. Docker Desktop on macOS) filesystem +confinement is silently dropped; production must use `hard_requirement` to fail closed. + +The agent only ever sees a `read_file`/`write_file`/`edit_file`/`execute` tool surface +plus `/shared/` for durable text. Binary artifacts are harvested host-side via +`download_files` and referenced in reports as `artifact://` (never base64). + +## Module map + +| File | Purpose | +|---|---| +| `base.py` | `SandboxProvider` ABC. Force only `execute` + `capabilities`; the base owns lazy single-flight creation, a serialization lock, idempotency-gated retry, `close()`, `terminate()`. | +| `registry.py` | `register_sandbox_provider` / `create_sandbox_backend` (config-driven dispatch + capability gate). | +| `config.py` | `SandboxConfig`: common fields + nested `providers.` + `artifact_capture` + `lifecycle_scope`; legacy flat-config shim; provider validated against the registry. | +| `capabilities.py` | `SandboxCapabilities` + `verify_capabilities` (fail-closed: refuse to run if a required guarantee like `block_network` is unsupported). | +| `providers/modal.py` | Modal provider (cloud). Create-fresh semantics (no silent attach-by-name). | +| `providers/openshell.py` | OpenShell provider (enterprise/on-prem). Lazy, ad-hoc deps; policy requires a named sandbox. | +| `artifacts/models.py` | `Artifact` record (id, mime, sha256, size, provenance, status). Metadata only. | +| `artifacts/manifest.py` | `manifest.json` schema + parser. | +| `artifacts/store.py` | `SqlArtifactStore` on the shared job `db_url` (metadata table + capped BLOB; pluggable to S3). | +| `artifacts/manager.py` | Harvest pipeline: manifest-first + scan, path-traversal confinement, MIME-from-bytes, SVG sanitize, render-gate, quotas, dedup, store-then-emit, `artifact://` resolution. | + +## Adding a provider (the whole surface) + +```python +from deepagents.backends.sandbox import BaseSandbox +from ..base import SandboxProvider +from ..capabilities import SandboxCapabilities +from ..registry import register_sandbox_provider + +class MySandboxProvider(SandboxProvider): + provider_name = "mybox" + + @classmethod + def _scoped_name(cls, job_id: str) -> str: + return job_id # apply provider naming rules + + @property + def capabilities(self) -> SandboxCapabilities: + return SandboxCapabilities(supports_network_policy=True, supports_artifact_download=True) + + def is_recoverable_error(self, exc: Exception) -> bool: + return False # classify stale/transient errors for retry + + def _create_session(self) -> BaseSandbox: + # lazy-import your SDK; create/attach a job-scoped sandbox; return a BaseSandbox. + ... + +register_sandbox_provider("mybox", MySandboxProvider) +``` + +Add a `MyProviderConfig` sub-model to `SandboxProvidersConfig` in `config.py`. You do +NOT implement `read_file`/`write_file`/`ls`/`glob` (inherited from `BaseSandbox` on top +of `execute`) or the retry/lock/lifecycle (the base owns those). Every provider must +pass the compliance suite (`tests/.../sandbox/test_provider_compliance.py`). + +### Out-of-tree providers (entry points) + +Third-party packages contribute providers without editing AI-Q by declaring the +`aiq.sandbox_providers` entry-point group (the same plug-in pattern as NAT and +deepagents Code). They are discovered lazily on first registry use: + +```toml +# in the third-party package's pyproject.toml +[project.entry-points."aiq.sandbox_providers"] +mybox = "my_pkg.provider:MySandboxProvider" +``` + +The entry-point name becomes the config `provider` key. A broken plugin is logged and +skipped — it can never break resolution of the built-in providers. + +## Config (sandbox block) + +```yaml +sandbox: + enabled: true + provider: openshell # registry key + workdir: /sandbox # injected into prompts + skills + network: # normalized, provider-neutral egress policy + mode: blocked # blocked | allowlist | open (legacy `block_network: true` => blocked) + # allow: [pypi.org] # required for mode: allowlist; needs supports_network_allowlist + timeout: 1200 + idle_timeout: 1800 + resources: # optional CPU/memory caps; omit for no limit + # cpu: 2 # cores; needs supports_resource_limits (Modal enforces; OpenShell does not) + # memory_mb: 4096 # a requested limit on a provider that can't enforce it fails closed + artifact_capture: + enabled: true # requires supports_artifact_download + max_file_bytes: 50000000 + allow_extensions: [.png, .jpg, .jpeg, .webp, .csv, .json, .md, .ipynb, .pdf] + providers: + modal: + app_name: aiq-deep-research + image: python:3.12-slim + python_packages: [matplotlib, numpy, pandas, pillow, tabulate] + openshell: + gateway: null # null = locally selected gateway + sandbox_name: aiq-openshell-demo + policy: configs/openshell/generated/aiq-openshell-policy.yaml +``` + +The legacy flat shape (top-level `app_name`/`image`/`python_packages`) still loads and +is lifted into `providers.modal`. + +## Artifact runtime + +- Generated code writes binaries + a `manifest.json` to `artifact_dir`. +- Once at the end of the agent run (`agent.run()` -> `ArtifactManager.final_harvest`), + the `ArtifactManager` pulls bytes via `download_files`, runs the validation pipeline + (path-traversal confinement -> extension allowlist -> size cap -> MIME-from-bytes/spoof + reject -> quota -> SVG sanitize -> sha256), stores via `SqlArtifactStore`, then emits an + `artifact` SSE event (`to_sse_payload`, metadata + `content_url`, never bytes). +- Failed or cancelled runs are not harvested in the current implementation. +- Reports reference artifacts as `![caption](artifact://)`; the report + postprocessor rewrites filename refs to durable ids and drops unknown/foreign refs. +- Endpoints: `GET /v1/jobs/async/job/{job_id}/artifacts` and `.../artifacts/{id}/content` + (auth-scoped via `authorize_job_access`). CLI: `python3 skills/aiq-research/scripts/aiq.py artifacts [--download-dir DIR]`. +- Render gate: only PNG/JPEG/WebP may render inline; SVG/notebook/PDF are download-only. +- Transfer guards (artifacts come from an untrusted sandbox): the OpenShell download + bootstrap fails closed BEFORE reading bytes - it rejects symlink escapes (`realpath` + differs from the lexical path: leaf or parent), directories, and files over + `max_file_bytes` - so a hostile sandbox cannot pull an out-of-tree or oversized file + into host memory. The harvest also count-gates before each download and bounds the + directory scan, and decoded bytes are base64-validated. SQL is fully parameterized + and the content endpoint is auth-scoped per job with `nosniff` + RFC 5987 filenames. + +### Report post-processing (host-side, in `agent.run`) + +Run once after the report is produced, reusing a single artifact fetch: +- `resolve_report_references` - rewrite `artifact://` to `artifact://`; drop unknown/foreign refs. +- `ensure_inline_artifacts_embedded` - append a `## Figures` section embedding any produced + inline image the model forgot to reference (so a generated chart always surfaces). +- `append_artifact_index` - append a `## Generated Artifacts` list crediting every harvested + file (charts and their backing CSVs), alongside the external `## Sources`. + +### Rendering surfaces + +The stored report keeps `artifact://`; each surface resolves it at its own edge (one +shared helper, `MarkdownRenderer/artifact-url.ts`, builds the content path): +- **UI report**: `MarkdownRenderer` preserves the `artifact://` scheme via a custom + `urlTransform` (react-markdown would otherwise blank a non-standard scheme), and an `img` + renderer rewrites it to the same-origin `/api/jobs/async/job/{job_id}/artifacts/{id}/content` + (the Next proxy streams bytes through). Job id comes from `selectResolvedDeepResearchJobId`. +- **PDF export**: `/api/generate-pdf` fetches each artifact server-side and inlines it as a + `data:` URI (<= 8 MiB); `ReactPdfDocument` renders raster images as block figures (paragraphs + and list items). Non-image refs are skipped. +- **Markdown download**: `artifact://` is rewritten to an absolute content URL so the `.md` + renders while the backend is reachable. +- **Skill/CLI**: `aiq.py report --out-dir DIR` writes `report.md` plus an + `artifacts/` folder and rewrites links to local files (portable, renders offline). + +## Providers + +### Modal (cloud) + +Requires `modal` + `langchain-modal` (in `pyproject`) and `modal setup`. See +`docs/source/examples/skills-sandbox/index.md`. + +### OpenShell (experimental, local single-operator) + +> OpenShell jobs currently attach to one pre-created named sandbox. Its policy is applied by +> the setup command and is not verified when AI-Q attaches. Job directories prevent accidental +> collisions but do not isolate mutually untrusted jobs. Do not use this path as a multi-tenant +> security boundary. + +Two ad-hoc deps (never in `pyproject`): the `openshell` SDK and the official +`langchain-nvidia-openshell` adapter (`OpenShellSandbox`), the OpenShell partner package in +[`langchain-ai/langchain-nvidia`](https://github.com/langchain-ai/langchain-nvidia/pull/303). +The adapter is published on PyPI as `langchain-nvidia-openshell` — `./scripts/setup_openshell.sh` +installs it for you (override the source with `LANGCHAIN_NVIDIA_REPO` to use a git spec or local +checkout). To install it into your `.venv` manually: + +```bash +uv pip install 'langchain-nvidia-openshell==0.1.0' +``` + +One-command setup: + +```bash +./scripts/setup_openshell.sh --policy offline +./scripts/start_e2e.sh --config_file configs/config_openshell.yml +``` + +The setup script prints the environment variables needed by any later shell that starts +AI-Q. If you start the backend in a different terminal/session, export the printed values +before running `start_e2e.sh` (or put them in your local env file): + +```bash +export AIQ_OPENSHELL_SANDBOX_NAME="aiq-openshell-demo" +export AIQ_OPENSHELL_POLICY_FILE="$PWD/configs/openshell/generated/aiq-openshell-policy.yaml" +``` + +Inference is routed host-side (e.g. NVIDIA Build or an internal inference hub set in the +config); the network-blocked sandbox never sees the key. + +**File-transfer gotcha:** the provider overrides file transfer with an env-free shim that +passes the path via `argv`. OpenShell 0.0.57-0.0.67 strip +`OPENSHELL_`-prefixed env before exec, so the adapter's env-based file transfer silently +fails (masked host-side as `permission_denied`). Set `AIQ_OPENSHELL_ADAPTER_FILE_TRANSFER=1` +to delegate uploads to the official adapter and validate the upstream argv fix +([langchain-nvidia#303](https://github.com/langchain-ai/langchain-nvidia/pull/303)). Downloads +always use AI-Q's bounded shim so realpath confinement and pre-transfer size checks remain +in force. Once the upstream adapter provides equivalent guards, drop the shim and toggle. + +## Operational knobs + +- `AIQ_MAX_SANDBOXES_PER_PRINCIPAL` / `AIQ_MAX_SANDBOXES_GLOBAL` (default-off): submit-path + concurrency/cost caps for sandbox-enabled jobs. +- `AIQ_OPENSHELL_ADAPTER_FILE_TRANSFER` (default-off): route OpenShell uploads through the + official adapter instead of the env-free shim (see OpenShell gotcha above). +- Artifact retention reuses the job-expiry periodic cleanup (`expiry_seconds`). +- In-container OpenShell log verbosity (opt-in): `agent.execute()` calls and their output are + already logged on the AI-Q side (the `execute` tool-call events). To also see what runs + inside the OpenShell container, rebuild the sandbox image with a higher `RUST_LOG`: + `./scripts/setup_openshell.sh --sandbox-log-level debug` (or `--build-arg + OPENSHELL_SANDBOX_LOG_LEVEL=debug`). Default `warn` keeps OpenShell's stock behavior. + Read the container logs with `openshell logs `, the OpenShell TUI, or inside + the sandbox at `/var/log/openshell.*.log` (e.g. `grep "OCSF PROC:"` for process activity). + +## Testing + +```bash +pytest tests/aiq_agent/agents/deep_researcher/sandbox/ -q +``` + +All provider/artifact tests run without a live Modal/OpenShell gateway (OpenShell +compliance auto-skips when the SDK is absent). + +## Troubleshooting + +- **`Input tag 'tavily_web_search' ... does not match` / `Unknown field name front_end`**: + the workspace plugin packages aren't installed. Install them (don't re-run `setup.sh`, + which recreates `.venv`): + `uv pip install -e ./frontends/aiq_api -e ./sources/tavily_web_search -e "./sources/knowledge_layer[llamaindex,foundational_rag]" -e ./sources/exa_web_search -e ./sources/google_scholar_paper_search` +- **`langchain-nvidia-openshell was not found in the package registry`**: the adapter is + the OpenShell partner package in `langchain-ai/langchain-nvidia` (PR #303), not yet on + PyPI. The setup script installs it from a git spec by default; override with + `LANGCHAIN_NVIDIA_REPO=` or `--langchain-nvidia /path/to/checkout`. +- **`unbound variable` in `setup_openshell.sh` on macOS**: the system bash is 3.2; run under + bash 5 (`brew install bash` then `/opt/homebrew/bin/bash ./scripts/setup_openshell.sh ...`). +- **`network.mode` rejected at startup**: the selected provider doesn't declare the + matching capability (`supports_network_policy` for `blocked`, `supports_network_allowlist` + for `allowlist`). Choose a capable provider or relax `network.mode` (e.g. to `open`). +- **Chart shows as text / blank instead of an image in the report or PDF**: the stored + report carries `artifact://`; rendering needs all of (a) a resolved job id + (`selectResolvedDeepResearchJobId`), (b) the `MarkdownRenderer` `urlTransform` preserving + the `artifact://` scheme, and (c) for PDF, an explicit image width (react-pdf draws + intrinsic pixel size otherwise, overflowing the page). Re-export after the dev server + recompiles and check the `[PDF] inline:` lines on `/api/generate-pdf`. CSVs are not images + and never embed as pictures - they appear in `## Generated Artifacts` and download links. +- **Job harvested 0 artifacts though the report describes a chart**: the model wrote a + sandbox path as prose without embedding `![caption](artifact://)`, or wrote outside + `artifact_dir`. The skill mandates the embed token and writing to `artifact_dir`; the + `final_harvest` scan + manifest union and `ensure_inline_artifacts_embedded` are the + backstops. Confirm `artifact_dir` in the prompt matches `workdir`. diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/__init__.py b/src/aiq_agent/agents/deep_researcher/sandbox/__init__.py new file mode 100644 index 000000000..8b250a288 --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/__init__.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Provider-neutral sandbox + artifact runtime for deep research. + +Layout: + base.py SandboxProvider contract (execute required; optional hooks/capabilities) + registry.py register_sandbox_provider / create_sandbox_backend (the config-driven seam) + config.py SandboxConfig (common + providers. + artifact_capture + lifecycle_scope) + capabilities.py SandboxCapabilities + fail-closed verification gate + providers/ one module per provider (modal, openshell, ...) — each self-registers + artifacts/ durable artifact records, manifest parsing, store, and harvester + +Importing this package registers the built-in providers, so the registry is +populated before any config is validated against it. +""" + +from __future__ import annotations + +# Import providers for their registration side effects (built-ins self-register). +from . import providers as _providers # noqa: E402,F401 +from .artifacts import Artifact +from .artifacts import ArtifactManager +from .artifacts import ArtifactStore +from .artifacts import LocalArtifactStore +from .base import SandboxProvider +from .base import SandboxTerminatedError +from .capabilities import CapabilityError +from .capabilities import SandboxCapabilities +from .capabilities import verify_capabilities +from .config import NetworkPolicy +from .config import SandboxConfig +from .registry import SANDBOX_PROVIDER_ENTRY_POINT_GROUP +from .registry import create_sandbox_backend +from .registry import register_sandbox_provider +from .registry import registered_providers + +__all__ = [ + "SandboxProvider", + "SandboxTerminatedError", + "SandboxConfig", + "NetworkPolicy", + "SandboxCapabilities", + "CapabilityError", + "verify_capabilities", + "register_sandbox_provider", + "create_sandbox_backend", + "registered_providers", + "SANDBOX_PROVIDER_ENTRY_POINT_GROUP", + "Artifact", + "ArtifactStore", + "LocalArtifactStore", + "ArtifactManager", +] diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/__init__.py b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/__init__.py new file mode 100644 index 000000000..9bf1578ad --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/__init__.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Durable artifact runtime: records, manifest parsing, storage, and harvesting.""" + +from __future__ import annotations + +from .manager import ArtifactManager +from .manifest import Manifest +from .manifest import ManifestEntry +from .manifest import parse_manifest +from .models import Artifact +from .models import ArtifactKind +from .models import ArtifactProvenance +from .models import ArtifactStatus +from .store import ArtifactStore +from .store import LocalArtifactStore +from .store import SqlArtifactStore + +__all__ = [ + "Artifact", + "ArtifactKind", + "ArtifactStatus", + "ArtifactProvenance", + "Manifest", + "ManifestEntry", + "parse_manifest", + "ArtifactStore", + "LocalArtifactStore", + "SqlArtifactStore", + "ArtifactManager", +] diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/manager.py b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/manager.py new file mode 100644 index 000000000..5cae44941 --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/manager.py @@ -0,0 +1,518 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ArtifactManager - the host-side harvest + validation envelope. + +Harvests bytes via the backend's ``download_files`` (never a CLI), confines paths +to ``artifact_dir``, validates them (MIME-from-bytes, extension allowlist, +size/count quotas), hashes, dedups per-job by digest, persists via the +``ArtifactStore``, and emits SSE artifact events. The agent only ever sees +``artifact://`` references; bytes never enter its context. +""" + +from __future__ import annotations + +import hashlib +import logging +import re +import shlex +import threading +import uuid +from collections.abc import Callable +from pathlib import PurePosixPath +from typing import TYPE_CHECKING +from typing import Any + +from ..config import ArtifactCaptureConfig +from .manifest import ManifestEntry +from .manifest import parse_manifest +from .models import Artifact +from .models import ArtifactKind +from .models import ArtifactProvenance +from .models import ArtifactStatus +from .store import ArtifactStore + +if TYPE_CHECKING: + from deepagents.backends.sandbox import BaseSandbox + +logger = logging.getLogger(__name__) + +_MANIFEST_NAME = "manifest.json" + +# Markdown image references the agent writes as ![caption](artifact://). +_ARTIFACT_REF_RE = re.compile(r"!\[([^\]]*)\]\(artifact://([^)]+)\)") + +# Magic-number sniffing for the formats we inline-render or commonly produce. +_MAGIC_SIGNATURES: tuple[tuple[bytes, str], ...] = ( + (b"\x89PNG\r\n\x1a\n", "image/png"), + (b"\xff\xd8\xff", "image/jpeg"), + (b"GIF87a", "image/gif"), + (b"GIF89a", "image/gif"), + (b"%PDF-", "application/pdf"), +) + +_EXT_MIME: dict[str, str] = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".webp": "image/webp", + ".gif": "image/gif", + ".svg": "image/svg+xml", + ".csv": "text/csv", + ".json": "application/json", + ".md": "text/markdown", + ".ipynb": "application/x-ipynb+json", + ".pdf": "application/pdf", +} + +_MIME_KIND: dict[str, ArtifactKind] = { + "image/png": ArtifactKind.IMAGE, + "image/jpeg": ArtifactKind.IMAGE, + "image/webp": ArtifactKind.IMAGE, + "image/gif": ArtifactKind.IMAGE, + "image/svg+xml": ArtifactKind.IMAGE, + "text/csv": ArtifactKind.TABLE, + "application/json": ArtifactKind.DATASET, + "text/markdown": ArtifactKind.TEXT, + "application/x-ipynb+json": ArtifactKind.NOTEBOOK, + "application/pdf": ArtifactKind.DOCUMENT, +} + + +# Raster images must be magic-confirmed; only these may render inline (the rest are +# download-only until/unless sanitized) to prevent stored-XSS via SVG/HTML/notebooks. +_RASTER_IMAGE_MIMES = frozenset({"image/png", "image/jpeg", "image/webp", "image/gif"}) +_INLINE_SAFE_MIMES = frozenset({"image/png", "image/jpeg", "image/webp"}) + + +def _magic_mime(data: bytes) -> str | None: + """Return the MIME implied by content magic bytes, or None if unrecognized.""" + for signature, mime in _MAGIC_SIGNATURES: + if data.startswith(signature): + return mime + if data[:4] == b"RIFF" and data[8:12] == b"WEBP": + return "image/webp" + return None + + +def _sniff_mime(data: bytes, filename: str) -> str: + """Resolve MIME from content magic bytes, falling back to the extension.""" + magic = _magic_mime(data) + if magic is not None: + return magic + return _EXT_MIME.get(PurePosixPath(filename).suffix.lower(), "application/octet-stream") + + +def _resolve_mime(data: bytes, filename: str) -> str | None: + """Resolve a trusted MIME or return None when content/extension indicate spoofing. + + Raster image extensions must be magic-confirmed; a mismatch between confident + magic and the extension (e.g. a PDF named ``.png``) is rejected. + """ + magic = _magic_mime(data) + ext_mime = _EXT_MIME.get(PurePosixPath(filename).suffix.lower()) + if magic is not None: + both_raster = ext_mime in _RASTER_IMAGE_MIMES and magic in _RASTER_IMAGE_MIMES + if ext_mime is not None and ext_mime != magic and not both_raster: + return None + return magic + if ext_mime in _RASTER_IMAGE_MIMES: + # Claims to be a raster image but has no matching magic bytes -> spoof/corrupt. + return None + return ext_mime or "application/octet-stream" + + +def _sanitize(data: bytes, mime: str) -> bytes | None: + """Validate active-content formats; return None to reject what we cannot make safe.""" + if mime == "image/svg+xml": + # Regex stripping cannot fully neutralize SVG (javascript: URIs, , + # external references, CSS payloads), and the content endpoint serves bytes as the + # stored MIME, so a partial clean still leaves a stored-XSS vector. Fail closed and + # reject SVG until a vetted allowlist sanitizer (e.g. DOMPurify-equivalent) exists. + return None + return data + + +class ArtifactManager: + """Harvests, validates, and persists artifacts produced inside the sandbox.""" + + def __init__( + self, + *, + job_id: str, + backend: BaseSandbox, + store: ArtifactStore, + config: ArtifactCaptureConfig, + artifact_dir: str, + emit: Callable[[dict[str, Any]], None] | None = None, + content_url_template: str = "/v1/jobs/async/job/{job_id}/artifacts/{artifact_id}/content", + ) -> None: + """Configure harvesting for one job against a backend, store, and artifact dir. + + Args: + job_id: Owning job id used to scope and key artifacts. + backend: Sandbox backend used to download/enumerate artifact files. + store: Durable store for persisting metadata and bytes. + config: Capture policy (quotas and allowed extensions). + artifact_dir: Sandbox directory that confines harvestable paths. + emit: Optional SSE emitter for artifact/warning events. + content_url_template: Template for the artifact content endpoint URL. + """ + self.job_id = job_id + self.backend = backend + self.store = store + self.config = config + self.artifact_dir = artifact_dir.rstrip("/") + self._emit = emit + self._content_url_template = content_url_template + self._lock = threading.Lock() + self._seen: set[tuple[str, str]] = set() + self._total_bytes = 0 + self._count = 0 + + def final_harvest(self) -> list[Artifact]: + """Harvest at the end of a successful agent run, with a directory scan fallback.""" + if not self.config.enabled: + return [] + return self._harvest(scan=True) + + def resolve_report_references(self, markdown: str, artifacts: list[Artifact] | None = None) -> str: + """Validate ``artifact://`` image references against this job's artifacts. + + The agent references artifacts by filename (``artifact://``) since + it does not know the host-assigned id. Known references are rewritten to the + durable artifact id and preserved (the logical scheme is kept for the UI/PDF/ + CLI to resolve at their edge); unknown or foreign references are dropped. + + Args: + markdown: The report body to rewrite. + artifacts: Pre-fetched artifacts for this job; loaded from the store when omitted. + """ + if artifacts is None: + try: + artifacts = self.store.list(self.job_id) + except Exception: # noqa: BLE001 - report resolution must not fail the job + logger.warning("Could not load artifacts to resolve report references", exc_info=True) + return markdown + + by_id = {a.artifact_id: a for a in artifacts} + by_name = {a.filename: a for a in artifacts} + + def _replace(match: re.Match[str]) -> str: + """Rewrite a known reference to its durable id; drop unknown ones.""" + caption = match.group(1) + token = match.group(2).strip() + artifact = by_id.get(token) or by_name.get(token) or by_name.get(PurePosixPath(token).name) + if artifact is None: + logger.warning("Dropping report reference to unknown artifact: %s", token) + return "" + return f"![{caption}](artifact://{artifact.artifact_id})" + + return _ARTIFACT_REF_RE.sub(_replace, markdown) + + def ensure_inline_artifacts_embedded(self, markdown: str, artifacts: list[Artifact] | None = None) -> str: + """Append any harvested inline image not already embedded under a ``## Figures`` section. + + Safety net for when the model produces a chart but forgets to embed it: every + magic-verified raster image flagged ``inline`` is guaranteed to surface in the report. + Artifacts already referenced (by durable id) are left untouched and never duplicated. + + Args: + markdown: The report body to augment. + artifacts: Pre-fetched artifacts for this job; loaded from the store when omitted. + """ + if artifacts is None: + try: + artifacts = self.store.list(self.job_id) + except Exception: # noqa: BLE001 - embedding must not fail the job + logger.warning("Could not load artifacts to embed inline figures", exc_info=True) + return markdown + + orphans = [ + a + for a in artifacts + if a.inline and a.kind == ArtifactKind.IMAGE and f"artifact://{a.artifact_id}" not in markdown + ] + if not orphans: + return markdown + + lines = ["", "## Figures", ""] + for artifact in orphans: + caption = artifact.caption or artifact.title or artifact.filename + lines.append(f"![{caption}](artifact://{artifact.artifact_id})") + lines.append("") + logger.info( + "Auto-embedded %d inline figure(s) the report did not reference (job=%s)", + len(orphans), + self.job_id, + ) + return markdown.rstrip() + "\n" + "\n".join(lines) + + def append_artifact_index(self, markdown: str, artifacts: list[Artifact] | None = None) -> str: + """Append a ``## Generated Artifacts`` section crediting sandbox-produced outputs. + + Lists every harvested artifact (charts, CSVs, etc.) so figures and their backing data + are credited alongside the report's external sources. Harvest is job-scoped (each job + writes to its own ``artifact_dir/``), so this lists only the current job's + outputs - not leftovers from other jobs sharing a persistent sandbox. + + Args: + markdown: The report body to augment. + artifacts: Pre-fetched artifacts for this job; loaded from the store when omitted. + """ + if artifacts is None: + try: + artifacts = self.store.list(self.job_id) + except Exception: # noqa: BLE001 - indexing must not fail the job + logger.warning("Could not load artifacts to index generated outputs", exc_info=True) + return markdown + + if not artifacts: + return markdown + + lines = ["", "## Generated Artifacts", ""] + for artifact in artifacts: + descriptor = artifact.caption or artifact.title or artifact.kind.value + lines.append(f"- `{artifact.filename}` - {descriptor} (generated in the analysis sandbox)") + return markdown.rstrip() + "\n" + "\n".join(lines) + "\n" + + # ------------------------------------------------------------------ # + # Internal + # ------------------------------------------------------------------ # + def _harvest(self, *, scan: bool) -> list[Artifact]: + """Discover and capture artifacts under the lock; optionally scan the directory.""" + with self._lock: + entries = self._discover(scan=scan) + captured: list[Artifact] = [] + for entry in entries: + artifact = self._capture(entry) + if artifact is not None: + captured.append(artifact) + if captured: + # Structured lifecycle log (no secrets): provider tokens never appear here. + logger.info( + "Artifact harvest: job=%s scan=%s captured=%d total_bytes=%d count=%d", + self.job_id, + scan, + len(captured), + self._total_bytes, + self._count, + ) + return captured + + def _discover(self, *, scan: bool) -> list[ManifestEntry]: + """Return manifest entries, unioned with a directory scan when ``scan`` is set.""" + entries: list[ManifestEntry] = list(self._read_manifest()) + if not scan: + return entries + # Final harvest: union the manifest with a directory scan so allowed outputs the + # manifest omitted (e.g. a CSV written alongside a declared PNG) are still captured. + # Dedup by path; manifest entries win since they carry title/caption/inline metadata. + seen = {entry.path for entry in entries} + for scanned in self._scan_dir(): + if scanned.path not in seen: + entries.append(scanned) + seen.add(scanned.path) + return entries + + def _read_manifest(self) -> list[ManifestEntry]: + """Download and parse ``manifest.json``; return [] if absent or invalid.""" + manifest_path = f"{self.artifact_dir}/{_MANIFEST_NAME}" + try: + responses = self.backend.download_files([manifest_path]) + except Exception: # noqa: BLE001 - missing manifest is normal + return [] + for resp in responses: + _, content, error = _extract_download(resp) + if error or content is None: + continue + manifest = parse_manifest(content.decode("utf-8", errors="replace")) + if manifest is not None: + return list(manifest.artifacts) + return [] + + def _scan_dir(self) -> list[ManifestEntry]: + """Enumerate allowed files in the artifact dir (bounded, best-effort fallback).""" + try: + response = self.backend.execute(f"find {shlex.quote(self.artifact_dir)} -type f") + except Exception: # noqa: BLE001 - scan is best-effort + logger.warning("Artifact scan failed for job %s", self.job_id, exc_info=True) + return [] + output = getattr(response, "output", "") or "" + entries: list[ManifestEntry] = [] + # Bound the scan so a flood of files can't drive one download round-trip each; + # generous relative to the count quota, which is the real ceiling on stored files. + max_scan = max(self.config.max_file_count * 5, 100) + for line in output.splitlines(): + if len(entries) >= max_scan: + logger.warning("Artifact scan truncated at %d files for job %s", max_scan, self.job_id) + break + path = line.strip() + if not path or path.endswith(f"/{_MANIFEST_NAME}"): + continue + ext = PurePosixPath(path).suffix.lower() + if ext not in self.config.allow_extensions: + continue + entries.append(ManifestEntry(path=path)) + return entries + + def _capture(self, entry: ManifestEntry) -> Artifact | None: + """Validate, download, sanitize, and persist one entry; return the stored artifact. + + Returns ``None`` when the entry is rejected (confinement, allowlist, quota, size, + MIME spoofing, sanitization, or per-run dedup). + """ + # 0. Path-traversal confinement. + if not self._is_confined(entry.path): + logger.warning("Rejecting artifact outside artifact_dir: %s", entry.path) + return None + + # 1. Extension allowlist. + ext = PurePosixPath(entry.path).suffix.lower() + if ext not in self.config.allow_extensions: + self._emit_warning(entry.path, f"extension {ext} not allowed") + return None + + # 1b. Count quota, checked BEFORE download so a flood of files cannot drive one + # transfer round-trip per file before the quota stops storing. + if self._count >= self.config.max_file_count: + self._emit_warning(entry.path, "artifact count quota exceeded; summarize remaining outputs in text") + return None + + # 2. Download bytes. + try: + responses = self.backend.download_files([entry.path]) + except Exception: # noqa: BLE001 - per-file failure must not fail the job + self._emit_warning(entry.path, "download failed") + return None + if not responses: + return None + _, data, error = _extract_download(responses[0]) + if error or data is None: + self._emit_warning(entry.path, error or "no content") + return None + + # 3. Size cap. + if len(data) > self.config.max_file_bytes: + self._emit_warning(entry.path, f"exceeds max_file_bytes ({len(data)})") + return None + + # 4. Quota (count + cumulative bytes). + if self._count >= self.config.max_file_count or self._total_bytes + len(data) > self.config.max_total_bytes: + self._emit_warning(entry.path, "artifact quota exceeded; summarize remaining outputs in text") + return None + + # 5. MIME from bytes; reject content/extension mismatch (spoofing). + filename = PurePosixPath(entry.path).name + mime = _resolve_mime(data, filename) + if mime is None: + self._emit_warning(entry.path, "content does not match its declared type") + return None + + # 6. Sanitize active content (e.g. SVG scripts) before persisting. + sanitized = _sanitize(data, mime) + if sanitized is None: + self._emit_warning(entry.path, "failed sanitization") + return None + data = sanitized + + # 7. Hash + per-run dedup (after sanitization so the digest matches stored bytes). + digest = hashlib.sha256(data).hexdigest() + if (entry.path, digest) in self._seen: + return None + + kind = entry.kind if entry.kind != ArtifactKind.OTHER else _MIME_KIND.get(mime, ArtifactKind.OTHER) + # Render gate: only magic-verified raster images may be embedded inline; SVG, + # notebooks, PDFs, etc. are download-only until/unless deeper sanitization exists. + inline = bool(entry.inline) and mime in _INLINE_SAFE_MIMES + + artifact = Artifact( + artifact_id=f"art_{uuid.uuid4().hex}", + job_id=self.job_id, + kind=kind, + mime_type=mime, + filename=filename, + sandbox_path=entry.path, + storage_uri="", # assigned by the store + sha256=digest, + size_bytes=len(data), + title=entry.title, + caption=entry.caption, + inline=inline, + provenance=ArtifactProvenance(), + status=ArtifactStatus.PENDING, + ) + + # 8. Store first (durable), then emit (outbox discipline). + stored = self.store.put(artifact, data) + self._seen.add((entry.path, digest)) + # A dedup hit returns a pre-existing artifact (different id); it was already + # accounted for and emitted on first capture, so don't charge quota or emit again. + if stored.artifact_id == artifact.artifact_id: + self._total_bytes += len(data) + self._count += 1 + self._emit_artifact(stored) + return stored + + def _is_confined(self, path: str) -> bool: + """Return whether the normalized path stays within ``artifact_dir``.""" + try: + resolved = PurePosixPath(path) + if not resolved.is_absolute(): + resolved = PurePosixPath(self.artifact_dir) / resolved + normalized = _normalize_posix(resolved) + base = _normalize_posix(PurePosixPath(self.artifact_dir)) + return normalized == base or normalized.startswith(base + "/") + except Exception: # noqa: BLE001 + return False + + def _emit_artifact(self, artifact: Artifact) -> None: + """Emit an ``artifact.update`` SSE event with the artifact's content URL.""" + if self._emit is None: + return + content_url = self._content_url_template.format(job_id=self.job_id, artifact_id=artifact.artifact_id) + self._emit(artifact.to_sse_payload(content_url)) + + def _emit_warning(self, path: str, reason: str) -> None: + """Log and emit an ``artifact.warning`` SSE event for a rejected file.""" + logger.warning("Artifact rejected (%s): %s", reason, path) + if self._emit is None: + return + self._emit({"type": "artifact.warning", "data": {"path": path, "reason": reason}}) + + +def _normalize_posix(path: PurePosixPath) -> str: + """Collapse ``.`` / ``..`` segments without touching the filesystem.""" + parts: list[str] = [] + for part in path.parts: + if part == "..": + if parts and parts[-1] not in ("", "/"): + parts.pop() + elif part not in (".", "", "/"): + parts.append(part) + prefix = "/" if path.is_absolute() else "" + return prefix + "/".join(parts) + + +def _extract_download(resp: Any) -> tuple[str, bytes | None, str | None]: + """Defensively read (path, bytes, error) from a FileDownloadResponse.""" + path = getattr(resp, "path", "") or "" + error = getattr(resp, "error", None) + content = getattr(resp, "content", None) + if content is None: + content = getattr(resp, "data", None) + if isinstance(content, str): + content = content.encode("utf-8") + return path, content, error diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/manifest.py b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/manifest.py new file mode 100644 index 000000000..3027c8c8a --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/manifest.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Artifact manifest parsing. + +Generated code may write a ``manifest.json`` into the sandbox artifact directory to +declare its outputs explicitly (reliable). Directory scanning for changed allowed +files is the safety-net fallback for agents that forget to write a manifest. +""" + +from __future__ import annotations + +import json +import logging + +from pydantic import BaseModel +from pydantic import Field +from pydantic import ValidationError + +from .models import ArtifactKind + +logger = logging.getLogger(__name__) + + +class ManifestEntry(BaseModel): + """A single declared artifact in a manifest.""" + + path: str = Field(..., description="Absolute path inside the sandbox") + kind: ArtifactKind = Field(default=ArtifactKind.OTHER) + title: str | None = Field(default=None) + caption: str | None = Field(default=None) + inline: bool = Field(default=False) + source_files: tuple[str, ...] = Field(default=(), description="Inputs used to produce this artifact") + + +class Manifest(BaseModel): + """Top-level manifest schema written by generated code.""" + + version: int = Field(default=1) + artifacts: tuple[ManifestEntry, ...] = Field(default=()) + + +def parse_manifest(raw: str) -> Manifest | None: + """Parse a manifest JSON string into a :class:`Manifest`. + + Args: + raw: The manifest file contents. + + Returns: + A parsed ``Manifest``, or ``None`` if the content is invalid (a warning is + logged; callers fall back to directory scanning). + """ + try: + data = json.loads(raw) + except (json.JSONDecodeError, TypeError): + logger.warning("Artifact manifest is not valid JSON; falling back to scan") + return None + try: + return Manifest.model_validate(data) + except ValidationError: + logger.warning("Artifact manifest failed schema validation; falling back to scan") + return None diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/models.py b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/models.py new file mode 100644 index 000000000..0ca080d05 --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/models.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Durable artifact record. + +This is the metadata record only. Bytes are stored out-of-band by the +``ArtifactStore`` and fetched server-side on demand; they never enter the agent's +context or conversation history. Reports reference artifacts by ``artifact_id`` +(``artifact://``), keeping prompt cost independent of artifact size. +""" + +from __future__ import annotations + +from datetime import UTC +from datetime import datetime +from enum import StrEnum + +from pydantic import BaseModel +from pydantic import Field + + +class ArtifactKind(StrEnum): + """High-level artifact category used for rendering and grouping.""" + + IMAGE = "image" + TABLE = "table" + DATASET = "dataset" + NOTEBOOK = "notebook" + DOCUMENT = "document" + TEXT = "text" + ARCHIVE = "archive" + OTHER = "other" + + +class ArtifactStatus(StrEnum): + """Lifecycle status of an artifact record.""" + + PENDING = "pending" + AVAILABLE = "available" + REJECTED = "rejected" + DELETED = "deleted" + + +class ArtifactProvenance(BaseModel): + """Reproducibility metadata for a generated artifact.""" + + command: str | None = Field(default=None, description="Command that produced the artifact") + script_sha256: str | None = Field(default=None, description="Digest of the generating script") + input_file_hashes: dict[str, str] = Field(default_factory=dict, description="Path -> sha256 of input files") + package_snapshot: tuple[str, ...] = Field(default=(), description="Installed package versions at run time") + + +class Artifact(BaseModel): + """Durable record for a single generated artifact (metadata, not bytes).""" + + artifact_id: str = Field(..., max_length=64, description="Stable ID (UUID, optionally suffixed with a digest)") + job_id: str = Field(..., max_length=64, description="Owning async job (retention + authorization boundary)") + kind: ArtifactKind = Field(default=ArtifactKind.OTHER) + mime_type: str = Field(..., description="MIME type validated from bytes, not just filename") + filename: str = Field(..., description="User-facing filename") + sandbox_path: str = Field(..., description="Original path inside the sandbox") + storage_uri: str = Field(..., description="Durable location of the bytes (local path or object-store URI)") + sha256: str = Field(..., min_length=64, max_length=64, description="Content digest for integrity and deduplication") + size_bytes: int = Field(..., ge=0, description="Byte size for quota and UI display") + title: str | None = Field(default=None, description="Optional display title") + caption: str | None = Field(default=None, description="Optional report caption") + inline: bool = Field(default=False, description="Whether the report may embed it inline") + workflow: str | None = Field(default=None, description="orchestrator | planner-agent | researcher-agent | skill") + source_tool_call_id: str | None = Field(default=None, description="Tool call that created or registered it") + provenance: ArtifactProvenance = Field(default_factory=ArtifactProvenance) + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC), description="Event ordering timestamp") + status: ArtifactStatus = Field(default=ArtifactStatus.PENDING) + + def to_sse_payload(self, content_url: str) -> dict[str, object]: + """Build the richer ``artifact`` SSE payload for live UI updates. + + Args: + content_url: Authenticated URL where the bytes can be fetched. + + Returns: + A JSON-serializable payload (no bytes) matching the design's artifact event. + """ + return { + "type": "artifact", + "artifact_id": self.artifact_id, + "kind": self.kind.value, + "filename": self.filename, + "mime_type": self.mime_type, + "size_bytes": self.size_bytes, + "sha256": self.sha256, + "title": self.title, + "caption": self.caption, + "inline": self.inline, + "content_url": content_url, + } diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/store.py b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/store.py new file mode 100644 index 000000000..ba89c566d --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/store.py @@ -0,0 +1,376 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Durable artifact storage. + +Metadata and bytes are stored together in an ``artifacts`` table on the shared job +``db_url`` (the same database used by the job/event stores), keyed by +``artifact_id``. This is the byte-locality fix: the Dask worker (writer) and the API +process (reader) both reach artifacts with no shared filesystem. + +- Metadata columns are queryable and listed in the UI/CLI. +- Bytes live in a size-capped ``content`` BLOB column for milestone 1. +- ``storage_uri`` stays abstract so a production deployment can move bytes to + object storage (S3) without changing callers. +""" + +from __future__ import annotations + +import json +import logging +import threading +from abc import ABC +from abc import abstractmethod +from collections.abc import Iterator +from typing import Any + +from .models import Artifact +from .models import ArtifactProvenance +from .models import ArtifactStatus + +logger = logging.getLogger(__name__) + +_READ_CHUNK_BYTES = 1 << 20 # 1 MiB streaming chunk + + +def _normalize_db_url(db_url: str) -> str: + """Normalize a SQLAlchemy URL to a sync driver (mirrors the event store). + + Artifacts use sync SQLAlchemy; the async API path wraps reads in an executor. + """ + if db_url.startswith("postgresql") or db_url.startswith("postgres"): + base = db_url.replace("+asyncpg", "").replace("+psycopg2", "").replace("+psycopg", "") + base = base.replace("postgres://", "postgresql://") + return base.replace("postgresql://", "postgresql+psycopg://") + if db_url.startswith("sqlite"): + return db_url.replace("+aiosqlite", "") + return db_url + + +class ArtifactStore(ABC): + """Pluggable durable store for artifact metadata and bytes.""" + + @abstractmethod + def put(self, artifact: Artifact, data: bytes) -> Artifact: + """Persist bytes and metadata, returning the stored (possibly deduped) record.""" + + @abstractmethod + def open_bytes(self, job_id: str, artifact_id: str) -> Iterator[bytes]: + """Stream an artifact's bytes for the content endpoint or PDF embedding.""" + + @abstractmethod + def get(self, job_id: str, artifact_id: str) -> Artifact | None: + """Return artifact metadata, or ``None`` if not found for this job.""" + + @abstractmethod + def find_by_digest(self, job_id: str, sha256: str) -> Artifact | None: + """Return an existing artifact with the same content digest for this job.""" + + @abstractmethod + def list(self, job_id: str) -> list[Artifact]: + """List all artifacts owned by a job.""" + + @abstractmethod + def delete_job(self, job_id: str) -> int: + """Delete all artifacts for a job (retention/expiry). Returns count removed.""" + + @abstractmethod + def cleanup_old_artifacts(self, retention_seconds: int) -> int: + """Delete artifacts older than the retention period. Returns count removed.""" + + +class SqlArtifactStore(ArtifactStore): + """SQL-backed store (SQLite/Postgres) on the shared job ``db_url``. + + Bytes live in a capped BLOB column. Dedup is per-job by content digest so + harvest is idempotent across job retries. + """ + + _engines: dict[str, Any] = {} + _initialized: set[str] = set() + _lock = threading.Lock() + + def __init__(self, db_url: str = "sqlite+aiosqlite:///./jobs.db") -> None: + """Bind to the shared job database and ensure the artifacts table exists. + + Args: + db_url: SQLAlchemy URL of the shared job/event database. + """ + self.db_url = db_url + self._engine = self._get_engine(db_url) + self._ensure_table() + + @classmethod + def _get_engine(cls, db_url: str) -> Any: + """Return a process-wide engine for the URL, creating it once (thread-safe).""" + from sqlalchemy import create_engine + from sqlalchemy import event + + with cls._lock: + if db_url in cls._engines: + return cls._engines[db_url] + normalized = _normalize_db_url(db_url) + connect_args = {"check_same_thread": False, "timeout": 30} if normalized.startswith("sqlite") else {} + engine = create_engine(normalized, pool_pre_ping=True, pool_recycle=1800, connect_args=connect_args) + if normalized.startswith("sqlite"): + # WAL improves concurrency between the worker writer and API reader. + @event.listens_for(engine, "connect") + def _set_wal(dbapi_conn: Any, _record: Any) -> None: # pragma: no cover - driver callback + """Enable SQLite WAL mode on connect for reader/writer concurrency.""" + cursor = dbapi_conn.cursor() + cursor.execute("PRAGMA journal_mode=WAL") + cursor.close() + + cls._engines[db_url] = engine + return engine + + def _ensure_table(self) -> None: + """Create the ``artifacts`` table on first use for this database URL.""" + if self.db_url in SqlArtifactStore._initialized: + return + from sqlalchemy import BigInteger + from sqlalchemy import Boolean + from sqlalchemy import Column + from sqlalchemy import DateTime + from sqlalchemy import Index + from sqlalchemy import LargeBinary + from sqlalchemy import MetaData + from sqlalchemy import String + from sqlalchemy import Table + from sqlalchemy import Text + from sqlalchemy import inspect + from sqlalchemy.sql import func + + metadata = MetaData() + Table( + "artifacts", + metadata, + Column("artifact_id", String(64), primary_key=True), + Column("job_id", String(64), nullable=False, index=True), + Column("kind", String(32), nullable=False), + Column("mime_type", String(128), nullable=False), + Column("filename", String(512), nullable=False), + Column("sandbox_path", Text, nullable=False), + Column("storage_uri", Text, nullable=False), + Column("sha256", String(64), nullable=False), + Column("size_bytes", BigInteger, nullable=False), + Column("title", Text, nullable=True), + Column("caption", Text, nullable=True), + Column("inline", Boolean, nullable=False, default=False), + Column("workflow", String(64), nullable=True), + Column("source_tool_call_id", String(128), nullable=True), + Column("provenance", Text, nullable=True), + Column("status", String(16), nullable=False), + Column("content", LargeBinary, nullable=True), + Column("created_at", DateTime(timezone=True), server_default=func.now()), + Index("idx_artifacts_job_sha", "job_id", "sha256"), + ) + inspector = inspect(self._engine) + if not inspector.has_table("artifacts"): + metadata.create_all(self._engine) + logger.info("Created artifacts table (backend=%s)", self._engine.dialect.name) + SqlArtifactStore._initialized.add(self.db_url) + + def put(self, artifact: Artifact, data: bytes) -> Artifact: + """Persist bytes and metadata, deduping per job by content digest. + + Args: + artifact: Metadata for the artifact to store. + data: Raw artifact bytes. + + Returns: + The stored record, or the existing record on a digest dedup hit. + """ + existing = self.find_by_digest(artifact.job_id, artifact.sha256) + if existing is not None: + logger.debug("Artifact dedup hit for job=%s sha=%s", artifact.job_id, artifact.sha256[:12]) + return existing + + from sqlalchemy import text + + stored = artifact.model_copy( + update={ + # Logical location only — never embed db_url (it may carry credentials). + "storage_uri": f"db://artifacts/{artifact.artifact_id}", + "status": ArtifactStatus.AVAILABLE, + } + ) + with self._engine.connect() as conn: + conn.execute( + text( + "INSERT INTO artifacts (artifact_id, job_id, kind, mime_type, filename, sandbox_path, " + "storage_uri, sha256, size_bytes, title, caption, inline, workflow, source_tool_call_id, " + "provenance, status, content) VALUES (:artifact_id, :job_id, :kind, :mime_type, :filename, " + ":sandbox_path, :storage_uri, :sha256, :size_bytes, :title, :caption, :inline, :workflow, " + ":source_tool_call_id, :provenance, :status, :content)" + ), + { + "artifact_id": stored.artifact_id, + "job_id": stored.job_id, + "kind": stored.kind.value, + "mime_type": stored.mime_type, + "filename": stored.filename, + "sandbox_path": stored.sandbox_path, + "storage_uri": stored.storage_uri, + "sha256": stored.sha256, + "size_bytes": stored.size_bytes, + "title": stored.title, + "caption": stored.caption, + "inline": stored.inline, + "workflow": stored.workflow, + "source_tool_call_id": stored.source_tool_call_id, + "provenance": stored.provenance.model_dump_json(), + "status": stored.status.value, + "content": data, + }, + ) + conn.commit() + return stored + + def open_bytes(self, job_id: str, artifact_id: str) -> Iterator[bytes]: + """Yield an artifact's bytes in 1 MiB chunks, or nothing if not found.""" + from sqlalchemy import text + + with self._engine.connect() as conn: + row = conn.execute( + text("SELECT content FROM artifacts WHERE job_id = :job_id AND artifact_id = :artifact_id"), + {"job_id": job_id, "artifact_id": artifact_id}, + ).fetchone() + if row is None or row[0] is None: + return + data: bytes = row[0] + for start in range(0, len(data), _READ_CHUNK_BYTES): + yield data[start : start + _READ_CHUNK_BYTES] + + def get(self, job_id: str, artifact_id: str) -> Artifact | None: + """Return artifact metadata for the job, or ``None`` if not found.""" + from sqlalchemy import text + + with self._engine.connect() as conn: + row = ( + conn.execute( + text( + f"SELECT {_META_COLUMNS} FROM artifacts WHERE job_id = :job_id AND artifact_id = :artifact_id" + ), + {"job_id": job_id, "artifact_id": artifact_id}, + ) + .mappings() + .fetchone() + ) + return _row_to_artifact(row) if row else None + + def find_by_digest(self, job_id: str, sha256: str) -> Artifact | None: + """Return an existing artifact with the same digest for the job, if any.""" + from sqlalchemy import text + + with self._engine.connect() as conn: + row = ( + conn.execute( + text(f"SELECT {_META_COLUMNS} FROM artifacts WHERE job_id = :job_id AND sha256 = :sha256 LIMIT 1"), + {"job_id": job_id, "sha256": sha256}, + ) + .mappings() + .fetchone() + ) + return _row_to_artifact(row) if row else None + + def list(self, job_id: str) -> list[Artifact]: + """Return all artifacts owned by the job, ordered by creation time.""" + from sqlalchemy import text + + with self._engine.connect() as conn: + rows = ( + conn.execute( + text(f"SELECT {_META_COLUMNS} FROM artifacts WHERE job_id = :job_id ORDER BY created_at"), + {"job_id": job_id}, + ) + .mappings() + .fetchall() + ) + return [_row_to_artifact(row) for row in rows] + + def delete_job(self, job_id: str) -> int: + """Delete all artifacts for the job and return the number removed.""" + from sqlalchemy import text + + with self._engine.connect() as conn: + result = conn.execute(text("DELETE FROM artifacts WHERE job_id = :job_id"), {"job_id": job_id}) + conn.commit() + return result.rowcount + + def cleanup_old_artifacts(self, retention_seconds: int) -> int: + """Delete artifacts older than the retention window and return the count. + + A non-positive retention is refused (returns 0) to avoid deleting everything. + """ + from sqlalchemy import text + + # A non-positive retention would make the cutoff "now or later" and delete everything. + if retention_seconds <= 0: + logger.warning("Refusing artifact cleanup with non-positive retention_seconds=%s", retention_seconds) + return 0 + is_postgres = _normalize_db_url(self.db_url).startswith("postgresql") + with self._engine.connect() as conn: + if is_postgres: + result = conn.execute( + text("DELETE FROM artifacts WHERE created_at < NOW() - :seconds * INTERVAL '1 second'"), + {"seconds": retention_seconds}, + ) + else: + result = conn.execute( + text("DELETE FROM artifacts WHERE created_at < datetime('now', :interval)"), + {"interval": f"-{retention_seconds} seconds"}, + ) + conn.commit() + return result.rowcount + + +# Backwards-friendly alias; local development and single-node use the same SQL store. +LocalArtifactStore = SqlArtifactStore + +_META_COLUMNS = ( + "artifact_id, job_id, kind, mime_type, filename, sandbox_path, storage_uri, sha256, size_bytes, " + "title, caption, inline, workflow, source_tool_call_id, provenance, status, created_at" +) + + +def _row_to_artifact(row: Any) -> Artifact: + """Build an ``Artifact`` from a metadata row, tolerating bad provenance JSON.""" + provenance = ArtifactProvenance() + raw = row.get("provenance") + if raw: + try: + provenance = ArtifactProvenance.model_validate(json.loads(raw)) + except (ValueError, TypeError): + logger.warning("Failed to parse artifact provenance for %s", row.get("artifact_id")) + return Artifact( + artifact_id=row["artifact_id"], + job_id=row["job_id"], + kind=row["kind"], + mime_type=row["mime_type"], + filename=row["filename"], + sandbox_path=row["sandbox_path"], + storage_uri=row["storage_uri"], + sha256=row["sha256"], + size_bytes=row["size_bytes"], + title=row.get("title"), + caption=row.get("caption"), + inline=bool(row.get("inline")), + workflow=row.get("workflow"), + source_tool_call_id=row.get("source_tool_call_id"), + provenance=provenance, + status=row["status"], + created_at=row["created_at"], + ) diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/base.py b/src/aiq_agent/agents/deep_researcher/sandbox/base.py new file mode 100644 index 000000000..20d726f77 --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/base.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The ``SandboxProvider`` base — the thin contract a sandbox backend must satisfy. + +Design: force only the minimum (``_create_session`` + declared ``capabilities``). +Everything else — lazy creation, locking, idempotency-gated retry, and cleanup — is +shared here so a new provider implements just the SDK-specific parts. File tools +(``read_file``/``write_file``/``edit_file``/``ls``/``glob``) are inherited from +``BaseSandbox``, which builds them on top of ``execute``; providers never reimplement them. + +This mirrors the knowledge-layer adapter philosophy: a small required surface, optional +capabilities with safe defaults, and provider-owned error classification. +""" + +from __future__ import annotations + +import logging +import shlex +import threading +from abc import ABC +from abc import abstractmethod +from collections.abc import Callable +from typing import TYPE_CHECKING +from typing import TypeVar + +from deepagents.backends.protocol import ExecuteResponse +from deepagents.backends.protocol import FileDownloadResponse +from deepagents.backends.protocol import FileUploadResponse +from deepagents.backends.sandbox import BaseSandbox + +from .capabilities import SandboxCapabilities +from .config import job_scoped_artifact_dir +from .config import job_scoped_workdir + +if TYPE_CHECKING: + from .config import SandboxConfig + +logger = logging.getLogger(__name__) + +_T = TypeVar("_T") + +# Per-job workspace creation is a quick mkdir; bound it well under the sandbox lifetime. +_WORKSPACE_PREP_TIMEOUT = 60 + + +class SandboxTerminatedError(RuntimeError): + """Raised when an operation is attempted on a terminated (cancelled/closed) provider.""" + + +class SandboxProvider(BaseSandbox, ABC): + """Job-scoped, lazily-created sandbox backend behind a uniform contract. + + Subclasses implement provider-specific session creation and declare their + capabilities. The base provides shared resilience (single-flight creation, + a serialization lock around remote calls, and idempotency-gated retry driven + by the provider's own :meth:`is_recoverable_error`). + + Attributes: + provider_name: Registry key for this provider. + """ + + provider_name: str = "base" + + def __init__(self, config: SandboxConfig, job_id: str) -> None: + """Initialize the provider. + + Args: + config: Resolved sandbox configuration for the job. + job_id: Async job identifier used to scope the sandbox identity. + """ + self.config = config + self.job_id = job_id + self.sandbox_name = self._scoped_name(job_id) + # Per-job workspace roots inside the (possibly shared/reused) sandbox. Computed via + # the same helpers the runtime uses, so the directory the agent writes to and the + # directory the harvest scans always agree. + self.workdir = job_scoped_workdir(config.workdir, job_id) + self.artifact_dir = job_scoped_artifact_dir(config.workdir, job_id) + self._session: BaseSandbox | None = None + # Operation lock: serializes remote calls + gated retry (held across the call). + self._lock = threading.RLock() + # State lock: guards the session reference and the terminated flag only. Held for + # microseconds and NEVER across a remote call or session creation, so close()/ + # terminate() can tear down out-of-band without waiting on an in-flight execute. + # Lock order is strictly operation-lock -> state-lock; teardown takes only the + # state lock, so the two can never deadlock. + self._state_lock = threading.Lock() + self._terminated = False + + # ------------------------------------------------------------------ # + # Required surface (the only things a provider must implement) + # ------------------------------------------------------------------ # + @abstractmethod + def _create_session(self) -> BaseSandbox: + """Create the underlying provider-specific ``BaseSandbox`` session. + + Implementations own the SDK calls (gateway connect, create/attach, image + build, ready-wait). They must NOT silently attach to a sandbox owned by a + prior job; collisions should produce a fresh, job-scoped sandbox. + + Returns: + A concrete ``BaseSandbox`` (e.g. the langchain-modal / langchain-openshell adapter). + """ + + @property + @abstractmethod + def capabilities(self) -> SandboxCapabilities: + """Return the security/lifecycle guarantees this provider can enforce.""" + + # ------------------------------------------------------------------ # + # Optional hooks with safe, conservative defaults (override to opt in) + # ------------------------------------------------------------------ # + @classmethod + def _scoped_name(cls, job_id: str) -> str: + """Translate a job id into a provider-legal, job-scoped sandbox name. + + Providers override to apply their own naming rules (length, charset). + """ + return job_id + + def is_recoverable_error(self, exc: Exception) -> bool: + """Classify an exception as a transient/stale-sandbox error worth retrying. + + Conservative default returns ``False`` so unknown providers never recreate + and silently re-run against an empty sandbox. Providers override using their + own SDK's typed exceptions rather than fragile string matching. + """ + return False + + def close(self) -> None: + """Release the underlying sandbox session, if any (idempotent). + + Tears the session down out-of-band (under the short state lock, never the + operation lock) so cleanup never blocks behind an in-flight call. Unlike + :meth:`terminate`, this does not permanently terminate the provider: a later + operation may lazily recreate the session. Default delegates to the session's + ``close`` when present. + """ + with self._state_lock: + session = self._session + self._session = None + self._safe_close(session) + + def terminate(self) -> None: + """Forcibly stop any in-flight execution and release the sandbox (idempotent). + + Used on the cancellation/timeout path. Because teardown runs out-of-band (it does + not take the operation lock), closing the underlying session interrupts a + long-running ``execute`` rather than waiting for it to finish. Providers that can + hard-kill a remote process should override :meth:`_terminate_session`. + """ + with self._state_lock: + session = self._session + self._session = None + self._terminated = True + self._terminate_session(session) + + def _terminate_session(self, session: BaseSandbox | None) -> None: + """Forcibly stop a session. Default closes it; providers may override to hard-kill.""" + self._safe_close(session) + + def _prepare_workspace(self, session: BaseSandbox) -> None: + """Create the job-scoped workspace and artifact directories in a new session. + + Idempotent (``mkdir -p``); runs once per session creation so the per-job root + exists before the first ``write_file``, even when the underlying sandbox is shared + and reused across jobs. Best-effort: a failure here is logged rather than raised, + since a genuine filesystem problem resurfaces on the first real write. + """ + command = f"mkdir -p {shlex.quote(self.workdir)} {shlex.quote(self.artifact_dir)}" + try: + session.execute(command, timeout=self._clamp_timeout(_WORKSPACE_PREP_TIMEOUT)) + except Exception: # noqa: BLE001 - workspace prep is best-effort; real failures resurface on first write + logger.warning("Sandbox %s workspace prep failed (%s)", self.sandbox_name, command, exc_info=True) + + def _safe_close(self, session: BaseSandbox | None) -> None: + """Best-effort close of a session; never raises on the teardown path.""" + if session is not None and hasattr(session, "close"): + try: + session.close() + except Exception: # noqa: BLE001 - cleanup must never raise on the terminal path + logger.warning("Sandbox %s cleanup failed", self.sandbox_name, exc_info=True) + + @property + def id(self) -> str: + """Stable identifier: the live session id once created, else the scoped name.""" + with self._state_lock: + session = self._session + return session.id if session is not None else self.sandbox_name + + # ------------------------------------------------------------------ # + # Byte/exec surface — shared resilience, delegated to the session + # ------------------------------------------------------------------ # + def execute(self, command: str, *, timeout: int | None = None) -> ExecuteResponse: + """Run a command in the sandbox (non-idempotent; no recreate-and-retry). + + The per-call ``timeout`` is clamped to the configured sandbox lifetime + (``config.timeout``). Agent-supplied timeouts are unreliable (e.g. a tool + may pass milliseconds where the backend expects seconds), and a single + ``execute`` should never outlive the sandbox or exceed a provider's hard + cap, so we bound it rather than let the backend reject the call. + """ + timeout = self._clamp_timeout(timeout) + return self._call("execute", lambda s: s.execute(command, timeout=timeout), idempotent=False) + + def _clamp_timeout(self, timeout: int | None) -> int | None: + """Bound a per-call timeout to ``config.timeout`` (the sandbox max lifetime).""" + if timeout is None: + return None + ceiling = self.config.timeout + if timeout > ceiling: + logger.warning( + "Sandbox %s execute timeout %ss exceeds configured limit %ss; clamping", + self.sandbox_name, + timeout, + ceiling, + ) + return ceiling + return max(1, timeout) + + def upload_files(self, files: list[tuple[str, bytes]]) -> list[FileUploadResponse]: + """Upload input files into the sandbox (idempotent; safe to retry).""" + return self._call("upload_files", lambda s: s.upload_files(files), idempotent=True) + + def download_files(self, paths: list[str]) -> list[FileDownloadResponse]: + """Download files from the sandbox for artifact harvesting (idempotent).""" + return self._call("download_files", lambda s: s.download_files(paths), idempotent=True) + + # ------------------------------------------------------------------ # + # Internal lifecycle + # ------------------------------------------------------------------ # + def _session_or_create(self) -> BaseSandbox: + """Return the live session, creating it once (single-flight). + + Called while the operation lock is held, so creation is serialized. Only the + session-reference reads/writes take the short state lock (not the slow + ``_create_session`` call), so a concurrent terminate/close can swap the session + out without waiting for the in-flight remote call. + """ + with self._state_lock: + if self._terminated: + raise SandboxTerminatedError(f"Sandbox {self.sandbox_name} has been terminated") + if self._session is not None: + return self._session + logger.info("Sandbox session init: provider=%s name=%s", self.provider_name, self.sandbox_name) + created = self._create_session() + self._prepare_workspace(created) + with self._state_lock: + if not self._terminated: + self._session = created + return created + # Terminated mid-creation: discard the freshly created session rather than leak it. + self._safe_close(created) + raise SandboxTerminatedError(f"Sandbox {self.sandbox_name} has been terminated") + + def _reset_session(self) -> None: + """Drop and recreate the session (used only for idempotent recoverable retries).""" + logger.warning( + "Sandbox session RESET: provider=%s name=%s (prior in-sandbox files are lost)", + self.provider_name, + self.sandbox_name, + ) + with self._state_lock: + stale = self._session + self._session = None + self._safe_close(stale) + created = self._create_session() + self._prepare_workspace(created) + with self._state_lock: + if not self._terminated: + self._session = created + return + self._safe_close(created) + raise SandboxTerminatedError(f"Sandbox {self.sandbox_name} has been terminated") + + def _call(self, op_name: str, fn: Callable[[BaseSandbox], _T], *, idempotent: bool) -> _T: + """Run a remote call with the serialization lock and gated retry. + + The operation lock serializes calls into a single shared job sandbox to avoid + filesystem races and reset-during-execute hazards. A concurrent ``terminate()`` + does not take this lock: it closes the session out-of-band, which interrupts the + in-flight call; the resulting error is re-raised (never retried) once we observe + the terminated flag. Retry otherwise happens only when the operation is idempotent + AND the provider classifies the error as recoverable (fail-safe over fail-silent). + """ + with self._lock: + try: + result = fn(self._session_or_create()) + # A concurrent terminate() can flip _terminated while the call was in + # flight; if the close did not interrupt it (or the call won the race), + # surface cancellation rather than returning a result for a terminated job. + with self._state_lock: + if self._terminated: + raise SandboxTerminatedError(f"Sandbox {self.sandbox_name} has been terminated") + return result + except Exception as exc: + with self._state_lock: + terminated = self._terminated + if terminated: + raise + if idempotent and self.is_recoverable_error(exc): + logger.warning("Sandbox %s recoverable error on %s; recreating and retrying once", self.id, op_name) + self._reset_session() + return fn(self._session_or_create()) + raise diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/capabilities.py b/src/aiq_agent/agents/deep_researcher/sandbox/capabilities.py new file mode 100644 index 000000000..fc730b813 --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/capabilities.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sandbox capability declarations and the fail-closed verification gate. + +A provider declares what security/lifecycle guarantees it can enforce. The runtime +verifies the active :class:`SandboxConfig` against those declared capabilities +*before* creating a backend and refuses to run when a required guarantee is not +available (fail-closed). This keeps the provider interface thin while keeping the +security floor enforceable across every provider. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pydantic import BaseModel +from pydantic import Field + +if TYPE_CHECKING: + from .config import SandboxConfig + + +class SandboxCapabilities(BaseModel): + """Security and lifecycle guarantees a provider declares it can enforce. + + Defaults are conservative: an unknown provider is assumed to support nothing, + so the fail-closed gate refuses workloads that require guarantees the provider + has not explicitly claimed. + + Attributes: + supports_network_policy: Provider can enforce outbound network blocking. + supports_network_allowlist: Provider can enforce a per-host egress allowlist. + supports_filesystem_policy: Provider can restrict filesystem access at creation. + supports_process_policy: Provider can restrict process/exec behavior. + supports_resource_limits: Provider can cap CPU/memory/disk. + supports_artifact_download: Provider implements byte-accurate ``download_files``. + supports_cleanup: Provider implements ``close`` for deterministic teardown. + supports_terminate: Provider can forcibly terminate a running execution. + """ + + supports_network_policy: bool = Field(default=False) + supports_network_allowlist: bool = Field(default=False) + supports_filesystem_policy: bool = Field(default=False) + supports_process_policy: bool = Field(default=False) + supports_resource_limits: bool = Field(default=False) + supports_artifact_download: bool = Field(default=False) + supports_cleanup: bool = Field(default=False) + supports_terminate: bool = Field(default=False) + + +class CapabilityError(ValueError): + """Raised when a sandbox config requires a guarantee the provider cannot enforce.""" + + +def verify_capabilities(config: SandboxConfig, capabilities: SandboxCapabilities) -> None: + """Fail closed if the config demands a guarantee the provider does not declare. + + Args: + config: The resolved sandbox configuration for the job. + capabilities: The selected provider's declared capabilities. + + Raises: + CapabilityError: If a required guarantee (e.g. network policy, allowlist, or + artifact capture) is requested but unsupported by the provider. + """ + mode = config.network.mode + if mode == "blocked" and not capabilities.supports_network_policy: + raise CapabilityError( + f"Provider '{config.provider}' cannot enforce network.mode='blocked' (block_network). " + "Refusing to run code with un-enforceable network policy. " + "Choose a provider that declares supports_network_policy, or set network.mode='open' explicitly." + ) + + if mode == "allowlist" and not capabilities.supports_network_allowlist: + raise CapabilityError( + f"Provider '{config.provider}' cannot enforce network.mode='allowlist'. " + "This provider's network policy is all-or-nothing. " + "Choose a provider that declares supports_network_allowlist, or use network.mode='blocked'/'open'." + ) + + if config.resources.any_set() and not capabilities.supports_resource_limits: + raise CapabilityError( + f"Provider '{config.provider}' cannot enforce resource limits (CPU/memory), but " + "sandbox.resources requests one. Refusing to run with un-enforceable limits. " + "Choose a provider that declares supports_resource_limits, or remove sandbox.resources." + ) + + if config.artifact_capture.enabled and not capabilities.supports_artifact_download: + raise CapabilityError( + f"Provider '{config.provider}' cannot download artifacts (no download_files support), " + "but artifact_capture.enabled=true. Disable artifact capture or choose a provider that supports it." + ) diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/config.py b/src/aiq_agent/agents/deep_researcher/sandbox/config.py new file mode 100644 index 000000000..c6cca98cd --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/config.py @@ -0,0 +1,255 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Provider-neutral sandbox configuration models. + +Common fields apply to every provider; provider-specific settings live under +``providers.``. A backward-compatible validator lifts legacy flat Modal +fields into ``providers.modal`` so pre-existing configs keep loading. The +``provider`` field is validated against the registry, so any registered provider +is automatically accepted with no edits here. +""" + +from __future__ import annotations + +from typing import Any +from typing import Literal + +from pydantic import BaseModel +from pydantic import Field +from pydantic import field_validator +from pydantic import model_validator + +DEFAULT_WORKDIR = "/workspace" + + +def _safe_job_segment(job_id: str) -> str: + """Return a filename-safe single path segment for ``job_id``. + + The job id becomes a directory name, so keep only filename-safe characters: a + crafted id containing ``/`` or ``..`` must not be able to escape the base workdir. + """ + return "".join(c if (c.isalnum() or c in "-_") else "_" for c in job_id) or "job" + + +def job_scoped_workdir(base_workdir: str, job_id: str) -> str: + """Return the per-job working directory ``/``. + + Per-job directories prevent accidental filename collisions in a shared sandbox. + They are organization, not an access-control boundary: code running in one job can + still access sibling directories unless the provider enforces stronger isolation. + """ + return f"{base_workdir.rstrip('/')}/{_safe_job_segment(job_id)}" + + +def job_scoped_artifact_dir(base_workdir: str, job_id: str) -> str: + """Return the per-job artifact directory nested under the job working directory.""" + return f"{job_scoped_workdir(base_workdir, job_id)}/aiq-artifacts" + + +# Allowed artifact extensions for the first capture milestone (validated MIME-from-bytes +# happens in the ArtifactManager; this is the coarse filename allowlist). +_DEFAULT_ALLOW_EXTENSIONS: tuple[str, ...] = ( + ".png", + ".jpg", + ".jpeg", + ".webp", + ".svg", + ".csv", + ".json", + ".md", + ".ipynb", + ".pdf", +) + + +class ModalProviderConfig(BaseModel): + """Modal-specific sandbox settings.""" + + app_name: str = Field(default="aiq-deep-research", description="Modal app name for deep research sandboxes") + image: str = Field(default="python:3.12-slim", description="Container image for Modal sandboxes") + python_packages: tuple[str, ...] = Field( + default=(), + description="Python packages to install into the Modal sandbox image (e.g. matplotlib, pandas).", + ) + + +class OpenShellProviderConfig(BaseModel): + """OpenShell-specific sandbox settings (enterprise/on-prem example provider).""" + + gateway: str | None = Field(default=None, description="OpenShell gateway/cluster endpoint or name") + sandbox_name: str | None = Field( + default=None, + description="Existing named OpenShell sandbox to attach to (required when a policy file is used).", + ) + policy: str | None = Field( + default=None, + description="OpenShell policy file path. Requires a pre-created named sandbox (sandbox_name).", + ) + image: str = Field(default="base", description="OpenShell image identifier") + ready_timeout_seconds: float = Field(default=300.0, description="Seconds to wait for the sandbox to become ready") + delete_on_exit: bool = Field(default=True, description="Delete the sandbox when its session context closes") + shell: tuple[str, ...] = Field( + default=("bash", "-c"), + description="Shell argv prefix passed to the langchain-nvidia-openshell adapter.", + ) + + +class SandboxProvidersConfig(BaseModel): + """Per-provider configuration blocks. Add a provider by adding an optional field here.""" + + modal: ModalProviderConfig = Field(default_factory=ModalProviderConfig) + openshell: OpenShellProviderConfig = Field(default_factory=OpenShellProviderConfig) + + +class ArtifactCaptureConfig(BaseModel): + """Controls durable harvesting of generated binary/rich artifacts.""" + + enabled: bool = Field(default=False, description="Enable artifact harvesting from the sandbox") + max_file_bytes: int = Field(default=50_000_000, description="Maximum size of a single harvested artifact") + max_total_bytes: int = Field(default=500_000_000, description="Maximum total artifact bytes per job (quota)") + max_file_count: int = Field(default=200, description="Maximum number of artifacts per job") + allow_extensions: tuple[str, ...] = Field( + default=_DEFAULT_ALLOW_EXTENSIONS, + description="Filename extension allowlist for captured artifacts.", + ) + + +class NetworkPolicy(BaseModel): + """Provider-neutral outbound network policy. + + One normalized shape every provider maps to its native mechanism (Modal's + ``block_network`` flag, OpenShell's gateway policy file, etc.): + + * ``blocked`` - no outbound network (the safe default). + * ``allowlist`` - only the hosts in ``allow`` are reachable. + * ``open`` - unrestricted (use only for trusted workloads). + """ + + mode: Literal["blocked", "allowlist", "open"] = Field( + default="blocked", + description="Outbound network policy mode enforced inside the sandbox.", + ) + allow: tuple[str, ...] = Field( + default=(), + description="Allowed hostnames/domains; only used (and required) when mode='allowlist'.", + ) + + @model_validator(mode="after") + def _validate_allowlist(self) -> NetworkPolicy: + """An allowlist policy is meaningless without hosts; fail loudly at config time.""" + if self.mode == "allowlist" and not self.allow: + raise ValueError("network.mode='allowlist' requires a non-empty network.allow list of hosts.") + return self + + +class ResourceLimits(BaseModel): + """Provider-neutral CPU/memory caps for the sandbox (opt-in). + + Both default to ``None`` (no limit), so an unset ``resources`` block changes nothing. + When a limit is set, the fail-closed capability gate refuses to run on a provider that + cannot enforce it (``supports_resource_limits``), rather than silently ignoring it. + Disk quotas are intentionally omitted: no current provider can enforce them, so a disk + field would be unenforceable. + """ + + cpu: float | None = Field(default=None, gt=0, description="Max CPU cores (provider-enforced).") + memory_mb: int | None = Field(default=None, gt=0, description="Max memory in MB (provider-enforced).") + + def any_set(self) -> bool: + """Whether any limit is requested (so the capability gate applies).""" + return self.cpu is not None or self.memory_mb is not None + + +class SandboxConfig(BaseModel): + """Provider-neutral configuration for a DeepAgents sandbox backend. + + Swapping providers is a config-only change: set ``provider`` and the matching + ``providers.`` block. Common fields (workdir, network, timeouts, artifact + capture, lifecycle scope) apply to every provider. + """ + + enabled: bool = Field(default=True, description="Whether the sandbox is active for this agent") + provider: str = Field(default="modal", description="Sandbox backend provider (must be registered).") + lifecycle_scope: Literal["job", "skill", "subagent"] = Field( + default="job", + description="Isolation scope for the sandbox. 'job' shares one sandbox across subagents.", + ) + workdir: str = Field(default=DEFAULT_WORKDIR, description="Writable working directory inside the sandbox") + network: NetworkPolicy = Field( + default_factory=NetworkPolicy, + description="Normalized outbound network policy. Legacy `block_network: bool` is lifted into this.", + ) + timeout: int = Field(default=1200, description="Maximum sandbox lifetime in seconds") + idle_timeout: int = Field(default=1800, description="Sandbox idle timeout in seconds") + resources: ResourceLimits = Field( + default_factory=ResourceLimits, + description="Optional CPU/memory caps; enforced only by providers declaring supports_resource_limits.", + ) + artifact_capture: ArtifactCaptureConfig = Field(default_factory=ArtifactCaptureConfig) + providers: SandboxProvidersConfig = Field(default_factory=SandboxProvidersConfig) + + @model_validator(mode="before") + @classmethod + def _lift_legacy_block_network(cls, data: Any) -> Any: + """Lift legacy ``block_network: bool`` into the normalized ``network`` policy. + + Explicit ``network`` always wins; ``block_network`` is only consulted when + ``network`` is not provided, so old configs keep working unchanged. + """ + if not isinstance(data, dict) or "block_network" not in data: + return data + data = dict(data) + legacy_block = data.pop("block_network") + if isinstance(legacy_block, str): + # Env-interpolated values arrive as strings. Reject anything unrecognized rather + # than mapping it to falsy, so a typo (e.g. "flase") cannot silently open egress + # on what the operator intended to be a network-blocked sandbox. + raw = legacy_block.strip().lower() + if raw in {"1", "true", "yes", "on"}: + legacy_block = True + elif raw in {"0", "false", "no", "off", ""}: + legacy_block = False + else: + raise ValueError( + f"Invalid block_network value {legacy_block!r}; expected a boolean " + "(true/false). Use the 'network' policy for finer control." + ) + if "network" not in data or data.get("network") is None: + data["network"] = {"mode": "blocked" if legacy_block else "open"} + return data + + @field_validator("provider") + @classmethod + def _provider_must_be_registered(cls, value: str) -> str: + """Validate the provider name against the registry (the single source of truth).""" + from .registry import is_registered + from .registry import registered_providers + + provider = value.lower() + if not is_registered(provider): + registered = ", ".join(registered_providers()) or "(none registered)" + raise ValueError(f"Unsupported sandbox provider: {value}. Registered providers: {registered}") + return provider + + @property + def block_network(self) -> bool: + """Back-compat accessor: any non-``open`` network policy blocks unrestricted egress.""" + return self.network.mode != "open" + + @property + def python_packages(self) -> tuple[str, ...]: + """Active provider's package list (Modal). Empty for providers without one.""" + return self.providers.modal.python_packages diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/providers/__init__.py b/src/aiq_agent/agents/deep_researcher/sandbox/providers/__init__.py new file mode 100644 index 000000000..5439a6fa6 --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/providers/__init__.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Built-in sandbox providers. + +Importing this package registers the built-in providers with the registry. Each +provider self-registers at import; adding a new provider is one new module here. +""" + +from __future__ import annotations + +from .modal import ModalSandboxProvider +from .openshell import OpenShellSandboxProvider + +__all__ = [ + "ModalSandboxProvider", + "OpenShellSandboxProvider", +] diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/providers/modal.py b/src/aiq_agent/agents/deep_researcher/sandbox/providers/modal.py new file mode 100644 index 000000000..3df751ef3 --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/providers/modal.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Modal sandbox provider (cloud example).""" + +from __future__ import annotations + +import logging +import re +import shlex +from typing import TYPE_CHECKING + +from deepagents.backends.sandbox import BaseSandbox + +from ..base import SandboxProvider +from ..capabilities import SandboxCapabilities +from ..registry import register_sandbox_provider + +if TYPE_CHECKING: + from ..config import SandboxConfig + +logger = logging.getLogger(__name__) + +_IMPORT_HINT = ( + "The Modal sandbox backend requires the `langchain-modal` and `modal` packages. " + "Install the updated AIQ dependencies and run `modal setup` before enabling a Modal sandbox." +) + + +def _validate_modal_sandbox_name(job_id: str) -> str: + """Validate that ``job_id`` is a legal Modal object name. + + Args: + job_id: Candidate sandbox name. + + Returns: + The validated name. + + Raises: + ValueError: If the name is too long, has illegal characters, or matches a + reserved Modal app-id shape. + """ + if len(job_id) > 64 or re.match(r"^[a-zA-Z0-9-_.]+$", job_id) is None or re.match(r"^ap-[a-zA-Z0-9]{22}$", job_id): + raise ValueError( + "Deep research job_id must be a valid Modal sandbox name: 64 characters or fewer, using only " + "alphanumeric characters, dashes, periods, and underscores." + ) + return job_id + + +def _is_modal_not_found_error(exc: Exception) -> bool: + """Return whether ``exc`` is Modal's typed NotFoundError (stale container).""" + try: + import modal + + return isinstance(exc, modal.exception.NotFoundError) + except ImportError: + return exc.__class__.__name__ == "NotFoundError" and exc.__class__.__module__.startswith("modal") + + +class ModalSandboxProvider(SandboxProvider): + """Job-scoped Modal backend. + + Modal enforces network blocking via the ``block_network`` create flag and + supports deterministic termination, so it declares those capabilities. + """ + + provider_name = "modal" + + def __init__(self, config: SandboxConfig, job_id: str) -> None: + """Initialize the provider, requiring the Modal SDK and adapter to import.""" + super().__init__(config, job_id) + try: + import langchain_modal # noqa: F401 + import modal # noqa: F401 + except ImportError as exc: + raise ImportError(_IMPORT_HINT) from exc + + @classmethod + def _scoped_name(cls, job_id: str) -> str: + """Return the validated, job-scoped Modal sandbox name.""" + return _validate_modal_sandbox_name(job_id) + + @property + def capabilities(self) -> SandboxCapabilities: + """Declare the guarantees the Modal backend can enforce.""" + return SandboxCapabilities( + supports_network_policy=True, + supports_resource_limits=True, + supports_artifact_download=True, + supports_cleanup=True, + ) + + def is_recoverable_error(self, exc: Exception) -> bool: + """Return whether the error is a missing-sandbox condition worth one retry.""" + return _is_modal_not_found_error(exc) + + def _create_session(self) -> BaseSandbox: + """Create a fresh, job-scoped Modal sandbox. + + Create-first semantics: unlike the legacy backend, this does NOT attach to + an existing sandbox by name as its primary path (which risked binding a new + job to a prior job's workspace). It creates fresh; only an + ``AlreadyExistsError`` (this job's own sandbox from earlier in the run, since + the name is the unique job id) falls back to attach. + """ + try: + import modal + from langchain_modal import ModalSandbox + except ImportError as exc: + raise ImportError(_IMPORT_HINT) from exc + + cfg = self.config + modal_cfg = cfg.providers.modal + app = modal.App.lookup(name=modal_cfg.app_name, create_if_missing=True) + + image = modal.Image.from_registry(modal_cfg.image) + if modal_cfg.python_packages: + image = image.pip_install(*modal_cfg.python_packages) + if cfg.workdir: + image = image.run_commands(f"mkdir -p {shlex.quote(cfg.workdir)}") + + # Opt-in resource caps (None => Modal default). The capability gate has already + # refused limits on providers that cannot enforce them, so passing them here is safe. + resource_kwargs: dict[str, object] = {} + if cfg.resources.cpu is not None: + resource_kwargs["cpu"] = cfg.resources.cpu + if cfg.resources.memory_mb is not None: + resource_kwargs["memory"] = cfg.resources.memory_mb + + try: + sandbox = modal.Sandbox.create( + app=app, + image=image, + workdir=cfg.workdir, + name=self.sandbox_name, + timeout=cfg.timeout, + idle_timeout=cfg.idle_timeout, + block_network=cfg.block_network, + **resource_kwargs, + ) + logger.info( + "Modal sandbox CREATED: name=%s image=%s workdir=%s timeout=%ds", + self.sandbox_name, + modal_cfg.image, + cfg.workdir, + cfg.timeout, + ) + except modal.exception.AlreadyExistsError: + sandbox = modal.Sandbox.from_name(modal_cfg.app_name, self.sandbox_name) + logger.info("Modal sandbox attached to this job's existing instance: name=%s", self.sandbox_name) + return ModalSandbox(sandbox=sandbox) + + +register_sandbox_provider("modal", ModalSandboxProvider) diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/providers/openshell.py b/src/aiq_agent/agents/deep_researcher/sandbox/providers/openshell.py new file mode 100644 index 000000000..4d870402f --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/providers/openshell.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""OpenShell sandbox provider (enterprise/on-prem). + +Governed, policy-enforced execution on local Docker/Podman/Kubernetes/microVM via +the OpenShell gateway. The deepagents ``BaseSandbox`` adapter is the official +``langchain-nvidia-openshell`` partner package (``OpenShellSandbox``), the same +adapter AI-Q PR #274 integrates. Both the ``openshell`` SDK and the adapter are +intentionally NOT declared in ``pyproject``; they are optional, ad-hoc +dependencies imported lazily, so this provider is never force-installed. + +Until ``langchain-ai/langchain-nvidia`` PR #303 publishes the adapter to PyPI, +install it from a git spec (see ``scripts/setup_openshell.sh`` / +``LANGCHAIN_NVIDIA_REPO``). +""" + +from __future__ import annotations + +import base64 +import logging +import os +import re +from typing import TYPE_CHECKING + +from deepagents.backends.protocol import FileDownloadResponse +from deepagents.backends.protocol import FileUploadResponse +from deepagents.backends.sandbox import BaseSandbox + +from ..base import SandboxProvider +from ..capabilities import SandboxCapabilities +from ..registry import register_sandbox_provider + +if TYPE_CHECKING: + from ..config import SandboxConfig + +logger = logging.getLogger(__name__) + +# Migration switch: when set truthy, delegate file transfer to the official adapter's +# upload_files/download_files instead of the local env-free shim. Use this to validate the +# upstream argv fix (langchain-ai/langchain-nvidia PR #303); once that ships, the shim and +# this switch can be removed and the adapter used unconditionally. +_ADAPTER_FILE_TRANSFER_ENV = "AIQ_OPENSHELL_ADAPTER_FILE_TRANSFER" + + +def _adapter_file_transfer_enabled() -> bool: + """True only when the toggle env var is an explicit truthy value (not just any string).""" + return os.getenv(_ADAPTER_FILE_TRANSFER_ENV, "").strip().lower() in {"1", "true", "yes", "on"} + + +# File-transfer bootstraps pass the path via argv (not env): OpenShell <=0.0.67 strips +# OPENSHELL_* env before exec, breaking the adapter's env-based transfer. We keep the +# adapter for execute and override only these two methods until the SDK propagates env. +# The download bootstrap fails closed before reading untrusted bytes: reject a symlink +# leaf (exit 5) or directory (exit 3), and read at most cap+1 bytes (exit 4 if over) so +# an oversized/out-of-tree file is never pulled into host memory. +_UPLOAD_CODE = ( + "import base64,os,sys;" + "p=sys.argv[1];" + "d=os.path.dirname(p);" + "(os.makedirs(d,exist_ok=True) if d else None);" + "open(p,'wb').write(base64.b64decode(sys.stdin.buffer.read()))" +) +_DOWNLOAD_CODE = ( + "import base64,os,sys;" + "p=sys.argv[1];" + "limit=int(sys.argv[2]);" + "root=os.path.realpath(sys.argv[3]);" + "rp=os.path.realpath(p);" + "(sys.exit(5) if not (rp==root or rp.startswith(root+os.sep)) else None);" + "(sys.exit(3) if os.path.isdir(rp) else None);" + "b=open(rp,'rb').read(limit+1);" + "(sys.exit(4) if len(b)>limit else None);" + "sys.stdout.write(base64.b64encode(b).decode())" +) + +# Bootstrap exit codes mapped to a download error reason (see _DOWNLOAD_CODE). +_DOWNLOAD_EXIT_ERRORS = {3: "is_directory", 4: "too_large", 5: "symlink_rejected"} + + +def _classify_fs_error(text: str) -> str: + """Map sandbox-side stderr to a deepagents FileOperationError literal.""" + lowered = text.lower() + if "no such file" in lowered or "file not found" in lowered or "filenotfounderror" in lowered: + return "file_not_found" + if "is a directory" in lowered or "isadirectoryerror" in lowered: + return "is_directory" + if "invalid" in lowered and "path" in lowered: + return "invalid_path" + return "permission_denied" + + +_OPENSHELL_IMPORT_HINT = ( + "The OpenShell sandbox provider requires the `openshell>=0.0.72,<0.1` SDK and the " + "`langchain-nvidia-openshell` adapter (published on PyPI). They are optional, ad-hoc " + "dependencies. Install them with `./scripts/setup_openshell.sh` (which installs " + "`langchain-nvidia-openshell` from PyPI; override the source via `LANGCHAIN_NVIDIA_REPO`), " + "and configure an OpenShell gateway before enabling this provider." +) + + +def _normalize_openshell_name(job_id: str, prefix: str = "aiq-deep-research") -> str: + """Normalize a job id into a DNS-style, length-bounded OpenShell sandbox name.""" + raw = f"{prefix}-{job_id}" if prefix else job_id + normalized = re.sub(r"[^a-z0-9-]+", "-", raw.lower()) + normalized = re.sub(r"-+", "-", normalized).strip("-") + return (normalized[:63].rstrip("-")) or prefix + + +def _is_openshell_not_found_error(exc: Exception) -> bool: + """Best-effort classification of OpenShell stale-sandbox errors.""" + text = str(exc).lower() + return "not found" in text and ("sandbox" in text or exc.__class__.__module__.startswith("openshell")) + + +class OpenShellSandboxProvider(SandboxProvider): + """OpenShell backend that attaches to a configured sandbox. + + OpenShell enforces filesystem/process/network policy at the gateway, so this + provider declares those capabilities. The SDK cannot apply or verify a policy + file while attaching: ``policy`` requires a pre-created named sandbox whose + policy is managed externally. + """ + + provider_name = "openshell" + + def __init__(self, config: SandboxConfig, job_id: str) -> None: + """Initialize the provider, requiring the OpenShell SDK and adapter to import.""" + super().__init__(config, job_id) + self._os_context: object | None = None + try: + import langchain_nvidia_openshell # noqa: F401 + import openshell # noqa: F401 + except ImportError as exc: + raise ImportError(_OPENSHELL_IMPORT_HINT) from exc + + @classmethod + def _scoped_name(cls, job_id: str) -> str: + """Return the OpenShell-safe sandbox name derived from the job id.""" + return _normalize_openshell_name(job_id) + + @property + def capabilities(self) -> SandboxCapabilities: + """Declare the gateway-enforced guarantees this provider supports.""" + return SandboxCapabilities( + supports_network_policy=True, + supports_network_allowlist=True, + supports_filesystem_policy=True, + supports_process_policy=True, + supports_artifact_download=True, + supports_cleanup=True, + ) + + def is_recoverable_error(self, exc: Exception) -> bool: + """Return whether the error is a missing-sandbox condition worth one retry.""" + return _is_openshell_not_found_error(exc) + + def upload_files(self, files: list[tuple[str, bytes]]) -> list[FileUploadResponse]: + """Upload files. Uses the local env-free shim by default (OpenShell <=0.0.67 strips + ``OPENSHELL_*`` env); set ``AIQ_OPENSHELL_ADAPTER_FILE_TRANSFER`` to delegate to the + official adapter (validates the upstream argv fix).""" + if _adapter_file_transfer_enabled(): + return self._call("upload_files", lambda session: session.upload_files(files), idempotent=True) + return self._call("upload_files", lambda _s: self._upload_files_envfree(files), idempotent=True) + + def download_files(self, paths: list[str]) -> list[FileDownloadResponse]: + """Download artifacts through the bounded, job-confined local shim.""" + return self._call("download_files", lambda _s: self._download_files_envfree(paths), idempotent=True) + + def _upload_files_envfree(self, files: list[tuple[str, bytes]]) -> list[FileUploadResponse]: + """Upload files via argv + stdin so no ``OPENSHELL_*`` env is required.""" + sandbox = self._os_context + responses: list[FileUploadResponse] = [] + for path, content in files: + if not path.startswith("/"): + responses.append(FileUploadResponse(path=path, error="invalid_path")) + continue + result = sandbox.exec( # type: ignore[union-attr] + ["python3", "-c", _UPLOAD_CODE, path], + stdin=base64.b64encode(content), + timeout_seconds=self.config.timeout, + ) + exit_code = getattr(result, "exit_code", 1) + error = None if exit_code == 0 else _classify_fs_error(getattr(result, "stderr", "") or "") + responses.append(FileUploadResponse(path=path, error=error)) + return responses + + def _download_files_envfree(self, paths: list[str]) -> list[FileDownloadResponse]: + """Download files via an argv bootstrap that enforces size/symlink limits in-sandbox.""" + sandbox = self._os_context + # Cap passed to the bootstrap so oversized files are refused before transfer. + max_bytes = self.config.artifact_capture.max_file_bytes + responses: list[FileDownloadResponse] = [] + for path in paths: + if not path.startswith("/"): + responses.append(FileDownloadResponse(path=path, content=None, error="invalid_path")) + continue + result = sandbox.exec( # type: ignore[union-attr] + # Confine resolved paths to this job's artifact directory. The configured + # workdir may be shared by several jobs in an attached named sandbox. + ["python3", "-c", _DOWNLOAD_CODE, path, str(max_bytes), self.artifact_dir], + timeout_seconds=self.config.timeout, + ) + exit_code = getattr(result, "exit_code", 1) + if exit_code != 0: + error = _DOWNLOAD_EXIT_ERRORS.get(exit_code) or _classify_fs_error(getattr(result, "stderr", "") or "") + responses.append(FileDownloadResponse(path=path, content=None, error=error)) + continue + # Validate base64 so stray stdout fails closed rather than storing corrupt bytes. + try: + content = base64.b64decode((getattr(result, "stdout", "") or "").strip().encode("ascii"), validate=True) + except ValueError: + responses.append(FileDownloadResponse(path=path, content=None, error="invalid_content")) + continue + responses.append(FileDownloadResponse(path=path, content=content, error=None)) + return responses + + def _create_session(self) -> BaseSandbox: + """Create/attach the OpenShell sandbox and wrap it in the official adapter. + + A configured ``policy`` requires a pre-created named sandbox, because the + SDK cannot apply policy files to anonymous sandboxes. + """ + try: + import openshell + from langchain_nvidia_openshell import OpenShellSandbox + except ImportError as exc: + raise ImportError(_OPENSHELL_IMPORT_HINT) from exc + + cfg = self.config + oscfg = cfg.providers.openshell + if oscfg.policy and not oscfg.sandbox_name: + raise ValueError( + "OpenShell `policy` requires `sandbox_name`. The SDK cannot apply a policy file to an " + "anonymous sandbox. Create the named sandbox with `openshell sandbox create --policy ` " + "first, then set providers.openshell.sandbox_name." + ) + + # Release any prior context (covers the recoverable-error reset path). + self._exit_context() + + sandbox_kwargs: dict[str, object] = { + "cluster": oscfg.gateway, + "delete_on_exit": oscfg.delete_on_exit, + "ready_timeout_seconds": oscfg.ready_timeout_seconds, + } + if oscfg.sandbox_name: + sandbox_kwargs["sandbox"] = oscfg.sandbox_name + + os_sandbox = openshell.Sandbox(**sandbox_kwargs) + os_sandbox.__enter__() + self._os_context = os_sandbox + backend = OpenShellSandbox(sandbox=os_sandbox, timeout=cfg.timeout, shell=oscfg.shell) + logger.info( + "OpenShell sandbox READY: id=%s gateway=%s sandbox_name=%s policy=%s", + backend.id, + oscfg.gateway, + oscfg.sandbox_name, + oscfg.policy, + ) + return backend + + def close(self) -> None: + """Terminate the session and exit the OpenShell context manager.""" + super().close() + self._exit_context() + + def _terminate_session(self, session: BaseSandbox | None) -> None: + """Close the adapter session and the owning OpenShell context on cancellation.""" + super()._terminate_session(session) + self._exit_context() + + def _exit_context(self) -> None: + """Exit the OpenShell context once, swallowing cleanup errors on the terminal path.""" + ctx = self._os_context + self._os_context = None + if ctx is not None and hasattr(ctx, "__exit__"): + try: + ctx.__exit__(None, None, None) + except Exception: # noqa: BLE001 - cleanup must never raise on the terminal path + logger.warning("OpenShell sandbox %s context cleanup failed", self.sandbox_name, exc_info=True) + + +register_sandbox_provider("openshell", OpenShellSandboxProvider) diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/registry.py b/src/aiq_agent/agents/deep_researcher/sandbox/registry.py new file mode 100644 index 000000000..bfd107631 --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/sandbox/registry.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Provider registry — the single, config-driven swap point. + +A provider registers its class under a name; the config's ``provider`` field is +validated against this registry; and :func:`create_sandbox_backend` is the one +place the runtime resolves a name to a backend. Adding a provider never edits this +module — the provider self-registers at import. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from .capabilities import verify_capabilities + +if TYPE_CHECKING: + from .base import SandboxProvider + from .config import SandboxConfig + +logger = logging.getLogger(__name__) + +#: Python entry-point group third-party packages use to contribute providers. +#: Declare in a distribution's pyproject as, e.g.:: +#: +#: [project.entry-points."aiq.sandbox_providers"] +#: mybox = "my_pkg.provider:MySandboxProvider" +SANDBOX_PROVIDER_ENTRY_POINT_GROUP = "aiq.sandbox_providers" + +_SANDBOX_PROVIDERS: dict[str, type[SandboxProvider]] = {} +_entry_points_loaded = False + + +def register_sandbox_provider(name: str, provider_cls: type[SandboxProvider]) -> None: + """Register a provider class under a config-facing name. + + Args: + name: Provider key used in ``SandboxConfig.provider`` and YAML config. + provider_cls: A ``SandboxProvider`` subclass constructed as ``(config, job_id)``. + """ + key = name.lower() + if key in _SANDBOX_PROVIDERS and _SANDBOX_PROVIDERS[key] is not provider_cls: + logger.warning("Overriding already-registered sandbox provider '%s'", key) + _SANDBOX_PROVIDERS[key] = provider_cls + + +def _load_entry_point_providers() -> None: + """Discover and register third-party providers via the entry-point group. + + Idempotent and best-effort: a broken or missing plugin is logged and skipped so + it can never take down provider resolution for the built-ins. Built-in providers + register eagerly at package import; this only adds external contributions. + """ + global _entry_points_loaded + if _entry_points_loaded: + return + _entry_points_loaded = True + + from importlib.metadata import entry_points + + try: + discovered = entry_points(group=SANDBOX_PROVIDER_ENTRY_POINT_GROUP) + except TypeError: # pragma: no cover - Python < 3.10 select-by-group fallback + discovered = entry_points().get(SANDBOX_PROVIDER_ENTRY_POINT_GROUP, []) + + for entry_point in discovered: + try: + provider_cls = entry_point.load() + except Exception: # noqa: BLE001 - a bad plugin must not break built-in resolution + logger.warning("Failed to load sandbox provider entry point '%s'", entry_point.name, exc_info=True) + continue + register_sandbox_provider(entry_point.name, provider_cls) + logger.info("Registered sandbox provider '%s' from entry point", entry_point.name) + + +def is_registered(name: str) -> bool: + """Return whether a provider name is registered (loading entry points on first use).""" + _load_entry_point_providers() + return name.lower() in _SANDBOX_PROVIDERS + + +def registered_providers() -> list[str]: + """Return the sorted list of registered provider names (built-in + entry-point).""" + _load_entry_point_providers() + return sorted(_SANDBOX_PROVIDERS) + + +def create_sandbox_backend(config: SandboxConfig, job_id: str) -> SandboxProvider: + """Resolve the configured provider to a backend and verify its capabilities. + + Args: + config: Resolved sandbox configuration. + job_id: Async job identifier used to scope the sandbox. + + Returns: + A constructed, capability-verified ``SandboxProvider``. + + Raises: + ValueError: If the provider name is not registered. + CapabilityError: If the provider cannot enforce a required guarantee (fail-closed). + """ + _load_entry_point_providers() + provider_cls = _SANDBOX_PROVIDERS.get(config.provider) + if provider_cls is None: + registered = ", ".join(registered_providers()) or "(none registered)" + raise ValueError(f"Unsupported sandbox provider: {config.provider}. Registered providers: {registered}") + backend = provider_cls(config, job_id) + verify_capabilities(config, backend.capabilities) + return backend diff --git a/src/aiq_agent/agents/deep_researcher/skills/research/chart-generation/SKILL.md b/src/aiq_agent/agents/deep_researcher/skills/research/chart-generation/SKILL.md new file mode 100644 index 000000000..67aafdb68 --- /dev/null +++ b/src/aiq_agent/agents/deep_researcher/skills/research/chart-generation/SKILL.md @@ -0,0 +1,169 @@ +--- +name: chart-generation +description: > + Use this skill to turn researched or computed numeric data into source-grounded + charts (PNG) plus the underlying CSV, by writing Python/matplotlib code and running + it in the job-scoped sandbox. The chart is harvested as a durable artifact and + embedded in the final report. + Triggers: "chart", "plot", "graph", "bar chart", "line chart", "visualize", + "trend over time", "compare visually", "figure". + Outputs: a PNG chart artifact, a CSV of the plotted data, and a manifest describing them. +--- + +# Chart Generation Skill + +Produce accurate, source-grounded charts using Python/matplotlib, save them as durable +artifacts, and embed them in the report by reference (never by pasting image data). + +## Required Execution Standard + +1. **Ground the data:** build the plotted rows from researched facts or `/shared/...` + inputs. Keep source URLs/notes alongside the values. +2. **Normalize units** before plotting (currencies, magnitudes, periods). +3. **Render with code:** call `execute` to run Python/matplotlib. Do not hand-draw or + fabricate charts. +4. **Write to the artifact directory:** save the PNG and its CSV under the exact + `sandbox_artifact_dir` given in your instructions (a per-job path such as + `/sandbox//aiq-artifacts`). Use that value verbatim - do NOT write to a bare + `/sandbox/aiq-artifacts`; the runtime only harvests files under `sandbox_artifact_dir`. +5. **Write a manifest** so the chart is harvested reliably (see below). +6. **Reference, do not embed bytes:** in the report, link the chart with + `![caption](artifact://.png)`. The runtime resolves this to the durable + artifact; never paste base64 image data into the report. + +## Data sufficiency (earn the chart) + +A chart confers authority, so it must be earned - never give unreliable data a cleaner +outfit. A polished chart of wrong or sparse numbers misleads more than it informs. + +1. **Source-anchored points only:** every plotted value must trace to a specific source + (the as-reported figure and its URL). Never plot a fabricated, guessed, or inferred + number as if it were reported; mark genuine estimates as estimates. +2. **Suppress misleading charts:** if a series is mostly missing (a majority of periods + undisclosed) or mixes metric definitions (e.g. "cash capex" vs "capex including finance + leases"), do NOT produce a trend chart. Present the table (which shows the gaps) and + state the limitation in one sentence instead. +3. **Show gaps honestly:** never interpolate or connect across missing periods. Plot only + the periods a series actually reports, and render estimates distinctly (e.g. hollow or + dashed markers) so they do not read as reported values. +4. **Prefer gap-tolerant forms:** grouped bars show missing periods as absent bars; favor + them over a connected line when series are uneven, since a line drawn across gaps + implies a trend that the data does not support. + +## Execution Flow + +1. Assemble the normalized rows (prefer explicit records embedded in the script). If the + inputs live in `/shared/...`, `read_file` them first and embed the values; sandbox + code cannot open `/shared/...`. +2. Use `write_file` to create the chart script under the `sandbox_workdir` from your + instructions (e.g. `/make_chart.py`), then `execute` that exact path. + `sandbox_workdir` is already per-job, so scripts there cannot collide with another job's + leftovers. Only ever execute a script you wrote this session. The script must: + - import pandas and matplotlib (use the non-interactive `Agg` backend), + - build the DataFrame, compute any derived metrics, + - set a single `ARTIFACT_DIR` to your `sandbox_artifact_dir` and write the chart + (`.png`), its data (`.csv`), and `manifest.json` there (see the example). +3. Inspect the `execute` output; if it fails, fix the script and re-run (max 2 retries). +4. In the report, embed the chart with `![](artifact://.png)` and cite the + original data sources in the surrounding text. + +## Placement and description in the report + +Each figure must appear where it is discussed, not buried in a file list: + +1. **Embed once, in context:** place the `![](artifact://.png)` line inside + the section that analyzes the figure (e.g. Results, Findings, or a Visualization + subsection) - immediately after the paragraph that introduces it. +2. **Describe it:** precede the embed with one sentence stating what the chart shows and the + takeaway (e.g. "The chart below compares 2025 resident population across the top five + states; California leads at roughly 3x Pennsylvania."). +3. **Reference by filename, never a raw path:** the way to show a figure is the + `![caption](artifact://.png)` token. Do NOT instead write the sandbox path + (e.g. `/.png`) as prose and expect it to render - a bare path + is not an image. +4. **One embed per artifact:** list supporting files (CSVs, manifests) by name in an + appendix if useful, but the chart itself must be embedded inline as above. + +## Manifest + +Write a `manifest.json` in your `sandbox_artifact_dir` so the runtime captures the chart +with metadata. Manifest `path` values must be absolute and inside your `sandbox_artifact_dir` +(the per-job path from your instructions, e.g. `/sandbox//aiq-artifacts`): + +```json +{ + "version": 1, + "artifacts": [ + { + "path": "/revenue_chart.png", + "kind": "image", + "title": "2024 Semiconductor Revenue Comparison", + "caption": "Revenue normalized to USD billions.", + "inline": true, + "source_files": ["/shared/semiconductor_revenue_normalized.csv"] + } + ] +} +``` + +## Example Script + +```python +import json +import os + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import pandas as pd + +# ARTIFACT_DIR MUST be the exact sandbox_artifact_dir from your instructions - a per-job +# path such as /sandbox//aiq-artifacts. Copy that value here verbatim. Do NOT use a +# bare /sandbox/aiq-artifacts: the runtime only harvests files under sandbox_artifact_dir. +ARTIFACT_DIR = "" # replace with the path from your instructions +os.makedirs(ARTIFACT_DIR, exist_ok=True) + +rows = [ + {"company": "ExampleCo", "revenue_usd_billions": 12.4, "source": "https://example.com/filing"}, + {"company": "SampleInc", "revenue_usd_billions": 9.1, "source": "https://example.com/10k"}, +] +df = pd.DataFrame(rows).sort_values("revenue_usd_billions", ascending=False) + +fig, ax = plt.subplots(figsize=(8, 5)) +ax.bar(df["company"], df["revenue_usd_billions"]) +ax.set_ylabel("Revenue (USD billions)") +ax.set_title("2024 Revenue Comparison") +fig.tight_layout() + +png_path = f"{ARTIFACT_DIR}/revenue_chart.png" +csv_path = f"{ARTIFACT_DIR}/revenue_chart.csv" +fig.savefig(png_path, dpi=150) +df.to_csv(csv_path, index=False) + +manifest = { + "version": 1, + "artifacts": [ + { + "path": png_path, + "kind": "image", + "title": "2024 Revenue Comparison", + "caption": "Revenue normalized to USD billions.", + "inline": True, + "source_files": [r["source"] for r in rows], + } + ], +} +with open(f"{ARTIFACT_DIR}/manifest.json", "w") as handle: + json.dump(manifest, handle) + +print(f"wrote {png_path}") +``` + +## Notes and Limitations + +- Use the `Agg` backend; the sandbox has no display. +- Keep charts legible: labeled axes, a title, and a legend when multiple series are shown. +- If matplotlib or pandas is unavailable, report that the sandbox image needs them rather + than fabricating a chart. +- Reference charts only by `artifact://`; the runtime assigns the durable id and + rewrites the reference for the UI, PDF export, and the packaged skill CLI. diff --git a/src/aiq_agent/agents/deep_researcher/skills/research/data-table-analysis/SKILL.md b/src/aiq_agent/agents/deep_researcher/skills/research/data-table-analysis/SKILL.md index 654db0f4c..2a52988b7 100644 --- a/src/aiq_agent/agents/deep_researcher/skills/research/data-table-analysis/SKILL.md +++ b/src/aiq_agent/agents/deep_researcher/skills/research/data-table-analysis/SKILL.md @@ -18,11 +18,27 @@ To ensure the calculation is reproducible and useful, you MUST: 5. **Return Text Outputs:** Include the markdown, CSV, or JSON output in your returned `ResearchNotes` (e.g. a `ResearchFinding`'s `evidence` and/or `narrative_notes`). Do not call `write_file`; `run_research_batch` persists your returned notes. 6. **Report Caveats:** Include assumptions, missing values, restatements, estimated figures, or non-comparable metrics in the output notes. +## Data honesty + +The table is the trustworthy, gap-aware deliverable that any downstream chart depends on, +so it must be honest about what is and isn't known: + +1. **Per-cell status:** treat each value as reported, estimate, or not disclosed. Leave + undisclosed cells explicitly empty (e.g. `—`); never fabricate or infer a number to + fill a gap, and never carry a prior period forward to hide one. +2. **One metric definition:** compare like with like. If sources use different definitions + (e.g. "cash paid for property and equipment" vs "capital expenditures including finance + leases"), keep them in separate rows/columns or pick one and label it - do not silently + blend definitions into a single series. +3. **Surface coverage:** in the notes, state how many cells are reported vs estimated vs + undisclosed, so the reader (and any chart built from this table) can judge how much + weight it bears. + ## Execution Flow 1. Gather candidate facts from researcher outputs, user-provided data, or source excerpts. -2. Create a normalized input table with one row per comparable observation. Prefer explicit CSV or JSON records embedded in the Python script. If the source rows are in `/shared/...`, call `read_file` first and embed the returned content in the script, or write a sandbox-local input file under `/workspace`. Sandbox code cannot open `/shared/...` directly. +2. Create a normalized input table with one row per comparable observation. Prefer explicit CSV or JSON records embedded in the Python script. If the source rows are in `/shared/...`, call `read_file` first and embed the returned content in the script, or write a sandbox-local input file under your sandbox working directory (`sandbox_workdir`; e.g. `/sandbox` on OpenShell or `/workspace` on Modal). Sandbox code cannot open `/shared/...` directly. 3. Call the `execute` tool with a Python command or script that: - imports pandas, @@ -31,7 +47,7 @@ To ensure the calculation is reproducible and useful, you MUST: - standardizes units and period labels, - computes the requested metrics, - prints markdown, CSV, JSON, and data-quality notes as text. - - uses `/workspace` for any sandbox-local input or output files. + - uses your sandbox working directory (`sandbox_workdir`) for any sandbox-local input or output files, and writes any script file at the job-unique path your instructions specify (the `_.py` form) so a shared sandbox never reuses a stale leftover from another job. - does not read from or write to `/shared/...` inside the sandbox process. 4. Inspect the `execute` output. If the code fails, fix the code and call `execute` again. Do not continue with hand-computed fallback tables unless the sandbox or pandas is unavailable. diff --git a/src/aiq_agent/common/citation_verification.py b/src/aiq_agent/common/citation_verification.py index 96cd5610c..13670678d 100644 --- a/src/aiq_agent/common/citation_verification.py +++ b/src/aiq_agent/common/citation_verification.py @@ -81,6 +81,7 @@ def __init__( unavailable_tools: list[str] | None = None, available_count: int = 0, ) -> None: + """Build the empty-registry error with agent type and tool-availability context.""" self.agent_type = agent_type self.unavailable_tools = unavailable_tools or [] self.available_count = available_count @@ -171,6 +172,7 @@ class SourceRegistry: """ def __init__(self) -> None: + """Initialize empty URL, parsed-URL, and citation-key indexes.""" self._urls: dict[str, SourceEntry] = {} self._parsed_urls: dict[str, _ParsedURL] = {} self._citation_keys: list[SourceEntry] = [] @@ -667,6 +669,86 @@ def _format_registry_reference(num: int, entry: SourceEntry) -> str | None: return None +def _normalize_reference_title(text: str) -> str: + """Normalize a source-line title for exact-match backfill comparison. + + Strips any trailing URL, trailing parenthetical (e.g. ``(Internal)``), + markdown emphasis, and a trailing ``:`` separator, then lowercases and + collapses whitespace so a writer line and its registry entry compare equal. + """ + cleaned = _URL_IN_LINE_RE.sub("", text) + cleaned = re.sub(r"\s*\(.*?\)\s*$", "", cleaned) + cleaned = re.sub(r"\*+", "", cleaned) + cleaned = cleaned.strip().rstrip(":").strip() + return re.sub(r"\s+", " ", cleaned).lower() + + +def _build_reference_title_index( + reference_sources: Sequence[SourceEntry] | None, + registry: SourceRegistry, +) -> dict[str, tuple[str | None, str | None]]: + """Index writer-facing source titles to their registry-validated target. + + Used to repair URL-less ``[N] Title`` lines (the writer dropped the ``: url`` + suffix) from the same captured-source list the writer was shown. Each title + is validated against the registry exactly as an inline target would be, and + titles shared by more than one distinct source are dropped so an ambiguous + title is never auto-resolved. + + Args: + reference_sources: Writer-facing source list, or None. + registry: SourceRegistry the targets must resolve against. + + Returns: + Mapping of normalized title -> ``(canonical_url, citation_key)`` with + exactly one tuple value set. + """ + if not reference_sources: + return {} + index: dict[str, tuple[str | None, str | None]] = {} + ambiguous: set[str] = set() + for entry in reference_sources: + if not entry.title: + continue + key = _normalize_reference_title(entry.title) + if not key: + continue + target: tuple[str | None, str | None] | None = None + if entry.url: + canonical = registry.resolve_url(entry.url) + if canonical: + target = (canonical, None) + elif entry.citation_key and registry.has_citation_key(entry.citation_key): + target = (None, entry.citation_key) + if target is None: + continue + if key in index and index[key] != target: + ambiguous.add(key) + continue + index[key] = target + for key in ambiguous: + index.pop(key, None) + return index + + +def _backfill_reference_target( + ref_text: str, + reference_index: dict[str, tuple[str | None, str | None]], +) -> tuple[str | None, str | None] | None: + """Resolve a URL-less source line's title to a registry-backed target. + + Exact normalized-title match only. Aggregate-style labels (multiple sources + joined with ``;``) are never matched, so they keep stripping as before. + + Returns: + ``(canonical_url, citation_key)`` when the title uniquely maps to a + registry source, otherwise None. + """ + if ";" in ref_text: + return None + return reference_index.get(_normalize_reference_title(ref_text)) + + def _normalize_citation_syntax(report_text: str) -> str: """Normalize citation bracket variants before verification/sanitization.""" report_text = report_text.replace("【", "[").replace("】", "]") @@ -880,6 +962,12 @@ def verify_citations( valid_citations: list[dict] = [] removed_citations: list[dict] = [] url_replacements: dict[str, str] = {} # garbled_url -> canonical_url + line_replacements: dict[str, str] = {} # url-less line -> backfilled canonical line + + # Exact-title index over the writer-facing source list, so a "[N] Title" + # line whose ": url" suffix the writer dropped can be repaired from the same + # captured registry instead of stripped (recall fix, precision unchanged). + reference_index = _build_reference_title_index(reference_sources, registry) for line_match in _CITATION_LINE_RE.finditer(ref_section): num = int(line_match.group(1)) @@ -914,6 +1002,22 @@ def verify_citations( removed_citations.append({"number": num, "line": full_line, "reason": "citation_key_not_in_registry"}) continue + # Backfill: the writer emitted "[N] Title" but dropped the verified URL. + # Recover the target from the writer-facing source list (exact-title, + # unique match) and rewrite the line to canonical "[N] Title: url" form. + backfilled = _backfill_reference_target(ref_text, reference_index) + if backfilled: + canonical_url, citation_key = backfilled + if canonical_url: + rewritten = f"[{num}] {ref_text}: {canonical_url}" + line_replacements[full_line] = rewritten + logger.info("[CitationVerify] [%d] BACKFILL — %s", num, canonical_url) + valid_citations.append({"number": num, "url": canonical_url, "citation_key": None, "line": rewritten}) + else: + logger.info("[CitationVerify] [%d] BACKFILL — %s", num, citation_key) + valid_citations.append({"number": num, "url": None, "citation_key": citation_key, "line": full_line}) + continue + # Neither URL nor recognizable citation key logger.debug("[CitationVerify] [%d] REMOVE — unverifiable: %s", num, ref_text[:80]) removed_citations.append({"number": num, "line": full_line, "reason": "unverifiable"}) @@ -961,6 +1065,11 @@ def verify_citations( for garbled, canonical in url_replacements.items(): ref_section = ref_section.replace(garbled, canonical) + # Apply backfilled lines (url-less "[N] Title" -> "[N] Title: url"). + if line_replacements: + for original, rewritten in line_replacements.items(): + ref_section = ref_section.replace(original, rewritten) + removed_numbers = {c["number"] for c in removed_citations} # Remove invalid (and duplicate) source lines from the source section. @@ -1036,7 +1145,7 @@ def verify_citations( _BARE_URL_RE = re.compile(r"https?://[^\s<>\"'\]]+") # Body URL patterns (used by sanitize_report) -_MD_LINK_RE = re.compile(r"\[([^\]]*)\]\(\s*\w+://[^\s)]+\)") +_MD_LINK_RE = re.compile(r"\[([^\]]*)\]\(\s*(\w+://[^\s)]+)\)") _BODY_URL_RE = re.compile(r"\w+://[^\s<>\"'\]]+") @@ -1102,8 +1211,14 @@ def sanitize_report(report_text: str) -> ReportSanitizationResult: url_to_citation[_normalize_url(url_m.group(0).rstrip(_URL_TRIM_CHARS))] = num def _replace_body_url(match: re.Match) -> str: + """Replace a bare body URL with its citation number, or strip it; keep artifact refs.""" nonlocal body_urls_removed, body_urls_replaced url = match.group(0).rstrip(_URL_TRIM_CHARS) + # Preserve internal artifact references (e.g. embedded chart images). They use the + # non-http ``artifact://`` scheme and are validated/rewritten downstream by + # ArtifactManager.resolve_report_references, so they must survive sanitization. + if url.startswith("artifact://"): + return match.group(0) normalized = _normalize_url(url) if normalized in url_to_citation: body_urls_replaced += 1 @@ -1111,8 +1226,17 @@ def _replace_body_url(match: re.Match) -> str: body_urls_removed += 1 return "" - # Collapse markdown links to display text - cleaned_body = _MD_LINK_RE.sub(r"\1", body) + # Collapse markdown links to display text, but preserve internal artifact:// image + # references verbatim so embedded charts survive into the final report. + def _collapse_md_link(match: re.Match) -> str: + """Collapse a markdown link to its display text, preserving ``artifact://`` targets.""" + # Preserve only when the link destination is an artifact ref; a label that merely + # contains "artifact://" must still collapse so stray URLs don't leak into the body. + if match.group(2).startswith("artifact://"): + return match.group(0) + return match.group(1) + + cleaned_body = _MD_LINK_RE.sub(_collapse_md_link, body) # Replace matching bare URLs with [N], strip the rest cleaned_body = _BODY_URL_RE.sub(_replace_body_url, cleaned_body) # Clean up leftover empty parentheses and extra spaces diff --git a/tests/aiq_agent/agents/deep_researcher/sandbox/__init__.py b/tests/aiq_agent/agents/deep_researcher/sandbox/__init__.py new file mode 100644 index 000000000..3bcc1c39b --- /dev/null +++ b/tests/aiq_agent/agents/deep_researcher/sandbox/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/aiq_agent/agents/deep_researcher/sandbox/test_artifacts.py b/tests/aiq_agent/agents/deep_researcher/sandbox/test_artifacts.py new file mode 100644 index 000000000..6a5628ed2 --- /dev/null +++ b/tests/aiq_agent/agents/deep_researcher/sandbox/test_artifacts.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the artifact runtime: manifest parsing, MIME sniffing, the harvest +validation pipeline (traversal/extension/size/quota/dedup/scan), and the SQL store.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from typing import Any + +from aiq_agent.agents.deep_researcher.sandbox.artifacts import Artifact +from aiq_agent.agents.deep_researcher.sandbox.artifacts import ArtifactKind +from aiq_agent.agents.deep_researcher.sandbox.artifacts import ArtifactManager +from aiq_agent.agents.deep_researcher.sandbox.artifacts import SqlArtifactStore +from aiq_agent.agents.deep_researcher.sandbox.artifacts import parse_manifest +from aiq_agent.agents.deep_researcher.sandbox.artifacts.manager import _normalize_posix +from aiq_agent.agents.deep_researcher.sandbox.artifacts.manager import _sniff_mime +from aiq_agent.agents.deep_researcher.sandbox.config import ArtifactCaptureConfig + +_ARTIFACT_DIR = "/workspace/aiq-artifacts" +_PNG = b"\x89PNG\r\n\x1a\n" + b"\x00" * 64 + + +class _FakeBackend: + """Minimal BaseSandbox stand-in mapping sandbox paths to bytes.""" + + def __init__(self, files: dict[str, bytes]) -> None: + self.files = files + + def download_files(self, paths: list[str]) -> list[Any]: + return [ + SimpleNamespace(path=p, content=self.files.get(p), error=None if p in self.files else "not found") + for p in paths + ] + + def execute(self, command: str, *, timeout: int | None = None) -> Any: + return SimpleNamespace(output="\n".join(self.files), exit_code=0) + + +def _manifest_bytes(path: str, kind: str = "image") -> bytes: + return json.dumps({"version": 1, "artifacts": [{"path": path, "kind": kind, "inline": True}]}).encode("utf-8") + + +def _make_manager(store: SqlArtifactStore, files: dict[str, bytes], **capture: Any) -> tuple[ArtifactManager, list]: + emitted: list = [] + config = ArtifactCaptureConfig(enabled=True, **capture) + manager = ArtifactManager( + job_id="job-1", + backend=_FakeBackend(files), + store=store, + config=config, + artifact_dir=_ARTIFACT_DIR, + emit=emitted.append, + ) + return manager, emitted + + +class TestManifest: + def test_parse_valid(self) -> None: + manifest = parse_manifest(_manifest_bytes(f"{_ARTIFACT_DIR}/c.png").decode()) + assert manifest is not None + assert manifest.artifacts[0].path == f"{_ARTIFACT_DIR}/c.png" + + def test_parse_invalid_returns_none(self) -> None: + assert parse_manifest("not json{") is None + + +class TestSniffMime: + def test_png_by_magic(self) -> None: + assert _sniff_mime(_PNG, "x.bin") == "image/png" + + def test_csv_by_extension(self) -> None: + assert _sniff_mime(b"a,b\n1,2\n", "data.csv") == "text/csv" + + +class TestNormalizePosix: + def test_absolute_path_has_single_leading_slash(self) -> None: + from pathlib import PurePosixPath + + # The absolute-root sentinel must not be re-appended as a path segment. + assert _normalize_posix(PurePosixPath("/sandbox/aiq-artifacts")) == "/sandbox/aiq-artifacts" + + def test_collapses_dot_and_parent_segments(self) -> None: + from pathlib import PurePosixPath + + assert _normalize_posix(PurePosixPath("/sandbox/./sub/../aiq-artifacts")) == "/sandbox/aiq-artifacts" + + def test_relative_path_has_no_leading_slash(self) -> None: + from pathlib import PurePosixPath + + assert _normalize_posix(PurePosixPath("sub/aiq-artifacts")) == "sub/aiq-artifacts" + + +class TestHarvest: + def test_captures_manifest_artifact(self, tmp_path: Any) -> None: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + png_path = f"{_ARTIFACT_DIR}/chart.png" + files = {f"{_ARTIFACT_DIR}/manifest.json": _manifest_bytes(png_path), png_path: _PNG} + manager, emitted = _make_manager(store, files) + + captured = manager.final_harvest() + + assert len(captured) == 1 + assert captured[0].mime_type == "image/png" + assert captured[0].kind == ArtifactKind.IMAGE + assert store.list("job-1")[0].filename == "chart.png" + assert emitted and emitted[0]["type"] == "artifact" + assert "content" not in emitted[0] # bytes never in the event payload + + def test_rejects_path_traversal(self, tmp_path: Any) -> None: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + evil = "/etc/passwd.png" + files = {f"{_ARTIFACT_DIR}/manifest.json": _manifest_bytes(evil), evil: _PNG} + manager, _ = _make_manager(store, files) + assert manager.final_harvest() == [] + + def test_enforces_extension_allowlist(self, tmp_path: Any) -> None: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + exe = f"{_ARTIFACT_DIR}/evil.exe" + files = {f"{_ARTIFACT_DIR}/manifest.json": _manifest_bytes(exe), exe: _PNG} + manager, _ = _make_manager(store, files) + assert manager.final_harvest() == [] + + def test_enforces_size_cap(self, tmp_path: Any) -> None: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + png_path = f"{_ARTIFACT_DIR}/chart.png" + files = {f"{_ARTIFACT_DIR}/manifest.json": _manifest_bytes(png_path), png_path: _PNG} + manager, _ = _make_manager(store, files, max_file_bytes=8) + assert manager.final_harvest() == [] + + def test_enforces_quota(self, tmp_path: Any) -> None: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + a = f"{_ARTIFACT_DIR}/a.png" + b = f"{_ARTIFACT_DIR}/b.png" + manifest = json.dumps( + {"version": 1, "artifacts": [{"path": a, "kind": "image"}, {"path": b, "kind": "image"}]} + ).encode("utf-8") + files = {f"{_ARTIFACT_DIR}/manifest.json": manifest, a: _PNG, b: _PNG + b"x"} + manager, _ = _make_manager(store, files, max_file_count=1) + captured = manager.final_harvest() + assert len(captured) == 1 + + def test_dedups_identical_content(self, tmp_path: Any) -> None: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + png_path = f"{_ARTIFACT_DIR}/chart.png" + files = {f"{_ARTIFACT_DIR}/manifest.json": _manifest_bytes(png_path), png_path: _PNG} + manager, _ = _make_manager(store, files) + manager.final_harvest() + manager.final_harvest() # same bytes again + assert len(store.list("job-1")) == 1 + + def test_scan_fallback_without_manifest(self, tmp_path: Any) -> None: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + png_path = f"{_ARTIFACT_DIR}/chart.png" + files = {png_path: _PNG} # no manifest + manager, _ = _make_manager(store, files) + captured = manager.final_harvest() # job_end allows scan + assert len(captured) == 1 + assert captured[0].filename == "chart.png" + + def test_final_harvest_unions_manifest_and_scan(self, tmp_path: Any) -> None: + # A manifest that declares only the PNG must not hide a sibling CSV at job end. + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + png_path = f"{_ARTIFACT_DIR}/chart.png" + csv_path = f"{_ARTIFACT_DIR}/chart.csv" + files = { + f"{_ARTIFACT_DIR}/manifest.json": _manifest_bytes(png_path), + png_path: _PNG, + csv_path: b"state,pop\nCA,39431263\n", + } + manager, _ = _make_manager(store, files) + + captured = manager.final_harvest() + + names = sorted(a.filename for a in captured) + assert names == ["chart.csv", "chart.png"] + + def test_rejects_mime_spoof(self, tmp_path: Any) -> None: + # A file claiming .png but with non-image content must be rejected. + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + png_path = f"{_ARTIFACT_DIR}/chart.png" + files = {f"{_ARTIFACT_DIR}/manifest.json": _manifest_bytes(png_path), png_path: b"#!/bin/sh\nrm -rf /\n"} + manager, _ = _make_manager(store, files) + assert manager.final_harvest() == [] + + def test_rejects_svg_fail_closed(self, tmp_path: Any) -> None: + # SVG cannot be reliably sanitized (javascript: URIs, , CSS payloads), + # and the content endpoint serves bytes as the stored MIME, so SVG is rejected at + # harvest rather than partially cleaned and stored (stored-XSS prevention). + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + svg_path = f"{_ARTIFACT_DIR}/diagram.svg" + svg = b'' + files = {f"{_ARTIFACT_DIR}/manifest.json": _manifest_bytes(svg_path), svg_path: svg} + manager, _ = _make_manager(store, files) + + assert manager.final_harvest() == [] + + +class TestStore: + def _artifact(self) -> Artifact: + return Artifact( + artifact_id="art_" + "a" * 32, + job_id="job-1", + kind=ArtifactKind.IMAGE, + mime_type="image/png", + filename="chart.png", + sandbox_path=f"{_ARTIFACT_DIR}/chart.png", + storage_uri="", + sha256="d" * 64, + size_bytes=len(_PNG), + ) + + def test_put_get_list_open(self, tmp_path: Any) -> None: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + stored = store.put(self._artifact(), _PNG) + assert stored.status.value == "available" + assert store.get("job-1", stored.artifact_id) is not None + assert len(store.list("job-1")) == 1 + assert b"".join(store.open_bytes("job-1", stored.artifact_id)) == _PNG + + def test_dedup_by_digest(self, tmp_path: Any) -> None: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + first = store.put(self._artifact(), _PNG) + # Same content/digest but a DIFFERENT artifact_id: dedup must key on (job_id, sha256), + # not on artifact_id, so the second put returns the existing row instead of inserting + # a duplicate. Reusing the same id here would let an id-based regression pass silently. + duplicate = self._artifact().model_copy(update={"artifact_id": "art_" + "b" * 32}) + again = store.put(duplicate, _PNG) + assert again.artifact_id == first.artifact_id + assert len(store.list("job-1")) == 1 + + +class TestEnsureInlineArtifactsEmbedded: + """The safety net that surfaces produced inline figures the report forgot to embed.""" + + def _store_with(self, tmp_path: Any, *artifacts: Artifact) -> SqlArtifactStore: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + for index, artifact in enumerate(artifacts): + store.put(artifact, _PNG + bytes([index])) + return store + + def _image(self, artifact_id: str, *, inline: bool = True) -> Artifact: + return Artifact( + artifact_id=artifact_id, + job_id="job-1", + kind=ArtifactKind.IMAGE, + mime_type="image/png", + filename=f"{artifact_id}.png", + sandbox_path=f"{_ARTIFACT_DIR}/{artifact_id}.png", + storage_uri="", + sha256=artifact_id.ljust(64, "0"), + size_bytes=len(_PNG), + title="Chart Title", + caption="A caption", + inline=inline, + ) + + def test_appends_orphan_inline_image(self, tmp_path: Any) -> None: + artifact_id = "art_" + "a" * 32 + store = self._store_with(tmp_path, self._image(artifact_id)) + manager, _ = _make_manager(store, {}) + + result = manager.ensure_inline_artifacts_embedded("# Report\n\nNo figures here.\n") + + assert "## Figures" in result + assert f"![A caption](artifact://{artifact_id})" in result + + def test_does_not_duplicate_referenced_image(self, tmp_path: Any) -> None: + artifact_id = "art_" + "b" * 32 + store = self._store_with(tmp_path, self._image(artifact_id)) + manager, _ = _make_manager(store, {}) + markdown = f"# Report\n\n![Inline](artifact://{artifact_id})\n" + + result = manager.ensure_inline_artifacts_embedded(markdown) + + assert result == markdown + assert "## Figures" not in result + + def test_append_artifact_index_lists_all_artifacts(self, tmp_path: Any) -> None: + png = self._image("art_" + "e" * 32) + csv = Artifact( + artifact_id="art_" + "f" * 32, + job_id="job-1", + kind=ArtifactKind.TABLE, + mime_type="text/csv", + filename="chart.csv", + sandbox_path=f"{_ARTIFACT_DIR}/chart.csv", + storage_uri="", + sha256="f" * 64, + size_bytes=8, + caption="Plotted values", + ) + store = self._store_with(tmp_path, png, csv) + manager, _ = _make_manager(store, {}) + + result = manager.append_artifact_index("# Report\n\nBody.\n") + + assert "## Generated Artifacts" in result + assert "`art_eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.png`" in result + assert "`chart.csv` - Plotted values (generated in the analysis sandbox)" in result + + def test_append_artifact_index_noop_without_artifacts(self, tmp_path: Any) -> None: + store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") + manager, _ = _make_manager(store, {}) + assert manager.append_artifact_index("# Report\n") == "# Report\n" + + def test_skips_non_inline_and_non_image(self, tmp_path: Any) -> None: + non_inline = self._image("art_" + "c" * 32, inline=False) + csv = Artifact( + artifact_id="art_" + "d" * 32, + job_id="job-1", + kind=ArtifactKind.TABLE, + mime_type="text/csv", + filename="data.csv", + sandbox_path=f"{_ARTIFACT_DIR}/data.csv", + storage_uri="", + sha256="d" * 64, + size_bytes=8, + inline=True, + ) + store = self._store_with(tmp_path, non_inline, csv) + manager, _ = _make_manager(store, {}) + + result = manager.ensure_inline_artifacts_embedded("# Report\n") + + assert result == "# Report\n" diff --git a/tests/aiq_agent/agents/deep_researcher/sandbox/test_openshell_provider.py b/tests/aiq_agent/agents/deep_researcher/sandbox/test_openshell_provider.py new file mode 100644 index 000000000..d8a1fef7f --- /dev/null +++ b/tests/aiq_agent/agents/deep_researcher/sandbox/test_openshell_provider.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""OpenShell provider tests: the env-free upload/download shim. + +OpenShell 0.0.57's exec does not propagate ``env`` to the child process, which +breaks the official adapter's env-var file-transfer bootstraps. Our provider +overrides ``upload_files``/``download_files`` to pass the path via argv (and data +via stdin). These tests assert that contract with a fake sandbox, so they require +only the optional SDK + adapter to be importable. +""" + +from __future__ import annotations + +import base64 +from dataclasses import dataclass +from typing import Any +from unittest.mock import MagicMock + +import pytest + +pytest.importorskip("openshell") +pytest.importorskip("langchain_nvidia_openshell") + +from aiq_agent.agents.deep_researcher.sandbox.config import SandboxConfig # noqa: E402 +from aiq_agent.agents.deep_researcher.sandbox.providers.openshell import OpenShellSandboxProvider # noqa: E402 + + +@dataclass +class _ExecResult: + exit_code: int + stdout: str = "" + stderr: str = "" + + +class _FakeOpenShellSandbox: + """Records exec calls and returns scripted results.""" + + id = "fake-os-id" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + self.result = _ExecResult(exit_code=0) + self.exit_calls = 0 + + def exec(self, command, **kwargs): # noqa: ANN001 - mirrors openshell.Sandbox.exec + self.calls.append({"command": list(command), **kwargs}) + return self.result + + def __exit__(self, *_args: object) -> None: + self.exit_calls += 1 + + +def _provider() -> OpenShellSandboxProvider: + cfg = SandboxConfig( + provider="openshell", + network={"mode": "blocked"}, + providers={"openshell": {"sandbox_name": "demo", "delete_on_exit": False}}, + ) + provider = OpenShellSandboxProvider(cfg, "job-1") + # Avoid real session creation: a non-None _session short-circuits _session_or_create. + provider._session = MagicMock() # type: ignore[assignment] + return provider + + +def test_upload_passes_path_via_argv_and_data_via_stdin_no_env() -> None: + provider = _provider() + fake = _FakeOpenShellSandbox() + provider._os_context = fake + + result = provider.upload_files([("/sandbox/x.py", b"print('hi')")]) + + assert result[0].error is None + call = fake.calls[0] + assert call["command"][0] == "python3" and call["command"][-1] == "/sandbox/x.py" + assert call["stdin"] == base64.b64encode(b"print('hi')") + assert "env" not in call # the whole point: never rely on exec env propagation + + +def test_upload_classifies_failure() -> None: + provider = _provider() + fake = _FakeOpenShellSandbox() + fake.result = _ExecResult(exit_code=1, stderr="No such file or directory") + provider._os_context = fake + + result = provider.upload_files([("/sandbox/x.py", b"data")]) + assert result[0].error == "file_not_found" + + +def test_upload_rejects_relative_path() -> None: + provider = _provider() + provider._os_context = _FakeOpenShellSandbox() + result = provider.upload_files([("relative.py", b"data")]) + assert result[0].error == "invalid_path" + + +def test_download_passes_path_via_argv_and_decodes_base64() -> None: + provider = _provider() + fake = _FakeOpenShellSandbox() + fake.result = _ExecResult(exit_code=0, stdout=base64.b64encode(b"chart-bytes").decode()) + provider._os_context = fake + + artifact_path = f"{provider.artifact_dir}/chart.png" + result = provider.download_files([artifact_path]) + + assert result[0].error is None + assert result[0].content == b"chart-bytes" + call = fake.calls[0] + # Path is passed positionally via argv (with the size cap appended); never via env. + assert artifact_path in call["command"] + assert call["command"][-1] == provider.artifact_dir + assert "env" not in call + + +def test_download_uses_confined_shim_when_adapter_transfer_enabled(monkeypatch: pytest.MonkeyPatch) -> None: + provider = _provider() + fake = _FakeOpenShellSandbox() + fake.result = _ExecResult(exit_code=0, stdout=base64.b64encode(b"chart-bytes").decode()) + provider._os_context = fake + provider._session.download_files.side_effect = AssertionError("adapter download bypassed confinement") # type: ignore[union-attr] + monkeypatch.setenv("AIQ_OPENSHELL_ADAPTER_FILE_TRANSFER", "1") + + result = provider.download_files([f"{provider.artifact_dir}/chart.png"]) + + assert result[0].content == b"chart-bytes" + assert fake.calls[0]["command"][-1] == provider.artifact_dir + + +def test_download_is_directory_exit_code() -> None: + provider = _provider() + fake = _FakeOpenShellSandbox() + fake.result = _ExecResult(exit_code=3, stderr="") + provider._os_context = fake + + result = provider.download_files(["/sandbox"]) + assert result[0].content is None + assert result[0].error == "is_directory" # _DOWNLOAD_CODE exits 3 specifically for a directory + + +def test_download_passes_size_cap_via_argv() -> None: + provider = _provider() + fake = _FakeOpenShellSandbox() + fake.result = _ExecResult(exit_code=0, stdout=base64.b64encode(b"x").decode()) + provider._os_context = fake + + provider.download_files(["/sandbox/aiq-artifacts/chart.png"]) + + # The artifact size cap is passed to the bootstrap so it can refuse oversized files + # before reading them into host memory. + assert str(provider.config.artifact_capture.max_file_bytes) in fake.calls[0]["command"] + + +def test_download_rejects_oversized_and_symlink() -> None: + provider = _provider() + fake = _FakeOpenShellSandbox() + provider._os_context = fake + + fake.result = _ExecResult(exit_code=4, stderr="") + assert provider.download_files(["/sandbox/aiq-artifacts/huge.bin"])[0].error == "too_large" + + fake.result = _ExecResult(exit_code=5, stderr="") + assert provider.download_files(["/sandbox/aiq-artifacts/evil.png"])[0].error == "symlink_rejected" + + +def test_download_rejects_non_base64_stdout() -> None: + provider = _provider() + fake = _FakeOpenShellSandbox() + fake.result = _ExecResult(exit_code=0, stdout="not valid base64 !!!") + provider._os_context = fake + + result = provider.download_files(["/sandbox/aiq-artifacts/chart.png"]) + assert result[0].content is None + assert result[0].error == "invalid_content" + + +def test_terminate_exits_openshell_context_once() -> None: + provider = _provider() + fake = _FakeOpenShellSandbox() + provider._os_context = fake + + provider.terminate() + provider.terminate() + + assert fake.exit_calls == 1 + assert provider._os_context is None diff --git a/tests/aiq_agent/agents/deep_researcher/sandbox/test_provider_compliance.py b/tests/aiq_agent/agents/deep_researcher/sandbox/test_provider_compliance.py new file mode 100644 index 000000000..cb8129e84 --- /dev/null +++ b/tests/aiq_agent/agents/deep_researcher/sandbox/test_provider_compliance.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SandboxProvider compliance suite. + +Mirrors the knowledge-layer adapter compliance harness: every registered provider +must satisfy the same contract. Providers whose optional SDK is not installed +(e.g. OpenShell) are skipped rather than failed, so this runs without a live gateway. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from aiq_agent.agents.deep_researcher.sandbox import SandboxCapabilities +from aiq_agent.agents.deep_researcher.sandbox import SandboxConfig +from aiq_agent.agents.deep_researcher.sandbox import SandboxProvider +from aiq_agent.agents.deep_researcher.sandbox import create_sandbox_backend + +_BUILTIN_PROVIDERS = ("modal", "openshell") + + +def assert_provider_contract(provider: SandboxProvider) -> None: + """Assert a provider honors the SandboxProvider contract (no live session needed).""" + # Declared capabilities are a real SandboxCapabilities model. + assert isinstance(provider.capabilities, SandboxCapabilities) + + # Identity is a non-empty job-scoped name before any session exists. + assert isinstance(provider.sandbox_name, str) and provider.sandbox_name + assert provider.id == provider.sandbox_name + + # Error classification is conservative for unrelated errors. + assert provider.is_recoverable_error(ValueError("unrelated")) is False + + # close() is idempotent and safe with no live session. + provider.close() + provider.close() + + # The shared resilience path delegates to the session created by _create_session, and + # prepares the per-job workspace (idempotent mkdir -p) before the first real call. + session = MagicMock() + session.execute.return_value = "ok" + provider._create_session = lambda: session # type: ignore[method-assign] + assert provider.execute("echo ok", timeout=5) == "ok" + first_cmd = session.execute.call_args_list[0].args[0] + assert first_cmd.startswith("mkdir -p") and provider.workdir in first_cmd + session.execute.assert_called_with("echo ok", timeout=5) # most recent call is the real command + + +@pytest.mark.parametrize("provider_name", _BUILTIN_PROVIDERS) +def test_builtin_provider_compliance(provider_name: str) -> None: + config = SandboxConfig(provider=provider_name, block_network=False) + try: + provider = create_sandbox_backend(config, "compliance-job-123") + except ImportError: + pytest.skip(f"{provider_name} SDK/adapter not installed") + assert_provider_contract(provider) diff --git a/tests/aiq_agent/agents/deep_researcher/sandbox/test_sandbox_runtime.py b/tests/aiq_agent/agents/deep_researcher/sandbox/test_sandbox_runtime.py new file mode 100644 index 000000000..e20599b55 --- /dev/null +++ b/tests/aiq_agent/agents/deep_researcher/sandbox/test_sandbox_runtime.py @@ -0,0 +1,431 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the provider-neutral sandbox seam (registry, config, capabilities, base). + +These run without a live Modal/OpenShell gateway: provider behavior is exercised +through small fakes, so only the framework logic (dispatch, fail-closed gate, +lazy creation, idempotency-gated retry, cleanup) is under test. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import pytest +from pydantic import ValidationError + +from aiq_agent.agents.deep_researcher.sandbox import CapabilityError +from aiq_agent.agents.deep_researcher.sandbox import SandboxCapabilities +from aiq_agent.agents.deep_researcher.sandbox import SandboxConfig +from aiq_agent.agents.deep_researcher.sandbox import SandboxProvider +from aiq_agent.agents.deep_researcher.sandbox import SandboxTerminatedError +from aiq_agent.agents.deep_researcher.sandbox import create_sandbox_backend +from aiq_agent.agents.deep_researcher.sandbox import register_sandbox_provider +from aiq_agent.agents.deep_researcher.sandbox import registered_providers +from aiq_agent.agents.deep_researcher.sandbox import verify_capabilities +from aiq_agent.agents.deep_researcher.sandbox.config import job_scoped_artifact_dir +from aiq_agent.agents.deep_researcher.sandbox.config import job_scoped_workdir + + +class _RecoverableError(Exception): + """Stand-in for a provider's transient/stale-sandbox error.""" + + +class _RegisteredFake(SandboxProvider): + """Minimal registered provider with conservative (default) capabilities.""" + + provider_name = "registered-fake" + + @property + def capabilities(self) -> SandboxCapabilities: + return SandboxCapabilities() + + def _create_session(self) -> Any: + return MagicMock() + + def _prepare_workspace(self, session: Any) -> None: + # These fakes assert exact execute call counts in the lock/retry/timeout tests; + # per-job workspace prep is covered explicitly in TestWorkspacePreparation. + return None + + +class _ScriptedProvider(SandboxProvider): + """Provider that hands out caller-supplied sessions, for retry/lazy tests.""" + + provider_name = "scripted" + + def __init__(self, config: SandboxConfig, job_id: str, sessions: list[Any]) -> None: + super().__init__(config, job_id) + self._sessions = sessions + self.sessions_created = 0 + + @property + def capabilities(self) -> SandboxCapabilities: + return SandboxCapabilities(supports_network_policy=True) + + def is_recoverable_error(self, exc: Exception) -> bool: + return isinstance(exc, _RecoverableError) + + def _create_session(self) -> Any: + session = self._sessions[self.sessions_created] + self.sessions_created += 1 + return session + + def _prepare_workspace(self, session: Any) -> None: + # See _RegisteredFake: keep exact execute call counts under test. + return None + + +class _WorkspaceProvider(SandboxProvider): + """Uses the default ``_prepare_workspace`` so its mkdir-on-create is observable.""" + + provider_name = "workspace-fake" + + def __init__(self, config: SandboxConfig, job_id: str, session: Any) -> None: + super().__init__(config, job_id) + self._next = session + + @property + def capabilities(self) -> SandboxCapabilities: + return SandboxCapabilities(supports_network_policy=True) + + def _create_session(self) -> Any: + return self._next + + +register_sandbox_provider("registered-fake", _RegisteredFake) + + +def _fake_config(**overrides: Any) -> SandboxConfig: + base: dict[str, Any] = {"provider": "registered-fake", "block_network": False} + base.update(overrides) + return SandboxConfig(**base) + + +class TestRegistry: + def test_builtin_providers_registered(self) -> None: + assert "modal" in registered_providers() + assert "openshell" in registered_providers() + + def test_create_unknown_provider_raises(self) -> None: + config = _fake_config() + object.__setattr__(config, "provider", "ghost") # bypass field validation + with pytest.raises(ValueError, match="Registered providers"): + create_sandbox_backend(config, "job-1") + + def test_create_returns_provider_instance(self) -> None: + backend = create_sandbox_backend(_fake_config(), "job-1") + assert isinstance(backend, _RegisteredFake) + + +class TestSandboxConfig: + def test_nested_modal_provider_settings(self) -> None: + config = SandboxConfig(provider="modal", providers={"modal": {"image": "nested:tag"}}) + assert config.providers.modal.image == "nested:tag" + + def test_provider_normalized_lowercase(self) -> None: + assert SandboxConfig(provider="MODAL").provider == "modal" + + def test_default_workdir(self) -> None: + config = SandboxConfig() + assert config.workdir == "/workspace" + + def test_unknown_provider_rejected(self) -> None: + with pytest.raises(ValueError, match="Registered providers"): + SandboxConfig(provider="does-not-exist") + + +class TestCapabilityGate: + def test_block_network_requires_capability(self) -> None: + config = _fake_config(block_network=True) # _RegisteredFake declares no network policy + with pytest.raises(CapabilityError, match="block_network"): + create_sandbox_backend(config, "job-1") + + def test_passes_when_network_unblocked(self) -> None: + backend = create_sandbox_backend(_fake_config(block_network=False), "job-1") + assert isinstance(backend, _RegisteredFake) + + def test_artifact_capture_requires_download(self) -> None: + caps = SandboxCapabilities(supports_network_policy=True, supports_artifact_download=False) + config = SandboxConfig(provider="registered-fake", artifact_capture={"enabled": True}) + with pytest.raises(CapabilityError, match="download"): + verify_capabilities(config, caps) + + +class TestNetworkPolicy: + def test_legacy_block_network_true_maps_to_blocked(self) -> None: + config = SandboxConfig(provider="registered-fake", block_network=True) + assert config.network.mode == "blocked" + assert config.block_network is True + + def test_legacy_block_network_false_maps_to_open(self) -> None: + config = SandboxConfig(provider="registered-fake", block_network=False) + assert config.network.mode == "open" + assert config.block_network is False + + def test_legacy_block_network_rejects_unknown_string(self) -> None: + # A typo must fail loudly, not silently open egress on a network-blocked sandbox. + with pytest.raises(ValidationError): + SandboxConfig(provider="registered-fake", block_network="flase") + + def test_explicit_network_wins_over_legacy_block_network(self) -> None: + config = SandboxConfig(provider="registered-fake", block_network=True, network={"mode": "open"}) + assert config.network.mode == "open" + assert config.block_network is False + + def test_allowlist_requires_hosts(self) -> None: + with pytest.raises(ValueError, match="allowlist"): + SandboxConfig(provider="registered-fake", network={"mode": "allowlist"}) + + def test_allowlist_requires_capability(self) -> None: + # _RegisteredFake declares neither network policy nor allowlist support. + config = SandboxConfig(provider="registered-fake", network={"mode": "allowlist", "allow": ["pypi.org"]}) + with pytest.raises(CapabilityError, match="allowlist"): + create_sandbox_backend(config, "job-1") + + def test_allowlist_passes_when_capability_declared(self) -> None: + config = SandboxConfig(provider="registered-fake", network={"mode": "allowlist", "allow": ["pypi.org"]}) + verify_capabilities(config, SandboxCapabilities(supports_network_allowlist=True)) + + +class TestResourceLimits: + """Opt-in CPU/memory caps, gated fail-closed (SANDBOX-5).""" + + def test_unset_resources_run_on_any_provider(self) -> None: + # Non-breaking: no limits requested -> no resource gate, regardless of capability. + verify_capabilities(_fake_config(), SandboxCapabilities()) + + def test_resource_limit_requires_capability(self) -> None: + # A requested limit on a provider that can't enforce it must fail closed. + config = _fake_config(resources={"cpu": 2}) + with pytest.raises(CapabilityError, match="resource limits"): + verify_capabilities(config, SandboxCapabilities()) + + def test_resource_limit_passes_when_declared(self) -> None: + config = _fake_config(resources={"memory_mb": 2048}) + verify_capabilities(config, SandboxCapabilities(supports_resource_limits=True)) + + def test_resource_limit_rejects_non_positive(self) -> None: + with pytest.raises(ValidationError): + SandboxConfig(provider="registered-fake", resources={"cpu": 0}) + + +class TestEntryPointDiscovery: + def test_entry_point_provider_is_discovered(self, monkeypatch: Any) -> None: + from aiq_agent.agents.deep_researcher.sandbox import registry + + class _EntryPointProvider(_RegisteredFake): + provider_name = "ep-fake" + + class _FakeEntryPoint: + name = "ep-fake" + + def load(self) -> type[SandboxProvider]: + return _EntryPointProvider + + def _fake_entry_points(*, group: str) -> list[Any]: + assert group == registry.SANDBOX_PROVIDER_ENTRY_POINT_GROUP + return [_FakeEntryPoint()] + + monkeypatch.setattr(registry, "_entry_points_loaded", False) + monkeypatch.setattr("importlib.metadata.entry_points", _fake_entry_points) + registry._SANDBOX_PROVIDERS.pop("ep-fake", None) + try: + assert registry.is_registered("ep-fake") + assert "ep-fake" in registered_providers() + finally: + registry._SANDBOX_PROVIDERS.pop("ep-fake", None) + + def test_broken_entry_point_is_skipped(self, monkeypatch: Any) -> None: + from aiq_agent.agents.deep_researcher.sandbox import registry + + class _BrokenEntryPoint: + name = "broken" + + def load(self) -> type[SandboxProvider]: + raise RuntimeError("boom") + + monkeypatch.setattr(registry, "_entry_points_loaded", False) + monkeypatch.setattr("importlib.metadata.entry_points", lambda *, group: [_BrokenEntryPoint()]) + # Must not raise; built-in resolution stays intact. + registry._load_entry_point_providers() + assert "broken" not in registry._SANDBOX_PROVIDERS + assert "modal" in registered_providers() + + +class TestProviderLifecycle: + def test_session_created_lazily(self) -> None: + session = MagicMock() + provider = _ScriptedProvider(_fake_config(), "job-1", sessions=[session]) + assert provider.sessions_created == 0 + provider.execute("echo ok", timeout=5) + assert provider.sessions_created == 1 + session.execute.assert_called_once_with("echo ok", timeout=5) + + def test_idempotent_download_retries_on_recoverable_error(self) -> None: + first = MagicMock() + first.download_files.side_effect = _RecoverableError("gone") + second = MagicMock() + second.download_files.return_value = ["ok"] + provider = _ScriptedProvider(_fake_config(), "job-1", sessions=[first, second]) + + result = provider.download_files(["/workspace/a.png"]) + + assert result == ["ok"] + assert provider.sessions_created == 2 # recreated once + + def test_execute_does_not_retry_on_recoverable_error(self) -> None: + first = MagicMock() + first.execute.side_effect = _RecoverableError("gone") + second = MagicMock() + provider = _ScriptedProvider(_fake_config(), "job-1", sessions=[first, second]) + + with pytest.raises(_RecoverableError): + provider.execute("echo ok") + + # Non-idempotent op must NOT silently recreate + re-run on a fresh empty sandbox. + assert provider.sessions_created == 1 + + def test_execute_timeout_clamped_to_config_limit(self) -> None: + session = MagicMock() + session.execute.return_value = "ok" + provider = _ScriptedProvider(_fake_config(timeout=1200), "job-1", sessions=[session]) + + provider.execute("echo ok", timeout=120000) # e.g. a tool passing milliseconds + + session.execute.assert_called_once_with("echo ok", timeout=1200) + + def test_execute_timeout_passthrough_when_within_limit(self) -> None: + session = MagicMock() + session.execute.return_value = "ok" + provider = _ScriptedProvider(_fake_config(timeout=1200), "job-1", sessions=[session]) + + provider.execute("echo ok", timeout=30) + + session.execute.assert_called_once_with("echo ok", timeout=30) + + def test_close_releases_session(self) -> None: + session = MagicMock() + provider = _ScriptedProvider(_fake_config(), "job-1", sessions=[session]) + provider.execute("echo ok") + provider.close() + session.close.assert_called_once() + assert provider._session is None + + def test_terminate_releases_session_and_blocks_further_ops(self) -> None: + session = MagicMock() + provider = _ScriptedProvider(_fake_config(), "job-1", sessions=[session]) + provider.execute("echo ok") + + provider.terminate() + + session.close.assert_called_once() + assert provider._session is None + # A terminated provider refuses new work instead of silently recreating. + with pytest.raises(SandboxTerminatedError): + provider.execute("echo again") + assert provider.sessions_created == 1 + + def test_terminate_preempts_in_flight_execute(self) -> None: + # terminate() must not wait on the operation lock: while execute() is blocked in a + # remote call, a concurrent terminate() closes the session out-of-band, which + # interrupts the call. Without the two-lock split this test would deadlock/hang. + import threading + + entered = threading.Event() + released = threading.Event() + + session = MagicMock() + + def _blocking_execute(command: str, timeout: int | None = None) -> str: + entered.set() + if not released.wait(timeout=5): + raise AssertionError("execute was not interrupted by terminate()") + raise RuntimeError("session closed") + + session.execute.side_effect = _blocking_execute + session.close.side_effect = lambda: released.set() + provider = _ScriptedProvider(_fake_config(), "job-1", sessions=[session]) + + result: dict[str, BaseException] = {} + + def _run_execute() -> None: + try: + provider.execute("sleep") + except BaseException as exc: # noqa: BLE001 - capture for assertion + result["error"] = exc + + worker = threading.Thread(target=_run_execute) + worker.start() + assert entered.wait(timeout=5), "execute never started" + + provider.terminate() # must return without waiting for the in-flight execute + worker.join(timeout=5) + + assert not worker.is_alive(), "execute did not unblock after terminate()" + session.close.assert_called_once() + assert isinstance(result.get("error"), Exception) + + +class TestJobScopedPaths: + """The per-job workspace helper isolates jobs within a shared/reused sandbox.""" + + def test_workdir_and_artifact_dir_are_job_scoped(self) -> None: + assert job_scoped_workdir("/sandbox", "abc") == "/sandbox/abc" + assert job_scoped_artifact_dir("/sandbox", "abc") == "/sandbox/abc/aiq-artifacts" + + def test_trailing_slash_normalized(self) -> None: + assert job_scoped_workdir("/sandbox/", "abc") == "/sandbox/abc" + + def test_unsafe_job_id_cannot_escape_base(self) -> None: + # A crafted id with path separators must collapse to a single safe segment so it + # cannot move the workspace (and harvest confinement root) outside the base. + result = job_scoped_workdir("/sandbox", "../../etc") + assert result.startswith("/sandbox/") + assert ".." not in result + assert result.count("/") == 2 + + +class TestWorkspacePreparation: + """The provider base creates the per-job workspace on session start (mkdir -p).""" + + def test_provider_paths_are_job_scoped(self) -> None: + provider = _RegisteredFake(_fake_config(workdir="/sandbox"), "job-xyz") + assert provider.workdir == "/sandbox/job-xyz" + assert provider.artifact_dir == "/sandbox/job-xyz/aiq-artifacts" + + def test_workspace_created_on_session_start(self) -> None: + session = MagicMock() + session.execute.return_value = "ok" + provider = _WorkspaceProvider(_fake_config(workdir="/sandbox"), "job-xyz", session) + + provider.execute("echo hi") + + # The first execute after creation is the idempotent mkdir -p of the per-job roots. + mkdir_cmd = session.execute.call_args_list[0].args[0] + assert mkdir_cmd.startswith("mkdir -p") + assert "/sandbox/job-xyz" in mkdir_cmd + assert "/sandbox/job-xyz/aiq-artifacts" in mkdir_cmd + + def test_workspace_prep_failure_does_not_abort_the_call(self) -> None: + # Best-effort: a mkdir failure is swallowed so it cannot break session creation; + # a real filesystem problem resurfaces on the first actual write. + session = MagicMock() + session.execute.side_effect = [RuntimeError("mkdir boom"), "ok"] + provider = _WorkspaceProvider(_fake_config(workdir="/sandbox"), "job-1", session) + + assert provider.execute("echo hi") == "ok" diff --git a/tests/aiq_agent/agents/deep_researcher/test_agent.py b/tests/aiq_agent/agents/deep_researcher/test_agent.py index bb33faea8..8ebe646f2 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_agent.py +++ b/tests/aiq_agent/agents/deep_researcher/test_agent.py @@ -180,7 +180,15 @@ def test_init_with_defaults(self, mock_llm_provider, real_tool, mock_create_deep def test_init_with_custom_settings(self, mock_llm_provider, real_tool, mock_create_deep_agent): """Test DeepResearcherAgent initialization with custom settings.""" - with patch("aiq_agent.agents.deep_researcher.factory.create_deep_agent", return_value=mock_create_deep_agent): + with ( + patch("aiq_agent.agents.deep_researcher.factory.create_deep_agent", return_value=mock_create_deep_agent), + # Patch backend creation so the test does not require the optional OpenShell adapter + # (the default sandbox provider) to be installed. + patch( + "aiq_agent.agents.deep_researcher.deepagents_runtime._create_sandbox_backend", + return_value=MagicMock(), + ), + ): from aiq_agent.agents.deep_researcher.agent import DeepResearcherAgent from aiq_agent.agents.deep_researcher.deepagents_runtime import DeepResearchSandboxConfig from aiq_agent.agents.deep_researcher.deepagents_runtime import DeepResearchSkillsConfig @@ -213,13 +221,13 @@ def test_init_with_custom_settings(self, mock_llm_provider, real_tool, mock_crea assert agent.deepagents_runtime.skill_sources_for("researcher-agent") == ["/skills/research/"] def test_sandbox_config_rejects_unsupported_provider(self): - """Unsupported sandbox providers fail early with a clear error.""" + """Unsupported sandbox providers fail validation at config load (registry-backed).""" from pydantic import ValidationError - from aiq_agent.agents.deep_researcher.deepagents_runtime import DeepResearchSandboxConfig + from aiq_agent.agents.deep_researcher.sandbox.config import SandboxConfig - with pytest.raises(ValidationError, match="Input should be 'modal'"): - DeepResearchSandboxConfig(provider="not-modal") + with pytest.raises(ValidationError, match="Unsupported sandbox provider"): + SandboxConfig(provider="not-a-real-provider") def test_register_uses_runtime_config_models(self): """NAT config uses the same skills and sandbox models as runtime.""" @@ -252,7 +260,7 @@ def test_register_uses_runtime_config_models(self): assert config.max_source_tool_batch_size == 4 assert config.enable_source_router is False assert config.sandbox is not None - assert config.sandbox.provider == "modal" + assert config.sandbox.provider == "openshell" assert config.sandbox.app_name == "custom-aiq" assert config.sandbox.packages == ("matplotlib", "pillow") @@ -283,13 +291,13 @@ def test_register_resolves_named_runtime_config_refs(self): def test_modal_sandbox_name_is_job_id(self): """Modal sandbox names use the resolved job ID directly.""" - from aiq_agent.agents.deep_researcher.deepagents_runtime import _validate_modal_sandbox_name + from aiq_agent.agents.deep_researcher.sandbox.providers.modal import _validate_modal_sandbox_name assert _validate_modal_sandbox_name("job-123") == "job-123" def test_modal_sandbox_name_rejects_invalid_job_id(self): """Invalid custom job IDs fail before creating a Modal sandbox.""" - from aiq_agent.agents.deep_researcher.deepagents_runtime import _validate_modal_sandbox_name + from aiq_agent.agents.deep_researcher.sandbox.providers.modal import _validate_modal_sandbox_name with pytest.raises(ValueError, match="valid Modal sandbox name"): _validate_modal_sandbox_name("bad/job/id") @@ -960,86 +968,38 @@ def test_modal_backend_is_concrete_cached_and_routes_skills_locally(self, mock_l from aiq_agent.agents.deep_researcher.deepagents_runtime import DeepResearchSkillsConfig sandbox = DeepResearchSandboxConfig() - agent = DeepResearcherAgent( - llm_provider=mock_llm_provider, - tools=[real_tool], - skills=DeepResearchSkillsConfig(agents={"writer-agent": ("synthesis",)}), - sandbox=sandbox, - job_id="job-123", - ) fake_modal_backend = MagicMock() - with ( - patch( - "aiq_agent.agents.deep_researcher.deepagents_runtime._create_sandbox_backend", - return_value=fake_modal_backend, - ) as create_backend, - ): + # The runtime now builds the sandbox provider eagerly in __init__, so the patch + # must wrap construction, not just the later .backend access. + with patch( + "aiq_agent.agents.deep_researcher.deepagents_runtime._create_sandbox_backend", + return_value=fake_modal_backend, + ) as create_backend: + agent = DeepResearcherAgent( + llm_provider=mock_llm_provider, + tools=[real_tool], + skills=DeepResearchSkillsConfig(agents={"writer-agent": ("synthesis",)}), + sandbox=sandbox, + job_id="job-123", + ) backend_one = agent.deepagents_runtime.backend backend_two = agent.deepagents_runtime.backend assert backend_one is backend_two assert backend_one.default is fake_modal_backend - create_backend.assert_called_once_with( - sandbox, - "job-123", - ) + create_backend.assert_called_once_with(sandbox, "job-123") assert isinstance(backend_one.routes[BUILTIN_SKILL_SOURCE], FilesystemBackend) assert isinstance(backend_one.routes[SHARED_ROUTE], StateBackend) fake_modal_backend.ls.assert_not_called() fake_modal_backend.read.assert_not_called() - def test_modal_backend_creates_sandbox_lazily(self): - """Modal sandbox lifetime starts on first sandbox operation, not agent construction.""" - from deepagents.backends.protocol import ExecuteResponse - - from aiq_agent.agents.deep_researcher.deepagents_runtime import DeepResearchSandboxConfig - from aiq_agent.agents.deep_researcher.deepagents_runtime import _create_sandbox_backend - - fake_modal_backend = MagicMock() - fake_modal_backend.execute.return_value = ExecuteResponse(output="ok", exit_code=0) - - with patch( - "aiq_agent.agents.deep_researcher.deepagents_runtime._create_modal_backend_now", - return_value=fake_modal_backend, - ) as create_modal: - backend = _create_sandbox_backend(DeepResearchSandboxConfig(), "job-123") - - create_modal.assert_not_called() - result = backend.execute("echo ok", timeout=5) - - assert result.output == "ok" - create_modal.assert_called_once() - fake_modal_backend.execute.assert_called_once_with("echo ok", timeout=5) - - def test_modal_backend_recreates_and_retries_once_on_not_found(self): - """A disappeared Modal container is recreated once for the same job-scoped name.""" - import modal - from deepagents.backends.protocol import ExecuteResponse - - from aiq_agent.agents.deep_researcher.deepagents_runtime import DeepResearchSandboxConfig - from aiq_agent.agents.deep_researcher.deepagents_runtime import _create_sandbox_backend - - first_modal_backend = MagicMock() - first_modal_backend.execute.side_effect = modal.exception.NotFoundError("gone") - second_modal_backend = MagicMock() - second_modal_backend.execute.return_value = ExecuteResponse(output="ok", exit_code=0) - config = DeepResearchSandboxConfig() - - with patch( - "aiq_agent.agents.deep_researcher.deepagents_runtime._create_modal_backend_now", - side_effect=[first_modal_backend, second_modal_backend], - ) as create_modal: - backend = _create_sandbox_backend(config, "job-123") - result = backend.execute("echo ok", timeout=5) - - assert result.output == "ok" - assert create_modal.call_args_list[0].args == (config, "job-123") - assert create_modal.call_args_list[0].kwargs == {} - assert create_modal.call_args_list[1].args == (config, "job-123") - assert create_modal.call_args_list[1].kwargs == {"force_new": True} - first_modal_backend.execute.assert_called_once_with("echo ok", timeout=5) - second_modal_backend.execute.assert_called_once_with("echo ok", timeout=5) + # NOTE: the former test_modal_backend_creates_sandbox_lazily and + # test_modal_backend_recreates_and_retries_once_on_not_found tested 284's inline + # _create_modal_backend_now / _LazyModalSandboxBackend internals, which are now + # replaced by the provider-neutral sandbox package. Lazy session creation and + # idempotency-gated retry are covered by tests/.../sandbox/test_sandbox_runtime.py + # (note: our provider intentionally does NOT retry the non-idempotent execute()). def test_load_prompts_raises_when_missing(self, mock_llm_provider, real_tool, mock_create_deep_agent): """Missing prompts fail fast instead of silently using inline defaults.""" diff --git a/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py b/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py index dbccfaa7d..9571962e8 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py +++ b/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py @@ -138,7 +138,13 @@ def test_require_sandbox_collection_with_sandbox_is_allowed(self) -> None: agents={"researcher-agent": ("research",)}, require_sandbox=("research",), ) - runtime = DeepAgentsRuntime(skills=skills, sandbox=DeepResearchSandboxConfig()) + # Patch backend creation so the test does not require the optional OpenShell adapter + # (the default provider) to be installed. + with patch( + "aiq_agent.agents.deep_researcher.deepagents_runtime._create_sandbox_backend", + return_value=MagicMock(), + ): + runtime = DeepAgentsRuntime(skills=skills, sandbox=DeepResearchSandboxConfig()) assert runtime.skill_sources_for("researcher-agent") == ["/skills/research/"] @@ -223,6 +229,12 @@ def fake_create_agent(*_args: Any, **kwargs: Any) -> MagicMock: class TestDeepAgentsRuntimeJobId: """job_id should drive the sandbox name; a missing one falls back to uuid.""" + def test_default_sandbox_provider_is_openshell(self) -> None: + sandbox = DeepResearchSandboxConfig() + + assert sandbox.provider == "openshell" + assert sandbox.workdir is None + def test_explicit_job_id_is_kept(self) -> None: sandbox = DeepResearchSandboxConfig() with patch("aiq_agent.agents.deep_researcher.deepagents_runtime._create_sandbox_backend") as create_backend: @@ -257,4 +269,4 @@ def find_spec(module_name: str): ), pytest.raises(ImportError, match="langchain-modal"), ): - _ = DeepAgentsRuntime(sandbox=DeepResearchSandboxConfig()).backend + _ = DeepAgentsRuntime(sandbox=DeepResearchSandboxConfig(provider="modal")).backend diff --git a/tests/aiq_agent/common/test_citation_verification.py b/tests/aiq_agent/common/test_citation_verification.py index 962981c52..ee8ace7f5 100644 --- a/tests/aiq_agent/common/test_citation_verification.py +++ b/tests/aiq_agent/common/test_citation_verification.py @@ -1014,6 +1014,88 @@ def test_dedup_keeps_lowest_number_when_duplicates_appear_after_unique(self): assert ref_section.count("[3]") == 1 +class TestVerifyCitationsBackfill: + """Backfill of URL-less ``[N] Title`` lines from the writer-facing list. + + The writer sometimes emits ``[N] Title`` and drops the ``: url`` suffix. + These tests cover recovering the URL from ``reference_sources`` (same + captured registry) without weakening precision. + """ + + @pytest.fixture(name="registry") + def fixture_registry(self): + reg = SourceRegistry() + reg.add(SourceEntry(url="https://amd.com/q1-2024", title="AMD Q1 2024 Financial Results", source_type="tavily")) + reg.add(SourceEntry(url="https://meta.com/q1-2024", title="Meta Q1 2024 Results", source_type="tavily")) + reg.add(SourceEntry(citation_key="report.pdf, p.15", title="Internal Report", source_type="knowledge_layer")) + return reg + + def test_url_less_lines_backfilled_from_reference_sources(self, registry): + report = ( + "AMD spent more [1]. Meta followed [2].\n\n" + "## Sources\n" + "[1] AMD Q1 2024 Financial Results\n" + "[2] Meta Q1 2024 Results" + ) + result = verify_citations(report, registry, reference_sources=registry.all_sources()) + + assert "[1] AMD Q1 2024 Financial Results: https://amd.com/q1-2024" in result.verified_report + assert "[2] Meta Q1 2024 Results: https://meta.com/q1-2024" in result.verified_report + assert len(result.valid_citations) == 2 + assert not result.removed_citations + + def test_backfill_recovers_url_less_citation_key_source(self, registry): + report = "Per the internal report [1].\n\n## Sources\n[1] Internal Report" + result = verify_citations(report, registry, reference_sources=registry.all_sources()) + + assert len(result.valid_citations) == 1 + assert result.valid_citations[0]["citation_key"] == "report.pdf, p.15" + assert not result.removed_citations + assert "[1]" in result.verified_report.split("## Sources", 1)[1] + + def test_url_less_line_without_matching_reference_still_removed(self, registry): + report = "Claim [1].\n\n## Sources\n[1] Some Source The Writer Invented" + result = verify_citations(report, registry, reference_sources=registry.all_sources()) + + assert not result.valid_citations + assert len(result.removed_citations) == 1 + assert result.removed_citations[0]["reason"] == "unverifiable" + + def test_ambiguous_title_not_backfilled(self): + reg = SourceRegistry() + reg.add(SourceEntry(url="https://a.com/one", title="Quarterly Update", source_type="tavily")) + reg.add(SourceEntry(url="https://b.com/two", title="Quarterly Update", source_type="tavily")) + report = "Claim [1].\n\n## Sources\n[1] Quarterly Update" + result = verify_citations(report, reg, reference_sources=reg.all_sources()) + + assert not result.valid_citations + assert len(result.removed_citations) == 1 + assert result.removed_citations[0]["reason"] == "unverifiable" + + def test_aggregate_label_not_resurrected(self): + reg = SourceRegistry() + reg.add(SourceEntry(url="https://cnn.com/biz", title="CNN", source_type="tavily")) + report = "Markets moved [1].\n\n## Sources\n[1] CNN; Yahoo Finance; Barchart" + result = verify_citations(report, reg, reference_sources=reg.all_sources()) + + assert not result.valid_citations + assert len(result.removed_citations) == 1 + assert result.removed_citations[0]["reason"] == "unverifiable" + + def test_existing_url_and_key_paths_unchanged_without_reference_sources(self, registry): + report = ( + "AMD [1]. Doc [2].\n\n" + "## Sources\n" + "[1] AMD Q1 2024 Financial Results: https://amd.com/q1-2024\n" + "[2] report.pdf, p.15" + ) + result = verify_citations(report, registry) + + assert len(result.valid_citations) == 2 + assert not result.removed_citations + assert "https://amd.com/q1-2024" in result.verified_report + + # --------------------------------------------------------------------------- # sanitize_report tests # --------------------------------------------------------------------------- diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index 78721f56f..f4c78171d 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -1394,6 +1394,8 @@ def __init__( skills=None, sandbox=None, job_id=None, + artifact_db_url=None, + artifact_emit=None, max_research_concurrency=None, max_concurrent_source_tool_calls=None, max_source_tool_batch_size=None, @@ -1408,6 +1410,8 @@ def __init__( self.skills = skills self.sandbox = sandbox self.job_id = job_id + self.artifact_db_url = artifact_db_url + self.artifact_emit = artifact_emit self.max_research_concurrency = max_research_concurrency self.max_concurrent_source_tool_calls = max_concurrent_source_tool_calls self.max_source_tool_batch_size = max_source_tool_batch_size @@ -1649,6 +1653,8 @@ def __init__( skills=None, sandbox=None, job_id=None, + artifact_db_url=None, + artifact_emit=None, max_research_concurrency=None, max_concurrent_source_tool_calls=None, max_source_tool_batch_size=None, @@ -1672,3 +1678,56 @@ def __init__( callbacks=["callback"], job_id="job-123", ) + + +class TestTerminalTeardown: + """_teardown_sandbox routes close()/terminate() and never raises on the terminal path.""" + + def test_none_runtime_is_noop(self): + from aiq_api.jobs.runner import _teardown_sandbox + + # Must not raise when no sandbox runtime is present (non-sandbox agents). + _teardown_sandbox(None, job_id="job-1", interrupted=False) + + def test_normal_path_calls_close(self): + from aiq_api.jobs.runner import _teardown_sandbox + + runtime = MagicMock(spec=["close", "terminate"]) + _teardown_sandbox(runtime, job_id="job-1", interrupted=False) + + runtime.close.assert_called_once_with() + runtime.terminate.assert_not_called() + + def test_interrupted_path_calls_terminate(self): + from aiq_api.jobs.runner import _teardown_sandbox + + runtime = MagicMock(spec=["close", "terminate"]) + _teardown_sandbox(runtime, job_id="job-1", interrupted=True) + + runtime.terminate.assert_called_once_with() + runtime.close.assert_not_called() + + def test_interrupted_without_terminate_falls_back_to_close(self): + from aiq_api.jobs.runner import _teardown_sandbox + + runtime = MagicMock(spec=["close"]) # no terminate attribute + _teardown_sandbox(runtime, job_id="job-1", interrupted=True) + + runtime.close.assert_called_once_with() + + def test_never_raises_when_teardown_fails(self): + from aiq_api.jobs.runner import _teardown_sandbox + + runtime = MagicMock(spec=["close", "terminate"]) + runtime.close.side_effect = RuntimeError("sdk session close failed") + # Must swallow the error; teardown is best-effort on the terminal path. + _teardown_sandbox(runtime, job_id="job-1", interrupted=False) + + def test_does_not_harvest(self): + from aiq_api.jobs.runner import _teardown_sandbox + + # The single harvest happens in agent.run(); teardown must not call final_harvest. + runtime = MagicMock(spec=["close", "terminate", "final_harvest"]) + _teardown_sandbox(runtime, job_id="job-1", interrupted=False) + + runtime.final_harvest.assert_not_called()