From 4f13bc6a44091cd7fdb8bf3494f83a508c6f4490 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Sat, 29 Aug 2026 20:12:32 +0200 Subject: [PATCH 1/3] fix(cloud-agent): improve failure reporting and log retention --- .../CloudAgentNextOutcomesPage.tsx | 168 ++++- .../health-summary.test.ts | 12 + .../CloudAgentNextTelemetry/health-summary.ts | 6 + .../terminal-reason-from-failure.test.ts | 34 +- .../terminal-reason-from-failure.ts | 4 + .../admin-cloud-agent-next-router.test.ts | 578 +++++++++++++++++- .../routers/admin-cloud-agent-next-router.ts | 164 +++-- packages/db/src/schema.test.ts | 66 ++ packages/db/src/schema.ts | 6 + .../src/cloud-agent-failure.test.ts | 120 +++- .../worker-utils/src/cloud-agent-failure.ts | 29 +- .../src/cloud-agent-queue-report.test.ts | 58 ++ .../src/cloud-agent-queue-report.ts | 4 + .../src/persistence/CloudAgentSession.ts | 43 ++ .../src/persistence/session-metadata.test.ts | 49 +- .../src/persistence/session-metadata.ts | 1 + .../services/git-token-service-client.test.ts | 151 +++-- .../src/services/git-token-service-client.ts | 54 +- .../src/session-service.test.ts | 305 ++++++++- .../cloud-agent-next/src/session-service.ts | 36 +- .../session/safe-failure-projection.test.ts | 369 +++++++++++ .../src/session/safe-failure-projection.ts | 110 +--- .../src/session/session-message-queue.test.ts | 109 +++- .../src/session/session-message-queue.ts | 14 +- .../src/session/session-message-state.test.ts | 167 +++++ .../src/session/session-message-state.ts | 8 +- .../src/session/session-prepare.test.ts | 199 ++++-- .../src/session/session-registration.ts | 27 +- .../src/session/wrapper-supervisor.test.ts | 192 +++++- .../src/session/wrapper-supervisor.ts | 16 +- .../src/shared/assistant-failure.ts | 175 ++++++ .../src/shared/ingest-frame.test.ts | 182 ++++++ .../src/shared/ingest-frame.ts | 3 +- .../src/telemetry/queue-reports.test.ts | 198 ++++++ .../src/telemetry/queue-reports.ts | 11 +- .../src/telemetry/report-consumer.test.ts | 196 +++++- .../src/telemetry/report-consumer.ts | 35 +- .../src/telemetry/report-store.test.ts | 505 ++++++++++++++- .../src/telemetry/report-store.ts | 66 +- .../src/telemetry/session-reports.test.ts | 177 +++++- .../src/telemetry/session-reports.ts | 38 +- .../src/websocket/ingest.test.ts | 276 ++++++++- .../cloud-agent-next/src/websocket/ingest.ts | 14 +- .../session/admission-recovery.test.ts | 323 +++++++++- .../session/execute-directly-failure.test.ts | 11 +- .../session/message-terminalization.test.ts | 124 ++++ .../test/unit/wrapper/batch-admission.test.ts | 21 +- .../cloud-agent-next/vitest.workers.config.ts | 13 + .../wrapper/src/log-uploader.test.ts | 355 ++++++++--- .../wrapper/src/log-uploader.ts | 190 +++--- services/cloud-agent-next/wrapper/src/main.ts | 23 +- .../wrapper/src/server.test.ts | 369 ++++++++++- .../cloud-agent-next/wrapper/src/server.ts | 86 ++- .../wrapper/src/session-bootstrap.test.ts | 192 ++++++ .../wrapper/src/session-bootstrap.ts | 15 +- .../wrapper/src/shutdown.test.ts | 3 + .../cloud-agent-next/wrapper/src/state.ts | 4 +- .../src/github-token-service.test.ts | 120 +++- .../src/github-token-service.ts | 37 +- services/git-token-service/src/index.test.ts | 89 ++- services/git-token-service/src/index.ts | 74 ++- 61 files changed, 6217 insertions(+), 807 deletions(-) create mode 100644 services/cloud-agent-next/src/shared/assistant-failure.ts diff --git a/apps/web/src/app/admin/components/CloudAgentNextTelemetry/CloudAgentNextOutcomesPage.tsx b/apps/web/src/app/admin/components/CloudAgentNextTelemetry/CloudAgentNextOutcomesPage.tsx index 0c92edd640..ff71fef259 100644 --- a/apps/web/src/app/admin/components/CloudAgentNextTelemetry/CloudAgentNextOutcomesPage.tsx +++ b/apps/web/src/app/admin/components/CloudAgentNextTelemetry/CloudAgentNextOutcomesPage.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useEffect, useState } from 'react'; +import { Fragment, useEffect, useState } from 'react'; import { AlertCircle, Loader2, RefreshCw } from 'lucide-react'; import AdminPage from '@/app/admin/components/AdminPage'; import { @@ -85,6 +85,7 @@ const utcLongLabel = new Intl.DateTimeFormat('en-US', { day: 'numeric', hour: '2-digit', minute: '2-digit', + second: '2-digit', hourCycle: 'h23', }); @@ -249,17 +250,49 @@ function responsibilityBadge(responsibility: TopError['responsibility']) { function ErrorSessionsDialog({ error, interval, + trigger, onClose, }: { error: TopError; interval: CloudAgentNextHealthFilters; + trigger: HTMLButtonElement; onClose: () => void; }) { const sessions = useCloudAgentNextHealthErrorSessions(interval, error); - const rows = sessions.data?.rows ?? []; + const [diagnosticCheckTime, setDiagnosticCheckTime] = useState(Date.now); + const now = Date.now(); + const rows = (sessions.data?.rows ?? []).map(row => { + const retained = row.diagnosticExpiresAt && new Date(row.diagnosticExpiresAt).getTime() > now; + return { + ...row, + diagnostic: retained ? row.diagnostic : null, + }; + }); + const nextDiagnosticExpiry = rows.reduce((next, row) => { + const expiry = row.diagnosticExpiresAt ? new Date(row.diagnosticExpiresAt).getTime() : 0; + return row.diagnostic && expiry > diagnosticCheckTime ? Math.min(next, expiry) : next; + }, Infinity); + + useEffect(() => { + if (!Number.isFinite(nextDiagnosticExpiry)) return; + const timeout = window.setTimeout( + () => setDiagnosticCheckTime(Date.now()), + Math.min(nextDiagnosticExpiry - Date.now(), 2_147_483_647) + ); + return () => window.clearTimeout(timeout); + }, [nextDiagnosticExpiry, diagnosticCheckTime]); + return ( !open && onClose()}> - + { + if (trigger.isConnected) { + event.preventDefault(); + trigger.focus(); + } + }} + > Affected sessions @@ -304,43 +337,86 @@ function ErrorSessionsDialog({ showText /> +

+ Message and wrapper IDs, diagnostics, and last seen refer to the latest matching event + per session. Diagnostics expire after 30 days; stored text is redacted and may be + generic. No message history is loaded. +

- Sessions affected by the selected Cloud Agent error. + Sessions affected by the selected Cloud Agent error and their latest matching + event. Kilo session ID Cloud Agent ID - Latest occurrence (UTC) + Last seen (UTC) Events {rows.map(row => ( - - - - {row.kiloSessionId} - - - - - - {row.cloudAgentSessionId} - - - - - {row.occurredAt - ? `${utcLongLabel.format(new Date(row.occurredAt))} UTC` - : '--'} - - - {row.matchingEvents.toLocaleString()} - - + + + + + {row.kiloSessionId} + + + + + + {row.cloudAgentSessionId} + + + + + {row.lastSeen ? ( + + ) : ( + '--' + )} + + + {row.matchingEvents.toLocaleString()} + + + + +
+ {[ + { label: 'Sandbox ID', value: row.sandboxId }, + { + label: + error.source === 'setup' ? 'Initial message ID' : 'Message ID', + value: row.messageId, + }, + { label: 'Wrapper run ID', value: row.wrapperRunId }, + ].map(({ label, value }) => ( +
+
{label}
+
+ {value ?? 'Not recorded'} + {value && ( + + )} +
+
+ ))} +
+
Stored diagnostic
+
+ {row.diagnostic || 'Not available (missing or expired)'} +
+
+
+ + + ))}
@@ -354,18 +430,23 @@ function ErrorSessionsDialog({ function TopErrors({ errors, + totals, interval, responsibility, summary, onResponsibilityChange, }: { errors: TopError[]; + totals: HealthData['errorTotals']; interval: CloudAgentNextHealthFilters; responsibility: CloudAgentFailureResponsibilityFilter; summary: HealthData['summary']; onResponsibilityChange: (value: CloudAgentFailureResponsibilityFilter) => void; }) { - const [selectedError, setSelectedError] = useState(null); + const [selectedError, setSelectedError] = useState<{ + error: TopError; + trigger: HTMLButtonElement; + } | null>(null); const total = errors.reduce((count, error) => count + error.count, 0); return ( @@ -374,7 +455,11 @@ function TopErrors({
Top errors - Setup failures and failed runs only. {total.toLocaleString()} events in the top 10. + Setup failures and failed runs only. Showing {errors.length.toLocaleString()} of{' '} + {totals.groups.toLocaleString()} groups, covering {total.toLocaleString()} of{' '} + {totals.events.toLocaleString()} events for the selected responsibility. + {totals.groups > errors.length && + ` ${(totals.groups - errors.length).toLocaleString()} groups not shown (top 10 limit).`}{' '} Select an error to inspect sessions.
@@ -404,6 +489,11 @@ function TopErrors({
+

+ Sessions and known sandboxes are distinct within each group, not additive across groups. + Known sandboxes exclude sessions without a recorded sandbox ID; zero known does not mean + zero impact. +

{errors.length === 0 ? ( @@ -423,6 +513,8 @@ function TopErrors({ Reason Source Events + Affected sessions + Known sandboxes @@ -436,7 +528,7 @@ function TopErrors({ variant="ghost" className="h-auto w-full justify-start px-2 py-2 text-left" aria-label={`View affected sessions for ${RESPONSIBILITY_LABELS[error.responsibility]} ${failureReasonLabel(error.reason)}, ${error.count.toLocaleString()} events`} - onClick={() => setSelectedError(error)} + onClick={event => setSelectedError({ error, trigger: event.currentTarget })} > {failureReasonLabel(error.reason)} @@ -450,6 +542,18 @@ function TopErrors({ {error.count.toLocaleString()} + + {error.affectedSessions.toLocaleString()} + + + {error.knownSandboxes.toLocaleString()} + {error.sessionsWithoutSandbox > 0 && ( +

+ {error.sessionsWithoutSandbox.toLocaleString()}{' '} + {error.sessionsWithoutSandbox === 1 ? 'session' : 'sessions'} without ID +

+ )} + ))} @@ -458,8 +562,9 @@ function TopErrors({ )} {selectedError && ( setSelectedError(null)} /> )} @@ -571,6 +676,7 @@ export default function CloudAgentNextOutcomesPage() { { it('has a human-readable label for every shared reason', () => { expect(hasExhaustiveFailureReasonLabels()).toBe(true); expect(CLOUD_AGENT_FAILURE_REASONS.map(failureReasonLabel)).not.toContain(''); + expect(CLOUD_AGENT_FAILURE_REASONS.map(failureReasonLabel)).not.toContain(undefined); expect(failureReasonLabel('unclassified')).toBe('Unclassified'); }); + + it.each([ + ['context_limit', 'Context limit'], + ['output_limit', 'Output limit'], + ['content_filter', 'Content filter'], + ['structured_output', 'Invalid structured output'], + ['invalid_request', 'Model request rejected'], + ['request_timeout', 'Request timed out'], + ] as const)('labels %s as %s', (reason, label) => { + expect(failureReasonLabel(reason)).toBe(label); + }); }); diff --git a/apps/web/src/app/admin/components/CloudAgentNextTelemetry/health-summary.ts b/apps/web/src/app/admin/components/CloudAgentNextTelemetry/health-summary.ts index 2ca5970dcf..98c02c8fa2 100644 --- a/apps/web/src/app/admin/components/CloudAgentNextTelemetry/health-summary.ts +++ b/apps/web/src/app/admin/components/CloudAgentNextTelemetry/health-summary.ts @@ -35,6 +35,12 @@ const FAILURE_REASON_LABELS = { managed_provider_authentication: 'Managed provider authentication', managed_model_configuration: 'Managed model configuration', provider_unavailable: 'Provider unavailable', + request_timeout: 'Request timed out', + invalid_request: 'Model request rejected', + context_limit: 'Context limit', + output_limit: 'Output limit', + content_filter: 'Content filter', + structured_output: 'Invalid structured output', source_control_network: 'Source control network', assistant_unknown: 'Unknown assistant failure', workspace_unknown: 'Unknown workspace failure', diff --git a/apps/web/src/lib/code-reviews/terminal-reason-from-failure.test.ts b/apps/web/src/lib/code-reviews/terminal-reason-from-failure.test.ts index 22d38c5f4b..1a89eb3de8 100644 --- a/apps/web/src/lib/code-reviews/terminal-reason-from-failure.test.ts +++ b/apps/web/src/lib/code-reviews/terminal-reason-from-failure.test.ts @@ -82,14 +82,34 @@ describe('terminalReasonFromCloudAgentFailure', () => { ).toBe('assistant_unavailable'); }); - it('resolves every assistant reason to a valid terminal reason', () => { - const valid = new Set(CODE_REVIEW_TERMINAL_REASONS); - const resolved = CLOUD_AGENT_ASSISTANT_FAILURE_REASONS.map(assistantReason => [ - assistantReason, - terminalReasonFromCloudAgentFailure({ code: 'assistant_error', assistantReason }), - ]); + it.each(CLOUD_AGENT_ASSISTANT_FAILURE_REASONS)( + 'resolves assistant reason %s to a defined, valid terminal reason', + assistantReason => { + const reason = terminalReasonFromCloudAgentFailure({ + code: 'assistant_error', + assistantReason, + }); - expect(resolved.filter(([, reason]) => !valid.has(reason as string))).toEqual([]); + expect(reason).toBeDefined(); + expect(CODE_REVIEW_TERMINAL_REASONS).toContain(reason); + } + ); + + it.each([ + ['context_limit', 'assistant_failed'], + ['output_limit', 'assistant_failed'], + ['content_filter', 'assistant_failed'], + ['structured_output', 'assistant_failed'], + ['timeout', 'assistant_timeout'], + ['invalid_request', 'assistant_invalid_request'], + ] as const)('keeps %s mapped to the existing %s category', (assistantReason, expected) => { + expect( + terminalReasonFromCloudAgentFailure({ + code: 'assistant_error', + assistantReason, + message: 'Assistant request was rate limited', + }) + ).toBe(expected); }); it('splits assistant failures by their safe message', () => { diff --git a/apps/web/src/lib/code-reviews/terminal-reason-from-failure.ts b/apps/web/src/lib/code-reviews/terminal-reason-from-failure.ts index 537a390b4d..3b08ebd139 100644 --- a/apps/web/src/lib/code-reviews/terminal-reason-from-failure.ts +++ b/apps/web/src/lib/code-reviews/terminal-reason-from-failure.ts @@ -89,6 +89,10 @@ const ASSISTANT_REASON_REASONS = { timeout: 'assistant_timeout', provider_authentication: 'assistant_unauthorized', invalid_request: 'assistant_invalid_request', + context_limit: 'assistant_failed', + output_limit: 'assistant_failed', + content_filter: 'assistant_failed', + structured_output: 'assistant_failed', insufficient_credits: 'billing', model_unavailable: 'model_not_found', unknown: 'assistant_failed', diff --git a/apps/web/src/routers/admin-cloud-agent-next-router.test.ts b/apps/web/src/routers/admin-cloud-agent-next-router.test.ts index d737189373..4ed7067d51 100644 --- a/apps/web/src/routers/admin-cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/admin-cloud-agent-next-router.test.ts @@ -7,11 +7,13 @@ import { kilocode_users, type User, } from '@kilocode/db/schema'; -import { eq, inArray } from 'drizzle-orm'; +import { and, eq, inArray } from 'drizzle-orm'; const START_DATE = '2035-01-10T00:00:00.000Z'; const END_DATE = '2035-01-11T00:00:00.000Z'; const RAW_CREATED_TIME = '2035-01-10 00:00:00+00'; +const SHARED_SANDBOX_ID = 'usr_admin_outcomes_shared'; +const DIAGNOSTIC_EXPIRES_AT = '2035-02-01T00:00:00.000Z'; const ids = { mapped: 'agent_admin_outcomes_mapped', setupFailed: 'agent_admin_outcomes_setup_failed', @@ -28,6 +30,24 @@ function at(hours: number, minutes: number = 0, seconds: number = 0) { return new Date(Date.UTC(2035, 0, 10, hours, minutes, seconds)).toISOString(); } +const matchingRun = { + cloud_agent_session_id: ids.mapped, + status: 'failed', + terminal_at: at(8), + failure_stage: 'pre_dispatch', + failure_code: 'sandbox_connect_failed', + failure_responsibility: 'platform', + failure_reason: 'sandbox_connectivity', +} as const; + +const matchingError = { + source: 'run', + stage: matchingRun.failure_stage, + code: matchingRun.failure_code, + responsibility: matchingRun.failure_responsibility, + reason: matchingRun.failure_reason, +} as const; + describe('adminCloudAgentNextRouter', () => { let adminUser: User; let regularUser: User; @@ -48,18 +68,22 @@ describe('adminCloudAgentNextRouter', () => { cloud_agent_session_id: ids.mapped, kilo_session_id: 'ses_admin_outcomes_mapped', initial_message_id: 'msg_admin_initial', + sandbox_id: SHARED_SANDBOX_ID, created_at: RAW_CREATED_TIME, }, { cloud_agent_session_id: ids.setupFailed, kilo_session_id: 'ses_admin_setup_failed', initial_message_id: 'msg_setup_failed', + sandbox_id: SHARED_SANDBOX_ID, created_at: '2035-01-09T23:56:00.000Z', - failure_at: at(0, 6), + failure_at: '2035-01-10 00:06:00+00', failure_stage: 'initial_admission', failure_code: 'initial_admission_rejected', failure_responsibility: 'unknown', failure_reason: 'initial_admission_unknown', + error_message_redacted: 'Initial admission failed', + error_expires_at: DIAGNOSTIC_EXPIRES_AT, }, { cloud_agent_session_id: ids.setupFailedLater, @@ -82,6 +106,7 @@ describe('adminCloudAgentNextRouter', () => { cloud_agent_session_id: ids.expired, kilo_session_id: 'ses_admin_outcomes_expired', initial_message_id: 'msg_expired_initial', + sandbox_id: 'usr_admin_outcomes_expired', created_at: '2025-01-10T00:20:00.000Z', }, ]); @@ -101,6 +126,9 @@ describe('adminCloudAgentNextRouter', () => { failure_code: 'sandbox_connect_failed', failure_responsibility: 'platform', failure_reason: 'sandbox_connectivity', + wrapper_run_id: 'wrapper_admin_original', + error_message_redacted: 'Sandbox connection failed', + error_expires_at: DIAGNOSTIC_EXPIRES_AT, }, { cloud_agent_session_id: ids.setupFailed, @@ -158,18 +186,22 @@ describe('adminCloudAgentNextRouter', () => { reason: 'sandbox_connectivity', }) ).rejects.toThrow('Admin access required'); - await expect( - adminCaller.admin.cloudAgentNext.getHealthOverview({ - startDate: END_DATE, - endDate: END_DATE, - }) - ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); - await expect( - adminCaller.admin.cloudAgentNext.getHealthOverview({ - startDate: START_DATE, - endDate: '2035-04-11T00:00:00.000Z', - }) - ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + for (const invalidInterval of [ + { startDate: END_DATE, endDate: END_DATE }, + { startDate: END_DATE, endDate: START_DATE }, + { startDate: START_DATE, endDate: '2035-04-11T00:00:00.000Z' }, + { startDate: RAW_CREATED_TIME, endDate: END_DATE }, + ]) { + await expect( + adminCaller.admin.cloudAgentNext.getHealthOverview(invalidInterval) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + await expect( + adminCaller.admin.cloudAgentNext.listHealthErrorSessions({ + ...invalidInterval, + ...matchingError, + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + } }); it('summarizes health and ranks operational errors without interruptions', async () => { @@ -196,6 +228,20 @@ describe('adminCloudAgentNextRouter', () => { responsibility: 'unknown', reason: 'initial_admission_unknown', count: 1, + affectedSessions: 1, + knownSandboxes: 1, + sessionsWithoutSandbox: 0, + }, + { + source: 'setup', + stage: 'initial_admission', + code: 'invalid_initial_intent', + responsibility: 'user', + reason: 'initial_request_invalid', + count: 1, + affectedSessions: 1, + knownSandboxes: 0, + sessionsWithoutSandbox: 1, }, { source: 'run', @@ -204,6 +250,9 @@ describe('adminCloudAgentNextRouter', () => { responsibility: 'platform', reason: 'sandbox_connectivity', count: 1, + affectedSessions: 1, + knownSandboxes: 1, + sessionsWithoutSandbox: 0, }, { source: 'run', @@ -212,6 +261,9 @@ describe('adminCloudAgentNextRouter', () => { responsibility: 'user', reason: 'insufficient_credits', count: 1, + affectedSessions: 1, + knownSandboxes: 1, + sessionsWithoutSandbox: 0, }, ]) ); @@ -219,6 +271,76 @@ describe('adminCloudAgentNextRouter', () => { expect(JSON.stringify(health.topErrors)).not.toContain('wrapper_start_failed'); }); + it('counts distinct affected sessions and known sandboxes without multiplying setup failures', async () => { + await db.insert(cloud_agent_session_runs).values( + [ + ids.mapped, + ids.mapped, + ids.setupFailed, + ids.unmapped, + ids.unmapped, + ids.setupFailedLater, + ids.expired, + ].map((sessionId, index) => ({ + ...matchingRun, + cloud_agent_session_id: sessionId, + message_id: `msg_admin_repeated_${index}`, + })) + ); + const caller = await createCallerForUser(adminUser.id); + const health = await caller.admin.cloudAgentNext.getHealthOverview(interval()); + const sessions = await caller.admin.cloudAgentNext.listHealthErrorSessions({ + ...interval(), + ...matchingError, + }); + + expect(health.topErrors).toEqual( + expect.arrayContaining([ + { + ...matchingError, + count: 7, + affectedSessions: 4, + knownSandboxes: 1, + sessionsWithoutSandbox: 2, + }, + expect.objectContaining({ + source: 'setup', + reason: 'initial_admission_unknown', + count: 1, + affectedSessions: 1, + knownSandboxes: 1, + sessionsWithoutSandbox: 0, + }), + expect.objectContaining({ + source: 'setup', + reason: 'initial_request_invalid', + count: 1, + affectedSessions: 1, + knownSandboxes: 0, + sessionsWithoutSandbox: 1, + }), + ]) + ); + expect(sessions.totalSessions).toBe(4); + expect(sessions.rows).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + cloudAgentSessionId: ids.mapped, + matchingEvents: 3, + sandboxId: SHARED_SANDBOX_ID, + }), + expect.objectContaining({ + cloudAgentSessionId: ids.unmapped, + matchingEvents: 2, + sandboxId: null, + }), + ]) + ); + expect(sessions.rows).toHaveLength(4); + expect(health.summary.setupFailures).toBe(2); + expect(health.errorTotals).toEqual({ groups: 4, events: 10 }); + }); + it('excludes runs whose session falls outside the 90-day retention window', async () => { const caller = await createCallerForUser(adminUser.id); const health = await caller.admin.cloudAgentNext.getHealthOverview(interval()); @@ -228,24 +350,67 @@ describe('adminCloudAgentNextRouter', () => { expect(JSON.stringify(health.topErrors)).not.toContain('wrapper_start_failed'); }); - it('filters top errors by responsibility without changing visible summary counts', async () => { + it.each([ + { responsibility: 'platform', events: 1, groups: 1 }, + { responsibility: 'user', events: 2, groups: 2 }, + { responsibility: 'unknown', events: 1, groups: 1 }, + ] as const)( + 'filters top errors by $responsibility without changing the global summary', + async ({ responsibility, events, groups }) => { + const caller = await createCallerForUser(adminUser.id); + const allHealth = await caller.admin.cloudAgentNext.getHealthOverview(interval()); + const health = await caller.admin.cloudAgentNext.getHealthOverview({ + ...interval(), + responsibility, + }); + + expect(health.summary).toEqual(allHealth.summary); + expect(health.topErrors).toEqual( + allHealth.topErrors.filter(error => error.responsibility === responsibility) + ); + expect(health.errorTotals).toEqual({ events, groups }); + expect(allHealth.errorTotals).toEqual({ events: 4, groups: 4 }); + } + ); + + it('reports all matching groups and events even when only the top ten are returned', async () => { + const failures = [ + ['pre_dispatch', 'workspace_setup_failed', 'platform', 'sandbox_capacity'], + ['pre_dispatch', 'workspace_setup_failed', 'user', 'source_control_authentication'], + ['pre_dispatch', 'workspace_setup_failed', 'unknown', 'source_control_network'], + ['pre_dispatch', 'wrapper_start_failed', 'platform', 'runtime_startup'], + ['pre_dispatch', 'kilo_server_failed', 'platform', 'runtime_startup'], + ['pre_dispatch', 'delivery_failure_unknown', 'platform', 'delivery'], + ['post_dispatch_no_activity', 'wrapper_ping_timeout', 'platform', 'wrapper_liveness'], + ['agent_activity', 'assistant_error', 'platform', 'managed_provider_unavailable'], + ['agent_activity', 'assistant_error', 'unknown', 'assistant_unknown'], + ] as const; + await db.insert(cloud_agent_session_runs).values([ + ...failures.map(([stage, code, responsibility, reason], index) => ({ + ...matchingRun, + message_id: `msg_admin_group_${index}`, + failure_stage: stage, + failure_code: code, + failure_responsibility: responsibility, + failure_reason: reason, + })), + { ...matchingRun, message_id: 'msg_admin_top_group_repeat_1' }, + { ...matchingRun, message_id: 'msg_admin_top_group_repeat_2' }, + ]); const caller = await createCallerForUser(adminUser.id); - const health = await caller.admin.cloudAgentNext.getHealthOverview({ + const health = await caller.admin.cloudAgentNext.getHealthOverview(interval()); + const platformHealth = await caller.admin.cloudAgentNext.getHealthOverview({ ...interval(), responsibility: 'platform', }); - expect(health.summary).toMatchObject({ - platformFailures: 1, - userFailures: 2, - unknownFailures: 1, - }); - expect(health.topErrors).toEqual([ - expect.objectContaining({ - responsibility: 'platform', - reason: 'sandbox_connectivity', - }), - ]); + expect(health.errorTotals).toEqual({ groups: 13, events: 15 }); + expect(health.topErrors).toHaveLength(10); + expect(health.topErrors[0]).toMatchObject({ ...matchingError, count: 3 }); + expect(health.topErrors.reduce((total, error) => total + error.count, 0)).toBe(12); + expect(platformHealth.errorTotals).toEqual({ groups: 7, events: 9 }); + expect(platformHealth.topErrors).toHaveLength(7); + expect(platformHealth.summary).toEqual(health.summary); }); it('returns null rates when there are no assessed outcomes', async () => { @@ -256,8 +421,349 @@ describe('adminCloudAgentNextRouter', () => { }); expect(health.summary.platformFailureRate).toBeNull(); expect(health.summary.allFailureRate).toBeNull(); + expect(health.topErrors).toEqual([]); + expect(health.errorTotals).toEqual({ events: 0, groups: 0 }); }); + it('keeps inclusive start and exclusive end boundaries for counts and matching details', async () => { + await db + .update(cloud_agent_session_runs) + .set({ terminal_at: START_DATE }) + .where( + and( + eq(cloud_agent_session_runs.cloud_agent_session_id, ids.mapped), + eq(cloud_agent_session_runs.message_id, 'msg_admin_failed_predispatch') + ) + ); + await db.insert(cloud_agent_session_runs).values([ + { ...matchingRun, message_id: 'msg_at_interval_end', terminal_at: END_DATE }, + { + ...matchingRun, + message_id: 'msg_before_interval', + terminal_at: '2035-01-09T23:59:59.999Z', + }, + { ...matchingRun, message_id: 'msg_without_terminal_time', terminal_at: null }, + ]); + await db + .update(cloud_agent_sessions) + .set({ failure_at: END_DATE }) + .where(eq(cloud_agent_sessions.cloud_agent_session_id, ids.setupFailedLater)); + const caller = await createCallerForUser(adminUser.id); + const health = await caller.admin.cloudAgentNext.getHealthOverview(interval()); + const sessions = await caller.admin.cloudAgentNext.listHealthErrorSessions({ + ...interval(), + ...matchingError, + }); + + expect(health.topErrors).toContainEqual({ + ...matchingError, + count: 1, + affectedSessions: 1, + knownSandboxes: 1, + sessionsWithoutSandbox: 0, + }); + expect(health.summary.setupFailures).toBe(1); + expect(sessions.rows).toEqual([ + expect.objectContaining({ + messageId: 'msg_admin_failed_predispatch', + lastSeen: START_DATE, + matchingEvents: 1, + }), + ]); + }); + + it.each([ + [null, 'assistant_error', 'unknown', 'assistant_error'], + ['agent_activity', null, 'agent_activity', 'unclassified'], + [null, 'unclassified', 'unknown', 'unclassified'], + ['unknown', null, 'unknown', 'unclassified'], + ] as const)( + 'matches partially classified runs (%s, %s) with overview normalization', + async (storedStage, storedCode, stage, code) => { + await db.insert(cloud_agent_session_runs).values([ + { + ...matchingRun, + cloud_agent_session_id: ids.unmapped, + message_id: 'msg_partial_classification', + failure_stage: storedStage, + failure_code: storedCode, + failure_responsibility: null, + failure_reason: null, + }, + { + ...matchingRun, + cloud_agent_session_id: ids.unmapped, + message_id: 'msg_explicit_classification', + terminal_at: at(7), + failure_stage: stage, + failure_code: code, + failure_responsibility: 'unknown', + failure_reason: 'unclassified', + }, + ]); + const caller = await createCallerForUser(adminUser.id); + const error = { + source: 'run', + stage, + code, + responsibility: 'unknown', + reason: 'unclassified', + } as const; + const health = await caller.admin.cloudAgentNext.getHealthOverview(interval()); + const sessions = await caller.admin.cloudAgentNext.listHealthErrorSessions({ + ...interval(), + ...error, + }); + + expect(health.topErrors).toContainEqual({ + ...error, + count: 2, + affectedSessions: 1, + knownSandboxes: 0, + sessionsWithoutSandbox: 1, + }); + expect(sessions).toMatchObject({ + totalSessions: 1, + rows: [ + { + cloudAgentSessionId: ids.unmapped, + messageId: 'msg_partial_classification', + lastSeen: at(8), + matchingEvents: 2, + }, + ], + }); + } + ); + + it('limits the drilldown to the latest 100 sessions without truncating totals or event counts', async () => { + const sessionIds = Array.from( + { length: 101 }, + (_, index) => `agent_admin_outcomes_limit_${index.toString().padStart(3, '0')}` + ); + try { + await db.insert(cloud_agent_sessions).values( + sessionIds.map(sessionId => ({ + cloud_agent_session_id: sessionId, + kilo_session_id: `ses_${sessionId}`, + initial_message_id: `msg_${sessionId}`, + created_at: RAW_CREATED_TIME, + })) + ); + await db.insert(cloud_agent_session_runs).values([ + ...sessionIds.map(sessionId => ({ + ...matchingRun, + cloud_agent_session_id: sessionId, + message_id: `msg_${sessionId}`, + })), + { + ...matchingRun, + cloud_agent_session_id: sessionIds[100], + message_id: 'msg_limit_repeated', + terminal_at: at(9), + }, + ]); + const caller = await createCallerForUser(adminUser.id); + const sessions = await caller.admin.cloudAgentNext.listHealthErrorSessions({ + ...interval(), + ...matchingError, + }); + + expect(sessions.totalSessions).toBe(102); + expect(sessions.limit).toBe(100); + expect(sessions.rows.map(row => row.cloudAgentSessionId)).toEqual( + [...sessionIds].reverse().slice(0, 100) + ); + expect(sessions.rows[0]).toMatchObject({ + messageId: 'msg_limit_repeated', + lastSeen: at(9), + matchingEvents: 2, + }); + } finally { + await db + .delete(cloud_agent_sessions) + .where(inArray(cloud_agent_sessions.cloud_agent_session_id, sessionIds)); + } + }); + + it('pairs details from the latest matching run using terminal time and message ID ties', async () => { + await db.insert(cloud_agent_session_runs).values([ + { + ...matchingRun, + message_id: 'msg_zz_older', + terminal_at: at(7), + wrapper_run_id: 'wrapper_zz_older', + error_message_redacted: 'Z older diagnostic', + error_expires_at: DIAGNOSTIC_EXPIRES_AT, + }, + { + ...matchingRun, + message_id: 'msg_tied_a', + wrapper_run_id: 'wrapper_z_tied', + error_message_redacted: 'B tied diagnostic', + error_expires_at: DIAGNOSTIC_EXPIRES_AT, + }, + { + ...matchingRun, + message_id: 'msg_tied_b', + terminal_at: '2035-01-10 08:00:00+00', + wrapper_run_id: 'wrapper_a_selected', + error_message_redacted: 'A selected diagnostic', + error_expires_at: '2035-01-30 12:00:00+00', + }, + { + ...matchingRun, + message_id: 'msg_newer_other_reason', + terminal_at: at(9), + failure_reason: 'runtime_startup', + wrapper_run_id: 'wrapper_other_reason', + error_message_redacted: 'Other reason diagnostic', + error_expires_at: DIAGNOSTIC_EXPIRES_AT, + }, + { + ...matchingRun, + message_id: 'msg_newer_completed', + terminal_at: at(10), + status: 'completed', + }, + { + ...matchingRun, + message_id: 'msg_outside_interval', + terminal_at: END_DATE, + }, + ]); + const caller = await createCallerForUser(adminUser.id); + const sessions = await caller.admin.cloudAgentNext.listHealthErrorSessions({ + ...interval(), + ...matchingError, + }); + + expect(sessions).toEqual({ + totalSessions: 1, + limit: 100, + rows: [ + { + cloudAgentSessionId: ids.mapped, + kiloSessionId: 'ses_admin_outcomes_mapped', + sandboxId: SHARED_SANDBOX_ID, + messageId: 'msg_tied_b', + wrapperRunId: 'wrapper_a_selected', + diagnostic: 'A selected diagnostic', + diagnosticExpiresAt: '2035-01-30T12:00:00.000Z', + occurredAt: at(8), + lastSeen: at(8), + matchingEvents: 4, + }, + ], + }); + }); + + it.each([ + { + name: 'expired', + diagnostic: 'Expired run diagnostic', + expiresAt: '2000-01-01T00:00:00.000Z', + }, + { name: 'missing', diagnostic: null, expiresAt: null }, + ] as const)( + 'keeps the latest run when its diagnostic is $name without returning older text', + async ({ diagnostic, expiresAt }) => { + await db.insert(cloud_agent_session_runs).values({ + ...matchingRun, + message_id: 'msg_admin_latest_without_diagnostic', + error_message_redacted: diagnostic, + error_expires_at: expiresAt, + }); + const caller = await createCallerForUser(adminUser.id); + const sessions = await caller.admin.cloudAgentNext.listHealthErrorSessions({ + ...interval(), + ...matchingError, + }); + const health = await caller.admin.cloudAgentNext.getHealthOverview(interval()); + + expect(sessions.rows).toEqual([ + expect.objectContaining({ + messageId: 'msg_admin_latest_without_diagnostic', + wrapperRunId: null, + diagnostic: null, + diagnosticExpiresAt: expiresAt, + lastSeen: at(8), + matchingEvents: 2, + }), + ]); + expect(JSON.stringify(sessions)).not.toContain('Expired run diagnostic'); + expect(JSON.stringify(sessions)).not.toContain('Sandbox connection failed'); + expect(health.topErrors).toContainEqual({ + ...matchingError, + count: 2, + affectedSessions: 1, + knownSandboxes: 1, + sessionsWithoutSandbox: 0, + }); + expect(health.summary.failedRuns).toBe(3); + } + ); + + it('rejects retained diagnostic text without an expiry', async () => { + await expect( + db.insert(cloud_agent_session_runs).values({ + ...matchingRun, + message_id: 'msg_unexpiring_diagnostic', + error_message_redacted: 'Diagnostic requiring expiry', + error_expires_at: null, + }) + ).rejects.toMatchObject({ + cause: { constraint: 'cloud_agent_session_runs_error_expiry_check' }, + }); + }); + + it.each([ + { + name: 'expired', + diagnostic: 'Expired setup diagnostic', + expiresAt: '2000-01-01T00:00:00.000Z', + }, + { name: 'missing', diagnostic: null, expiresAt: null }, + ])( + 'keeps setup counts and initial message IDs when the diagnostic is $name', + async ({ diagnostic, expiresAt }) => { + await db + .update(cloud_agent_sessions) + .set({ error_message_redacted: diagnostic, error_expires_at: expiresAt }) + .where(eq(cloud_agent_sessions.cloud_agent_session_id, ids.setupFailed)); + const caller = await createCallerForUser(adminUser.id); + const sessions = await caller.admin.cloudAgentNext.listHealthErrorSessions({ + ...interval(), + source: 'setup', + stage: 'initial_admission', + code: 'initial_admission_rejected', + responsibility: 'unknown', + reason: 'initial_admission_unknown', + }); + const health = await caller.admin.cloudAgentNext.getHealthOverview(interval()); + + expect(sessions).toMatchObject({ + totalSessions: 1, + rows: [ + { + cloudAgentSessionId: ids.setupFailed, + messageId: 'msg_setup_failed', + wrapperRunId: null, + diagnostic: null, + diagnosticExpiresAt: expiresAt, + lastSeen: at(0, 6), + matchingEvents: 1, + }, + ], + }); + expect(JSON.stringify(sessions)).not.toContain('Expired setup diagnostic'); + expect(health.summary.setupFailures).toBe(2); + expect(health.topErrors).toContainEqual( + expect.objectContaining({ source: 'setup', reason: 'initial_admission_unknown', count: 1 }) + ); + } + ); + it('lists affected sessions for an exact top-error source and occurrence interval', async () => { await db.insert(cloud_agent_session_runs).values([ { @@ -335,7 +841,13 @@ describe('adminCloudAgentNextRouter', () => { expect.objectContaining({ cloudAgentSessionId: ids.setupFailed, kiloSessionId: 'ses_admin_setup_failed', + sandboxId: SHARED_SANDBOX_ID, + messageId: 'msg_setup_failed', + wrapperRunId: null, + diagnostic: 'Initial admission failed', + diagnosticExpiresAt: DIAGNOSTIC_EXPIRES_AT, occurredAt: at(0, 6), + lastSeen: at(0, 6), matchingEvents: 1, }), ]) @@ -347,7 +859,13 @@ describe('adminCloudAgentNextRouter', () => { expect.objectContaining({ cloudAgentSessionId: ids.mapped, kiloSessionId: 'ses_admin_outcomes_mapped', + sandboxId: SHARED_SANDBOX_ID, + messageId: 'msg_admin_failed_predispatch', + wrapperRunId: 'wrapper_admin_original', + diagnostic: 'Sandbox connection failed', + diagnosticExpiresAt: DIAGNOSTIC_EXPIRES_AT, occurredAt: at(2, 2), + lastSeen: at(2, 2), matchingEvents: 1, }), ], @@ -358,7 +876,13 @@ describe('adminCloudAgentNextRouter', () => { expect.arrayContaining([ expect.objectContaining({ cloudAgentSessionId: ids.unmapped, + sandboxId: null, + messageId: 'msg_admin_failed_unclassified', + wrapperRunId: null, + diagnostic: null, + diagnosticExpiresAt: null, occurredAt: at(6, 1), + lastSeen: at(6, 1), matchingEvents: 1, }), expect.objectContaining({ diff --git a/apps/web/src/routers/admin-cloud-agent-next-router.ts b/apps/web/src/routers/admin-cloud-agent-next-router.ts index cafa058830..cfa5ab17e8 100644 --- a/apps/web/src/routers/admin-cloud-agent-next-router.ts +++ b/apps/web/src/routers/admin-cloud-agent-next-router.ts @@ -1,7 +1,7 @@ import { adminProcedure, createTRPCRouter } from '@/lib/trpc/init'; import { db } from '@/lib/drizzle'; import { cloud_agent_session_runs, cloud_agent_sessions } from '@kilocode/db/schema'; -import { and, desc, eq, gte, isNotNull, isNull, lt, or, sql, type SQL } from 'drizzle-orm'; +import { and, desc, eq, gte, isNotNull, lt, sql, type SQL } from 'drizzle-orm'; import * as z from 'zod'; import { CloudAgentFailureReasonSchema, @@ -87,6 +87,9 @@ type HealthError = { responsibility: CloudAgentFailureResponsibility; reason: z.infer; count: number; + affectedSessions: number; + knownSandboxes: number; + sessionsWithoutSandbox: number; }; function failureRate(failures: number, completed: number): number | null { @@ -138,6 +141,8 @@ export const adminCloudAgentNextRouter = createTRPCRouter({ responsibility: sessionResponsibility, reason: sessionReason, count: sql`COUNT(*)`, + knownSandboxes: sql`COUNT(DISTINCT ${cloud_agent_sessions.sandbox_id})`, + sessionsWithoutSandbox: sql`COUNT(*) FILTER (WHERE ${cloud_agent_sessions.sandbox_id} IS NULL)`, }) .from(cloud_agent_sessions) .where( @@ -157,6 +162,9 @@ export const adminCloudAgentNextRouter = createTRPCRouter({ responsibility: runResponsibility, reason: runReason, count: sql`COUNT(*)`, + affectedSessions: sql`COUNT(DISTINCT ${cloud_agent_session_runs.cloud_agent_session_id})`, + knownSandboxes: sql`COUNT(DISTINCT ${cloud_agent_sessions.sandbox_id})`, + sessionsWithoutSandbox: sql`COUNT(DISTINCT ${cloud_agent_session_runs.cloud_agent_session_id}) FILTER (WHERE ${cloud_agent_sessions.sandbox_id} IS NULL)`, }) .from(cloud_agent_session_runs) .innerJoin( @@ -187,7 +195,7 @@ export const adminCloudAgentNextRouter = createTRPCRouter({ platformFailureRate: null as number | null, allFailureRate: null as number | null, }; - const setupErrorsByCode = new Map(); + const setupErrors: HealthError[] = []; for (const setupRow of setupRows) { const occurrences = count(setupRow.count); summary.setupFailures += occurrences; @@ -197,23 +205,20 @@ export const adminCloudAgentNextRouter = createTRPCRouter({ if (input.responsibility !== 'all' && setupRow.responsibility !== input.responsibility) { continue; } - const key = `${setupRow.responsibility}:${setupRow.reason}:${setupRow.stage}:${setupRow.code}`; - const existingError = setupErrorsByCode.get(key); - if (existingError) { - existingError.count += occurrences; - } else { - setupErrorsByCode.set(key, { - source: 'setup', - stage: setupRow.stage, - code: setupRow.code, - responsibility: setupRow.responsibility, - reason: setupRow.reason, - count: occurrences, - }); - } + setupErrors.push({ + source: 'setup', + stage: setupRow.stage, + code: setupRow.code, + responsibility: setupRow.responsibility, + reason: setupRow.reason, + count: occurrences, + affectedSessions: occurrences, + knownSandboxes: count(setupRow.knownSandboxes), + sessionsWithoutSandbox: count(setupRow.sessionsWithoutSandbox), + }); } - const topErrors = [ - ...setupErrorsByCode.values(), + const errors = [ + ...setupErrors, ...runErrorRows.map( runRow => ({ @@ -223,24 +228,30 @@ export const adminCloudAgentNextRouter = createTRPCRouter({ responsibility: runRow.responsibility, reason: runRow.reason, count: count(runRow.count), + affectedSessions: count(runRow.affectedSessions), + knownSandboxes: count(runRow.knownSandboxes), + sessionsWithoutSandbox: count(runRow.sessionsWithoutSandbox), }) satisfies HealthError ), - ] - .sort( - (left, right) => - right.count - left.count || - left.source.localeCompare(right.source) || - left.stage.localeCompare(right.stage) || - left.code.localeCompare(right.code) || - left.reason.localeCompare(right.reason) - ) - .slice(0, 10); + ].sort( + (left, right) => + right.count - left.count || + left.source.localeCompare(right.source) || + left.stage.localeCompare(right.stage) || + left.code.localeCompare(right.code) || + left.reason.localeCompare(right.reason) + ); + const topErrors = errors.slice(0, 10); + const errorTotals = { + groups: errors.length, + events: errors.reduce((total, error) => total + error.count, 0), + }; summary.platformFailureRate = failureRate(summary.platformFailures, summary.completedRuns); summary.allFailureRate = failureRate( summary.failedRuns + summary.setupFailures, summary.completedRuns ); - return { summary, topErrors }; + return { summary, topErrors, errorTotals }; }), listHealthErrorSessions: adminProcedure @@ -266,6 +277,12 @@ export const adminCloudAgentNextRouter = createTRPCRouter({ .select({ cloudAgentSessionId: cloud_agent_sessions.cloud_agent_session_id, kiloSessionId: cloud_agent_sessions.kilo_session_id, + sandboxId: cloud_agent_sessions.sandbox_id, + messageId: cloud_agent_sessions.initial_message_id, + diagnostic: sql< + string | null + >`CASE WHEN ${cloud_agent_sessions.error_expires_at} > now() THEN ${cloud_agent_sessions.error_message_redacted} ELSE NULL END`, + diagnosticExpiresAt: cloud_agent_sessions.error_expires_at, occurredAt: cloud_agent_sessions.failure_at, }) .from(cloud_agent_sessions) @@ -282,35 +299,60 @@ export const adminCloudAgentNextRouter = createTRPCRouter({ rows: rows.map(row => ({ cloudAgentSessionId: row.cloudAgentSessionId, kiloSessionId: row.kiloSessionId, + sandboxId: row.sandboxId, + messageId: row.messageId, + wrapperRunId: null, + diagnostic: row.diagnostic, + diagnosticExpiresAt: nullableIso(row.diagnosticExpiresAt), occurredAt: nullableIso(row.occurredAt), + lastSeen: nullableIso(row.occurredAt), matchingEvents: 1, })), }; } - const classifiedFailureCondition = - and( - sql`${cloud_agent_session_runs.failure_stage} = ${input.stage}`, - sql`${cloud_agent_session_runs.failure_code} = ${input.code}` - ) ?? sql`false`; - const selectedFailureCondition = - input.stage === 'unknown' && input.code === 'unclassified' - ? (or( - and( - isNull(cloud_agent_session_runs.failure_stage), - isNull(cloud_agent_session_runs.failure_code) - ), - classifiedFailureCondition - ) ?? sql`false`) - : classifiedFailureCondition; const where = and( eq(cloud_agent_session_runs.status, 'failed'), - selectedFailureCondition, + sql`COALESCE(${cloud_agent_session_runs.failure_stage}, 'unknown') = ${input.stage}`, + sql`COALESCE(${cloud_agent_session_runs.failure_code}, 'unclassified') = ${input.code}`, sql`COALESCE(${cloud_agent_session_runs.failure_responsibility}, 'unknown') = ${input.responsibility}`, sql`COALESCE(${cloud_agent_session_runs.failure_reason}, 'unclassified') = ${input.reason}`, ...terminalRunIntervalConditions(input) ); - const latestOccurredAt = sql`MAX(${cloud_agent_session_runs.terminal_at})`; + const latestMatchingRuns = db + .selectDistinctOn([cloud_agent_session_runs.cloud_agent_session_id], { + cloudAgentSessionId: cloud_agent_session_runs.cloud_agent_session_id, + kiloSessionId: cloud_agent_sessions.kilo_session_id, + sandboxId: cloud_agent_sessions.sandbox_id, + messageId: cloud_agent_session_runs.message_id, + wrapperRunId: cloud_agent_session_runs.wrapper_run_id, + diagnostic: sql< + string | null + >`CASE WHEN ${cloud_agent_session_runs.error_expires_at} > now() THEN ${cloud_agent_session_runs.error_message_redacted} ELSE NULL END`.as( + 'diagnostic' + ), + diagnosticExpiresAt: cloud_agent_session_runs.error_expires_at, + occurredAt: cloud_agent_session_runs.terminal_at, + matchingEvents: + sql`COUNT(*) OVER (PARTITION BY ${cloud_agent_session_runs.cloud_agent_session_id})`.as( + 'matching_events' + ), + }) + .from(cloud_agent_session_runs) + .innerJoin( + cloud_agent_sessions, + eq( + cloud_agent_session_runs.cloud_agent_session_id, + cloud_agent_sessions.cloud_agent_session_id + ) + ) + .where(where) + .orderBy( + cloud_agent_session_runs.cloud_agent_session_id, + desc(cloud_agent_session_runs.terminal_at), + desc(cloud_agent_session_runs.message_id) + ) + .as('latest_matching_runs'); const [totals, rows] = await Promise.all([ db .select({ @@ -326,26 +368,12 @@ export const adminCloudAgentNextRouter = createTRPCRouter({ ) .where(where), db - .select({ - cloudAgentSessionId: cloud_agent_sessions.cloud_agent_session_id, - kiloSessionId: cloud_agent_sessions.kilo_session_id, - occurredAt: latestOccurredAt, - matchingEvents: sql`COUNT(*)`, - }) - .from(cloud_agent_session_runs) - .innerJoin( - cloud_agent_sessions, - eq( - cloud_agent_session_runs.cloud_agent_session_id, - cloud_agent_sessions.cloud_agent_session_id - ) - ) - .where(where) - .groupBy( - cloud_agent_sessions.cloud_agent_session_id, - cloud_agent_sessions.kilo_session_id + .select() + .from(latestMatchingRuns) + .orderBy( + desc(latestMatchingRuns.occurredAt), + desc(latestMatchingRuns.cloudAgentSessionId) ) - .orderBy(desc(latestOccurredAt), desc(cloud_agent_sessions.cloud_agent_session_id)) .limit(HEALTH_ERROR_SESSION_LIMIT), ]); return { @@ -354,7 +382,13 @@ export const adminCloudAgentNextRouter = createTRPCRouter({ rows: rows.map(row => ({ cloudAgentSessionId: row.cloudAgentSessionId, kiloSessionId: row.kiloSessionId, + sandboxId: row.sandboxId, + messageId: row.messageId, + wrapperRunId: row.wrapperRunId, + diagnostic: row.diagnostic, + diagnosticExpiresAt: nullableIso(row.diagnosticExpiresAt), occurredAt: nullableIso(row.occurredAt), + lastSeen: nullableIso(row.occurredAt), matchingEvents: count(row.matchingEvents), })), }; diff --git a/packages/db/src/schema.test.ts b/packages/db/src/schema.test.ts index 501cae8d65..247c38f802 100644 --- a/packages/db/src/schema.test.ts +++ b/packages/db/src/schema.test.ts @@ -456,6 +456,72 @@ describe('database schema', () => { } }); + it('requires cloud agent run diagnostic text and expiry to be set and cleared together', async () => { + const sessionId = `schema-cloud-agent-${crypto.randomUUID()}`; + const now = new Date().toISOString(); + await schemaTestDb.db.insert(schema.cloud_agent_sessions).values({ + cloud_agent_session_id: sessionId, + kilo_session_id: `${sessionId}-kilo`, + initial_message_id: `${sessionId}-initial`, + created_at: now, + }); + + try { + const runs = await schemaTestDb.db + .insert(schema.cloud_agent_session_runs) + .values([ + { + cloud_agent_session_id: sessionId, + message_id: 'no-diagnostic', + status: 'failed', + }, + { + cloud_agent_session_id: sessionId, + message_id: 'retained-diagnostic', + status: 'failed', + error_message_redacted: 'The model context limit was exceeded', + error_expires_at: now, + }, + ]) + .returning(); + + expect(runs).toHaveLength(2); + const clearedRuns = await schemaTestDb.db + .update(schema.cloud_agent_session_runs) + .set({ error_message_redacted: null, error_expires_at: null }) + .where(eq(schema.cloud_agent_session_runs.cloud_agent_session_id, sessionId)) + .returning(); + expect(clearedRuns).toHaveLength(2); + for (const run of clearedRuns) { + expect(run).toMatchObject({ error_message_redacted: null, error_expires_at: null }); + } + await expect( + schemaTestDb.db.insert(schema.cloud_agent_session_runs).values({ + cloud_agent_session_id: sessionId, + message_id: 'missing-text', + status: 'failed', + error_expires_at: now, + }) + ).rejects.toMatchObject({ + cause: { constraint: 'cloud_agent_session_runs_error_expiry_check' }, + }); + await expect( + schemaTestDb.db.insert(schema.cloud_agent_session_runs).values({ + cloud_agent_session_id: sessionId, + message_id: 'missing-expiry', + status: 'failed', + error_message_redacted: 'The model context limit was exceeded', + }) + ).rejects.toMatchObject({ + cause: { constraint: 'cloud_agent_session_runs_error_expiry_check' }, + }); + } finally { + await schemaTestDb.db + .delete(schema.cloud_agent_sessions) + .where(eq(schema.cloud_agent_sessions.cloud_agent_session_id, sessionId)); + } + }); + /** * This test ensures that if someone adds/removes values from enums used in schema check constraints, * they are reminded to generate a migration. The check constraints in the database must match the diff --git a/packages/db/src/schema.ts b/packages/db/src/schema.ts index 9f4d976685..935d86276a 100644 --- a/packages/db/src/schema.ts +++ b/packages/db/src/schema.ts @@ -5965,6 +5965,12 @@ export type CloudAgentFailureReason = | 'managed_provider_authentication' | 'managed_model_configuration' | 'provider_unavailable' + | 'request_timeout' + | 'invalid_request' + | 'context_limit' + | 'output_limit' + | 'content_filter' + | 'structured_output' | 'source_control_network' | 'assistant_unknown' | 'workspace_unknown' diff --git a/packages/worker-utils/src/cloud-agent-failure.test.ts b/packages/worker-utils/src/cloud-agent-failure.test.ts index 1ce98df5e4..353e3060da 100644 --- a/packages/worker-utils/src/cloud-agent-failure.test.ts +++ b/packages/worker-utils/src/cloud-agent-failure.test.ts @@ -2,7 +2,9 @@ import { describe, expect, it } from 'vitest'; import { CLOUD_AGENT_FAILURE_CODES, CLOUD_AGENT_FAILURE_STAGES, + CLOUD_AGENT_PROVIDER_OWNERSHIPS, CloudAgentCallbackFailureSchema, + CloudAgentFailureReasonSchema, CloudAgentSafeFailureSchema, classifyCloudAgentFailure, isWorkspaceFailureSubtype, @@ -26,6 +28,9 @@ describe('CloudAgentCallbackFailureSchema', () => { { code: 'future_failure_code', message: 'Future failure' }, { code: 'workspace_setup_failed', subtype: 'future_workspace_failure' }, { code: 'assistant_error', futureField: true }, + { code: 'assistant_error', assistantReason: 'ContextOverflowError' }, + { code: 'assistant_error', assistantReason: null }, + { code: 'assistant_error', providerOwnership: 'future_ownership' }, { attempts: -1 }, { message: 'x'.repeat(4_097) }, ])('discards unsupported or malformed structured failures: %o', failure => { @@ -72,27 +77,110 @@ describe('classifyCloudAgentFailure', () => { ).toEqual({ responsibility: 'unknown', reason: 'source_control_network' }); }); - it('uses provider ownership for authentication and availability failures', () => { - expect( - classifyCloudAgentFailure({ - source: 'run', - stage: 'agent_activity', - code: 'assistant_error', - assistantReason: 'provider_authentication', - providerOwnership: 'byok', - }) - ).toEqual({ responsibility: 'user', reason: 'provider_authentication' }); - expect( - classifyCloudAgentFailure({ + it.each([ + ['provider_authentication', 'byok', 'user', 'provider_authentication'], + ['provider_authentication', 'managed', 'platform', 'managed_provider_authentication'], + ['provider_authentication', 'unknown', 'unknown', 'provider_authentication'], + ['provider_authentication', undefined, 'unknown', 'provider_authentication'], + ['provider_unavailable', 'byok', 'unknown', 'provider_unavailable'], + ['provider_unavailable', 'managed', 'platform', 'managed_provider_unavailable'], + ['provider_unavailable', 'unknown', 'unknown', 'provider_unavailable'], + ['provider_unavailable', undefined, 'unknown', 'provider_unavailable'], + ] as const)( + 'classifies %s with %s ownership without losing the known cause', + (assistantReason, providerOwnership, responsibility, reason) => { + expect( + classifyCloudAgentFailure({ + source: 'run', + stage: 'agent_activity', + code: 'assistant_error', + assistantReason, + providerOwnership, + }) + ).toEqual({ responsibility, reason }); + } + ); + + it.each([ + ['managed', 'platform'], + ['byok', 'unknown'], + ['unknown', 'unknown'], + [undefined, 'unknown'], + ] as const)('retains request_timeout with %s ownership', (providerOwnership, responsibility) => { + const failure = classifyCloudAgentFailure({ + source: 'run', + stage: 'agent_activity', + code: 'assistant_error', + assistantReason: 'timeout', + providerOwnership, + }); + + expect(failure).toEqual({ responsibility, reason: 'request_timeout' }); + expect(CloudAgentFailureReasonSchema.parse(failure.reason)).toBe('request_timeout'); + }); + + it.each([ + 'context_limit', + 'output_limit', + 'content_filter', + 'structured_output', + 'invalid_request', + ] as const)('retains %s without inferring responsibility from ownership', assistantReason => { + for (const providerOwnership of [...CLOUD_AGENT_PROVIDER_OWNERSHIPS, undefined]) { + const failure = classifyCloudAgentFailure({ source: 'run', stage: 'agent_activity', code: 'assistant_error', - assistantReason: 'provider_unavailable', - providerOwnership: 'managed', - }) - ).toEqual({ responsibility: 'platform', reason: 'managed_provider_unavailable' }); + assistantReason, + providerOwnership, + }); + + expect(failure).toEqual({ responsibility: 'unknown', reason: assistantReason }); + expect(CloudAgentFailureReasonSchema.parse(failure.reason)).toBe(assistantReason); + } }); + it.each(['insufficient_credits', 'rate_limited'] as const)( + 'keeps %s as user responsibility regardless of provider ownership', + assistantReason => { + for (const providerOwnership of [...CLOUD_AGENT_PROVIDER_OWNERSHIPS, undefined]) { + expect( + classifyCloudAgentFailure({ + source: 'run', + stage: 'agent_activity', + code: 'assistant_error', + assistantReason, + providerOwnership, + }) + ).toEqual({ responsibility: 'user', reason: assistantReason }); + } + } + ); + + it.each([true, false, undefined])( + 'uses managed model selection %s rather than provider ownership for model failures', + managedModelSelection => { + for (const providerOwnership of [...CLOUD_AGENT_PROVIDER_OWNERSHIPS, undefined]) { + for (const code of ['assistant_error', 'model_missing'] as const) { + expect( + classifyCloudAgentFailure({ + source: 'run', + stage: 'agent_activity', + code, + assistantReason: 'model_unavailable', + providerOwnership, + managedModelSelection, + }) + ).toEqual( + managedModelSelection + ? { responsibility: 'platform', reason: 'managed_model_configuration' } + : { responsibility: 'user', reason: 'model_unavailable' } + ); + } + } + } + ); + it('classifies setup failures from structured stage and code only', () => { expect( classifyCloudAgentFailure({ diff --git a/packages/worker-utils/src/cloud-agent-failure.ts b/packages/worker-utils/src/cloud-agent-failure.ts index bc41d1f110..4faec14808 100644 --- a/packages/worker-utils/src/cloud-agent-failure.ts +++ b/packages/worker-utils/src/cloud-agent-failure.ts @@ -78,6 +78,12 @@ export const CLOUD_AGENT_FAILURE_REASONS = [ 'managed_provider_authentication', 'managed_model_configuration', 'provider_unavailable', + 'request_timeout', + 'invalid_request', + 'context_limit', + 'output_limit', + 'content_filter', + 'structured_output', 'source_control_network', 'assistant_unknown', 'workspace_unknown', @@ -101,6 +107,10 @@ export const CLOUD_AGENT_ASSISTANT_FAILURE_REASONS = [ 'provider_unavailable', 'timeout', 'invalid_request', + 'context_limit', + 'output_limit', + 'content_filter', + 'structured_output', 'unknown', ] as const; export const CloudAgentAssistantFailureReasonSchema = z.enum(CLOUD_AGENT_ASSISTANT_FAILURE_REASONS); @@ -188,13 +198,28 @@ function classifyAssistantFailure(input: RunFailureFacts): CloudAgentFailureClas if (input.providerOwnership === 'managed') { return classified('platform', 'managed_provider_authentication'); } - return classified('unknown', 'assistant_unknown'); + return classified('unknown', 'provider_authentication'); + } + if (input.assistantReason === 'timeout') { + return classified( + input.providerOwnership === 'managed' ? 'platform' : 'unknown', + 'request_timeout' + ); } - if (input.assistantReason === 'provider_unavailable' || input.assistantReason === 'timeout') { + if (input.assistantReason === 'provider_unavailable') { return input.providerOwnership === 'managed' ? classified('platform', 'managed_provider_unavailable') : classified('unknown', 'provider_unavailable'); } + if ( + input.assistantReason === 'invalid_request' || + input.assistantReason === 'context_limit' || + input.assistantReason === 'output_limit' || + input.assistantReason === 'content_filter' || + input.assistantReason === 'structured_output' + ) { + return classified('unknown', input.assistantReason); + } return classified('unknown', 'assistant_unknown'); } diff --git a/packages/worker-utils/src/cloud-agent-queue-report.test.ts b/packages/worker-utils/src/cloud-agent-queue-report.test.ts index 8a7189bd4e..f6b3506115 100644 --- a/packages/worker-utils/src/cloud-agent-queue-report.test.ts +++ b/packages/worker-utils/src/cloud-agent-queue-report.test.ts @@ -42,6 +42,24 @@ describe('CloudAgentQueueReportSchema', () => { ).toBe(true); }); + it('rejects unknown diagnostic content at the strict producer boundary', () => { + expect( + CloudAgentQueueReportSchema.safeParse( + reportWithRun({ + status: 'failed', + terminalAt, + failureStage: 'unknown', + failureCode: 'unclassified', + diagnostic: { + errorMessageRedacted: 'The agent failed', + errorExpiresAt: '2026-06-25T08:04:00.000Z', + responseBody: 'private fixture output', + }, + }) + ).success + ).toBe(false); + }); + it('supports retained statuses and typed failure classifications', () => { const reportsByStatus = { queued: reportWithRun({ status: 'queued' }), @@ -79,6 +97,46 @@ describe('CloudAgentQueueReportSchema', () => { } }); + it.each(['wrapper_ping_timeout', 'wrapper_no_output', 'wrapper_disconnected'] as const)( + 'accepts %s before and after observed agent activity', + failureCode => { + for (const failureStage of ['post_dispatch_no_activity', 'agent_activity'] as const) { + const report = CloudAgentQueueReportSchema.parse( + reportWithRun({ + status: 'failed', + dispatchAcceptedAt: '2026-05-26T08:02:00.000Z', + ...(failureStage === 'agent_activity' + ? { agentActivityObservedAt: '2026-05-26T08:03:00.000Z' } + : {}), + terminalAt, + failureStage, + failureCode, + }) + ); + + expect(report.run).toMatchObject({ failureStage, failureCode }); + } + } + ); + + it('accepts payment_required before dispatch without acceptance or activity timestamps', () => { + const report = CloudAgentQueueReportSchema.parse( + reportWithRun({ + status: 'failed', + terminalAt, + failureStage: 'pre_dispatch', + failureCode: 'payment_required', + }) + ); + + expect(report.run).toMatchObject({ + failureStage: 'pre_dispatch', + failureCode: 'payment_required', + }); + expect(report.run).not.toHaveProperty('dispatchAcceptedAt'); + expect(report.run).not.toHaveProperty('agentActivityObservedAt'); + }); + it('rejects removed lifecycle fields and unsafe transport content', () => { const removedOrUnsafeFields = [ ['isInitialMessage', true], diff --git a/packages/worker-utils/src/cloud-agent-queue-report.ts b/packages/worker-utils/src/cloud-agent-queue-report.ts index 81595ca45a..a0d44e5bc1 100644 --- a/packages/worker-utils/src/cloud-agent-queue-report.ts +++ b/packages/worker-utils/src/cloud-agent-queue-report.ts @@ -23,6 +23,7 @@ export const CloudAgentRunFailureClassifications = [ { failureStage: 'pre_dispatch', failureCode: 'invalid_delivery_request' }, { failureStage: 'pre_dispatch', failureCode: 'session_metadata_missing' }, { failureStage: 'pre_dispatch', failureCode: 'model_missing' }, + { failureStage: 'pre_dispatch', failureCode: 'payment_required' }, { failureStage: 'pre_dispatch', failureCode: 'delivery_failure_unknown' }, { failureStage: 'post_dispatch_no_activity', failureCode: 'wrapper_disconnected' }, { failureStage: 'post_dispatch_no_activity', failureCode: 'wrapper_no_output' }, @@ -35,6 +36,9 @@ export const CloudAgentRunFailureClassifications = [ { failureStage: 'agent_activity', failureCode: 'payment_required' }, { failureStage: 'agent_activity', failureCode: 'model_missing' }, { failureStage: 'agent_activity', failureCode: 'wrapper_error_after_activity' }, + { failureStage: 'agent_activity', failureCode: 'wrapper_disconnected' }, + { failureStage: 'agent_activity', failureCode: 'wrapper_no_output' }, + { failureStage: 'agent_activity', failureCode: 'wrapper_ping_timeout' }, { failureStage: 'interruption', failureCode: 'user_interrupt' }, { failureStage: 'interruption', failureCode: 'container_shutdown' }, { failureStage: 'interruption', failureCode: 'system_interrupt' }, diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index 77a1b1679b..79a8267662 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -108,6 +108,7 @@ import { } from '../session/pending-messages.js'; import { createSessionMessageQueue, + enqueuePendingSessionMessageIntent, PENDING_FLUSH_DEBOUNCE_MS, type SessionMessageQueue, } from '../session/session-message-queue.js'; @@ -127,6 +128,7 @@ import { type KiloGlobalFeedValidationResult, } from '../session/wrapper-global-feed-validation.js'; import { + createQueuedSessionMessageState, getSessionMessageState, listNonTerminalAcceptedMessages, markAgentActivityObserved, @@ -155,6 +157,7 @@ import { type WrapperTerminalEvent, } from '../session/wrapper-supervisor.js'; import { emitRunStateReport } from '../telemetry/queue-reports.js'; +import { ensureCloneSessionReport } from '../telemetry/session-reports.js'; import { createAgentSandbox, createAgentSandboxLifecycle } from '../agent-sandbox/factory.js'; import { isCloudAgentContainerBillingEnabled } from '../container-billing-rollout.js'; import type { @@ -645,6 +648,13 @@ export class CloudAgentSession extends DurableObject { try { const sessionId = await this.resolveSessionId(); if (!sessionId) return; + try { + await ensureCloneSessionReport(await this.getMetadata(), this.env); + } catch { + logger + .withFields({ sessionId, messageId: state.messageId }) + .warn('Cloud Agent clone report anchor skipped'); + } await emitRunStateReport({ queue: { send: report => this.sendRunStateReport(report) }, cloudAgentSessionId: sessionId, @@ -1146,6 +1156,31 @@ export class CloudAgentSession extends DurableObject { reportQueuedState: state => { this.ctx.waitUntil(this.reportRunState(state)); }, + persistCloneQueuedMessage: (intent, callbackSnapshot) => + this.ctx.storage.transaction(async () => { + const metadata = await this.getMetadata(); + if (!metadata) throw new Error('Session metadata unavailable'); + const now = Date.now(); + await enqueuePendingSessionMessageIntent( + this.ctx.storage, + intent, + now, + callbackSnapshot + ); + await putSessionMessageState( + this.ctx.storage, + createQueuedSessionMessageState(intent, callbackSnapshot, now) + ); + if (metadata.clone?.reportingCreatedAt && !metadata.initialMessage?.id) { + await this.ctx.storage.put( + 'metadata', + serializeSessionMetadata({ + ...metadata, + initialMessage: { id: intent.turn.messageId }, + }) + ); + } + }), ensureAcceptedMessageEffects: messageId => this.ensureAcceptedMessageEffects(messageId), persistTerminalTransition: (messageId, params, options) => this.getMessageSettlementOutbox().persistTerminalTransition(messageId, params, options), @@ -1726,6 +1761,12 @@ export class CloudAgentSession extends DurableObject { const newMetadata = serializeSessionMetadata(parseSessionMetadata(data)); const existingMetadata = await this.getMetadata(); if (existingMetadata) { + if (existingMetadata.clone?.reportingCreatedAt !== newMetadata.clone?.reportingCreatedAt) { + throw new Error('Clone reporting creation time cannot be changed'); + } + if (existingMetadata.clone?.reportingCreatedAt && existingMetadata.initialMessage?.id) { + newMetadata.initialMessage = existingMetadata.initialMessage; + } if (getSandboxProvider(existingMetadata) !== getSandboxProvider(newMetadata)) { throw new Error('Registered sandbox provider cannot be changed'); } @@ -2086,6 +2127,8 @@ export class CloudAgentSession extends DurableObject { timestamp: Date.now(), }); } + const canceledState = await getSessionMessageState(this.ctx.storage, messageId); + if (canceledState) this.ctx.waitUntil(this.reportRunState(canceledState)); return { dropped: true }; } diff --git a/services/cloud-agent-next/src/persistence/session-metadata.test.ts b/services/cloud-agent-next/src/persistence/session-metadata.test.ts index 58eceab7f3..bec3806933 100644 --- a/services/cloud-agent-next/src/persistence/session-metadata.test.ts +++ b/services/cloud-agent-next/src/persistence/session-metadata.test.ts @@ -166,21 +166,40 @@ describe('session metadata boundary', () => { expect(parseSessionMetadata(current).repository).not.toHaveProperty('githubIntegrationId'); }); - it('parses and serializes clone source metadata', () => { - const current = { - metadataSchemaVersion: 2 as const, - identity: { sessionId: 'agent_clone', userId: 'user_clone' }, - auth: {}, - clone: { - cloneFromKiloSessionId: 'ses_aaaaaaaaaaaaaaaaaaaaaaaaaa', - }, - lifecycle: { version: 1, timestamp: 1 }, - }; - - expect(parseSessionMetadata(current)).toEqual(current); - expect(serializeSessionMetadata(current)).toEqual(current); - expect(CurrentSessionMetadataSchema.parse(current)).toEqual(current); - }); + it.each([undefined, '2026-08-29T10:00:00.000Z'])( + 'round-trips clone metadata without inventing a reporting creation time (%s)', + reportingCreatedAt => { + const current = { + metadataSchemaVersion: 2 as const, + identity: { sessionId: 'agent_clone', userId: 'user_clone' }, + auth: {}, + clone: { + cloneFromKiloSessionId: 'ses_aaaaaaaaaaaaaaaaaaaaaaaaaa', + ...(reportingCreatedAt ? { reportingCreatedAt } : {}), + }, + lifecycle: { version: 1, timestamp: 1 }, + }; + + expect(parseSessionMetadata(current)).toEqual(current); + expect(serializeSessionMetadata(current)).toEqual(current); + expect(CurrentSessionMetadataSchema.parse(current)).toEqual(current); + } + ); + + it.each(['invalid', '2026-08-29', '2026-02-30T10:00:00.000Z', 1_700_000_000_000, null])( + 'rejects an invalid clone reporting creation time (%s)', + reportingCreatedAt => { + expect(() => + parseSessionMetadata({ + metadataSchemaVersion: 2, + identity: { sessionId: 'agent_clone', userId: 'user_clone' }, + auth: {}, + clone: { cloneFromKiloSessionId: 'ses_aaaaaaaaaaaaaaaaaaaaaaaaaa', reportingCreatedAt }, + lifecycle: { version: 1, timestamp: 1 }, + }) + ).toThrow('Invalid current session metadata'); + } + ); it('keeps metadata without clone as an empty-session bootstrap', () => { const current = { diff --git a/services/cloud-agent-next/src/persistence/session-metadata.ts b/services/cloud-agent-next/src/persistence/session-metadata.ts index 2cf4cadd24..93b87402ca 100644 --- a/services/cloud-agent-next/src/persistence/session-metadata.ts +++ b/services/cloud-agent-next/src/persistence/session-metadata.ts @@ -315,6 +315,7 @@ const MetadataLifecycleSchema = z const MetadataCloneSchema = z .object({ cloneFromKiloSessionId: kiloSessionIdSchema, + reportingCreatedAt: z.string().datetime({ offset: true }).optional(), }) .strip(); diff --git a/services/cloud-agent-next/src/services/git-token-service-client.test.ts b/services/cloud-agent-next/src/services/git-token-service-client.test.ts index 32b37ddb20..3e6d171805 100644 --- a/services/cloud-agent-next/src/services/git-token-service-client.test.ts +++ b/services/cloud-agent-next/src/services/git-token-service-client.test.ts @@ -1,10 +1,11 @@ -import { describe, expect, it, vi } from 'vitest'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import { logger } from '../logger.js'; import type { GitTokenService } from '../types.js'; import { issueCloudAgentGitHubSessionCapability, issueCloudAgentGitLabSessionCapability, resolveCloudAgentGitHubAuthForRepo, + resolveGitHubTokenForRepo, resolveManagedBitbucketToken, resolveManagedGitLabToken, } from './git-token-service-client.js'; @@ -192,7 +193,7 @@ describe('resolveManagedGitLabToken', () => { }); describe('issueCloudAgentGitHubSessionCapability', () => { - it('falls back to installation authentication when the capability RPC is not deployed yet', async () => { + it('fails closed without raw authentication when the capability RPC is not deployed yet', async () => { const getTokenForRepo = vi.fn().mockResolvedValue({ success: true, token: 'installation-token', @@ -200,23 +201,25 @@ describe('issueCloudAgentGitHubSessionCapability', () => { accountLogin: 'acme', appType: 'standard', }); + const getCloudAgentAuthForRepo = vi.fn(); - const result = await issueCloudAgentGitHubSessionCapability(createEnv({ getTokenForRepo }), { - githubRepo: 'acme/repo', - userId: 'user_1', - outboundContainerId: 'container-test', - allowUserAuthorization: true, - }); + const result = await issueCloudAgentGitHubSessionCapability( + createEnv({ getTokenForRepo, getCloudAgentAuthForRepo }), + { + githubRepo: 'acme/repo', + userId: 'user_1', + outboundContainerId: 'container-test', + allowUserAuthorization: true, + } + ); - expect(getTokenForRepo).toHaveBeenCalledWith({ githubRepo: 'acme/repo', userId: 'user_1' }); + expect(getTokenForRepo).not.toHaveBeenCalled(); + expect(getCloudAgentAuthForRepo).not.toHaveBeenCalled(); expect(result).toEqual({ - success: true, - value: { - githubToken: 'installation-token', - installationId: '123', - accountLogin: 'acme', - appType: 'standard', - source: 'installation', + success: false, + error: { + reason: 'service_not_configured', + message: 'git-token-service capability issuance is not configured', }, }); }); @@ -321,7 +324,7 @@ describe('issueCloudAgentGitHubSessionCapability', () => { expect(getTokenForRepo).not.toHaveBeenCalled(); }); - it('falls back to direct authentication when the capability RPC rejects during rollout', async () => { + it('fails closed without direct authentication when the capability RPC rejects', async () => { const issueGitHubSessionCapability = vi .fn() .mockRejectedValue(new Error('service unavailable')); @@ -347,25 +350,84 @@ describe('issueCloudAgentGitHubSessionCapability', () => { ); expect(result).toEqual({ - success: true, - value: { - githubToken: 'user-token', - installationId: '123', - accountLogin: 'acme', - appType: 'standard', - source: 'user', - gitAuthor: { name: 'octocat', email: '101+octocat@users.noreply.github.com' }, - }, - }); - expect(getCloudAgentAuthForRepo).toHaveBeenCalledWith({ - githubRepo: 'acme/repo', - userId: 'user_1', - allowUserAuthorization: true, + success: false, + error: { reason: 'rpc_error', message: 'GitHub credential service is unavailable' }, }); + expect(getCloudAgentAuthForRepo).not.toHaveBeenCalled(); expect(getTokenForRepo).not.toHaveBeenCalled(); }); }); +describe.each([ + { method: 'getTokenForRepo', resolve: resolveGitHubTokenForRepo }, + { method: 'getCloudAgentAuthForRepo', resolve: resolveCloudAgentGitHubAuthForRepo }, + { method: 'issueGitHubSessionCapability', resolve: issueCloudAgentGitHubSessionCapability }, +] as const)('GitHub credential client $method failures', ({ method, resolve }) => { + const params = { + githubRepo: 'acme/repo', + userId: 'user_1', + outboundContainerId: 'container-test', + allowUserAuthorization: false, + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it.each(['no_installation_found', 'repository_not_installed', 'integration_mismatch'] as const)( + 'preserves %s without trying another authorization path', + async reason => { + const service = { ...createGitTokenService(), getCloudAgentAuthForRepo: vi.fn() }; + service[method].mockResolvedValue({ success: false, reason }); + + const result = await resolve({ GIT_TOKEN_SERVICE: service }, params); + + expect(result).toMatchObject({ success: false, error: { reason } }); + for (const rpc of [ + 'getTokenForRepo', + 'getCloudAgentAuthForRepo', + 'issueGitHubSessionCapability', + ] as const) { + expect(service[rpc]).toHaveBeenCalledTimes(rpc === method ? 1 : 0); + } + } + ); + + it('keeps RPC exceptions distinct and excludes private errors from returned or logged data', async () => { + const secret = 'sensitive-credential-or-query-parameter'; + const logFields = vi.spyOn(logger, 'withFields'); + const service = { ...createGitTokenService(), getCloudAgentAuthForRepo: vi.fn() }; + service[method].mockRejectedValue( + Object.assign(new Error(`query or transport failed: ${secret}`), { + request: { headers: { authorization: `Bearer ${secret}` } }, + response: { data: { token: secret } }, + }) + ); + + const result = await resolve({ GIT_TOKEN_SERVICE: service }, params); + + expect(result).toEqual({ + success: false, + error: { reason: 'rpc_error', message: 'GitHub credential service is unavailable' }, + }); + for (const rpc of [ + 'getTokenForRepo', + 'getCloudAgentAuthForRepo', + 'issueGitHubSessionCapability', + ] as const) { + expect(service[rpc]).toHaveBeenCalledTimes(rpc === method ? 1 : 0); + } + expect( + JSON.stringify({ + result, + fields: logFields.mock.calls, + errors: vi.mocked(logger.error).mock.calls, + warnings: vi.mocked(logger.warn).mock.calls, + }) + ).not.toContain(secret); + }); +}); + describe('issueCloudAgentGitLabSessionCapability', () => { it('returns an opaque code-review project capability and preserves CLI mode metadata', async () => { const issueGitLabSessionCapability = vi.fn().mockResolvedValue({ @@ -527,7 +589,7 @@ describe('resolveCloudAgentGitHubAuthForRepo', () => { }); }); - it('falls back to installation authentication when an older service rejects the managed RPC', async () => { + it('does not fall back to a different authorization path when the managed RPC rejects', async () => { const getCloudAgentAuthForRepo = vi .fn() .mockRejectedValue(new Error('RPC method getCloudAgentAuthForRepo is not available')); @@ -553,19 +615,14 @@ describe('resolveCloudAgentGitHubAuthForRepo', () => { userId: 'user_1', allowUserAuthorization: true, }); - expect(getTokenForRepo).toHaveBeenCalledWith({ githubRepo: 'acme/repo', userId: 'user_1' }); - expect(result).toMatchObject({ - success: true, - value: { - githubToken: 'installation-token', - installationId: '123', - appType: 'standard', - source: 'installation', - }, + expect(getTokenForRepo).not.toHaveBeenCalled(); + expect(result).toEqual({ + success: false, + error: { reason: 'rpc_error', message: 'GitHub credential service is unavailable' }, }); }); - it('preserves the expected integration id through direct and legacy fallbacks', async () => { + it('preserves the expected integration id through direct and legacy authentication', async () => { const getCloudAgentAuthForRepo = vi .fn() .mockRejectedValue(new Error('RPC method getCloudAgentAuthForRepo is not available')); @@ -594,6 +651,16 @@ describe('resolveCloudAgentGitHubAuthForRepo', () => { expectedIntegrationId, allowUserAuthorization: false, }); + expect(getTokenForRepo).not.toHaveBeenCalled(); + + await expect( + resolveCloudAgentGitHubAuthForRepo(createEnv({ getTokenForRepo }), { + githubRepo: 'acme/repo', + userId: 'user_1', + expectedIntegrationId, + allowUserAuthorization: false, + }) + ).resolves.toMatchObject({ success: true, value: { source: 'installation' } }); expect(getTokenForRepo).toHaveBeenCalledWith({ githubRepo: 'acme/repo', userId: 'user_1', diff --git a/services/cloud-agent-next/src/services/git-token-service-client.ts b/services/cloud-agent-next/src/services/git-token-service-client.ts index e21d84a8d3..d9d5161e8b 100644 --- a/services/cloud-agent-next/src/services/git-token-service-client.ts +++ b/services/cloud-agent-next/src/services/git-token-service-client.ts @@ -70,12 +70,11 @@ export async function resolveGitHubTokenForRepo( message: `GitHub token lookup failed (${result.reason})`, }, }; - } catch (error) { - const message = error instanceof Error ? error.message : String(error); - logger.withFields({ error: message }).error('Failed to call git-token-service getTokenForRepo'); + } catch { + logger.error('Failed to call git-token-service getTokenForRepo'); return { success: false, - error: { reason: 'rpc_error', message: `git-token-service RPC failed: ${message}` }, + error: { reason: 'rpc_error', message: 'GitHub credential service is unavailable' }, }; } } @@ -112,7 +111,7 @@ type IssueCloudAgentGitHubSessionCapabilityParams = { }; type IssueCloudAgentGitHubSessionCapabilityResult = - | { success: true; value: ResolvedCloudAgentGitHubCapability | ResolvedCloudAgentGitHubAuth } + | { success: true; value: ResolvedCloudAgentGitHubCapability } | { success: false; error: ResolveGitHubTokenError }; type CloudAgentGitHubAuthResult = @@ -201,35 +200,20 @@ export async function resolveCloudAgentGitHubAuthForRepo( ...(result.fallbackReason ? { fallbackReason: result.fallbackReason } : {}), }, }; - } catch (error) { - const message = error instanceof Error ? error.message : String(error); - logger - .withFields({ error: message }) - .warn('Managed GitHub auth RPC unavailable; using installation authentication fallback'); - return resolveLegacyInstallationAuthForRepo(env, params); + } catch { + logger.error('Failed to call git-token-service getCloudAgentAuthForRepo'); + return { + success: false, + error: { reason: 'rpc_error', message: 'GitHub credential service is unavailable' }, + }; } } -function resolveGitHubAuthFallbackForCapability( - env: GitTokenServiceEnv, - params: IssueCloudAgentGitHubSessionCapabilityParams -): Promise { - return resolveCloudAgentGitHubAuthForRepo(env, { - githubRepo: params.githubRepo, - userId: params.userId, - ...(params.orgId !== undefined ? { orgId: params.orgId } : {}), - ...(params.expectedIntegrationId !== undefined - ? { expectedIntegrationId: params.expectedIntegrationId } - : {}), - allowUserAuthorization: params.allowUserAuthorization, - }); -} - export async function issueCloudAgentGitHubSessionCapability( env: GitTokenServiceEnv, params: IssueCloudAgentGitHubSessionCapabilityParams ): Promise { - if (!env.GIT_TOKEN_SERVICE) { + if (typeof env.GIT_TOKEN_SERVICE?.issueGitHubSessionCapability !== 'function') { return { success: false, error: { @@ -238,10 +222,6 @@ export async function issueCloudAgentGitHubSessionCapability( }, }; } - if (typeof env.GIT_TOKEN_SERVICE.issueGitHubSessionCapability !== 'function') { - logger.warn('Managed GitHub capability RPC unavailable; using direct authentication fallback'); - return resolveGitHubAuthFallbackForCapability(env, params); - } try { const result = await env.GIT_TOKEN_SERVICE.issueGitHubSessionCapability(params); @@ -276,12 +256,12 @@ export async function issueCloudAgentGitHubSessionCapability( ...(result.fallbackReason ? { fallbackReason: result.fallbackReason } : {}), }, }; - } catch (error) { - const message = error instanceof Error ? error.message : String(error); - logger - .withFields({ error: message }) - .warn('Managed GitHub capability RPC unavailable; using direct authentication fallback'); - return resolveGitHubAuthFallbackForCapability(env, params); + } catch { + logger.error('Failed to call git-token-service issueGitHubSessionCapability'); + return { + success: false, + error: { reason: 'rpc_error', message: 'GitHub credential service is unavailable' }, + }; } } diff --git a/services/cloud-agent-next/src/session-service.test.ts b/services/cloud-agent-next/src/session-service.test.ts index a74b452cab..bb0f7ac3b4 100644 --- a/services/cloud-agent-next/src/session-service.test.ts +++ b/services/cloud-agent-next/src/session-service.test.ts @@ -1,9 +1,15 @@ import { dirname, relative } from 'node:path'; -import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { beforeEach, describe, expect, it, vi, type Mock } from 'vitest'; import type * as DevContainerModule from './kilo/devcontainer.js'; import type * as GitTokenServiceClientModule from './services/git-token-service-client.js'; import { validateWrapperDispatchTicket } from './auth.js'; import { deriveKiloSandboxTargets } from './kilo/kilo-targets.js'; +import { ExecutionError } from './execution/errors.js'; +import { + createPendingSessionMessage, + recordPendingFlushFailure, + type SessionQueueStorage, +} from './session/pending-messages.js'; vi.mock('./logger.js', () => ({ logger: { @@ -591,6 +597,297 @@ describe('SessionService.resolveWorkspaceTokens', () => { }); }); + describe.each([false, true])('GitHub credential containment=%s', contained => { + const metadata = createMetadata({ + githubRepo: 'acme/repo', + gitUrl: undefined, + gitToken: undefined, + platform: 'github', + credentialContainment: { github: contained, gitlab: false, kilocode: false }, + }); + + beforeEach(async () => { + const client = await vi.importActual( + './services/git-token-service-client.js' + ); + tokenMocks.resolveCloudAgentGitHubAuthForRepo.mockImplementation( + client.resolveCloudAgentGitHubAuthForRepo + ); + tokenMocks.issueCloudAgentGitHubSessionCapability.mockImplementation( + client.issueCloudAgentGitHubSessionCapability + ); + }); + + function createGitHubEnv(rpc: Mock): PersistenceEnv { + const env = createEnv(); + if (!env.GIT_TOKEN_SERVICE) throw new Error('Expected git-token-service fixture'); + if (contained) env.GIT_TOKEN_SERVICE.issueGitHubSessionCapability = rpc; + else env.GIT_TOKEN_SERVICE.getCloudAgentAuthForRepo = rpc; + return env; + } + + function createQueueStorage(): SessionQueueStorage { + const entries = new Map(); + return { + async get(key: string) { + return entries.get(key) as T | undefined; + }, + async put(key, value) { + entries.set(key, value); + }, + async delete(keys) { + for (const key of typeof keys === 'string' ? [keys] : keys) entries.delete(key); + }, + async list({ prefix }: { prefix: string }) { + return new Map([...entries].filter(([key]) => key.startsWith(prefix)) as [string, T][]); + }, + }; + } + + it.each(['no_installation_found', 'repository_not_installed'] as const)( + 'terminalizes %s on the first preparation attempt', + async reason => { + const rpc = vi.fn().mockResolvedValue({ success: false, reason }); + const error = await new SessionService() + .resolveWorkspaceTokens(createGitHubEnv(rpc), metadata, 'ses-abcdef' as SandboxId) + .catch((error: unknown) => error); + + expect(error).toMatchObject({ + code: 'WORKSPACE_SETUP_FAILED', + workspaceFailureSubtype: 'git_authentication_failed', + retryable: false, + message: + 'GitHub repository authentication failed. Check that the GitHub App is installed and has access to this repository.', + }); + if (!(error instanceof ExecutionError) || error.code !== 'WORKSPACE_SETUP_FAILED') { + throw new Error('Expected a workspace setup error'); + } + expect(error.safeFailureMessage).toBe(error.message); + const result = await recordPendingFlushFailure( + createQueueStorage(), + createPendingSessionMessage({ + messageId: 'msg_018f1e2d3c4bGitHubFailureA', + role: 'user', + content: 'prepare', + createdAt: 1, + }), + error.message, + 100_000, + { + policy: 'cold-init', + code: error.code, + subtype: error.workspaceFailureSubtype, + safeFailureMessage: error.safeFailureMessage, + retryable: error.retryable, + } + ); + expect(result).toMatchObject({ + attempts: 1, + exhausted: true, + nextFlushAttemptAt: undefined, + message: { + lastFlushFailureSubtype: 'git_authentication_failed', + deliveryDisposition: 'terminalization-pending', + }, + }); + expect(rpc).toHaveBeenCalledOnce(); + } + ); + + it.each([ + new Error('database query failed with secret-query-parameter'), + new TypeError('transport failed with secret-credential'), + ])('uses the existing bounded preparation retry for an RPC exception', async rpcError => { + const rpc = vi.fn().mockRejectedValue(rpcError); + const env = createGitHubEnv(rpc); + const error = await new SessionService() + .resolveWorkspaceTokens(env, metadata, 'ses-abcdef' as SandboxId) + .catch((error: unknown) => error); + + expect(error).toMatchObject({ + code: 'WORKSPACE_SETUP_FAILED', + retryable: true, + workspaceFailureSubtype: undefined, + message: 'GitHub credential service is unavailable. Please try again.', + }); + if (!(error instanceof ExecutionError) || error.code !== 'WORKSPACE_SETUP_FAILED') { + throw new Error('Expected a workspace setup error'); + } + expect(error.safeFailureMessage).toBe(error.message); + expect(JSON.stringify({ error, message: error.message, cause: error.cause })).not.toContain( + 'secret-' + ); + const storage = createQueueStorage(); + const options = { + policy: 'cold-init', + code: error.code, + subtype: error.workspaceFailureSubtype, + safeFailureMessage: error.safeFailureMessage, + retryable: error.retryable, + } as const; + const first = await recordPendingFlushFailure( + storage, + createPendingSessionMessage({ + messageId: 'msg_018f1e2d3c4bGitHubFailureA', + role: 'user', + content: 'prepare', + createdAt: 1, + }), + error.message, + 100_000, + options + ); + expect(first).toMatchObject({ + attempts: 1, + exhausted: false, + nextFlushAttemptAt: 102_000, + }); + const second = await recordPendingFlushFailure( + storage, + first.message, + error.message, + 102_000, + options + ); + expect(second).toMatchObject({ + attempts: 2, + exhausted: true, + nextFlushAttemptAt: undefined, + }); + expect(rpc).toHaveBeenCalledOnce(); + expect( + contained + ? tokenMocks.resolveCloudAgentGitHubAuthForRepo + : tokenMocks.issueCloudAgentGitHubSessionCapability + ).not.toHaveBeenCalled(); + }); + + it.each(['invalid_repo_format', 'invalid_org_id', 'integration_mismatch'] as const)( + 'keeps %s as a permanent validation or authorization rejection', + async reason => { + const rpc = vi.fn().mockResolvedValue({ success: false, reason }); + + await expect( + new SessionService().resolveWorkspaceTokens( + createGitHubEnv(rpc), + metadata, + 'ses-abcdef' as SandboxId + ) + ).rejects.toMatchObject({ + code: 'INVALID_REQUEST', + retryable: false, + workspaceFailureSubtype: undefined, + message: `GitHub repository authorization failed (${reason})`, + }); + expect(rpc).toHaveBeenCalledOnce(); + } + ); + + it('keeps an unconfigured database retryable without installation guidance', async () => { + const rpc = vi.fn().mockResolvedValue({ success: false, reason: 'database_not_configured' }); + + await expect( + new SessionService().resolveWorkspaceTokens( + createGitHubEnv(rpc), + metadata, + 'ses-abcdef' as SandboxId + ) + ).rejects.toMatchObject({ + code: 'WORKSPACE_SETUP_FAILED', + retryable: true, + message: 'GitHub credential service is unavailable. Please try again.', + }); + }); + + it('keeps an unavailable binding retryable without falling back to raw credentials', async () => { + const env = createEnv(); + delete env.GIT_TOKEN_SERVICE; + + await expect( + new SessionService().resolveWorkspaceTokens(env, metadata, 'ses-abcdef' as SandboxId) + ).rejects.toMatchObject({ + code: 'WORKSPACE_SETUP_FAILED', + retryable: true, + message: 'GitHub credential service is unavailable. Please try again.', + }); + expect( + contained + ? tokenMocks.resolveCloudAgentGitHubAuthForRepo + : tokenMocks.issueCloudAgentGitHubSessionCapability + ).not.toHaveBeenCalled(); + }); + + it('does not infer an access failure or expose details from an unknown reason', async () => { + const rpc = vi.fn().mockResolvedValue({ + success: false, + reason: 'unknown-secret-credential', + }); + + await expect( + new SessionService().resolveWorkspaceTokens( + createGitHubEnv(rpc), + metadata, + 'ses-abcdef' as SandboxId + ) + ).rejects.toMatchObject({ + code: 'WORKSPACE_SETUP_FAILED', + retryable: true, + workspaceFailureSubtype: undefined, + message: 'GitHub credential resolution failed. Please try again.', + safeFailureMessage: 'GitHub credential resolution failed. Please try again.', + cause: undefined, + }); + }); + }); + + it('preserves a permanent capability configuration rejection without a raw token fallback', async () => { + tokenMocks.issueCloudAgentGitHubSessionCapability.mockResolvedValueOnce({ + success: false, + error: { reason: 'capability_configuration_error', message: 'private-key-details' }, + }); + + await expect( + new SessionService().resolveWorkspaceTokens( + createEnv(), + createMetadata({ + githubRepo: 'acme/repo', + gitUrl: undefined, + gitToken: undefined, + platform: 'github', + }), + 'ses-abcdef' as SandboxId + ) + ).rejects.toMatchObject({ + code: 'INVALID_REQUEST', + retryable: false, + message: 'GitHub repository authorization failed (capability_configuration_error)', + cause: undefined, + }); + expect(tokenMocks.resolveCloudAgentGitHubAuthForRepo).not.toHaveBeenCalled(); + }); + + it('rejects a DIND sandbox before attempting contained GitHub credential issuance', async () => { + await expect( + new SessionService().resolveWorkspaceTokens( + createEnv(), + createMetadata({ + githubRepo: 'acme/repo', + gitUrl: undefined, + gitToken: undefined, + platform: 'github', + credentialContainment: { github: true, gitlab: false, kilocode: false }, + }), + 'dind-test' as SandboxId + ) + ).rejects.toMatchObject({ + code: 'INVALID_REQUEST', + retryable: false, + message: 'Managed SCM containment is not supported for DIND sandboxes', + }); + expect(tokenMocks.issueCloudAgentGitHubSessionCapability).not.toHaveBeenCalled(); + expect(tokenMocks.resolveCloudAgentGitHubAuthForRepo).not.toHaveBeenCalled(); + }); + it('fails closed for replayed Bitbucket metadata without an organization', async () => { await expect( new SessionService().resolveWorkspaceTokens( @@ -1964,7 +2261,7 @@ describe('SessionService.prepareWorkspace', () => { platform: 'github', }), }) - ).rejects.toThrow('GitHub token or active app installation required'); + ).rejects.toThrow('GitHub credential service is unavailable. Please try again.'); expect(tokenMocks.resolveCloudAgentGitHubAuthForRepo).not.toHaveBeenCalled(); }); @@ -2394,7 +2691,7 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { platform: 'github', }) ) - ).rejects.toThrow('GitHub token or active app installation required'); + ).rejects.toThrow('GitHub credential service is unavailable. Please try again.'); expect(tokenMocks.resolveCloudAgentGitHubAuthForRepo).not.toHaveBeenCalled(); }); @@ -2975,7 +3272,7 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { }); await expect(buildPromptWrapperRequests(metadata)).rejects.toThrow( - 'GitHub token or active app installation required' + 'GitHub credential service is unavailable. Please try again.' ); expect(tokenMocks.issueCloudAgentGitHubSessionCapability).toHaveBeenCalled(); diff --git a/services/cloud-agent-next/src/session-service.ts b/services/cloud-agent-next/src/session-service.ts index c71e16b239..2d874f151c 100644 --- a/services/cloud-agent-next/src/session-service.ts +++ b/services/cloud-agent-next/src/session-service.ts @@ -1748,9 +1748,39 @@ export class SessionService { }) : await resolveCloudAgentGitHubAuthForRepo(env, authParams); if (!result.success) { - throw ExecutionError.invalidRequest( - `GitHub token or active app installation required for this repository (${result.error.reason})` - ); + switch (result.error.reason) { + case 'no_installation_found': + case 'repository_not_installed': { + const message = + 'GitHub repository authentication failed. Check that the GitHub App is installed and has access to this repository.'; + throw ExecutionError.workspaceSetupFailed(message, undefined, { + subtype: 'git_authentication_failed', + safeFailureMessage: message, + retryable: false, + }); + } + case 'rpc_error': + case 'service_not_configured': + case 'database_not_configured': { + const message = 'GitHub credential service is unavailable. Please try again.'; + throw ExecutionError.workspaceSetupFailed(message, undefined, { + safeFailureMessage: message, + }); + } + case 'invalid_repo_format': + case 'invalid_org_id': + case 'integration_mismatch': + case 'capability_configuration_error': + throw ExecutionError.invalidRequest( + `GitHub repository authorization failed (${result.error.reason})` + ); + default: { + const message = 'GitHub credential resolution failed. Please try again.'; + throw ExecutionError.workspaceSetupFailed(message, undefined, { + safeFailureMessage: message, + }); + } + } } githubToken = 'capability' in result.value ? result.value.capability : result.value.githubToken; diff --git a/services/cloud-agent-next/src/session/safe-failure-projection.test.ts b/services/cloud-agent-next/src/session/safe-failure-projection.test.ts index 46132c078c..5e3489cd43 100644 --- a/services/cloud-agent-next/src/session/safe-failure-projection.test.ts +++ b/services/cloud-agent-next/src/session/safe-failure-projection.test.ts @@ -8,13 +8,121 @@ import { describe, expect, it } from 'vitest'; import { SAFE_FAILURE_MESSAGE_MAX_LENGTH, SafeFailureProjectionSchema, + assistantFailureMessage, classifyAssistantFailureMessage, classifyAssistantFailure, genericFailureMessage, isAssistantInterrupt, + projectSafeAssistantError, projectSafeFailure, } from './safe-failure-projection.js'; +describe('projectSafeAssistantError', () => { + it.each([ + 'Payment required: insufficient credits', + 'Unknown model', + 'Rate limit exceeded', + 'Provider timeout', + 'Provider authentication failed', + 'Invalid request', + 'Service unavailable', + 'Poolside: Tool calls cutoff by max_tokens', + 'Unrecognized failure', + ])('preserves classification and BYOK semantics through safe projection: %s', text => { + for (const message of [text, `[BYOK] ${text}`]) { + for (const error of [message, { name: 'APIError', data: { message } }]) { + const failure = classifyAssistantFailure(error); + const safeError = projectSafeAssistantError(error); + + expect(safeError).toBe( + failure.providerOwnership === 'byok' + ? `[BYOK] ${failure.safeMessage}` + : failure.safeMessage + ); + expect(classifyAssistantFailure(safeError)).toEqual(failure); + expect(projectSafeAssistantError(safeError)).toBe(safeError); + } + } + }); + + it.each(CLOUD_AGENT_ASSISTANT_FAILURE_REASONS)( + 'round-trips canonical %s messages without losing cause or ownership', + reason => { + const message = assistantFailureMessage(reason); + for (const prefix of ['', '[BYOK] ']) { + const safeError = projectSafeAssistantError(`${prefix}${message}`); + + expect(safeError).toBe(`${prefix}${message}`); + expect(projectSafeAssistantError(safeError)).toBe(safeError); + for (const providerOwnership of CLOUD_AGENT_PROVIDER_OWNERSHIPS) { + expect(classifyAssistantFailure(safeError, providerOwnership)).toMatchObject({ + reason, + safeMessage: message, + providerOwnership: prefix ? 'byok' : providerOwnership, + }); + } + } + } + ); + + it.each([ + ['APIError', 503, 'Assistant service is unavailable'], + ['APIError', undefined, 'Assistant request failed'], + ['FutureError', 402, 'Assistant request failed'], + ['ContextOverflowError', 402, 'The model context limit was exceeded'], + ] as const)('does not inspect or retain private fields from %s', (name, statusCode, expected) => { + const error = { + name, + message: 'outer poison-message', + data: { + message: 'Unrecognized failure token=poison-token', + statusCode, + responseBody: JSON.stringify({ + error: { message: '[BYOK] insufficient credits', statusCode: 402 }, + prompt: 'poison-prompt', + }), + responseHeaders: { + authorization: 'Bearer poison-header', + cookie: 'poison-cookie', + 'x-error': 'Unknown model', + }, + metadata: { message: 'Rate limit exceeded', token: 'poison-metadata' }, + }, + }; + const safeError = projectSafeAssistantError(error); + + expect(safeError).toBe(expected); + expect(safeError).not.toContain('poison'); + expect(classifyAssistantFailure(safeError).providerOwnership).toBe('unknown'); + }); + + it.each([ + [{ name: 'MessageAbortedError', data: { responseBody: 'poison-body' } }, ''], + ['user-interrupt token=poison-token', ''], + ['[BYOK] The message was interrupted by the user', '[BYOK] '], + [ + { + name: 'MessageAbortedError', + data: { message: '[BYOK] 401 token=poison-token', responseBody: 'poison-body' }, + }, + '[BYOK] ', + ], + ] as const)('preserves interrupt meaning as a safe string: %j', (error, prefix) => { + const safeError = projectSafeAssistantError(error); + + expect(safeError).toBe(`${prefix}The message was interrupted by the user`); + expect(isAssistantInterrupt(safeError)).toBe(true); + expect(classifyAssistantFailureMessage(safeError)).toBe( + 'The message was interrupted by the user' + ); + expect(projectSafeAssistantError(safeError)).toBe(safeError); + }); + + it.each([null, undefined])('omits absent errors: %s', error => { + expect(projectSafeAssistantError(error)).toBeUndefined(); + }); +}); + describe('projectSafeFailure', () => { it('projects structured fields while omitting raw failure text', () => { const durableState = { @@ -94,6 +202,12 @@ describe('projectSafeFailure', () => { expect(SafeFailureProjectionSchema.parse(failure)).toEqual(failure); }); + it('describes no-output failures as a lack of execution progress during the watchdog window', () => { + expect(projectSafeFailure({ failureCode: 'wrapper_no_output' })?.message).toBe( + 'Agent wrapper made no execution progress during the watchdog window' + ); + }); + it.each([ ['git_clone_timeout', 'Repository clone timed out'], ['git_authentication_failed', 'Repository authentication failed'], @@ -159,6 +273,28 @@ describe('projectSafeFailure', () => { }); }); +describe('assistantFailureMessage', () => { + it.each([ + ['insufficient_credits', 'Assistant request failed: insufficient credits'], + ['rate_limited', 'Assistant request was rate limited'], + ['model_unavailable', 'Assistant request failed: model not found'], + ['provider_authentication', 'Assistant request was not authorized'], + ['provider_unavailable', 'Assistant service is unavailable'], + ['timeout', 'Assistant request timed out'], + ['invalid_request', 'Assistant request was invalid'], + ['context_limit', 'The model context limit was exceeded'], + ['output_limit', 'The model output limit was reached'], + ['content_filter', 'The model provider blocked the response under its content policy'], + ['structured_output', 'The model response did not match the required format'], + ['unknown', 'Assistant request failed'], + ] as const)('returns bounded safe wording for %s', (reason, expected) => { + expect(assistantFailureMessage(reason)).toBe(expected); + expect(assistantFailureMessage(reason).length).toBeLessThanOrEqual( + SAFE_FAILURE_MESSAGE_MAX_LENGTH + ); + }); +}); + describe('classifyAssistantFailureMessage', () => { it.each([ ['Payment Required: token=secret', 'Assistant request failed: insufficient credits'], @@ -169,6 +305,11 @@ describe('classifyAssistantFailureMessage', () => { ['403 Forbidden: private policy', 'Assistant request was not authorized'], ['400 invalid request: prompt secret', 'Assistant request was invalid'], ['503 Service Unavailable: internal host', 'Assistant service is unavailable'], + [ + 'Poolside: Tool calls cutoff by max_tokens token=secret', + 'The model output limit was reached', + ], + ['Tool calls cut off by max_tokens: private body', 'The model output limit was reached'], ['provider exploded with token=secret', 'Assistant request failed'], ])('maps raw assistant text to allowlisted wording', (source, expected) => { const result = classifyAssistantFailureMessage(source); @@ -188,6 +329,158 @@ describe('classifyAssistantFailureMessage', () => { }); describe('classifyAssistantFailure', () => { + it.each([ + ['ContextOverflowError', 'context_limit', 'The model context limit was exceeded'], + ['MessageOutputLengthError', 'output_limit', 'The model output limit was reached'], + [ + 'ContentFilterError', + 'content_filter', + 'The model provider blocked the response under its content policy', + ], + [ + 'StructuredOutputError', + 'structured_output', + 'The model response did not match the required format', + ], + ['ProviderAuthError', 'provider_authentication', 'Assistant request was not authorized'], + ] as const)( + 'retains the precise %s cause through safe projection', + (name, reason, safeMessage) => { + expect(classifyAssistantFailure({ name })).toEqual({ + reason, + safeMessage, + providerOwnership: 'unknown', + }); + + for (const message of [ + 'Unrecognized failure', + '400 invalid request', + '503 Service unavailable', + ]) { + for (const prefix of ['', '[BYOK] ']) { + const error = { name, data: { message: `${prefix}${message}`, statusCode: 402 } }; + const failure = classifyAssistantFailure(error); + const safeError = projectSafeAssistantError(error); + + expect(failure).toEqual({ + reason, + safeMessage, + providerOwnership: prefix ? 'byok' : 'unknown', + }); + expect(safeError).toBe(`${prefix}${safeMessage}`); + expect(classifyAssistantFailure(safeError)).toEqual(failure); + } + } + } + ); + + it.each([ + [100, 'unknown', undefined], + [200, 'unknown', undefined], + [399, 'unknown', undefined], + [400, 'invalid_request', undefined], + [401, 'provider_authentication', undefined], + [402, 'insufficient_credits', 'payment_required'], + [403, 'provider_authentication', undefined], + [404, 'invalid_request', undefined], + [408, 'timeout', undefined], + [422, 'invalid_request', undefined], + [429, 'rate_limited', undefined], + [499, 'invalid_request', undefined], + [500, 'provider_unavailable', undefined], + [503, 'provider_unavailable', undefined], + [504, 'timeout', undefined], + [599, 'provider_unavailable', undefined], + ] as const)( + 'uses numeric APIError status %s only when no message cause is known', + (statusCode, reason, terminalCode) => { + for (const message of [undefined, 'Unrecognized failure token=poison-token']) { + const error = { name: 'APIError', data: { statusCode, message } }; + const failure = classifyAssistantFailure(error); + const safeError = projectSafeAssistantError(error); + + expect(failure).toEqual({ + reason, + safeMessage: assistantFailureMessage(reason), + providerOwnership: 'unknown', + ...(terminalCode === undefined ? {} : { terminalCode }), + }); + expect(safeError).toBe(assistantFailureMessage(reason)); + expect(classifyAssistantFailure(safeError)).toEqual(failure); + } + } + ); + + it.each([undefined, null, true, '402', '503', 99, 600, 503.5, NaN, Infinity, -Infinity, {}, []])( + 'ignores malformed APIError status %s without losing a known message cause', + statusCode => { + expect(classifyAssistantFailure({ name: 'APIError', data: { statusCode } })).toEqual({ + reason: 'unknown', + safeMessage: 'Assistant request failed', + providerOwnership: 'unknown', + }); + expect( + classifyAssistantFailure({ + name: 'APIError', + data: { statusCode, message: 'Rate limit exceeded' }, + }).reason + ).toBe('rate_limited'); + } + ); + + it.each([undefined, 'UnknownError', 'FutureError', 'apierror', 402])( + 'does not infer a status cause from unrecognized SDK kind %s', + name => { + for (const statusCode of [400, 401, 402, 403, 408, 429, 503, 504]) { + const error = { name, data: { statusCode } }; + + expect(classifyAssistantFailure(error)).toEqual({ + reason: 'unknown', + safeMessage: 'Assistant request failed', + providerOwnership: 'unknown', + }); + expect(projectSafeAssistantError(error)).toBe('Assistant request failed'); + } + } + ); + + it.each([ + false, + 0, + '', + [], + {}, + { name: 'APIError', data: null }, + { name: 'APIError', data: 402 }, + { name: 'APIError', data: '[BYOK] payment required' }, + { name: 'APIError', data: [{ message: '[BYOK] payment required', statusCode: 402 }] }, + { name: { name: 'ProviderAuthError' }, data: { statusCode: 401 } }, + { name: 'APIError', statusCode: 402, data: { message: 401 } }, + ])('falls back safely for malformed structured errors: %j', error => { + expect(classifyAssistantFailure(error)).toEqual({ + reason: 'unknown', + safeMessage: 'Assistant request failed', + providerOwnership: 'unknown', + }); + expect(projectSafeAssistantError(error)).toBe('Assistant request failed'); + }); + + it.each([null, 402, {}, ['Payment required']])( + 'falls back from non-string SDK message %j to a valid message or status', + message => { + const error = { name: 'APIError', data: { message, statusCode: 408 } }; + + expect(classifyAssistantFailure(error).reason).toBe('timeout'); + expect(projectSafeAssistantError(error)).toBe('Assistant request timed out'); + expect(classifyAssistantFailure({ ...error, message: 'Unknown model' })).toEqual({ + reason: 'model_unavailable', + safeMessage: 'Assistant request failed: model not found', + providerOwnership: 'unknown', + terminalCode: 'model_missing', + }); + } + ); + it('retains safe structured reason and explicit BYOK ownership without source text', () => { expect(classifyAssistantFailure('[BYOK] 401 token=secret')).toEqual({ reason: 'provider_authentication', @@ -207,6 +500,82 @@ describe('classifyAssistantFailure', () => { }); }); + it.each([ + { + source: '[BYOK] insufficient credits; unknown model; 429 timeout; token=poisoned-secret', + reason: 'insufficient_credits', + terminalCode: 'payment_required', + safeMessage: 'Assistant request failed: insufficient credits', + }, + { + source: '[BYOK] unknown model; 429 timeout; token=poisoned-secret', + reason: 'model_unavailable', + terminalCode: 'model_missing', + safeMessage: 'Assistant request failed: model not found', + }, + { + source: + '[BYOK] Poolside: Tool calls cutoff by max_tokens; 429 timeout; token=poisoned-secret', + reason: 'output_limit', + safeMessage: 'The model output limit was reached', + }, + { + source: '[BYOK] 429 Too Many Requests; timeout; token=poisoned-secret', + reason: 'rate_limited', + safeMessage: 'Assistant request was rate limited', + }, + { + source: '[BYOK] deadline exceeded; 403 Forbidden; token=poisoned-secret', + reason: 'timeout', + safeMessage: 'Assistant request timed out', + }, + { + source: '[BYOK] 401 Unauthorized; token=poisoned-secret', + reason: 'provider_authentication', + safeMessage: 'Assistant request was not authorized', + }, + ])( + 'preserves $reason message precedence over SDK kinds and status', + ({ source, ...expected }) => { + const errors = [ + source, + ...[ + 'APIError', + 'ContextOverflowError', + 'MessageOutputLengthError', + 'ContentFilterError', + 'StructuredOutputError', + 'ProviderAuthError', + ].map(name => ({ name, data: { message: source, statusCode: 500 } })), + ]; + for (const error of errors) { + const failure = classifyAssistantFailure(error, 'managed'); + + expect(failure).toEqual({ ...expected, providerOwnership: 'byok' }); + expect(JSON.stringify(failure)).not.toContain('poisoned-secret'); + expect(classifyAssistantFailure(projectSafeAssistantError(error), 'managed')).toEqual( + failure + ); + } + } + ); + + it.each([ + ['400 invalid request', 503, 'invalid_request'], + ['503 Service unavailable', 400, 'provider_unavailable'], + ] as const)( + 'keeps the known message %s ahead of APIError status %s', + (message, statusCode, reason) => { + expect(classifyAssistantFailure({ name: 'APIError', data: { message, statusCode } })).toEqual( + { + reason, + safeMessage: assistantFailureMessage(reason), + providerOwnership: 'unknown', + } + ); + } + ); + it('does not guess ownership for an unmarked provider outage', () => { expect(classifyAssistantFailure('503 Service Unavailable')).toMatchObject({ reason: 'provider_unavailable', diff --git a/services/cloud-agent-next/src/session/safe-failure-projection.ts b/services/cloud-agent-next/src/session/safe-failure-projection.ts index d9b6a0a5ab..a4e8c5cf6d 100644 --- a/services/cloud-agent-next/src/session/safe-failure-projection.ts +++ b/services/cloud-agent-next/src/session/safe-failure-projection.ts @@ -12,6 +12,15 @@ import type { SessionMessageFailureStage, } from './session-message-state.js'; +export { + assistantFailureMessage, + classifyAssistantFailure, + classifyAssistantFailureMessage, + isAssistantInterrupt, + projectSafeAssistantError, + type AssistantFailureClassification, +} from '../shared/assistant-failure.js'; + export const SAFE_FAILURE_MESSAGE_MAX_LENGTH = CLOUD_AGENT_SAFE_FAILURE_MESSAGE_MAX_LENGTH; export const SafeFailureProjectionSchema = CloudAgentSafeFailureSchema; export type SafeFailureProjection = CloudAgentSafeFailure; @@ -43,7 +52,7 @@ const GENERIC_FAILURE_MESSAGES = { model_missing: 'No model was selected', delivery_failure_unknown: 'The message could not be delivered', wrapper_disconnected: 'Agent wrapper disconnected', - wrapper_no_output: 'Agent wrapper produced no output', + wrapper_no_output: 'Agent wrapper made no execution progress during the watchdog window', wrapper_ping_timeout: 'Agent wrapper stopped responding', wrapper_error_before_activity: 'Agent wrapper failed before processing the message', assistant_error: 'Assistant request failed', @@ -81,105 +90,6 @@ export function workspaceFailureMessage(subtype: WorkspaceFailureSubtype): strin return WORKSPACE_FAILURE_MESSAGES[subtype]; } -export type AssistantFailureClassification = { - reason: CloudAgentAssistantFailureReason; - safeMessage: string; - providerOwnership: CloudAgentProviderOwnership; - terminalCode?: 'payment_required' | 'model_missing'; -}; - -export function isAssistantInterrupt(source: unknown): boolean { - if (typeof source === 'object' && source !== null && 'name' in source) { - if (source.name === 'MessageAbortedError') return true; - } - return /messageabortederror|user[_ -]?interrupt|interrupted by the user/.test( - extractErrorMessage(source).toLocaleLowerCase() - ); -} - -export function classifyAssistantFailure( - source: unknown, - defaultProviderOwnership: CloudAgentProviderOwnership = 'unknown' -): AssistantFailureClassification { - const message = extractErrorMessage(source).toLocaleLowerCase(); - const providerOwnership = /\[byok\]/i.test(message) ? 'byok' : defaultProviderOwnership; - if (/\b(payment required|insufficient (?:credits?|balance|funds))\b/.test(message)) { - return { - reason: 'insufficient_credits', - safeMessage: 'Assistant request failed: insufficient credits', - providerOwnership, - terminalCode: 'payment_required', - }; - } - if (/\b(model (?:was )?not found|unknown model|invalid model)\b/.test(message)) { - return { - reason: 'model_unavailable', - safeMessage: 'Assistant request failed: model not found', - providerOwnership, - terminalCode: 'model_missing', - }; - } - if ( - /\b(rate limit|rate_limit|usage[_ -]?limit[_ -]?exceeded|too many requests|429)\b/.test(message) - ) { - return { - reason: 'rate_limited', - safeMessage: 'Assistant request was rate limited', - providerOwnership, - }; - } - if (/\b(timed? out|timeout|deadline exceeded)\b/.test(message)) { - return { - reason: 'timeout', - safeMessage: 'Assistant request timed out', - providerOwnership, - }; - } - if (/\b(unauthorized|forbidden|authorization|authentication|401|403)\b/.test(message)) { - return { - reason: 'provider_authentication', - safeMessage: 'Assistant request was not authorized', - providerOwnership, - }; - } - if (/\b(invalid request|bad request|malformed request|400)\b/.test(message)) { - return { - reason: 'invalid_request', - safeMessage: 'Assistant request was invalid', - providerOwnership, - }; - } - if (/\b(service unavailable|temporarily unavailable|overloaded|502|503|504)\b/.test(message)) { - return { - reason: 'provider_unavailable', - safeMessage: 'Assistant service is unavailable', - providerOwnership, - }; - } - return { - reason: 'unknown', - safeMessage: GENERIC_FAILURE_MESSAGES.assistant_error, - providerOwnership, - }; -} - -export function classifyAssistantFailureMessage(source: unknown): string { - if (isAssistantInterrupt(source)) return GENERIC_FAILURE_MESSAGES.user_interrupt; - return classifyAssistantFailure(source).safeMessage; -} - -function extractErrorMessage(source: unknown): string { - if (typeof source === 'string') return source; - if (typeof source !== 'object' || source === null) return ''; - if ('data' in source && typeof source.data === 'object' && source.data !== null) { - if ('message' in source.data && typeof source.data.message === 'string') { - return source.data.message; - } - } - if ('message' in source && typeof source.message === 'string') return source.message; - return ''; -} - function boundedWorkspaceMessage(subtype: WorkspaceFailureSubtype, safeDetail?: string): string { const genericMessage = workspaceFailureMessage(subtype); const detail = safeDetail?.trim(); diff --git a/services/cloud-agent-next/src/session/session-message-queue.test.ts b/services/cloud-agent-next/src/session/session-message-queue.test.ts index 133d23b804..d29f9314f5 100644 --- a/services/cloud-agent-next/src/session/session-message-queue.test.ts +++ b/services/cloud-agent-next/src/session/session-message-queue.test.ts @@ -132,6 +132,7 @@ function createQueueHarness(options?: { getDeliveryBlock?: () => Promise; recoverExhaustedDeliveryBlock?: () => Promise; checkBillingAdmission?: SessionMessageQueueDependencies['checkBillingAdmission']; + persistCloneQueuedMessage?: SessionMessageQueueDependencies['persistCloneQueuedMessage']; }) { const storage = options?.storage ?? createMemoryStorage(); const events: QueueEvent[] = []; @@ -182,6 +183,19 @@ function createQueueHarness(options?: { events.push(event); if (messageId) admittedEventMessageIds.add(messageId); }, + persistCloneQueuedMessage: + options?.persistCloneQueuedMessage ?? + (async (intent, callbackSnapshot) => { + const now = Date.now(); + await storePendingSessionMessage( + storage, + createPendingSessionMessageFromIntent(intent, now, callbackSnapshot) + ); + await putSessionMessageState( + storage, + createQueuedSessionMessageState(intent, callbackSnapshot, now) + ); + }), ensureAcceptedMessageEffects: options?.ensureAcceptedMessageEffects ?? (async () => undefined), persistTerminalTransition: async (messageId, params, options) => { @@ -1136,35 +1150,76 @@ describe('SessionMessageQueue', () => { await expect(listPendingSessionMessages(harness.storage)).resolves.toHaveLength(0); }); - it('admits a durable queued message once and replays the original acknowledgement', async () => { - const harness = createQueueHarness(); - const request = { - userId: 'user_test' as UserId, - turn: { type: 'prompt' as const, id: FIRST_MESSAGE_ID, prompt: 'queue this prompt' }, - }; + it.each([false, true])( + 'admits a durable queued message once and replays its acknowledgement (marked clone: %s)', + async markedClone => { + const harness = createQueueHarness({ + metadata: createMetadata({ + clone: markedClone + ? { + cloneFromKiloSessionId: 'ses_aaaaaaaaaaaaaaaaaaaaaaaaaa', + reportingCreatedAt: '2026-08-01T10:00:00.000Z', + } + : undefined, + }), + }); + const request = { + userId: 'user_test' as UserId, + turn: { type: 'prompt' as const, id: FIRST_MESSAGE_ID, prompt: 'queue this prompt' }, + }; - const admitted = await harness.queue.admitSubmittedMessage(request); - const replay = await harness.queue.admitSubmittedMessage(request); - const pending = await listPendingSessionMessages(harness.storage); - const messageState = await getSessionMessageState(harness.storage, FIRST_MESSAGE_ID); + const admitted = await harness.queue.admitSubmittedMessage(request); + const replay = await harness.queue.admitSubmittedMessage(request); + const pending = await listPendingSessionMessages(harness.storage); + const messageState = await getSessionMessageState(harness.storage, FIRST_MESSAGE_ID); - expect(admitted).toEqual({ - success: true, - outcome: 'queued', - compatibilityDelivery: 'queued', - messageId: FIRST_MESSAGE_ID, - }); - expect(replay).toEqual(admitted); - expect(pending.map(message => message.messageId)).toEqual([FIRST_MESSAGE_ID]); - expect(messageState?.status).toBe('queued'); - expect(harness.events.map(event => event.streamEventType)).toEqual(['cloud.message.queued']); - expect(JSON.parse(harness.events[0]?.payload ?? '{}')).toMatchObject({ - messageId: FIRST_MESSAGE_ID, - content: 'queue this prompt', - delivery: 'queued', - }); - expect(harness.alarmDeadlines).toHaveLength(2); - }); + expect(admitted).toEqual({ + success: true, + outcome: 'queued', + compatibilityDelivery: 'queued', + messageId: FIRST_MESSAGE_ID, + }); + expect(replay).toEqual(admitted); + expect(pending.map(message => message.messageId)).toEqual([FIRST_MESSAGE_ID]); + expect(messageState?.status).toBe('queued'); + expect(harness.events.map(event => event.streamEventType)).toEqual(['cloud.message.queued']); + expect(JSON.parse(harness.events[0]?.payload ?? '{}')).toMatchObject({ + messageId: FIRST_MESSAGE_ID, + content: 'queue this prompt', + delivery: 'queued', + }); + expect(harness.alarmDeadlines).toHaveLength(2); + } + ); + + it.each([false, true])( + 'uses DO-owned queued persistence only for marked clones (%s)', + async markedClone => { + const persistCloneQueuedMessage = vi.fn().mockRejectedValue(new Error('transaction failed')); + const harness = createQueueHarness({ + metadata: createMetadata({ + clone: { + cloneFromKiloSessionId: 'ses_aaaaaaaaaaaaaaaaaaaaaaaaaa', + ...(markedClone ? { reportingCreatedAt: '2026-08-01T10:00:00.000Z' } : {}), + }, + }), + persistCloneQueuedMessage, + }); + + const result = await harness.queue.admitSubmittedMessage({ + userId: 'user_test' as UserId, + turn: { type: 'prompt', id: FIRST_MESSAGE_ID, prompt: 'clone first message' }, + }); + + expect(result.success).toBe(!markedClone); + expect(await listPendingSessionMessages(harness.storage)).toHaveLength(markedClone ? 0 : 1); + expect(await getSessionMessageState(harness.storage, FIRST_MESSAGE_ID)).toEqual( + markedClone ? undefined : expect.objectContaining({ status: 'queued' }) + ); + expect(harness.events).toHaveLength(markedClone ? 0 : 1); + expect(harness.alarmDeadlines).toHaveLength(markedClone ? 0 : 1); + } + ); it('repairs submitted admission event/drain effects after an event persistence failure', async () => { const harness = createQueueHarness({ failQueuedEventOnce: true }); diff --git a/services/cloud-agent-next/src/session/session-message-queue.ts b/services/cloud-agent-next/src/session/session-message-queue.ts index 950d6e5174..b142a5d6d8 100644 --- a/services/cloud-agent-next/src/session/session-message-queue.ts +++ b/services/cloud-agent-next/src/session/session-message-queue.ts @@ -164,6 +164,10 @@ export type SessionMessageQueueDependencies = { > | null>; ensureQueuedMessageEvent: (event: PersistedQueuedMessageEvent & { entityId: string }) => void; reportQueuedState?: (state: SessionMessageState) => void; + persistCloneQueuedMessage: ( + intent: SessionMessageIntent, + callbackSnapshot?: PendingSessionMessage['callbackSnapshot'] + ) => Promise; ensureAcceptedMessageEffects: (messageId: string) => Promise; persistTerminalTransition: ( messageId: string, @@ -861,9 +865,13 @@ export function createSessionMessageQueue( ? { required: true, target: callbackTarget } : undefined; - await enqueuePendingSessionMessageIntent(storage, intent, Date.now(), callbackSnapshot); - const messageState = createQueuedSessionMessageState(intent, callbackSnapshot); - await putSessionMessageState(storage, messageState); + if (metadata?.clone?.reportingCreatedAt) { + await dependencies.persistCloneQueuedMessage(intent, callbackSnapshot); + } else { + await enqueuePendingSessionMessageIntent(storage, intent, Date.now(), callbackSnapshot); + const messageState = createQueuedSessionMessageState(intent, callbackSnapshot); + await putSessionMessageState(storage, messageState); + } await completeQueuedAdmissionEffects(intent); return buildAdmissionAck(turn.messageId); } diff --git a/services/cloud-agent-next/src/session/session-message-state.test.ts b/services/cloud-agent-next/src/session/session-message-state.test.ts index 887f148d9a..faf59d9e4e 100644 --- a/services/cloud-agent-next/src/session/session-message-state.test.ts +++ b/services/cloud-agent-next/src/session/session-message-state.test.ts @@ -15,6 +15,7 @@ import { isTerminalMessageState, type SessionMessageState, type SessionMessageStorage, + type MarkMessageFailedParams, } from './session-message-state.js'; import type { SessionMessageIntent } from '../execution/types.js'; @@ -255,6 +256,93 @@ describe('getSessionMessageState / putSessionMessageState', () => { expect(loaded?.agentActivityObservedAt).toBeUndefined(); }); + it.each([ + ['legacy', undefined], + ['current', createIntent(VALID_MESSAGE_ID, 'stored prompt')], + ] as const)( + 'strips retired failureFacts on %s state reads and writes', + async (_name, admissionSnapshot) => { + const storage = createFakeStorage(); + const storedState = { + messageId: VALID_MESSAGE_ID, + status: 'failed' as const, + prompt: 'stored prompt', + createdAt: 1000, + terminalAt: 2000, + assistantMessageId: 'asst_original', + assistantFailureReason: 'provider_unavailable' as const, + safeFailureMessage: 'Assistant service is unavailable', + failureFacts: { + version: 1, + sdkModel: 'vendor/model', + sdkProvider: 'kilo', + errorName: 'APIError', + sdkStatusCode: 503, + isRetryable: false, + }, + admissionSnapshot, + agent: { model: 'legacy-model' }, + callbackRequired: true, + }; + await storage.put(`session_message:${VALID_MESSAGE_ID}`, storedState); + + const loaded = await getSessionMessageState(storage, VALID_MESSAGE_ID); + expect(loaded).toMatchObject({ + status: 'failed', + terminalAt: 2000, + assistantMessageId: 'asst_original', + assistantFailureReason: 'provider_unavailable', + safeFailureMessage: 'Assistant service is unavailable', + callbackRequired: true, + }); + expect(loaded).not.toHaveProperty('failureFacts'); + expect(loaded?.admissionSnapshot).toEqual(admissionSnapshot); + expect(loaded?.legacyAdmissionConstraints).toEqual( + admissionSnapshot === undefined ? { agent: { model: 'legacy-model' } } : undefined + ); + expect(await listMessagesWithPendingCallbacks(storage)).toEqual([loaded]); + + await putSessionMessageState(storage, storedState); + expect(storage.store.get(`session_message:${VALID_MESSAGE_ID}`)).not.toHaveProperty( + 'failureFacts' + ); + expect(await getSessionMessageState(storage, VALID_MESSAGE_ID)).toEqual(loaded); + expect(storedState).toHaveProperty('failureFacts'); + } + ); + + it.each(['future_assistant_reason', null, 42, {}])( + 'ignores an unknown optional assistant reason without losing lifecycle or callback state: %j', + async assistantFailureReason => { + const storage = createFakeStorage(); + await storage.put(`session_message:${VALID_MESSAGE_ID}`, { + messageId: VALID_MESSAGE_ID, + status: 'failed', + prompt: 'stored prompt', + createdAt: 1000, + terminalAt: 2000, + assistantMessageId: 'asst_original', + assistantFailureReason, + safeFailureMessage: 'Assistant request failed', + callbackRequired: true, + }); + + const loaded = await getSessionMessageState(storage, VALID_MESSAGE_ID); + expect(loaded).toMatchObject({ + status: 'failed', + terminalAt: 2000, + assistantMessageId: 'asst_original', + safeFailureMessage: 'Assistant request failed', + callbackRequired: true, + }); + expect(loaded?.assistantFailureReason).toBeUndefined(); + expect(await listMessagesWithPendingCallbacks(storage)).toEqual([loaded]); + if (!loaded) throw new Error('Expected stored lifecycle state'); + await putSessionMessageState(storage, loaded); + expect(await getSessionMessageState(storage, VALID_MESSAGE_ID)).toEqual(loaded); + } + ); + it('returns undefined for unknown messageId', async () => { const storage = createFakeStorage(); const loaded = await getSessionMessageState(storage, 'msg_unknown00000000ABCDEFGHIJKLMN'); @@ -448,6 +536,85 @@ describe('markMessageFailed', () => { }); }); +describe.each(['markMessageFailed', 'terminalizeMessageOnce'] as const)( + '%s failure correlation', + path => { + async function fail( + storage: SessionMessageStorage, + params: MarkMessageFailedParams, + now: number + ) { + return path === 'markMessageFailed' + ? markMessageFailed(storage, VALID_MESSAGE_ID, params, now) + : ( + await terminalizeMessageOnce( + storage, + VALID_MESSAGE_ID, + { kind: 'failed', ...params }, + now + ) + ).state; + } + + it('round-trips the failure reason and assistant identity without accepting a conflicting late failure', async () => { + const storage = createFakeStorage(); + await putSessionMessageState( + storage, + createQueuedSessionMessageState(createIntent(VALID_MESSAGE_ID, 'hello')) + ); + await markMessageAccepted(storage, VALID_MESSAGE_ID, 'wr_failure', 1000); + const params: MarkMessageFailedParams = { + reason: 'assistant_error', + completionSource: 'assistant_message_event', + assistantMessageId: 'asst_original', + assistantFailureReason: 'provider_unavailable', + providerOwnership: 'managed', + safeFailureMessage: 'Assistant service is unavailable', + }; + + const failed = await fail(storage, params, 2000); + expect(failed).toMatchObject({ + status: 'failed', + terminalAt: 2000, + assistantMessageId: params.assistantMessageId, + assistantFailureReason: params.assistantFailureReason, + safeFailureMessage: params.safeFailureMessage, + }); + expect(await getSessionMessageState(storage, VALID_MESSAGE_ID)).toEqual(failed); + + await fail(storage, params, 3000); + await fail( + storage, + { + ...params, + assistantMessageId: 'asst_late_conflict', + assistantFailureReason: 'rate_limited', + safeFailureMessage: 'Assistant request was rate limited', + }, + 4000 + ); + expect(await getSessionMessageState(storage, VALID_MESSAGE_ID)).toEqual(failed); + }); + + it('keeps the stored assistant identity when the failure omits it', async () => { + const storage = createFakeStorage(); + await putSessionMessageState(storage, { + ...createQueuedSessionMessageState(createIntent(VALID_MESSAGE_ID, 'hello')), + assistantMessageId: 'asst_original', + }); + await fail( + storage, + { reason: 'assistant_error', completionSource: 'idle_reconciliation' }, + 2000 + ); + expect(await getSessionMessageState(storage, VALID_MESSAGE_ID)).toMatchObject({ + status: 'failed', + assistantMessageId: 'asst_original', + }); + }); + } +); + describe('markMessageInterrupted', () => { it('transitions queued to interrupted', async () => { const storage = createFakeStorage(); diff --git a/services/cloud-agent-next/src/session/session-message-state.ts b/services/cloud-agent-next/src/session/session-message-state.ts index 370b29db28..3507df9ae1 100644 --- a/services/cloud-agent-next/src/session/session-message-state.ts +++ b/services/cloud-agent-next/src/session/session-message-state.ts @@ -222,7 +222,7 @@ export const SessionMessageStateSchema = z failureStage: SessionMessageFailureStageSchema.optional(), failureCode: SessionMessageFailureCodeSchema.optional(), failureSubtype: WorkspaceFailureSubtypeSchema.optional(), - assistantFailureReason: CloudAgentAssistantFailureReasonSchema.optional(), + assistantFailureReason: CloudAgentAssistantFailureReasonSchema.optional().catch(undefined), providerOwnership: CloudAgentProviderOwnershipSchema.optional(), safeFailureMessage: z.string().max(WRAPPER_READY_ERROR_DETAIL_MAX_LENGTH).optional(), modelNotFoundRuntimeDiagnostics: ModelNotFoundRuntimeDiagnosticsSchema.optional(), @@ -321,6 +321,7 @@ function normalizeParsedSessionMessageState( ): SessionMessageState { const currentState = { ...state }; + delete currentState.failureFacts; delete currentState.turn; delete currentState.images; delete currentState.agent; @@ -434,6 +435,7 @@ export async function putSessionMessageState( state: SessionMessageState ): Promise { const parsedState = SessionMessageStateSchema.parse(state); + delete parsedState.failureFacts; if (parsedState.wrapperRunId) { await storage.put( wrapperRunMessageIndexKey(parsedState.wrapperRunId, parsedState.messageId), @@ -523,6 +525,7 @@ export async function markMessageCompleted( } export type MarkMessageFailedParams = { + assistantMessageId?: string; reason: string; error?: string; completionSource: SessionMessageCompletionSource; @@ -558,6 +561,7 @@ export async function markMessageFailed( assistantFailureReason: params.assistantFailureReason, providerOwnership: params.providerOwnership, safeFailureMessage: params.safeFailureMessage, + assistantMessageId: params.assistantMessageId ?? state.assistantMessageId, modelNotFoundRuntimeDiagnostics: params.modelNotFoundRuntimeDiagnostics, attempts: params.attempts, }; @@ -770,6 +774,7 @@ export type TerminalizeParams = } | { kind: 'failed'; + assistantMessageId?: string; reason: string; error?: string; completionSource: SessionMessageCompletionSource; @@ -844,6 +849,7 @@ export async function terminalizeMessageOnce( assistantFailureReason: params.assistantFailureReason, providerOwnership: resolveTerminalProviderOwnership(params, state), safeFailureMessage: params.safeFailureMessage, + assistantMessageId: params.assistantMessageId ?? state.assistantMessageId, modelNotFoundRuntimeDiagnostics: params.modelNotFoundRuntimeDiagnostics, attempts: params.attempts, terminalEffects, diff --git a/services/cloud-agent-next/src/session/session-prepare.test.ts b/services/cloud-agent-next/src/session/session-prepare.test.ts index 15476eed5b..023792c499 100644 --- a/services/cloud-agent-next/src/session/session-prepare.test.ts +++ b/services/cloud-agent-next/src/session/session-prepare.test.ts @@ -7,7 +7,7 @@ * Durable Object RPC transport are mocked so each ladder branch is exercised * deterministically; `startNewSession` and the reconcile ladder run real code. */ -import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import type { WorkerDb } from '@kilocode/db/client'; import type { OperationLedgerRow } from '@kilocode/db/schema'; @@ -21,6 +21,7 @@ import type * as SharedSandboxRouteModule from '../shared-sandbox-route.js'; import type * as MessageIdModule from './message-id.js'; import { createSessionWithLedger, + registerNewSession, sessionCreateIntentFingerprint, SESSION_CREATE_ABANDON_AFTER_SECONDS, SESSION_CREATE_ABANDONED_OUTCOME_CODE, @@ -235,6 +236,10 @@ function makeRequest(overrides: Partial = {}): SessionCrea const CREATE_OPTIONS = { operationKey: OPERATION_KEY, startedAt: 1_700_000_000_000 }; +afterEach(() => { + vi.useRealTimers(); +}); + /** Runs the ledger-guarded create with the standard operation options. */ function runCreate( ctx: SessionRegistrationContext, @@ -1965,6 +1970,10 @@ describe('createSessionWithLedger clone allocation outcomes', () => { expect(sandboxSessionGet).toHaveBeenCalledTimes(1); expect(cloudAgentSessionGet).not.toHaveBeenCalled(); expect(doStub.registerSession).toHaveBeenCalledTimes(1); + expect(doStub.registerSession.mock.calls[0]?.[0].clone).toEqual({ + cloneFromKiloSessionId: SOURCE_KILO_SESSION_ID, + }); + expect(recordOperationProgressMock.mock.calls[0]?.[2]).not.toHaveProperty('reportingCreatedAt'); }); it('rejects isolated Standard allocation for workspace sessions before external effects', async () => { @@ -2221,12 +2230,65 @@ describe('createSessionWithLedger clone allocation outcomes', () => { }); }); - it('records a none initial-turn fingerprint and no initialMessageId for a clone-only create', async () => { - createCliSessionMock.mockResolvedValue({ + it('uses server time rather than supplied clone fields for registration-only allocations', async () => { + const reportingCreatedAt = '2026-08-01T10:00:00.000Z'; + vi.useFakeTimers(); + vi.setSystemTime(reportingCreatedAt); + createCliSessionMock.mockResolvedValueOnce({ status: 'ready', clone: { sessionId: KILO_SESSION_ID, copiedItemCount: 1 }, }); const doStub = makeDoStub(); + const request = { + ...cloneRequest(), + clone: { + cloneFromKiloSessionId: SOURCE_KILO_SESSION_ID, + reportingCreatedAt: '2099-01-01T00:00:00.000Z', + }, + }; + + await registerNewSession(request, makeContext(doStub)); + + expect(doStub.registerSession).toHaveBeenCalledWith( + expect.objectContaining({ + clone: { cloneFromKiloSessionId: SOURCE_KILO_SESSION_ID, reportingCreatedAt }, + }) + ); + expect(recordOperationProgressMock).not.toHaveBeenCalled(); + expect(createSessionReportMock).not.toHaveBeenCalled(); + }); + + it('does not mark clones that already have an initial turn', async () => { + createCliSessionMock.mockResolvedValueOnce({ + status: 'ready', + clone: { sessionId: KILO_SESSION_ID, copiedItemCount: 1 }, + }); + const doStub = makeDoStub(); + + await runCreate(makeContext(doStub), { + ...cloneRequest(), + initialTurn: { type: 'prompt', prompt: 'continue the clone' }, + }); + + expect(recordOperationProgressMock.mock.calls[0]?.[2]).not.toHaveProperty('reportingCreatedAt'); + expect(doStub.createSessionWithInitialAdmission.mock.calls[0]?.[0].clone).toEqual({ + cloneFromKiloSessionId: SOURCE_KILO_SESSION_ID, + }); + expect(createSessionReportMock).toHaveBeenCalledTimes(1); + }); + + it('checkpoints clone-only reporting age before ownership allocation without an initial message', async () => { + const reportingCreatedAt = '2026-08-01T10:00:00.000Z'; + vi.useFakeTimers(); + vi.setSystemTime(reportingCreatedAt); + createCliSessionMock.mockImplementationOnce(async () => { + vi.setSystemTime('2026-08-01T11:00:00.000Z'); + return { + status: 'ready', + clone: { sessionId: KILO_SESSION_ID, copiedItemCount: 1 }, + }; + }); + const doStub = makeDoStub(); const ctx = makeContext(doStub); await runCreate(ctx, cloneRequest()); @@ -2234,8 +2296,15 @@ describe('createSessionWithLedger clone allocation outcomes', () => { expect(recordOperationProgressMock).toHaveBeenNthCalledWith(1, expect.any(Object), ROW_ID, { cloudAgentSessionId: CLOUD_AGENT_SESSION_ID, kiloSessionId: KILO_SESSION_ID, + reportingCreatedAt, createIntentFingerprint: await sessionCreateIntentFingerprint(cloneRequest()), }); + expect(doStub.registerSession).toHaveBeenCalledWith( + expect.objectContaining({ + clone: { cloneFromKiloSessionId: SOURCE_KILO_SESSION_ID, reportingCreatedAt }, + }) + ); + expect(createSessionReportMock).not.toHaveBeenCalled(); // A clone-only create has no synthetic initial turn, so no initialMessageId // is recorded and the fingerprint differs from a prompt create. expect(recordOperationProgressMock.mock.calls[0]?.[2]).not.toHaveProperty('initialMessageId'); @@ -2291,52 +2360,83 @@ describe('createSessionWithLedger clone reconciliation', () => { }); } - it('resumes a clone with the stored IDs when the ownership row is absent', async () => { - const request = cloneRequest(); + it.each([undefined, '2026-08-01T10:00:00.000Z'])( + 'resumes stored clone IDs without changing or inventing reporting age (%s)', + async reportingCreatedAt => { + vi.useFakeTimers(); + vi.setSystemTime('2026-08-29T10:00:00.000Z'); + const request = cloneRequest(); + admitOperationMock.mockResolvedValueOnce({ + admission: 'takeover', + row: await cloneRow(reportingCreatedAt ? { reportingCreatedAt } : {}), + }); + createCliSessionMock.mockResolvedValue({ + status: 'ready', + clone: { sessionId: KILO_SESSION_ID, copiedItemCount: 3 }, + }); + // First query: ownership absent. Second: distinct-id email. + getPgDbMock.mockReturnValue(makeDb([[], [{ email: 'test@example.com' }]])); + const doStub = makeDoStub(); + const ctx = makeContext(doStub); + + const result = await runCreate(ctx, request); + + expect(createCliSessionMock).toHaveBeenCalledWith( + KILO_SESSION_ID, + CLOUD_AGENT_SESSION_ID, + USER_ID, + expect.any(Object), + undefined, + 'cloud-agent', + expect.any(String), + 'https://github.com/acme/repo', + SOURCE_KILO_SESSION_ID + ); + expect(doStub.registerSession).toHaveBeenCalledTimes(1); + expect(settleOperationMock).toHaveBeenCalledWith( + expect.any(Object), + expect.objectContaining({ + rowId: ROW_ID, + status: 'completed', + outcomeCode: 'ok', + canonicalResult: { + cloudAgentSessionId: CLOUD_AGENT_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + }, + }) + ); + expect(result).toEqual({ + cloudAgentSessionId: CLOUD_AGENT_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + replayed: true, + }); + expect(doStub.registerSession.mock.calls[0]?.[0].clone).toEqual({ + cloneFromKiloSessionId: SOURCE_KILO_SESSION_ID, + ...(reportingCreatedAt ? { reportingCreatedAt } : {}), + }); + expect(createSessionReportMock).not.toHaveBeenCalled(); + expect(recordOperationProgressMock).not.toHaveBeenCalled(); + } + ); + + it('rejects a clone checkpoint with an invalid reporting creation time', async () => { + createCliSessionMock.mockResolvedValueOnce({ + status: 'ready', + clone: { sessionId: KILO_SESSION_ID, copiedItemCount: 1 }, + }); admitOperationMock.mockResolvedValueOnce({ admission: 'takeover', - row: await cloneRow(), - }); - createCliSessionMock.mockResolvedValue({ - status: 'ready', - clone: { sessionId: KILO_SESSION_ID, copiedItemCount: 3 }, + row: await cloneRow({ reportingCreatedAt: 'invalid' }), }); - // First query: ownership absent. Second: distinct-id email. getPgDbMock.mockReturnValue(makeDb([[], [{ email: 'test@example.com' }]])); const doStub = makeDoStub(); - const ctx = makeContext(doStub); - const result = await runCreate(ctx, request); - - expect(createCliSessionMock).toHaveBeenCalledWith( - KILO_SESSION_ID, - CLOUD_AGENT_SESSION_ID, - USER_ID, - expect.any(Object), - undefined, - 'cloud-agent', - expect.any(String), - 'https://github.com/acme/repo', - SOURCE_KILO_SESSION_ID - ); - expect(doStub.registerSession).toHaveBeenCalledTimes(1); - expect(settleOperationMock).toHaveBeenCalledWith( - expect.any(Object), - expect.objectContaining({ - rowId: ROW_ID, - status: 'completed', - outcomeCode: 'ok', - canonicalResult: { - cloudAgentSessionId: CLOUD_AGENT_SESSION_ID, - kiloSessionId: KILO_SESSION_ID, - }, - }) - ); - expect(result).toEqual({ - cloudAgentSessionId: CLOUD_AGENT_SESSION_ID, - kiloSessionId: KILO_SESSION_ID, - replayed: true, + await expect(runCreate(makeContext(doStub), cloneRequest())).rejects.toMatchObject({ + code: 'CONFLICT', + message: 'creation_in_progress', }); + expect(doStub.registerSession).not.toHaveBeenCalled(); + expect(createSessionReportMock).not.toHaveBeenCalled(); }); it('resumes a workspace clone with its persisted Vercel provider', async () => { @@ -2433,7 +2533,10 @@ describe('createSessionWithLedger clone reconciliation', () => { }); }); - it('keeps a clone unknown outcome reconcile-pending and resumes the stored IDs on a same-key retry', async () => { + it('keeps a clone unknown outcome reconcile-pending and resumes the stored IDs and age on a same-key retry', async () => { + const reportingCreatedAt = '2026-08-01T10:00:00.000Z'; + vi.useFakeTimers(); + vi.setSystemTime(reportingCreatedAt); const request = cloneRequest(); const doStub = makeDoStub(); const ctx = makeContext(doStub); @@ -2450,10 +2553,12 @@ describe('createSessionWithLedger clone reconciliation', () => { expect(markReconcilePendingMock).toHaveBeenCalledWith(expect.any(Object), { rowId: ROW_ID }); expect(settleOperationMock).not.toHaveBeenCalled(); - // Attempt 2: same-key retry reconciles the stored IDs via the clone resume. + const checkpoint = recordOperationProgressMock.mock.calls[0]?.[2]; + expect(checkpoint).toMatchObject({ reportingCreatedAt }); + vi.setSystemTime('2026-08-29T10:00:00.000Z'); admitOperationMock.mockResolvedValueOnce({ admission: 'duplicate_reconcile_pending', - row: await cloneRow(), + row: await cloneRow(checkpoint), }); createCliSessionMock.mockResolvedValue({ status: 'ready', @@ -2488,6 +2593,10 @@ describe('createSessionWithLedger clone reconciliation', () => { kiloSessionId: KILO_SESSION_ID, replayed: true, }); + expect(doStub.registerSession.mock.calls[0]?.[0].clone).toEqual({ + cloneFromKiloSessionId: SOURCE_KILO_SESSION_ID, + reportingCreatedAt, + }); }); it('surfaces SERVICE_UNAVAILABLE session_clone_unavailable when the resume ingest sends no acknowledgement', async () => { diff --git a/services/cloud-agent-next/src/session/session-registration.ts b/services/cloud-agent-next/src/session/session-registration.ts index 1c3193d935..5a2e74da0b 100644 --- a/services/cloud-agent-next/src/session/session-registration.ts +++ b/services/cloud-agent-next/src/session/session-registration.ts @@ -16,6 +16,7 @@ */ import { TRPCError } from '@trpc/server'; import { and, eq } from 'drizzle-orm'; +import { z } from 'zod'; import type { WorkerDb } from '@kilocode/db/client'; import { cli_sessions_v2, kilocode_users } from '@kilocode/db/schema'; import { @@ -165,6 +166,7 @@ type SessionEstablishmentFailure = | { stage: 'transport'; code: 'do_rpc_outcome_unknown' }; type NewSessionAllocation = SessionRegistrationResult & { + reportingCreatedAt?: string; credentialContainment: CredentialContainment; sessionService: SessionService; rollbackCliSession: () => Promise; @@ -449,6 +451,10 @@ async function allocateNewSession( sessionPlaneForNewOwner(ctx.env, { userId: ctx.userId, orgId }) ); const kiloSessionId = generateKiloSessionId(); + const reportingCreatedAt = + input.clone && !initialTurn && cloudAgentSessionId.startsWith('agent_') + ? new Date().toISOString() + : undefined; const createdOnPlatform = input.options?.createdOnPlatform ?? 'cloud-agent'; try { @@ -465,6 +471,7 @@ async function allocateNewSession( cloudAgentSessionId, kiloSessionId, ...(initialTurn ? { initialMessageId: initialTurn.messageId } : {}), + ...(reportingCreatedAt ? { reportingCreatedAt } : {}), [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint(input), }); } @@ -643,6 +650,7 @@ async function allocateNewSession( sandboxRoute, sandboxProvider, initialTurn, + reportingCreatedAt, credentialContainment, sessionService, rollbackCliSession: async () => { @@ -687,8 +695,14 @@ function rebuildCloneAllocation( const sandboxId = canonical.sandboxId; const sandboxProvider = canonical.sandboxProvider; const sandboxRoute = canonical.sandboxRoute; + const reportingCreatedAt = z + .string() + .datetime({ offset: true }) + .optional() + .safeParse(canonical.reportingCreatedAt); if ( + !reportingCreatedAt.success || typeof cloudAgentSessionId !== 'string' || cloudAgentSessionId.length === 0 || typeof kiloSessionId !== 'string' || @@ -732,6 +746,10 @@ function rebuildCloneAllocation( sandboxRoute: route, sandboxProvider, initialTurn, + reportingCreatedAt: + !initialTurn && cloudAgentSessionId.startsWith('agent_') + ? reportingCreatedAt.data + : undefined, credentialContainment: computeCredentialContainment(input, ctx), sessionService, rollbackCliSession: async () => { @@ -774,7 +792,14 @@ function buildSessionRegistrationCommand( kiloSessionId: allocation.kiloSessionId, kilocodeToken: ctx.authToken, }, - clone: input.clone, + clone: input.clone + ? { + cloneFromKiloSessionId: input.clone.cloneFromKiloSessionId, + ...(allocation.reportingCreatedAt + ? { reportingCreatedAt: allocation.reportingCreatedAt } + : {}), + } + : undefined, ...(allocation.initialTurn ? { message: { diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts index d47c719ff3..65437579ff 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts @@ -963,6 +963,89 @@ describe('WrapperSupervisor', () => { }); }); + it.each(['wrapper_ping_timeout', 'wrapper_no_output', 'wrapper_disconnected'] as const)( + 'preserves %s with per-message activity stages and positive terminal reconciliation', + async failureCode => { + const deadlineAt = 100_000; + const harness = createHarness( + [ + liveRuntimeState( + failureCode === 'wrapper_no_output' + ? { noOutputDeadlineAt: deadlineAt, nextPingAt: deadlineAt + 1 } + : { pingDeadlineAt: deadlineAt, noOutputDeadlineAt: deadlineAt + 1 } + ), + OWNED_WRAPPER_LEASE, + ...(failureCode === 'wrapper_disconnected' + ? [disconnectGraceForCurrentConnection(deadlineAt - 10_000)] + : []), + ], + { + getAssistantMessageForUserMessage: (_sessionId, _kiloSessionId, parentMessageId) => + parentMessageId === MESSAGE_ID + ? null + : { + eventId: 1, + timestamp: 2_500, + info: { + id: `assistant_${parentMessageId}`, + role: 'assistant', + time: { + created: 2_500, + ...(parentMessageId === NEWEST_MESSAGE_ID ? { completed: 90_000 } : {}), + }, + }, + parts: [], + }, + } + ); + await putSessionMessageState(harness.storage, acceptedMessage()); + for (const messageId of [NEWER_MESSAGE_ID, NEWEST_MESSAGE_ID]) { + await putSessionMessageState(harness.storage, { + ...acceptedMessage(messageId), + agentActivityObservedAt: 2_500, + }); + } + + await harness.supervisor.runMaintenance(deadlineAt - 1); + expect(harness.events).toHaveLength(0); + await harness.supervisor.runMaintenance(deadlineAt); + + for (const [messageId, failureStage] of [ + [MESSAGE_ID, 'post_dispatch_no_activity'], + [NEWER_MESSAGE_ID, 'agent_activity'], + ] as const) { + await expect(getSessionMessageState(harness.storage, messageId)).resolves.toMatchObject({ + status: 'failed', + completionSource: 'wrapper_failure', + failureStage, + failureCode, + }); + expect(harness.events.map(event => JSON.parse(event.payload))).toContainEqual( + expect.objectContaining({ + messageId, + status: 'failed', + failure: expect.objectContaining({ stage: failureStage, code: failureCode }), + }) + ); + } + await expect( + getSessionMessageState(harness.storage, NEWEST_MESSAGE_ID) + ).resolves.toMatchObject({ + status: 'completed', + completionSource: 'idle_reconciliation', + assistantMessageId: `assistant_${NEWEST_MESSAGE_ID}`, + }); + expect( + harness.events.filter(event => event.streamEventType === 'cloud.message.completed') + ).toHaveLength(1); + expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'unhealthy-wrapper', + }); + } + ); + it('fails genuinely silent accepted work at the no-output deadline', async () => { const acceptedAt = 2_000; const noOutputDeadlineAt = acceptedAt + WRAPPER_NO_OUTPUT_TIMEOUT_MS; @@ -984,7 +1067,7 @@ describe('WrapperSupervisor', () => { expect(state).toMatchObject({ status: 'failed', failureReason: 'wrapper_failure', - error: 'Wrapper accepted the message but produced no output', + error: 'Wrapper made no execution progress during the watchdog window', completionSource: 'wrapper_failure', failureStage: 'post_dispatch_no_activity', failureCode: 'wrapper_no_output', @@ -1100,7 +1183,16 @@ describe('WrapperSupervisor', () => { info: { id: 'ase_terminal_error', role: 'assistant', - error: { data: { message: 'Payment required: insufficient credits' } }, + modelID: 'vendor/model', + providerID: 'kilo', + error: { + name: 'APIError', + data: { + message: 'Payment required: insufficient credits', + statusCode: 402, + isRetryable: false, + }, + }, }, parts: [], }) as unknown as LatestAssistantMessage, @@ -1108,6 +1200,10 @@ describe('WrapperSupervisor', () => { ); await putSessionMessageState(harness.storage, acceptedMessage()); + await harness.supervisor.runMaintenance(pingDeadlineAt - 1); + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'accepted', + }); await harness.supervisor.runMaintenance(pingDeadlineAt); await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ @@ -1115,6 +1211,9 @@ describe('WrapperSupervisor', () => { failureReason: 'assistant_error', completionSource: 'idle_reconciliation', failureCode: 'payment_required', + assistantMessageId: 'ase_terminal_error', + assistantFailureReason: 'insufficient_credits', + safeFailureMessage: 'Assistant request failed: insufficient credits', }); expect(harness.events.map(event => event.streamEventType)).toEqual(['cloud.message.failed']); }); @@ -1876,40 +1975,75 @@ describe('WrapperSupervisor', () => { } ); - it('fails a still-accepted errored reply when the wrapper completes', async () => { - const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE], { - getAssistantMessageForUserMessage: () => - ({ + it.each([ + ['APIError', 'provider_unavailable', 'Assistant service is unavailable'], + ['ContextOverflowError', 'context_limit', 'The model context limit was exceeded'], + ['MessageOutputLengthError', 'output_limit', 'The model output limit was reached'], + [ + 'ContentFilterError', + 'content_filter', + 'The model provider blocked the response under its content policy', + ], + [ + 'StructuredOutputError', + 'structured_output', + 'The model response did not match the required format', + ], + ] as const)( + 'fails a still-accepted %s reply when the wrapper completes', + async (name, assistantFailureReason, safeFailureMessage) => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE], { + getAssistantMessageForUserMessage: () => ({ + eventId: 1, + timestamp: 2_500, info: { id: 'ase_complete_error', role: 'assistant', - error: { data: { message: 'provider failed during completion' } }, + modelID: 'vendor/model', + providerID: 'kilo', + error: { + name, + data: { + message: 'provider failed during completion', + statusCode: 503, + isRetryable: false, + responseBody: 'poison-body', + responseHeaders: { authorization: 'poison-header' }, + }, + }, }, parts: [], - }) as unknown as LatestAssistantMessage, - }); - await putSessionMessageState(harness.storage, acceptedMessage()); + }), + }); + await putSessionMessageState(harness.storage, acceptedMessage()); - await harness.supervisor.onTerminalEvent({ - wrapperRunId: WRAPPER_RUN_ID, - status: 'completed', - messageIds: [MESSAGE_ID], - }); + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + messageIds: [MESSAGE_ID], + }); - await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ - status: 'failed', - failureReason: 'assistant_error', - error: 'provider failed during completion', - completionSource: 'idle_reconciliation', - }); - await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ - state: 'stop_needed', - reason: 'terminal-failed', - }); - await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ - wrapperGeneration: 4, - }); - }); + const failed = await getSessionMessageState(harness.storage, MESSAGE_ID); + expect(failed).toMatchObject({ + status: 'failed', + failureReason: 'assistant_error', + error: 'provider failed during completion', + completionSource: 'idle_reconciliation', + assistantMessageId: 'ase_complete_error', + assistantFailureReason, + safeFailureMessage, + }); + expect(failed).not.toHaveProperty('failureFacts'); + expect(JSON.stringify(failed)).not.toMatch(/poison|vendor\/model|statusCode|isRetryable/); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'terminal-failed', + }); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ + wrapperGeneration: 4, + }); + } + ); it('includes the gate result when wrapper completion releases a reconciled callback', async () => { const assistantMessageId = 'ase_complete_gate'; diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.ts index 01c8dc2bdc..75dc062944 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.ts @@ -246,7 +246,8 @@ function getAssistantErrorMessage(error: unknown): string | undefined { return 'Assistant message failed'; } -function assistantErrorTerminalizeParams(assistantError: unknown): TerminalizeParams { +function assistantErrorTerminalizeParams(info: LatestAssistantMessage['info']): TerminalizeParams { + const assistantError = info.error; if (isAssistantInterrupt(assistantError)) { return { kind: 'interrupted', @@ -268,6 +269,7 @@ function assistantErrorTerminalizeParams(assistantError: unknown): TerminalizePa assistantFailureReason: assistantFailure.reason, providerOwnership: assistantFailure.providerOwnership, safeFailureMessage: assistantFailure.safeMessage, + assistantMessageId: info.id, }; } @@ -289,7 +291,7 @@ function projectWrapperDeathReconciliation( ): TerminalizeParams | null { if (!assistantMessage) return null; if (assistantMessage.info.error !== undefined && assistantMessage.info.error !== null) { - return assistantErrorTerminalizeParams(assistantMessage.info.error); + return assistantErrorTerminalizeParams(assistantMessage.info); } if (!hasAssistantCompletionMarker(assistantMessage.info)) return null; return { @@ -783,7 +785,7 @@ export function createWrapperSupervisor( error, completionSource: 'wrapper_failure', failureStage: activityObserved ? 'agent_activity' : 'post_dispatch_no_activity', - failureCode: activityObserved ? 'wrapper_error_after_activity' : failureCode, + failureCode, }; }); await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch(); @@ -862,7 +864,7 @@ export function createWrapperSupervisor( error: 'Wrapper disconnected', completionSource: 'wrapper_failure', failureStage: activityObserved ? 'agent_activity' : 'post_dispatch_no_activity', - failureCode: activityObserved ? 'wrapper_error_after_activity' : 'wrapper_disconnected', + failureCode: 'wrapper_disconnected', }; }); await clearWrapperRuntimeIdentity( @@ -969,7 +971,7 @@ export function createWrapperSupervisor( .warn('Wrapper liveness no-output deadline expired'); await handleUnhealthyWrapper( state, - 'Wrapper accepted the message but produced no output', + 'Wrapper made no execution progress during the watchdog window', 'wrapper_no_output' ); return true; @@ -1104,11 +1106,11 @@ export function createWrapperSupervisor( ? getAssistantMessageForUserMessage(metadata.identity.sessionId, kiloSessionId, messageId) : null; const assistantError = assistantMessage?.info.error; - if (assistantError !== undefined && assistantError !== null) { + if (assistantMessage && assistantError !== undefined && assistantError !== null) { projectedSettlements.push({ message, observeCorrelatedActivity: true, - params: assistantErrorTerminalizeParams(assistantError), + params: assistantErrorTerminalizeParams(assistantMessage.info), }); } else if (assistantMessage) { projectedSettlements.push({ diff --git a/services/cloud-agent-next/src/shared/assistant-failure.ts b/services/cloud-agent-next/src/shared/assistant-failure.ts new file mode 100644 index 0000000000..58d97ec13f --- /dev/null +++ b/services/cloud-agent-next/src/shared/assistant-failure.ts @@ -0,0 +1,175 @@ +import type { + CloudAgentAssistantFailureReason, + CloudAgentProviderOwnership, +} from '@kilocode/worker-utils/cloud-agent-failure'; + +const ASSISTANT_FAILURE_MESSAGES = { + insufficient_credits: 'Assistant request failed: insufficient credits', + rate_limited: 'Assistant request was rate limited', + model_unavailable: 'Assistant request failed: model not found', + provider_authentication: 'Assistant request was not authorized', + provider_unavailable: 'Assistant service is unavailable', + timeout: 'Assistant request timed out', + invalid_request: 'Assistant request was invalid', + context_limit: 'The model context limit was exceeded', + output_limit: 'The model output limit was reached', + content_filter: 'The model provider blocked the response under its content policy', + structured_output: 'The model response did not match the required format', + unknown: 'Assistant request failed', +} as const satisfies Record; +const ASSISTANT_INTERRUPT_MESSAGE = 'The message was interrupted by the user'; +const ASSISTANT_FAILURE_REASONS = Object.keys( + ASSISTANT_FAILURE_MESSAGES +) as CloudAgentAssistantFailureReason[]; + +export function assistantFailureMessage(reason: CloudAgentAssistantFailureReason): string { + return ASSISTANT_FAILURE_MESSAGES[reason]; +} + +export type AssistantFailureClassification = { + reason: CloudAgentAssistantFailureReason; + safeMessage: string; + providerOwnership: CloudAgentProviderOwnership; + terminalCode?: 'payment_required' | 'model_missing'; +}; + +export function projectSafeAssistantError(source: unknown): string | undefined { + if (source === undefined || source === null) return undefined; + const failure = classifyAssistantFailure(source); + const message = isAssistantInterrupt(source) ? ASSISTANT_INTERRUPT_MESSAGE : failure.safeMessage; + return failure.providerOwnership === 'byok' ? `[BYOK] ${message}` : message; +} + +export function isAssistantInterrupt(source: unknown): boolean { + if (typeof source === 'object' && source !== null && 'name' in source) { + if (source.name === 'MessageAbortedError') return true; + } + return /messageabortederror|user[_ -]?interrupt|interrupted by the user/.test( + extractErrorMessage(source).toLocaleLowerCase() + ); +} + +export function classifyAssistantFailure( + source: unknown, + defaultProviderOwnership: CloudAgentProviderOwnership = 'unknown' +): AssistantFailureClassification { + const message = extractErrorMessage(source).toLocaleLowerCase(); + const providerOwnership = /\[byok\]/i.test(message) ? 'byok' : defaultProviderOwnership; + const messageReason = classifyAssistantFailureText(message); + const specificMessageReason = + messageReason !== 'unknown' && + messageReason !== 'invalid_request' && + messageReason !== 'provider_unavailable'; + const reason = specificMessageReason + ? messageReason + : (classifySdkErrorName(source) ?? + (messageReason !== 'unknown' ? messageReason : classifySdkStatus(source)) ?? + 'unknown'); + const terminalCode = + reason === 'insufficient_credits' + ? 'payment_required' + : reason === 'model_unavailable' + ? 'model_missing' + : undefined; + + return { + reason, + safeMessage: assistantFailureMessage(reason), + providerOwnership, + ...(terminalCode === undefined ? {} : { terminalCode }), + }; +} + +function classifySdkErrorName(source: unknown): CloudAgentAssistantFailureReason | undefined { + if (typeof source !== 'object' || source === null || !('name' in source)) return undefined; + switch (source.name) { + case 'ProviderAuthError': + return 'provider_authentication'; + case 'ContextOverflowError': + return 'context_limit'; + case 'MessageOutputLengthError': + return 'output_limit'; + case 'ContentFilterError': + return 'content_filter'; + case 'StructuredOutputError': + return 'structured_output'; + default: + return undefined; + } +} + +function classifySdkStatus(source: unknown): CloudAgentAssistantFailureReason | undefined { + if ( + typeof source !== 'object' || + source === null || + !('name' in source) || + source.name !== 'APIError' || + !('data' in source) || + typeof source.data !== 'object' || + source.data === null || + !('statusCode' in source.data) + ) { + return undefined; + } + const status = source.data.statusCode; + if (typeof status !== 'number' || !Number.isInteger(status) || status < 100 || status > 599) { + return undefined; + } + if (status === 402) return 'insufficient_credits'; + if (status === 429) return 'rate_limited'; + if (status === 401 || status === 403) return 'provider_authentication'; + if (status === 408 || status === 504) return 'timeout'; + if (status >= 500) return 'provider_unavailable'; + if (status >= 400) return 'invalid_request'; + return undefined; +} + +function classifyAssistantFailureText(message: string): CloudAgentAssistantFailureReason { + const canonicalMessage = message.replace(/^\[byok\] /, ''); + const canonicalReason = ASSISTANT_FAILURE_REASONS.find( + reason => ASSISTANT_FAILURE_MESSAGES[reason].toLocaleLowerCase() === canonicalMessage + ); + if (canonicalReason !== undefined) return canonicalReason; + if (/\b(payment required|insufficient (?:credits?|balance|funds))\b/.test(message)) { + return 'insufficient_credits'; + } + if (/\b(model (?:was )?not found|unknown model|invalid model)\b/.test(message)) { + return 'model_unavailable'; + } + if (/\btool calls (?:cutoff|cut off) by max_tokens\b/.test(message)) { + return 'output_limit'; + } + if ( + /\b(rate limit|rate_limit|usage[_ -]?limit[_ -]?exceeded|too many requests|429)\b/.test(message) + ) { + return 'rate_limited'; + } + if (/\b(timed? out|timeout|deadline exceeded)\b/.test(message)) return 'timeout'; + if (/\b(unauthorized|forbidden|authorization|authentication|401|403)\b/.test(message)) { + return 'provider_authentication'; + } + if (/\b(invalid request|bad request|malformed request|400)\b/.test(message)) { + return 'invalid_request'; + } + if (/\b(service unavailable|temporarily unavailable|overloaded|502|503|504)\b/.test(message)) { + return 'provider_unavailable'; + } + return 'unknown'; +} + +export function classifyAssistantFailureMessage(source: unknown): string { + if (isAssistantInterrupt(source)) return ASSISTANT_INTERRUPT_MESSAGE; + return classifyAssistantFailure(source).safeMessage; +} + +function extractErrorMessage(source: unknown): string { + if (typeof source === 'string') return source; + if (typeof source !== 'object' || source === null) return ''; + if ('data' in source && typeof source.data === 'object' && source.data !== null) { + if ('message' in source.data && typeof source.data.message === 'string') { + return source.data.message; + } + } + if ('message' in source && typeof source.message === 'string') return source.message; + return ''; +} diff --git a/services/cloud-agent-next/src/shared/ingest-frame.test.ts b/services/cloud-agent-next/src/shared/ingest-frame.test.ts index ac46f583fb..a4b36ffdab 100644 --- a/services/cloud-agent-next/src/shared/ingest-frame.test.ts +++ b/services/cloud-agent-next/src/shared/ingest-frame.test.ts @@ -9,6 +9,11 @@ import { } from './ingest-frame.js'; import { MAX_INLINE_FILE_URL_LENGTH } from './trim-payload.js'; import type { IngestEvent } from './protocol.js'; +import { + classifyAssistantFailure, + isAssistantInterrupt, + projectSafeAssistantError, +} from './assistant-failure.js'; describe('prepareIngestFrame', () => { it('passes small events through unchanged', () => { @@ -161,6 +166,183 @@ describe('prepareIngestFrame', () => { expect(info.parts).toBeUndefined(); }); + it.each([ + [undefined, false], + [null, false], + ['', true], + [false, true], + [0, true], + [true, true], + [1, true], + [[], true], + [{}, true], + [{ message: '' }, true], + [{ data: {} }, true], + [{ data: { message: '' } }, true], + [ + { + responseBody: 'poison-body', + headers: { authorization: 'poison-header' }, + metadata: 'poison-metadata', + }, + true, + ], + ...[ + 'ProviderAuthError', + 'UnknownError', + 'MessageOutputLengthError', + 'MessageAbortedError', + 'StructuredOutputError', + 'ContextOverflowError', + 'ContentFilterError', + 'APIError', + ].map(name => [{ name, data: {} }, true] as const), + ] as const)( + 'preserves non-null error presence through projection and compaction: %j', + (error, hasError) => { + const projected = projectSafeAssistantError(error); + expect(projected !== undefined).toBe(hasError); + const frame = prepareIngestFrame({ + streamEventType: 'kilocode', + timestamp: '2026-04-14T08:00:00.000Z', + data: { + event: 'message.updated', + properties: { + info: { + id: 'asst_presence', + role: 'assistant', + parentID: 'msg_parent', + error, + metadata: 'poison-metadata'.repeat(MAX_INGEST_EVENT_BYTES), + }, + }, + }, + }); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(true); + expect(frame.serialized).not.toContain('poison'); + const compactInfo = JSON.parse(frame.serialized).data.properties.info; + expect(compactInfo.error !== undefined && compactInfo.error !== null).toBe(hasError); + expect(compactInfo.error).toEqual(projected); + expect(isAssistantInterrupt(compactInfo.error)).toBe(isAssistantInterrupt(error)); + if (!isAssistantInterrupt(error)) { + expect(classifyAssistantFailure(compactInfo.error)).toEqual( + classifyAssistantFailure(error) + ); + } + } + ); + + it.each([ + ['APIError', 'Payment required: insufficient credits'], + ['APIError', 'Unknown model'], + ['APIError', '[BYOK] Rate limit exceeded'], + ['APIError', 'Provider timeout'], + ['ProviderAuthError', '[BYOK] Provider authentication failed'], + ['APIError', 'Invalid request'], + ['APIError', 'Service unavailable'], + ['APIError', 'Unrecognized failure'], + ['ContextOverflowError', 'Unrecognized failure'], + ['MessageOutputLengthError', '[BYOK] Unrecognized failure'], + ['ContentFilterError', 'Unrecognized failure'], + ['StructuredOutputError', 'Unrecognized failure'], + ['UnknownError', 'Unrecognized failure'], + ['MessageAbortedError', 'aborted'], + ['MessageAbortedError', '[BYOK] aborted'], + ['FutureError', 'Unrecognized failure'], + ])( + 'round-trips canonical error and interruption meaning through compaction for %s: %s', + (name, message) => { + const error = { + name, + data: { + message, + statusCode: 503, + isRetryable: false, + responseBody: 'poison-body'.repeat(MAX_INGEST_EVENT_BYTES), + responseHeaders: { authorization: 'poison-header' }, + metadata: { url: 'https://poison.example' }, + }, + }; + const safeError = projectSafeAssistantError(error); + for (const source of [error, safeError]) { + const info = { + id: 'asst_compacted', + parentID: 'msg_parent', + sessionID: 'session_root', + role: 'assistant', + modelID: 'vendor/model', + providerID: 'kilo', + error: source, + time: { created: 1, completed: 2 }, + metadata: { + prompt: 'poison-prompt'.repeat(MAX_INGEST_EVENT_BYTES), + toolArguments: 'poison-arguments', + }, + }; + const frame = prepareIngestFrame({ + streamEventType: 'kilocode', + timestamp: '2026-04-14T08:00:00.000Z', + data: { event: 'message.updated', properties: { info } }, + }); + + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(true); + expect(frame.bytes).toBeLessThanOrEqual(MAX_INGEST_EVENT_BYTES); + expect(frame.serialized).not.toContain('poison'); + const compactInfo = JSON.parse(frame.serialized).data.properties.info; + expect(compactInfo).toEqual({ + id: info.id, + parentID: info.parentID, + sessionID: info.sessionID, + role: 'assistant', + time: info.time, + error: safeError, + }); + expect(compactInfo.error).toBeTypeOf('string'); + expect(isAssistantInterrupt(compactInfo.error)).toBe(isAssistantInterrupt(error)); + expect(projectSafeAssistantError(compactInfo.error)).toBe(safeError); + if (!isAssistantInterrupt(error)) { + expect(classifyAssistantFailure(compactInfo.error)).toEqual( + classifyAssistantFailure(error) + ); + } + } + } + ); + + it.each(['legacy error with poison-message', { metadata: 'poison-metadata' }])( + 'preserves terminal error presence without SDK metadata for %j', + error => { + const frame = prepareIngestFrame({ + streamEventType: 'kilocode', + timestamp: '2026-04-14T08:00:00.000Z', + data: { + event: 'message.updated', + properties: { + info: { + id: 'asst_legacy', + role: 'assistant', + modelID: 'https://poison.example/model', + providerID: 'kilo', + error, + parts: ['poison-prompt'.repeat(MAX_INGEST_EVENT_BYTES)], + }, + }, + }, + }); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.serialized).not.toContain('poison'); + const info = JSON.parse(frame.serialized).data.properties.info; + expect(info.error).toBe('Assistant request failed'); + expect(info).not.toHaveProperty('modelID'); + expect(info).not.toHaveProperty('providerID'); + } + ); + it('still sends a compact fatal event for a terminal error with huge diagnostics', () => { const event: IngestEvent = { streamEventType: 'error', diff --git a/services/cloud-agent-next/src/shared/ingest-frame.ts b/services/cloud-agent-next/src/shared/ingest-frame.ts index 5e1c6cfbc8..717c67f8f0 100644 --- a/services/cloud-agent-next/src/shared/ingest-frame.ts +++ b/services/cloud-agent-next/src/shared/ingest-frame.ts @@ -19,6 +19,7 @@ */ import { trimPayload } from './trim-payload.js'; +import { projectSafeAssistantError } from './assistant-failure.js'; import type { IngestEvent, StreamEventType, WrapperEventTruncatedData } from './protocol.js'; /** Cloudflare's per-WebSocket-message receive limit. Documentation/logging only. */ @@ -185,7 +186,7 @@ function compactMessageUpdated(data: Record): Record { }, ]); expect(JSON.stringify(reports)).not.toContain('never report'); + expect(JSON.stringify(reports)).not.toContain('secret'); expect(JSON.stringify(reports)).not.toContain('model/test'); }); + it('keeps diagnostic expiry tied to terminal time when a failed run is reported again later', async () => { + const reports: CloudAgentQueueReport[] = []; + for (const occurredAt of [6, 60 * 24 * 60 * 60 * 1000]) { + await emitRunStateReport({ + queue: { send: async report => void reports.push(report) }, + cloudAgentSessionId: 'agent_report', + state: { + ...state, + failureCode: 'assistant_error', + assistantFailureReason: 'timeout', + providerOwnership: 'managed', + }, + occurredAt, + }); + } + + expect(reports).toHaveLength(2); + expect(reports[1].run).toEqual(reports[0].run); + expect(reports[1].run).toMatchObject({ + failureReason: 'request_timeout', + diagnostic: { + errorMessageRedacted: 'Assistant request timed out', + errorExpiresAt: new Date(5 + 30 * 24 * 60 * 60 * 1000).toISOString(), + }, + }); + }); + it.each([ ['agent_activity', 'payment_required', 'insufficient_credits'], ['agent_activity', 'model_missing', 'model_unavailable'], @@ -106,6 +134,176 @@ describe('Cloud Agent report emitter', () => { } ); + it('reports pre-dispatch payment_required without acceptance or activity timestamps', async () => { + const reports: CloudAgentQueueReport[] = []; + await emitRunStateReport({ + queue: { send: async report => void reports.push(report) }, + cloudAgentSessionId: 'agent_report', + state: { + ...state, + acceptedAt: undefined, + dispatchAcceptanceKind: undefined, + agentActivityObservedAt: undefined, + wrapperRunId: undefined, + completionSource: 'delivery_failure', + failureStage: 'pre_dispatch', + failureCode: 'payment_required', + }, + }); + + expect(reports).toHaveLength(1); + expect(reports[0].run).toMatchObject({ + status: 'failed', + failureStage: 'pre_dispatch', + failureCode: 'payment_required', + failureResponsibility: 'user', + failureReason: 'insufficient_credits', + diagnostic: { errorMessageRedacted: 'Model request failed: insufficient credits' }, + }); + expect(reports[0].run).not.toHaveProperty('dispatchAcceptedAt'); + expect(reports[0].run).not.toHaveProperty('agentActivityObservedAt'); + }); + + it.each([ + ['wrapper_ping_timeout', 'Wrapper health check timed out'], + ['wrapper_no_output', 'Wrapper made no execution progress during the watchdog window'], + ['wrapper_disconnected', 'Wrapper disconnected before completion'], + ] as const)('reports %s before and after activity', async (failureCode, expectedDiagnostic) => { + for (const failureStage of ['post_dispatch_no_activity', 'agent_activity'] as const) { + const reports: CloudAgentQueueReport[] = []; + await emitRunStateReport({ + queue: { send: async report => void reports.push(report) }, + cloudAgentSessionId: 'agent_report', + state: { + ...state, + failureStage, + failureCode, + agentActivityObservedAt: failureStage === 'agent_activity' ? 4 : undefined, + }, + }); + + expect(reports).toHaveLength(1); + expect(reports[0].run).toMatchObject({ + failureStage, + failureCode, + failureResponsibility: 'platform', + failureReason: 'wrapper_liveness', + diagnostic: { errorMessageRedacted: expectedDiagnostic }, + }); + if (failureStage === 'agent_activity') { + expect(reports[0].run.agentActivityObservedAt).toBe(new Date(4).toISOString()); + } else { + expect(reports[0].run).not.toHaveProperty('agentActivityObservedAt'); + } + } + }); + + it.each([ + [ + 'insufficient_credits', + 'Assistant request failed: insufficient credits', + 'insufficient_credits', + ], + ['rate_limited', 'Assistant request was rate limited', 'rate_limited'], + ['model_unavailable', 'Assistant request failed: model not found', 'model_unavailable'], + [ + 'provider_authentication', + 'Assistant request was not authorized', + 'managed_provider_authentication', + ], + ['provider_unavailable', 'Assistant service is unavailable', 'managed_provider_unavailable'], + ['timeout', 'Assistant request timed out', 'request_timeout'], + ['invalid_request', 'Assistant request was invalid', 'invalid_request'], + ['context_limit', 'The model context limit was exceeded', 'context_limit'], + ['output_limit', 'The model output limit was reached', 'output_limit'], + [ + 'content_filter', + 'The model provider blocked the response under its content policy', + 'content_filter', + ], + [ + 'structured_output', + 'The model response did not match the required format', + 'structured_output', + ], + ['unknown', 'Assistant request failed', 'assistant_unknown'], + [undefined, 'Assistant request failed', 'assistant_unknown'], + ] as const)( + 'uses only fixed diagnostic wording for assistant reason %s', + async (assistantFailureReason, expectedDiagnostic, expectedFailureReason) => { + const reports: CloudAgentQueueReport[] = []; + await emitRunStateReport({ + queue: { send: async report => void reports.push(report) }, + cloudAgentSessionId: 'agent_report', + state: { + ...state, + failureCode: 'assistant_error', + assistantFailureReason, + providerOwnership: 'managed', + error: 'Payment Required: Authorization: Bearer raw-secret-token', + safeFailureMessage: 'model not found: api_key=safe-message-secret', + }, + }); + + expect(reports).toHaveLength(1); + expect(reports[0].run.diagnostic).toEqual({ + errorMessageRedacted: expectedDiagnostic, + errorExpiresAt: new Date(5 + 30 * 24 * 60 * 60 * 1000).toISOString(), + }); + expect(reports[0].run.failureReason).toBe(expectedFailureReason); + expect(JSON.stringify(reports)).not.toContain('raw-secret-token'); + expect(JSON.stringify(reports)).not.toContain('safe-message-secret'); + expect(JSON.stringify(reports)).not.toContain('model/test'); + } + ); + + it.each([ + ['payment_required', 'Model request failed: insufficient credits'], + ['model_missing', 'No model is available for this run'], + ['wrapper_error_after_activity', 'Wrapper failed after agent activity'], + ] as const)( + 'keeps the %s diagnostic ahead of assistant reason wording', + async (failureCode, expectedDiagnostic) => { + const reports: CloudAgentQueueReport[] = []; + await emitRunStateReport({ + queue: { send: async report => void reports.push(report) }, + cloudAgentSessionId: 'agent_report', + state: { + ...state, + failureCode, + assistantFailureReason: 'timeout', + safeFailureMessage: 'Authorization: Bearer safe-message-secret', + }, + }); + + expect(reports[0]?.run.diagnostic?.errorMessageRedacted).toBe(expectedDiagnostic); + expect(JSON.stringify(reports)).not.toContain('safe-message-secret'); + } + ); + + it.each([ + 'Insufficient credits', + ' Insufficient credits: PAYMENT_REQUIRED ', + 'Insufficient credits: insufficient_funds', + 'Payment Required', + ])('keeps the exact legacy credit fallback ahead of assistant reasons: %s', async error => { + const reports: CloudAgentQueueReport[] = []; + await emitRunStateReport({ + queue: { send: async report => void reports.push(report) }, + cloudAgentSessionId: 'agent_report', + state: { + ...state, + failureCode: 'assistant_error', + assistantFailureReason: 'timeout', + error, + }, + }); + + expect(reports[0]?.run.diagnostic?.errorMessageRedacted).toBe( + 'Model request failed: insufficient credits' + ); + }); + it('attributes an absent model chosen by managed auto-routing to platform configuration', async () => { const reports: CloudAgentQueueReport[] = []; await emitRunStateReport({ diff --git a/services/cloud-agent-next/src/telemetry/queue-reports.ts b/services/cloud-agent-next/src/telemetry/queue-reports.ts index 76cfb5ff69..53b9402c4b 100644 --- a/services/cloud-agent-next/src/telemetry/queue-reports.ts +++ b/services/cloud-agent-next/src/telemetry/queue-reports.ts @@ -5,7 +5,10 @@ import { type CloudAgentRunStateReport, } from '@kilocode/worker-utils/cloud-agent-queue-report'; import { logger } from '../logger.js'; -import { workspaceFailureMessage } from '../session/safe-failure-projection.js'; +import { + assistantFailureMessage, + workspaceFailureMessage, +} from '../session/safe-failure-projection.js'; import { admittedAgentModel, type SessionMessageState } from '../session/session-message-state.js'; import { classifyCloudAgentFailure, @@ -40,7 +43,7 @@ const FAILED_RUN_DIAGNOSTIC_MESSAGES: Partial< model_missing: 'No model is available for this run', delivery_failure_unknown: 'Message delivery outcome is unknown', wrapper_disconnected: 'Wrapper disconnected before completion', - wrapper_no_output: 'Wrapper produced no output before timeout', + wrapper_no_output: 'Wrapper made no execution progress during the watchdog window', wrapper_ping_timeout: 'Wrapper health check timed out', wrapper_error_before_activity: 'Wrapper failed before agent activity', assistant_error: 'Assistant request failed', @@ -84,8 +87,8 @@ function diagnosticForFailedRun( : 'Workspace setup failed'; } else if (isKnownInsufficientCreditFailure(state)) { errorMessageRedacted = 'Model request failed: insufficient credits'; - } else if (state.assistantFailureReason === 'rate_limited') { - errorMessageRedacted = 'Assistant request was rate limited'; + } else if (state.failureCode === 'assistant_error') { + errorMessageRedacted = assistantFailureMessage(state.assistantFailureReason ?? 'unknown'); } if (errorMessageRedacted === undefined) return undefined; diff --git a/services/cloud-agent-next/src/telemetry/report-consumer.test.ts b/services/cloud-agent-next/src/telemetry/report-consumer.test.ts index 1ff88b37b9..81ed7bd032 100644 --- a/services/cloud-agent-next/src/telemetry/report-consumer.test.ts +++ b/services/cloud-agent-next/src/telemetry/report-consumer.test.ts @@ -25,8 +25,14 @@ const report = { }, } as const; -function makeMessage(body: unknown) { - return { body, ack: vi.fn(), retry: vi.fn() }; +const diagnostic = { + errorMessageRedacted: 'The model provider is unavailable', + errorExpiresAt: '2026-06-25T08:04:00.000Z', +}; +const retiredFacts = { version: 1, sdkStatusCode: 503 }; + +function makeMessage(body: unknown, attempts = 1) { + return { body, attempts, ack: vi.fn(), retry: vi.fn() }; } const env = { @@ -57,34 +63,87 @@ describe('Cloud Agent report consumer', () => { expect(message.retry).not.toHaveBeenCalled(); }); - it('acks an expired saved report', async () => { + it('acks an expired saved report after discarding retired diagnostic data', async () => { const saveReport = vi.fn().mockResolvedValueOnce({ outcome: 'expired' }); vi.mocked(createCloudAgentReportStore).mockReturnValue({ saveReport } as never); - const expired = makeMessage(report); + const expired = makeMessage({ + ...report, + run: { ...report.run, diagnostic: { ...diagnostic, facts: retiredFacts } }, + }); await consumeCloudAgentReportBatch( { messages: [expired] } as unknown as MessageBatch, env ); + expect(saveReport).toHaveBeenCalledExactlyOnceWith({ + ...report, + run: { ...report.run, diagnostic }, + }); expect(expired.ack).toHaveBeenCalledOnce(); + expect(expired.retry).not.toHaveBeenCalled(); }); - it('logs and acknowledges reports whose session anchor is absent', async () => { - const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); - const saveReport = vi.fn().mockResolvedValueOnce({ outcome: 'missing_parent' }); + it.each([1, 4])( + 'retries missing parents on delivery attempt %s using the queue policy', + async attempt => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + const saveReport = vi.fn().mockResolvedValueOnce({ outcome: 'missing_parent' }); + vi.mocked(createCloudAgentReportStore).mockReturnValue({ saveReport } as never); + const message = makeMessage({ ...report, run: { ...report.run, diagnostic } }, attempt); + + await consumeCloudAgentReportBatch( + { messages: [message] } as unknown as MessageBatch, + env + ); + + expect(warn).toHaveBeenCalledExactlyOnceWith( + 'Retrying Cloud Agent run report without a session anchor', + { + cloudAgentSessionId: report.session.cloudAgentSessionId, + messageId: report.run.messageId, + status: report.run.status, + attempt, + } + ); + expect(message.retry).toHaveBeenCalledExactlyOnceWith(); + expect(message.ack).not.toHaveBeenCalled(); + } + ); + + it('continues the batch after a missing parent and acknowledges a later successful redelivery', async () => { + vi.spyOn(console, 'warn').mockImplementation(() => undefined); + const saveReport = vi + .fn() + .mockResolvedValueOnce({ outcome: 'missing_parent' }) + .mockResolvedValueOnce({ outcome: 'applied' }) + .mockResolvedValueOnce({ outcome: 'expired' }) + .mockResolvedValueOnce({ outcome: 'applied' }); vi.mocked(createCloudAgentReportStore).mockReturnValue({ saveReport } as never); - const message = makeMessage(report); + const missing = makeMessage(report); + const applied = makeMessage(report); + const expired = makeMessage(report); + const malformed = makeMessage({ version: 99 }); await consumeCloudAgentReportBatch( - { messages: [message] } as unknown as MessageBatch, + { messages: [missing, applied, expired, malformed] } as unknown as MessageBatch, env ); - expect(warn).toHaveBeenCalledWith('Dropping Cloud Agent run report without a session anchor', { - cloudAgentSessionId: report.session.cloudAgentSessionId, - }); - expect(message.ack).toHaveBeenCalledOnce(); + expect(missing.retry).toHaveBeenCalledOnce(); + expect(missing.ack).not.toHaveBeenCalled(); + for (const message of [applied, expired, malformed]) { + expect(message.ack).toHaveBeenCalledOnce(); + expect(message.retry).not.toHaveBeenCalled(); + } + const redelivery = makeMessage(report, 2); + await consumeCloudAgentReportBatch( + { messages: [redelivery] } as unknown as MessageBatch, + env + ); + expect(redelivery.ack).toHaveBeenCalledOnce(); + expect(redelivery.retry).not.toHaveBeenCalled(); + expect(saveReport).toHaveBeenCalledTimes(4); }); it('drops malformed messages without logging their body', async () => { @@ -106,19 +165,103 @@ describe('Cloud Agent report consumer', () => { }); }); - it('discards an invalid optional diagnostic while saving a typed failed outcome', async () => { - vi.spyOn(console, 'warn').mockImplementation(() => undefined); - const saveReport = vi.fn(async () => ({ outcome: 'applied' as const })); - vi.mocked(createCloudAgentReportStore).mockReturnValue({ saveReport } as never); - const message = makeMessage({ - ...report, - run: { - ...report.run, - diagnostic: { - errorMessageRedacted: 'too late but not saved', - errorExpiresAt: report.run.terminalAt, + it.each([ + { name: 'current', diagnostic }, + { name: 'older', diagnostic: { ...diagnostic, facts: retiredFacts } }, + { + name: 'opaque retired data', + diagnostic: { ...diagnostic, facts: 'discarded fixture payload' }, + }, + { name: 'null retired data', diagnostic: { ...diagnostic, facts: null } }, + ])( + 'preserves safe diagnostic text and typed outcomes for $name reports without retaining retired data', + async ({ diagnostic: incomingDiagnostic }) => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + const saveReport = vi.fn(async () => ({ outcome: 'applied' as const })); + vi.mocked(createCloudAgentReportStore).mockReturnValue({ saveReport } as never); + const message = makeMessage({ + ...report, + run: { ...report.run, diagnostic: incomingDiagnostic }, + }); + + await consumeCloudAgentReportBatch( + { messages: [message] } as unknown as MessageBatch, + env + ); + + expect(saveReport).toHaveBeenCalledExactlyOnceWith({ + ...report, + run: { ...report.run, diagnostic }, + }); + expect(warn).not.toHaveBeenCalled(); + expect(message.ack).toHaveBeenCalledOnce(); + expect(message.retry).not.toHaveBeenCalled(); + } + ); + + it.each([ + { errorMessageRedacted: '' }, + { errorMessageRedacted: 'm'.repeat(4097) }, + { errorExpiresAt: 'invalid timestamp' }, + { errorExpiresAt: report.run.terminalAt }, + { errorExpiresAt: '2026-06-26T08:04:00.000Z' }, + { responseBody: 'private fixture output' }, + ])( + 'retains whole-diagnostic fallback for invalid message, expiry, or unknown content: %j', + async invalidDiagnostic => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + const saveReport = vi.fn(async () => ({ outcome: 'applied' as const })); + vi.mocked(createCloudAgentReportStore).mockReturnValue({ saveReport } as never); + const message = makeMessage({ + ...report, + run: { + ...report.run, + diagnostic: { ...diagnostic, facts: retiredFacts, ...invalidDiagnostic }, }, + }); + + await consumeCloudAgentReportBatch( + { messages: [message] } as unknown as MessageBatch, + env + ); + + expect(saveReport).toHaveBeenCalledExactlyOnceWith(report); + expect(warn).toHaveBeenCalledExactlyOnceWith( + 'Dropping invalid Cloud Agent report diagnostic' + ); + expect(message.ack).toHaveBeenCalledOnce(); + expect(message.retry).not.toHaveBeenCalled(); + } + ); + + it.each([ + { + name: 'missing terminal time', + body: { ...report, run: { ...report.run, terminalAt: undefined } }, + }, + { + name: 'invalid classification', + body: { + ...report, + run: { ...report.run, failureStage: 'pre_dispatch', failureCode: 'assistant_error' }, }, + }, + { name: 'unknown envelope field', body: { ...report, metadata: 'private fixture' } }, + { + name: 'unknown session field', + body: { ...report, session: { ...report.session, metadata: 'private fixture' } }, + }, + { + name: 'unknown run field', + body: { ...report, run: { ...report.run, metadata: 'private fixture' } }, + }, + ])('rejects $name even with retired diagnostic data', async ({ body }) => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + const saveReport = vi.fn(); + vi.mocked(createCloudAgentReportStore).mockReturnValue({ saveReport } as never); + const message = makeMessage({ + ...body, + run: { ...body.run, diagnostic: { ...diagnostic, facts: retiredFacts } }, }); await consumeCloudAgentReportBatch( @@ -126,7 +269,10 @@ describe('Cloud Agent report consumer', () => { env ); - expect(saveReport).toHaveBeenCalledWith(report); + expect(saveReport).not.toHaveBeenCalled(); + expect(warn).toHaveBeenCalledExactlyOnceWith('Dropping malformed Cloud Agent report message', { + issueCount: expect.any(Number), + }); expect(message.ack).toHaveBeenCalledOnce(); expect(message.retry).not.toHaveBeenCalled(); }); diff --git a/services/cloud-agent-next/src/telemetry/report-consumer.ts b/services/cloud-agent-next/src/telemetry/report-consumer.ts index 677223630b..64844067f2 100644 --- a/services/cloud-agent-next/src/telemetry/report-consumer.ts +++ b/services/cloud-agent-next/src/telemetry/report-consumer.ts @@ -11,10 +11,6 @@ export const CLOUD_AGENT_REPORT_QUEUE_NAMES = new Set([ ]); function parseReportWithoutInvalidDiagnostic(body: unknown) { - const parsed = CloudAgentQueueReportSchema.safeParse(body); - if (parsed.success) { - return parsed; - } if ( typeof body !== 'object' || body === null || @@ -25,12 +21,28 @@ function parseReportWithoutInvalidDiagnostic(body: unknown) { body.run === null || !('diagnostic' in body.run) ) { + return CloudAgentQueueReportSchema.safeParse(body); + } + const diagnostic = body.run.diagnostic; + const run = + typeof diagnostic === 'object' && diagnostic !== null && 'facts' in diagnostic + ? { + ...body.run, + diagnostic: Object.fromEntries( + Object.entries(diagnostic).filter(([fieldName]) => fieldName !== 'facts') + ), + } + : body.run; + const parsed = CloudAgentQueueReportSchema.safeParse({ ...body, run }); + if (parsed.success) { return parsed; } - const run = Object.fromEntries( - Object.entries(body.run).filter(([fieldName]) => fieldName !== 'diagnostic') - ); - const typedOnly = CloudAgentQueueReportSchema.safeParse({ ...body, run }); + const typedOnly = CloudAgentQueueReportSchema.safeParse({ + ...body, + run: Object.fromEntries( + Object.entries(run).filter(([fieldName]) => fieldName !== 'diagnostic') + ), + }); if (typedOnly.success) { console.warn('Dropping invalid Cloud Agent report diagnostic'); } @@ -55,9 +67,14 @@ export async function consumeCloudAgentReportBatch( try { const result = await reportStore.saveReport(parsed.data); if (result.outcome === 'missing_parent') { - console.warn('Dropping Cloud Agent run report without a session anchor', { + console.warn('Retrying Cloud Agent run report without a session anchor', { cloudAgentSessionId: parsed.data.session.cloudAgentSessionId, + messageId: parsed.data.run.messageId, + status: parsed.data.run.status, + attempt: message.attempts, }); + message.retry(); + continue; } message.ack(); } catch { diff --git a/services/cloud-agent-next/src/telemetry/report-store.test.ts b/services/cloud-agent-next/src/telemetry/report-store.test.ts index de3b5491b0..c1d4d997a1 100644 --- a/services/cloud-agent-next/src/telemetry/report-store.test.ts +++ b/services/cloud-agent-next/src/telemetry/report-store.test.ts @@ -2,10 +2,46 @@ import { describe, expect, it, vi } from 'vitest'; import type { SQL } from 'drizzle-orm'; import { getWorkerDb } from '@kilocode/db/client'; import { cloud_agent_session_runs, cloud_agent_sessions } from '@kilocode/db/schema'; +import type { CloudAgentRunStateReport } from '@kilocode/worker-utils/cloud-agent-queue-report'; import { createCloudAgentReportStore } from './report-store.js'; const cloudAgentSessionId = 'agent_12345678-1234-4234-8234-123456789abc'; const occurredAt = '2026-05-25T12:00:00.000Z'; +const diagnostic = { + errorMessageRedacted: 'The model provider is unavailable', + errorExpiresAt: '2026-06-01T12:00:00.000Z', +}; +const failedReport = { + version: 1, + type: 'run.state', + occurredAt, + session: { cloudAgentSessionId }, + run: { + messageId: 'msg_failed', + wrapperRunId: 'wr_first', + status: 'failed', + terminalAt: occurredAt, + failureStage: 'agent_activity', + failureCode: 'assistant_error', + failureResponsibility: 'unknown', + failureReason: 'provider_unavailable', + diagnostic, + }, +} satisfies CloudAgentRunStateReport; +const storedFailedRun = { + status: 'failed', + wrapperRunId: 'wr_first', + queuedAt: occurredAt, + dispatchAcceptedAt: null, + agentActivityObservedAt: null, + terminalAt: occurredAt, + failureStage: 'agent_activity', + failureCode: 'assistant_error', + failureResponsibility: 'unknown', + failureReason: 'provider_unavailable', + errorMessageRedacted: null, + errorExpiresAt: null, +}; function makeDb( selectResults: unknown[][] = [], @@ -20,6 +56,8 @@ function makeDb( }> = []; const updates: Array<{ table: unknown; values?: Record; where?: SQL }> = []; const deletes: unknown[] = []; + const deleteConditions: SQL[] = []; + const selects: Array<{ fields: Record; table?: unknown; where?: SQL }> = []; const operations: string[] = []; const execute = vi.fn(async () => operations.push('execute')); @@ -82,7 +120,8 @@ function makeDb( operations.push('delete'); deletes.push(table); const chain = { - where() { + where(condition: SQL) { + deleteConditions.push(condition); return chain; }, returning() { @@ -96,13 +135,17 @@ function makeDb( return chain; } - function select() { + function select(fields: Record) { + const call: { fields: Record; table?: unknown; where?: SQL } = { fields }; + selects.push(call); const result = selectResults.shift() ?? []; const chain = { - from() { + from(table: unknown) { + call.table = table; return chain; }, - where() { + where(condition: SQL) { + call.where = condition; return chain; }, limit() { @@ -131,6 +174,8 @@ function makeDb( inserts, updates, deletes, + deleteConditions, + selects, operations, }; } @@ -302,6 +347,8 @@ describe('cloud agent reporting store', () => { failureCode: 'unclassified', failureResponsibility: null, failureReason: null, + errorMessageRedacted: null, + errorExpiresAt: null, }, ], ]); @@ -348,6 +395,8 @@ describe('cloud agent reporting store', () => { failureCode: 'unclassified', failureResponsibility: 'unknown', failureReason: 'unclassified', + errorMessageRedacted: diagnostic.errorMessageRedacted, + errorExpiresAt: diagnostic.errorExpiresAt, }, ], ]); @@ -375,10 +424,438 @@ describe('cloud agent reporting store', () => { dispatch_accepted_at: '2026-05-25T12:02:00.000Z', terminal_at: '2026-05-25T12:04:00.000Z', failure_code: 'unclassified', + error_message_redacted: diagnostic.errorMessageRedacted, + error_expires_at: diagnostic.errorExpiresAt, + }); + }); + + it('stores diagnostics using only the matching session and message', async () => { + const fake = makeDb([[{ createdAt: occurredAt }], []]); + const store = createCloudAgentReportStore(fake.db as never); + + expect(await store.saveReport(failedReport, occurredAt)).toEqual({ outcome: 'applied' }); + expect(fake.inserts[0]?.values).toMatchObject({ + cloud_agent_session_id: cloudAgentSessionId, + message_id: failedReport.run.messageId, + status: 'failed', + error_message_redacted: diagnostic.errorMessageRedacted, + error_expires_at: diagnostic.errorExpiresAt, + }); + const selection = fake.selects.find(call => call.table === cloud_agent_session_runs); + expect(selection?.fields.errorMessageRedacted).toBe( + cloud_agent_session_runs.error_message_redacted + ); + expect(selection?.fields.errorExpiresAt).toBe(cloud_agent_session_runs.error_expires_at); + if (!selection?.where) throw new Error('expected run selection predicate'); + const query = getWorkerDb('postgres://unused:unused@localhost:0/unused') + .select() + .from(cloud_agent_session_runs) + .where(selection.where) + .toSQL(); + expect(query.sql).toMatch( + /"cloud_agent_session_runs"\."cloud_agent_session_id" = \$1 and "cloud_agent_session_runs"\."message_id" = \$2/ + ); + expect(query.params).toEqual([cloudAgentSessionId, failedReport.run.messageId]); + }); + + it.each([ + ['missing', undefined], + ['empty message', { ...diagnostic, errorMessageRedacted: '' }], + ['oversized message', { ...diagnostic, errorMessageRedacted: 'm'.repeat(4097) }], + ['invalid expiry', { ...diagnostic, errorExpiresAt: 'invalid timestamp' }], + ['expiry at terminal time', { ...diagnostic, errorExpiresAt: occurredAt }], + ['expiry beyond 30 days', { ...diagnostic, errorExpiresAt: '2026-06-25T12:00:00.000Z' }], + ])('does not store a %s diagnostic or lose the outcome', async (_name, invalidDiagnostic) => { + const fake = makeDb([[{ createdAt: occurredAt }], []]); + const store = createCloudAgentReportStore(fake.db as never); + const report: CloudAgentRunStateReport = { + ...failedReport, + run: { + ...failedReport.run, + diagnostic: invalidDiagnostic as CloudAgentRunStateReport['run']['diagnostic'], + }, + }; + + expect(await store.saveReport(report, occurredAt)).toEqual({ outcome: 'applied' }); + expect(fake.inserts[0]?.values).toMatchObject({ status: 'failed' }); + expect(fake.inserts[0]?.values).not.toHaveProperty('error_message_redacted'); + expect(fake.inserts[0]?.values).not.toHaveProperty('error_expires_at'); + }); + + it.each(['queued', 'accepted', 'completed', 'interrupted'] as const)( + 'never stores diagnostics on a %s run', + async status => { + const fake = makeDb([[{ createdAt: occurredAt }], []]); + const store = createCloudAgentReportStore(fake.db as never); + await store.saveReport( + { + ...failedReport, + run: { + messageId: failedReport.run.messageId, + status, + dispatchAcceptedAt: status === 'accepted' ? occurredAt : undefined, + terminalAt: status === 'completed' || status === 'interrupted' ? occurredAt : undefined, + diagnostic, + }, + }, + occurredAt + ); + + expect(fake.inserts[0]?.values).toMatchObject({ status }); + expect(fake.inserts[0]?.values).not.toHaveProperty('error_message_redacted'); + expect(fake.inserts[0]?.values).not.toHaveProperty('error_expires_at'); + } + ); + + it.each([null, '2026-05-29 12:00:00+00'])( + 'fills missing metadata without losing established diagnostic content or extending expiry %s', + async errorExpiresAt => { + const fake = makeDb([ + [{ createdAt: occurredAt }], + [ + { + ...storedFailedRun, + wrapperRunId: null, + terminalAt: '2026-05-25 12:00:00+00', + failureResponsibility: null, + failureReason: null, + errorMessageRedacted: errorExpiresAt ? 'Original diagnostic' : null, + errorExpiresAt, + }, + ], + ]); + const store = createCloudAgentReportStore(fake.db as never); + await store.saveReport(failedReport, occurredAt); + + expect(fake.updates[0]?.values).toMatchObject({ + status: 'failed', + wrapper_run_id: 'wr_first', + terminal_at: '2026-05-25 12:00:00+00', + failure_responsibility: failedReport.run.failureResponsibility, + failure_reason: failedReport.run.failureReason, + error_message_redacted: errorExpiresAt + ? 'Original diagnostic' + : diagnostic.errorMessageRedacted, + error_expires_at: errorExpiresAt ?? diagnostic.errorExpiresAt, + }); + } + ); + + it.each(['queued', 'accepted'] as const)( + 'adds diagnostics on the first terminal report after %s', + async status => { + const fake = makeDb([ + [{ createdAt: occurredAt }], + [ + { + ...storedFailedRun, + status, + terminalAt: null, + failureStage: null, + failureCode: null, + failureResponsibility: null, + failureReason: null, + }, + ], + ]); + const store = createCloudAgentReportStore(fake.db as never); + await store.saveReport(failedReport, occurredAt); + expect(fake.updates[0]?.values).toMatchObject({ + status: 'failed', + terminal_at: occurredAt, + error_message_redacted: diagnostic.errorMessageRedacted, + error_expires_at: diagnostic.errorExpiresAt, + }); + } + ); + + it.each([ + { name: 'terminal status', existing: { status: 'completed' } }, + { name: 'failure code', existing: { failureCode: 'wrapper_no_output' } }, + { name: 'responsibility', existing: { failureResponsibility: 'platform' } }, + { name: 'reason', existing: { failureReason: 'assistant_unknown' } }, + { name: 'wrapper identity', existing: { wrapperRunId: 'wr_other' } }, + { name: 'earlier terminal timestamp', existing: { terminalAt: '2026-05-25T11:59:00.000Z' } }, + { name: 'later terminal timestamp', existing: { terminalAt: '2026-05-25T12:01:00.000Z' } }, + { + name: 'unidentified terminal timestamp', + existing: { wrapperRunId: null, terminalAt: '2026-05-25T11:59:00.000Z' }, + }, + ])('does not fill diagnostics from a replay with conflicting $name', async ({ existing }) => { + const fake = makeDb([[{ createdAt: occurredAt }], [{ ...storedFailedRun, ...existing }]]); + const store = createCloudAgentReportStore(fake.db as never); + expect(await store.saveReport(failedReport, occurredAt)).toEqual({ outcome: 'applied' }); + expect(fake.updates[0]?.values).toMatchObject({ + error_message_redacted: null, + error_expires_at: null, + }); + }); + + it('keeps the first terminal timestamp across two mismatched replays without admitting rejected diagnostics', async () => { + const existing = { + ...storedFailedRun, + queuedAt: '2026-05-25T11:57:30.000Z', + dispatchAcceptedAt: '2026-05-25T11:58:30.000Z', + agentActivityObservedAt: '2026-05-25T11:59:30.000Z', + terminalAt: '2026-05-25T12:01:00.000Z', + }; + const report = { + ...failedReport, + run: { + ...failedReport.run, + queuedAt: '2026-05-25T11:57:00.000Z', + dispatchAcceptedAt: '2026-05-25T11:58:00.000Z', + agentActivityObservedAt: '2026-05-25T11:59:00.000Z', + }, + }; + const selectResults: unknown[][] = [[{ createdAt: occurredAt }], [existing]]; + const fake = makeDb(selectResults); + const store = createCloudAgentReportStore(fake.db as never); + await store.saveReport(report, '2026-05-25T12:02:00.000Z'); + const firstUpdate = fake.updates[0]?.values; + if (!firstUpdate) throw new Error('expected first replay update'); + selectResults.push( + [{ createdAt: occurredAt }], + [ + { + ...existing, + queuedAt: firstUpdate.queued_at, + dispatchAcceptedAt: firstUpdate.dispatch_accepted_at, + agentActivityObservedAt: firstUpdate.agent_activity_observed_at, + terminalAt: firstUpdate.terminal_at, + errorMessageRedacted: firstUpdate.error_message_redacted, + errorExpiresAt: firstUpdate.error_expires_at, + }, + ] + ); + await store.saveReport(report, '2026-05-25T12:02:00.000Z'); + + expect(fake.updates).toHaveLength(2); + for (const update of fake.updates) { + expect(update.values).toMatchObject({ + status: 'failed', + queued_at: report.run.queuedAt, + dispatch_accepted_at: report.run.dispatchAcceptedAt, + agent_activity_observed_at: report.run.agentActivityObservedAt, + terminal_at: existing.terminalAt, + error_message_redacted: null, + error_expires_at: null, + }); + } + }); + + it.each(['failed', 'completed', 'interrupted'] as const)( + 'fills a null legacy terminal time only when the failed replay matches status %s', + async status => { + const fake = makeDb([ + [{ createdAt: occurredAt }], + [{ ...storedFailedRun, status, terminalAt: null }], + ]); + const store = createCloudAgentReportStore(fake.db as never); + await store.saveReport(failedReport, occurredAt); + expect(fake.updates[0]?.values).toMatchObject({ + status, + terminal_at: status === 'failed' ? occurredAt : null, + error_message_redacted: status === 'failed' ? diagnostic.errorMessageRedacted : null, + error_expires_at: status === 'failed' ? diagnostic.errorExpiresAt : null, + }); + } + ); + + it('does not fill diagnostics for the same failure code at a different stage', async () => { + const fake = makeDb([ + [{ createdAt: occurredAt }], + [{ ...storedFailedRun, failureStage: 'pre_dispatch', failureCode: 'payment_required' }], + ]); + const store = createCloudAgentReportStore(fake.db as never); + await store.saveReport( + { + ...failedReport, + run: { ...failedReport.run, failureCode: 'payment_required' }, + }, + occurredAt + ); + expect(fake.updates[0]?.values).toMatchObject({ + failure_stage: 'pre_dispatch', + failure_code: 'payment_required', + error_message_redacted: null, + error_expires_at: null, }); }); - it('clears expired sanitized detail and purges rows older than 90 days', async () => { + it.each([ + undefined, + { ...diagnostic, errorMessageRedacted: 'The agent failed' }, + { ...diagnostic, errorMessageRedacted: '' }, + { ...diagnostic, errorExpiresAt: 'invalid timestamp' }, + { ...diagnostic, errorExpiresAt: '2026-06-24T12:00:00.000Z' }, + ])( + 'preserves the first diagnostic on absent, conflicting, invalid, or later-expiry replays: %j', + async incomingDiagnostic => { + const fake = makeDb([ + [{ createdAt: occurredAt }], + [ + { + ...storedFailedRun, + errorMessageRedacted: diagnostic.errorMessageRedacted, + errorExpiresAt: diagnostic.errorExpiresAt, + }, + ], + ]); + const store = createCloudAgentReportStore(fake.db as never); + await store.saveReport( + { + ...failedReport, + run: { ...failedReport.run, diagnostic: incomingDiagnostic }, + }, + occurredAt + ); + expect(fake.updates[0]?.values).toMatchObject({ + error_message_redacted: diagnostic.errorMessageRedacted, + error_expires_at: diagnostic.errorExpiresAt, + }); + } + ); + + it('does not persist diagnostic text that has already expired', async () => { + const fake = makeDb([[{ createdAt: occurredAt }], []]); + const store = createCloudAgentReportStore(fake.db as never); + expect(await store.saveReport(failedReport, diagnostic.errorExpiresAt)).toEqual({ + outcome: 'applied', + }); + expect(fake.inserts[0]?.values).toMatchObject({ status: 'failed' }); + expect(fake.inserts[0]?.values).not.toHaveProperty('error_message_redacted'); + expect(fake.inserts[0]?.values).not.toHaveProperty('error_expires_at'); + }); + + it('clears expired diagnostic text and expiry across repeated reports', async () => { + const existing = { + ...storedFailedRun, + errorMessageRedacted: diagnostic.errorMessageRedacted, + errorExpiresAt: diagnostic.errorExpiresAt, + }; + const selectResults: unknown[][] = [[{ createdAt: occurredAt }], [existing]]; + const fake = makeDb(selectResults); + const store = createCloudAgentReportStore(fake.db as never); + await store.saveReport(failedReport, diagnostic.errorExpiresAt); + const firstUpdate = fake.updates[0]?.values; + if (!firstUpdate) throw new Error('expected first replay update'); + selectResults.push( + [{ createdAt: occurredAt }], + [ + { + ...existing, + errorMessageRedacted: firstUpdate.error_message_redacted, + errorExpiresAt: firstUpdate.error_expires_at, + }, + ] + ); + await store.saveReport(failedReport, diagnostic.errorExpiresAt); + + expect(fake.updates).toHaveLength(2); + for (const update of fake.updates) { + expect(update.values).toMatchObject({ + status: 'failed', + error_message_redacted: null, + error_expires_at: null, + }); + } + }); + + it('does not revive expired diagnostics after cleanup clears their expiry', async () => { + const existing = { + ...storedFailedRun, + errorMessageRedacted: diagnostic.errorMessageRedacted, + errorExpiresAt: '2026-06-01 12:00:00+00', + }; + const selectResults: unknown[][] = []; + const fake = makeDb(selectResults); + const store = createCloudAgentReportStore(fake.db as never); + await store.removeExpiredData(diagnostic.errorExpiresAt); + const cleanup = fake.updates.find(update => update.table === cloud_agent_session_runs)?.values; + if (!cleanup) throw new Error('expected expired run cleanup'); + selectResults.push( + [{ createdAt: occurredAt }], + [ + { + ...existing, + errorMessageRedacted: cleanup.error_message_redacted, + errorExpiresAt: cleanup.error_expires_at, + }, + ] + ); + await store.saveReport(failedReport, diagnostic.errorExpiresAt); + + expect(fake.updates.at(-1)?.values).toMatchObject({ + status: 'failed', + error_message_redacted: null, + error_expires_at: null, + }); + expect(cleanup).toEqual({ error_message_redacted: null, error_expires_at: null }); + }); + + it('stores a later turn diagnostic after an earlier turn diagnostic expires', async () => { + const fake = makeDb([ + [{ createdAt: occurredAt }], + [ + { + ...storedFailedRun, + errorMessageRedacted: diagnostic.errorMessageRedacted, + errorExpiresAt: diagnostic.errorExpiresAt, + }, + ], + [{ createdAt: occurredAt }], + [], + ]); + const store = createCloudAgentReportStore(fake.db as never); + await store.saveReport(failedReport, diagnostic.errorExpiresAt); + const nextReport = { + ...failedReport, + occurredAt: diagnostic.errorExpiresAt, + run: { + ...failedReport.run, + messageId: 'msg_next_failed', + terminalAt: diagnostic.errorExpiresAt, + diagnostic: { ...diagnostic, errorExpiresAt: '2026-07-01T12:00:00.000Z' }, + }, + }; + await store.saveReport(nextReport, diagnostic.errorExpiresAt); + + expect(fake.updates[0]?.values).toMatchObject({ + error_message_redacted: null, + error_expires_at: null, + }); + expect(fake.inserts[0]?.values).toMatchObject({ + cloud_agent_session_id: cloudAgentSessionId, + message_id: nextReport.run.messageId, + status: 'failed', + error_message_redacted: diagnostic.errorMessageRedacted, + error_expires_at: nextReport.run.diagnostic.errorExpiresAt, + }); + }); + + it('does not revive cleaned diagnostic text after 30 days', async () => { + const fake = makeDb([[{ createdAt: occurredAt }], [{ ...storedFailedRun }]]); + const store = createCloudAgentReportStore(fake.db as never); + await store.saveReport(failedReport, '2026-06-24T12:00:00.000Z'); + expect(fake.updates[0]?.values).toMatchObject({ + status: 'failed', + error_message_redacted: null, + error_expires_at: null, + }); + }); + + it.each([89, 90, 91])('keeps the parent retention boundary unchanged at %s days', async age => { + const createdAt = new Date(Date.parse(occurredAt) - age * 24 * 60 * 60 * 1000).toISOString(); + const fake = makeDb([[{ createdAt }], []]); + const store = createCloudAgentReportStore(fake.db as never); + expect(await store.saveReport(failedReport, occurredAt)).toEqual({ + outcome: age < 90 ? 'applied' : 'expired', + }); + expect(fake.inserts).toHaveLength(age < 90 ? 1 : 0); + }); + + it('clears expired diagnostic pairs and purges rows older than 90 days', async () => { const fake = makeDb(); const store = createCloudAgentReportStore(fake.db as never); await store.removeExpiredData(occurredAt); @@ -386,10 +863,26 @@ describe('cloud agent reporting store', () => { error_message_redacted: null, error_expires_at: null, }); - expect(fake.updates.find(call => call.table === cloud_agent_session_runs)?.values).toEqual({ + const runCleanup = fake.updates.find(call => call.table === cloud_agent_session_runs); + expect(runCleanup?.values).toEqual({ error_message_redacted: null, error_expires_at: null, }); + const db = getWorkerDb('postgres://unused:unused@localhost:0/unused'); + const diagnosticQuery = db + .update(cloud_agent_session_runs) + .set({ error_message_redacted: null, error_expires_at: null }) + .where(runCleanup?.where) + .toSQL(); + expect(diagnosticQuery.sql).toMatch( + /"error_expires_at" is not null and "cloud_agent_session_runs"\."error_expires_at" <=/ + ); + expect(diagnosticQuery.params.at(-1)).toBe(occurredAt); + const retentionQuery = db.delete(cloud_agent_sessions).where(fake.deleteConditions[0]).toSQL(); + expect(retentionQuery.sql).toMatch(/"cloud_agent_sessions"\."created_at" <=/); + expect(retentionQuery.params).toEqual([ + new Date(Date.parse(occurredAt) - 90 * 24 * 60 * 60 * 1000).toISOString(), + ]); expect(fake.deletes).toContain(cloud_agent_sessions); expect(fake.db.transaction).not.toHaveBeenCalled(); }); diff --git a/services/cloud-agent-next/src/telemetry/report-store.ts b/services/cloud-agent-next/src/telemetry/report-store.ts index 4134f2e8a0..ea6d415b8b 100644 --- a/services/cloud-agent-next/src/telemetry/report-store.ts +++ b/services/cloud-agent-next/src/telemetry/report-store.ts @@ -80,6 +80,8 @@ type StoredRunRow = { failureCode: string | null; failureResponsibility: string | null; failureReason: string | null; + errorMessageRedacted: string | null; + errorExpiresAt: string | null; }; function retentionCutoff(now: string): string { @@ -117,14 +119,22 @@ function isTerminalStatus(status: StoredRunRow['status']): boolean { return status === 'completed' || status === 'failed' || status === 'interrupted'; } -function validDiagnostic( - diagnostic: { errorMessageRedacted: string; errorExpiresAt: string } | undefined, - now: string -): { error_message_redacted: string; error_expires_at: string } | undefined { - if (!diagnostic || Date.parse(diagnostic.errorExpiresAt) <= Date.parse(now)) return undefined; +function validDiagnostic(run: CloudAgentRunStateReport['run'], now: string) { + if (run.status !== 'failed' || run.terminalAt === undefined) return undefined; + const diagnostic = diagnosticSchema.safeParse(run.diagnostic); + if (!diagnostic.success) return undefined; + const expiresAt = Date.parse(diagnostic.data.errorExpiresAt); + const terminalAt = Date.parse(run.terminalAt); + if ( + expiresAt <= Date.parse(now) || + expiresAt <= terminalAt || + expiresAt - terminalAt > CLOUD_AGENT_ERROR_RETENTION_DAYS * 24 * 60 * 60 * 1000 + ) { + return undefined; + } return { - error_message_redacted: diagnostic.errorMessageRedacted, - error_expires_at: diagnostic.errorExpiresAt, + error_message_redacted: diagnostic.data.errorMessageRedacted, + error_expires_at: diagnostic.data.errorExpiresAt, }; } @@ -158,6 +168,8 @@ export function createCloudAgentReportStore(db: WorkerDb) { failureCode: cloud_agent_session_runs.failure_code, failureResponsibility: cloud_agent_session_runs.failure_responsibility, failureReason: cloud_agent_session_runs.failure_reason, + errorMessageRedacted: cloud_agent_session_runs.error_message_redacted, + errorExpiresAt: cloud_agent_session_runs.error_expires_at, }) .from(cloud_agent_session_runs) .where( @@ -169,7 +181,7 @@ export function createCloudAgentReportStore(db: WorkerDb) { .limit(1); const existing = rows[0]; const incoming = report.run; - const diagnostic = validDiagnostic(incoming.diagnostic, now); + const diagnostic = validDiagnostic(incoming, now); if (!existing) { await tx.insert(cloud_agent_session_runs).values({ @@ -233,6 +245,40 @@ export function createCloudAgentReportStore(db: WorkerDb) { const mayFillFailureResponsibility = mayFillTerminalFacts && (existing.failureCode === null || existing.failureCode === incoming.failureCode); + const mayFillDiagnostic = + mayFillTerminalFacts && + incoming.status === 'failed' && + (existing.failureStage === null || existing.failureStage === incoming.failureStage) && + (existing.failureCode === null || existing.failureCode === incoming.failureCode) && + (existing.failureResponsibility === null || + incoming.failureResponsibility === undefined || + existing.failureResponsibility === incoming.failureResponsibility) && + (existing.failureReason === null || + incoming.failureReason === undefined || + existing.failureReason === incoming.failureReason) && + (existing.wrapperRunId === null || + incoming.wrapperRunId === undefined || + existing.wrapperRunId === incoming.wrapperRunId) && + (existing.terminalAt === null || + (incoming.terminalAt !== undefined && + Date.parse(existing.terminalAt) === Date.parse(incoming.terminalAt))); + const retainedDiagnostic = { + error_message_redacted: existing.errorMessageRedacted, + error_expires_at: existing.errorExpiresAt, + }; + if ( + existing.errorExpiresAt !== null && + Date.parse(existing.errorExpiresAt) <= Date.parse(now) + ) { + retainedDiagnostic.error_message_redacted = null; + retainedDiagnostic.error_expires_at = null; + } else if (mayFillDiagnostic && diagnostic) { + retainedDiagnostic.error_message_redacted ??= diagnostic.error_message_redacted; + retainedDiagnostic.error_expires_at = earliestTimestamp( + existing.errorExpiresAt, + diagnostic.error_expires_at + ); + } const status = establishedTerminal ? existing.status : incomingTerminal || incoming.status === 'accepted' @@ -255,7 +301,7 @@ export function createCloudAgentReportStore(db: WorkerDb) { terminal_at: mayApplyTerminalFacts ? (incoming.terminalAt ?? null) : incomingSameTerminal - ? earliestTimestamp(existing.terminalAt, incoming.terminalAt) + ? (existing.terminalAt ?? incoming.terminalAt ?? null) : existing.terminalAt, failure_stage: existing.failureStage ?? (mayFillTerminalFacts ? (incoming.failureStage ?? null) : null), @@ -267,7 +313,7 @@ export function createCloudAgentReportStore(db: WorkerDb) { failure_reason: existing.failureReason ?? (mayFillFailureResponsibility ? (incoming.failureReason ?? null) : null), - ...(mayFillTerminalFacts ? (diagnostic ?? {}) : {}), + ...retainedDiagnostic, }) .where( and( diff --git a/services/cloud-agent-next/src/telemetry/session-reports.test.ts b/services/cloud-agent-next/src/telemetry/session-reports.test.ts index 5b578f1284..be80c8387b 100644 --- a/services/cloud-agent-next/src/telemetry/session-reports.test.ts +++ b/services/cloud-agent-next/src/telemetry/session-reports.test.ts @@ -1,4 +1,6 @@ -import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type { SessionMetadata } from '../persistence/session-metadata.js'; +import type * as ReportStoreModule from './report-store.js'; const reportStore = vi.hoisted(() => ({ createSessionReport: vi.fn().mockResolvedValue(undefined), @@ -7,31 +9,56 @@ const reportStore = vi.hoisted(() => ({ })); vi.mock('../db/pg.js', () => ({ getPgDb: vi.fn(() => ({})) })); -vi.mock('./report-store.js', () => ({ +vi.mock('./report-store.js', async importOriginal => ({ + ...(await importOriginal()), createCloudAgentReportStore: vi.fn(() => reportStore), })); import { createCloudAgentSessionReport, + ensureCloneSessionReport, recordCloudAgentSandboxIdentity, recordCloudAgentSessionFailure, } from './session-reports.js'; const env = { HYPERDRIVE: { connectionString: 'postgres://test' } } as never; const cloudAgentSessionId = 'agent_12345678-1234-1234-1234-123456789abc'; +const kiloSessionId = 'ses_12345678901234567890123456'; +const cloneFromKiloSessionId = 'ses_aaaaaaaaaaaaaaaaaaaaaaaaaa'; +const initialMessageId = 'msg_018f1e2d3c4bAbCdEfGhIjKlMn'; +const reportingCreatedAt = '2026-08-01T10:00:00.000Z'; +const now = Date.parse('2026-08-29T10:00:00.000Z'); +const retentionCutoff = now - 90 * 24 * 60 * 60 * 1000; + +function cloneMetadata(overrides: Partial = {}): SessionMetadata { + return { + metadataSchemaVersion: 2, + identity: { sessionId: cloudAgentSessionId, userId: 'user_clone' }, + auth: { kiloSessionId }, + clone: { cloneFromKiloSessionId, reportingCreatedAt }, + initialMessage: { id: initialMessageId }, + workspace: { sandboxId: 'usr-123456789abc' }, + lifecycle: { version: now, timestamp: now }, + ...overrides, + }; +} describe('Cloud Agent session report writes', () => { beforeEach(() => { vi.clearAllMocks(); + reportStore.createSessionReport.mockResolvedValue(undefined); + reportStore.recordSandboxIdentity.mockResolvedValue({}); + vi.useFakeTimers(); + vi.setSystemTime(now); + }); + + afterEach(() => { + vi.useRealTimers(); }); it('writes setup facts through the Cloud Agent report store', async () => { await createCloudAgentSessionReport( - { - cloudAgentSessionId, - kiloSessionId: 'ses_12345678901234567890123456', - initialMessageId: 'msg_initial', - }, + { cloudAgentSessionId, kiloSessionId, initialMessageId }, env ); await recordCloudAgentSandboxIdentity( @@ -43,15 +70,143 @@ describe('Cloud Agent session report writes', () => { env ); - expect(reportStore.createSessionReport).toHaveBeenCalledWith( - expect.objectContaining({ occurredAt: expect.stringMatching(/^\d{4}-\d{2}-\d{2}T/) }) - ); + expect(reportStore.createSessionReport).toHaveBeenCalledWith({ + cloudAgentSessionId, + kiloSessionId, + initialMessageId, + occurredAt: new Date(now).toISOString(), + }); expect(reportStore.recordSandboxIdentity).toHaveBeenCalledWith({ cloudAgentSessionId, sandboxId: 'ses-sandbox-id', }); expect(reportStore.recordSessionFailure).toHaveBeenCalledWith( - expect.objectContaining({ occurredAt: expect.stringMatching(/^\d{4}-\d{2}-\d{2}T/) }) + expect.objectContaining({ occurredAt: new Date(now).toISOString() }) ); }); + + it('preserves an explicit reporting creation time instead of using report delivery time', async () => { + await createCloudAgentSessionReport( + { cloudAgentSessionId, kiloSessionId, initialMessageId, occurredAt: reportingCreatedAt }, + env + ); + + expect(reportStore.createSessionReport).toHaveBeenCalledWith({ + cloudAgentSessionId, + kiloSessionId, + initialMessageId, + occurredAt: reportingCreatedAt, + }); + }); + + it('awaits creation for the destination and persisted first message before sandbox identity', async () => { + const creation = Promise.withResolvers(); + const sandbox = Promise.withResolvers(); + reportStore.createSessionReport.mockReturnValueOnce(creation.promise); + reportStore.recordSandboxIdentity.mockReturnValueOnce(sandbox.promise); + let finished = false; + const write = ensureCloneSessionReport(cloneMetadata(), env).then(() => { + finished = true; + }); + + expect(reportStore.createSessionReport).toHaveBeenCalledWith({ + cloudAgentSessionId, + kiloSessionId, + initialMessageId, + occurredAt: reportingCreatedAt, + }); + expect(reportStore.recordSandboxIdentity).not.toHaveBeenCalled(); + creation.resolve(); + await creation.promise; + await Promise.resolve(); + expect(reportStore.recordSandboxIdentity).toHaveBeenCalledWith({ + cloudAgentSessionId, + sandboxId: 'usr-123456789abc', + }); + expect(finished).toBe(false); + sandbox.resolve(); + await write; + expect(finished).toBe(true); + }); + + it('retries idempotent creation on later reports without changing clone age', async () => { + const metadata = cloneMetadata(); + await ensureCloneSessionReport(metadata, env); + vi.setSystemTime(now + 60_000); + await ensureCloneSessionReport(metadata, env); + + expect(reportStore.createSessionReport.mock.calls.map(([input]) => input)).toEqual([ + { cloudAgentSessionId, kiloSessionId, initialMessageId, occurredAt: reportingCreatedAt }, + { cloudAgentSessionId, kiloSessionId, initialMessageId, occurredAt: reportingCreatedAt }, + ]); + }); + + it.each([ + { name: 'missing metadata', metadata: null }, + { name: 'unmarked old clone', metadata: cloneMetadata({ clone: { cloneFromKiloSessionId } }) }, + { name: 'non-clone session', metadata: cloneMetadata({ clone: undefined }) }, + { name: 'unadmitted clone', metadata: cloneMetadata({ initialMessage: undefined }) }, + { name: 'missing destination Kilo identity', metadata: cloneMetadata({ auth: {} }) }, + { + name: 'control-plane session', + metadata: cloneMetadata({ + identity: { + sessionId: 'workspace_12345678-1234-1234-1234-123456789abc', + userId: 'user_clone', + }, + }), + }, + ...[0, -1].map(offset => ({ + name: offset === 0 ? 'exactly expired clone' : 'older expired clone', + metadata: cloneMetadata({ + clone: { + cloneFromKiloSessionId, + reportingCreatedAt: new Date(retentionCutoff + offset).toISOString(), + }, + }), + })), + ])('does not create an anchor for $name', async ({ metadata }) => { + await ensureCloneSessionReport(metadata, env); + + expect(reportStore.createSessionReport).not.toHaveBeenCalled(); + expect(reportStore.recordSandboxIdentity).not.toHaveBeenCalled(); + }); + + it('anchors a clone one millisecond inside the reporting retention window', async () => { + const occurredAt = new Date(retentionCutoff + 1).toISOString(); + await ensureCloneSessionReport( + cloneMetadata({ clone: { cloneFromKiloSessionId, reportingCreatedAt: occurredAt } }), + env + ); + + expect(reportStore.createSessionReport).toHaveBeenCalledWith({ + cloudAgentSessionId, + kiloSessionId, + initialMessageId, + occurredAt, + }); + }); + + it.each(['createSessionReport', 'recordSandboxIdentity'] as const)( + 'allows a later report to retry after %s fails', + async method => { + reportStore[method].mockRejectedValueOnce(new Error('report storage unavailable')); + await expect(ensureCloneSessionReport(cloneMetadata(), env)).rejects.toThrow( + 'report storage unavailable' + ); + await ensureCloneSessionReport(cloneMetadata(), env); + + expect(reportStore.createSessionReport).toHaveBeenCalledTimes(2); + expect(reportStore.createSessionReport).toHaveBeenLastCalledWith({ + cloudAgentSessionId, + kiloSessionId, + initialMessageId, + occurredAt: reportingCreatedAt, + }); + expect(reportStore.recordSandboxIdentity).toHaveBeenLastCalledWith({ + cloudAgentSessionId, + sandboxId: 'usr-123456789abc', + }); + } + ); }); diff --git a/services/cloud-agent-next/src/telemetry/session-reports.ts b/services/cloud-agent-next/src/telemetry/session-reports.ts index c4d519e150..519f32dc39 100644 --- a/services/cloud-agent-next/src/telemetry/session-reports.ts +++ b/services/cloud-agent-next/src/telemetry/session-reports.ts @@ -1,6 +1,7 @@ import type { Env } from '../types.js'; +import type { SessionMetadata } from '../persistence/session-metadata.js'; import { getPgDb } from '../db/pg.js'; -import { createCloudAgentReportStore } from './report-store.js'; +import { CLOUD_AGENT_REPORT_RETENTION_DAYS, createCloudAgentReportStore } from './report-store.js'; export type CloudAgentSessionFailure = | { stage: 'sandbox_identity'; code: 'sandbox_id_derivation_failed' } @@ -14,15 +15,46 @@ export type CloudAgentSessionFailure = type ReportingEnv = Pick; export async function createCloudAgentSessionReport( - params: { cloudAgentSessionId: string; kiloSessionId: string; initialMessageId: string }, + params: { + cloudAgentSessionId: string; + kiloSessionId: string; + initialMessageId: string; + occurredAt?: string; + }, env: ReportingEnv ): Promise { await createCloudAgentReportStore(getPgDb(env)).createSessionReport({ ...params, - occurredAt: new Date().toISOString(), + occurredAt: params.occurredAt ?? new Date().toISOString(), }); } +export async function ensureCloneSessionReport( + metadata: SessionMetadata | null, + env: ReportingEnv +): Promise { + const occurredAt = metadata?.clone?.reportingCreatedAt; + const kiloSessionId = metadata?.auth.kiloSessionId; + const initialMessageId = metadata?.initialMessage?.id; + if ( + !metadata?.identity.sessionId.startsWith('agent_') || + !occurredAt || + !kiloSessionId || + !initialMessageId || + Date.parse(occurredAt) <= Date.now() - CLOUD_AGENT_REPORT_RETENTION_DAYS * 24 * 60 * 60 * 1000 + ) { + return; + } + + const cloudAgentSessionId = metadata.identity.sessionId; + await createCloudAgentSessionReport( + { cloudAgentSessionId, kiloSessionId, initialMessageId, occurredAt }, + env + ); + const sandboxId = metadata.workspace?.sandboxId; + if (sandboxId) await recordCloudAgentSandboxIdentity({ cloudAgentSessionId, sandboxId }, env); +} + export async function recordCloudAgentSandboxIdentity( params: { cloudAgentSessionId: string; sandboxId: string }, env: ReportingEnv diff --git a/services/cloud-agent-next/src/websocket/ingest.test.ts b/services/cloud-agent-next/src/websocket/ingest.test.ts index 945d1f2535..58e2c69d65 100644 --- a/services/cloud-agent-next/src/websocket/ingest.test.ts +++ b/services/cloud-agent-next/src/websocket/ingest.test.ts @@ -3,6 +3,13 @@ import { describe, expect, it, vi } from 'vitest'; import { createIngestHandler, type IngestDOContext, type IngestAttachment } from './ingest.js'; import type { EventQueries } from '../session/queries/index.js'; import type { SessionId } from '../types/ids.js'; +import { + getSessionMessageState, + putSessionMessageState, + terminalizeMessageOnce, + type SessionMessageStorage, +} from '../session/session-message-state.js'; +import { prepareIngestFrame } from '../shared/ingest-frame.js'; const SESSION_ID = 'sess_test' as SessionId; const WRAPPER_RUN_ID = 'wr_test_basic'; @@ -952,6 +959,150 @@ describe('createIngestHandler', () => { expect(doContext.terminalizeSessionMessageOnce).not.toHaveBeenCalled(); }); + it.each([ + [undefined, undefined], + [null, undefined], + ['', 'failed'], + [false, 'failed'], + [0, 'failed'], + [[], 'failed'], + [{}, 'failed'], + [{ message: '' }, 'failed'], + [{ data: {} }, 'failed'], + [{ data: { message: '' } }, 'failed'], + [{ name: 'MessageOutputLengthError', data: {} }, 'failed'], + [{ name: 'APIError', data: {} }, 'failed'], + [{ name: 'MessageAbortedError', data: {} }, 'interrupted'], + ] as const)( + 'keeps the existing assistant-error terminalization semantics before and after compaction: %j', + async (error, kind) => { + for (const oversized of [false, true]) { + const doContext = createNewPathDOContext(); + const handler = createIngestHandler( + createFakeState(), + createFakeEventQueries(), + SESSION_ID, + vi.fn(), + doContext + ); + const ws = createFakeWebSocket(makeNewPathAttachment()); + const frame = prepareIngestFrame({ + streamEventType: 'kilocode', + timestamp: new Date().toISOString(), + data: { + event: 'message.updated', + properties: { + info: { + id: 'asst_presence', + role: 'assistant', + parentID: 'msg_user_presence', + error, + ...(oversized ? { metadata: 'poison-metadata'.repeat(100_000) } : {}), + }, + }, + }, + }); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(oversized); + await handler.handleIngestMessage(ws, frame.serialized); + if (kind === undefined) { + expect(doContext.terminalizeSessionMessageOnce).not.toHaveBeenCalled(); + } else { + expect(doContext.terminalizeSessionMessageOnce).toHaveBeenCalledWith( + 'msg_user_presence', + expect.objectContaining({ kind }), + WRAPPER_RUN_ID + ); + } + } + } + ); + + it.each([ + ['ContextOverflowError', undefined, 'context_limit', 'The model context limit was exceeded'], + ['MessageOutputLengthError', undefined, 'output_limit', 'The model output limit was reached'], + [ + 'ContentFilterError', + undefined, + 'content_filter', + 'The model provider blocked the response under its content policy', + ], + [ + 'StructuredOutputError', + undefined, + 'structured_output', + 'The model response did not match the required format', + ], + ['APIError', 504, 'timeout', 'Assistant request timed out'], + ['APIError', 400, 'invalid_request', 'Assistant request was invalid'], + ] as const)( + 'classifies %s/%s transiently before and after compaction', + async (name, statusCode, assistantFailureReason, safeFailureMessage) => { + for (const oversized of [false, true]) { + const doContext = createNewPathDOContext(); + const eventQueries = createFakeEventQueries(); + const handler = createIngestHandler( + createFakeState(), + eventQueries, + SESSION_ID, + vi.fn(), + doContext + ); + const frame = prepareIngestFrame({ + streamEventType: 'kilocode', + timestamp: new Date().toISOString(), + data: { + event: 'message.updated', + properties: { + info: { + id: 'asst_classified', + role: 'assistant', + parentID: 'msg_user_classified', + error: { + name, + data: { + message: '[BYOK] opaque failure', + statusCode, + isRetryable: false, + responseBody: 'poison-body'.repeat(oversized ? 200_000 : 1), + responseHeaders: { authorization: 'poison-header' }, + }, + }, + }, + }, + }, + }); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(oversized); + await handler.handleIngestMessage( + createFakeWebSocket(makeNewPathAttachment()), + frame.serialized + ); + + const payload = vi.mocked(eventQueries.upsert).mock.calls[0][0].payload; + expect(JSON.parse(payload).properties.info.error).toBe(`[BYOK] ${safeFailureMessage}`); + expect(payload).not.toMatch(/poison|statusCode|isRetryable|responseBody|responseHeaders/); + expect(doContext.terminalizeSessionMessageOnce).toHaveBeenCalledWith( + 'msg_user_classified', + expect.objectContaining({ + kind: 'failed', + assistantMessageId: 'asst_classified', + failureCode: 'assistant_error', + assistantFailureReason, + providerOwnership: 'byok', + safeFailureMessage, + }), + WRAPPER_RUN_ID + ); + expect(doContext.terminalizeSessionMessageOnce.mock.calls[0][1]).not.toHaveProperty( + 'failureFacts' + ); + } + } + ); + it('terminalizes on wrapper cloud.message.completed control event', async () => { const state = createFakeState(); const doContext = createNewPathDOContext(); @@ -1031,7 +1182,19 @@ describe('createIngestHandler', () => { sessionID: 'kilo_session_333', role: 'assistant', parentID: 'msg_user_333', - error: { data: { message: secretError, responseBody: 'secret-response' } }, + modelID: 'vendor/model', + providerID: 'kilo', + error: { + name: 'APIError', + data: { + message: secretError, + statusCode: 429, + isRetryable: false, + responseBody: 'secret-response', + responseHeaders: { authorization: 'secret-header' }, + metadata: { url: 'https://secret.example', toolArguments: 'secret-arguments' }, + }, + }, }, }, }, @@ -1048,6 +1211,8 @@ describe('createIngestHandler', () => { sessionID: 'kilo_session_333', role: 'assistant', parentID: 'msg_user_333', + modelID: 'vendor/model', + providerID: 'kilo', error: 'Assistant request was rate limited', }, }, @@ -1075,6 +1240,115 @@ describe('createIngestHandler', () => { }), WRAPPER_RUN_ID ); + expect(doContext.terminalizeSessionMessageOnce.mock.calls[0][1]).not.toHaveProperty( + 'failureFacts' + ); + expect(JSON.stringify(vi.mocked(eventQueries.upsert).mock.calls)).not.toMatch( + /statusCode|isRetryable|responseBody|responseHeaders|secret-header|secret-arguments/ + ); + }); + + it('does not attribute compacted child, replayed, late, or stale failures to another message', async () => { + const records = new Map(); + const storage: SessionMessageStorage = { + async get(key: string) { + return records.get(key) as T | undefined; + }, + async put(key, value) { + records.set(key, value); + }, + async list({ prefix }: { prefix: string }) { + return new Map( + [...records].filter(([key]) => key.startsWith(prefix)) as Array<[string, T]> + ); + }, + }; + const parentId = 'msg_0123456789abABCDEFGHIJKLMN'; + const nextId = 'msg_0123456789abABCDEFGHIJKLMO'; + for (const messageId of [parentId, nextId]) { + await putSessionMessageState(storage, { + messageId, + status: 'accepted', + prompt: 'do not report this', + createdAt: 1, + acceptedAt: 2, + wrapperRunId: WRAPPER_RUN_ID, + }); + } + const doContext = createNewPathDOContext(); + doContext.terminalizeSessionMessageOnce.mockImplementation(async (messageId, params) => { + await terminalizeMessageOnce(storage, messageId, params); + }); + const handler = createIngestHandler( + createFakeState(), + createFakeEventQueries(), + SESSION_ID, + vi.fn(), + doContext + ); + const ws = createFakeWebSocket(makeNewPathAttachment()); + async function sendFailure( + parentID: string, + id: string, + message = 'Rate limit exceeded', + sessionID = 'root_session' + ) { + const frame = prepareIngestFrame({ + streamEventType: 'kilocode', + timestamp: new Date().toISOString(), + data: { + event: 'message.updated', + properties: { + info: { + id, + parentID, + sessionID, + role: 'assistant', + error: { + name: 'APIError', + data: { + message, + statusCode: 429, + isRetryable: false, + responseBody: 'poison-body'.repeat(200_000), + }, + }, + }, + }, + }, + }); + expect(frame.kind).toBe('send'); + if (frame.kind !== 'send') return; + expect(frame.compacted).toBe(true); + expect(frame.serialized).not.toContain('poison'); + await handler.handleIngestMessage(ws, frame.serialized); + } + + await sendFailure('msg_child_only', 'asst_child', 'Rate limit exceeded', 'child_session'); + expect((await getSessionMessageState(storage, parentId))?.status).toBe('accepted'); + expect((await getSessionMessageState(storage, nextId))?.status).toBe('accepted'); + + await sendFailure(parentId, 'asst_original'); + const failed = await getSessionMessageState(storage, parentId); + expect(failed).toMatchObject({ + status: 'failed', + assistantMessageId: 'asst_original', + assistantFailureReason: 'rate_limited', + safeFailureMessage: 'Assistant request was rate limited', + }); + expect(failed).not.toHaveProperty('failureFacts'); + await sendFailure(parentId, 'asst_original'); + await sendFailure(parentId, 'asst_late', 'Provider timeout'); + expect(await getSessionMessageState(storage, parentId)).toEqual(failed); + expect( + (await getSessionMessageState(storage, nextId))?.assistantFailureReason + ).toBeUndefined(); + expect((await getSessionMessageState(storage, nextId))?.assistantMessageId).toBeUndefined(); + + vi.mocked(doContext.wrapperSupervisor.isCurrentConnection).mockResolvedValue(false); + await sendFailure(nextId, 'asst_stale'); + expect((await getSessionMessageState(storage, nextId))?.status).toBe('accepted'); + expect(ws.close).toHaveBeenCalledWith(4401, 'Stale wrapper connection'); }); it('publishes and persists a safe session error with session correlation intact', async () => { diff --git a/services/cloud-agent-next/src/websocket/ingest.ts b/services/cloud-agent-next/src/websocket/ingest.ts index bcbaf08d92..b17b142427 100644 --- a/services/cloud-agent-next/src/websocket/ingest.ts +++ b/services/cloud-agent-next/src/websocket/ingest.ts @@ -38,6 +38,7 @@ import { classifyAssistantFailure, classifyAssistantFailureMessage, isAssistantInterrupt, + projectSafeAssistantError, } from '../session/safe-failure-projection.js'; import { parseModelNotFoundRuntimeDiagnostics } from '../shared/runtime-model-diagnostics.js'; import { @@ -135,7 +136,7 @@ function sanitizeKilocodeEventData(data: unknown): unknown { ...properties, info: { ...info, - error: classifyAssistantFailureMessage(info.error), + error: projectSafeAssistantError(info.error), }, }, }; @@ -796,27 +797,28 @@ export function createIngestHandler( if (eventName === 'message.updated') { const properties = data.properties as Record | undefined; const info = properties?.info as Record | undefined; - const assistantError = getAssistantErrorMessage(info?.error); + const assistantError = info?.error; + const assistantErrorMessage = getAssistantErrorMessage(assistantError); const parentMessageId = info?.role === 'assistant' && typeof info.parentID === 'string' ? info.parentID : undefined; if (parentMessageId !== undefined) { await doContext.observeCorrelatedAgentActivity?.(parentMessageId); - if (info?.error !== undefined && isAssistantInterrupt(info.error)) { + if (assistantError !== undefined && isAssistantInterrupt(assistantError)) { await doContext.terminalizeSessionMessageOnce( parentMessageId, { kind: 'interrupted', assistantMessageId: typeof info?.id === 'string' ? info.id : undefined, - error: assistantError ?? 'The message was interrupted by the user', + error: assistantErrorMessage ?? 'The message was interrupted by the user', failureStage: 'interruption', failureCode: 'user_interrupt', completionSource: 'interrupt', }, wrapperRunId ); - } else if (assistantError !== undefined) { + } else if (assistantErrorMessage !== undefined) { const assistantFailure = classifyAssistantFailure(assistantError); await doContext.terminalizeSessionMessageOnce( parentMessageId, @@ -824,7 +826,7 @@ export function createIngestHandler( kind: 'failed', assistantMessageId: typeof info?.id === 'string' ? info.id : undefined, reason: 'assistant_error', - error: assistantError, + error: assistantErrorMessage, failureStage: 'agent_activity', failureCode: assistantFailure.terminalCode ?? 'assistant_error', assistantFailureReason: assistantFailure.reason, diff --git a/services/cloud-agent-next/test/integration/session/admission-recovery.test.ts b/services/cloud-agent-next/test/integration/session/admission-recovery.test.ts index 16c47c2eb2..9c6c7af168 100644 --- a/services/cloud-agent-next/test/integration/session/admission-recovery.test.ts +++ b/services/cloud-agent-next/test/integration/session/admission-recovery.test.ts @@ -1,9 +1,17 @@ import { env, runInDurableObject } from 'cloudflare:test'; -import { describe, expect, it } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type { CloudAgentQueueReport } from '@kilocode/worker-utils/cloud-agent-queue-report'; import type { CallbackJob } from '../../../src/callbacks/types.js'; +import type { CloudAgentSession } from '../../../src/persistence/CloudAgentSession.js'; +import { parseSessionMetadata } from '../../../src/persistence/session-metadata.js'; import { listPendingSessionMessages } from '../../../src/session/pending-messages.js'; import { getSessionMessageState } from '../../../src/session/session-message-state.js'; -import { queueUserMessageInput, registerReadySession } from '../../helpers/session-setup.js'; +import * as sessionReports from '../../../src/telemetry/session-reports.js'; +import { + groupedRegisterSessionInput, + queueUserMessageInput, + registerReadySession, +} from '../../helpers/session-setup.js'; describe('partial admission callback snapshot recovery', () => { it('retains the admission-time callback target when delivery accepts before state repair', async () => { @@ -49,7 +57,7 @@ describe('partial admission callback snapshot recovery', () => { const realPut = instance.ctx.storage.put.bind(instance.ctx.storage); let failedQueuedState = false; instance.ctx.storage.put = async (key, value) => { - if (!failedQueuedState && String(key).startsWith('session_message:')) { + if (!failedQueuedState && typeof key === 'string' && key.startsWith('session_message:')) { failedQueuedState = true; throw new Error('queued state unavailable'); } @@ -89,3 +97,312 @@ describe('partial admission callback snapshot recovery', () => { expect(result.captured[0]?.target.url).toBe('https://callback.example.com/original'); }); }); + +const cloneFromKiloSessionId = 'ses_aaaaaaaaaaaaaaaaaaaaaaaaaa'; +const destinationKiloSessionId = 'ses_bbbbbbbbbbbbbbbbbbbbbbbbbb'; +const firstMessageId = 'msg_018f1e2d3c4bAbCdEfGhIjKlMn'; +const secondMessageId = 'msg_018f1e2d3c4bBBBBBBBBBBBBBB'; + +function cloneRegistrationInput( + reportingCreatedAt?: string +): Parameters[0] { + return { + ...groupedRegisterSessionInput({ + sessionId: `agent_${crypto.randomUUID()}`, + userId: 'user_clone_anchor', + kiloSessionId: destinationKiloSessionId, + prompt: '', + mode: 'code', + model: 'test-model', + }), + message: undefined, + clone: { + cloneFromKiloSessionId, + ...(reportingCreatedAt ? { reportingCreatedAt } : {}), + }, + workspace: { sandboxId: 'usr-123456789abc' }, + }; +} + +function cloneStub(input: ReturnType) { + return env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${input.identity.userId}:${input.identity.sessionId}`) + ); +} + +describe('forward-only clone reporting admission', () => { + beforeEach(() => { + vi.spyOn(sessionReports, 'ensureCloneSessionReport').mockResolvedValue(undefined); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it('retains the first admitted ID through rejection, replay, later sends and stale metadata updates', async () => { + const reportingCreatedAt = new Date().toISOString(); + const input = cloneRegistrationInput(reportingCreatedAt); + await runInDurableObject(cloneStub(input), async instance => { + expect(await instance.registerSession(input)).toEqual({ success: true }); + const originalMetadata = await instance.getMetadata(); + expect(originalMetadata?.initialMessage).toBeUndefined(); + const request = queueUserMessageInput({ + userId: input.identity.userId, + messageId: firstMessageId, + prompt: 'continue the clone', + }); + expect( + await instance.admitSubmittedMessage({ + ...request, + turn: { type: 'prompt', id: secondMessageId, prompt: '' }, + }) + ).toMatchObject({ success: false, code: 'BAD_REQUEST' }); + expect((await instance.getMetadata())?.initialMessage).toBeUndefined(); + expect(await listPendingSessionMessages(instance.ctx.storage)).toHaveLength(0); + + const admission = await instance.admitSubmittedMessage(request); + expect(admission).toMatchObject({ success: true, messageId: firstMessageId }); + expect((await instance.getMetadata())?.initialMessage).toEqual({ id: firstMessageId }); + expect(await instance.admitSubmittedMessage(request)).toEqual(admission); + expect( + await instance.admitSubmittedMessage({ + ...request, + turn: { type: 'prompt', id: secondMessageId, prompt: 'second message' }, + }) + ).toMatchObject({ success: true, messageId: secondMessageId }); + expect( + await instance.admitSubmittedMessage({ + ...request, + turn: { type: 'prompt', id: firstMessageId, prompt: 'changed replay' }, + }) + ).toMatchObject({ success: false, code: 'BAD_REQUEST' }); + await instance.updateMetadata(originalMetadata); + expect((await instance.getMetadata())?.initialMessage).toEqual({ id: firstMessageId }); + await expect( + instance.updateMetadata({ + ...originalMetadata, + clone: { cloneFromKiloSessionId, reportingCreatedAt: '2099-01-01T00:00:00.000Z' }, + }) + ).rejects.toThrow('Clone reporting creation time cannot be changed'); + expect((await instance.getMetadata())?.clone?.reportingCreatedAt).toBe(reportingCreatedAt); + expect(await listPendingSessionMessages(instance.ctx.storage)).toHaveLength(2); + }); + }); + + it('chooses one first ID across concurrent first sends', async () => { + const input = cloneRegistrationInput(new Date().toISOString()); + await runInDurableObject(cloneStub(input), async instance => { + expect(await instance.registerSession(input)).toEqual({ success: true }); + const firstIdWrites: Array = []; + const put = instance.ctx.storage.put.bind(instance.ctx.storage); + instance.ctx.storage.put = async (key, value) => { + if (key === 'metadata') { + firstIdWrites.push(parseSessionMetadata(value).initialMessage?.id); + } + return put(key, value); + }; + try { + const results = await Promise.all( + [firstMessageId, secondMessageId].map(messageId => + instance.admitSubmittedMessage( + queueUserMessageInput({ + userId: input.identity.userId, + messageId, + prompt: 'concurrent clone send', + }) + ) + ) + ); + expect(results.every(result => result.success)).toBe(true); + const initialMessage = (await instance.getMetadata())?.initialMessage; + expect([firstMessageId, secondMessageId]).toContain(initialMessage?.id); + expect(firstIdWrites).toEqual([initialMessage?.id]); + expect(await listPendingSessionMessages(instance.ctx.storage)).toHaveLength(2); + for (const messageId of [firstMessageId, secondMessageId]) { + expect(await getSessionMessageState(instance.ctx.storage, messageId)).toMatchObject({ + status: 'queued', + }); + } + } finally { + instance.ctx.storage.put = put; + } + }); + }); + + it.each(['session_message:', 'metadata'])( + 'rolls back all queued admission writes when writing %s fails', + async failingKey => { + const input = cloneRegistrationInput(new Date().toISOString()); + await runInDurableObject(cloneStub(input), async instance => { + expect(await instance.registerSession(input)).toEqual({ success: true }); + const put = instance.ctx.storage.put.bind(instance.ctx.storage); + instance.ctx.storage.put = async (key, value) => { + await put(key, value); + if (typeof key === 'string' && key.startsWith(failingKey)) { + throw new Error('queued admission write failed'); + } + }; + try { + expect( + await instance.admitSubmittedMessage( + queueUserMessageInput({ + userId: input.identity.userId, + messageId: firstMessageId, + prompt: 'failed first attempt', + }) + ) + ).toMatchObject({ success: false, code: 'INTERNAL' }); + } finally { + instance.ctx.storage.put = put; + } + expect(await listPendingSessionMessages(instance.ctx.storage)).toEqual([]); + expect(await getSessionMessageState(instance.ctx.storage, firstMessageId)).toBeUndefined(); + expect((await instance.getMetadata())?.initialMessage).toBeUndefined(); + expect( + instance['eventQueries'].findByFilters({ eventTypes: ['cloud.message.queued'] }) + ).toEqual([]); + expect( + await instance.admitSubmittedMessage( + queueUserMessageInput({ + userId: input.identity.userId, + messageId: secondMessageId, + prompt: 'successful first admission', + }) + ) + ).toMatchObject({ success: true }); + expect((await instance.getMetadata())?.initialMessage).toEqual({ id: secondMessageId }); + }); + } + ); + + it('does not enroll old clones on first send or metadata update', async () => { + const input = cloneRegistrationInput(); + await runInDurableObject(cloneStub(input), async instance => { + expect(await instance.registerSession(input)).toEqual({ success: true }); + expect( + await instance.admitSubmittedMessage( + queueUserMessageInput({ + userId: input.identity.userId, + messageId: firstMessageId, + prompt: 'old clone send', + }) + ) + ).toMatchObject({ success: true }); + const metadata = await instance.getMetadata(); + expect(metadata?.initialMessage).toBeUndefined(); + expect(metadata?.clone).toEqual({ cloneFromKiloSessionId }); + await expect( + instance.updateMetadata({ + ...metadata, + clone: { cloneFromKiloSessionId, reportingCreatedAt: new Date().toISOString() }, + }) + ).rejects.toThrow('Clone reporting creation time cannot be changed'); + }); + }); + + it.each([false, true])( + 'does not block admission or drain on background anchor writes (first write fails: %s)', + async failFirstWrite => { + const input = cloneRegistrationInput(new Date().toISOString()); + const release = Promise.withResolvers(); + let shouldFail = failFirstWrite; + vi.mocked(sessionReports.ensureCloneSessionReport).mockImplementation(async () => { + await release.promise; + if (shouldFail) { + shouldFail = false; + throw new Error('PostgreSQL unavailable'); + } + }); + await runInDurableObject(cloneStub(input), async instance => { + const reports: CloudAgentQueueReport[] = []; + const delivered: string[] = []; + instance['sendRunStateReport'] = async report => { + reports.push(report); + }; + instance['executeDirectly'] = async plan => { + delivered.push(plan.turn.messageId); + return { + success: true, + outcome: 'accepted', + messageId: plan.turn.messageId, + wrapperRunId: 'wr_clone_anchor', + }; + }; + expect(await instance.registerSession(input)).toEqual({ success: true }); + const progress = (async () => { + expect( + await instance.admitSubmittedMessage( + queueUserMessageInput({ + userId: input.identity.userId, + messageId: firstMessageId, + prompt: 'drain without reporting', + }) + ) + ).toMatchObject({ success: true }); + expect(await instance['getSessionMessageQueue']().drainNextPendingMessage()).toEqual({ + remainingPendingCount: 0, + }); + })(); + try { + await vi.waitFor(() => expect(delivered).toEqual([firstMessageId])); + await progress; + expect(reports).toEqual([]); + expect(await listPendingSessionMessages(instance.ctx.storage)).toEqual([]); + expect(sessionReports.ensureCloneSessionReport).toHaveBeenCalledWith( + expect.objectContaining({ + auth: expect.objectContaining({ kiloSessionId: destinationKiloSessionId }), + clone: input.clone, + initialMessage: { id: firstMessageId }, + }), + expect.anything() + ); + release.resolve(); + await vi.waitFor(() => expect(reports).toHaveLength(2)); + const state = await getSessionMessageState(instance.ctx.storage, firstMessageId); + if (!state) throw new Error('Expected persisted first message'); + await instance['reportRunState'](state); + expect(reports).toHaveLength(3); + expect(sessionReports.ensureCloneSessionReport).toHaveBeenCalledTimes(3); + } finally { + release.resolve(); + await progress; + } + }); + } + ); + + it.each(['missing', 'deletion-pending'] as const)( + 'does not anchor from cached identity when live metadata is %s', + async missingState => { + const input = cloneRegistrationInput(new Date().toISOString()); + await runInDurableObject(cloneStub(input), async instance => { + expect(await instance.registerSession(input)).toEqual({ success: true }); + expect( + await instance.admitSubmittedMessage( + queueUserMessageInput({ + userId: input.identity.userId, + messageId: firstMessageId, + prompt: 'before deletion', + }) + ) + ).toMatchObject({ success: true }); + const state = await getSessionMessageState(instance.ctx.storage, firstMessageId); + if (!state) throw new Error('Expected persisted first message'); + if (missingState === 'missing') { + await instance.ctx.storage.delete('metadata'); + } else { + await instance.ctx.storage.put('session_deletion_intent', { + reason: 'explicit', + requestedAt: Date.now(), + }); + } + await instance['reportRunState'](state); + expect(sessionReports.ensureCloneSessionReport).toHaveBeenLastCalledWith( + null, + expect.anything() + ); + expect(await instance.getMetadata()).toBeNull(); + }); + } + ); +}); diff --git a/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts b/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts index 5d089f1754..0809992b60 100644 --- a/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts +++ b/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts @@ -7,7 +7,7 @@ */ import { env, runInDurableObject, listDurableObjectIds } from 'cloudflare:test'; -import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { describe, it, expect, beforeEach } from 'vitest'; import { drizzle } from 'drizzle-orm/durable-sqlite'; import { createEventQueries } from '../../../src/session/queries/events.js'; import type { FencedWrapperDispatchRequest } from '../../../src/execution/types.js'; @@ -15,7 +15,6 @@ import { listPendingSessionMessages } from '../../../src/session/pending-message import { getWrapperLease, getWrapperRuntimeState, - recordWrapperPong, allocateWrapperRuntimeState, recordWrapperAcceptedMessage, } from '../../../src/session/wrapper-runtime-state.js'; @@ -572,13 +571,13 @@ describe('new-path liveness without executionId', () => { expect(failedPayload).toMatchObject({ messageId: 'msg_018f1e2d3c4bnewlivabcdefgh', status: 'failed', - error: 'Agent wrapper produced no output', + error: 'Agent wrapper made no execution progress during the watchdog window', delivery: 'sent', accepted: true, failure: { stage: 'post_dispatch_no_activity', code: 'wrapper_no_output', - message: 'Agent wrapper produced no output', + message: 'Agent wrapper made no execution progress during the watchdog window', }, }); @@ -689,7 +688,7 @@ describe('hot delivery failure preserves existing wrapper identity', () => { const result = await runInDurableObject(stub, async instance => { (instance as any).orchestrator = { - execute: async (plan: FencedWrapperDispatchRequest) => { + execute: async () => { throw new Error('Sandbox connect failed'); }, }; @@ -792,7 +791,7 @@ describe('hot delivery failure preserves existing wrapper identity', () => { const result = await runInDurableObject(stub, async instance => { (instance as any).orchestrator = { - execute: async (plan: FencedWrapperDispatchRequest) => { + execute: async () => { throw new Error('Sandbox connect failed'); }, }; diff --git a/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts b/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts index 00051748a4..426e015728 100644 --- a/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts +++ b/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts @@ -115,6 +115,130 @@ describe('message terminalization and stream events', () => { ); }); + it.each([false, true])( + 'reports queued cancellation and successful retries with one client cancel event (first send fails: %s)', + async failFirstReport => { + const userId = `user_cancel_report_${failFirstReport}`; + const sessionId = `agent_cancel_report_${failFirstReport}`; + const messageId = 'msg_018f1e2d3c4bCancelReportAB'; + const acceptedMessageId = 'msg_018f1e2d3c4bCancelActiveAB'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + const reports: CloudAgentQueueReport[] = []; + const callbackQueue = createCapturedQueue(); + const broadcastTypes: string[] = []; + let interruptedReportAttempts = 0; + + const first = await runInDurableObject(stub, async instance => { + injectReportQueue(instance, reports); + injectCallbackQueue(instance, callbackQueue); + await registerReadySession(instance, { + sessionId, + userId, + kiloSessionId, + prompt: 'queued cancellation', + mode: 'code', + model: 'test-model', + callbackTarget: { url: 'https://example.com/callback' }, + }); + for (const id of [messageId, acceptedMessageId]) { + await instance.admitSubmittedMessage( + queueUserMessageInput({ userId, messageId: id, prompt: 'queued cancellation' }) + ); + } + const admitted = await getSessionMessageState(instance.ctx.storage, acceptedMessageId); + if (!admitted) throw new Error('Expected admitted message'); + await putSessionMessageState(instance.ctx.storage, { + ...admitted, + status: 'accepted', + acceptedAt: Date.now(), + dispatchAcceptanceKind: 'observed', + wrapperRunId: 'wr_cancel_report', + }); + const accepted = await getSessionMessageState(instance.ctx.storage, acceptedMessageId); + const sendReport = instance['sendRunStateReport'].bind(instance); + instance['sendRunStateReport'] = async report => { + if (report.run.status === 'interrupted') { + interruptedReportAttempts += 1; + if (failFirstReport && interruptedReportAttempts === 1) { + throw new Error('Report queue temporarily unavailable'); + } + } + return sendReport(report); + }; + const broadcastEvent = instance['broadcastEvent'].bind(instance); + instance['broadcastEvent'] = event => { + broadcastTypes.push(event.stream_event_type); + broadcastEvent(event); + }; + + expect(await instance.cancelQueuedMessage(messageId)).toEqual({ dropped: true }); + expect(await instance.cancelQueuedMessage(acceptedMessageId)).toEqual({ dropped: false }); + expect(await instance.cancelQueuedMessage('msg_missing_cancel')).toEqual({ + dropped: false, + }); + expect(await getSessionMessageState(instance.ctx.storage, acceptedMessageId)).toEqual( + accepted + ); + return { + canceled: await getSessionMessageState(instance.ctx.storage, messageId), + accepted, + }; + }); + + expect(interruptedReportAttempts).toBe(1); + expect(reports.filter(report => report.run.status === 'interrupted')).toHaveLength( + failFirstReport ? 0 : 1 + ); + expect(first.canceled).toMatchObject({ + status: 'interrupted', + completionSource: 'canceled', + failureStage: 'interruption', + failureCode: 'user_interrupt', + terminalAt: expect.any(Number), + }); + expect(first.canceled?.terminalEffects).toBeUndefined(); + + const retry = await runInDurableObject(stub, async (instance, state) => { + expect(await instance.cancelQueuedMessage(messageId)).toEqual({ dropped: true }); + const events = createEventQueries( + drizzle(state.storage, { logger: false }), + state.storage.sql + ).findByFilters({ eventTypes: ['cloud.message.canceled', 'cloud.message.failed'] }); + return { + canceled: await getSessionMessageState(instance.ctx.storage, messageId), + accepted: await getSessionMessageState(instance.ctx.storage, acceptedMessageId), + events, + }; + }); + + expect(interruptedReportAttempts).toBe(2); + expect(retry.canceled).toEqual(first.canceled); + expect(retry.accepted).toEqual(first.accepted); + expect(retry.events).toHaveLength(1); + expect(retry.events[0].stream_event_type).toBe('cloud.message.canceled'); + expect(JSON.parse(retry.events[0].payload)).toEqual({ messageId }); + expect(broadcastTypes).toEqual(['cloud.message.canceled']); + expect(callbackQueue.captured).toEqual([]); + const interruptedReports = reports.filter(report => report.run.status === 'interrupted'); + expect(interruptedReports).toHaveLength(failFirstReport ? 1 : 2); + for (const report of interruptedReports) { + expect(report.run).toMatchObject({ + messageId, + status: 'interrupted', + failureStage: 'interruption', + failureCode: 'user_interrupt', + queuedAt: new Date(first.canceled?.queuedAt ?? 0).toISOString(), + terminalAt: new Date(first.canceled?.terminalAt ?? 0).toISOString(), + }); + expect(report.run).not.toHaveProperty('dispatchAcceptedAt'); + expect(report.run).not.toHaveProperty('agentActivityObservedAt'); + expect(report.run).not.toHaveProperty('diagnostic'); + } + } + ); + it('alarm repairs terminal effects without duplicating a durable terminal event', async () => { const userId = 'user_term_repair_alarm'; const sessionId = 'agent_term_repair_alarm'; diff --git a/services/cloud-agent-next/test/unit/wrapper/batch-admission.test.ts b/services/cloud-agent-next/test/unit/wrapper/batch-admission.test.ts index 4ec8b2de5d..fb0b23d496 100644 --- a/services/cloud-agent-next/test/unit/wrapper/batch-admission.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/batch-admission.test.ts @@ -15,13 +15,20 @@ vi.mock('../../../wrapper/src/utils.js', () => ({ logToFile: vi.fn(), })); -vi.mock('../../../wrapper/src/log-uploader.js', () => ({ - createLogUploader: vi.fn(() => ({ - start: vi.fn(), - uploadNow: vi.fn().mockResolvedValue(undefined), - stop: vi.fn(), - })), -})); +vi.mock(import('../../../wrapper/src/log-uploader.js'), async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + createLogUploader: vi.fn(options => ({ + archiveId: options.archiveId, + start: vi.fn(), + updateContext: vi.fn(), + uploadNow: vi.fn().mockResolvedValue(undefined), + finalize: vi.fn().mockResolvedValue(undefined), + stop: vi.fn(), + })), + }; +}); const config: ServerConfig = { port: 5000, diff --git a/services/cloud-agent-next/vitest.workers.config.ts b/services/cloud-agent-next/vitest.workers.config.ts index a4f73f007f..988c011e06 100644 --- a/services/cloud-agent-next/vitest.workers.config.ts +++ b/services/cloud-agent-next/vitest.workers.config.ts @@ -1,3 +1,4 @@ +import { createRequire } from 'node:module'; import { cloudflareTest } from '@cloudflare/vitest-pool-workers'; import { defineConfig } from 'vitest/config'; @@ -5,6 +6,18 @@ import { defineConfig } from 'vitest/config'; // Use cloudflare:test utilities: env, runInDurableObject, createMessageBatch, etc. export default defineConfig({ plugins: [ + { + name: 'fix-pg-cjs-dependencies', + enforce: 'pre', + resolveId(source: string, importer?: string) { + if (importer === undefined) return undefined; + if (source === 'pg-protocol') { + return createRequire(importer).resolve('pg-protocol/dist/index.js'); + } + if (source === 'pg-pool') return createRequire(importer).resolve(source); + return undefined; + }, + }, cloudflareTest({ wrangler: { // Use test-specific wrangler config that excludes Sandbox DO diff --git a/services/cloud-agent-next/wrapper/src/log-uploader.test.ts b/services/cloud-agent-next/wrapper/src/log-uploader.test.ts index e7c2af03dd..2effa27297 100644 --- a/services/cloud-agent-next/wrapper/src/log-uploader.test.ts +++ b/services/cloud-agent-next/wrapper/src/log-uploader.test.ts @@ -3,13 +3,63 @@ import fs from 'node:fs'; import fsp from 'node:fs/promises'; import os from 'node:os'; import path from 'node:path'; -import { createLogUploader } from './log-uploader'; +import { gunzipSync } from 'node:zlib'; +import { createLogArchiveId, createLogUploader, type LogUploader } from './log-uploader'; const originalFetch = globalThis.fetch; +const originalLogPath = process.env.WRAPPER_LOG_PATH; const temporaryDirectories: string[] = []; +const uploaders: LogUploader[] = []; + +async function createFixture() { + const directory = await fsp.mkdtemp(path.join(os.tmpdir(), 'log-uploader-test-')); + temporaryDirectories.push(directory); + const cliLogDir = path.join(directory, 'cli-logs'); + const wrapperLogPath = path.join(directory, 'wrapper.log'); + await fsp.mkdir(cliLogDir); + await fsp.writeFile(path.join(cliLogDir, 'kilo.log'), 'kilo log'); + await fsp.writeFile(wrapperLogPath, 'wrapper log\n'); + process.env.WRAPPER_LOG_PATH = wrapperLogPath; + return { cliLogDir, wrapperLogPath }; +} + +function createUploader( + files: { cliLogDir: string; wrapperLogPath: string }, + archiveId = createLogArchiveId('run_1') +): LogUploader { + const uploader = createLogUploader({ + archiveId, + context: { + workerBaseUrl: 'https://worker.example.com', + kiloSessionId: 'kilo-session', + workerAuthToken: 'kka1.opaque', + }, + sessionId: 'agent-session', + userId: 'user', + ...files, + }); + uploaders.push(uploader); + return uploader; +} + +function mockFetch(handler: (url: URL, init: RequestInit) => Promise): void { + globalThis.fetch = Object.assign( + (input: string | URL | Request, init?: RequestInit) => + handler(new URL(input instanceof Request ? input.url : input.toString()), init ?? {}), + { preconnect: originalFetch.preconnect } + ); +} + +async function readArchive(init: RequestInit): Promise { + const bytes = await new Response(init.body).arrayBuffer(); + return gunzipSync(bytes).toString(); +} afterEach(async () => { + for (const uploader of uploaders.splice(0)) uploader.stop(); globalThis.fetch = originalFetch; + if (originalLogPath === undefined) delete process.env.WRAPPER_LOG_PATH; + else process.env.WRAPPER_LOG_PATH = originalLogPath; await Promise.all( temporaryDirectories .splice(0) @@ -18,37 +68,18 @@ afterEach(async () => { }); describe('createLogUploader', () => { - it('aborts an active upload when stopped', async () => { - const directory = await fsp.mkdtemp(path.join(os.tmpdir(), 'log-uploader-stop-test-')); - temporaryDirectories.push(directory); - const wrapperLogPath = path.join(directory, 'wrapper.log'); - await fsp.writeFile(wrapperLogPath, 'wrapper log'); - - let requestSignal: AbortSignal | undefined; - globalThis.fetch = Object.assign( - async (_input: string | URL | Request, init?: RequestInit) => { - requestSignal = init?.signal ?? undefined; - return new Promise((_resolve, reject) => { - requestSignal?.addEventListener('abort', () => reject(requestSignal?.reason), { - once: true, - }); - }); - }, - { preconnect: originalFetch.preconnect } - ); - const uploader = createLogUploader({ - workerBaseUrl: 'https://worker.example.com', - sessionId: 'agent-session', - getKiloSessionId: () => 'kilo-session', - executionId: 'session', - userId: 'user', - getWorkerAuthToken: () => 'kka1.opaque', - cliLogDir: path.join(directory, 'missing-cli-logs'), - wrapperLogPath, + it('settles an active upload when stopped even if fetch ignores abort', async () => { + const files = await createFixture(); + const requestStarted = Promise.withResolvers(); + mockFetch(async (_url, init) => { + if (!init.signal) throw new Error('Expected upload signal'); + requestStarted.resolve(init.signal); + return new Promise(() => {}); }); + const uploader = createUploader(files); const upload = uploader.uploadNow(); - while (!requestSignal) await Bun.sleep(1); + const requestSignal = await requestStarted.promise; uploader.stop(); const settled = await Promise.race([upload.then(() => true), Bun.sleep(100).then(() => false)]); @@ -56,85 +87,237 @@ describe('createLogUploader', () => { expect(settled).toBe(true); }); - it('binds log uploads to the Kilo session and sends the opaque worker credential', async () => { - const directory = await fsp.mkdtemp(path.join(os.tmpdir(), 'log-uploader-test-')); - temporaryDirectories.push(directory); - const cliLogDir = path.join(directory, 'cli-logs'); - const wrapperLogPath = path.join(directory, 'wrapper.log'); - await fsp.mkdir(cliLogDir); - await fsp.writeFile(path.join(cliLogDir, 'kilo.log'), 'kilo log'); - await fsp.writeFile(wrapperLogPath, 'wrapper log'); - + it('preserves the upload route, filename, Kilo session query and opaque credential', async () => { + const files = await createFixture(); let capturedUrl: URL | undefined; let capturedInit: RequestInit | undefined; - globalThis.fetch = Object.assign( - async (input: string | URL | Request, init?: RequestInit) => { - capturedUrl = new URL(input instanceof Request ? input.url : input.toString()); - capturedInit = init; - return new Response(null, { status: 204 }); - }, - { preconnect: originalFetch.preconnect } - ); - + let capturedArchive: string | undefined; + mockFetch(async (url, init) => { + capturedUrl = url; + capturedInit = init; + capturedArchive = await readArchive(init); + return new Response(null, { status: 204 }); + }); const uploader = createLogUploader({ - workerBaseUrl: 'https://worker.example.com', + archiveId: 'run_1--nonce', + context: { + workerBaseUrl: 'https://worker.example.com', + kiloSessionId: 'kilo/session?one', + workerAuthToken: 'kka1.opaque', + }, sessionId: 'agent/session', - getKiloSessionId: () => 'kilo/session?one', - executionId: 'session', userId: 'user@example.com', - getWorkerAuthToken: () => 'kka1.opaque', - cliLogDir, - wrapperLogPath, + ...files, }); + uploaders.push(uploader); await uploader.uploadNow(); expect(capturedUrl?.pathname).toBe( - '/sessions/user%40example.com/agent%2Fsession/logs/session/logs.tar.gz' + '/sessions/user%40example.com/agent%2Fsession/logs/run_1--nonce/logs.tar.gz' ); expect(capturedUrl?.searchParams.get('kiloSessionId')).toBe('kilo/session?one'); expect(new Headers(capturedInit?.headers).get('Authorization')).toBe('Bearer kka1.opaque'); expect(capturedInit?.body).toBeInstanceOf(ReadableStream); - expect(fs.existsSync(wrapperLogPath)).toBe(true); + expect(capturedArchive).toContain('wrapper log'); + expect(fs.existsSync(files.wrapperLogPath)).toBe(true); }); - it('reads the worker auth token fresh on every upload instead of a value captured at creation', async () => { - const directory = await fsp.mkdtemp(path.join(os.tmpdir(), 'log-uploader-refresh-test-')); - temporaryDirectories.push(directory); - const cliLogDir = path.join(directory, 'cli-logs'); - const wrapperLogPath = path.join(directory, 'wrapper.log'); - await fsp.mkdir(cliLogDir); - await fsp.writeFile(path.join(cliLogDir, 'kilo.log'), 'kilo log'); - await fsp.writeFile(wrapperLogPath, 'wrapper log'); - - const capturedAuthHeaders: Array = []; - globalThis.fetch = Object.assign( - async (_input: string | URL | Request, init?: RequestInit) => { - capturedAuthHeaders.push(new Headers(init?.headers).get('Authorization')); - return new Response(null, { status: 204 }); - }, - { preconnect: originalFetch.preconnect } - ); - - let currentToken = 'kka1.first-ticket'; - const uploader = createLogUploader({ - workerBaseUrl: 'https://worker.example.com', - sessionId: 'agent-session', - getKiloSessionId: () => 'kilo-session', - executionId: 'session', - userId: 'user', - getWorkerAuthToken: () => currentToken, - cliLogDir, - wrapperLogPath, + it('retains an earlier archive when a later wrapper uses the same run and session IDs', async () => { + const files = await createFixture(); + const archives = new Map(); + mockFetch(async (url, init) => { + archives.set(url.pathname, await readArchive(init)); + return new Response(null, { status: 204 }); }); + const first = createUploader(files); + await fsp.appendFile(files.wrapperLogPath, 'original failed bootstrap\n'); + await first.finalize(); + const firstPath = `/sessions/user/agent-session/logs/${first.archiveId}/logs.tar.gz`; + const firstArchive = archives.get(firstPath); + + await fsp.writeFile(files.wrapperLogPath, 'later successful bootstrap\n'); + const second = createUploader(files); + await second.uploadNow(); + await second.uploadNow(); + expect(second.archiveId).not.toBe(first.archiveId); + expect(archives.size).toBe(2); + expect(archives.get(firstPath)).toBe(firstArchive); + expect(firstArchive).toContain('original failed bootstrap'); + expect(firstArchive).not.toContain('later successful bootstrap'); + expect( + archives.get(`/sessions/user/agent-session/logs/${second.archiveId}/logs.tar.gz`) + ).toContain('later successful bootstrap'); + }); + + it('refreshes the complete upload context without changing the archive', async () => { + const files = await createFixture(); + const captured: Array<{ url: URL; authorization: string | null }> = []; + mockFetch(async (url, init) => { + captured.push({ url, authorization: new Headers(init.headers).get('Authorization') }); + await readArchive(init); + return new Response(null, { status: 204 }); + }); + const uploader = createUploader(files); await uploader.uploadNow(); - currentToken = 'kka1.refreshed-ticket'; + uploader.updateContext({ + workerBaseUrl: 'https://refreshed.example.com', + kiloSessionId: 'kilo-session-refreshed', + workerAuthToken: 'kka1.refreshed-ticket', + }); await uploader.uploadNow(); - expect(capturedAuthHeaders).toEqual([ - 'Bearer kka1.first-ticket', + expect(captured).toHaveLength(2); + expect(captured[1]?.url.pathname).toBe(captured[0]?.url.pathname); + expect(captured[1]?.url.origin).toBe('https://refreshed.example.com'); + expect(captured[1]?.url.searchParams.get('kiloSessionId')).toBe('kilo-session-refreshed'); + expect(captured.map(upload => upload.authorization)).toEqual([ + 'Bearer kka1.opaque', 'Bearer kka1.refreshed-ticket', ]); }); + + it.each(['uploadNow', 'finalize'] as const)( + '%s waits for a periodic upload and then uploads a fresh final snapshot', + async method => { + const files = await createFixture(); + const periodicStarted = Promise.withResolvers(); + const releasePeriodic = Promise.withResolvers(); + const captured: Array<{ url: URL; authorization: string | null; archive: string }> = []; + mockFetch(async (url, init) => { + const archive = await readArchive(init); + captured.push({ + url, + authorization: new Headers(init.headers).get('Authorization'), + archive, + }); + if (captured.length === 1) { + if (!init.signal) throw new Error('Expected upload signal'); + periodicStarted.resolve(init.signal); + await releasePeriodic.promise; + } + return new Response(null, { status: 204 }); + }); + const uploader = createUploader(files); + uploader.start(5); + const periodicSignal = await periodicStarted.promise; + await fsp.appendFile(files.wrapperLogPath, 'final failure evidence\n'); + uploader.updateContext({ + workerBaseUrl: 'https://refreshed.example.com', + kiloSessionId: 'kilo-session-refreshed', + workerAuthToken: 'kka1.refreshed-ticket', + }); + let settled = false; + const finalUpload = uploader[method]().then(() => { + settled = true; + uploader.stop(); + }); + uploader.updateContext({ + workerBaseUrl: 'https://later.example.com', + kiloSessionId: 'kilo-session-later', + workerAuthToken: 'kka1.later-ticket', + }); + await Bun.sleep(20); + expect(settled).toBe(false); + expect(periodicSignal.aborted).toBe(false); + expect(captured).toHaveLength(1); + + releasePeriodic.resolve(); + await finalUpload; + await Bun.sleep(20); + + expect(captured).toHaveLength(2); + expect(captured[0]?.archive).not.toContain('final failure evidence'); + expect(captured[1]?.archive).toContain('final failure evidence'); + expect(captured[1]?.url.pathname).toBe(captured[0]?.url.pathname); + expect(captured[0]?.url.searchParams.get('kiloSessionId')).toBe('kilo-session'); + expect(captured[0]?.authorization).toBe('Bearer kka1.opaque'); + expect(captured[1]?.url.origin).toBe('https://refreshed.example.com'); + expect(captured[1]?.url.searchParams.get('kiloSessionId')).toBe('kilo-session-refreshed'); + expect(captured[1]?.authorization).toBe('Bearer kka1.refreshed-ticket'); + } + ); + + it('bounds finalization and cancels queued work when a periodic upload never settles', async () => { + const files = await createFixture(); + const requestStarted = Promise.withResolvers(); + let fetchCalls = 0; + mockFetch(async (_url, init) => { + fetchCalls++; + if (!init.signal) throw new Error('Expected upload signal'); + requestStarted.resolve(init.signal); + return new Promise(() => {}); + }); + const uploader = createUploader(files); + uploader.start(5); + const requestSignal = await requestStarted.promise; + const finalUpload = uploader.finalize(25); + expect(uploader.finalize()).toBe(finalUpload); + const settled = await Promise.race([ + finalUpload.then(() => true), + Bun.sleep(250).then(() => false), + ]); + await Bun.sleep(20); + + expect(settled).toBe(true); + expect(requestSignal.aborted).toBe(true); + expect(fetchCalls).toBe(1); + }); + + it('swallows HTTP and fetch failures without logging credentials or getting stuck', async () => { + const files = await createFixture(); + const secret = 'kka1.do-not-log-this-ticket'; + let fetchCalls = 0; + mockFetch(async (_url, init) => { + await readArchive(init); + fetchCalls++; + if (fetchCalls === 1) return new Response(null, { status: 403, statusText: secret }); + if (fetchCalls === 2) throw new Error(`Authorization: Bearer ${secret}`); + return new Response(null, { status: 204 }); + }); + const uploader = createUploader(files); + uploader.updateContext({ + workerBaseUrl: 'https://worker.example.com', + kiloSessionId: 'kilo-session', + workerAuthToken: secret, + }); + + expect(await uploader.uploadNow()).toBeUndefined(); + expect(await uploader.uploadNow()).toBeUndefined(); + expect(await uploader.finalize()).toBeUndefined(); + const logs = await fsp.readFile(files.wrapperLogPath, 'utf8'); + + expect(fetchCalls).toBe(3); + expect(logs).toContain('Log upload failed: 403'); + expect(logs).toContain('Log upload did not complete'); + expect(logs).not.toContain(secret); + expect(logs).not.toContain('Authorization'); + }); + + it('swallows setup failures and permits subsequent uploads', async () => { + const files = await createFixture(); + let fetchCalls = 0; + mockFetch(async (_url, init) => { + fetchCalls++; + await readArchive(init); + return new Response(null, { status: 204 }); + }); + const uploader = createUploader(files); + uploader.updateContext({ + workerBaseUrl: 'invalid URL kka1.secret', + kiloSessionId: 'kilo-session', + workerAuthToken: 'kka1.secret', + }); + expect(await uploader.uploadNow()).toBeUndefined(); + uploader.updateContext({ + workerBaseUrl: 'https://worker.example.com', + kiloSessionId: 'kilo-session', + workerAuthToken: 'kka1.refreshed-ticket', + }); + await uploader.uploadNow(); + + expect(fetchCalls).toBe(1); + expect(await fsp.readFile(files.wrapperLogPath, 'utf8')).not.toContain('kka1.secret'); + }); }); diff --git a/services/cloud-agent-next/wrapper/src/log-uploader.ts b/services/cloud-agent-next/wrapper/src/log-uploader.ts index 8c7b0b1966..e1238da9dc 100644 --- a/services/cloud-agent-next/wrapper/src/log-uploader.ts +++ b/services/cloud-agent-next/wrapper/src/log-uploader.ts @@ -1,38 +1,40 @@ +import { randomUUID } from 'node:crypto'; import { existsSync } from 'node:fs'; import { basename, dirname } from 'node:path'; import { spawn } from 'node:child_process'; -import { logToFile } from './utils.js'; +import { logToFile, withTimeoutAndAbort } from './utils.js'; -type LogUploaderOpts = { +type LogUploadContext = { workerBaseUrl: string; + kiloSessionId: string; + workerAuthToken: string; +}; + +type LogUploaderOpts = { + archiveId: string; + context: LogUploadContext; sessionId: string; - /** - * Read at upload time rather than captured once at creation: a later session - * bind can carry a different kiloSessionId, and the Worker rejects uploads - * whose query param disagrees with the ticket's kiloSessionId claim. - */ - getKiloSessionId: () => string; - executionId: string; userId: string; - /** - * Reads the current wrapper dispatch ticket at upload time rather than a - * value captured once at creation. The uploader is created once per wrapper - * process and outlives the ticket's own (shorter) lifetime, which is - * refreshed on every subsequent session bind. - */ - getWorkerAuthToken: () => string; /** Directory containing CLI log files (e.g. ~/.local/share/kilo/log/) */ cliLogDir: string; wrapperLogPath: string; }; export type LogUploader = { + readonly archiveId: string; start: (intervalMs?: number) => void; + updateContext: (context: LogUploadContext) => void; uploadNow: () => Promise; + finalize: (timeoutMs?: number) => Promise; stop: () => void; }; const UPLOAD_TIMEOUT_MS = 15_000; +const FINAL_UPLOAD_TIMEOUT_MS = 5_000; + +export function createLogArchiveId(wrapperRunId: string): string { + return `${wrapperRunId}--${randomUUID()}`; +} type TarStream = { stream: ReadableStream; @@ -53,12 +55,9 @@ function createTarStream(paths: Array): TarStream | undefined { const { stdout, stderr: stderrStream } = proc; if (!stdout || !stderrStream) return undefined; - let stderr = ''; - stderrStream.on('data', (chunk: Buffer) => { - stderr += chunk.toString(); - }); + stderrStream.resume(); proc.on('close', code => { - if (code !== 0) logToFile(`tar exited with code ${code}: ${stderr}`); + if (code !== 0) logToFile(`tar exited with code ${code}`); }); const stream = new ReadableStream({ @@ -82,7 +81,7 @@ function createTarStream(paths: Array): TarStream | undefined { stdout.on('end', close); stdout.on('error', error); proc.on('error', err => { - logToFile(`tar spawn error: ${err.message}`); + logToFile('tar spawn error'); error(err); }); }, @@ -92,67 +91,120 @@ function createTarStream(paths: Array): TarStream | undefined { } export function createLogUploader(opts: LogUploaderOpts): LogUploader { + type Upload = { promise: Promise; abort: AbortController }; + + const { archiveId, sessionId, userId, cliLogDir, wrapperLogPath } = opts; + let context = { ...opts.context }; let intervalId: ReturnType | undefined; - let isUploading = false; - let activeAbort: AbortController | undefined; - - async function uploadNow(): Promise { - if (isUploading) return; - isUploading = true; - const tar = createTarStream([opts.cliLogDir, opts.wrapperLogPath]); - if (!tar) { - isUploading = false; - return; - } + let activeUpload: Upload | undefined; + let queuedUpload: Upload | undefined; + let finalUpload: Promise | undefined; + let stopped = false; + function uploadNow(): Promise { + if (finalUpload) return finalUpload; + if (stopped) return Promise.resolve(); + if (queuedUpload) return queuedUpload.promise; + + const uploadContext = { ...context }; + const previousUpload = activeUpload; const abort = new AbortController(); - activeAbort = abort; - const timer = setTimeout(() => abort.abort(), UPLOAD_TIMEOUT_MS); - try { - const url = new URL( - `${opts.workerBaseUrl}/sessions/${encodeURIComponent(opts.userId)}/${encodeURIComponent(opts.sessionId)}/logs/${encodeURIComponent(opts.executionId)}/logs.tar.gz` - ); - url.searchParams.set('kiloSessionId', opts.getKiloSessionId()); - const response = await fetch(url, { - method: 'PUT', - headers: { Authorization: `Bearer ${opts.getWorkerAuthToken()}` }, - body: tar.stream, - // @ts-expect-error -- Node/Bun fetch supports duplex for streaming request bodies - duplex: 'half', - signal: abort.signal, - }); - if (!response.ok) { - logToFile(`Log upload failed: ${response.status} ${response.statusText}`); - } - } catch (error) { - if (abort.signal.aborted) { - logToFile(`Log upload timed out after ${UPLOAD_TIMEOUT_MS}ms`); - } else { - logToFile(`Log upload error: ${error instanceof Error ? error.message : String(error)}`); + const upload: Upload = { promise: performUpload(), abort }; + if (previousUpload) queuedUpload = upload; + else activeUpload = upload; + return upload.promise; + + async function performUpload(): Promise { + let tar: TarStream | undefined; + try { + await withTimeoutAndAbort( + (async () => { + await previousUpload?.promise; + abort.signal.throwIfAborted(); + if (queuedUpload === upload) queuedUpload = undefined; + activeUpload = upload; + + const url = new URL( + `${uploadContext.workerBaseUrl}/sessions/${encodeURIComponent(userId)}/${encodeURIComponent(sessionId)}/logs/${encodeURIComponent(archiveId)}/logs.tar.gz` + ); + url.searchParams.set('kiloSessionId', uploadContext.kiloSessionId); + tar = createTarStream([cliLogDir, wrapperLogPath]); + if (!tar) return; + + const response = await fetch(url, { + method: 'PUT', + headers: { Authorization: `Bearer ${uploadContext.workerAuthToken}` }, + body: tar.stream, + // @ts-expect-error -- Node/Bun fetch supports duplex for streaming request bodies + duplex: 'half', + signal: abort.signal, + }); + if (!abort.signal.aborted && !response.ok) { + logToFile(`Log upload failed: ${response.status}`); + } + })(), + { + timeoutMs: UPLOAD_TIMEOUT_MS, + timeoutMessage: 'Log upload timed out', + signal: abort.signal, + abortMessage: 'Log upload aborted', + } + ); + } catch { + logToFile('Log upload did not complete'); + } finally { + abort.abort(); + tar?.kill(); + if (activeUpload === upload) activeUpload = undefined; + if (queuedUpload === upload) queuedUpload = undefined; } - } finally { - clearTimeout(timer); - tar.kill(); - isUploading = false; - if (activeAbort === abort) activeAbort = undefined; + } + } + + function clearUploadInterval(): void { + if (intervalId !== undefined) { + clearInterval(intervalId); + intervalId = undefined; } } function start(intervalMs = 30_000): void { stop(); + stopped = false; + finalUpload = undefined; intervalId = setInterval(() => { - uploadNow().catch(() => {}); + if (!activeUpload && !queuedUpload) void uploadNow(); }, intervalMs); } + function finalize(timeoutMs = FINAL_UPLOAD_TIMEOUT_MS): Promise { + if (finalUpload) return finalUpload; + clearUploadInterval(); + finalUpload = withTimeoutAndAbort(uploadNow(), { + timeoutMs, + timeoutMessage: 'Final log upload timed out', + abortMessage: 'Final log upload aborted', + }) + .catch(() => logToFile('Final log upload timed out')) + .finally(stop); + return finalUpload; + } + function stop(): void { - if (intervalId !== undefined) { - clearInterval(intervalId); - intervalId = undefined; - } - activeAbort?.abort(); - activeAbort = undefined; + stopped = true; + clearUploadInterval(); + activeUpload?.abort.abort(); + queuedUpload?.abort.abort(); } - return { start, uploadNow, stop }; + return { + archiveId, + start, + updateContext: nextContext => { + context = { ...nextContext }; + }, + uploadNow, + finalize, + stop, + }; } diff --git a/services/cloud-agent-next/wrapper/src/main.ts b/services/cloud-agent-next/wrapper/src/main.ts index 286c750775..d7a7c14c19 100644 --- a/services/cloud-agent-next/wrapper/src/main.ts +++ b/services/cloud-agent-next/wrapper/src/main.ts @@ -588,7 +588,8 @@ async function main() { } async function readySession( - request: WrapperSessionReadyRequest + request: WrapperSessionReadyRequest, + logArchiveId: string ): Promise { if (isShuttingDown) return wrapperFinalizingResponse(); @@ -610,7 +611,8 @@ async function main() { serverConfig, serverDeps, 'close-until-runtime-ready', - request.kiloSessionId + request.kiloSessionId, + logArchiveId ); if (bindError) { const error = (await bindError.json()) as { @@ -634,6 +636,7 @@ async function main() { }; } + const bootstrapLogUploader = state.logUploader; serverConfig.workspacePath = request.workspace.workspacePath; serverConfig.sessionId = request.kiloSessionId; serverConfig.platform = request.materialized.env.KILO_PLATFORM ?? process.env.KILO_PLATFORM; @@ -785,7 +788,13 @@ async function main() { const workspaceBootstrap = prepareWrapperBootstrapWorkspace( request, emitBootstrapProgress, - {}, + { + beforeFailureCleanup: async () => { + if (bootstrapLogUploader?.archiveId === logArchiveId) { + await bootstrapLogUploader.finalize(); + } + }, + }, workspaceBootstrapController.signal ); activeWorkspaceBootstraps.add(workspaceBootstrap); @@ -986,9 +995,7 @@ async function main() { // Best-effort final log upload const uploader = state.logUploader; if (uploader) { - const uploadTimeout = new Promise(resolve => setTimeout(resolve, 5_000)); - await Promise.race([uploader.uploadNow().catch(() => {}), uploadTimeout]); - uploader.stop(); + await uploader.finalize(); } // Close connections @@ -1028,9 +1035,7 @@ async function main() { const uploader = state.logUploader; if (uploader) { - const timeout = new Promise(resolve => setTimeout(resolve, 5_000)); - void Promise.race([uploader.uploadNow().catch(() => {}), timeout]).finally(() => { - uploader.stop(); + void uploader.finalize().finally(() => { process.exit(1); }); } else { diff --git a/services/cloud-agent-next/wrapper/src/server.test.ts b/services/cloud-agent-next/wrapper/src/server.test.ts index acdd135cdf..0cfa75db25 100644 --- a/services/cloud-agent-next/wrapper/src/server.test.ts +++ b/services/cloud-agent-next/wrapper/src/server.test.ts @@ -1,5 +1,9 @@ import { afterEach, describe, expect, it } from 'bun:test'; +import fsp from 'node:fs/promises'; import { createServer as createNetServer } from 'node:net'; +import os from 'node:os'; +import path from 'node:path'; +import { gunzipSync } from 'node:zlib'; import { WrapperState } from './state'; import { bindSessionContext, @@ -9,10 +13,12 @@ import { createServer, createSessionReadyHandler, resolvePtyClientClose, + type ServerDependencies, type WrapperServer, } from './server'; import type { WrapperKiloClient, WrapperPty, WrapperPtySize } from './kilo-api'; import { PNPM_STORE_DIR, PNPM_STORE_ENV_VAR } from '../../src/shared/runtime-environment.js'; +import type { WrapperSessionReadyRequest } from '../../src/shared/wrapper-bootstrap.js'; type PtyCall = { cwd: string; @@ -21,6 +27,16 @@ type PtyCall = { }; const servers: WrapperServer[] = []; +const states: WrapperState[] = []; +const originalFetch = globalThis.fetch; +const originalLogPath = process.env.WRAPPER_LOG_PATH; +const temporaryDirectories: string[] = []; + +function createTestState(): WrapperState { + const state = new WrapperState(); + states.push(state); + return state; +} function getFreePort(): Promise { return new Promise((resolve, reject) => { @@ -89,7 +105,7 @@ function createTestFetch(overrides?: { wrapperInstanceGeneration: 8, }, { - state: new WrapperState(), + state: createTestState(), kiloClient, openConnection: async () => {}, closeConnection: async () => {}, @@ -105,7 +121,16 @@ function createTestFetch(overrides?: { } afterEach(async () => { + for (const state of states.splice(0)) state.clearSession(); await Promise.all(servers.splice(0).map(server => server.stop())); + globalThis.fetch = originalFetch; + if (originalLogPath === undefined) delete process.env.WRAPPER_LOG_PATH; + else process.env.WRAPPER_LOG_PATH = originalLogPath; + await Promise.all( + temporaryDirectories + .splice(0) + .map(directory => fsp.rm(directory, { recursive: true, force: true })) + ); }); describe('kilo server unreachable recovery', () => { @@ -120,7 +145,7 @@ describe('kilo server unreachable recovery', () => { }; function boundState(): WrapperState { - const state = new WrapperState(); + const state = createTestState(); state.bindSession(sessionBinding); return state; } @@ -271,7 +296,7 @@ describe('session readiness errors', () => { it('forwards validated workspace subtype and safe diagnostic fields', async () => { const { fetchHandler } = createTestFetch(); const handler = createSessionReadyHandler({ - state: new WrapperState(), + state: createTestState(), kiloClient: {} as WrapperKiloClient, openConnection: async () => {}, closeConnection: async () => {}, @@ -430,7 +455,7 @@ describe('wrapper PTY routes', () => { userId: 'user_test', }, { - state: new WrapperState(), + state: createTestState(), kiloClient, openConnection: async () => {}, closeConnection: async () => {}, @@ -524,7 +549,7 @@ describe('wrapper Kilo proxy route', () => { userId: 'user_test', }, { - state: new WrapperState(), + state: createTestState(), kiloClient, openConnection: async () => {}, closeConnection: async () => {}, @@ -548,6 +573,328 @@ describe('wrapper Kilo proxy route', () => { }); }); +describe('wrapper log archive retention', () => { + const binding = { + ingestUrl: 'wss://worker.test/ingest', + workerAuthToken: 'kka1.first-ticket', + wrapperRunId: 'run_1', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_1', + }; + const config = { + port: 5000, + workspacePath: '/workspace/repo', + version: 'test', + sessionId: 'kilo_sess_test', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + userId: 'user_test', + }; + const readyRequest: WrapperSessionReadyRequest = { + agentSessionId: config.agentSessionId, + userId: config.userId, + sandboxId: 'sandbox_test', + kiloSessionId: config.sessionId, + workspace: { + workspacePath: config.workspacePath, + sessionHome: '/home/session', + branchName: 'main', + }, + materialized: { env: {} }, + preparation: { attemptId: 'preparation_1', triggerMessageId: 'message_1' }, + session: binding, + }; + + function requestReady(): Request { + return new Request('http://wrapper.test/session/ready', { + method: 'POST', + body: JSON.stringify(readyRequest), + }); + } + + async function createArchiveFixture() { + const directory = await fsp.mkdtemp(path.join(os.tmpdir(), 'wrapper-archive-test-')); + temporaryDirectories.push(directory); + const wrapperLogPath = path.join(directory, 'wrapper.log'); + await fsp.writeFile(wrapperLogPath, 'wrapper started\n'); + process.env.WRAPPER_LOG_PATH = wrapperLogPath; + const archives = new Map(); + const uploads: Array<{ url: URL; authorization: string | null }> = []; + const state = createTestState(); + const deps: ServerDependencies = { + state, + kiloClient: {} as WrapperKiloClient, + openConnection: async () => {}, + closeConnection: async () => {}, + setAborted: () => {}, + resetLifecycle: () => {}, + }; + globalThis.fetch = Object.assign( + async (input: string | URL | Request, init?: RequestInit) => { + const url = new URL(input instanceof Request ? input.url : input.toString()); + const bytes = await new Response(init?.body).arrayBuffer(); + uploads.push({ url, authorization: new Headers(init?.headers).get('Authorization') }); + archives.set(url.pathname, gunzipSync(bytes).toString()); + return new Response(null, { status: 204 }); + }, + { preconnect: originalFetch.preconnect } + ); + return { state, deps, wrapperLogPath, archives, uploads }; + } + + it('retains the old archive and captured credentials when a warm wrapper changes runs', async () => { + const { state, deps, wrapperLogPath, archives, uploads } = await createArchiveFixture(); + await bindSessionContext(binding, config, deps); + const first = state.logUploader; + if (!first) throw new Error('Expected first archive'); + const recordUpload = globalThis.fetch; + const periodicStarted = Promise.withResolvers(); + const releasePeriodic = Promise.withResolvers(); + let firstUpload = true; + globalThis.fetch = Object.assign( + async (input: string | URL | Request, init?: RequestInit) => { + const response = await recordUpload(input, init); + if (firstUpload) { + firstUpload = false; + periodicStarted.resolve(); + await releasePeriodic.promise; + } + return response; + }, + { preconnect: originalFetch.preconnect } + ); + first.start(5); + await periodicStarted.promise; + await fsp.appendFile(wrapperLogPath, 'original run failed\n'); + + const rebind = bindSessionContext( + { + ...binding, + ingestUrl: 'wss://next-worker.test/ingest', + workerAuthToken: 'kka1.next-run-ticket', + wrapperRunId: 'run_2', + wrapperGeneration: 2, + wrapperConnectionId: 'conn_2', + }, + config, + deps, + 'restart', + 'kilo_sess_next' + ); + expect(await Promise.race([rebind.then(() => true), Bun.sleep(100).then(() => false)])).toBe( + true + ); + const second = state.logUploader; + if (!second) throw new Error('Expected second archive'); + expect(second.archiveId).toMatch(/^run_2--[a-f0-9-]+$/); + expect(second.archiveId).not.toBe(first.archiveId); + expect(state.currentSession?.wrapperRunId).toBe('run_2'); + + await second.uploadNow(); + releasePeriodic.resolve(); + await first.finalize(); + const firstPath = `/sessions/user_test/${config.agentSessionId}/logs/${first.archiveId}/logs.tar.gz`; + const firstArchive = archives.get(firstPath); + await fsp.appendFile(wrapperLogPath, 'later run completed\n'); + await second.uploadNow(); + + expect(archives.size).toBe(2); + expect(firstArchive).toContain('original run failed'); + expect(firstArchive).not.toContain('later run completed'); + expect(archives.get(firstPath)).toBe(firstArchive); + for (const upload of uploads) { + if (upload.url.pathname === firstPath) { + expect(upload.url.origin).toBe('https://worker.test'); + expect(upload.url.searchParams.get('kiloSessionId')).toBe(config.sessionId); + expect(upload.authorization).toBe('Bearer kka1.first-ticket'); + } else { + expect(upload.url.origin).toBe('https://next-worker.test'); + expect(upload.url.searchParams.get('kiloSessionId')).toBe('kilo_sess_next'); + expect(upload.authorization).toBe('Bearer kka1.next-run-ticket'); + } + } + }); + + it('keeps the archive for same-run credential refreshes and refreshes the Kilo session', async () => { + const { state, deps, uploads } = await createArchiveFixture(); + await bindSessionContext(binding, config, deps); + const uploader = state.logUploader; + if (!uploader) throw new Error('Expected archive'); + await uploader.uploadNow(); + const refreshedBinding = { + ...binding, + ingestUrl: 'wss://refreshed-worker.test/ingest', + workerAuthToken: 'kka1.refreshed-ticket', + ingestToken: 'refreshed-ingest-ticket', + wrapperConnectionId: 'conn_refreshed', + }; + await bindSessionContext(refreshedBinding, config, deps, 'restart', 'kilo_sess_refreshed'); + await state.logUploader?.uploadNow(); + await bindSessionContext(refreshedBinding, config, deps, 'restart', 'kilo_sess_latest'); + await state.logUploader?.uploadNow(); + + expect(state.logUploader).toBe(uploader); + expect(state.currentSession?.workerAuthToken).toBe(refreshedBinding.workerAuthToken); + expect(state.currentSession?.kiloSessionId).toBe('kilo_sess_latest'); + expect(new Set(uploads.map(upload => upload.url.pathname)).size).toBe(1); + expect(uploads.map(upload => upload.url.searchParams.get('kiloSessionId'))).toEqual([ + config.sessionId, + 'kilo_sess_refreshed', + 'kilo_sess_latest', + ]); + expect(uploads.map(upload => upload.authorization)).toEqual([ + 'Bearer kka1.first-ticket', + 'Bearer kka1.refreshed-ticket', + 'Bearer kka1.refreshed-ticket', + ]); + expect(uploads[1]?.url.origin).toBe('https://refreshed-worker.test'); + }); + + it('retains a failed bootstrap archive when retrying the same run and preparation attempt', async () => { + const { state, deps, wrapperLogPath, archives, uploads } = await createArchiveFixture(); + let attempts = 0; + deps.readySession = async (request, archiveId) => { + attempts++; + await bindSessionContext( + request.session, + config, + deps, + 'close-until-runtime-ready', + request.kiloSessionId, + archiveId + ); + if (attempts === 1) { + await fsp.appendFile(wrapperLogPath, 'failed bootstrap evidence\n'); + return { + status: 'error', + error: { + code: 'WORKSPACE_SETUP_FAILED', + message: 'Workspace setup failed', + retryable: true, + }, + }; + } + await fsp.appendFile(wrapperLogPath, 'bootstrap retry ready\n'); + return { + status: 'ready', + kiloSessionId: request.kiloSessionId, + workspaceReady: { + ...request.workspace, + sandboxId: request.sandboxId, + kiloSessionId: request.kiloSessionId, + }, + }; + }; + const handler = createSessionReadyHandler(deps); + const firstResponse = await handler(requestReady()); + const firstPath = uploads[0]?.url.pathname; + if (!firstPath) throw new Error('Expected failed bootstrap upload'); + const failedArchive = archives.get(firstPath); + expect(state.logUploader).toBeNull(); + + const retryResponse = await handler(requestReady()); + await state.logUploader?.uploadNow(); + await state.logUploader?.uploadNow(); + + expect(firstResponse.status).toBe(503); + expect(retryResponse.status).toBe(200); + expect(attempts).toBe(2); + expect(archives.size).toBe(2); + expect(failedArchive).toContain('failed bootstrap evidence'); + expect(failedArchive).not.toContain('bootstrap retry ready'); + expect(archives.get(firstPath)).toBe(failedArchive); + expect(uploads[1]?.url.pathname).not.toBe(firstPath); + expect(uploads[2]?.url.pathname).toBe(uploads[1]?.url.pathname); + expect(state.logUploader?.archiveId).toMatch(/^run_1--[a-f0-9-]+$/); + }); + + it('shares an archive for duplicate in-flight ready requests but rotates for the next bootstrap', async () => { + const { state, deps, archives } = await createArchiveFixture(); + const bootstrapStarted = Promise.withResolvers(); + const releaseBootstrap = Promise.withResolvers(); + let attempts = 0; + deps.readySession = async (request, archiveId) => { + attempts++; + await bindSessionContext( + request.session, + config, + deps, + 'close-until-runtime-ready', + request.kiloSessionId, + archiveId + ); + bootstrapStarted.resolve(); + await releaseBootstrap.promise; + return { + status: 'ready', + kiloSessionId: request.kiloSessionId, + workspaceReady: { + ...request.workspace, + sandboxId: request.sandboxId, + kiloSessionId: request.kiloSessionId, + }, + }; + }; + const handler = createSessionReadyHandler(deps); + const firstRequest = handler(requestReady()); + await bootstrapStarted.promise; + const first = state.logUploader; + if (!first) throw new Error('Expected bootstrap archive'); + const duplicate = handler(requestReady()); + await Bun.sleep(10); + await first.uploadNow(); + + expect(attempts).toBe(1); + expect(state.logUploader).toBe(first); + expect(archives.size).toBe(1); + releaseBootstrap.resolve(); + const responses = await Promise.all([firstRequest, duplicate]); + expect(responses.map(response => response.status)).toEqual([200, 200]); + + const nextResponse = await handler(requestReady()); + await first.finalize(); + await state.logUploader?.uploadNow(); + expect(nextResponse.status).toBe(200); + expect(attempts).toBe(2); + expect(state.logUploader?.archiveId).not.toBe(first.archiveId); + expect(archives.size).toBe(2); + }); + + it('does not replace a bootstrap failure with a log upload failure or log its credential', async () => { + const { deps, state, wrapperLogPath } = await createArchiveFixture(); + globalThis.fetch = Object.assign( + async () => { + throw new Error(`Authorization: Bearer ${binding.workerAuthToken}`); + }, + { preconnect: originalFetch.preconnect } + ); + deps.readySession = async (request, archiveId) => { + await bindSessionContext( + request.session, + config, + deps, + 'close-until-runtime-ready', + request.kiloSessionId, + archiveId + ); + return { + status: 'error', + error: { code: 'KILO_SERVER_FAILED', message: 'Kilo server failed', retryable: true }, + }; + }; + + const response = await createSessionReadyHandler(deps)(requestReady()); + + expect(response.status).toBe(503); + expect(await response.json()).toEqual({ + error: 'KILO_SERVER_FAILED', + message: 'Kilo server failed', + retryable: true, + }); + expect(state.logUploader).toBeNull(); + expect(await fsp.readFile(wrapperLogPath, 'utf8')).not.toContain(binding.workerAuthToken); + }); +}); + describe('wrapper session binding', () => { it('binds the kiloSessionId supplied by the caller even when config has not been updated yet', async () => { // A freshly bootstrapped wrapper's ServerConfig.sessionId starts out empty @@ -555,7 +902,7 @@ describe('wrapper session binding', () => { // processed. bindSessionContext must use the id the caller already knows // (from the ready request), not a stale config value, since that id also // seeds the log uploader's kiloSessionId for the wrapper's entire life. - const state = new WrapperState(); + const state = createTestState(); const response = await bindSessionContext( { @@ -590,7 +937,7 @@ describe('wrapper session binding', () => { }); it('rejects even the current binding while the wrapper is finalizing', async () => { - const state = new WrapperState(); + const state = createTestState(); const sessionBinding = { kiloSessionId: 'kilo_sess_test', ingestUrl: 'ws://worker.test/ingest', @@ -633,7 +980,7 @@ describe('wrapper session binding', () => { }); it('keeps bootstrap rebindings close-only until runtime readiness is verified', async () => { - const state = new WrapperState(); + const state = createTestState(); state.bindSession({ kiloSessionId: 'kilo_sess_test', ingestUrl: 'ws://worker.test/ingest', @@ -683,7 +1030,7 @@ describe('wrapper session binding', () => { }); it('closes the bootstrap feed for an unchanged binding until runtime readiness is verified', async () => { - const state = new WrapperState(); + const state = createTestState(); const sessionBinding = { kiloSessionId: 'kilo_sess_test', ingestUrl: 'ws://worker.test/ingest', @@ -733,7 +1080,7 @@ describe('wrapper session binding', () => { }); it('keeps restart behavior for legacy direct rebindings', async () => { - const state = new WrapperState(); + const state = createTestState(); state.bindSession({ kiloSessionId: 'kilo_sess_test', ingestUrl: 'ws://worker.test/ingest', @@ -779,7 +1126,7 @@ describe('wrapper session binding', () => { }); it('resets lifecycle state when warm rebinding an existing connected session', async () => { - const state = new WrapperState(); + const state = createTestState(); state.bindSession({ kiloSessionId: 'kilo_sess_test', ingestUrl: 'ws://worker.test/ingest', diff --git a/services/cloud-agent-next/wrapper/src/server.ts b/services/cloud-agent-next/wrapper/src/server.ts index 5df81bf126..a19798d9f3 100644 --- a/services/cloud-agent-next/wrapper/src/server.ts +++ b/services/cloud-agent-next/wrapper/src/server.ts @@ -13,13 +13,14 @@ * - POST /job/abort - Abort the current session */ +import { isDeepStrictEqual } from 'node:util'; import type { WrapperState, SessionContext } from './state.js'; import { isKiloServerUnreachableError, type WrapperKiloClient, type WrapperPtySize, } from './kilo-api.js'; -import { createLogUploader } from './log-uploader.js'; +import { createLogArchiveId, createLogUploader } from './log-uploader.js'; import { configureCommitCoAuthorHook } from './commit-co-author-hook.js'; import { logToFile } from './utils.js'; import { materializePromptAttachments as defaultMaterializePromptAttachments } from './session-bootstrap.js'; @@ -70,7 +71,10 @@ export type ServerDependencies = { /** Notify lifecycle after an acknowledgement guard clears. */ onDeliveryAcknowledged?: (kind: 'async-prompt' | 'sync-command' | 'failed') => void; /** Workspace/Kilo readiness path */ - readySession?: (request: WrapperSessionReadyRequest) => Promise; + readySession?: ( + request: WrapperSessionReadyRequest, + logArchiveId: string + ) => Promise; /** Apply refreshed runtime variables to the active Kilo runtime. */ updateRuntimeEnvironment?: (env: Record) => Promise; /** Materialize signed prompt attachments into local file parts. */ @@ -277,7 +281,8 @@ export async function bindSessionContext( * Falls back to `config.sessionId` for callers that rebind an * already-bootstrapped wrapper, where `config.sessionId` is already current. */ - kiloSessionIdOverride?: string + kiloSessionIdOverride?: string, + logArchiveId?: string ): Promise { const { state } = deps; const kiloSessionId = kiloSessionIdOverride ?? config.sessionId; @@ -325,6 +330,29 @@ export async function bindSessionContext( } const existingSession = state.currentSession; + const uploadContext = { + workerBaseUrl, + kiloSessionId, + workerAuthToken: binding.workerAuthToken, + }; + if ( + !state.logUploader || + existingSession?.wrapperRunId !== binding.wrapperRunId || + (logArchiveId !== undefined && state.logUploader.archiveId !== logArchiveId) + ) { + const logUploader = createLogUploader({ + archiveId: logArchiveId ?? createLogArchiveId(binding.wrapperRunId), + context: uploadContext, + sessionId: config.agentSessionId, + userId: config.userId, + cliLogDir: `/home/${config.agentSessionId}/.local/share/kilo/log`, + wrapperLogPath: process.env.WRAPPER_LOG_PATH ?? '/tmp/kilocode-wrapper.log', + }); + state.setLogUploader(logUploader); + logUploader.start(); + } else { + state.logUploader.updateContext(uploadContext); + } if (!existingSession) { if (state.isConnected) { @@ -345,20 +373,6 @@ export async function bindSessionContext( }; state.bindSession(sessionContext); - const cliLogDir = `/home/${config.agentSessionId}/.local/share/kilo/log`; - const wrapperLogPath = process.env.WRAPPER_LOG_PATH ?? '/tmp/kilocode-wrapper.log'; - const logUploader = createLogUploader({ - workerBaseUrl, - sessionId: config.agentSessionId, - getKiloSessionId: () => state.currentSession?.kiloSessionId ?? kiloSessionId, - executionId: 'session', - userId: config.userId, - getWorkerAuthToken: () => state.currentSession?.workerAuthToken ?? binding.workerAuthToken, - cliLogDir, - wrapperLogPath, - }); - state.setLogUploader(logUploader); - logUploader.start(); logToFile(`session bound: sessionId=${kiloSessionId}`); await notifySessionBound(deps, feedPolicy); return null; @@ -1026,8 +1040,16 @@ function createWebSocketHandlers(config: ServerConfig, deps: ServerDependencies) } export function createSessionReadyHandler(deps: ServerDependencies) { + let inFlight: + | { + request: WrapperSessionReadyRequest; + promise: Promise; + } + | undefined; + return async (req: Request): Promise => { - if (!deps.readySession) { + const readySession = deps.readySession; + if (!readySession) { return errorResponse('NOT_READY', 'Wrapper readiness executor is not configured', 503); } @@ -1042,7 +1064,33 @@ export function createSessionReadyHandler(deps: ServerDependencies) { return errorResponse('INVALID_REQUEST', 'Invalid session ready request', 400); } - const result = await deps.readySession(body); + let pending = inFlight; + if (!pending || !isDeepStrictEqual(pending.request, body)) { + const archiveId = createLogArchiveId(body.session.wrapperRunId); + const promise = (async () => { + let ready = false; + try { + const result = await readySession(body, archiveId); + ready = result.status === 'ready'; + return result; + } finally { + const uploader = deps.state.logUploader; + if (!ready && uploader?.archiveId === archiveId) { + await uploader.finalize(); + if (deps.state.logUploader === uploader) deps.state.setLogUploader(null); + } + } + })(); + pending = { request: body, promise }; + inFlight = pending; + } + + let result: WrapperSessionReadyResponse; + try { + result = await pending.promise; + } finally { + if (inFlight === pending) inFlight = undefined; + } if (result.status === 'error') { const status = result.error.code === 'INVALID_REQUEST' ? 400 : 503; return jsonResponse( diff --git a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts index 3df046ed69..cb39152714 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts @@ -3,6 +3,11 @@ import fs from 'node:fs'; import fsp from 'node:fs/promises'; import os from 'node:os'; import path from 'node:path'; +import { gunzipSync } from 'node:zlib'; +import { createLogUploader, type LogUploader } from './log-uploader'; +import { createSessionReadyHandler, type ServerDependencies } from './server'; +import { WrapperState } from './state'; +import type { WrapperKiloClient } from './kilo-api'; import { materializePromptAttachments, prepareWrapperBootstrapWorkspace, @@ -96,11 +101,15 @@ function gitCredentialsPath(sessionHome: string): string { describe('prepareWrapperBootstrapWorkspace', () => { let tmpDir: string; let originalEnv: Record; + let originalFetch: typeof fetch; + const uploaders: LogUploader[] = []; beforeEach(() => { tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'wrapper-bootstrap-')); + originalFetch = globalThis.fetch; originalEnv = { HOME: process.env.HOME, + WRAPPER_LOG_PATH: process.env.WRAPPER_LOG_PATH, KILOCODE_TOKEN: process.env.KILOCODE_TOKEN, GH_TOKEN: process.env.GH_TOKEN, GITLAB_TOKEN: process.env.GITLAB_TOKEN, @@ -110,6 +119,8 @@ describe('prepareWrapperBootstrapWorkspace', () => { }); afterEach(() => { + for (const uploader of uploaders.splice(0)) uploader.stop(); + globalThis.fetch = originalFetch; for (const [key, value] of Object.entries(originalEnv)) { if (value === undefined) { delete process.env[key]; @@ -1027,6 +1038,187 @@ describe('prepareWrapperBootstrapWorkspace', () => { expect(fs.existsSync(request.workspace.sessionHome)).toBe(false); }); + it.each([ + { stage: 'cold', uploadResult: 'success' }, + { stage: 'cold', uploadResult: 'failure' }, + { stage: 'cold', uploadResult: 'timeout' }, + { stage: 'restored', uploadResult: 'success' }, + { stage: 'restored', uploadResult: 'failure' }, + { stage: 'restored', uploadResult: 'timeout' }, + ])( + 'retains CLI evidence before $stage cleanup with a $uploadResult final upload', + async ({ stage, uploadResult }) => { + const request = makeRequest(tmpDir); + request.workspace.upstreamBranch = 'main'; + request.workspace.restoredFromBackup = stage === 'restored'; + if (stage === 'restored') await createCompleteGitWorkspace(request.workspace.workspacePath); + const cliLogDir = path.join(request.workspace.sessionHome, '.local/share/kilo/log'); + const cliLogPath = path.join(cliLogDir, 'kilo.log'); + const wrapperLogPath = path.join(tmpDir, 'wrapper.log'); + process.env.WRAPPER_LOG_PATH = wrapperLogPath; + await fsp.mkdir(cliLogDir, { recursive: true }); + await fsp.writeFile(cliLogPath, 'earlier CLI evidence\n'); + await fsp.writeFile(wrapperLogPath, 'wrapper started\n'); + + const archives = new Map(); + let uploadCalls = 0; + let finalUploadSignal: AbortSignal | null | undefined; + let cliPresentDuringFinalUpload = false; + let cliPresentAfterFinalUpload = false; + let beforeCleanupCalls = 0; + globalThis.fetch = asFetch(async (input, init) => { + uploadCalls++; + const url = new URL(input instanceof Request ? input.url : input.toString()); + const archive = gunzipSync(await new Response(init?.body).arrayBuffer()).toString(); + if (uploadCalls === 2) { + cliPresentDuringFinalUpload = fs.existsSync(cliLogPath); + finalUploadSignal = init?.signal; + if (uploadResult === 'failure') { + throw new Error(`Authorization: Bearer ${request.session.workerAuthToken}`); + } + if (uploadResult === 'timeout') return new Promise(() => {}); + await Bun.sleep(20); + } + archives.set(url.pathname, archive); + return new Response(null, { status: 204 }); + }); + + const state = new WrapperState(); + let uploader: LogUploader | undefined; + let bootstrapError: unknown; + const deps: ServerDependencies = { + state, + kiloClient: {} as WrapperKiloClient, + openConnection: async () => {}, + closeConnection: async () => {}, + setAborted: () => {}, + resetLifecycle: () => {}, + readySession: async (readyRequest, archiveId) => { + const attemptUploader = createLogUploader({ + archiveId, + context: { + workerBaseUrl: 'https://worker.example.com', + kiloSessionId: readyRequest.kiloSessionId, + workerAuthToken: readyRequest.session.workerAuthToken, + }, + sessionId: readyRequest.agentSessionId, + userId: readyRequest.userId, + cliLogDir, + wrapperLogPath, + }); + uploader = attemptUploader; + uploaders.push(attemptUploader); + state.setLogUploader(attemptUploader); + attemptUploader.start(); + await attemptUploader.uploadNow(); + try { + await prepareWrapperBootstrapWorkspace(readyRequest, undefined, { + git: async args => { + if (args[0] === 'clone') { + await fsp.mkdir(path.join(readyRequest.workspace.workspacePath, '.git'), { + recursive: true, + }); + } + if (args[0] === 'fetch') { + await fsp.appendFile(cliLogPath, 'final CLI failure evidence\n'); + return { stdout: '', stderr: '', exitCode: 1 }; + } + return { stdout: '', stderr: '', exitCode: 0 }; + }, + beforeFailureCleanup: async () => { + beforeCleanupCalls++; + await attemptUploader.finalize(uploadResult === 'timeout' ? 100 : 5_000); + cliPresentAfterFinalUpload = fs.existsSync(cliLogPath); + }, + }); + } catch (error) { + bootstrapError = error; + return { + status: 'error', + error: { + code: workspaceBootstrapErrorCode(error), + message: 'Workspace preparation failed', + }, + }; + } + throw new Error('Expected bootstrap failure'); + }, + }; + + const startedAt = Date.now(); + const response = await createSessionReadyHandler(deps)( + new Request('http://wrapper.test/session/ready', { + method: 'POST', + body: JSON.stringify(request), + }) + ); + const elapsedMs = Date.now() - startedAt; + if (!uploader) throw new Error('Expected attempt uploader'); + const archivePath = `/sessions/${request.userId}/${request.agentSessionId}/logs/${uploader.archiveId}/logs.tar.gz`; + const retainedArchive = archives.get(archivePath); + await uploader.finalize(); + await uploader.uploadNow(); + + expect(response.status).toBe(503); + expect(workspaceBootstrapErrorCode(bootstrapError)).toBe( + stage === 'restored' ? 'WORKSPACE_RECONCILIATION_FAILED' : 'WORKSPACE_SETUP_FAILED' + ); + expect(retainedArchive?.includes('earlier CLI evidence')).toBe(true); + expect(beforeCleanupCalls).toBe(1); + expect(cliPresentDuringFinalUpload).toBe(true); + expect(cliPresentAfterFinalUpload).toBe(true); + expect(fs.existsSync(request.workspace.workspacePath)).toBe(false); + expect(fs.existsSync(request.workspace.sessionHome)).toBe(false); + expect(state.logUploader).toBeNull(); + expect(uploadCalls).toBe(2); + expect(archives.size).toBe(1); + expect(archives.get(archivePath)).toBe(retainedArchive); + if (uploadResult === 'success') { + expect(retainedArchive).toContain('final CLI failure evidence'); + } else { + expect(retainedArchive).not.toContain('final CLI failure evidence'); + } + if (uploadResult === 'timeout') { + expect(finalUploadSignal?.aborted).toBe(true); + expect(elapsedMs).toBeLessThan(1_500); + } + const wrapperLogs = await fsp.readFile(wrapperLogPath, 'utf8'); + expect(wrapperLogs).not.toContain(request.session.workerAuthToken); + expect(retainedArchive).not.toContain(request.session.workerAuthToken); + expect(retainedArchive).not.toContain(request.materialized.env.KILOCODE_TOKEN); + } + ); + + it('still cleans up and preserves the bootstrap error if the pre-cleanup callback throws', async () => { + const request = makeRequest(tmpDir); + const wrapperLogPath = path.join(tmpDir, 'wrapper.log'); + process.env.WRAPPER_LOG_PATH = wrapperLogPath; + let callbackCalls = 0; + let bootstrapError: unknown; + try { + await prepareWrapperBootstrapWorkspace(request, undefined, { + git: async () => ({ stdout: '', stderr: '', exitCode: 128 }), + beforeFailureCleanup: () => { + callbackCalls++; + throw new Error(`Authorization: Bearer ${request.session.workerAuthToken}`); + }, + }); + } catch (error) { + bootstrapError = error; + } + + expect(callbackCalls).toBe(1); + expect(bootstrapError).toMatchObject({ + code: 'WORKSPACE_SETUP_FAILED', + message: 'Repository clone failed', + }); + expect(fs.existsSync(request.workspace.workspacePath)).toBe(false); + expect(fs.existsSync(request.workspace.sessionHome)).toBe(false); + expect(await fsp.readFile(wrapperLogPath, 'utf8')).not.toContain( + request.session.workerAuthToken + ); + }); + it('aborts active work and cleans up when the shared workspace deadline expires', async () => { const request = makeRequest(tmpDir); request.materialized.setupCommands = []; diff --git a/services/cloud-agent-next/wrapper/src/session-bootstrap.ts b/services/cloud-agent-next/wrapper/src/session-bootstrap.ts index 711af46689..65f0627477 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.ts @@ -166,6 +166,7 @@ export type WrapperBootstrapDeps = { git?: GitRunner; runProcess?: ProcessRunner; restoreSession?: typeof restoreSession; + beforeFailureCleanup?: () => Promise; workspacePreparationTimeoutMs?: number; }; @@ -427,7 +428,15 @@ async function removePath(filePath: string, signal?: AbortSignal): Promise } } -async function cleanupWorkspace(request: WrapperSessionReadyRequest): Promise { +async function cleanupWorkspace( + request: WrapperSessionReadyRequest, + beforeFailureCleanup: WrapperBootstrapDeps['beforeFailureCleanup'] +): Promise { + try { + await beforeFailureCleanup?.(); + } catch { + logToFile('Failed to finalize logs before workspace cleanup'); + } await Promise.allSettled([ removePath(request.workspace.workspacePath), removePath(request.workspace.sessionHome), @@ -1310,7 +1319,7 @@ async function prepareWrapperBootstrapWorkspaceWithinDeadline( } catch (error) { if (error instanceof RestoredWorkspaceReconciliationError) { if (workspaceNeedsBootstrap) { - await cleanupWorkspace(request); + await cleanupWorkspace(request, deps.beforeFailureCleanup); } throw error; } @@ -1322,7 +1331,7 @@ async function prepareWrapperBootstrapWorkspaceWithinDeadline( `bootstrap workspace failed kiloSessionId=${request.kiloSessionId} workspaceWasWarm=${workspaceWasWarm} workspaceNeedsBootstrap=${workspaceNeedsBootstrap} willCleanup=${workspaceNeedsBootstrap} code=${bootstrapError.code} subtype=${bootstrapError.subtype ?? '(none)'}` ); if (workspaceNeedsBootstrap) { - await cleanupWorkspace(request); + await cleanupWorkspace(request, deps.beforeFailureCleanup); logToFile(`bootstrap workspace cleanup finished kiloSessionId=${request.kiloSessionId}`); } throw bootstrapError; diff --git a/services/cloud-agent-next/wrapper/src/shutdown.test.ts b/services/cloud-agent-next/wrapper/src/shutdown.test.ts index da4fc06c61..5fcefe06b1 100644 --- a/services/cloud-agent-next/wrapper/src/shutdown.test.ts +++ b/services/cloud-agent-next/wrapper/src/shutdown.test.ts @@ -69,7 +69,10 @@ describe('abortKiloSessionForShutdown', () => { }); let uploaderStopped = false; state.setLogUploader({ + archiveId: 'run_1--test', start: () => {}, + updateContext: () => {}, + finalize: async () => {}, uploadNow: async () => { signalUploadStarted?.(); await new Promise(resolve => { diff --git a/services/cloud-agent-next/wrapper/src/state.ts b/services/cloud-agent-next/wrapper/src/state.ts index 79a64f6001..f3bafc8c1e 100644 --- a/services/cloud-agent-next/wrapper/src/state.ts +++ b/services/cloud-agent-next/wrapper/src/state.ts @@ -137,8 +137,9 @@ export class WrapperState { } setLogUploader(uploader: LogUploader | null): void { - this._logUploader?.stop(); + const previousUploader = this._logUploader; this._logUploader = uploader; + if (previousUploader) void previousUploader.finalize().catch(() => {}); } updateActivity(): void { @@ -210,6 +211,7 @@ export class WrapperState { return { changed: true }; } const changed = + this.session.kiloSessionId !== context.kiloSessionId || this.session.ingestUrl !== context.ingestUrl || this.session.ingestToken !== context.ingestToken || this.session.workerAuthToken !== context.workerAuthToken || diff --git a/services/git-token-service/src/github-token-service.test.ts b/services/git-token-service/src/github-token-service.test.ts index 1198d05951..51100105c6 100644 --- a/services/git-token-service/src/github-token-service.test.ts +++ b/services/git-token-service/src/github-token-service.test.ts @@ -1,7 +1,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import { createAppAuth } from '@octokit/auth-app'; import { Octokit } from '@octokit/rest'; -import { GitHubTokenService } from './github-token-service.js'; +import { GitHubTokenGenerationError, GitHubTokenService } from './github-token-service.js'; vi.mock('@octokit/auth-app', () => ({ createAppAuth: vi.fn(), @@ -123,29 +123,101 @@ describe('GitHubTokenService', () => { expect(JSON.stringify(consoleWarn.mock.calls)).not.toContain('sensitive-metadata-response'); }); - it('does not log authenticated upstream response data when scoped token minting fails', async () => { - const upstreamError = Object.assign(new Error('repository unavailable'), { - response: { data: { token: 'sensitive-upstream-data' } }, - }); - vi.mocked(createAppAuth).mockReturnValue( - vi.fn().mockRejectedValue(upstreamError) as unknown as ReturnType - ); - const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}); - const service = new GitHubTokenService({ - GITHUB_APP_ID: 'app-id', - GITHUB_APP_PRIVATE_KEY: 'private-key', - } as CloudflareEnv); - - await expect(service.getTokenForRepo('123', 'repository')).rejects.toThrow( - 'Failed to generate GitHub installation token: repository unavailable' - ); - - expect(consoleError).toHaveBeenCalledWith( - JSON.stringify({ + it.each([ + [404, 'Not Found', 'no_installation_found'], + [403, 'Resource not accessible by integration', 'repository_not_installed'], + [ + 422, + 'There is at least one repository that does not exist or is not accessible to the parent installation.', + 'repository_not_installed', + ], + ] as const)( + 'preserves the expected GitHub %s token failure as %s without exposing credentials', + async (status, message, reason) => { + const secret = 'sensitive-upstream-credential'; + const upstreamError = Object.assign(new Error(`provider detail: ${secret}`), { + status, + request: { headers: { authorization: `Bearer ${secret}` } }, + response: { data: { message, token: secret } }, + }); + const auth = vi.fn().mockRejectedValue(upstreamError); + vi.mocked(createAppAuth).mockReturnValue(Object.assign(auth, { hook: vi.fn() })); + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}); + const tokenCache = createTokenCache(); + const service = new GitHubTokenService({ + GITHUB_APP_ID: 'app-id', + GITHUB_APP_PRIVATE_KEY: 'private-key', + TOKEN_CACHE: tokenCache, + } as unknown as CloudflareEnv); + + const error = await service + .getTokenForRepo('123', 'repository') + .catch((error: unknown) => error); + + expect(error).toBeInstanceOf(GitHubTokenGenerationError); + expect(error).toMatchObject({ + status, + reason, message: 'Failed to generate GitHub installation token', - errorType: 'Error', - }) - ); - expect(JSON.stringify(consoleError.mock.calls)).not.toContain('sensitive-upstream-data'); + }); + expect(error).not.toHaveProperty('cause'); + expect(JSON.stringify({ error, logs: consoleError.mock.calls })).not.toContain(secret); + expect(consoleError).toHaveBeenCalledWith( + JSON.stringify({ message: 'Failed to generate GitHub installation token', status, reason }) + ); + expect(auth).toHaveBeenCalledOnce(); + expect(tokenCache.put).not.toHaveBeenCalled(); + } + ); + + it.each([ + [403, 'API rate limit exceeded'], + [403, 'You have exceeded a secondary rate limit.'], + [403, 'Forbidden'], + [429, 'Too Many Requests'], + [401, 'Bad credentials'], + [422, 'The permissions requested are not granted to this installation.'], + [500, 'Resource not accessible by integration'], + [undefined, 'Not Found'], + [undefined, 'Resource not accessible by integration'], + ])( + 'does not infer repository access from unrecognized %s token errors', + async (status, message) => { + const upstreamError = Object.assign(new Error('secret-provider-detail'), { + status, + response: { data: { message, token: 'secret-token' } }, + }); + vi.mocked(createAppAuth).mockReturnValue( + Object.assign(vi.fn().mockRejectedValue(upstreamError), { hook: vi.fn() }) + ); + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}); + const service = new GitHubTokenService({ + GITHUB_APP_ID: 'app-id', + GITHUB_APP_PRIVATE_KEY: 'private-key', + } as CloudflareEnv); + + const error = await service + .getTokenForRepo('123', 'repository') + .catch((error: unknown) => error); + + expect(error).toBeInstanceOf(GitHubTokenGenerationError); + expect(error).toMatchObject({ + status, + reason: undefined, + message: 'Failed to generate GitHub installation token', + }); + expect(error).not.toHaveProperty('cause'); + expect(JSON.stringify({ error, logs: consoleError.mock.calls })).not.toContain('secret-'); + } + ); + + it('does not classify token-cache failures as GitHub access failures', async () => { + const cacheError = Object.assign(new Error('cache unavailable'), { status: 404 }); + const tokenCache = createTokenCache(); + tokenCache.get.mockRejectedValueOnce(cacheError); + const service = new GitHubTokenService({ TOKEN_CACHE: tokenCache } as unknown as CloudflareEnv); + + await expect(service.getTokenForRepo('123', 'repository')).rejects.toBe(cacheError); + expect(createAppAuth).not.toHaveBeenCalled(); }); }); diff --git a/services/git-token-service/src/github-token-service.ts b/services/git-token-service/src/github-token-service.ts index 1d5d49837b..3e5b2974eb 100644 --- a/services/git-token-service/src/github-token-service.ts +++ b/services/git-token-service/src/github-token-service.ts @@ -33,6 +33,25 @@ const GitHubInstallationAccountSchema = z.object({ }), }); +const GitHubTokenErrorSchema = z.object({ + status: z.number().int().min(400).max(599), + response: z + .object({ + data: z.object({ message: z.string() }), + }) + .optional(), +}); + +export class GitHubTokenGenerationError extends Error { + constructor( + readonly status: number | undefined, + readonly reason?: 'no_installation_found' | 'repository_not_installed' + ) { + super('Failed to generate GitHub installation token'); + this.name = 'GitHubTokenGenerationError'; + } +} + export class GitHubTokenService { constructor(private env: CloudflareEnv) {} @@ -215,14 +234,26 @@ export class GitHubTokenService { expiresAt: new Date(result.expiresAt).getTime(), }; } catch (error) { + const parsed = GitHubTokenErrorSchema.safeParse(error); + const status = parsed.success ? parsed.data.status : undefined; + const message = parsed.success ? parsed.data.response?.data.message : undefined; + const reason = + status === 404 + ? 'no_installation_found' + : (status === 403 && message === 'Resource not accessible by integration') || + (status === 422 && + message === + 'There is at least one repository that does not exist or is not accessible to the parent installation.') + ? 'repository_not_installed' + : undefined; console.error( JSON.stringify({ message: 'Failed to generate GitHub installation token', - errorType: error instanceof Error ? error.name : 'UnknownError', + status, + reason, }) ); - const message = error instanceof Error ? error.message : 'Unknown error'; - throw new Error(`Failed to generate GitHub installation token: ${message}`); + throw new GitHubTokenGenerationError(status, reason); } } diff --git a/services/git-token-service/src/index.test.ts b/services/git-token-service/src/index.test.ts index b1735c6993..d4463c462d 100644 --- a/services/git-token-service/src/index.test.ts +++ b/services/git-token-service/src/index.test.ts @@ -2,6 +2,7 @@ import { signKiloToken } from '@kilocode/worker-utils'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import type * as GitLabCredentialBrokerHandlerModule from './gitlab-credential-broker-handler.js'; import type * as GitLabLookupServiceModule from './gitlab-lookup-service.js'; +import type * as GitHubTokenServiceModule from './github-token-service.js'; const serviceMocks = vi.hoisted(() => ({ findInstallationId: vi.fn(), @@ -32,7 +33,8 @@ vi.mock('cloudflare:workers', () => ({ }, })); -vi.mock('./github-token-service.js', () => ({ +vi.mock('./github-token-service.js', async importOriginal => ({ + ...(await importOriginal()), GitHubTokenService: class GitHubTokenService { getToken = serviceMocks.getToken; getTokenForRepo = serviceMocks.getTokenForRepo; @@ -98,6 +100,7 @@ vi.mock('./bitbucket-runtime-token-resolver.js', () => ({ })); import gitTokenServiceWorker, { GitTokenRPCEntrypoint } from './index.js'; +import { GitHubTokenGenerationError } from './github-token-service.js'; beforeEach(() => { serviceMocks.hasGitLabProjectCredentialCandidates.mockReset().mockResolvedValue(false); @@ -750,6 +753,90 @@ describe('GitTokenRPCEntrypoint.getTokenForRepo', () => { }); }); +describe.each([ + 'getTokenForRepo', + 'getCloudAgentAuthForRepo', + 'issueGitHubSessionCapability', +] as const)('GitHub credential failures through %s', method => { + const params = { githubRepo: 'acme/repo', userId: 'user_1' }; + + beforeEach(() => { + vi.clearAllMocks(); + const installation = { + success: true, + installationId: '123', + accountLogin: 'acme', + githubAppType: 'standard', + repoName: 'repo', + permissions: { contents: 'write', pull_requests: 'write' }, + }; + serviceMocks.findInstallationId.mockReset().mockResolvedValue(installation); + serviceMocks.findManagedInstallationForRepo.mockReset().mockResolvedValue(installation); + serviceMocks.findRefreshCandidates + .mockReset() + .mockResolvedValue({ success: true, candidates: [] }); + serviceMocks.getTokenForRepo.mockReset().mockResolvedValue('installation-token'); + }); + + it.each([ + [404, 'no_installation_found'], + [403, 'repository_not_installed'], + [422, 'repository_not_installed'], + ] as const)( + 'returns the structured %s token failure without a token or capability', + async (status, reason) => { + serviceMocks.getTokenForRepo.mockRejectedValueOnce( + new GitHubTokenGenerationError(status, reason) + ); + + await expect(createService()[method](params)).resolves.toEqual({ success: false, reason }); + expect(serviceMocks.getTokenForRepo).toHaveBeenCalledOnce(); + expect(serviceMocks.getToken).not.toHaveBeenCalled(); + expect(serviceMocks.selectUserAuthorization).not.toHaveBeenCalled(); + } + ); + + it.each([ + 'no_installation_found', + 'invalid_repo_format', + 'invalid_org_id', + 'integration_mismatch', + ] as const)('preserves the %s lookup gate before token issuance', async reason => { + serviceMocks.findInstallationId.mockResolvedValue({ success: false, reason }); + serviceMocks.findManagedInstallationForRepo.mockResolvedValue({ success: false, reason }); + + await expect(createService()[method](params)).resolves.toEqual({ success: false, reason }); + expect(serviceMocks.getTokenForRepo).not.toHaveBeenCalled(); + expect(serviceMocks.getToken).not.toHaveBeenCalled(); + expect(serviceMocks.selectUserAuthorization).not.toHaveBeenCalled(); + }); + + it('does not turn a database lookup exception into a missing installation', async () => { + const databaseError = new Error('query failed with private SQL parameters', { + cause: new Error('database connection closed'), + }); + serviceMocks.findInstallationId.mockRejectedValueOnce(databaseError); + serviceMocks.findManagedInstallationForRepo.mockRejectedValueOnce(databaseError); + + await expect(createService()[method](params)).rejects.toBe(databaseError); + expect(serviceMocks.getTokenForRepo).not.toHaveBeenCalled(); + expect(serviceMocks.getToken).not.toHaveBeenCalled(); + expect(serviceMocks.selectUserAuthorization).not.toHaveBeenCalled(); + }); + + it.each([403, 429, 500, undefined])( + 'leaves unclassified %s token failures as exceptions', + async status => { + const error = new GitHubTokenGenerationError(status); + serviceMocks.getTokenForRepo.mockRejectedValueOnce(error); + + await expect(createService()[method](params)).rejects.toBe(error); + expect(serviceMocks.getTokenForRepo).toHaveBeenCalledOnce(); + expect(serviceMocks.getToken).not.toHaveBeenCalled(); + } + ); +}); + const outboundContainerId = 'outbound-container-1'; describe('GitTokenRPCEntrypoint GitHub session capability RPCs', () => { diff --git a/services/git-token-service/src/index.ts b/services/git-token-service/src/index.ts index aff0480b1b..952b7d6425 100644 --- a/services/git-token-service/src/index.ts +++ b/services/git-token-service/src/index.ts @@ -13,7 +13,11 @@ import { } from '@kilocode/worker-utils/internal-service-token-audiences'; import { WorkerEntrypoint } from 'cloudflare:workers'; import { z } from 'zod'; -import { GitHubTokenService, type GitHubAppType } from './github-token-service.js'; +import { + GitHubTokenGenerationError, + GitHubTokenService, + type GitHubAppType, +} from './github-token-service.js'; import { GitLabLookupService, type GitLabLookupSuccess } from './gitlab-lookup-service.js'; import { resolveGitLabRuntimeToken, @@ -825,19 +829,26 @@ export class GitTokenRPCEntrypoint extends WorkerEntrypoint { return { success: false, reason: 'invalid_repo_format' }; } - const token = await this.githubService.getTokenForRepo( - installation.installationId, - repoName, - installation.githubAppType - ); + try { + const token = await this.githubService.getTokenForRepo( + installation.installationId, + repoName, + installation.githubAppType + ); - return { - success: true, - token, - installationId: installation.installationId, - accountLogin: installation.accountLogin, - appType: installation.githubAppType, - }; + return { + success: true, + token, + installationId: installation.installationId, + accountLogin: installation.accountLogin, + appType: installation.githubAppType, + }; + } catch (error) { + if (error instanceof GitHubTokenGenerationError && error.reason !== undefined) { + return { success: false, reason: error.reason }; + } + throw error; + } } async getCloudAgentAuthForRepo( @@ -861,20 +872,29 @@ export class GitTokenRPCEntrypoint extends WorkerEntrypoint { const installationAuthor = this.getInstallationAuthor(installation.githubAppType); const installationAuth = async ( fallbackReason?: ManagedGitHubFallbackReason - ): Promise => ({ - success: true, - githubToken: await this.githubService.getTokenForRepo( - installation.installationId, - installation.repoName, - installation.githubAppType - ), - installationId: installation.installationId, - accountLogin: installation.accountLogin, - appType: installation.githubAppType, - source: 'installation', - gitAuthor: installationAuthor, - ...(fallbackReason !== undefined ? { fallbackReason } : {}), - }); + ): Promise => { + try { + return { + success: true, + githubToken: await this.githubService.getTokenForRepo( + installation.installationId, + installation.repoName, + installation.githubAppType + ), + installationId: installation.installationId, + accountLogin: installation.accountLogin, + appType: installation.githubAppType, + source: 'installation', + gitAuthor: installationAuthor, + ...(fallbackReason !== undefined ? { fallbackReason } : {}), + }; + } catch (error) { + if (error instanceof GitHubTokenGenerationError && error.reason !== undefined) { + return { success: false, reason: error.reason }; + } + throw error; + } + }; if (params.allowUserAuthorization !== true) return installationAuth(); if (installation.githubAppType === 'lite') return installationAuth('lite_installation'); From 6ea5d13d45b684b64772938d330b4a87491131fb Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Mon, 31 Aug 2026 16:27:12 +0200 Subject: [PATCH 2/3] fix(cloud-agent): archive bootstrap diagnostics before cleanup --- .../wrapper/src/session-bootstrap.test.ts | 48 ++++++++++++++++++- .../wrapper/src/session-bootstrap.ts | 2 +- 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts index cb39152714..8242984ea6 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts @@ -1738,10 +1738,40 @@ describe('prepareWrapperBootstrapWorkspace', () => { }); }); - it('exposes redacted setup command and stderr on failure but redacts secrets', async () => { + it('archives sanitized setup failure details before workspace cleanup', async () => { const request = makeRequest(tmpDir); request.materialized.setupCommands = ['private-tool --token argv-secret']; + const cliLogDir = path.join(request.workspace.sessionHome, '.local/share/kilo/log'); + const cliLogPath = path.join(cliLogDir, 'kilo.log'); + const wrapperLogPath = path.join(tmpDir, 'wrapper.log'); + process.env.WRAPPER_LOG_PATH = wrapperLogPath; + await fsp.mkdir(cliLogDir, { recursive: true }); + await fsp.writeFile(cliLogPath, 'earlier CLI evidence\n'); + + let archivedLogs = ''; + let uploadCalls = 0; + let cliPresentDuringUpload = false; + globalThis.fetch = asFetch(async (_input, init) => { + uploadCalls++; + cliPresentDuringUpload = fs.existsSync(cliLogPath); + archivedLogs = gunzipSync(await new Response(init?.body).arrayBuffer()).toString(); + return new Response(null, { status: 204 }); + }); + const uploader = createLogUploader({ + archiveId: 'wr_test--setup-failure', + context: { + workerBaseUrl: 'https://worker.example.com', + kiloSessionId: request.kiloSessionId, + workerAuthToken: request.session.workerAuthToken, + }, + sessionId: request.agentSessionId, + userId: request.userId, + cliLogDir, + wrapperLogPath, + }); + uploaders.push(uploader); const deps: WrapperBootstrapDeps = { + beforeFailureCleanup: () => uploader.finalize(), git: async args => { if (args[0] === 'clone') { await fsp.mkdir(path.join(request.workspace.workspacePath, '.git'), { recursive: true }); @@ -1754,6 +1784,7 @@ describe('prepareWrapperBootstrapWorkspace', () => { runProcess: async () => ({ stdout: 'private-file-content', stderr: [ + '\u001b[31mDependency resolution failed\u001b[0m', 'bare-unlabeled-token', 'https://user:url-secret@example.com/repo.git', 'Authorization: Bearer bearer-secret', @@ -1800,6 +1831,18 @@ describe('prepareWrapperBootstrapWorkspace', () => { expect(detail).toContain('SECRET_VALUE=[REDACTED]'); expect(detail).toContain('private-file-content'); expect(detail).toContain('bare-unlabeled-token'); + await uploader.finalize(); + await uploader.uploadNow(); + expect(uploadCalls).toBe(1); + expect(cliPresentDuringUpload).toBe(true); + expect(fs.existsSync(request.workspace.workspacePath)).toBe(false); + expect(fs.existsSync(request.workspace.sessionHome)).toBe(false); + expect(archivedLogs).toContain('earlier CLI evidence'); + expect(archivedLogs).toContain('subtype=setup_command_failed'); + expect(archivedLogs).toContain(`error=${setupError.message}`); + expect(archivedLogs).toContain(`detail=${detail}`); + expect(archivedLogs).toContain('Dependency resolution failed'); + expect(archivedLogs).not.toContain('\u001b'); const projectedError = JSON.stringify(setupError); for (const sensitiveValue of [ 'argv-secret', @@ -1807,8 +1850,11 @@ describe('prepareWrapperBootstrapWorkspace', () => { 'bearer-secret', 'cookie-secret', 'env-secret', + request.session.workerAuthToken, + request.materialized.env.KILOCODE_TOKEN, ]) { expect(projectedError).not.toContain(sensitiveValue); + expect(archivedLogs).not.toContain(sensitiveValue); } }); diff --git a/services/cloud-agent-next/wrapper/src/session-bootstrap.ts b/services/cloud-agent-next/wrapper/src/session-bootstrap.ts index 65f0627477..b19a1e4f85 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.ts @@ -1328,7 +1328,7 @@ async function prepareWrapperBootstrapWorkspaceWithinDeadline( ? error : workspaceBootstrapError('workspace_setup_unknown', 'Workspace setup failed'); logToFile( - `bootstrap workspace failed kiloSessionId=${request.kiloSessionId} workspaceWasWarm=${workspaceWasWarm} workspaceNeedsBootstrap=${workspaceNeedsBootstrap} willCleanup=${workspaceNeedsBootstrap} code=${bootstrapError.code} subtype=${bootstrapError.subtype ?? '(none)'}` + `bootstrap workspace failed kiloSessionId=${request.kiloSessionId} workspaceWasWarm=${workspaceWasWarm} workspaceNeedsBootstrap=${workspaceNeedsBootstrap} willCleanup=${workspaceNeedsBootstrap} code=${bootstrapError.code} subtype=${bootstrapError.subtype ?? '(none)'} error=${bootstrapError.message}${bootstrapError.detail ? ` detail=${bootstrapError.detail}` : ''}` ); if (workspaceNeedsBootstrap) { await cleanupWorkspace(request, deps.beforeFailureCleanup); From c98393ec0ac05f5ce865812549646bbf7632f140 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Tue, 1 Sep 2026 09:27:06 +0200 Subject: [PATCH 3/3] fix(cloud-agent): preserve terminal bootstrap archive diagnostics --- .../wrapper/src/session-bootstrap.test.ts | 118 +++++++++++++++++- .../wrapper/src/session-bootstrap.ts | 30 ++--- 2 files changed, 132 insertions(+), 16 deletions(-) diff --git a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts index 8242984ea6..11801d4119 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts @@ -1175,6 +1175,12 @@ describe('prepareWrapperBootstrapWorkspace', () => { expect(archives.get(archivePath)).toBe(retainedArchive); if (uploadResult === 'success') { expect(retainedArchive).toContain('final CLI failure evidence'); + expect(retainedArchive).toContain(`code=${workspaceBootstrapErrorCode(bootstrapError)}`); + expect(retainedArchive).toContain( + stage === 'restored' + ? 'error=Failed to fetch authoritative remote state' + : 'error=Repository checkout failed' + ); } else { expect(retainedArchive).not.toContain('final CLI failure evidence'); } @@ -1189,6 +1195,77 @@ describe('prepareWrapperBootstrapWorkspace', () => { } ); + it('redacts reconciliation diagnostics before cleanup while preserving the error type', async () => { + const request = makeRequest(tmpDir); + request.workspace.upstreamBranch = 'main'; + request.workspace.restoredFromBackup = true; + await createCompleteGitWorkspace(request.workspace.workspacePath); + const wrapperLogPath = path.join(tmpDir, 'wrapper.log'); + process.env.WRAPPER_LOG_PATH = wrapperLogPath; + const fetchError = new Error( + '\u001b[31mFetch failed\u001b[0m Authorization: Bearer fetch-secret' + ); + let logsBeforeCleanup = ''; + let bootstrapError: unknown; + try { + await prepareWrapperBootstrapWorkspace(request, undefined, { + git: async args => { + if (args[0] === 'fetch') throw fetchError; + return { stdout: '', stderr: '', exitCode: 0 }; + }, + beforeFailureCleanup: async () => { + logsBeforeCleanup = await fsp.readFile(wrapperLogPath, 'utf8'); + }, + }); + } catch (error) { + bootstrapError = error; + } + + expect(bootstrapError).toBeInstanceOf(RestoredWorkspaceReconciliationError); + expect(bootstrapError).toMatchObject({ cause: fetchError }); + expect(logsBeforeCleanup).toContain('code=WORKSPACE_RECONCILIATION_FAILED'); + expect(logsBeforeCleanup).toContain('error=Fetch failed Authorization: Bearer [REDACTED]'); + expect(logsBeforeCleanup).not.toContain('fetch-secret'); + expect(logsBeforeCleanup).not.toContain('\u001b'); + expect(fs.existsSync(request.workspace.workspacePath)).toBe(false); + expect(fs.existsSync(request.workspace.sessionHome)).toBe(false); + }); + + it('preserves the recorded failure when finalization crosses the workspace deadline', async () => { + const request = makeRequest(tmpDir); + const wrapperLogPath = path.join(tmpDir, 'wrapper.log'); + process.env.WRAPPER_LOG_PATH = wrapperLogPath; + let commandSignal: AbortSignal | undefined; + let logsBeforeCleanup = ''; + let bootstrapError: unknown; + try { + await prepareWrapperBootstrapWorkspace(request, undefined, { + workspacePreparationTimeoutMs: 100, + git: async (_args, options) => { + commandSignal = options?.signal; + return { stdout: '', stderr: '', exitCode: 1 }; + }, + beforeFailureCleanup: async () => { + logsBeforeCleanup = await fsp.readFile(wrapperLogPath, 'utf8'); + if (commandSignal && !commandSignal.aborted) { + await new Promise(resolve => + commandSignal?.addEventListener('abort', () => resolve(), { once: true }) + ); + } + }, + }); + } catch (error) { + bootstrapError = error; + } + + expect(commandSignal?.aborted).toBe(true); + expect(bootstrapError).toMatchObject({ message: 'Repository clone failed' }); + expect(logsBeforeCleanup).toContain('error=Repository clone failed'); + expect(logsBeforeCleanup).not.toContain('Workspace preparation timed out'); + expect(fs.existsSync(request.workspace.workspacePath)).toBe(false); + expect(fs.existsSync(request.workspace.sessionHome)).toBe(false); + }); + it('still cleans up and preserves the bootstrap error if the pre-cleanup callback throws', async () => { const request = makeRequest(tmpDir); const wrapperLogPath = path.join(tmpDir, 'wrapper.log'); @@ -1219,15 +1296,45 @@ describe('prepareWrapperBootstrapWorkspace', () => { ); }); - it('aborts active work and cleans up when the shared workspace deadline expires', async () => { + it('archives the shared workspace deadline failure before cleaning up aborted work', async () => { const request = makeRequest(tmpDir); request.materialized.setupCommands = []; + const cliLogDir = path.join(request.workspace.sessionHome, '.local/share/kilo/log'); + const cliLogPath = path.join(cliLogDir, 'kilo.log'); + const wrapperLogPath = path.join(tmpDir, 'wrapper.log'); + process.env.WRAPPER_LOG_PATH = wrapperLogPath; + await fsp.mkdir(cliLogDir, { recursive: true }); + await fsp.writeFile(cliLogPath, 'earlier CLI evidence\n'); + + let archivedLogs = ''; + let uploadCalls = 0; + let cliPresentDuringUpload = false; + globalThis.fetch = asFetch(async (_input, init) => { + uploadCalls++; + cliPresentDuringUpload = fs.existsSync(cliLogPath); + archivedLogs = gunzipSync(await new Response(init?.body).arrayBuffer()).toString(); + return new Response(null, { status: 204 }); + }); + const uploader = createLogUploader({ + archiveId: 'wr_test--workspace-timeout', + context: { + workerBaseUrl: 'https://worker.example.com', + kiloSessionId: request.kiloSessionId, + workerAuthToken: request.session.workerAuthToken, + }, + sessionId: request.agentSessionId, + userId: request.userId, + cliLogDir, + wrapperLogPath, + }); + uploaders.push(uploader); let commandSignal: AbortSignal | undefined; let caughtError: unknown; try { await prepareWrapperBootstrapWorkspace(request, undefined, { workspacePreparationTimeoutMs: 100, + beforeFailureCleanup: () => uploader.finalize(), git: async (args, opts) => { if (args[0] !== 'clone') { return { stdout: '', stderr: '', exitCode: 0 }; @@ -1262,6 +1369,15 @@ describe('prepareWrapperBootstrapWorkspace', () => { retryable: true, message: expect.stringContaining('Workspace preparation timed out'), }); + await uploader.finalize(); + await uploader.uploadNow(); + expect(uploadCalls).toBe(1); + expect(cliPresentDuringUpload).toBe(true); + expect(archivedLogs).toContain('earlier CLI evidence'); + expect(archivedLogs).toContain('code=WORKSPACE_SETUP_FAILED subtype=workspace_setup_unknown'); + expect(archivedLogs).toContain('error=Workspace preparation timed out after 0.1s'); + expect(archivedLogs).not.toContain('error=Repository clone failed'); + expect(archivedLogs).not.toContain(request.session.workerAuthToken); expect(fs.existsSync(request.workspace.workspacePath)).toBe(false); expect(fs.existsSync(request.workspace.sessionHome)).toBe(false); }); diff --git a/services/cloud-agent-next/wrapper/src/session-bootstrap.ts b/services/cloud-agent-next/wrapper/src/session-bootstrap.ts index b19a1e4f85..5df96c55e4 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.ts @@ -1198,7 +1198,8 @@ async function prepareWrapperBootstrapWorkspaceWithinDeadline( request: WrapperSessionReadyRequest, progress: BootstrapProgress | undefined, deps: WrapperBootstrapDeps, - signal: AbortSignal + signal: AbortSignal, + timeoutError: WrapperBootstrapError ): Promise { const runGit = deps.git ?? git; const run = deps.runProcess ?? runProcess; @@ -1317,16 +1318,17 @@ async function prepareWrapperBootstrapWorkspaceWithinDeadline( ...(restoreTelemetry ? { restore: restoreTelemetry } : {}), }; } catch (error) { - if (error instanceof RestoredWorkspaceReconciliationError) { - if (workspaceNeedsBootstrap) { - await cleanupWorkspace(request, deps.beforeFailureCleanup); - } - throw error; - } + const failure = signal.reason === timeoutError ? timeoutError : error; const bootstrapError = - error instanceof WrapperBootstrapError - ? error - : workspaceBootstrapError('workspace_setup_unknown', 'Workspace setup failed'); + failure instanceof WrapperBootstrapError + ? failure + : failure instanceof RestoredWorkspaceReconciliationError + ? new WrapperBootstrapError({ + code: 'WORKSPACE_RECONCILIATION_FAILED', + message: redactSecrets(cleanTerminalOutput(failure.message)), + retryable: true, + }) + : workspaceBootstrapError('workspace_setup_unknown', 'Workspace setup failed'); logToFile( `bootstrap workspace failed kiloSessionId=${request.kiloSessionId} workspaceWasWarm=${workspaceWasWarm} workspaceNeedsBootstrap=${workspaceNeedsBootstrap} willCleanup=${workspaceNeedsBootstrap} code=${bootstrapError.code} subtype=${bootstrapError.subtype ?? '(none)'} error=${bootstrapError.message}${bootstrapError.detail ? ` detail=${bootstrapError.detail}` : ''}` ); @@ -1334,7 +1336,7 @@ async function prepareWrapperBootstrapWorkspaceWithinDeadline( await cleanupWorkspace(request, deps.beforeFailureCleanup); logToFile(`bootstrap workspace cleanup finished kiloSessionId=${request.kiloSessionId}`); } - throw bootstrapError; + throw failure instanceof RestoredWorkspaceReconciliationError ? failure : bootstrapError; } } @@ -1386,11 +1388,9 @@ export async function prepareWrapperBootstrapWorkspace( : workspaceSignal, }), }, - workspaceSignal + workspaceSignal, + timeoutError ); - } catch (error) { - if (workspaceSignal.reason === timeoutError) throw timeoutError; - throw error; } finally { clearTimeout(timeout); }