From 71e5bba55150ba724972829f4ed1b1b7dfa05479 Mon Sep 17 00:00:00 2001 From: somew1nd Date: Fri, 21 Aug 2026 00:34:44 -0700 Subject: [PATCH] runtime: harden LLM retry handling and operator workflows --- .gitignore | 1 + .../database/repositories/step_repository.py | 4 + .../agent-stream/agent-stream.component.html | 81 +++---- .../agent-stream/agent-stream.component.scss | 199 ++++++++---------- .../agent-stream/agent-stream.component.ts | 87 +++++--- .../src/app/services/agent.service.spec.ts | 110 ++++++++++ .../src/app/services/agent.service.ts | 183 +++++++++++----- .../src/app/utils/tool-formatter.util.spec.ts | 34 ++- .../src/app/utils/tool-formatter.util.ts | 97 ++++++--- artemis/agents/operator/operator.py | 24 +-- artemis/agents/operator/prompt_builder.py | 34 +-- artemis/agents/operator/prompts.py | 61 ++++-- artemis/context.py | 32 ++- artemis/services/llm.py | 109 ++++++++-- .../admin_console/test_step_repository.py | 39 ++++ tests/unit/agents/test_operator_prompts.py | 55 +++++ tests/unit/test_context.py | 33 +++ tests/unit/test_llm_pause.py | 102 +++++++-- 18 files changed, 937 insertions(+), 348 deletions(-) create mode 100644 apps/showcase_ui/src/app/services/agent.service.spec.ts create mode 100644 tests/unit/agents/test_operator_prompts.py diff --git a/.gitignore b/.gitignore index e97851c..a2e647c 100644 --- a/.gitignore +++ b/.gitignore @@ -54,6 +54,7 @@ tests/**/outputs/ dmypy.json # IDEs & Editors +.vs/ .vscode/ .idea/ .cursor/ diff --git a/apps/admin_console/database/repositories/step_repository.py b/apps/admin_console/database/repositories/step_repository.py index 9b689c3..a9557ee 100644 --- a/apps/admin_console/database/repositories/step_repository.py +++ b/apps/admin_console/database/repositories/step_repository.py @@ -186,6 +186,10 @@ class StepRepository: "source", "recoverable", "pause", + "request_id", + "scheduled_at", + "waited_seconds", + "retries", ) cleaned = { key: self._clean_value(payload_obj.get(key)) diff --git a/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.html b/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.html index c8263d3..f27b331 100644 --- a/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.html +++ b/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.html @@ -528,50 +528,36 @@ - -
-
- cloud_sync - AI service retry - - {{ getLLMRetryCount(item.data) }}× retried - - - {{ isActionExpanded('retry-' + block.id + '-' + itemIdx) ? 'expand_less' : 'expand_more' }} - -
-
-
-
- Attempts - {{ getLLMRetryCount(item.data) }} -
-
- Extra wait - {{ totalDelay }} -
-
- Provider - {{ retryProvider }} -
-
-
-
-
- Retry {{ retryIndex + 1 }} - +{{ retryDelay }} -
-
{{ getLLMRetryEntryError(retry) }}
-
+ +
+
+ +
+ + + LLM temporarily unavailable + + · Retrying in {{ retryDelay }}… + + + + Frequent retries may indicate a low API rate limit. Consider increasing the limit or reducing request concurrency. + + + + Retrying LLM request · Waited {{ getLLMRetryEntryWaited(retry) }} +
-
+
error LLM Request Failed @@ -596,12 +582,29 @@ play_arrow Resume + + {{ isActionExpanded('llm-failure-' + block.id + '-' + itemIdx) ? 'expand_less' : 'expand_more' }} +
Reason: {{ getLLMErrorText(item.data) }}
+
+ Waited: + {{ waited }} +
+
+
Retry details
+
+ Retry {{ retryIndex + 1 }} + + Planned {{ getLLMRetryEntryDelay(retry) || '0s' }} + · Waited {{ getLLMRetryEntryWaited(retry) }} + +
+
diff --git a/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.scss b/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.scss index c22b0ba..6fe71b7 100644 --- a/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.scss +++ b/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.scss @@ -1955,6 +1955,55 @@ $uniform-gap: 8px; margin-top: 0; margin-bottom: 0; + .llm-retry-notes { + display: flex; + flex-direction: column; + gap: $uniform-gap; + padding: 0; + + .llm-retry-note { + min-height: 20px; + display: flex; + align-items: center; + gap: 8px; + color: #71717a; + font-family: 'Google Sans', 'Google Sans Text', 'Inter', Arial, sans-serif; + font-size: 12px; + line-height: 1.5; + letter-spacing: 0.2px; + + .stream-status-dot { + width: 8px; + height: 8px; + border-radius: 50%; + background: #d4d4d8; + flex: 0 0 auto; + + &.streaming { + background: #3b82f6; + box-shadow: 0 0 0 0 rgba(59, 130, 246, 0.45); + animation: pulse-stream-dot 1.8s cubic-bezier(0.4, 0, 0.6, 1) infinite; + } + } + + .llm-retry-note-text { + color: #71717a; + display: flex; + flex-direction: column; + min-width: 0; + + .llm-retry-status-line, + .llm-retry-guidance { + display: block; + } + + .llm-retry-guidance { + color: #a1a1aa; + } + } + } + } + .step-generic-tool { border: 1px solid transparent; @@ -2245,6 +2294,10 @@ $uniform-gap: 8px; background: #f8fafc; border-color: #fca5a5; + &.has-details { + cursor: pointer; + } + &:hover { border-color: #f87171; } @@ -2255,19 +2308,6 @@ $uniform-gap: 8px; } } - &.is-retrying { - background: #fffbeb; - border-color: #fcd34d; - - &:hover { - border-color: #f59e0b; - } - - .action-card-body .detail-value { - color: #78350f; - } - } - &.task-report-card { background: #f0fdf4; border: 1px solid #bbf7d0; @@ -2373,9 +2413,6 @@ $uniform-gap: 8px; color: #ef4444; } - &.is-retrying { - color: #d97706; - } } .action-title { @@ -2442,27 +2479,6 @@ $uniform-gap: 8px; } } - &.is-retrying { - background: #fef3c7; - color: #92400e; - border: 1px solid #fde68a; - - .status-icon { - color: #d97706; - - &.is-active { - animation: spin 1.5s linear infinite; - } - } - } - } - - &.llm-retry-header { - border-bottom-color: #fde68a; - - .action-title { - color: #78350f; - } } &.llm-failure-header { @@ -2528,83 +2544,6 @@ $uniform-gap: 8px; flex-direction: column; gap: 6px; - .retry-summary-grid { - display: grid; - grid-template-columns: repeat(3, minmax(0, 1fr)); - gap: 6px; - } - - .retry-summary-item { - min-width: 0; - padding: 7px 8px; - border: 1px solid rgba(245, 158, 11, 0.18); - border-radius: 7px; - background: rgba(255, 255, 255, 0.68); - display: flex; - flex-direction: column; - gap: 2px; - - .retry-summary-label { - color: #a16207; - font-size: 9.5px; - font-weight: 600; - letter-spacing: 0.02em; - text-transform: uppercase; - } - - strong { - color: #78350f; - font-size: 11.5px; - font-weight: 700; - overflow: hidden; - text-overflow: ellipsis; - white-space: nowrap; - } - } - - .retry-attempt-list { - border-top: 1px solid rgba(245, 158, 11, 0.22); - display: flex; - flex-direction: column; - gap: 6px; - margin-top: 2px; - padding-top: 8px; - animation: fadeInDetails 0.2s ease-out; - - .retry-attempt-item { - padding: 8px 9px; - border: 1px solid rgba(245, 158, 11, 0.2); - border-radius: 7px; - background: rgba(255, 255, 255, 0.76); - } - - .retry-attempt-heading { - display: flex; - align-items: center; - justify-content: space-between; - gap: 8px; - margin-bottom: 4px; - color: #92400e; - font-size: 10px; - font-weight: 700; - } - - .retry-wait-pill { - padding: 1px 5px; - border-radius: 999px; - background: #fef3c7; - color: #b45309; - } - - .retry-attempt-error { - color: #78350f; - font-size: 10.5px; - line-height: 1.45; - overflow-wrap: anywhere; - user-select: text; - } - } - .action-detail-row { display: flex; align-items: baseline; @@ -2654,6 +2593,38 @@ $uniform-gap: 8px; } } } + + .failure-retry-details { + margin-top: 4px; + padding-top: 9px; + border-top: 1px solid #fecaca; + display: flex; + flex-direction: column; + gap: 7px; + animation: fadeInDetails 0.2s ease-out; + + .failure-retry-title { + color: #991b1b; + font-size: 10px; + font-weight: 700; + text-transform: uppercase; + letter-spacing: 0.03em; + } + + .failure-retry-row { + display: flex; + justify-content: space-between; + gap: 12px; + color: #7f1d1d; + font-size: 10.5px; + line-height: 1.4; + + span:last-child { + color: #64748b; + text-align: right; + } + } + } } .action-card-expanded-details { diff --git a/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.ts b/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.ts index 78014fb..ae4be5d 100644 --- a/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.ts +++ b/apps/showcase_ui/src/app/components/agent-stream/agent-stream.component.ts @@ -158,6 +158,7 @@ export class AgentStreamComponent implements AfterViewInit { // Set to track expanded and collapsed state of action cards public expandedActionCards = signal>(new Set()); public collapsedActionCards = signal>(new Set()); + public retryClock = signal(Date.now()); // Performance caches for parameter and event extractions private actionParamsCache = new WeakMap(); @@ -280,6 +281,9 @@ export class AgentStreamComponent implements AfterViewInit { } constructor() { + const retryClockInterval = setInterval(() => this.retryClock.set(Date.now()), 1000); + this.destroyRef.onDestroy(() => clearInterval(retryClockInterval)); + // Auto scroll stream box during active text streaming effect(() => { const logs = this.filteredLogs(); @@ -827,7 +831,12 @@ export class AgentStreamComponent implements AfterViewInit { if (!block) return false; const hasNative = Boolean(this.getNativeThinking(block)); const hasRaw = Boolean(this.getRawThinking(block)); - const hasVisibleTools = Boolean(block.data?.generic_tools) && block.data.generic_tools.some((t: any) => this.shouldShowTool(t, block.data) || (t.type === 'llm_call' && (t.status === 'failed' || t.status === 'retrying')) || this.isReportStatusAction(t)); + const hasVisibleTools = Boolean(block.data?.generic_tools) && block.data.generic_tools.some((t: any) => + this.shouldShowTool(t, block.data) + || this.isDisplayableLLMFailure(t) + || this.isLLMRetry(t) + || this.isReportStatusAction(t) + ); const hasAndroidActions = Boolean(block.data?.action_taken) && (this.isAndroidAction(block.data.action_taken) || this.isReportStatusAction(block.data.action_taken)); return hasNative || hasRaw || hasVisibleTools || hasAndroidActions; } @@ -852,6 +861,12 @@ export class AgentStreamComponent implements AfterViewInit { return !!tool && this.retryablePauseTrace() === tool; } + public isDisplayableLLMFailure(tool: any): boolean { + return tool?.type === 'llm_call' + && tool?.status === 'failed' + && (tool?.name === 'llm_pause' || tool?.payload?.pause === true); + } + public resumePausedTask(event: Event): void { event.stopPropagation(); this.agentService.resumeTask(); @@ -865,7 +880,10 @@ export class AgentStreamComponent implements AfterViewInit { } public isLLMRetry(tool: any): boolean { - return tool?.type === 'llm_call' && tool?.status === 'retrying'; + return tool?.type === 'llm_call' + && tool?.status === 'retrying' + && tool?.payload?.source === 'provider_sdk' + && ['google', 'gemini'].includes(String(tool?.payload?.provider || '').toLowerCase()); } public getLLMRetryEntries(tool: any): any[] { @@ -873,25 +891,17 @@ export class AgentStreamComponent implements AfterViewInit { return Array.isArray(retries) && retries.length > 0 ? retries : [tool?.payload || tool]; } - public getLLMRetryCount(tool: any): number { - const count = Number(tool?.payload?.retry_count); - return Number.isFinite(count) && count > 0 ? count : this.getLLMRetryEntries(tool).length; + public getLLMFailureRetryEntries(tool: any): any[] { + const retries = tool?.payload?.retries; + if (!Array.isArray(retries)) return []; + return retries.filter((retry: any) => + retry?.source === 'provider_sdk' + && ['google', 'gemini'].includes(String(retry?.provider || '').toLowerCase()) + ); } - public getLLMRetryTotalDelay(tool: any): string | null { - const configuredTotal = Number(tool?.payload?.total_delay); - const total = Number.isFinite(configuredTotal) - ? configuredTotal - : this.getLLMRetryEntries(tool).reduce( - (sum: number, retry: any) => sum + (Number(retry?.delay) || 0), - 0 - ); - if (total <= 0) return null; - return `${total.toFixed(2).replace(/\.00$/, '')}s`; - } - - public getLLMRetryEntryError(entry: any): string { - return this.cleanErrorMessage(entry?.error || 'Unknown error'); + public hasLLMFailureRetryEntries(tool: any): boolean { + return this.getLLMFailureRetryEntries(tool).length > 0; } public getLLMRetryEntryDelay(entry: any): string | null { @@ -900,14 +910,37 @@ export class AgentStreamComponent implements AfterViewInit { return `${delay.toFixed(2).replace(/\.00$/, '')}s`; } - public getLLMRetryProvider(tool: any): string | null { - const providers = Array.isArray(tool?.payload?.providers) - ? tool.payload.providers - : [tool?.payload?.provider]; - const labels = providers - .filter((provider: any) => typeof provider === 'string' && provider.trim()) - .map((provider: string) => provider.trim().replace(/^./, (value: string) => value.toUpperCase())); - return labels.length > 0 ? Array.from(new Set(labels)).join(', ') : null; + public getLLMRetryEntryWaited(entry: any): string { + const delay = Number(entry?.delay); + if (!Number.isFinite(delay) || delay <= 0) return '0s'; + + const rawStartedAt = Number(entry?.scheduled_at ?? entry?.timestamp); + if (!Number.isFinite(rawStartedAt) || rawStartedAt <= 0) { + return this.formatLLMDuration(delay); + } + const startedAtMs = rawStartedAt < 1e11 ? rawStartedAt * 1000 : rawStartedAt; + const elapsed = Math.min(delay, Math.max(0, (this.retryClock() - startedAtMs) / 1000)); + return this.formatLLMDuration(elapsed); + } + + public isLLMRetryEntryWaiting(entry: any): boolean { + const delay = Number(entry?.delay); + const rawStartedAt = Number(entry?.scheduled_at ?? entry?.timestamp); + if (!Number.isFinite(delay) || delay <= 0 || !Number.isFinite(rawStartedAt)) return false; + const startedAtMs = rawStartedAt < 1e11 ? rawStartedAt * 1000 : rawStartedAt; + return this.retryClock() < startedAtMs + delay * 1000; + } + + public getLLMFailureWaited(tool: any): string | null { + const waited = Number(tool?.payload?.waited_seconds); + if (!Number.isFinite(waited) || waited < 0) return null; + return this.formatLLMDuration(waited); + } + + public formatLLMDuration(seconds: number): string { + if (seconds < 0.05) return '0s'; + if (seconds < 10) return `${seconds.toFixed(1).replace(/\.0$/, '')}s`; + return `${Math.round(seconds)}s`; } public cleanErrorMessage(rawError: any): string { diff --git a/apps/showcase_ui/src/app/services/agent.service.spec.ts b/apps/showcase_ui/src/app/services/agent.service.spec.ts new file mode 100644 index 0000000..342b65e --- /dev/null +++ b/apps/showcase_ui/src/app/services/agent.service.spec.ts @@ -0,0 +1,110 @@ +import { signal } from '@angular/core'; +import { of } from 'rxjs'; + +import { AgentService } from './agent.service'; + +describe('AgentService live LLM retry timeline', () => { + function createServiceWithoutPolling(): AgentService { + const service = Object.create(AgentService.prototype) as AgentService; + service.sessionLogs = signal([]); + return service; + } + + it('normalizes a live retry event into the historical trace contract', () => { + const service = createServiceWithoutPolling(); + const event = { + trace_id: 'retry-1', + step_id: 'step-1', + timestamp: 100, + error: '503 high demand', + delay: 1.5, + provider: 'google', + source: 'provider_sdk', + recoverable: true, + request_id: 'request-1', + scheduled_at: 100 + }; + + (service as any).appendLiveLLMRetryTrace(event, 'session-1'); + + const [log] = service.sessionLogs(); + expect(log.type).toBe('trace_recorded'); + expect(log.data).toEqual(jasmine.objectContaining({ + trace_id: 'retry-1', + session_id: 'session-1', + step_id: 'step-1', + type: 'llm_call', + name: 'llm_retry', + status: 'retrying' + })); + expect(log.data.payload).toEqual(jasmine.objectContaining({ + error: '503 high demand', + delay: 1.5, + provider: 'google', + source: 'provider_sdk', + request_id: 'request-1' + })); + }); + + it('upserts duplicate live retry events instead of adding another row', () => { + const service = createServiceWithoutPolling(); + const event = { + trace_id: 'retry-1', + timestamp: 100, + delay: 1.5, + provider: 'google', + source: 'provider_sdk', + request_id: 'request-1', + scheduled_at: 100 + }; + + (service as any).appendLiveLLMRetryTrace(event, 'session-1'); + (service as any).appendLiveLLMRetryTrace({ ...event, error: 'updated 503' }, 'session-1'); + + expect(service.sessionLogs().length).toBe(1); + expect(service.sessionLogs()[0].data.payload.error).toBe('updated 503'); + }); + + it('replaces an old history snapshot while preserving live retry events', () => { + const service = createServiceWithoutPolling(); + (service as any).currentSessionId = signal('session-1'); + (service as any).http = { + get: () => of([{ + step_id: 'step-1', + timestamp: 100, + generic_tools: [{ trace_id: 'persisted-retry', name: 'llm_retry' }] + }]) + }; + service.sessionLogs.set([ + { type: 'step_updated', history_snapshot: true, data: { step_id: 'old-step' } }, + { type: 'trace_recorded', data: { trace_id: 'live-retry', name: 'llm_retry' } } + ]); + + (service as any).backfillSessionSteps('session-1'); + + const logs = service.sessionLogs(); + expect(logs.length).toBe(2); + expect(logs[0].history_snapshot).toBeTrue(); + expect(logs[0].data.step_id).toBe('step-1'); + expect(logs[1].data.trace_id).toBe('live-retry'); + }); + + it('follows a just-started task even when status polling saw it first', () => { + const service = createServiceWithoutPolling(); + (service as any).http = { + post: () => of({ tasks: [{ session_id: 'new-session' }] }) + }; + service.agentStatus = signal('running'); + service.runningSessionId = signal('new-session'); + service.userPinnedSessionId = signal(null); + (service as any).sessions = signal([{ + session_id: 'new-session', + status: 'running' + }]); + const selectSpy = spyOn(service, 'selectSession'); + + service.runTask('test goal').subscribe(); + + expect(selectSpy).toHaveBeenCalledWith('new-session', false); + }); +}); diff --git a/apps/showcase_ui/src/app/services/agent.service.ts b/apps/showcase_ui/src/app/services/agent.service.ts index 8df14f0..129801a 100644 --- a/apps/showcase_ui/src/app/services/agent.service.ts +++ b/apps/showcase_ui/src/app/services/agent.service.ts @@ -184,7 +184,12 @@ export class AgentService { const isCurrentlyRunning = this.agentStatus() === 'running' || this.sessions().some((s) => s.status === 'running'); - if (!isCurrentlyRunning) { + const activeSessionId = this.runningSessionId() + || this.sessions().find((s) => s.status === 'running' || s.status === 'paused')?.session_id; + // The status poll can observe the new runner before /api/run + // returns. In that case it is still the task we just submitted, + // not an older task that should remain selected. + if (!isCurrentlyRunning || activeSessionId === newSessionId) { this.selectSession(newSessionId, false); } } @@ -401,45 +406,16 @@ export class AgentService { } console.log(`Connecting to session: ${sessionId}`); - this.http.get(`/api/sessions/${sessionId}/steps`).subscribe({ - next: (steps) => { - if (this.currentSessionId() !== sessionId) return; - const historicalLogs: any[] = []; - steps.forEach((step) => { - const stepTimeIso = new Date((step.timestamp || Date.now() / 1000) * 1000).toISOString(); - historicalLogs.push({ - type: 'step_updated', - session_id: sessionId, - timestamp: stepTimeIso, - data: step - }); - }); - // The history request and EventSource start together. On slower Windows - // hosts, live events can arrive first; replacing the signal here would - // erase them until a browser refresh reloads the persisted steps. - this.sessionLogs.update((liveLogs) => { - if (liveLogs.length === 0) return historicalLogs; - - const liveStepIds = new Set( - liveLogs - .filter((log) => log.type === 'step_recorded' || log.type === 'step_updated') - .map((log) => String(log.data?.step_id || log.data?.step_number || '')) - .filter(Boolean) - ); - const missingHistorical = historicalLogs.filter((log) => { - const stepKey = String(log.data?.step_id || log.data?.step_number || ''); - return !stepKey || !liveStepIds.has(stepKey); - }); - return [...missingHistorical, ...liveLogs]; - }); - }, - error: (err) => { - console.error('Failed to pre-load historical steps:', err); - } - }); - this.eventSource = new EventSource(`/api/stream/${sessionId}`); + // The server emits info only after this client has been registered as an + // SSE subscriber. Fetching the persisted snapshot after that acknowledgement + // closes the snapshot-to-stream gap: earlier events are in history and later + // events are already queued for this connection. + this.eventSource.addEventListener('info', () => { + this.backfillSessionSteps(sessionId); + }); + // Listen to standard SSE keep-alive this.eventSource.addEventListener('keep-alive', () => { // No-op keep-alive tick @@ -524,6 +500,7 @@ export class AgentService { const attemptText = attempt && max ? ` (Attempt ${attempt}/${max})` : ''; const delayText = delay > 0 ? `; retrying in ${delay.toFixed(2).replace(/\.00$/, '')}s` : ''; this.retryMessage.set(`AI service is temporarily busy${attemptText}${delayText}...`); + this.appendLiveLLMRetryTrace(parsedData, sessionId); return; } @@ -536,7 +513,8 @@ export class AgentService { pauseError, sessionId, parsedData.timestamp, - parsedData.step_id + parsedData.step_id, + parsedData ); return; } @@ -628,14 +606,28 @@ export class AgentService { ? new Date(parsedData.timestamp * 1000).toISOString() : new Date().toISOString(); - return [ - ...updatedLogs, - { - type: eventType, - timestamp: evtTime, - data: parsedData + const nextLog = { + type: eventType, + timestamp: evtTime, + data: parsedData + }; + + // A retry is published both as trace_recorded and as the + // purpose-built llm_retrying event. Replace the optimistic live + // copy when the persisted trace arrives instead of duplicating it. + if (eventType === 'trace_recorded' && parsedData?.trace_id) { + const existingTraceIndex = updatedLogs.findIndex((log) => + log.type === 'trace_recorded' + && log.data?.trace_id === parsedData.trace_id + ); + if (existingTraceIndex > -1) { + const deduplicatedLogs = [...updatedLogs]; + deduplicatedLogs[existingTraceIndex] = nextLog; + return deduplicatedLogs; } - ]; + } + + return [...updatedLogs, nextLog]; }); } @@ -668,6 +660,93 @@ export class AgentService { }; } + /** + * Load a fresh persisted snapshot without replacing events already received + * from SSE. Snapshot logs are replaceable, while live logs remain append-only; + * the stream aggregator merges matching steps and trace IDs within each step. + */ + private backfillSessionSteps(sessionId: string): void { + this.http.get(`/api/sessions/${sessionId}/steps`).subscribe({ + next: (steps) => { + if (this.currentSessionId() !== sessionId) return; + const historicalLogs = steps.map((step) => ({ + type: 'step_updated', + session_id: sessionId, + timestamp: new Date((step.timestamp || Date.now() / 1000) * 1000).toISOString(), + data: step, + history_snapshot: true + })); + + this.sessionLogs.update((logs) => [ + ...historicalLogs, + ...logs.filter((log) => !log.history_snapshot) + ]); + }, + error: (err) => { + console.error('Failed to backfill session steps:', err); + } + }); + } + + /** + * Convert the dedicated live retry event into the same trace shape used by + * history APIs. This keeps live streaming and post-refresh rendering on one + * data contract instead of maintaining a separate retry UI path. + */ + private appendLiveLLMRetryTrace(data: any, sessionId: string): void { + const timestampSeconds = typeof data?.timestamp === 'number' + ? (data.timestamp > 1e11 ? data.timestamp / 1000 : data.timestamp) + : Date.now() / 1000; + const traceId = data?.trace_id + || `llm-retry-live-${data?.request_id || 'unknown'}-${timestampSeconds}`; + const payload = { + error: data?.error, + delay: data?.delay, + attempt: data?.attempt, + max_retries: data?.max_retries, + provider: data?.provider, + source: data?.source, + recoverable: data?.recoverable, + request_id: data?.request_id, + scheduled_at: data?.scheduled_at ?? timestampSeconds + }; + const retryTrace = { + trace_id: traceId, + session_id: sessionId, + step_id: data?.step_id || null, + type: 'llm_call', + name: 'llm_retry', + timestamp: timestampSeconds, + status: 'retrying', + payload + }; + + this.sessionLogs.update((logs) => { + const existingIndex = logs.findIndex((log) => + log.type === 'trace_recorded' + && ( + log.data?.trace_id === traceId + || ( + data?.request_id + && log.data?.payload?.request_id === data.request_id + && Number(log.data?.payload?.scheduled_at) === Number(payload.scheduled_at) + ) + ) + ); + const retryLog = { + type: 'trace_recorded', + session_id: sessionId, + timestamp: new Date(timestampSeconds * 1000).toISOString(), + data: retryTrace + }; + if (existingIndex < 0) return [...logs, retryLog]; + + const updatedLogs = [...logs]; + updatedLogs[existingIndex] = retryLog; + return updatedLogs; + }); + } + /** * Preserve pause failures in the normal failed-LLM card stream. New runners * emit a persisted failed trace first; this fallback also supports older @@ -677,7 +756,8 @@ export class AgentService { error: unknown, sessionId: string, timestamp?: number | string, - stepId?: string | null + stepId?: string | null, + details?: any ): void { const errorText = typeof error === 'string' ? error : JSON.stringify(error); const pauseKey = `${sessionId}:${errorText}`; @@ -717,7 +797,14 @@ export class AgentService { name: 'llm_pause', status: 'failed', timestamp: timestampMs / 1000, - payload: { error: errorText, pause: true } + payload: { + error: errorText, + pause: true, + request_id: details?.request_id, + provider: details?.provider, + waited_seconds: details?.waited_seconds, + retries: Array.isArray(details?.retries) ? details.retries : [] + } } } ]; diff --git a/apps/showcase_ui/src/app/utils/tool-formatter.util.spec.ts b/apps/showcase_ui/src/app/utils/tool-formatter.util.spec.ts index a22ac4a..502b565 100644 --- a/apps/showcase_ui/src/app/utils/tool-formatter.util.spec.ts +++ b/apps/showcase_ui/src/app/utils/tool-formatter.util.spec.ts @@ -9,7 +9,7 @@ describe('getUniqueGenericTools retry aggregation', () => { name: 'llm_retry', status: 'retrying', timestamp: 100, - payload: { error: '503 first', delay: 1.18, provider: 'google' } + payload: { error: '503 first', delay: 1.18, provider: 'google', source: 'provider_sdk', request_id: 'request-1' } }, { trace_id: 'retry-2', @@ -17,7 +17,7 @@ describe('getUniqueGenericTools retry aggregation', () => { name: 'llm_retry', status: 'retrying', timestamp: 101, - payload: { error: '503 second', delay: 2.96, provider: 'google' } + payload: { error: '503 second', delay: 2.96, provider: 'google', source: 'provider_sdk', request_id: 'request-1' } } ]; @@ -27,31 +27,51 @@ describe('getUniqueGenericTools retry aggregation', () => { expect(result[0].name).toBe('llm_retry_group'); expect(result[0].payload.retry_count).toBe(2); expect(result[0].payload.total_delay).toBeCloseTo(4.14); - expect(result[0].payload.providers).toEqual(['google']); expect(result[0].payload.retries.map((retry: any) => retry.error)).toEqual([ '503 first', '503 second' ]); }); - it('keeps a terminal LLM failure separate from its preceding retries', () => { + it('folds SDK retry rows into their matching terminal failure', () => { const result = getUniqueGenericTools([ { trace_id: 'retry-1', type: 'llm_call', name: 'llm_retry', status: 'retrying', - payload: { error: '503', delay: 1 } + payload: { error: '503', delay: 1, provider: 'google', source: 'provider_sdk', request_id: 'request-1' } }, { trace_id: 'failure-1', type: 'llm_call', name: 'llm_pause', status: 'failed', - payload: { error: '503 exhausted', pause: true } + payload: { error: '503 exhausted', pause: true, request_id: 'request-1', retries: [{ delay: 1 }] } } ]); - expect(result.map(tool => tool.name)).toEqual(['llm_retry_group', 'llm_pause']); + expect(result.map(tool => tool.name)).toEqual(['llm_pause']); + }); + + it('does not expose opaque retries from non-Gemini providers', () => { + const result = getUniqueGenericTools([ + { + trace_id: 'retry-openai', + type: 'llm_call', + name: 'llm_retry', + status: 'retrying', + payload: { error: '429', delay: 2, provider: 'openai', source: 'artemis_wrapper' } + }, + { + trace_id: 'failure-openai', + type: 'llm_call', + name: 'llm_pause', + status: 'failed', + payload: { error: '429 exhausted', pause: true, waited_seconds: 12 } + } + ]); + + expect(result.map(tool => tool.name)).toEqual(['llm_pause']); }); }); diff --git a/apps/showcase_ui/src/app/utils/tool-formatter.util.ts b/apps/showcase_ui/src/app/utils/tool-formatter.util.ts index 73b9845..d3a4dc9 100644 --- a/apps/showcase_ui/src/app/utils/tool-formatter.util.ts +++ b/apps/showcase_ui/src/app/utils/tool-formatter.util.ts @@ -185,41 +185,57 @@ export function getToolAgentName(tool: any): string | null { export function getUniqueGenericTools(tools: any[] | undefined): any[] { if (!tools) return []; + // Only Google GenAI SDK retries are observable. Artemis wrapper retries and + // opaque provider retries must never be presented as provider internals. const retryTools = tools.filter(tool => tool?.type === 'llm_call' && tool?.status === 'retrying' && (tool?.name === 'llm_retry' || tool?.name === 'llm_retry_group') + && tool?.payload?.source === 'provider_sdk' + && ['google', 'gemini'].includes(String(tool?.payload?.provider || '').toLowerCase()) ); - const firstRetry = retryTools[0]; - const latestRetry = retryTools[retryTools.length - 1]; - const retryAggregate = firstRetry ? { - ...firstRetry, - trace_id: `llm-retry-group-${firstRetry.trace_id || firstRetry.timestamp || 'unknown'}`, - name: 'llm_retry_group', - payload: { - ...(latestRetry?.payload || {}), - retry_count: retryTools.length, - total_delay: retryTools.reduce( - (total, retry) => total + (Number(retry?.payload?.delay) || 0), - 0 - ), - providers: Array.from(new Set( - retryTools - .map(retry => retry?.payload?.provider) - .filter((provider): provider is string => typeof provider === 'string' && !!provider) - )), - retries: retryTools.map(retry => ({ - trace_id: retry.trace_id, - timestamp: retry.timestamp, - error: retry?.payload?.error || retry.error, - delay: Number(retry?.payload?.delay) || 0, - provider: retry?.payload?.provider, - source: retry?.payload?.source, - attempt: retry?.payload?.attempt, - max_retries: retry?.payload?.max_retries - })) - } - } : null; + const retryKey = (tool: any) => String( + tool?.payload?.request_id || `legacy-${tool?.step_id || 'unassigned'}` + ); + const retryGroups = new Map(); + for (const retry of retryTools) { + const key = retryKey(retry); + retryGroups.set(key, [...(retryGroups.get(key) || []), retry]); + } + const terminalRequestIds = new Set( + tools + .filter(tool => tool?.type === 'llm_call' && tool?.status === 'failed') + .map(tool => tool?.payload?.request_id) + .filter((requestId): requestId is string => typeof requestId === 'string' && !!requestId) + ); + + const buildRetryAggregate = (group: any[]) => { + const firstRetry = group[0]; + const latestRetry = group[group.length - 1]; + return { + ...firstRetry, + trace_id: `llm-retry-group-${firstRetry?.payload?.request_id || firstRetry?.trace_id || firstRetry?.timestamp || 'unknown'}`, + name: 'llm_retry_group', + payload: { + ...(latestRetry?.payload || {}), + retry_count: group.length, + total_delay: group.reduce( + (total, retry) => total + (Number(retry?.payload?.delay) || 0), + 0 + ), + retries: group.map(retry => ({ + trace_id: retry.trace_id, + timestamp: retry.timestamp, + error: retry?.payload?.error || retry.error, + delay: Number(retry?.payload?.delay) || 0, + provider: retry?.payload?.provider, + source: retry?.payload?.source, + request_id: retry?.payload?.request_id, + scheduled_at: retry?.payload?.scheduled_at + })) + } + }; + }; const toolMap = new Map(); for (const t of tools) { @@ -232,22 +248,33 @@ export function getUniqueGenericTools(tools: any[] | undefined): any[] { const result: any[] = []; const seenTraces = new Set(); - let retryAggregateAdded = false; + const addedRetryGroups = new Set(); for (const tool of tools) { if (!tool || !tool.name) continue; if (tool.type === 'agent') continue; if (tool.type === 'llm_call' && tool.status === 'retrying') { - if (!retryAggregateAdded && retryAggregate) { - result.push(retryAggregate); - retryAggregateAdded = true; + const groupKey = retryKey(tool); + const group = retryGroups.get(groupKey); + const requestId = tool?.payload?.request_id; + if ( + group + && !addedRetryGroups.has(groupKey) + && !(requestId && terminalRequestIds.has(requestId)) + ) { + result.push(buildRetryAggregate(group)); + addedRetryGroups.add(groupKey); } continue; } if (tool.type === 'llm_call' && tool.status === 'failed') { - result.push(tool); + // Callback-level attempt errors are implementation details. The wrapper + // emits one terminal llm_pause trace after retries are exhausted. + if (tool.name === 'llm_pause' || tool?.payload?.pause === true) { + result.push(tool); + } continue; } diff --git a/artemis/agents/operator/operator.py b/artemis/agents/operator/operator.py index 9f3aaf4..ee3f3d9 100644 --- a/artemis/agents/operator/operator.py +++ b/artemis/agents/operator/operator.py @@ -93,6 +93,7 @@ from artemis.agents.operator.prompts import ( ToolLimitWarningPromptComponent, InjectedInstructionPromptComponent, ) +from artemis.agents.operator.prompt_builder import load_operator_prompts class OperatorNode: @@ -107,13 +108,11 @@ class OperatorNode: self.tools = tools or [] self.prompt_components = prompt_components or [] self.last_n_detailed = last_n_detailed - self.prompts = {} - prompts_path = Path(__file__).parent / "operator.json" - if prompts_path.exists(): - try: - self.prompts = json.loads(prompts_path.read_text(encoding="utf-8")) - except Exception as e: - logger.error(f"Failed to load operator prompts: {e}") + try: + self.prompts = load_operator_prompts() + except (OSError, ValueError, json.JSONDecodeError) as e: + logger.error(f"Failed to load operator prompts: {e}") + self.prompts = {} def _check_infinite_loop(self, state: State): subagent_calls = state.subagent_calls or [] @@ -594,8 +593,9 @@ class OperatorNode: " screen action you attempted to" " execute was rejected because you are" " not allowed to simultaneously" - " read/write memory and execute" - " physical actions in the same turn." + " use result-dependent pre-decision" + " tools and execute physical actions in" + " the same turn." " Currently, no screen actions have" " been executed. Please review your" " updated context and re-output your" @@ -614,9 +614,9 @@ class OperatorNode: " successfully processed. However, your" " response to the checker was rejected" " because you are not allowed to" - " simultaneously read/write memory and" - " reply to the checker in the same" - " turn. Currently, no response has been" + " simultaneously use result-dependent" + " pre-decision tools and reply to the" + " checker. Currently, no response has been" " sent to the checker. Please review" " your updated context and re-output" " your reply." diff --git a/artemis/agents/operator/prompt_builder.py b/artemis/agents/operator/prompt_builder.py index 8114958..3b719eb 100644 --- a/artemis/agents/operator/prompt_builder.py +++ b/artemis/agents/operator/prompt_builder.py @@ -12,21 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Dynamic prompt template builder for Operator Agent.""" +"""Compatibility prompt builder backed by the canonical Operator templates.""" + +import json +from functools import lru_cache +from pathlib import Path from jinja2 import Template -OPERATOR_SYSTEM_PROMPT = """You are ARTEMIS Operator, an expert autonomous mobile device interaction agent. -Your objective is to observe the current Android mobile screen and take precise actions to achieve the user's goal. - -Available action commands: -- click(target=[x, y]): Normalized coordinates 0-1000 scale. -- swipe(action="up"|"down"|"left"|"right"|[x1, y1, x2, y2]): Drag across screen. -- input_text(text="...", target=[x, y]): Click field and type. -- press_key(key="home"|"back"|"enter"|"delete"): Press hardware key. -- wait_for_delay(seconds=...): Pause execution. -""" - OPERATOR_HUMAN_TEMPLATE = """Goal: {{ goal }} Active Sub-Goal: {{ sub_goal }} Current Turn: {{ current_turn }} @@ -35,12 +28,25 @@ Execution History: """ +@lru_cache(maxsize=1) +def load_operator_prompts() -> dict[str, str]: + """Load the single canonical prompt source shared by all Operator entry points.""" + prompts_path = Path(__file__).with_name("operator.json") + return json.loads(prompts_path.read_text(encoding="utf-8")) + + class OperatorPromptBuilder: """Builds prompt messages and multimodal contents for the Operator Agent.""" @classmethod - def build_system_message(cls) -> str: - return OPERATOR_SYSTEM_PROMPT + def build_system_message(cls, template_name: str = "main_template") -> str: + from artemis.agents.operator.prompts import apply_operator_prompt_contract + + prompts = load_operator_prompts() + try: + return apply_operator_prompt_contract(prompts[template_name]) + except KeyError as exc: + raise ValueError(f"Unknown Operator template: {template_name}") from exc @classmethod def build_human_message( diff --git a/artemis/agents/operator/prompts.py b/artemis/agents/operator/prompts.py index c363db3..4b328c4 100644 --- a/artemis/agents/operator/prompts.py +++ b/artemis/agents/operator/prompts.py @@ -27,14 +27,59 @@ from artemis.tools.command_tool import ( _format_long_output_response, _is_output_long, ) -from artemis.agents.explorer.constants import EXPLORE_DESCRIPTIONS -from artemis.config import resolve_explorer_version from artemis.utils.logger import get_logger from artemis.utils.task_tree import get_active_subgoal_hashes logger = get_logger(__name__) +_LEGACY_TOOL_CALLING_RULE = """- **Tool Calling Iron Rule**: Do NOT submit a Turn-Ending Action at the same time (in the same tool-call list) as a Pre-Decision Exploratory Tool. You must first gather information (e.g., finding coordinates or diagnosing a state), manage ADB background tasks, or update memory notes. You can use pre-decision exploratory tools to complete this task. You will be prompted again with the tool's result, at which point you can output your final physical/turn-ending action.""" + +_TOOL_CALLING_RULE = """- **Tool Calling Contract**: Tools that return information needed for the decision (`ask_explorer`, `ask_diagnoser`, `video_analyzer`, `run_adb_command`, `manage_task`, `analyze_task_output`, `read_note`, `list_notes`, and `save_note`) are result-dependent Pre-Decision Tools. Do not submit them in the same tool-call list as a Turn-Ending Action; inspect their result first. The write-through tools `update_note` and `append_note` may accompany at most one Turn-Ending Action when their content is based entirely on evidence already available before that action. Never write an action's expected outcome as if it had already been observed.""" + +_LEGACY_TOOL_PROTOCOL = """1. **Pre-Decision Exploratory Tools**: + - *What they are*: Helper/Subagent tools (such as `ask_explorer`, `ask_diagnoser`, `video_analyzer`), ADB command tools (`run_adb_command`, `manage_task`), and memory note tools (`read_note`, `list_notes`, `save_note`, `update_note`, `append_note`). + - *How they work*: These tools gather details or update/read memory. When you invoke these tools, the framework will immediately execute them and return the results, allowing you to continue thinking and make your final decision. +2. **Turn-Ending Actions**:""" + +_TOOL_PROTOCOL = """1. **Result-Dependent Pre-Decision Tools**: + - *What they are*: Helper/Subagent tools (`ask_explorer`, `ask_diagnoser`, `video_analyzer`), ADB/task tools (`run_adb_command`, `manage_task`, `analyze_task_output`), and memory tools whose result must be inspected (`read_note`, `list_notes`, `save_note`). + - *How they work*: Invoke these without a Turn-Ending Action, inspect the returned result, and then decide. +2. **Write-Through Memory Tools**: + - `update_note` and `append_note` may run alongside at most one Turn-Ending Action only when recording facts already observed in the current context. If the note content depends on the action's result, wait for the next observation before writing it. +3. **Turn-Ending Actions**:""" + +_LEGACY_CHECKER_REJECTION_TRIGGER = """ - **Validation/Checker Rejection**: The verification agent (Checker) rejected your subgoal completion (i.e., you are in troubleshooter mode and this is a retry).""" + +_AMBIGUOUS_CHECKER_REJECTION_TRIGGER = """ - **Ambiguous Validation/Checker Rejection**: Use `ask_diagnoser` only when the rejection cause is unclear, the evidence conflicts, or no safe local correction is evident. A Checker rejection alone is not a mandatory diagnosis trigger.""" + +_LEGACY_TROUBLESHOOTER_CHALLENGE = """ - *Challenge Checker*: If you visually confirm the subgoal's target state is *already achieved*, you **MUST** invoke the `reply_to_checker` tool to state your observation.""" + +_TROUBLESHOOTER_CHALLENGE = """ - *Challenge Checker*: If current evidence clearly proves the subgoal is already achieved, invoke `reply_to_checker` directly. This branch takes precedence over diagnosis; do not call `ask_diagnoser` first. + - *Clear Local Correction*: If the rejection cause is explicit and a safe correction is evident from the current screen, perform that correction directly. This branch also takes precedence over generic diagnosis triggers. Use `ask_diagnoser` only for ambiguity, conflicting evidence, or repeated failure.""" + + +def apply_operator_prompt_contract(prompt_template: str) -> str: + """Align prompt-level tool and recovery rules with Operator runtime semantics.""" + prompt_template = prompt_template.replace( + _LEGACY_TOOL_CALLING_RULE, + _TOOL_CALLING_RULE, + ) + prompt_template = prompt_template.replace( + _LEGACY_TOOL_PROTOCOL, + _TOOL_PROTOCOL, + ) + prompt_template = prompt_template.replace( + _LEGACY_CHECKER_REJECTION_TRIGGER, + _AMBIGUOUS_CHECKER_REJECTION_TRIGGER, + ) + prompt_template = prompt_template.replace( + _LEGACY_TROUBLESHOOTER_CHALLENGE, + _TROUBLESHOOTER_CHALLENGE, + ) + return prompt_template + + class PromptBuilder: def __init__(self): self.system_parts = [] @@ -83,17 +128,7 @@ class TemplatePromptComponent(PromptComponent): "Failed to format prompt, template not found in operator prompts config." ) - op_version = resolve_explorer_version(ctx, agent_or_profile_name="operator") - rule_prompt_info = EXPLORE_DESCRIPTIONS.get(op_version, EXPLORE_DESCRIPTIONS["pro"]) - rule_prompt = rule_prompt_info["rule_prompt"] - if "{max_iterations}" in rule_prompt: - max_iterations = 8 if op_version == "ultra" else 3 - rule_prompt = rule_prompt.format(max_iterations=max_iterations) - - prompt_template = prompt_template.replace( - "# OPERATIONAL PRINCIPLES\n", - f"# OPERATIONAL PRINCIPLES\n- **Visual Explorer Rule**: {rule_prompt}\n", - ) + prompt_template = apply_operator_prompt_contract(prompt_template) plan_and_history = kwargs.get("plan_and_history", "No plan or history yet.") diff --git a/artemis/context.py b/artemis/context.py index 0bffcff..229d10e 100644 --- a/artemis/context.py +++ b/artemis/context.py @@ -163,6 +163,7 @@ class ArtemisContext(BaseModel): checker_task: asyncio.Task[Any] | None = None planner_task: asyncio.Task[Any] | None = None background_tasks: list[asyncio.Task[Any]] = Field(default_factory=list) + background_task_grace_period_seconds: float = 2.0 package_cache: dict[str, str | None] = Field(default_factory=dict) background_jobs: dict[str, dict[str, Any]] = Field(default_factory=dict) @@ -178,9 +179,34 @@ class ArtemisContext(BaseModel): async def __aexit__(self, exc_type, exc_val, exc_tb): if self.background_tasks: - logger.info(f"Waiting for {len(self.background_tasks)} background tasks to complete...") - results = await asyncio.gather(*self.background_tasks, return_exceptions=True) - for task, result in zip(self.background_tasks, results): + tasks = list(self.background_tasks) + pending_tasks = [task for task in tasks if not task.done()] + grace_period = ( + max(0.0, self.background_task_grace_period_seconds) + if exc_type is None + else 0.0 + ) + + if pending_tasks and grace_period > 0: + logger.info( + f"Draining {len(pending_tasks)} background tasks for up to " + f"{grace_period:.1f}s..." + ) + _, pending_tasks = await asyncio.wait( + pending_tasks, + timeout=grace_period, + ) + + if pending_tasks: + logger.info( + f"Cancelling {len(pending_tasks)} unfinished background tasks; " + "primary automation is already complete." + ) + for task in pending_tasks: + task.cancel() + + results = await asyncio.gather(*tasks, return_exceptions=True) + for task, result in zip(tasks, results): if isinstance(result, Exception): logger.error( f"Background task {task.get_name()} failed: {result}", diff --git a/artemis/services/llm.py b/artemis/services/llm.py index c9b00ff..0be9c1e 100644 --- a/artemis/services/llm.py +++ b/artemis/services/llm.py @@ -20,6 +20,7 @@ thought stream recording, and role-based dynamic dispatching. import asyncio from collections.abc import Awaitable, Callable, Coroutine +from contextvars import ContextVar, Token from dataclasses import replace import functools import logging @@ -59,6 +60,33 @@ T = TypeVar("T") # global after this module has already been imported. _CURRENT_DATA_ENGINE = None +# Provider retry logs do not carry an Artemis trace id. Keep the active model +# request in async-local state so observable SDK retries and the terminal error +# are emitted as one lifecycle, even when several model calls run concurrently. +_ACTIVE_LLM_REQUEST: ContextVar[dict[str, Any] | None] = ContextVar( + "active_llm_request", default=None +) + + +def _provider_name_from_call(args: tuple[Any, ...]) -> str | None: + wrapper = args[0] if args else None + endpoint = getattr(wrapper, "endpoint", None) + provider = getattr(endpoint, "provider", None) + if provider is None: + return None + return str(getattr(provider, "value", provider)) + + +def _begin_llm_request(args: tuple[Any, ...]) -> Token: + return _ACTIVE_LLM_REQUEST.set( + { + "request_id": str(uuid4()), + "provider": _provider_name_from_call(args), + "started_at": time.time(), + "retries": [], + } + ) + def _get_current_data_engine(): if _CURRENT_DATA_ENGINE is not None: @@ -77,6 +105,11 @@ def _record_llm_retry( source: str = "artemis", ) -> None: """Persist and publish one recoverable LLM retry for UI transparency.""" + request = _ACTIVE_LLM_REQUEST.get() + timestamp = time.time() + request_id = request.get("request_id") if request else None + provider = provider or (request.get("provider") if request else None) + engine = _get_current_data_engine() if not engine or not getattr(engine, "current_session_id", None): return @@ -93,12 +126,12 @@ def _record_llm_retry( "provider": provider, "source": source, "recoverable": True, + "request_id": request_id, + "scheduled_at": timestamp, } payload = {key: value for key, value in payload.items() if value is not None} step_id = getattr(engine, "current_step_id", None) parent_trace_id = CURRENT_TRACE_ID.get() - timestamp = time.time() - try: trace_id = engine.record_trace( type="llm_call", @@ -112,6 +145,15 @@ def _record_llm_retry( llm_logger.warning("Failed to persist LLM retry trace: %s", trace_error) trace_id = None + if request is not None: + request["retries"].append( + { + **payload, + "trace_id": str(trace_id) if trace_id else None, + "timestamp": timestamp, + } + ) + try: engine._publish( "llm_retrying", @@ -145,6 +187,10 @@ class _ProviderRetryTelemetryHandler(logging.Handler): def emit(self, record: logging.LogRecord) -> None: try: + request = _ACTIVE_LLM_REQUEST.get() + if not request or request.get("provider") != ModelProvider.GOOGLE.value: + return + message = record.getMessage() lowered = message.lower() if "retrying" not in lowered or not any( @@ -159,7 +205,7 @@ class _ProviderRetryTelemetryHandler(logging.Handler): _record_llm_retry( error, delay, - provider="google", + provider=request["provider"], source="provider_sdk", ) except Exception: @@ -197,13 +243,29 @@ def _handle_llm_pause_and_resume(last_error: Exception) -> Path: if current_engine: step_id = getattr(current_engine, "current_step_id", None) timestamp = time.time() + request = _ACTIVE_LLM_REQUEST.get() + failure_payload: dict[str, Any] = {"error": err_msg, "pause": True} + if request: + failure_payload.update( + { + "request_id": request["request_id"], + "provider": request.get("provider"), + "waited_seconds": max(0.0, timestamp - request["started_at"]), + # Only SDK-observed retries are included. Other providers + # may retry internally, but their attempt data is opaque. + "retries": list(request["retries"]), + } + ) + failure_payload = { + key: value for key, value in failure_payload.items() if value is not None + } try: # Persist exhausted LLM retries as a failed call so the same error # card is available both in the live stream and after a refresh. current_engine.record_trace( type="llm_call", name="llm_pause", - payload={"error": err_msg, "pause": True}, + payload=failure_payload, step_id=step_id, status="failed", ) @@ -213,7 +275,7 @@ def _handle_llm_pause_and_resume(last_error: Exception) -> Path: current_engine._publish( "task_paused", { - "error": err_msg, + **failure_payload, "step_id": str(step_id) if step_id else None, "timestamp": timestamp, }, @@ -228,12 +290,15 @@ def robust_retry_async(func: Callable) -> Callable: @functools.wraps(func) async def wrapper(*args, **kwargs): while True: + request_token = _begin_llm_request(args) max_retries = 3 retry_delay = 0.0 last_error = None for attempt in range(max_retries): try: - return await func(*args, **kwargs) + result = await func(*args, **kwargs) + _ACTIVE_LLM_REQUEST.reset(request_token) + return result except Exception as e: last_error = e err_str = str(e).lower() @@ -256,16 +321,15 @@ def robust_retry_async(func: Callable) -> Callable: if retry_delay > 0: llm_logger.warning(f"LLM call failed, retrying in {retry_delay:.2f}s...") - _record_llm_retry( - str(e), - retry_delay, - attempt=attempt + 1, - max_retries=max_retries, - source="artemis_wrapper", - ) await asyncio.sleep(retry_delay) + except BaseException: + _ACTIVE_LLM_REQUEST.reset(request_token) + raise - pause_file = _handle_llm_pause_and_resume(last_error) + try: + pause_file = _handle_llm_pause_and_resume(last_error) + finally: + _ACTIVE_LLM_REQUEST.reset(request_token) while pause_file.exists(): await asyncio.sleep(1) @@ -284,6 +348,7 @@ def robust_retry_astream(func: Callable) -> Callable: @functools.wraps(func) async def wrapper(*args, **kwargs): while True: + request_token = _begin_llm_request(args) max_retries = 3 retry_delay = 0.0 last_error = None @@ -294,6 +359,7 @@ def robust_retry_astream(func: Callable) -> Callable: yield first async for chunk in gen: yield chunk + _ACTIVE_LLM_REQUEST.reset(request_token) return except Exception as e: last_error = e @@ -319,16 +385,15 @@ def robust_retry_astream(func: Callable) -> Callable: f"LLM stream handshake throttled, retrying in {retry_delay:.2f}s..." ) - _record_llm_retry( - str(e), - retry_delay, - attempt=attempt + 1, - max_retries=max_retries, - source="artemis_wrapper", - ) await asyncio.sleep(retry_delay) + except BaseException: + _ACTIVE_LLM_REQUEST.reset(request_token) + raise - pause_file = _handle_llm_pause_and_resume(last_error) + try: + pause_file = _handle_llm_pause_and_resume(last_error) + finally: + _ACTIVE_LLM_REQUEST.reset(request_token) while pause_file.exists(): await asyncio.sleep(1) diff --git a/tests/unit/admin_console/test_step_repository.py b/tests/unit/admin_console/test_step_repository.py index 6fb872b..da866a7 100644 --- a/tests/unit/admin_console/test_step_repository.py +++ b/tests/unit/admin_console/test_step_repository.py @@ -73,3 +73,42 @@ def test_retrying_llm_payload_keeps_user_visible_retry_details(): "source": "provider_sdk", "recoverable": True, } + + +def test_terminal_llm_payload_keeps_wait_and_retry_aggregation(): + repository = StepRepository() + retries = [ + { + "delay": 5.0, + "scheduled_at": 100.0, + "provider": "google", + "source": "provider_sdk", + } + ] + trace = { + "type": "llm_call", + "name": "llm_pause", + "status": "failed", + "payload": json.dumps( + { + "error": "503 UNAVAILABLE", + "pause": True, + "request_id": "request-1", + "provider": "google", + "waited_seconds": 18.5, + "retries": retries, + "messages": ["private prompt"], + } + ), + } + + normalized = repository._normalize_display_trace(trace) + + assert normalized["payload"] == { + "error": "503 UNAVAILABLE", + "provider": "google", + "pause": True, + "request_id": "request-1", + "waited_seconds": 18.5, + "retries": retries, + } diff --git a/tests/unit/agents/test_operator_prompts.py b/tests/unit/agents/test_operator_prompts.py new file mode 100644 index 0000000..5e4acfe --- /dev/null +++ b/tests/unit/agents/test_operator_prompts.py @@ -0,0 +1,55 @@ +# Copyright 2026 Google LLC +# +# 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 +# +# https://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. + +from artemis.agents.operator.prompt_builder import ( + OperatorPromptBuilder, + load_operator_prompts, +) +from artemis.agents.operator.prompts import apply_operator_prompt_contract + + +def test_operator_prompt_contract_matches_note_runtime_semantics(): + prompt = apply_operator_prompt_contract(load_operator_prompts()["main_template"]) + + assert "write-through tools `update_note` and `append_note`" in prompt + assert "may accompany at most one Turn-Ending Action" in prompt + assert "`read_note`, `list_notes`, and `save_note`" in prompt + assert "Do NOT submit a Turn-Ending Action at the same time" not in prompt + assert "memory note tools (`read_note`, `list_notes`, `save_note`, `update_note`" not in prompt + + +def test_troubleshooter_prompt_has_mutually_exclusive_recovery_branches(): + prompt = apply_operator_prompt_contract(load_operator_prompts()["troubleshooter_template"]) + + assert "invoke `reply_to_checker` directly" in prompt + assert "This branch takes precedence over diagnosis" in prompt + assert "perform that correction directly" in prompt + assert "also takes precedence over generic diagnosis triggers" in prompt + assert "A Checker rejection alone is not a mandatory diagnosis trigger" in prompt + assert "**Validation/Checker Rejection**" not in prompt + + +def test_operator_prompt_omits_environment_trust_and_explorer_directives(): + for template in load_operator_prompts().values(): + prompt = apply_operator_prompt_contract(template) + assert "Untrusted Screen Content & Instruction Priority" not in prompt + assert "Visual Explorer Rule" not in prompt + + +def test_compatibility_builder_uses_canonical_operator_template(): + expected = apply_operator_prompt_contract(load_operator_prompts()["main_template"]) + + assert OperatorPromptBuilder.build_system_message() == expected + assert "wait_for_delay(seconds=" not in expected + assert 'press_key(key="home"' not in expected diff --git a/tests/unit/test_context.py b/tests/unit/test_context.py index 0669632..4149de2 100644 --- a/tests/unit/test_context.py +++ b/tests/unit/test_context.py @@ -42,3 +42,36 @@ async def test_artemis_context_waits_for_background_tasks(): # After exiting the context, the task should have run to completion assert task_run is True + + +@pytest.mark.asyncio +async def test_artemis_context_cancels_background_tasks_after_grace_period(): + device = DeviceContext( + host_platform="LINUX", + mobile_platform=DevicePlatform.ANDROID, + device_id="dummy", + device_width=1080, + device_height=2400, + ) + context = ArtemisContext( + device=device, + background_task_grace_period_seconds=0.01, + ) + + task_cancelled = False + + async def stuck_summary(): + nonlocal task_cancelled + try: + await asyncio.Event().wait() + finally: + task_cancelled = True + + async with context: + task = asyncio.create_task(stuck_summary(), name="stuck_summary") + context.background_tasks.append(task) + + assert task.done() + assert task.cancelled() + assert task_cancelled is True + assert context.background_tasks == [] diff --git a/tests/unit/test_llm_pause.py b/tests/unit/test_llm_pause.py index e995181..bc4f288 100644 --- a/tests/unit/test_llm_pause.py +++ b/tests/unit/test_llm_pause.py @@ -57,24 +57,98 @@ def test_provider_sdk_retry_is_persisted_and_published(monkeypatch): exc_info=None, ) - llm._ProviderRetryTelemetryHandler().emit(record) - - assert engine.trace_kwargs == { - "type": "llm_call", - "name": "llm_retry", - "payload": { - "error": "ServerError: 503 UNAVAILABLE: high demand", - "delay": 1.18, + request_token = llm._ACTIVE_LLM_REQUEST.set( + { + "request_id": "request-1", "provider": "google", - "source": "provider_sdk", - "recoverable": True, - }, - "step_id": "step-7", - "parent_trace_id": None, - "status": "retrying", + "started_at": 100.0, + "retries": [], + } + ) + try: + llm._ProviderRetryTelemetryHandler().emit(record) + finally: + llm._ACTIVE_LLM_REQUEST.reset(request_token) + + assert engine.trace_kwargs["type"] == "llm_call" + assert engine.trace_kwargs["name"] == "llm_retry" + assert engine.trace_kwargs["status"] == "retrying" + assert engine.trace_kwargs["payload"] == { + "error": "ServerError: 503 UNAVAILABLE: high demand", + "delay": 1.18, + "provider": "google", + "source": "provider_sdk", + "recoverable": True, + "request_id": "request-1", + "scheduled_at": engine.trace_kwargs["payload"]["scheduled_at"], } event_type, payload = engine.published assert event_type == "llm_retrying" assert payload["trace_id"] == "trace-1" assert payload["delay"] == 1.18 assert payload["error"].startswith("ServerError: 503") + + +def test_provider_sdk_retry_is_hidden_for_non_gemini_request(monkeypatch): + engine = _RecordingEngine() + monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine) + record = logging.LogRecord( + name="google_genai._api_client", + level=logging.INFO, + pathname=__file__, + lineno=1, + msg="Retrying request in 2 seconds as it raised 503 UNAVAILABLE", + args=(), + exc_info=None, + ) + request_token = llm._ACTIVE_LLM_REQUEST.set( + { + "request_id": "request-openai", + "provider": "openai", + "started_at": 100.0, + "retries": [], + } + ) + try: + llm._ProviderRetryTelemetryHandler().emit(record) + finally: + llm._ACTIVE_LLM_REQUEST.reset(request_token) + + assert engine.trace_kwargs is None + assert engine.published is None + + +def test_terminal_failure_aggregates_observed_retries_and_total_wait(monkeypatch, tmp_path): + engine = _RecordingEngine() + monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine) + monkeypatch.setattr(llm.settings, "TRACES_PATH", str(tmp_path / "traces")) + monkeypatch.setattr(llm.time, "time", lambda: 120.0) + retries = [ + { + "error": "503 high demand", + "delay": 10.0, + "provider": "google", + "source": "provider_sdk", + "request_id": "request-2", + "scheduled_at": 101.0, + "timestamp": 101.0, + } + ] + request_token = llm._ACTIVE_LLM_REQUEST.set( + { + "request_id": "request-2", + "provider": "google", + "started_at": 100.0, + "retries": retries, + } + ) + try: + llm._handle_llm_pause_and_resume(RuntimeError("503 high demand")) + finally: + llm._ACTIVE_LLM_REQUEST.reset(request_token) + + payload = engine.trace_kwargs["payload"] + assert payload["request_id"] == "request-2" + assert payload["provider"] == "google" + assert payload["waited_seconds"] == 20.0 + assert payload["retries"] == retries