From 673e72c4df42f7c8402b2ef8fa0a33276efac539 Mon Sep 17 00:00:00 2001 From: Peter Kirkham Date: Tue, 5 May 2026 10:42:30 +0100 Subject: [PATCH] feat(code): include check for terminated for the instrumentation agent --- .../features/setup/components/SetupView.tsx | 3 +- .../features/setup/hooks/useSetupRun.ts | 2 + .../setup/services/setupRunService.ts | 98 +++++++++++++------ .../features/setup/stores/setupStore.ts | 8 ++ apps/code/src/shared/types/analytics.ts | 3 +- 5 files changed, 82 insertions(+), 32 deletions(-) diff --git a/apps/code/src/renderer/features/setup/components/SetupView.tsx b/apps/code/src/renderer/features/setup/components/SetupView.tsx index f302f3754..0aae187ef 100644 --- a/apps/code/src/renderer/features/setup/components/SetupView.tsx +++ b/apps/code/src/renderer/features/setup/components/SetupView.tsx @@ -22,6 +22,7 @@ export function SetupView() { wizardFeed, isDiscoveryDone, isWizardStarted, + isWizardDone, wizardSkipped, discoveredTasks, error, @@ -132,7 +133,7 @@ export function SetupView() { color="blue" currentTool={wizardFeed.currentTool} recentEntries={wizardFeed.recentEntries} - isDone={false} + isDone={isWizardDone} doneLabel="Integration ready" /> diff --git a/apps/code/src/renderer/features/setup/hooks/useSetupRun.ts b/apps/code/src/renderer/features/setup/hooks/useSetupRun.ts index ef229907b..b18638e15 100644 --- a/apps/code/src/renderer/features/setup/hooks/useSetupRun.ts +++ b/apps/code/src/renderer/features/setup/hooks/useSetupRun.ts @@ -11,6 +11,7 @@ export function useSetupRun() { const discoveredTasks = useSetupStore((s) => s.discoveredTasks); const wizardTaskId = useSetupStore((s) => s.wizardTaskId); const wizardSkipped = useSetupStore((s) => s.wizardSkipped); + const wizardCompleted = useSetupStore((s) => s.wizardCompleted); const discoveryFeed = useSetupStore((s) => s.discoveryFeed); const wizardFeed = useSetupStore((s) => s.wizardFeed); const error = useSetupStore((s) => s.error); @@ -34,6 +35,7 @@ export function useSetupRun() { wizardFeed, isDiscoveryDone: discoveryStatus === "done", isWizardStarted: !!wizardTaskId, + isWizardDone: wizardCompleted, wizardSkipped, discoveredTasks, wizardTaskId, diff --git a/apps/code/src/renderer/features/setup/services/setupRunService.ts b/apps/code/src/renderer/features/setup/services/setupRunService.ts index 7cd34853d..1ae496e36 100644 --- a/apps/code/src/renderer/features/setup/services/setupRunService.ts +++ b/apps/code/src/renderer/features/setup/services/setupRunService.ts @@ -346,10 +346,11 @@ export class SetupRunService { } private subscribeToWizardEvents(taskId: string, signal: AbortSignal): void { - const checkForRun = async () => { + const run = async () => { const client = await getAuthenticatedClient(); if (!client || signal.aborted) return; + let runId: string | null = null; for (let i = 0; i < 30; i++) { try { await sleep(2000, signal); @@ -359,41 +360,78 @@ export class SetupRunService { try { const taskData = (await client.getTask(taskId)) as unknown as Task; if (signal.aborted) return; - const runId = taskData.latest_run?.id; - if (runId) { - log.debug("Wizard run found, subscribing", { taskId, runId }); - const sub = trpcClient.agent.onSessionEvent.subscribe( - { taskRunId: runId }, - { - onData: (payload: unknown) => { - handleSessionUpdate(payload, (entry) => { - useSetupStore.getState().pushWizardActivity(entry); - }); - }, - onError: (err) => { - log.error("Wizard subscription error", { error: err }); - }, - }, - ); - this.wizardSubscription = sub; - signal.addEventListener( - "abort", - () => { - sub.unsubscribe(); - if (this.wizardSubscription === sub) { - this.wizardSubscription = null; - } - }, - { once: true }, - ); - return; + if (taskData.latest_run?.id) { + runId = taskData.latest_run.id; + break; } } catch { // keep polling } } + + if (!runId || signal.aborted) return; + + log.debug("Wizard run found, subscribing", { taskId, runId }); + const sub = trpcClient.agent.onSessionEvent.subscribe( + { taskRunId: runId }, + { + onData: (payload: unknown) => { + handleSessionUpdate(payload, (entry) => { + useSetupStore.getState().pushWizardActivity(entry); + }); + }, + onError: (err) => { + log.error("Wizard subscription error", { error: err }); + }, + }, + ); + this.wizardSubscription = sub; + signal.addEventListener( + "abort", + () => { + sub.unsubscribe(); + if (this.wizardSubscription === sub) { + this.wizardSubscription = null; + } + }, + { once: true }, + ); + + const POLL_INTERVAL_MS = 5000; + const MAX_ATTEMPTS = 360; // 30 minutes + for (let i = 0; i < MAX_ATTEMPTS; i++) { + try { + await sleep(POLL_INTERVAL_MS, signal); + } catch { + return; // aborted + } + if (useSetupStore.getState().wizardCompleted) return; + try { + const taskRun = await client.getTaskRun(taskId, runId); + if (signal.aborted) return; + if (isTerminalStatus(taskRun.status)) { + log.info("Wizard task reached terminal status", { + taskId, + runId, + status: taskRun.status, + }); + if (taskRun.status === "completed") { + useSetupStore.getState().completeWizard(); + } else { + useSetupStore.getState().skipWizard(); + track(ANALYTICS_EVENTS.SETUP_WIZARD_FAILED, { + reason: "task_run_terminal", + error_message: taskRun.status, + }); + } + return; + } + } catch (err) { + log.warn("Failed to poll wizard task run", { error: err }); + } + } }; - checkForRun().catch((err) => + run().catch((err) => log.error("Wizard event subscribe failed", { error: err }), ); } diff --git a/apps/code/src/renderer/features/setup/stores/setupStore.ts b/apps/code/src/renderer/features/setup/stores/setupStore.ts index bdd8bdd22..d4722024b 100644 --- a/apps/code/src/renderer/features/setup/stores/setupStore.ts +++ b/apps/code/src/renderer/features/setup/stores/setupStore.ts @@ -34,6 +34,7 @@ interface SetupStoreState { discoveryTaskRunId: string | null; wizardTaskId: string | null; wizardSkipped: boolean; + wizardCompleted: boolean; discoveryFeed: AgentFeedState; wizardFeed: AgentFeedState; error: string | null; @@ -49,6 +50,7 @@ interface SetupStoreActions { selectDiscoveredTask: (taskId: string | null) => void; setWizardTaskId: (taskId: string) => void; skipWizard: () => void; + completeWizard: () => void; pushDiscoveryActivity: (entry: ActivityEntry) => void; pushWizardActivity: (entry: ActivityEntry) => void; /** Wipes all setup state — discovered tasks, wizard, feeds, selection. */ @@ -64,6 +66,7 @@ const initialState: SetupStoreState = { discoveryTaskRunId: null, wizardTaskId: null, wizardSkipped: false, + wizardCompleted: false, discoveryFeed: EMPTY_FEED, wizardFeed: EMPTY_FEED, error: null, @@ -163,6 +166,11 @@ export const useSetupStore = create()( set({ wizardSkipped: true }); }, + completeWizard: () => { + log.info("Wizard task reached terminal status"); + set({ wizardCompleted: true }); + }, + pushDiscoveryActivity: (entry) => { set((state) => ({ discoveryFeed: pushEntry(state.discoveryFeed, entry), diff --git a/apps/code/src/shared/types/analytics.ts b/apps/code/src/shared/types/analytics.ts index d49022f89..8373c07a5 100644 --- a/apps/code/src/shared/types/analytics.ts +++ b/apps/code/src/shared/types/analytics.ts @@ -335,7 +335,8 @@ export interface SetupWizardFailedProperties { | "unauthenticated_client" | "missing_directory" | "startup_error" - | "already_installed"; + | "already_installed" + | "task_run_terminal"; error_message?: string; }