diff --git a/docs/source/architecture/agents/sandbox.md b/docs/source/architecture/agents/sandbox.md index 1a09cfb51..6e33128e5 100644 --- a/docs/source/architecture/agents/sandbox.md +++ b/docs/source/architecture/agents/sandbox.md @@ -51,6 +51,9 @@ and [Production Considerations](../../deployment/production.md#artifact-storage) 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. +- Manifest checkpoints preserve completed artifacts after successful sandbox commands. The + terminal finalizer harvests once before cleanup on success/failure; cancellation harvests + only when the provider is idle and otherwise terminates immediately. - The runtime closes provider sessions on success, failure, cancellation, and timeout. A named OpenShell sandbox persists when `delete_on_exit` is disabled. diff --git a/frontends/aiq_api/README.md b/frontends/aiq_api/README.md index abc157b41..5e959c49b 100644 --- a/frontends/aiq_api/README.md +++ b/frontends/aiq_api/README.md @@ -149,7 +149,7 @@ Events streamed during job execution: | `workflow.start` / `workflow.end` | Workflow lifecycle | | `llm.start` / `llm.chunk` / `llm.end` | LLM inference progress | | `tool.start` / `tool.end` | Tool invocations | -| `artifact.update` | Todos, files, citations, output updates | +| `artifact.update` | Todos, citations, output, and generated-file metadata (`content_url` for durable bytes) | | `job.error` | Error occurred | ## Configuration diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index d2127590a..932a6c937 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -801,7 +801,16 @@ async def run_agent_job( # Signal event stream completion event_stream.on_complete() - # Flush any buffered events before updating status + # Harvest artifacts (durable, idempotent) before SUCCESS so clients cannot + # stop streaming before the terminal metadata is persisted. Resource release + # is deferred to the finally block: the provider's close() is unbounded, so + # awaiting it here could strand a finished job in RUNNING if SDK cleanup hangs. + await asyncio.to_thread( + _harvest_sandbox_artifacts, + sandbox_runtime, + job_id=job_id, + interrupted=False, + ) if hasattr(event_store, "flush"): event_store.flush() @@ -835,17 +844,10 @@ 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) - if job and job.status != JobStatus.INTERRUPTED.value: - await job_store.update_status(job_id, JobStatus.INTERRUPTED, error="cancelled by user") - except (ConnectionError, TimeoutError, RuntimeError): - pass - if event_store is None: event_store = BatchingEventStore(EventStore(db_url, job_id)) + await asyncio.to_thread(_teardown_sandbox, sandbox_runtime, job_id=job_id, interrupted=True) _store_terminal_event_best_effort( event_store, { @@ -854,14 +856,20 @@ async def run_agent_job( }, ) - except Exception as e: - logger.exception("Job %s failed: %s", job_id, type(e).__name__) if job_store: - await job_store.update_status(job_id, JobStatus.FAILURE, error=str(e)) + try: + job = await job_store.get_job(job_id) + if job and job.status != JobStatus.INTERRUPTED.value: + await job_store.update_status(job_id, JobStatus.INTERRUPTED, error="cancelled by user") + except (ConnectionError, TimeoutError, RuntimeError): + pass + except Exception as e: + logger.exception("Job %s failed: %s", job_id, type(e).__name__) if event_store is None: event_store = BatchingEventStore(EventStore(db_url, job_id)) + await asyncio.to_thread(_harvest_sandbox_artifacts, sandbox_runtime, job_id=job_id, interrupted=False) _store_terminal_event_best_effort( event_store, { @@ -872,6 +880,8 @@ async def run_agent_job( }, }, ) + if job_store: + await job_store.update_status(job_id, JobStatus.FAILURE, error=str(e)) finally: # Ensure terminal-path events are not left in the batch buffer. @@ -886,9 +896,7 @@ async def run_agent_job( ) 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. + # Idempotent fallback for failures before a terminal branch finalized the runtime. 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: @@ -912,8 +920,30 @@ def _store_terminal_event_best_effort(event_store, event: dict) -> None: ) +def _harvest_sandbox_artifacts(sandbox_runtime: Any | None, *, job_id: str, interrupted: bool) -> None: + """Persist captured artifacts on a terminal path without releasing the sandbox. + + Callable before the terminal job status so artifact metadata is durable, yet it never invokes + the provider's unbounded ``close()``/``terminate()``. Resource release stays in + ``_teardown_sandbox`` (run from ``finally``) so a hanging SDK cleanup cannot strand a finished + job in ``RUNNING`` with a stream that never terminates. The harvest is idempotent. + """ + if sandbox_runtime is None: + return + finalize_artifacts = getattr(sandbox_runtime, "finalize_artifacts", None) + if callable(finalize_artifacts): + try: + finalize_artifacts(interrupted=interrupted) + except Exception as exc: # noqa: BLE001 - artifact capture cannot replace the job result + logger.warning( + "Terminal artifact harvest failed for job %s exception=%s", + job_id, + exc.__class__.__name__, + ) + + def _teardown_sandbox(sandbox_runtime: Any | None, *, job_id: str, interrupted: bool) -> None: - """Release sandbox resources on a terminal path (best-effort, never raises). + """Harvest artifacts and release sandbox resources on a terminal path. Interrupted jobs (cancel/timeout) call ``terminate()`` so a still-running ``execute`` is forcibly preempted; normal paths call ``close()`` gracefully. Both are idempotent. This runs @@ -921,6 +951,7 @@ def _teardown_sandbox(sandbox_runtime: Any | None, *, job_id: str, interrupted: """ if sandbox_runtime is None: return + _harvest_sandbox_artifacts(sandbox_runtime, job_id=job_id, interrupted=interrupted) teardown = getattr(sandbox_runtime, "terminate", None) if interrupted else None if teardown is None: teardown = getattr(sandbox_runtime, "close", None) @@ -928,8 +959,11 @@ def _teardown_sandbox(sandbox_runtime: Any | None, *, job_id: str, interrupted: 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) + except Exception as exc: # noqa: BLE001 - cleanup must never raise on the terminal path + # Secret-safe: log only the exception type. A provider cleanup error can carry a + # credential or internal hostname, which must never reach the logs (matches the + # finalize_artifacts handler above). + logger.warning("Sandbox cleanup failed for job %s exception=%s", job_id, exc.__class__.__name__) def _create_agent_instance( diff --git a/frontends/ui/src/adapters/api/deep-research-client.spec.ts b/frontends/ui/src/adapters/api/deep-research-client.spec.ts index a57b0f035..4df12cd1e 100644 --- a/frontends/ui/src/adapters/api/deep-research-client.spec.ts +++ b/frontends/ui/src/adapters/api/deep-research-client.spec.ts @@ -2,7 +2,29 @@ // SPDX-License-Identifier: Apache-2.0 import { afterEach, describe, expect, test, vi } from 'vitest' -import { getJobStatus } from './deep-research-client' +import { createDeepResearchClient, getJobStatus } from './deep-research-client' + +class FakeEventSource { + static latest: FakeEventSource | null = null + onopen: (() => void) | null = null + onmessage: ((event: MessageEvent) => void) | null = null + onerror: (() => void) | null = null + private listeners = new Map void>() + + constructor(_url: string) { + FakeEventSource.latest = this + } + + addEventListener(type: string, listener: EventListener): void { + this.listeners.set(type, listener as (event: MessageEvent) => void) + } + + close(): void {} + + emit(type: string, data: unknown): void { + this.listeners.get(type)?.(new MessageEvent(type, { data: JSON.stringify(data) })) + } +} describe('deep research REST client', () => { afterEach(() => { @@ -32,4 +54,82 @@ describe('deep research REST client', () => { 'Failed to get job status: 500 - PROXY_ERROR: fetch failed' ) }) + + test('maps generated binary file events to durable artifact metadata', () => { + vi.stubGlobal('EventSource', FakeEventSource) + const onFileUpdate = vi.fn() + const client = createDeepResearchClient({ jobId: 'job-1', callbacks: { onFileUpdate } }) + client.connect() + + FakeEventSource.latest?.emit('artifact.update', { + data: { + type: 'file', + file_path: 'chart.png', + artifact_id: 'art_123', + content_url: '/v1/jobs/async/job/job-1/artifacts/art_123/content', + kind: 'image', + mime_type: 'image/png', + size_bytes: 2048, + sha256: 'a'.repeat(64), + title: 'Quarterly CapEx', + caption: 'Comparison chart', + inline: true, + }, + }) + + expect(onFileUpdate).toHaveBeenCalledWith({ + filename: 'chart.png', + content: undefined, + artifactId: 'art_123', + contentUrl: '/api/jobs/async/job/job-1/artifacts/art_123/content', + kind: 'image', + mimeType: 'image/png', + sizeBytes: 2048, + sha256: 'a'.repeat(64), + title: 'Quarterly CapEx', + caption: 'Comparison chart', + inline: true, + }) + }) + + test('preserves legacy text file events', () => { + vi.stubGlobal('EventSource', FakeEventSource) + const onFileUpdate = vi.fn() + const client = createDeepResearchClient({ jobId: 'job-1', callbacks: { onFileUpdate } }) + client.connect() + + FakeEventSource.latest?.emit('artifact.update', { + data: { type: 'file', file_path: 'report.md', content: '# Report' }, + }) + + expect(onFileUpdate).toHaveBeenCalledWith( + expect.objectContaining({ filename: 'report.md', content: '# Report' }) + ) + }) + + test('maps url-only artifacts (no artifact_id) via the content_url/url fallback', () => { + vi.stubGlobal('EventSource', FakeEventSource) + const onFileUpdate = vi.fn() + const client = createDeepResearchClient({ jobId: 'job-1', callbacks: { onFileUpdate } }) + client.connect() + + // No artifact_id: filename derives from the url and contentUrl falls back to content_url. + FakeEventSource.latest?.emit('artifact.update', { + data: { + type: 'file', + url: 'https://example.com/exports/data.csv', + content_url: 'https://example.com/exports/data.csv', + mime_type: 'text/csv', + }, + }) + + expect(onFileUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + filename: 'data.csv', + artifactId: undefined, + contentUrl: 'https://example.com/exports/data.csv', + mimeType: 'text/csv', + }) + ) + }) }) diff --git a/frontends/ui/src/adapters/api/deep-research-client.ts b/frontends/ui/src/adapters/api/deep-research-client.ts index 5377e3d7c..4a62d3ee8 100644 --- a/frontends/ui/src/adapters/api/deep-research-client.ts +++ b/frontends/ui/src/adapters/api/deep-research-client.ts @@ -12,6 +12,7 @@ */ import { apiConfig } from './config' +import { artifactContentPath } from '@/shared/utils/artifact-url' // ============================================================ // Types @@ -149,6 +150,21 @@ export interface TodoItem { status: 'pending' | 'in_progress' | 'completed' | 'cancelled' } +/** File update emitted by artifact.update. Durable generated files carry metadata, not bytes. */ +export interface FileArtifactUpdate { + filename: string + content?: string + artifactId?: string + contentUrl?: string + kind?: string + mimeType?: string + sizeBytes?: number + sha256?: string + title?: string + caption?: string + inline?: boolean +} + /** artifact.update event */ export interface ArtifactUpdateEvent extends DeepResearchSSEEvent { event: 'artifact.update' @@ -157,8 +173,19 @@ export interface ArtifactUpdateEvent extends DeepResearchSSEEvent { timestamp: string data: { type: ArtifactType - content: string | TodoItem[] + content?: string | TodoItem[] url?: string // For citation_source and citation_use types + content_url?: string + file_path?: string + artifact_id?: string + job_id?: string + kind?: string + mime_type?: string + size_bytes?: number + sha256?: string + title?: string + caption?: string + inline?: boolean } metadata?: { workflow?: string @@ -208,7 +235,7 @@ export interface DeepResearchCallbacks { /** Called on artifact updates */ onTodoUpdate?: (todos: TodoItem[], workflow?: string) => void onCitationUpdate?: (url: string, content: string, isCited?: boolean) => void - onFileUpdate?: (filename: string, content: string) => void + onFileUpdate?: (file: FileArtifactUpdate) => void onOutputUpdate?: (content: string, outputCategory?: string, workflow?: string) => void /** Called on job heartbeat (confirms job is alive during long operations) */ onHeartbeat?: (uptimeSeconds: number) => void @@ -484,16 +511,13 @@ export const createDeepResearchClient = (options: DeepResearchStreamOptions): De case 'artifact.update': { // artifact.update has nested structure: { id, timestamp, data: { type, content, url?, output_category? }, metadata?: { workflow } } - const artifactWrapper = rawData as { - data?: { type: ArtifactType; content: string | TodoItem[]; url?: string; output_category?: string } - type?: ArtifactType - content?: string | TodoItem[] - url?: string - output_category?: string - metadata?: { workflow?: string } - } + // Reuse the exported ArtifactUpdateEvent data contract (durable metadata fields + // included) instead of re-declaring a narrower shape; `path`/`output_category` are + // the two extra keys this parser also reads. + type ArtifactData = ArtifactUpdateEvent['data']['data'] & { path?: string; output_category?: string } + const artifactWrapper = rawData as { data?: ArtifactData; metadata?: { workflow?: string } } & Partial // Handle both nested (data.type) and flat (type) structures - const artifactData = artifactWrapper.data || artifactWrapper + const artifactData = (artifactWrapper.data || artifactWrapper) as ArtifactData const artifactWorkflow = artifactWrapper.metadata?.workflow switch (artifactData.type) { @@ -509,11 +533,26 @@ export const createDeepResearchClient = (options: DeepResearchStreamOptions): De callbacks.onCitationUpdate?.(artifactData.url || '', artifactData.content as string, true) break case 'file': { - // file artifacts are written during research — extract filename from path - const raw = artifactData as Record - const filePath = (raw.file_path || raw.path || artifactData.url || 'unknown') as string - const fileName = filePath.split('/').pop() || filePath - callbacks.onFileUpdate?.(fileName, artifactData.content as string) + // Generated artifacts carry durable metadata; legacy text-file events carry content. + const artifactId = artifactData.artifact_id + // Fall back to artifactId (not a shared 'unknown') when no path/url is present, so two + // distinct pathless artifacts don't collapse onto one filename key downstream. + const filePath = artifactData.file_path || artifactData.path || artifactData.url + const fileName = filePath ? filePath.split('/').pop() || filePath : artifactId || 'unknown' + callbacks.onFileUpdate?.({ + filename: fileName, + // content is typed as string | TodoItem[]; only a string is a real file body. + content: typeof artifactData.content === 'string' ? artifactData.content : undefined, + artifactId, + contentUrl: artifactId ? artifactContentPath(jobId, artifactId) : artifactData.content_url || artifactData.url, + kind: artifactData.kind, + mimeType: artifactData.mime_type, + sizeBytes: artifactData.size_bytes, + sha256: artifactData.sha256, + title: artifactData.title, + caption: artifactData.caption, + inline: artifactData.inline, + }) break } case 'output': diff --git a/frontends/ui/src/adapters/api/index.ts b/frontends/ui/src/adapters/api/index.ts index 0300b8d3b..855e9f593 100644 --- a/frontends/ui/src/adapters/api/index.ts +++ b/frontends/ui/src/adapters/api/index.ts @@ -122,6 +122,7 @@ export type { ToolStartEvent, ToolEndEvent, TodoItem, + FileArtifactUpdate, ArtifactUpdateEvent, DeepResearchEvent, DeepResearchCallbacks, diff --git a/frontends/ui/src/features/chat/hooks/use-deep-research.spec.ts b/frontends/ui/src/features/chat/hooks/use-deep-research.spec.ts index 5fbce086e..9df9621b4 100644 --- a/frontends/ui/src/features/chat/hooks/use-deep-research.spec.ts +++ b/frontends/ui/src/features/chat/hooks/use-deep-research.spec.ts @@ -919,6 +919,33 @@ describe('useDeepResearch', () => { expect(updates).not.toHaveProperty('deepResearchTodos') }) + test('replay buffer merges file content with a later metadata-only update', async () => { + await setupBufferedHook() + + act(() => { + // Legacy content event, then a durable metadata-only event for the same file. + // Replay must merge (not overwrite) so the earlier content survives. + mockClient?.callbacks.onFileUpdate?.({ filename: 'chart.png', content: 'BASE64DATA' }) + mockClient?.callbacks.onFileUpdate?.({ filename: 'chart.png', artifactId: 'art-1', mimeType: 'image/png' }) + mockClient?.callbacks.onStreamMode?.('live') + }) + + const replayCommit = vi.mocked(useChatStore.setState).mock.calls[0]?.[0] + expect(replayCommit).toEqual(expect.any(Function)) + + const updates = (replayCommit as unknown as (state: { currentStatus: string }) => Record)({ + currentStatus: 'researching', + }) + const files = updates.deepResearchFiles as Array> + expect(files).toHaveLength(1) + expect(files[0]).toMatchObject({ + filename: 'chart.png', + content: 'BASE64DATA', + artifactId: 'art-1', + mimeType: 'image/png', + }) + }) + test('onCitationUpdate adds citation to store', async () => { await setupConnectedHook() @@ -937,7 +964,7 @@ describe('useDeepResearch', () => { await setupConnectedHook() act(() => { - mockClient?.callbacks.onFileUpdate?.('report.md', '# Report content') + mockClient?.callbacks.onFileUpdate?.({ filename: 'report.md', content: '# Report content' }) }) expect(mockAddDeepResearchFile).toHaveBeenCalledWith({ @@ -950,7 +977,7 @@ describe('useDeepResearch', () => { await setupConnectedHook() act(() => { - mockClient?.callbacks.onFileUpdate?.('report.md', '# Final report') + mockClient?.callbacks.onFileUpdate?.({ filename: 'report.md', content: '# Final report' }) }) expect(mockSetCurrentStatus).toHaveBeenCalledWith('writing') @@ -960,7 +987,7 @@ describe('useDeepResearch', () => { await setupConnectedHook() act(() => { - mockClient?.callbacks.onFileUpdate?.('artifacts/report.md', '# Final report') + mockClient?.callbacks.onFileUpdate?.({ filename: 'artifacts/report.md', content: '# Final report' }) }) expect(mockSetCurrentStatus).toHaveBeenCalledWith('writing') @@ -970,7 +997,7 @@ describe('useDeepResearch', () => { await setupConnectedHook() act(() => { - mockClient?.callbacks.onFileUpdate?.('notes.md', '# Some notes') + mockClient?.callbacks.onFileUpdate?.({ filename: 'notes.md', content: '# Some notes' }) }) expect(mockAddDeepResearchFile).toHaveBeenCalledWith({ 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 4cbcae37c..6de2a06f8 100644 --- a/frontends/ui/src/features/chat/hooks/use-deep-research.ts +++ b/frontends/ui/src/features/chat/hooks/use-deep-research.ts @@ -19,6 +19,7 @@ import { cancelJob, type DeepResearchClient, type DeepResearchJobStatus, + type FileArtifactUpdate, type TodoItem, } from '@/adapters/api' import { useChatStore } from '../store' @@ -212,7 +213,7 @@ export const useDeepResearch = (): UseDeepResearchReturn => { 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(), + files: new Map(), reportContent: null as string | null, } @@ -238,7 +239,7 @@ export const useDeepResearch = (): UseDeepResearchReturn => { 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, 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 files = Array.from(buf.files.values()).map((file, i) => ({ id: `file-${i}`, ...file, timestamp: now })) const todos = buf.todos ? normalizeDeepResearchTodos(buf.todos) : undefined useChatStore.setState((state) => ({ @@ -486,13 +487,19 @@ export const useDeepResearch = (): UseDeepResearchReturn => { resetTimeout(); addDeepResearchCitation(url, content, isCited) }, - onFileUpdate: (filename, content) => { - if (buf.active) { buf.files.set(filename, content); return } + onFileUpdate: (file) => { + if (buf.active) { + // Merge like the live store (addDeepResearchFile): a later metadata-only + // event must not drop content from an earlier event for the same filename. + const prev = buf.files.get(file.filename) + buf.files.set(file.filename, prev ? { ...prev, ...file } : file) + return + } if (!isActiveJob()) return - resetTimeout(); addDeepResearchFile({ filename, content }) + resetTimeout(); addDeepResearchFile(file) // report.md artifact arrives 1-2 min before the final_report output event — // use it as an early signal to switch the UI to "writing" status. - if (filename.endsWith('report.md')) { + if (file.filename.endsWith('report.md')) { setCurrentStatus('writing') } }, 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 b39b76543..2d1690c52 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 @@ -28,6 +28,7 @@ import { createDeepResearchClient, type DeepResearchClient, type DeepResearchJobStatus, + type FileArtifactUpdate, type TodoItem, } from '@/adapters/api' import { useChatStore } from '../store' @@ -371,7 +372,7 @@ export const useLoadJobData = (): UseLoadJobDataReturn => { >(), todos: null as TodoItem[] | null, citations: [] as Array<{ url: string; content: string; isCited: boolean }>, - files: new Map(), // filename -> latest content (deduped) + files: new Map(), // filename -> latest event (deduped) reportContent: null as string | null, } @@ -427,10 +428,9 @@ export const useLoadJobData = (): UseLoadJobDataReturn => { timestamp: now, })) - const files = Array.from(buffer.files.entries()).map(([filename, content], idx) => ({ + const files = Array.from(buffer.files.values()).map((file, idx) => ({ id: `file-${idx}`, - filename, - content, + ...file, timestamp: now, })) @@ -574,8 +574,11 @@ export const useLoadJobData = (): UseLoadJobDataReturn => { buffer.citations.push({ url, content, isCited: isCited ?? false }) }, - onFileUpdate: (filename, content) => { - buffer.files.set(filename, content) + onFileUpdate: (file) => { + // Merge like the live store: a later metadata-only event must not drop + // content from an earlier event for the same filename during replay. + const prev = buffer.files.get(file.filename) + buffer.files.set(file.filename, prev ? { ...prev, ...file } : file) }, onOutputUpdate: (content, outputCategory) => { diff --git a/frontends/ui/src/features/chat/store.ts b/frontends/ui/src/features/chat/store.ts index 9a617d6b2..49d26141b 100644 --- a/frontends/ui/src/features/chat/store.ts +++ b/frontends/ui/src/features/chat/store.ts @@ -2913,9 +2913,9 @@ export const useChatStore = create()( const existingIndex = deepResearchFiles.findIndex((f) => f.filename === file.filename) if (existingIndex >= 0) { - // Update existing file with latest content + // Update existing text or durable generated-artifact metadata. const updatedFiles = deepResearchFiles.map((f, i) => - i === existingIndex ? { ...f, content: file.content, timestamp: new Date() } : f + i === existingIndex ? { ...f, ...file, timestamp: new Date() } : f ) set({ deepResearchFiles: updatedFiles }, false, 'addDeepResearchFile:update') return deepResearchFiles[existingIndex].id diff --git a/frontends/ui/src/features/chat/types.ts b/frontends/ui/src/features/chat/types.ts index fc72bb435..f1d9b0db9 100644 --- a/frontends/ui/src/features/chat/types.ts +++ b/frontends/ui/src/features/chat/types.ts @@ -360,8 +360,19 @@ export interface DeepResearchFile { id: string /** File name/path */ filename: string - /** File content */ - content: string + /** Inline text content for legacy text-file events */ + content?: string + /** Durable generated-artifact identifier */ + artifactId?: string + /** Authenticated endpoint for the generated bytes */ + contentUrl?: string + kind?: string + mimeType?: string + sizeBytes?: number + sha256?: string + title?: string + caption?: string + inline?: boolean /** When file was created/updated */ timestamp: Date } diff --git a/frontends/ui/src/features/layout/components/FileCard.spec.tsx b/frontends/ui/src/features/layout/components/FileCard.spec.tsx index 6dc785f14..e49ac2b32 100644 --- a/frontends/ui/src/features/layout/components/FileCard.spec.tsx +++ b/frontends/ui/src/features/layout/components/FileCard.spec.tsx @@ -62,6 +62,30 @@ describe('FileCard', () => { const lineCount = screen.getByText(/lines/) expect(lineCount).toBeInTheDocument() }) + + test('renders durable artifact metadata and opens its authenticated content URL', async () => { + const user = userEvent.setup() + const open = vi.spyOn(window, 'open').mockImplementation(() => null) + render( + + ) + + expect(screen.getByText('image/png · 2.0 KiB')).toBeInTheDocument() + await user.click(screen.getByRole('button', { name: 'Open file' })) + expect(open).toHaveBeenCalledWith( + '/api/jobs/async/job/job-1/artifacts/art-1/content', + '_blank', + 'noopener,noreferrer' + ) + }) }) describe('expand/collapse behavior', () => { diff --git a/frontends/ui/src/features/layout/components/FileCard.tsx b/frontends/ui/src/features/layout/components/FileCard.tsx index 0d370173d..703702bf0 100644 --- a/frontends/ui/src/features/layout/components/FileCard.tsx +++ b/frontends/ui/src/features/layout/components/FileCard.tsx @@ -24,8 +24,17 @@ export interface FileInfo { id: string /** File name/path */ filename: string - /** File content */ - content: string + /** Inline content for legacy text-file events */ + content?: string + artifactId?: string + contentUrl?: string + kind?: string + mimeType?: string + sizeBytes?: number + sha256?: string + title?: string + caption?: string + inline?: boolean /** When file was created/updated */ timestamp?: Date | string } @@ -59,6 +68,12 @@ const isMarkdownFile = (filename: string): boolean => { return ['md', 'markdown'].includes(ext) } +const formatBytes = (size: number): string => { + if (size < 1024) return `${size} B` + if (size < 1024 * 1024) return `${(size / 1024).toFixed(1)} KiB` + return `${(size / (1024 * 1024)).toFixed(1)} MiB` +} + /** * Expandable card showing a file artifact's details. */ @@ -81,7 +96,7 @@ export const FileCard: FC = ({ file }) => { + {file.contentUrl && ( + + + + )} + {/* Collapsed preview */} {!isExpanded && contentPreview && ( diff --git a/frontends/ui/src/pages/api/generate-pdf.ts b/frontends/ui/src/pages/api/generate-pdf.ts index b18455afb..ba1237293 100644 --- a/frontends/ui/src/pages/api/generate-pdf.ts +++ b/frontends/ui/src/pages/api/generate-pdf.ts @@ -5,7 +5,7 @@ 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 { extractArtifactIds, replaceArtifactImages } from '@/shared/utils/artifact-url' import { isAuthRequired } from '@/adapters/auth/config' // Cap the bytes we embed per image. Charts are tiny; this guards against base64-inflating a diff --git a/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.tsx b/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.tsx index efb852fd3..222250ba2 100644 --- a/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.tsx +++ b/frontends/ui/src/shared/components/MarkdownRenderer/MarkdownRenderer.tsx @@ -9,7 +9,7 @@ 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' +import { ARTIFACT_SCHEME, isArtifactRef, resolveArtifactUrl } from '@/shared/utils/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 diff --git a/frontends/ui/src/shared/components/MarkdownRenderer/index.ts b/frontends/ui/src/shared/components/MarkdownRenderer/index.ts index cab48f1a5..2dcab2475 100644 --- a/frontends/ui/src/shared/components/MarkdownRenderer/index.ts +++ b/frontends/ui/src/shared/components/MarkdownRenderer/index.ts @@ -12,4 +12,4 @@ export { replaceArtifactImages, extractArtifactIds, rewriteArtifactRefs, -} from './artifact-url' +} from '@/shared/utils/artifact-url' diff --git a/frontends/ui/src/shared/components/MarkdownRenderer/artifact-url.spec.ts b/frontends/ui/src/shared/utils/artifact-url.spec.ts similarity index 100% rename from frontends/ui/src/shared/components/MarkdownRenderer/artifact-url.spec.ts rename to frontends/ui/src/shared/utils/artifact-url.spec.ts diff --git a/frontends/ui/src/shared/components/MarkdownRenderer/artifact-url.ts b/frontends/ui/src/shared/utils/artifact-url.ts similarity index 100% rename from frontends/ui/src/shared/components/MarkdownRenderer/artifact-url.ts rename to frontends/ui/src/shared/utils/artifact-url.ts diff --git a/src/aiq_agent/agents/deep_researcher/agent.py b/src/aiq_agent/agents/deep_researcher/agent.py index 057806f70..811995065 100644 --- a/src/aiq_agent/agents/deep_researcher/agent.py +++ b/src/aiq_agent/agents/deep_researcher/agent.py @@ -139,6 +139,7 @@ def __init__( tool_set=self.tool_set, source_registry_middleware=self.source_registry_middleware, enable_source_router=self.enable_source_router, + artifact_manager=self.deepagents_runtime.artifact_manager, ) self.source_tool_names = self.tool_set.source_tool_names @@ -183,7 +184,12 @@ def _build_orchestrator_agent(self, state: DeepResearchAgentState) -> Any: def _extract_final_markdown(self, result: dict | Any, files: dict[str, Any] | None = None) -> str | None: """Extract final Markdown from output files.""" output_paths = ("/shared/output.md", "/output.md") - files = result.get("files", None) if isinstance(result, dict) else getattr(result, "files", None) or files or {} + # Resolve result files first, then fall back to the passed-in files (state.files) and + # finally an empty dict. Without the explicit grouping, `or files or {}` bound only to + # the else branch, so a dict result lacking a usable "files" key silently discarded the + # fallback and skipped straight to inline salvage even when output files existed. + result_files = result.get("files", None) if isinstance(result, dict) else getattr(result, "files", None) + files = result_files or files or {} if isinstance(files, dict): for output_path in output_paths: output_entry = files.get(output_path) diff --git a/src/aiq_agent/agents/deep_researcher/custom_middleware.py b/src/aiq_agent/agents/deep_researcher/custom_middleware.py index 23b66d552..944ef39f5 100644 --- a/src/aiq_agent/agents/deep_researcher/custom_middleware.py +++ b/src/aiq_agent/agents/deep_researcher/custom_middleware.py @@ -116,6 +116,47 @@ async def awrap_model_call(self, request, handler): return await handler(request.override(messages=fixed_messages)) +class ExecuteTimeoutClampMiddleware(AgentMiddleware): + """Clamp the sandbox ``execute`` tool's per-call timeout to a configured ceiling. + + The deepagents ``execute`` tool forwards a model-supplied ``timeout`` straight to the + sandbox backend, and providers cap it (OpenShell rejects an ``exec`` timeout above the + gateway maximum). LLMs routinely pass an oversized value -- e.g. milliseconds, where the + backend expects seconds, or an arbitrarily large round number -- so an unclamped timeout + makes every ``execute`` fail with a "timeout exceeds maximum" error and no sandbox code + ever runs. Bound the argument to the configured sandbox lifetime (seconds). + + This guards a different boundary than ``SandboxProvider._clamp_timeout`` in + ``sandbox/base.py``: that clamp covers AI-Q's own provider-mediated calls (e.g. workspace + prep), whereas the deepagents ``execute`` tool reaches the backend without passing through + it, so the untrusted agent argument must be sanitized here at the tool-call boundary. + """ + + def __init__(self, *, max_timeout_seconds: int) -> None: + """Store the ceiling (in seconds) that a single ``execute`` call may request.""" + self.max_timeout_seconds = max(1, int(max_timeout_seconds)) + + async def awrap_tool_call(self, request, handler): + """Clamp an oversized ``timeout`` argument on ``execute`` tool calls.""" + tool_call = request.tool_call + if tool_call.get("name") != "execute": + return await handler(request) + args = tool_call.get("args") + if not isinstance(args, dict) or not isinstance(args.get("timeout"), (int, float)): + return await handler(request) + requested = int(args["timeout"]) + # A non-positive value means "no timeout" to the backend; leave it alone. + if requested <= 0 or requested <= self.max_timeout_seconds: + return await handler(request) + logger.warning( + "Clamping execute timeout %ss -> %ss (agent-supplied value exceeds the sandbox ceiling)", + requested, + self.max_timeout_seconds, + ) + modified = {**tool_call, "args": {**args, "timeout": self.max_timeout_seconds}} + return await handler(request.override(tool_call=modified)) + + # Common hallucinated tool name mappings _TOOL_NAME_ALIASES: dict[str, str] = { "open_file": "read_file", @@ -506,6 +547,28 @@ def get_source_list_text(self, mode: str = "compact") -> str | None: return self._render_source_list_text(self.get_source_entries(mode=mode)) +class ArtifactHarvestMiddleware(AgentMiddleware): + """Checkpoint durable artifacts after successful sandbox execute calls.""" + + def __init__(self, artifact_manager: object) -> None: + """Store the artifact manager used for best-effort checkpoints.""" + self.artifact_manager = artifact_manager + + async def awrap_tool_call(self, request, handler): + """Run the tool, then checkpoint manifest-declared artifacts after execute.""" + result = await handler(request) + tool_name = "" + if hasattr(request, "tool_call") and isinstance(request.tool_call, dict): + tool_name = request.tool_call.get("name", "") + result_status = result.get("status") if isinstance(result, dict) else getattr(result, "status", None) + if tool_name == "execute" and result_status != "error": + try: + await asyncio.to_thread(self.artifact_manager.harvest_after_execute) + except Exception as exc: # noqa: BLE001 - artifact capture must not fail the agent + logger.warning("Artifact checkpoint harvest failed (%s)", type(exc).__name__) + return result + + class PlanPersistenceMiddleware(AgentMiddleware): """Persists the planner's structured ResearchPlan to the shared filesystem. diff --git a/src/aiq_agent/agents/deep_researcher/deepagents_runtime.py b/src/aiq_agent/agents/deep_researcher/deepagents_runtime.py index 9081391b0..793946414 100644 --- a/src/aiq_agent/agents/deep_researcher/deepagents_runtime.py +++ b/src/aiq_agent/agents/deep_researcher/deepagents_runtime.py @@ -19,6 +19,7 @@ import importlib.util import logging +import threading from collections.abc import Callable from pathlib import Path from typing import Any @@ -148,6 +149,8 @@ def __init__( self._backend: Any | None = None self._sandbox_provider: Any | None = None self.artifact_manager: Any | None = None + self._artifact_finalize_lock = threading.Lock() + self._artifact_finalize_attempted = False 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) @@ -178,6 +181,19 @@ def skills_enabled(self) -> bool: """Return true when any agent has configured skill sources.""" return bool(self._skill_sources) + @property + def execute_timeout_seconds(self) -> int | None: + """Per-call ``execute`` timeout ceiling (seconds), sourced from the sandbox config. + + Agent-supplied execute timeouts are unreliable (LLMs pass milliseconds where the + backend expects seconds, or an arbitrary large value), so callers clamp to this + configured sandbox lifetime to keep a single execute under the provider's hard cap. + Returns None when no sandbox is configured (execute is unavailable anyway). + """ + if self._sandbox is None: + return None + return getattr(self._sandbox, "timeout", None) + def skill_sources_for(self, agent_name: str) -> list[str] | None: """Return DeepAgents source paths for an agent/subagent name.""" sources = self._skill_sources_by_agent.get(agent_name) @@ -214,8 +230,8 @@ def final_harvest(self) -> 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) + except Exception as exc: # noqa: BLE001 - harvest is best-effort on the terminal path + logger.warning("Final artifact harvest failed for job %s (%s)", self._job_id, type(exc).__name__) def close(self) -> None: """Release the sandbox provider on a normal terminal job path (idempotent).""" @@ -229,6 +245,33 @@ def terminate(self) -> None: if provider is not None and hasattr(provider, "terminate"): provider.terminate() + def finalize_artifacts(self, *, interrupted: bool) -> bool: + """Run the terminal artifact scan without delaying cancellation. + + Normal success/failure paths perform the idempotent final scan. On cancellation, + only scan when the provider operation lock is immediately available; otherwise + execution teardown takes priority. Successful execute calls are checkpointed by + middleware, so already completed artifacts remain durable in either case. + """ + with self._artifact_finalize_lock: + if self._artifact_finalize_attempted: + return False + self._artifact_finalize_attempted = True + if self.artifact_manager is None or self._sandbox_provider is None: + return False + if not interrupted: + self.final_harvest() + return True + lease = getattr(self._sandbox_provider, "try_operation_lease", None) + if not callable(lease): + return False + with lease() as acquired: + if not acquired: + logger.info("Skipping terminal harvest for busy cancelled sandbox job %s", self._job_id) + return False + self.final_harvest() + return True + def prepare_state_files(self, files: dict[str, Any]) -> dict[str, Any]: """Normalize seeded virtual filesystem files for the configured backend.""" return _normalize_state_files(files, strip_shared_route=self._sandbox is not None) diff --git a/src/aiq_agent/agents/deep_researcher/factory.py b/src/aiq_agent/agents/deep_researcher/factory.py index 833c2d46e..1556f7478 100644 --- a/src/aiq_agent/agents/deep_researcher/factory.py +++ b/src/aiq_agent/agents/deep_researcher/factory.py @@ -41,7 +41,9 @@ from aiq_agent.common import LLMRole from aiq_agent.common import render_prompt_template +from .custom_middleware import ArtifactHarvestMiddleware from .custom_middleware import EmptyContentFixMiddleware +from .custom_middleware import ExecuteTimeoutClampMiddleware from .custom_middleware import PlanPersistenceMiddleware from .custom_middleware import SourceRegistryMiddleware from .custom_middleware import SourceRoutingGuardMiddleware @@ -200,13 +202,14 @@ def build_common_middleware( *, tool_set: DeepResearchToolSet, source_registry_middleware: SourceRegistryMiddleware, + artifact_manager: object | None = None, extra_valid_tool_names: Sequence[str] = (), ) -> list[Any]: """Build the shared middleware stack with agent-specific valid tool names.""" valid_tool_names = {tool.name for tool in [*tool_set.all_tools, *tool_set.researcher_tools]} valid_tool_names.update(FILESYSTEM_TOOL_NAMES) valid_tool_names.update(extra_valid_tool_names) - return [ + middleware: list[Any] = [ EmptyContentFixMiddleware(), ToolNameSanitizationMiddleware(valid_tool_names=sorted(valid_tool_names)), ToolRetryMiddleware(max_retries=3, backoff_factor=2.0, initial_delay=1.0), @@ -214,6 +217,9 @@ def build_common_middleware( ToolResultPruningMiddleware(keep_last_n=10, max_chars=2000), ModelRetryMiddleware(max_retries=2, backoff_factor=2.0, initial_delay=1.0), ] + if artifact_manager is not None: + middleware.append(ArtifactHarvestMiddleware(artifact_manager)) + return middleware def build_orchestrator_middleware( @@ -262,6 +268,7 @@ def build_deep_research_middleware_set( tool_set: DeepResearchToolSet, source_registry_middleware: SourceRegistryMiddleware, enable_source_router: bool = True, + artifact_manager: object | None = None, ) -> DeepResearchMiddlewareSet: """Build researcher, writer, and orchestrator middleware stacks.""" @@ -270,6 +277,7 @@ def common(extra_valid_tool_names: Sequence[str] = ()) -> list[Any]: return build_common_middleware( tool_set=tool_set, source_registry_middleware=source_registry_middleware, + artifact_manager=artifact_manager, extra_valid_tool_names=extra_valid_tool_names, ) @@ -485,6 +493,17 @@ def build_deep_research_graph( enable_source_router: bool = True, ) -> Any: """Build the full DeepAgents graph for one deep research run.""" + # Cross-cutting middleware applied to every agent (researcher, subagents, orchestrator). + # Agent-supplied execute timeouts are unreliable (LLMs pass milliseconds or arbitrarily + # large values); clamp them to the configured sandbox lifetime so a single execute never + # exceeds the provider's hard cap and silently fails every code run. + cross_cutting_middleware = runtime_visibility_middleware(runtime) + execute_ceiling = runtime.execute_timeout_seconds + if execute_ceiling: + cross_cutting_middleware = [ + ExecuteTimeoutClampMiddleware(max_timeout_seconds=execute_ceiling), + *cross_cutting_middleware, + ] context = DeepResearchGraphContext( llm_provider=llm_provider, state=state, @@ -498,7 +517,7 @@ def build_deep_research_graph( max_research_concurrency=max_research_concurrency, enable_source_router=enable_source_router, backend=runtime.backend, - visibility_middleware=runtime_visibility_middleware(runtime), + visibility_middleware=cross_cutting_middleware, ) researcher_model = context.llm_provider.get(LLMRole.RESEARCHER) researcher_skill_sources = context.skill_sources(RESEARCHER_AGENT) diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/README.md b/src/aiq_agent/agents/deep_researcher/sandbox/README.md index baa502b4a..a83e01a3e 100644 --- a/src/aiq_agent/agents/deep_researcher/sandbox/README.md +++ b/src/aiq_agent/agents/deep_researcher/sandbox/README.md @@ -153,12 +153,15 @@ 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 +- Successful `execute` calls trigger a manifest-only checkpoint. Terminal finalization runs + one manifest + directory scan on success or failure. On cancellation, that scan runs only + when the provider operation lease is immediately available; a busy sandbox is terminated + immediately, while completed execute outputs remain preserved by earlier checkpoints. +- 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 `ArtifactStore`, 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. + reject -> quota -> SVG sanitize -> sha256), stores metadata in SQL and bytes through the + configured artifact blob provider, then emits an + `artifact.update` event (durable metadata + `content_url`, never bytes or URL-as-text). - 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` diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/manager.py b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/manager.py index 5cae44941..66dbd1291 100644 --- a/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/manager.py +++ b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/manager.py @@ -178,15 +178,30 @@ def __init__( self._emit = emit self._content_url_template = content_url_template self._lock = threading.Lock() + self._final_harvest_lock = threading.Lock() + self._final_harvest_done = False + self._final_harvest_result: list[Artifact] = [] 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.""" + """Harvest at a terminal boundary, with a directory scan fallback.""" if not self.config.enabled: return [] - return self._harvest(scan=True) + with self._final_harvest_lock: + if self._final_harvest_done: + return list(self._final_harvest_result) + captured = self._harvest(scan=True) + self._final_harvest_result = list(captured) + self._final_harvest_done = True + return captured + + def harvest_after_execute(self) -> list[Artifact]: + """Checkpoint manifest-declared artifacts after a sandbox execute call.""" + if not self.config.enabled: + return [] + return self._harvest(scan=False) def resolve_report_references(self, markdown: str, artifacts: list[Artifact] | None = None) -> str: """Validate ``artifact://`` image references against this job's artifacts. diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/models.py b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/models.py index 0ca080d05..95c3a4bf3 100644 --- a/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/models.py +++ b/src/aiq_agent/agents/deep_researcher/sandbox/artifacts/models.py @@ -84,24 +84,30 @@ class Artifact(BaseModel): 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. + """Build the canonical ``artifact.update`` 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. + A JSON-serializable payload (no bytes) matching the API/UI event contract. """ 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, + "type": "artifact.update", + "name": self.filename, + "data": { + "type": "file", + "url": content_url, + "content_url": content_url, + "file_path": self.filename, + "artifact_id": self.artifact_id, + "job_id": self.job_id, + "kind": self.kind.value, + "mime_type": self.mime_type, + "size_bytes": self.size_bytes, + "sha256": self.sha256, + "title": self.title, + "caption": self.caption, + "inline": self.inline, + }, } diff --git a/src/aiq_agent/agents/deep_researcher/sandbox/base.py b/src/aiq_agent/agents/deep_researcher/sandbox/base.py index 20d726f77..4331c2329 100644 --- a/src/aiq_agent/agents/deep_researcher/sandbox/base.py +++ b/src/aiq_agent/agents/deep_researcher/sandbox/base.py @@ -33,6 +33,7 @@ from abc import ABC from abc import abstractmethod from collections.abc import Callable +from contextlib import contextmanager from typing import TYPE_CHECKING from typing import TypeVar @@ -140,6 +141,16 @@ def is_recoverable_error(self, exc: Exception) -> bool: """ return False + @contextmanager + def try_operation_lease(self): + """Yield whether the provider is idle without waiting behind an in-flight execute.""" + acquired = self._lock.acquire(blocking=False) + try: + yield acquired + finally: + if acquired: + self._lock.release() + def close(self) -> None: """Release the underlying sandbox session, if any (idempotent). diff --git a/tests/aiq_agent/agents/deep_researcher/sandbox/test_artifacts.py b/tests/aiq_agent/agents/deep_researcher/sandbox/test_artifacts.py index eadccea91..ce33bb176 100644 --- a/tests/aiq_agent/agents/deep_researcher/sandbox/test_artifacts.py +++ b/tests/aiq_agent/agents/deep_researcher/sandbox/test_artifacts.py @@ -48,14 +48,18 @@ class _FakeBackend: def __init__(self, files: dict[str, bytes]) -> None: self.files = files + self.execute_calls: list[str] = [] + self.download_calls: list[list[str]] = [] def download_files(self, paths: list[str]) -> list[Any]: + self.download_calls.append(list(paths)) 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: + self.execute_calls.append(command) return SimpleNamespace(output="\n".join(self.files), exit_code=0) @@ -126,8 +130,12 @@ def test_captures_manifest_artifact(self, tmp_path: Any) -> None: 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 + assert emitted and emitted[0]["type"] == "artifact.update" + assert emitted[0]["name"] == "chart.png" + assert emitted[0]["data"]["type"] == "file" + assert emitted[0]["data"]["artifact_id"] == captured[0].artifact_id + assert emitted[0]["data"]["content_url"].endswith(f"/{captured[0].artifact_id}/content") + assert "content" not in emitted[0]["data"] # bytes and URL-as-text never enter the payload def test_rejects_path_traversal(self, tmp_path: Any) -> None: store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") @@ -168,8 +176,40 @@ def test_dedups_identical_content(self, tmp_path: Any) -> None: files = {f"{_ARTIFACT_DIR}/manifest.json": _manifest_bytes(png_path), png_path: _PNG} manager, _ = _make_manager(store, files) manager.final_harvest() + first_downloads = list(manager.backend.download_calls) manager.final_harvest() # same bytes again assert len(store.list("job-1")) == 1 + assert manager.backend.download_calls == first_downloads + + def test_checkpoint_harvest_uses_manifest_without_directory_scan(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) + + captured = manager.harvest_after_execute() + + assert [artifact.filename for artifact in captured] == ["chart.png"] + assert manager.backend.execute_calls == [] + + def test_final_harvest_after_checkpoint_does_not_reemit(self, tmp_path: Any) -> None: + 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, emitted = _make_manager(store, files) + + checkpointed = manager.harvest_after_execute() + assert [artifact.filename for artifact in checkpointed] == ["chart.png"] + + finalized = manager.final_harvest() + # chart.png was already captured at checkpoint; only the scan-discovered CSV is new. + assert [artifact.filename for artifact in finalized] == ["chart.csv"] + assert len(emitted) == 2 def test_scan_fallback_without_manifest(self, tmp_path: Any) -> None: store = SqlArtifactStore(f"sqlite:///{tmp_path}/jobs.db") diff --git a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py index 0ca4117ad..bd56260f1 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py +++ b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py @@ -25,6 +25,8 @@ from langchain_core.messages import SystemMessage from langchain_core.messages import ToolMessage +from aiq_agent.agents.deep_researcher.custom_middleware import ArtifactHarvestMiddleware +from aiq_agent.agents.deep_researcher.custom_middleware import ExecuteTimeoutClampMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import PlanPersistenceMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import SourceRegistryMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import SourceRoutingGuardMiddleware @@ -116,6 +118,85 @@ async def test_disabled_guard_is_noop(self): assert result is expected +class TestExecuteTimeoutClampMiddleware: + """Tests for clamping the sandbox execute tool's per-call timeout.""" + + @staticmethod + def _request(tool_name: str, *, args: dict | None = None) -> MagicMock: + request = MagicMock() + request.tool_call = {"name": tool_name, "args": args if args is not None else {}, "id": "tc1"} + + def _override(*, tool_call): + overridden = MagicMock() + overridden.tool_call = tool_call + return overridden + + request.override.side_effect = _override + return request + + @pytest.mark.asyncio + async def test_clamps_oversized_timeout(self): + """An agent timeout above the ceiling is reduced to the configured maximum.""" + middleware = ExecuteTimeoutClampMiddleware(max_timeout_seconds=1200) + handler = AsyncMock(return_value=ToolMessage(content="ok", tool_call_id="tc1")) + request = self._request("execute", args={"command": "python x.py", "timeout": 120000}) + + await middleware.awrap_tool_call(request, handler) + + request.override.assert_called_once() + forwarded = handler.await_args.args[0] + assert forwarded.tool_call["args"]["timeout"] == 1200 + assert forwarded.tool_call["args"]["command"] == "python x.py" + + @pytest.mark.asyncio + async def test_timeout_within_ceiling_passthrough(self): + """A reasonable timeout is left untouched (no override).""" + middleware = ExecuteTimeoutClampMiddleware(max_timeout_seconds=1200) + handler = AsyncMock(return_value=ToolMessage(content="ok", tool_call_id="tc1")) + request = self._request("execute", args={"command": "python x.py", "timeout": 60}) + + await middleware.awrap_tool_call(request, handler) + + request.override.assert_not_called() + handler.assert_awaited_once_with(request) + + @pytest.mark.asyncio + async def test_nonpositive_timeout_passthrough(self): + """A non-positive timeout means 'no timeout' to the backend and is not clamped.""" + middleware = ExecuteTimeoutClampMiddleware(max_timeout_seconds=1200) + handler = AsyncMock(return_value=ToolMessage(content="ok", tool_call_id="tc1")) + request = self._request("execute", args={"command": "python x.py", "timeout": 0}) + + await middleware.awrap_tool_call(request, handler) + + request.override.assert_not_called() + handler.assert_awaited_once_with(request) + + @pytest.mark.asyncio + async def test_missing_timeout_passthrough(self): + """execute calls without a timeout arg are forwarded unchanged.""" + middleware = ExecuteTimeoutClampMiddleware(max_timeout_seconds=1200) + handler = AsyncMock(return_value=ToolMessage(content="ok", tool_call_id="tc1")) + request = self._request("execute", args={"command": "python x.py"}) + + await middleware.awrap_tool_call(request, handler) + + request.override.assert_not_called() + handler.assert_awaited_once_with(request) + + @pytest.mark.asyncio + async def test_non_execute_tool_passthrough(self): + """A large timeout on a non-execute tool is ignored by this middleware.""" + middleware = ExecuteTimeoutClampMiddleware(max_timeout_seconds=1200) + handler = AsyncMock(return_value=ToolMessage(content="ok", tool_call_id="tc1")) + request = self._request("ls", args={"path": "/shared", "timeout": 120000}) + + await middleware.awrap_tool_call(request, handler) + + request.override.assert_not_called() + handler.assert_awaited_once_with(request) + + class TestToolNameSanitizationMiddleware: """Tests for ToolNameSanitizationMiddleware.""" @@ -626,6 +707,75 @@ async def test_content_returned_unchanged(self, middleware): assert result.content == content +class TestArtifactHarvestMiddleware: + """Checkpoint harvesting runs only after successful execute tool calls.""" + + @pytest.mark.asyncio + async def test_execute_checkpoints_after_handler(self) -> None: + manager = MagicMock() + middleware = ArtifactHarvestMiddleware(manager) + request = MagicMock() + request.tool_call = {"name": "execute"} + handler = AsyncMock(return_value="ok") + + result = await middleware.awrap_tool_call(request, handler) + + assert result == "ok" + manager.harvest_after_execute.assert_called_once_with() + + @pytest.mark.asyncio + async def test_non_execute_tool_does_not_harvest(self) -> None: + manager = MagicMock() + middleware = ArtifactHarvestMiddleware(manager) + request = MagicMock() + request.tool_call = {"name": "read_file"} + + await middleware.awrap_tool_call(request, AsyncMock(return_value="ok")) + + manager.harvest_after_execute.assert_not_called() + + @pytest.mark.asyncio + async def test_handler_failure_does_not_harvest(self) -> None: + manager = MagicMock() + middleware = ArtifactHarvestMiddleware(manager) + request = MagicMock() + request.tool_call = {"name": "execute"} + + with pytest.raises(RuntimeError, match="tool failed"): + await middleware.awrap_tool_call(request, AsyncMock(side_effect=RuntimeError("tool failed"))) + + manager.harvest_after_execute.assert_not_called() + + @pytest.mark.asyncio + async def test_execute_error_result_does_not_harvest(self) -> None: + manager = MagicMock() + middleware = ArtifactHarvestMiddleware(manager) + request = MagicMock() + request.tool_call = {"name": "execute"} + + await middleware.awrap_tool_call( + request, + AsyncMock(return_value=ToolMessage(content="failed", tool_call_id="tc1", status="error")), + ) + + manager.harvest_after_execute.assert_not_called() + + @pytest.mark.asyncio + async def test_checkpoint_failure_logs_only_exception_type(self, caplog: pytest.LogCaptureFixture) -> None: + manager = MagicMock() + manager.harvest_after_execute.side_effect = RuntimeError("credential=do-not-log") + middleware = ArtifactHarvestMiddleware(manager) + request = MagicMock() + request.tool_call = {"name": "execute"} + + with caplog.at_level(logging.WARNING): + result = await middleware.awrap_tool_call(request, AsyncMock(return_value="ok")) + + assert result == "ok" + assert "RuntimeError" in caplog.text + assert "credential=do-not-log" not in caplog.text + + class _RecordingBackend: """Minimal backend stub capturing upload_files calls (overwrite-safe).""" 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 5e98ef489..009d68523 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py +++ b/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py @@ -17,6 +17,8 @@ from __future__ import annotations +import logging +from contextlib import nullcontext from typing import Any from unittest.mock import MagicMock from unittest.mock import patch @@ -301,3 +303,54 @@ def find_spec(module_name: str): pytest.raises(ImportError, match="langchain-modal"), ): _ = DeepAgentsRuntime(sandbox=DeepResearchSandboxConfig(provider="modal")).backend + + +class TestDeepAgentsRuntimeArtifacts: + """Terminal artifact harvesting is safe on normal and interrupted paths.""" + + def test_final_harvest_logs_only_exception_type(self, caplog: pytest.LogCaptureFixture) -> None: + provider = MagicMock() + with patch( + "aiq_agent.agents.deep_researcher.deepagents_runtime._create_sandbox_backend", + return_value=provider, + ): + runtime = DeepAgentsRuntime(sandbox=DeepResearchSandboxConfig()) + runtime.artifact_manager = MagicMock() + runtime.artifact_manager.final_harvest.side_effect = RuntimeError("credential=do-not-log") + + with caplog.at_level(logging.WARNING): + runtime.final_harvest() + + assert "RuntimeError" in caplog.text + assert "credential=do-not-log" not in caplog.text + + def test_normal_finalize_artifacts_harvests(self) -> None: + provider = MagicMock() + with patch( + "aiq_agent.agents.deep_researcher.deepagents_runtime._create_sandbox_backend", + return_value=provider, + ): + runtime = DeepAgentsRuntime(sandbox=DeepResearchSandboxConfig()) + runtime.artifact_manager = MagicMock() + + assert runtime.finalize_artifacts(interrupted=False) is True + assert runtime.finalize_artifacts(interrupted=False) is False + runtime.artifact_manager.final_harvest.assert_called_once_with() + + @pytest.mark.parametrize(("lease_acquired", "expected"), [(True, True), (False, False)]) + def test_interrupted_finalize_harvests_only_when_provider_is_idle( + self, + lease_acquired: bool, + expected: bool, + ) -> None: + provider = MagicMock() + provider.try_operation_lease.return_value = nullcontext(lease_acquired) + with patch( + "aiq_agent.agents.deep_researcher.deepagents_runtime._create_sandbox_backend", + return_value=provider, + ): + runtime = DeepAgentsRuntime(sandbox=DeepResearchSandboxConfig()) + runtime.artifact_manager = MagicMock() + + assert runtime.finalize_artifacts(interrupted=True) is expected + assert runtime.artifact_manager.final_harvest.call_count == int(lease_acquired) diff --git a/tests/aiq_agent/agents/deep_researcher/test_factory.py b/tests/aiq_agent/agents/deep_researcher/test_factory.py index a35c6bf78..13fe705e8 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_factory.py +++ b/tests/aiq_agent/agents/deep_researcher/test_factory.py @@ -23,6 +23,7 @@ from langchain.agents.middleware import AgentMiddleware from langchain_core.tools import tool +from aiq_agent.agents.deep_researcher.custom_middleware import ArtifactHarvestMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import SourceRegistryMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import SourceRoutingGuardMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import TodoSuppressionMiddleware @@ -186,6 +187,28 @@ def test_middleware_set_configures_orchestrator_source_routing_guard(): assert not any(isinstance(item, SourceRoutingGuardMiddleware) for item in enabled.writer) +def test_middleware_set_wires_artifact_checkpoint_when_manager_present(): + """Artifact harvesting is attached to every execute-capable middleware stack.""" + registry = SourceRegistryMiddleware(source_tool_names={web_search_tool.name}) + tool_set = build_deep_research_tool_set( + [web_search_tool], + source_registry_middleware=registry, + max_concurrent_source_tool_calls=2, + max_source_tool_batch_size=3, + ) + manager = MagicMock() + + middleware_set = build_deep_research_middleware_set( + tool_set=tool_set, + source_registry_middleware=registry, + artifact_manager=manager, + ) + + for middleware in (middleware_set.researcher, middleware_set.planner, middleware_set.writer): + checkpoint = next(item for item in middleware if isinstance(item, ArtifactHarvestMiddleware)) + assert checkpoint.artifact_manager is manager + + def test_subagents_route_tools_and_writer_skills(): """Source-router excludes source tools, planner receives source tools, and writer receives configured skills.""" runtime = DeepAgentsRuntime( diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index 6f77d4476..080355b71 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -910,6 +910,37 @@ def test_store_event(self, tmp_path): assert len(events) == 1 assert events[0]["type"] == "test.event" + def test_artifact_update_survives_event_store_round_trip(self, tmp_path): + """Generated-file metadata remains reconstructable after durable storage.""" + from aiq_agent.agents.deep_researcher.sandbox.artifacts import Artifact + from aiq_agent.agents.deep_researcher.sandbox.artifacts import ArtifactKind + from aiq_api.jobs.event_store import EventStore + + db_url = f"sqlite:///{tmp_path / 'test.db'}" + content_url = "/v1/jobs/async/job/job-1/artifacts/artifact-1/content" + artifact = Artifact( + artifact_id="artifact-1", + job_id="job-1", + kind=ArtifactKind.IMAGE, + mime_type="image/png", + filename="chart.png", + sandbox_path="/sandbox/job-1/aiq-artifacts/chart.png", + storage_uri="db://artifacts/artifact-1", + sha256="a" * 64, + size_bytes=128, + inline=True, + ) + + EventStore(db_url, "job-1").store(artifact.to_sse_payload(content_url)) + + event = EventStore.get_events(db_url, "job-1")[0] + assert event["type"] == "artifact.update" + assert event["name"] == "chart.png" + assert event["data"]["type"] == "file" + assert event["data"]["content_url"] == content_url + assert "content" not in event["data"] + assert event["data"]["artifact_id"] == "artifact-1" + def test_get_events_empty(self, tmp_path): """Test get_events returns empty list for unknown job.""" from aiq_api.jobs.event_store import EventStore @@ -2447,11 +2478,41 @@ def test_never_raises_when_teardown_fails(self): # 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): + def test_finalizes_artifacts_before_close(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"]) + order: list[str] = [] + runtime = MagicMock(spec=["close", "terminate", "finalize_artifacts"]) + runtime.finalize_artifacts.side_effect = lambda **_kwargs: order.append("harvest") + runtime.close.side_effect = lambda: order.append("close") + _teardown_sandbox(runtime, job_id="job-1", interrupted=False) - runtime.final_harvest.assert_not_called() + runtime.finalize_artifacts.assert_called_once_with(interrupted=False) + assert order == ["harvest", "close"] + + def test_harvest_persists_artifacts_without_releasing_sandbox(self): + from aiq_api.jobs.runner import _harvest_sandbox_artifacts + + # Runs before the terminal status: artifacts must be persisted, but the + # unbounded close()/terminate() must NOT run here (deferred to finally), + # so a hanging SDK cleanup cannot strand a finished job in RUNNING. + runtime = MagicMock(spec=["close", "terminate", "finalize_artifacts"]) + _harvest_sandbox_artifacts(runtime, job_id="job-1", interrupted=False) + + runtime.finalize_artifacts.assert_called_once_with(interrupted=False) + runtime.close.assert_not_called() + runtime.terminate.assert_not_called() + + def test_harvest_none_runtime_is_noop(self): + from aiq_api.jobs.runner import _harvest_sandbox_artifacts + + _harvest_sandbox_artifacts(None, job_id="job-1", interrupted=False) + + def test_harvest_never_raises_when_finalize_fails(self): + from aiq_api.jobs.runner import _harvest_sandbox_artifacts + + runtime = MagicMock(spec=["finalize_artifacts"]) + runtime.finalize_artifacts.side_effect = RuntimeError("artifact scan failed") + # Artifact capture cannot replace or block the job result. + _harvest_sandbox_artifacts(runtime, job_id="job-1", interrupted=False)