Skip to content
3 changes: 3 additions & 0 deletions docs/source/architecture/agents/sandbox.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
2 changes: 1 addition & 1 deletion frontends/aiq_api/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
48 changes: 32 additions & 16 deletions frontends/aiq_api/src/aiq_api/jobs/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -798,7 +798,14 @@ async def run_agent_job(
# Signal event stream completion
event_stream.on_complete()

# Flush any buffered events before updating status
# Finalize before SUCCESS so clients cannot stop streaming before
# the terminal artifact scan and emitted metadata are durable.
await asyncio.to_thread(
Comment thread
KyleZheng1284 marked this conversation as resolved.
_teardown_sandbox,
sandbox_runtime,
job_id=job_id,
interrupted=False,
)
if hasattr(event_store, "flush"):
event_store.flush()

Expand Down Expand Up @@ -832,17 +839,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,
{
Expand All @@ -851,14 +851,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(_teardown_sandbox, sandbox_runtime, job_id=job_id, interrupted=False)
_store_terminal_event_best_effort(
event_store,
{
Expand All @@ -869,6 +875,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.
Expand All @@ -883,9 +891,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:
Expand All @@ -910,14 +916,24 @@ def _store_terminal_event_best_effort(event_store, event: dict) -> None:


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
off the event loop (``asyncio.to_thread``) so the SDK session close cannot block the worker.
"""
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__,
)
teardown = getattr(sandbox_runtime, "terminate", None) if interrupted else None
if teardown is None:
teardown = getattr(sandbox_runtime, "close", None)
Expand Down
76 changes: 75 additions & 1 deletion frontends/ui/src/adapters/api/deep-research-client.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, (event: MessageEvent) => 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(() => {
Expand Down Expand Up @@ -32,4 +54,56 @@ 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' })
)
})
Comment thread
coderabbitai[bot] marked this conversation as resolved.
})
63 changes: 58 additions & 5 deletions frontends/ui/src/adapters/api/deep-research-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
*/

import { apiConfig } from './config'
import { artifactContentPath } from '@/shared/components/MarkdownRenderer/artifact-url'
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

// ============================================================
// Types
Expand Down Expand Up @@ -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'
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -485,10 +512,17 @@ 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 }
data?: {
type: ArtifactType
content?: string | TodoItem[]
url?: string
content_url?: string
output_category?: string
}
type?: ArtifactType
content?: string | TodoItem[]
url?: string
content_url?: string
output_category?: string
metadata?: { workflow?: string }
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated
Expand All @@ -509,11 +543,30 @@ 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
// Generated artifacts carry durable metadata; legacy text-file events carry content.
const raw = artifactData as Record<string, unknown>
const filePath = (raw.file_path || raw.path || artifactData.url || 'unknown') as string
const fileName = filePath.split('/').pop() || filePath
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated
callbacks.onFileUpdate?.(fileName, artifactData.content as string)
const artifactId = typeof raw.artifact_id === 'string' ? raw.artifact_id : undefined
callbacks.onFileUpdate?.({
filename: fileName,
content: typeof artifactData.content === 'string' ? artifactData.content : undefined,
artifactId,
contentUrl: artifactId
? artifactContentPath(jobId, artifactId)
: typeof raw.content_url === 'string'
? raw.content_url
: typeof raw.url === 'string'
? raw.url
: undefined,
kind: typeof raw.kind === 'string' ? raw.kind : undefined,
mimeType: typeof raw.mime_type === 'string' ? raw.mime_type : undefined,
sizeBytes: typeof raw.size_bytes === 'number' ? raw.size_bytes : undefined,
sha256: typeof raw.sha256 === 'string' ? raw.sha256 : undefined,
title: typeof raw.title === 'string' ? raw.title : undefined,
caption: typeof raw.caption === 'string' ? raw.caption : undefined,
inline: typeof raw.inline === 'boolean' ? raw.inline : undefined,
})
break
}
case 'output':
Expand Down
1 change: 1 addition & 0 deletions frontends/ui/src/adapters/api/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@ export type {
ToolStartEvent,
ToolEndEvent,
TodoItem,
FileArtifactUpdate,
ArtifactUpdateEvent,
DeepResearchEvent,
DeepResearchCallbacks,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -937,7 +937,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({
Expand All @@ -950,7 +950,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')
Expand All @@ -960,7 +960,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')
Expand All @@ -970,7 +970,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({
Expand Down
Loading
Loading