diff --git a/src/agent/gateway-session.ts b/src/agent/gateway-session.ts index 1726013c3..b4c0e24f5 100644 --- a/src/agent/gateway-session.ts +++ b/src/agent/gateway-session.ts @@ -73,9 +73,9 @@ export function configureGatewayCredentialsForCI( }; } -// TODO(B2): CI credential loading belongs to the headless provider, not the -// agent. Leaves with the rest of this module once RunInput carries resolved -// inference auth. +// TODO: CI credential loading belongs outside the agent. It leaves with the rest +// of this module once RunInput carries resolved inference auth, later in the +// refactor. export function configureGatewayFromCIEnvironment( projectId: number, region: CloudRegion, diff --git a/src/agent/index.ts b/src/agent/index.ts index 3d5c77d46..c2fe909db 100644 --- a/src/agent/index.ts +++ b/src/agent/index.ts @@ -30,22 +30,12 @@ export { resolveBinding, shouldDisableAsk } from './runner'; export { LONGER_ASK_TIMEOUT_MS } from './wizard-ask-bridge'; /** - * Leaves in B2. Programs own credentials and the legacy adapter dies. - * initializeAgent, executeAgent and buildRunTags are the pre-runAgent surface - * that detection/agentic.ts and run-agent-legacy.ts still call; they go - * through runAgent or leave with detection, and AgentErrorType, which - * classifies executeAgent's failures, goes with them. - * configureGatewayFromCIEnvironment is CI inference auth the headless provider - * owns. flushScanReport becomes a progress event rather than a call. - * downloadSkill leaves once the skill scan runs at load and skill install - * becomes shared. + * Leaves later in the refactor. buildRunTags builds the trace tags runProgram + * and agentic detection send. configureGatewayFromCIEnvironment loads the CI + * gateway token. flushScanReport becomes a progress event rather than a call. + * downloadSkill leaves once skill install becomes shared. */ -export { - AgentErrorType, - buildRunTags, - initializeAgent, - runAgent as executeAgent, -} from './agent-interface'; +export { buildRunTags } from './agent-interface'; export { configureGatewayFromCIEnvironment } from './gateway-session'; export { flushScanReport } from './yara-hooks'; export { downloadSkill } from './tools'; diff --git a/src/agent/progress.ts b/src/agent/progress.ts index ccaca5aca..3f5bbafa3 100644 --- a/src/agent/progress.ts +++ b/src/agent/progress.ts @@ -156,7 +156,7 @@ export interface SpinnerHandle { } /** - * Context the agent attaches to a 401 so the host can pick the right copy. + * Context the agent attaches to a 401 so the caller can pick the right copy. * * `hasSettingsConflict` is true when a Claude Code settings file (project, * project-local, the user's global config, or managed) actually overrides the @@ -185,7 +185,7 @@ export interface AuthErrorDetail { logFilePath: string; } -/** One task as the host renders it. The same shape `WizardUI.syncTodos` takes. */ +/** One task as the caller renders it. The same shape `WizardUI.syncTodos` takes. */ export interface TaskSnapshot { id?: string; source?: string; @@ -207,7 +207,7 @@ export type ProgressLogLevel = 'info' | 'warn' | 'error' | 'success' | 'step'; export type AgentProgress = /** The run's main work has started (`WizardUI.startRun`). */ | { kind: 'lifecycle'; phase: 'started' } - /** The run finished and the host may show its outro (`WizardUI.outro`). */ + /** The run finished and the caller may show its outro (`WizardUI.outro`). */ | { kind: 'lifecycle'; phase: 'completed'; message: string } /** The run spinner (`WizardUI.spinner()`), one handle per run. */ | { @@ -246,7 +246,7 @@ export type ProgressEmitter = (event: AgentProgress) => void; * so `wizard_ask` returns its existing "not available" error, and an optional * task notice is declined — the same path a `--ci` run takes today. * Each request's `signal` aborts when that request times out, the run's - * signal aborts, or another task fails the run. On that abort the host + * signal aborts, or another task fails the run. On that abort the caller * dismisses that request alone, and that dismissal must not throw: abort * listeners run where the agent cannot catch them, so Node would rethrow the * error as an uncaught exception. diff --git a/src/agent/runner/harness/anthropic/index.ts b/src/agent/runner/harness/anthropic/index.ts index f9a99e6c0..9dca0e8c4 100644 --- a/src/agent/runner/harness/anthropic/index.ts +++ b/src/agent/runner/harness/anthropic/index.ts @@ -90,6 +90,7 @@ export const anthropicBackend: AgentHarness = { abortCases: config.abortCases, emitStepEvents: config.trackStepProgress ?? false, resolveStepKey: config.resolveStepKey, + requestRemark: config.requestRemark, triageProvider: boot.triageProvider, }, middleware, diff --git a/src/agent/runner/index.ts b/src/agent/runner/index.ts index 91d35dedc..2114ae761 100644 --- a/src/agent/runner/index.ts +++ b/src/agent/runner/index.ts @@ -36,6 +36,10 @@ import type { } from './shared/types'; import { prepareRun } from './shared/bootstrap'; import { createProgressCollector } from './shared/progress-collector'; +import { + createTranscriptTail, + type TranscriptTail, +} from './shared/transcript-tail'; import { getSequence } from './switchboard'; import { flushScanReport } from '@agent/yara-hooks'; @@ -80,9 +84,15 @@ export async function runAgent( options: RunAgentOptions = {}, ): Promise { let collector: ReturnType | undefined; + let transcript: TranscriptTail | undefined; const snapshot = (): RunResult['snapshot'] => { try { - if (collector) return collector.snapshot(); + if (collector) { + const collected = collector.snapshot(); + return transcript + ? { ...collected, transcriptTail: transcript.text() } + : collected; + } } catch { // A partial snapshot must not replace the run's primary failure. } @@ -98,6 +108,8 @@ export async function runAgent( }; }; const flushReport = (): void => { + // A deferred report keeps counting this run's scans toward the program run's. + if (config.scanReport === 'defer') return; try { const report = flushScanReport({ yaraReport: input.flags.yaraReport }); if (report) @@ -110,6 +122,7 @@ export async function runAgent( try { collector = createProgressCollector(options.onProgress); const { emit } = collector; + if (config.run.collectTranscript) transcript = createTranscriptTail(emit); const log = (message: string) => emit({ kind: 'log', level: 'info', message }); if (options.signal?.aborted) { @@ -143,6 +156,7 @@ export async function runAgent( emit, interaction: options.interaction, signal: options.signal, + transcript, }); result = { ...(options.signal?.aborted && diff --git a/src/agent/runner/sequence/linear.ts b/src/agent/runner/sequence/linear.ts index 1e5c651ae..2724f3a86 100644 --- a/src/agent/runner/sequence/linear.ts +++ b/src/agent/runner/sequence/linear.ts @@ -19,19 +19,20 @@ import { AGENT_ERROR_CODE } from '@agent/error-map'; import { analytics } from '@utils/analytics'; import { formatYaraAbortMessage } from '@agent/yara-hooks'; import { installSkillById } from '@agent/tools'; -import { assemblePrompt } from '../../agent-prompt'; +import { assemblePrompt, type PromptContext } from '../../agent-prompt'; import type { SequenceResult, SequenceContext } from '../shared/types'; import { failed, installFailure } from '../shared/errors'; import { RunOutcome } from '../shared/types'; import { shouldDisableAsk, runOptions } from '../shared/bootstrap'; import { createEmitSpinner } from '../shared/progress-collector'; import { createAskBridge } from '../shared/ask'; +import { withTranscript } from '../shared/transcript-tail'; import { getHarness } from '../switchboard'; export async function runLinearProgram( context: SequenceContext, ): Promise { - // Aborts on the host's signal or when the run ends, so no ask outlives it. + // Aborts on the caller's signal or when the run ends, so no ask outlives it. const controller = new AbortController(); const abortFromHost = () => controller.abort(); context.signal?.addEventListener('abort', abortFromHost, { once: true }); @@ -44,9 +45,17 @@ export async function runLinearProgram( } } -/** The host's `signal` decides the outcome; `runSignal` also ends with the run. */ +/** The caller's `signal` decides the outcome; `runSignal` also ends with the run. */ async function executeLinear( - { config, input, boot, emit, interaction, signal }: SequenceContext, + { + config, + input, + boot, + emit, + interaction, + signal, + transcript, + }: SequenceContext, runSignal: AbortSignal, ): Promise { const { run, composed } = config; @@ -96,12 +105,15 @@ async function executeLinear( signal: runSignal, }); - const middleware = input.flags.benchmark - ? createBenchmarkPipeline(emit, spinner, runOptions(input)) - : undefined; + const middleware = withTranscript( + input.flags.benchmark + ? createBenchmarkPipeline(emit, spinner, runOptions(input)) + : undefined, + transcript, + ); // 7. Build prompt - const prompt = assemblePrompt(run, { + const promptContext: PromptContext = { projectId, projectApiKey, host, @@ -115,7 +127,10 @@ async function executeLinear( surveys: project.surveys_opt_in ?? null, } : null, - }); + }; + const prompt = run.prompt + ? run.prompt(promptContext) + : assemblePrompt(run, promptContext); logToFile(`[agent-runner] prompt assembled (${prompt.length} chars)`); if (signal?.aborted) return aborted(); @@ -178,7 +193,7 @@ async function executeLinear( matched: matched?.message ?? null, }); return { - // An agent that stops itself failed the run; only the host's signal cancels it. + // An agent that stops itself failed the run; only the caller's signal cancels it. outcome: signal?.aborted ? RunOutcome.Aborted : RunOutcome.Failed, failure: { message: matched?.message ?? `${run.integrationLabel} aborted`, diff --git a/src/agent/runner/shared/bootstrap.ts b/src/agent/runner/shared/bootstrap.ts index 95986b2d7..734b70cf6 100644 --- a/src/agent/runner/shared/bootstrap.ts +++ b/src/agent/runner/shared/bootstrap.ts @@ -85,7 +85,7 @@ export async function prepareRun( // readers re-resolve through the cache, which re-mints past the refresh // point. const currentGatewayAuth = () => - // TODO(B2): the agent must not mint inference auth. It receives the + // TODO: the agent must not mint inference auth. It receives the // PostHog token here and derives a gateway token from it, re-minting near // expiry. Programs own credentials (stack plan 4.5): pass a resolved // inference-auth provider on RunInput.credentials and move diff --git a/src/agent/runner/shared/transcript-tail.ts b/src/agent/runner/shared/transcript-tail.ts new file mode 100644 index 000000000..874e4169e --- /dev/null +++ b/src/agent/runner/shared/transcript-tail.ts @@ -0,0 +1,89 @@ +/** The SDK message observer behind `collectTranscript`: a capped text tail plus one `activity` line per step. */ + +import type { ProgressEmitter } from '@agent/progress'; +import type { RunMiddleware } from '../harness/types'; + +/** Only the tail is kept: a caller's report is the run's last output. */ +export const TRANSCRIPT_TAIL_CHARS = 256 * 1024; +const ACTIVITY_LINE_CHARS = 100; + +export interface TranscriptTail extends RunMiddleware { + /** The kept assistant text, one block per line, then the final result. */ + text(): string; +} + +export function createTranscriptTail(emit: ProgressEmitter): TranscriptTail { + const collected: string[] = []; + let collectedChars = 0; + let resultText = ''; + + const collect = (text: string): void => { + collected.push(text); + collectedChars += text.length; + while (collectedChars > TRANSCRIPT_TAIL_CHARS && collected.length > 1) { + collectedChars -= collected.shift()?.length ?? 0; + } + }; + const activity = (line: string): void => emit({ kind: 'activity', line }); + + return { + onMessage(message: any): void { + if (message?.type === 'assistant') { + const content = message.message?.content; + if (!Array.isArray(content)) return; + for (const block of content) { + if (block?.type === 'text' && typeof block.text === 'string') { + collect(block.text); + const line = block.text.trim(); + if (line) { + activity( + line.length > ACTIVITY_LINE_CHARS + ? `${line.slice(0, ACTIVITY_LINE_CHARS)}…` + : line, + ); + } + } else if (block?.type === 'tool_use') { + activity(formatToolUse(block)); + } + } + } else if ( + message?.type === 'result' && + typeof message.result === 'string' + ) { + resultText = message.result; + } + }, + finalize: () => undefined, + text: () => `${collected.join('\n')}\n${resultText}`, + }; +} + +/** Put the tail in front of any other middleware, so both see every message. */ +export function withTranscript( + middleware: RunMiddleware | undefined, + transcript: TranscriptTail | undefined, +): RunMiddleware | undefined { + if (!transcript) return middleware; + if (!middleware) return transcript; + return { + onMessage(message) { + transcript.onMessage(message); + middleware.onMessage(message); + }, + finalize(resultMessage, totalDurationMs) { + transcript.finalize(resultMessage, totalDurationMs); + return middleware.finalize(resultMessage, totalDurationMs); + }, + }; +} + +function formatToolUse(block: any): string { + const name = typeof block?.name === 'string' ? block.name : 'tool'; + const input = (block?.input ?? {}) as Record; + const detail = + (input.file_path as string) || + (input.pattern as string) || + (input.path as string) || + ''; + return detail ? `${name} ${detail}` : name; +} diff --git a/src/agent/runner/shared/types.ts b/src/agent/runner/shared/types.ts index dd029b275..d466bf010 100644 --- a/src/agent/runner/shared/types.ts +++ b/src/agent/runner/shared/types.ts @@ -21,6 +21,7 @@ import type { LLMProvider } from '@posthog/warlock'; import type { AgentInteraction, ProgressEmitter } from '@agent/progress'; import type { EffortLevel } from '../switchboard/models'; import type { SwitchboardCtx } from '../switchboard'; +import type { TranscriptTail } from './transcript-tail'; export type { PromptContext, Credentials }; @@ -152,7 +153,7 @@ export interface RunConfig { programId: string; /** The run definition. A program's session-taking hooks are the caller's, see `hooks`. */ run: AgentRunDefinition; - /** A composed sub-run leaves the terminal outro to its host. */ + /** A composed sub-run leaves the terminal outro to its caller. */ composed: boolean; /** Run-level sequence, harness and model. */ binding: ResolvedBinding; @@ -332,4 +333,6 @@ export interface SequenceContext { emit: ProgressEmitter; interaction: AgentInteraction | undefined; signal?: AbortSignal; + /** Present when the run definition sets `collectTranscript`. */ + transcript?: TranscriptTail; } diff --git a/src/agent/types.ts b/src/agent/types.ts index 179401617..2a8e5d30f 100644 --- a/src/agent/types.ts +++ b/src/agent/types.ts @@ -33,10 +33,10 @@ export type { /** Leaves when the bindings table moves to programs. */ export type { ProgramBinding, SwitchboardCtx } from './runner'; -/** Leaves in B2 with downloadSkill. */ +/** Leaves with downloadSkill, later in the refactor. */ export type { InstallSkillResult } from './tools'; -/** Leaves in B2 with the legacy adapter that records it. */ +/** Leaves with the legacy adapter that records it. */ export type { TaskOutcome } from './runner'; /** Leaves in C2 with runMcpPromptViaSdk. */ diff --git a/src/lib/wizard-session.ts b/src/lib/wizard-session.ts index 71d8b2ff0..f3325299f 100644 --- a/src/lib/wizard-session.ts +++ b/src/lib/wizard-session.ts @@ -26,7 +26,7 @@ import type { } from '@agent/types'; // Leaf module on purpose: shared analytics imports this file, so the agent // entry would form a module cycle here. -// eslint-disable-next-line @typescript-eslint/no-restricted-imports -- B2: the session becomes a TUI projection +// eslint-disable-next-line @typescript-eslint/no-restricted-imports -- the session becomes a TUI projection later in the refactor import { OutroKind } from '@agent/progress'; import { DiscoveredFeature } from '@shared/discovered-feature'; diff --git a/src/programs/__tests__/program-store.test.ts b/src/programs/__tests__/program-store.test.ts new file mode 100644 index 000000000..1383114dc --- /dev/null +++ b/src/programs/__tests__/program-store.test.ts @@ -0,0 +1,116 @@ +import { RunOutcome } from '@agent'; +import type { RunResult } from '@agent/types'; +import type { Credentials } from '@shared/api'; +import { Harness, Sequence } from '@shared/constants'; +import { ErrorCodes } from '@shared/errors'; +import { + ProgramStore, + type ProgramDataProgress, + type ProgramRunProgress, +} from '../program-store'; + +it('forwards attributed copies of run events and settles the original result', () => { + const store = new ProgramStore(); + const observed: ProgramRunProgress[] = []; + const run = store.beginRun('run-1', (progress) => { + observed.push(progress); + if (progress.event.kind === 'tasks') { + progress.event.tasks[0].content = 'observer changed this'; + } + }); + const tasks = [{ content: 'Install', status: 'completed' as const }]; + + run.onProgress({ kind: 'tasks', tasks }); + run.onProgress({ + kind: 'url', + which: 'notebook', + url: 'https://us.posthog.com/notebook/7', + }); + + expect(tasks[0].content).toBe('Install'); + expect(observed.map(({ runId, event }) => [runId, event.kind])).toEqual([ + ['run-1', 'tasks'], + ['run-1', 'url'], + ]); + + const result = { + outcome: RunOutcome.Crashed, + failure: { + code: ErrorCodes.InternalUnhandled, + message: 'Connection failed', + }, + } as RunResult; + run.finish(result); + run.onProgress({ kind: 'status', message: 'late' }); + + expect(store.settledRuns()).toEqual([{ runId: 'run-1', result }]); + expect(store.settledRuns()[0].result).toBe(result); + expect(observed).toHaveLength(2); + expect(store.readDiagnostics()).toEqual([ + { runId: 'run-1', eventKind: 'status', message: 'progress after finish' }, + ]); +}); + +it('does not wait for an observer, and records its failures as diagnostics', async () => { + let rejectObserver!: (reason: Error) => void; + const observer = vi.fn( + () => + new Promise((_resolve, reject) => { + rejectObserver = reject; + }), + ); + const onData = vi.fn().mockImplementationOnce(() => { + throw new Error('observer threw'); + }); + const store = new ProgramStore({ onData }); + const run = store.beginRun( + 'run-1', + observer as (progress: ProgramRunProgress) => void, + ); + + expect(run.onProgress({ kind: 'status', message: 'First' })).toBeUndefined(); + store.setAiSdkStampReported(); + rejectObserver(new Error('delivery failed')); + + expect(store.readData().aiSdkStampReported).toBe(true); + await vi.waitFor(() => { + expect(store.readDiagnostics()).toEqual([ + { eventKind: 'data', message: 'observer threw' }, + { runId: 'run-1', eventKind: 'status', message: 'delivery failed' }, + ]); + }); +}); + +it('copies invocation data on write, on read and in each emitted snapshot', () => { + const observed: ProgramDataProgress[] = []; + const store = new ProgramStore({ + onData: (progress) => observed.push(progress), + }); + const credentials = { + accessToken: 'test-access-token', + projectApiKey: 'test-project-key', + projectId: 42, + host: { region: 'us', apiHost: 'https://example.test' }, + } as Credentials; + const binding = { + sequence: Sequence.linear, + harness: Harness.anthropic, + model: 'claude-test', + }; + + store.setAuthenticated({ credentials, apiProject: null, apiUser: null }); + store.setBinding(binding); + const written = store.readData(); + credentials.accessToken = 'changed input'; + binding.model = 'changed input'; + observed[1].data.binding!.model = 'changed by observer'; + store.readData().credentials!.accessToken = 'changed output'; + + expect(observed).toHaveLength(2); + expect(observed[0].data.binding).toBeNull(); + expect(store.readData()).toEqual(written); + expect(written).toMatchObject({ + credentials: { accessToken: 'test-access-token' }, + binding: { model: 'claude-test' }, + }); +}); diff --git a/src/programs/__tests__/refresh-access-token-if-needed.test.ts b/src/programs/__tests__/refresh-access-token-if-needed.test.ts index 413398dc1..11b55334d 100644 --- a/src/programs/__tests__/refresh-access-token-if-needed.test.ts +++ b/src/programs/__tests__/refresh-access-token-if-needed.test.ts @@ -1,12 +1,16 @@ -import { refreshAccessTokenIfNeeded } from '../authenticate'; +import { rotateCredentials } from '../credentials'; import { refreshAccessToken } from '@utils/oauth'; import { OAuthError } from '@utils/oauth-errors'; import { isGrantRevoked, resetAuthSessionState, } from '@shared/auth-session-state'; -import { resetOAuthSession } from '@shared/oauth-session'; -import type { WizardSession, Credentials } from '@lib/wizard-session'; +import { + configureOAuthSession, + oauthCredentials, + resetOAuthSession, +} from '@shared/oauth-session'; +import type { Credentials } from '@shared/api'; vi.mock('@utils/oauth', async (original) => ({ ...(await original()), @@ -17,16 +21,20 @@ vi.mock('@utils/analytics', () => ({ analytics: { wizardCapture: vi.fn() }, groupsFromUser: vi.fn(), })); - -const setAccessToken = vi.fn(); -vi.mock('@ui', () => ({ - getUI: () => ({ setAccessToken }), -})); +// The real @utils/oauth loads the UI module. +vi.mock('@ui', () => ({ getUI: vi.fn() })); const mockedRefresh = refreshAccessToken as Mock; -function sessionWith(credentials: Partial | null): WizardSession { - return { credentials: credentials as Credentials | null } as WizardSession; +/** The pre-run refresh runProgram does: the OAuth session decides, this grant rotates. */ +async function refresh( + credentials: Partial, + baseUrl?: string, +): Promise { + configureOAuthSession(credentials as Credentials, { + rotate: (held) => rotateCredentials(held, baseUrl), + }); + return (await oauthCredentials())!; } /** Aging enough to be under the 50-minute threshold. */ @@ -37,7 +45,7 @@ const aging = (over: Partial = {}): Partial => ({ ...over, }); -describe('refreshAccessTokenIfNeeded', () => { +describe('rotateCredentials through the OAuth session', () => { beforeEach(() => { vi.clearAllMocks(); resetAuthSessionState(); @@ -45,25 +53,18 @@ describe('refreshAccessTokenIfNeeded', () => { }); it('is a no-op without a refresh token (CI api-key runs, refresh-less grants)', async () => { - await refreshAccessTokenIfNeeded(sessionWith(null)); - await refreshAccessTokenIfNeeded( - sessionWith({ accessToken: 'pha_ci_key', expiresAt: 0 }), - ); + await refresh({ accessToken: 'pha_ci_key', expiresAt: 0 }); expect(mockedRefresh).not.toHaveBeenCalled(); }); it('skips a token that still has most of its lifetime left', async () => { - await refreshAccessTokenIfNeeded( - sessionWith(aging({ expiresAt: Date.now() + 59 * 60 * 1000 })), - ); + await refresh(aging({ expiresAt: Date.now() + 59 * 60 * 1000 })); expect(mockedRefresh).not.toHaveBeenCalled(); }); // `?? 0` would read as "expired" and spend a rotation on every run. it('skips a credential carrying a refresh token but no expiry', async () => { - await refreshAccessTokenIfNeeded( - sessionWith({ accessToken: 'pha_old', refreshToken: 'phr_old' }), - ); + await refresh({ accessToken: 'pha_old', refreshToken: 'phr_old' }); expect(mockedRefresh).not.toHaveBeenCalled(); }); @@ -75,16 +76,20 @@ describe('refreshAccessTokenIfNeeded', () => { token_type: 'Bearer', scope: 'project:read', }); - const session = sessionWith(aging({ projectId: 7 })); - - await refreshAccessTokenIfNeeded(session); + const refreshed = await refresh( + aging({ projectId: 7 }), + 'https://posthog.example', + ); - expect(mockedRefresh).toHaveBeenCalledWith('phr_old', undefined, undefined); - expect(session.credentials!.accessToken).toBe('pha_new'); - expect(session.credentials!.refreshToken).toBe('phr_rotated'); + expect(mockedRefresh).toHaveBeenCalledWith( + 'phr_old', + 'https://posthog.example', + undefined, + ); + expect(refreshed.accessToken).toBe('pha_new'); + expect(refreshed.refreshToken).toBe('phr_rotated'); // Unrelated fields survive the swap. - expect(session.credentials!.projectId).toBe(7); - expect(setAccessToken).toHaveBeenCalledWith(session.credentials); + expect(refreshed.projectId).toBe(7); }); it('refreshes under the minting client id when the credential carries one (provisioning signups)', async () => { @@ -95,9 +100,7 @@ describe('refreshAccessTokenIfNeeded', () => { scope: 'project:read', }); - await refreshAccessTokenIfNeeded( - sessionWith(aging({ oauthClientId: 'client_us_provisioning' })), - ); + await refresh(aging({ oauthClientId: 'client_us_provisioning' })); expect(mockedRefresh).toHaveBeenCalledWith( 'phr_old', @@ -113,30 +116,27 @@ describe('refreshAccessTokenIfNeeded', () => { token_type: 'Bearer', scope: 'project:read', }); - const session = sessionWith(aging()); - const before = session.credentials; + const before = aging(); - await refreshAccessTokenIfNeeded(session); + const refreshed = await refresh(before); - expect(session.credentials).not.toBe(before); - expect(before!.accessToken).toBe('pha_old'); + expect(refreshed).not.toBe(before); + expect(before.accessToken).toBe('pha_old'); // No rotation in the response: the old refresh token has to carry over. - expect(session.credentials!.refreshToken).toBe('phr_old'); + expect(refreshed.refreshToken).toBe('phr_old'); }); it('keeps the existing token and does not throw when the refresh fails', async () => { mockedRefresh.mockRejectedValueOnce(new Error('network down')); - const session = sessionWith(aging()); + const before = aging(); - await expect(refreshAccessTokenIfNeeded(session)).resolves.toBeUndefined(); - expect(session.credentials!.accessToken).toBe('pha_old'); - expect(setAccessToken).not.toHaveBeenCalled(); + await expect(refresh(before)).resolves.toBe(before); }); it('marks the grant revoked on invalid_grant, so a later 401 can name the cause', async () => { mockedRefresh.mockRejectedValueOnce(new OAuthError('invalid_grant')); - await refreshAccessTokenIfNeeded(sessionWith(aging())); + await refresh(aging()); expect(isGrantRevoked()).toBe(true); }); @@ -144,7 +144,7 @@ describe('refreshAccessTokenIfNeeded', () => { it('leaves the grant unmarked for a transport failure, which says nothing about the login', async () => { mockedRefresh.mockRejectedValueOnce(new Error('ETIMEDOUT')); - await refreshAccessTokenIfNeeded(sessionWith(aging())); + await refresh(aging()); expect(isGrantRevoked()).toBe(false); }); diff --git a/src/programs/__tests__/run-agent-legacy.test.ts b/src/programs/__tests__/run-agent-legacy.test.ts index b33f278e6..4773e1df7 100644 --- a/src/programs/__tests__/run-agent-legacy.test.ts +++ b/src/programs/__tests__/run-agent-legacy.test.ts @@ -4,6 +4,8 @@ import path from 'path'; import { runNonInteractive } from '@lib/runners/run-non-interactive'; import { runWizard } from '@lib/runners/run-wizard'; import { authenticate } from '@programs/authenticate'; +import { rotateCredentials } from '@programs/credentials'; +import { resetOAuthSession } from '@shared/oauth-session'; import { runProgramAgent } from '../run-agent-legacy'; import { runAgent, RunOutcome, type RunResult } from '@agent/runner'; import { Harness, Sequence } from '@shared/constants'; @@ -56,10 +58,14 @@ vi.mock('@utils/analytics', () => ({ wizardCapture: vi.fn(), captureException: vi.fn(), setTag: vi.fn(), + identifyUser: vi.fn(), + setGroups: vi.fn(), + groupIdentify: vi.fn(), getAllFlagsForWizard: vi.fn().mockResolvedValue({}), getWizardFlagPayloads: vi.fn().mockReturnValue({}), shutdown: vi.fn().mockResolvedValue(undefined), }, + groupsFromUser: () => ({}), sessionProperties: () => ({}), })); vi.mock('@agent/runner', async (original) => ({ @@ -68,7 +74,11 @@ vi.mock('@agent/runner', async (original) => ({ })); vi.mock('@programs/authenticate', () => ({ authenticate: vi.fn().mockResolvedValue(undefined), - refreshAccessTokenIfNeeded: vi.fn().mockResolvedValue(undefined), +})); +vi.mock('@programs/credentials', () => ({ + rotateCredentials: vi.fn((credentials: unknown) => + Promise.resolve(credentials), + ), })); vi.mock('@shared/claude-settings', () => ({ checkAllSettingsConflicts: vi.fn().mockReturnValue([]), @@ -82,6 +92,9 @@ vi.mock('@utils/wizard-abort', async (original) => ({ vi.mock('../posthog-integration/detect', () => ({ maybeStampAiSdkDetected: vi.fn(), })); +vi.mock('../posthog-integration/ai-sdk-stamp', () => ({ + stampAiSdkDetected: vi.fn(), +})); const program = (id: ProgramConfig['id'] = 'metrics'): ProgramConfig => ({ id, @@ -135,6 +148,7 @@ const finishRun: typeof runAgent = (_config, _input, options) => { beforeEach(() => { vi.clearAllMocks(); + resetOAuthSession(); vi.mocked(authenticate).mockImplementation((sess) => { sess.credentials = session().credentials; return Promise.resolve(); @@ -170,7 +184,7 @@ it.each([ }, ); -it('sends terminal analytics after the outro and the run, before the host goes on', async () => { +it('sends terminal analytics after the outro and the run, before the caller goes on', async () => { const order: string[] = []; logSpy.mockImplementation((line) => { if (line === '└ Done') order.push('outro'); @@ -195,7 +209,7 @@ it('sends terminal analytics after the outro and the run, before the host goes o expect(analytics.shutdown).toHaveBeenCalledExactlyOnceWith('success'); }); -it('clamps a composed program to linear and keeps host analytics alive', async () => { +it('clamps a composed program to linear and keeps the caller analytics alive', async () => { await runProgramAgent(program(), session(), { composed: true }); expect(runAgent).toHaveBeenCalledWith( expect.objectContaining({ @@ -319,6 +333,54 @@ it('rethrows the original crash for the outer runner', async () => { expect(analytics.shutdown).not.toHaveBeenCalled(); }); +it('rethrows a login failure for the CLI roots, before the agent starts', async () => { + const error = new Error('OAuth cancelled'); + vi.mocked(authenticate).mockRejectedValueOnce(error); + await expect(runProgramAgent(program(), session())).rejects.toBe(error); + expect(runAgent).not.toHaveBeenCalled(); + expect(wizardAbort).not.toHaveBeenCalled(); +}); + +it('projects a refreshed token and the AI SDK stamp back onto the session', async () => { + const base = session(); + // Near expiry, so the pre-run refresh rotates it. + const current = { + ...base, + credentials: { + ...base.credentials, + refreshToken: 'phr_test', + expiresAt: Date.now() + 60_000, + }, + }; + // The real login is a no-op when the session already holds credentials. + vi.mocked(authenticate).mockImplementationOnce(() => Promise.resolve()); + const setAccessToken = vi.spyOn(getUI(), 'setAccessToken'); + vi.mocked(rotateCredentials).mockImplementationOnce((credentials) => + Promise.resolve({ ...credentials, accessToken: 'pha_refreshed' }), + ); + await runProgramAgent(program(), current); + expect(vi.mocked(runAgent).mock.calls[0][1].credentials.accessToken).toBe( + 'pha_refreshed', + ); + expect(current.credentials?.accessToken).toBe('pha_refreshed'); + // Only the token fields change; the login keeps its host. + expect(current.credentials?.host).toBeInstanceOf(HostResolution); + expect(setAccessToken).toHaveBeenCalledExactlyOnceWith(current.credentials); + expect(current.aiSdkStampReported).toBe(true); +}); + +it('logs a progress handler that throws instead of dropping it', async () => { + vi.spyOn(getUI(), 'pushStatus').mockImplementationOnce(() => { + throw new Error('screen gone'); + }); + await runProgramAgent(program(), session()); + expect(logToFile).toHaveBeenCalledWith( + expect.stringMatching( + /^\[agent-runner\] progress diagnostic \(status run=.+\): screen gone$/, + ), + ); +}); + it.each([ [Harness.pi, Sequence.linear], [Harness.pi, Sequence.orchestrator], diff --git a/src/programs/__tests__/run-program.test.ts b/src/programs/__tests__/run-program.test.ts new file mode 100644 index 000000000..190f3fa24 --- /dev/null +++ b/src/programs/__tests__/run-program.test.ts @@ -0,0 +1,466 @@ +import { runAgent, RunOutcome } from '@agent'; +import { Harness, Sequence } from '@shared/constants'; +import { HostResolution } from '@shared/host-resolution'; +import type { ApiUser } from '@shared/api'; +import type { RunResult } from '@agent/types'; +import { ErrorCodes } from '@shared/errors'; +import { DiscoveredFeature } from '@lib/wizard-session'; +import { analytics } from '@utils/analytics'; +import { refreshAccessToken } from '@utils/oauth'; +import { oauthCredentials, resetOAuthSession } from '@shared/oauth-session'; +import type { ResolvedProgramCredentials } from '../credentials'; +import type { ProgramInput, ProgramOptions } from '../run-program'; +import { runProgram } from '@programs'; + +vi.mock('@agent', async (importOriginal) => ({ + ...(await importOriginal()), + runAgent: vi.fn(), +})); +vi.mock('@utils/analytics', async (importOriginal) => ({ + ...(await importOriginal()), + analytics: { + runId: 'analytics-run-id', + build: 'test', + setTag: vi.fn(), + wizardCapture: vi.fn(), + captureException: vi.fn(), + identifyUser: vi.fn(), + setGroups: vi.fn(), + groupIdentify: vi.fn(), + }, +})); +vi.mock('@utils/oauth', () => ({ + refreshAccessToken: vi.fn(), + missingOAuthScopes: vi.fn(() => []), +})); +vi.mock('@utils/debug'); + +const run = { + integrationLabel: 'metrics', + spinnerMessage: 'Configuring metrics', + successMessage: 'Metrics configured', + estimatedDurationMinutes: 5, + reportFile: 'posthog-metrics-report.md', + docsUrl: 'https://posthog.com/docs/metrics', +}; + +const snapshot = { + tasks: [], + statusMessages: ['Metrics configured'], + usage: { + inputTokens: 1, + outputTokens: 2, + cacheReadTokens: 0, + cacheCreationTokens: 0, + }, +}; + +const credentials: ResolvedProgramCredentials = { + posthog: { + accessToken: 'phx_access_test', + projectApiKey: 'phx_test', + projectId: 42, + host: HostResolution.fromRegion('us'), + }, + project: null, + apiUser: { + organization: { is_ai_data_processing_approved: true }, + } as ApiUser, +}; + +/** A token close enough to expiry that runProgram refreshes it. */ +const aging = () => ({ + ...credentials.posthog, + accessToken: 'pha_aging', + refreshToken: 'phr_aging', + expiresAt: Date.now() + 10 * 60 * 1000, +}); + +const refreshedToken = { + access_token: 'pha_refreshed', + refresh_token: 'phr_rotated', + expires_in: 3600, + token_type: 'Bearer', + scope: 'project:read', +}; + +const login = (apiUser: ApiUser | null = credentials.apiUser) => ({ + resolve: () => Promise.resolve({ ...credentials, apiUser }), +}); + +/** A program with one post-auth gate, like the source-maps project picker. */ +const gated: ProgramInput = { + installDir: '/project', + run, + program: { postAuthGates: ['detect'] }, +}; + +describe('runProgram', () => { + beforeEach(() => { + vi.clearAllMocks(); + resetOAuthSession(); + vi.mocked(runAgent).mockResolvedValue({ + outcome: RunOutcome.Success, + snapshot, + }); + }); + + it("runs the caller's run definition, settings and route, and returns its final results", async () => { + const excludedTaskTypes = () => ['logs']; + const { signal } = new AbortController(); + const outcome = await runProgram( + 'metrics', + { + installDir: '/project', + runId: 'run-1', + run, + program: { + agentFlow: 'metrics-flow', + allowedTools: ['Agent'], + disallowedTools: ['wizard_ask'], + excludedTaskTypes, + }, + credentials, + overrides: { harness: Harness.anthropic, sequence: Sequence.linear }, + }, + { signal }, + ); + + const [config, input, agentOptions] = vi.mocked(runAgent).mock.calls[0]; + expect(config).toMatchObject({ + programId: 'metrics', + run, + agentFlow: 'metrics-flow', + allowedTools: ['Agent'], + disallowedTools: ['wizard_ask'], + binding: { sequence: Sequence.linear, harness: Harness.anthropic }, + switchboard: { + program: 'metrics', + cliHarness: Harness.anthropic, + cliSequence: Sequence.linear, + }, + wizardMetadata: { + program_id: 'metrics', + integration: 'metrics', + run_id: 'analytics-run-id', + build: 'test', + call_type: 'agent', + SEQUENCE: Sequence.linear, + HARNESS: Harness.anthropic, + }, + }); + expect(config.excludedTaskTypes).toBe(excludedTaskTypes); + expect(input.credentials).toBe(credentials.posthog); + expect(agentOptions?.signal).toBe(signal); + expect(analytics.setTag).toHaveBeenCalledWith('harness', Harness.anthropic); + expect(analytics.wizardCapture).toHaveBeenCalledWith( + 'switchboard resolved', + expect.objectContaining({ program: 'metrics', cli_harness: 'anthropic' }), + ); + expect(outcome).toMatchObject({ + programId: 'metrics', + outcome: RunOutcome.Success, + data: { credentials: { projectId: 42 }, binding: config.binding }, + settledRuns: [ + { runId: 'run-1', result: { outcome: RunOutcome.Success } }, + ], + diagnostics: [], + artifacts: { reportFile: '/project/posthog-metrics-report.md' }, + }); + expect(outcome.failure).toBeUndefined(); + }); + + const closed = () => Promise.reject(new Error('caller closed')); + it.each<[string, () => ProgramOptions, RunOutcome, string]>([ + [ + 'no credentials', + () => ({}), + RunOutcome.Failed, + 'Credentials are required to run metrics.', + ], + [ + 'a rejecting credential provider', + () => ({ credentials: { resolve: closed } }), + RunOutcome.Failed, + 'caller closed', + ], + [ + 'no org AI approval and no caller approval capability', + () => ({ credentials: login(null) }), + RunOutcome.Failed, + 'AI processing approval is required before this program can run.', + ], + [ + 'a declined caller AI approval', + () => ({ + credentials: login(null), + awaitAiApproval: () => Promise.resolve(false), + }), + RunOutcome.Aborted, + 'AI processing approval declined.', + ], + ])( + '%s is a decided result before the agent it guards', + async (_case, options, outcome, message) => { + const result = await runProgram( + 'metrics', + { installDir: '/project', run }, + options(), + ); + + expect(result).toMatchObject({ outcome, failure: { message } }); + expect(result.settledRuns).toHaveLength(0); + expect(runAgent).not.toHaveBeenCalled(); + }, + ); + + it('a program that needs no AI runs without an approval', async () => { + const awaitAiApproval = vi.fn(); + + const result = await runProgram( + 'metrics', + { installDir: '/project', run, program: { requiresAi: false } }, + { credentials: login(null), awaitAiApproval }, + ); + + expect(result.outcome).toBe(RunOutcome.Success); + expect(awaitAiApproval).not.toHaveBeenCalled(); + }); + + it.each<[RunOutcome, RunResult['failure']]>([ + [ + RunOutcome.Failed, + { code: ErrorCodes.AgentApiError, message: 'API Error' }, + ], + [ + RunOutcome.Aborted, + { code: ErrorCodes.AgentAbort, message: 'Agent run cancelled' }, + ], + [ + RunOutcome.Crashed, + { + code: ErrorCodes.InternalUnhandled, + message: 'mint refused', + error: new Error('mint refused'), + }, + ], + ])( + 'a %s agent run settles with its failure and its run', + async (outcome, failure) => { + const result = { outcome, failure, snapshot } as RunResult; + vi.mocked(runAgent).mockResolvedValueOnce(result); + + const settled = await runProgram('metrics', { + installDir: '/project', + runId: 'run-1', + run, + credentials, + }); + + expect(settled).toMatchObject({ outcome, failure }); + expect(settled.settledRuns).toEqual([{ runId: 'run-1', result }]); + }, + ); + + it.each([ + 'credential resolution', + 'AI approval', + 'a post-auth gate', + ] as const)( + 'a caller abort during %s reaches the capability, returns Aborted and starts nothing else', + async (gate) => { + const controller = new AbortController(); + // Each capability takes the invocation signal last and rejects when it aborts. + const park = vi.fn( + (...args: unknown[]) => + new Promise((_resolve, reject) => { + const { signal } = args.at(-1) as { signal: AbortSignal }; + signal.addEventListener('abort', () => + reject(new Error('screen closed')), + ); + }), + ); + const options: ProgramOptions = { + 'credential resolution': { credentials: { resolve: park } }, + 'AI approval': { credentials: login(null), awaitAiApproval: park }, + 'a post-auth gate': { credentials: login(), awaitPostAuthGates: park }, + }[gate]; + + const pending = runProgram('gated', gated, { + ...options, + signal: controller.signal, + }); + await vi.waitFor(() => expect(park).toHaveBeenCalledOnce()); + controller.abort(); + + expect(await pending).toMatchObject({ + outcome: RunOutcome.Aborted, + failure: { + code: ErrorCodes.AgentAbort, + message: 'Run cancelled by the caller.', + }, + }); + expect(runAgent).not.toHaveBeenCalled(); + }, + ); + + it('a caller abort during the token refresh keeps the rotated refresh token', async () => { + const controller = new AbortController(); + vi.mocked(refreshAccessToken).mockImplementationOnce(() => { + controller.abort(); + return Promise.resolve(refreshedToken); + }); + + const result = await runProgram( + 'metrics', + { + installDir: '/project', + run, + credentials: { ...credentials, posthog: aging() }, + }, + { signal: controller.signal }, + ); + + expect(result.outcome).toBe(RunOutcome.Aborted); + const rotated = { + accessToken: 'pha_refreshed', + refreshToken: 'phr_rotated', + }; + expect(result.data.credentials).toMatchObject(rotated); + expect(await oauthCredentials()).toMatchObject(rotated); + expect(runAgent).not.toHaveBeenCalled(); + }); + + it('runs in order: agent started, credentials, approval, post-auth, flags, refresh, route, agent', async () => { + const order: string[] = []; + const answer = (name: string, value: T) => + vi.fn(() => { + order.push(name); + return Promise.resolve(value); + }); + vi.mocked(analytics.wizardCapture).mockImplementationOnce((event) => { + order.push(event); + }); + vi.mocked(analytics.setTag).mockImplementationOnce(() => { + order.push('route'); + }); + vi.mocked(refreshAccessToken).mockImplementationOnce( + answer('refresh', refreshedToken), + ); + vi.mocked(runAgent).mockImplementationOnce( + answer('runAgent', { outcome: RunOutcome.Success, snapshot }), + ); + const flags = { + flags: { 'wizard-test-flag': 'on' }, + payloads: { 'wizard-test-flag': { variant: 'b' } }, + }; + const awaitPostAuthGates = answer('post-auth', undefined); + + const result = await runProgram( + 'gated', + { + ...gated, + run: { ...run, integrationLabel: 'custom-label', skillId: 'skill-x' }, + }, + { + credentials: { + resolve: answer('credentials', { + ...credentials, + posthog: aging(), + apiUser: null, + }), + }, + awaitAiApproval: answer('approval', true), + awaitPostAuthGates, + featureFlags: answer('flags', flags), + }, + ); + + expect(result.outcome).toBe(RunOutcome.Success); + expect(order).toEqual([ + 'agent started', + 'credentials', + 'approval', + 'post-auth', + 'flags', + 'refresh', + 'route', + 'runAgent', + ]); + expect(analytics.wizardCapture).toHaveBeenCalledWith('agent started', { + integration: 'custom-label', + program_id: 'gated', + skill_id: 'skill-x', + }); + expect(awaitPostAuthGates).toHaveBeenCalledWith({ + programId: 'gated', + gates: ['detect'], + signal: expect.objectContaining({ aborted: false }), + }); + expect(vi.mocked(runAgent).mock.calls[0][0]).toMatchObject({ + wizardFlags: flags.flags, + wizardFlagPayloads: flags.payloads, + }); + }); + + it('a caller mutation after the call does not reach the run', async () => { + const flags = { ci: false }; + const host: NonNullable = { region: 'us' }; + + const pending = runProgram( + 'metrics', + { installDir: '/project', run, flags, host }, + { credentials: login() }, + ); + flags.ci = true; + host.region = 'eu'; + await pending; + + const [, runInput] = vi.mocked(runAgent).mock.calls[0]; + expect(runInput.flags.ci).toBe(false); + expect(runInput.host.region).toBe('us'); + }); + + it('a provider is resolved once, then identified and stamped, and refreshed before the agent starts', async () => { + vi.mocked(refreshAccessToken).mockResolvedValueOnce(refreshedToken); + const apiUser = { + distinct_id: 'user-1', + organization: { id: 'org-1', is_ai_data_processing_approved: true }, + } as ApiUser; + const resolve = vi + .fn() + .mockResolvedValue({ posthog: aging(), project: null, apiUser }); + + const result = await runProgram( + 'metrics', + { + installDir: '/project', + run, + host: { baseUrl: 'https://posthog.example' }, + mayReportScanResults: true, + discoveredFeatures: [DiscoveredFeature.LLM], + }, + { credentials: { resolve } }, + ); + + expect(resolve).toHaveBeenCalledOnce(); + expect(analytics.identifyUser).toHaveBeenCalledExactlyOnceWith(apiUser); + expect(analytics.groupIdentify).toHaveBeenCalledExactlyOnceWith( + 'organization', + 'org-1', + { wizard_ai_sdk_detected: true }, + ); + expect(refreshAccessToken).toHaveBeenCalledExactlyOnceWith( + 'phr_aging', + 'https://posthog.example', + undefined, + ); + expect(vi.mocked(runAgent).mock.calls[0][1].credentials.accessToken).toBe( + 'pha_refreshed', + ); + expect(result.data).toMatchObject({ + credentials: { refreshToken: 'phr_rotated' }, + aiSdkStampReported: true, + }); + }); +}); diff --git a/src/programs/authenticate.ts b/src/programs/authenticate.ts index 09aedd452..94a3fa846 100644 --- a/src/programs/authenticate.ts +++ b/src/programs/authenticate.ts @@ -10,14 +10,9 @@ * back rather than fetching again. */ -import type { Credentials, WizardSession } from '@lib/wizard-session'; +import type { WizardSession } from '@lib/wizard-session'; import type { ProgramId } from '@programs/program-registry'; import { getOrAskForProjectData } from '@utils/setup-utils'; -import { refreshAccessToken, missingOAuthScopes } from '@utils/oauth'; -import { WIZARD_OAUTH_SCOPES } from '@shared/constants'; -import { OAuthError } from '@utils/oauth-errors'; -import { markGrantRevoked } from '@shared/auth-session-state'; -import { configureOAuthSession, oauthCredentials } from '@shared/oauth-session'; import { analytics, groupsFromUser } from '@utils/analytics'; import { getUI } from '@ui'; import { logToFile } from '@utils/debug'; @@ -76,62 +71,3 @@ export async function authenticate( if (user) analytics.identifyUser(user); analytics.setGroups(groupsFromUser(user, host.apiHost)); } - -/** - * Grants the token endpoint refuses permanently. A dead grant means the login - * is gone, not that the network blipped, so only these mark the session. - */ -const DEAD_GRANT_CODES = new Set(['invalid_grant', 'invalid_client']); - -// Pre-run refresh through the run's OAuth session, so every later reader shares its rotations. -export async function refreshAccessTokenIfNeeded( - session: WizardSession, -): Promise { - if (!session.credentials) return; - configureOAuthSession(session.credentials, { - rotate: (credentials) => rotateCredentials(credentials, session.baseUrl), - onRefreshed: (refreshed) => { - // Replaced, not mutated: keeps the store and the (possibly shallow-copied) session in step. - session.credentials = refreshed; - getUI().setAccessToken(refreshed); - }, - }); - await oauthCredentials(); -} - -// One refresh-token grant. A failure returns the same credentials: the old token may still have minutes of life. -async function rotateCredentials( - credentials: Credentials, - baseUrl: string | undefined, -): Promise { - if (!credentials.refreshToken) return credentials; - try { - const token = await refreshAccessToken( - credentials.refreshToken, - baseUrl, - credentials.oauthClientId, - ); - return { - ...credentials, - accessToken: token.access_token, - missingScopes: missingOAuthScopes( - [...WIZARD_OAUTH_SCOPES, ...(credentials.missingScopes ?? [])], - token.scope, - ), - // Rotation: keep the returned refresh token or the old one stops working. - refreshToken: token.refresh_token ?? credentials.refreshToken, - expiresAt: Date.now() + token.expires_in * 1000, - }; - } catch (error) { - // A dead grant is recorded but not thrown. If a 401 does follow, the auth-error screen can name the cause. - if (error instanceof OAuthError && DEAD_GRANT_CODES.has(error.code)) { - markGrantRevoked(); - analytics.wizardCapture('auth session expired', { reason: error.code }); - } - logToFile( - '[oauth] token refresh failed, continuing with the existing token:', - error instanceof Error ? error.message : error, - ); - return credentials; - } -} diff --git a/src/programs/credentials.ts b/src/programs/credentials.ts index a344b6a0b..f5320bbed 100644 --- a/src/programs/credentials.ts +++ b/src/programs/credentials.ts @@ -1,6 +1,12 @@ -/** Resolved credentials runProgram passes to one agent run. */ +/** Resolved credentials runProgram passes to one agent run, and the grant that rotates them. */ import type { ApiProject, ApiUser, Credentials } from '@shared/api'; +import { WIZARD_OAUTH_SCOPES } from '@shared/constants'; +import { markGrantRevoked } from '@shared/auth-session-state'; +import { missingOAuthScopes, refreshAccessToken } from '@utils/oauth'; +import { OAuthError } from '@utils/oauth-errors'; +import { analytics } from '@utils/analytics'; +import { logToFile } from '@utils/debug'; export type ResolvedProgramCredentials = { posthog: Credentials; // token, project API key, project ID and host @@ -15,3 +21,46 @@ export type CredentialsProvider = { context: { signal: AbortSignal }, ): Promise; }; + +/** + * Grants the token endpoint refuses permanently. A dead grant means the login + * is gone, not that the network blipped, so only these mark the session. + */ +const DEAD_GRANT_CODES = new Set(['invalid_grant', 'invalid_client']); + +// One refresh-token grant. A failure returns the same credentials: the old token may still have minutes of life. +export async function rotateCredentials( + credentials: Credentials, + baseUrl: string | undefined, +): Promise { + if (!credentials.refreshToken) return credentials; + try { + const token = await refreshAccessToken( + credentials.refreshToken, + baseUrl, + credentials.oauthClientId, + ); + return { + ...credentials, + accessToken: token.access_token, + missingScopes: missingOAuthScopes( + [...WIZARD_OAUTH_SCOPES, ...(credentials.missingScopes ?? [])], + token.scope, + ), + // Rotation: keep the returned refresh token or the old one stops working. + refreshToken: token.refresh_token ?? credentials.refreshToken, + expiresAt: Date.now() + token.expires_in * 1000, + }; + } catch (error) { + // A dead grant is recorded but not thrown. If a 401 does follow, the auth-error screen can name the cause. + if (error instanceof OAuthError && DEAD_GRANT_CODES.has(error.code)) { + markGrantRevoked(); + analytics.wizardCapture('auth session expired', { reason: error.code }); + } + logToFile( + '[oauth] token refresh failed, continuing with the existing token:', + error instanceof Error ? error.message : error, + ); + return credentials; + } +} diff --git a/src/programs/detection/__tests__/agentic-progress.test.ts b/src/programs/detection/__tests__/agentic-progress.test.ts index 424c98414..40d7a4c38 100644 --- a/src/programs/detection/__tests__/agentic-progress.test.ts +++ b/src/programs/detection/__tests__/agentic-progress.test.ts @@ -1,21 +1,35 @@ import { detectProjectsWithAgent } from '../agentic'; import { + AgentErrorType, initializeAgent, runAgent as executeAgent, } from '@agent/agent-interface'; import { buildSession } from '@lib/wizard-session'; import { HostResolution } from '@shared/host-resolution'; -import { getUI } from '@ui'; import { ErrorCodes } from '@shared/errors'; +import { getUI } from '@ui'; vi.mock('@utils/debug'); +// Detection runs the real runAgent pipeline: no analytics or gateway mint may leave the process. +vi.mock('@utils/analytics'); +vi.mock('@agent/gateway-session', async (original) => ({ + ...(await original()), + gatewayAuth: vi.fn(() => + Promise.resolve({ + gatewayUrl: 'https://gateway.test', + token: 'phe_test', + refreshAtMs: Infinity, + }), + ), +})); vi.mock('@ui', () => ({ getUI: () => ui })); const ui = vi.hoisted(() => ({ addTokenUsage: vi.fn(), setStage: vi.fn(), pushStatus: vi.fn(), showAuthError: vi.fn(), - log: { error: vi.fn() }, + startRun: vi.fn(), + log: { error: vi.fn(), warn: vi.fn(), info: vi.fn() }, })); vi.mock('@agent/agent-interface', async (original) => ({ ...(await original()), @@ -23,6 +37,23 @@ vi.mock('@agent/agent-interface', async (original) => ({ runAgent: vi.fn(), })); +function detectionSession() { + const session = buildSession({ installDir: '/tmp/detection-test' }); + session.credentials = { + accessToken: 'test', + projectApiKey: 'phc_test', + projectId: 1, + host: HostResolution.fromApiHost('https://us.posthog.com'), + }; + return session; +} + +beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(initializeAgent).mockReset(); + vi.mocked(executeAgent).mockReset(); +}); + it('keeps initialization and execution progress visible during detection', async () => { const delta = { inputTokens: 5, @@ -60,14 +91,7 @@ it('keeps initialization and execution progress visible during detection', async return Promise.resolve({ kind: 'success' }); }, ); - const session = buildSession({ installDir: '/tmp/detection-test' }); - session.credentials = { - accessToken: 'test', - projectApiKey: 'phc_test', - projectId: 1, - host: HostResolution.fromApiHost('https://us.posthog.com'), - }; - const report = await detectProjectsWithAgent(session, { + const report = await detectProjectsWithAgent(detectionSession(), { programId: 'posthog-integration', targets: [{ id: 'node', name: 'Node.js' }], }); @@ -81,6 +105,54 @@ it('keeps initialization and execution progress visible during detection', async ]); }); +it('sends each agent step to onEvent and the UI only the progress it saw before runAgent', async () => { + vi.mocked(initializeAgent).mockImplementation((config) => + Promise.resolve({ emit: config.emit } as Awaited< + ReturnType + >), + ); + vi.mocked(executeAgent).mockImplementation( + (config, _prompt, _options, _spinner, _messages, middleware) => { + config.emit?.({ kind: 'status', message: 'Scanning' }); + config.emit?.({ kind: 'log', level: 'info', message: 'Info line' }); + config.emit?.({ kind: 'log', level: 'warn', message: 'Warn line' }); + middleware?.onMessage({ + type: 'assistant', + message: { + content: [ + { type: 'text', text: 'Reading the root manifest.' }, + { + type: 'tool_use', + name: 'Read', + input: { file_path: 'package.json' }, + }, + ], + }, + }); + middleware?.onMessage({ + type: 'result', + result: '{"path":".","targetId":"node","framework":"Node.js"}', + }); + return Promise.resolve({ kind: 'success' }); + }, + ); + const lines: string[] = []; + + const report = await detectProjectsWithAgent(detectionSession(), { + programId: 'posthog-integration', + targets: [{ id: 'node', name: 'Node.js' }], + onEvent: (line) => lines.push(line), + }); + + expect(report.projects[0].targetId).toBe('node'); + expect(lines).toEqual(['Reading the root manifest.', 'Read package.json']); + expect(ui.pushStatus).toHaveBeenCalledWith('Scanning'); + expect(ui.log.warn).toHaveBeenCalledWith('Warn line'); + // The scan's run lifecycle and setup logs never reach the program's UI. + expect(ui.log.info).not.toHaveBeenCalled(); + expect(ui.startRun).not.toHaveBeenCalled(); +}); + it('stops optional detection on a data-only 401 before parsing partial JSON', async () => { vi.mocked(initializeAgent).mockResolvedValue( {} as Awaited>, @@ -100,18 +172,51 @@ it('stops optional detection on a data-only 401 before parsing partial JSON', as }); }, ); - const session = buildSession({ installDir: '/tmp/detection-test' }); - session.credentials = { - accessToken: 'test', - projectApiKey: 'phc_test', - projectId: 1, - host: HostResolution.fromApiHost('https://us.posthog.com'), - }; await expect( - detectProjectsWithAgent(session, { + detectProjectsWithAgent(detectionSession(), { programId: 'posthog-integration', targets: [{ id: 'node', name: 'Node.js' }], }), ).rejects.toThrow('Authentication failed (401)'); expect(ui.showAuthError).not.toHaveBeenCalled(); }); + +it('preserves the original error from a decided failure', async () => { + const original = new Error('Gateway bearer rejected'); + vi.mocked(initializeAgent).mockResolvedValue( + {} as Awaited>, + ); + vi.mocked(executeAgent).mockResolvedValue({ + kind: 'decided_failure', + failure: { + code: ErrorCodes.GatewayMintRefused, + message: original.message, + error: original, + }, + }); + + await expect( + detectProjectsWithAgent(detectionSession(), { + programId: 'posthog-integration', + targets: [{ id: 'node', name: 'Node.js' }], + }), + ).rejects.toBe(original); +}); + +it('rejects classified agent failures', async () => { + vi.mocked(initializeAgent).mockResolvedValue( + {} as Awaited>, + ); + vi.mocked(executeAgent).mockResolvedValue({ + kind: 'failure', + classification: AgentErrorType.API_ERROR, + message: 'Agent API unavailable', + }); + + await expect( + detectProjectsWithAgent(detectionSession(), { + programId: 'posthog-integration', + targets: [{ id: 'node', name: 'Node.js' }], + }), + ).rejects.toThrow('Agent API unavailable'); +}); diff --git a/src/programs/detection/__tests__/agentic-retry.test.ts b/src/programs/detection/__tests__/agentic-retry.test.ts index 2ebd05bea..f0454c4eb 100644 --- a/src/programs/detection/__tests__/agentic-retry.test.ts +++ b/src/programs/detection/__tests__/agentic-retry.test.ts @@ -2,6 +2,7 @@ import { AgenticDetectionTimeoutError, detectProjectsWithAgent, } from '@programs/detection/agentic'; +import * as agentEntry from '@agent'; import { AgentErrorType, initializeAgent, @@ -9,12 +10,35 @@ import { } from '@agent/agent-interface'; import { buildSession } from '@lib/wizard-session'; import { HostResolution } from '@shared/host-resolution'; +import { flushScanReport } from '@agent/yara-hooks'; +import { Harness, HAIKU_MODEL, Sequence } from '@shared/constants'; +vi.mock('@utils/analytics'); vi.mock('@agent/agent-interface', async (importOriginal) => ({ ...(await importOriginal()), initializeAgent: vi.fn(), runAgent: vi.fn(), })); +// The entry's runAgent is the real one, spied so each attempt is visible. +vi.mock('@agent', async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, runAgent: vi.fn(actual.runAgent) }; +}); +vi.mock('@agent/yara-hooks', async (importOriginal) => ({ + ...(await importOriginal()), + flushScanReport: vi.fn(), +})); +// The runner mints before each attempt; no mint may leave the process. +vi.mock('@agent/gateway-session', async (importOriginal) => ({ + ...(await importOriginal()), + gatewayAuth: vi.fn(() => + Promise.resolve({ + gatewayUrl: 'https://gateway.test', + token: 'phe_test', + refreshAtMs: Infinity, + }), + ), +})); const init = vi.mocked(initializeAgent); const execute = vi.mocked(runAgent); @@ -22,6 +46,8 @@ const options = { programId: 'posthog-integration', targets: [{ id: 'nextjs', name: 'Next.js' }], }; +const verdict = + '{"path":".","framework":"Next.js","targetId":"nextjs","hasPostHog":false}'; function session() { const value = buildSession({ installDir: '/repo' }); @@ -41,6 +67,24 @@ function emitResult(text: string) { }); } +let deadlines: AbortController[]; +let sdkSawDeadline: boolean[]; + +/** The attempt's deadline fires while its SDK run is active. */ +function timeOut() { + return execute.mockImplementationOnce((agent) => { + deadlines + .at(-1) + ?.abort(new DOMException('The operation timed out.', 'TimeoutError')); + sdkSawDeadline.push(agent.signal?.aborted === true); + return Promise.resolve({ + kind: 'abort', + classification: AgentErrorType.ABORT, + message: 'Agent run cancelled', + }); + }); +} + describe('agentic detection retry', () => { beforeEach(() => { vi.resetAllMocks(); @@ -49,13 +93,50 @@ describe('agentic detection retry', () => { ReturnType >), ); + deadlines = []; + sdkSawDeadline = []; + vi.spyOn(AbortSignal, 'timeout').mockImplementation(() => { + const deadline = new AbortController(); + deadlines.push(deadline); + return deadline.signal; + }); + }); + + afterEach(() => vi.restoreAllMocks()); + + it('runs both attempts through runAgent on linear Haiku with read-only tools, its own prompt, no remark and a deferred scan report', async () => { + timeOut(); + emitResult(verdict); + + await detectProjectsWithAgent(session(), options); + + const calls = vi.mocked(agentEntry.runAgent).mock.calls; + expect(calls).toHaveLength(2); + for (const [config] of calls) { + expect(config.binding).toEqual({ + sequence: Sequence.linear, + harness: Harness.anthropic, + model: HAIKU_MODEL, + }); + expect(config.allowedTools).toEqual(['Read', 'Grep', 'Glob']); + expect(config.scanReport).toBe('defer'); + expect(config.run).toMatchObject({ + collectTranscript: true, + requestRemark: false, + }); + } + // The run definition's prompt replaces the assembled program prompt. + expect(execute.mock.calls[0][1]).toContain( + 'You are scanning a code repository', + ); + expect(execute.mock.calls[0][4]).toMatchObject({ requestRemark: false }); + // The program run's report counts the scan's scans. + expect(flushScanReport).not.toHaveBeenCalled(); }); it('restarts the scan once when the first result has no JSON', async () => { emitResult('I found a Next.js project.'); - emitResult( - '{"path":".","framework":"Next.js","targetId":"nextjs","hasPostHog":false}', - ); + emitResult(verdict); const report = await detectProjectsWithAgent(session(), options); @@ -69,18 +150,14 @@ describe('agentic detection retry', () => { ]); expect(init).toHaveBeenCalledTimes(2); expect(execute).toHaveBeenCalledTimes(2); - expect(execute.mock.calls[0][4]).toEqual( - expect.objectContaining({ timeoutMs: 60_000 }), - ); - expect(execute.mock.calls[1][4]).toEqual( - expect.objectContaining({ timeoutMs: 90_000 }), - ); + expect(vi.mocked(AbortSignal.timeout).mock.calls).toEqual([ + [60_000], + [90_000], + ]); }); it('returns the first valid report without starting a retry', async () => { - emitResult( - '{"path":".","framework":"Next.js","targetId":"nextjs","hasPostHog":false}', - ); + emitResult(verdict); const report = await detectProjectsWithAgent(session(), options); @@ -91,13 +168,8 @@ describe('agentic detection retry', () => { it('retries a timed-out first run with a fresh Haiku session', async () => { const events: string[] = []; - execute.mockResolvedValueOnce({ - kind: 'failure', - classification: AgentErrorType.AGENTIC_DETECTION_TIMEOUT, - }); - emitResult( - '{"path":".","framework":"Next.js","targetId":"nextjs","hasPostHog":false}', - ); + timeOut(); + emitResult(verdict); const report = await detectProjectsWithAgent(session(), { ...options, @@ -107,24 +179,24 @@ describe('agentic detection retry', () => { expect(report.projects).toHaveLength(1); expect(events).toContain('Project scan timed out; retrying...'); expect(execute.mock.calls[0][0]).not.toBe(execute.mock.calls[1][0]); - expect(execute.mock.calls[0][4]).toEqual( - expect.objectContaining({ timeoutMs: 60_000 }), - ); - expect(execute.mock.calls[1][4]).toEqual( - expect.objectContaining({ timeoutMs: 90_000 }), - ); + expect(vi.mocked(AbortSignal.timeout).mock.calls).toEqual([ + [60_000], + [90_000], + ]); }); it('reports a typed timeout when the retry also times out', async () => { - execute.mockResolvedValue({ - kind: 'failure', - classification: AgentErrorType.AGENTIC_DETECTION_TIMEOUT, - }); + timeOut(); + timeOut(); - await expect(detectProjectsWithAgent(session(), options)).rejects.toThrow( - AgenticDetectionTimeoutError, + const scan = detectProjectsWithAgent(session(), options); + + await expect(scan).rejects.toThrow(AgenticDetectionTimeoutError); + await expect(scan).rejects.toThrow( + 'Project scan attempt 2 timed out after 90s', ); expect(execute).toHaveBeenCalledTimes(2); + expect(sdkSawDeadline).toEqual([true, true]); }); it('accepts a streamed verdict after a no-JSON result', async () => { @@ -133,14 +205,7 @@ describe('agentic detection retry', () => { execute.mockImplementationOnce((...args) => { args[5]?.onMessage({ type: 'assistant', - message: { - content: [ - { - type: 'text', - text: '{"path":".","framework":"Next.js","targetId":"nextjs","hasPostHog":false}', - }, - ], - }, + message: { content: [{ type: 'text', text: verdict }] }, }); args[5]?.onMessage({ type: 'result', result: 'Done.' }); return Promise.resolve({ kind: 'success' }); diff --git a/src/programs/detection/agentic.ts b/src/programs/detection/agentic.ts index 33c94a341..f9ad6ec88 100644 --- a/src/programs/detection/agentic.ts +++ b/src/programs/detection/agentic.ts @@ -9,17 +9,18 @@ * Product-knowledge-free by design — the caller passes the targets to classify * into (e.g. source-map skill variants) and maps the result back. Sits next to * the other detection tools (framework, features, package-manager) so it's - * discoverable. Runs AFTER auth: it uses the same agent loop every program - * uses, which needs credentials. + * discoverable. Runs AFTER auth: it goes through the same `runAgent` every + * program uses, which needs credentials. */ -import { - initializeAgent, - executeAgent, - buildRunTags, - AgentSignals, - AgentErrorType, -} from '@agent'; +import { AgentSignals, buildRunTags, runAgent, RunOutcome } from '@agent'; +import type { + AgentProgress, + AgentRunDefinition, + ResolvedBinding, + RunConfig, + RunInput, +} from '@agent/types'; import { isAbsolute, resolve, sep } from 'path'; import { detectNodePackageManagers } from './package-manager.js'; import { @@ -27,12 +28,14 @@ import { AGENTIC_DETECTION_RETRY_TIMEOUT_MS, CallType, getSkillsBaseUrl, + Harness, HAIKU_MODEL, + POSTHOG_DOCS_URL, + Sequence, } from '@shared/constants'; import { analytics } from '@utils/analytics'; import type { WizardSession } from '@lib/wizard-session'; -import type { WizardRunOptions } from '@utils/types'; -import { getUI, type SpinnerHandle } from '@ui'; +import { getUI } from '@ui'; import { createUiReducer } from '@ui/agent-progress'; /** A category the agent classifies each project into (id the agent returns). */ @@ -308,42 +311,45 @@ export function coerceAgenticReport( return { repoType, projects }; } -// eslint-disable-next-line @typescript-eslint/no-explicit-any -function formatToolUse(block: any): string { - const name = typeof block?.name === 'string' ? block.name : 'tool'; - const input = (block?.input ?? {}) as Record; - const detail = - (input.file_path as string) || - (input.pattern as string) || - (input.path as string) || - ''; - return detail ? `${name} ${detail}` : name; -} +/** A fast mechanical scan: linear Haiku on the Anthropic harness. */ +const AGENTIC_DETECTION_BINDING: ResolvedBinding = { + sequence: Sequence.linear, + harness: Harness.anthropic, + model: HAIKU_MODEL, +}; -function sessionToWizardOptions(session: WizardSession): WizardRunOptions { +/** No skill and no remark; the report is read back from the transcript tail. */ +function detectionRunDefinition(prompt: string): AgentRunDefinition { return { - installDir: session.installDir, - ci: session.ci, - debug: session.debug, - benchmark: session.benchmark, - yaraReport: session.yaraReport, - signup: session.signup, - apiKey: session.apiKey, - projectId: session.projectId, + integrationLabel: 'agentic-detect', + prompt: () => prompt, + collectTranscript: true, + requestRemark: false, + detectPackageManager: detectNodePackageManagers, + spinnerMessage: 'Scanning the repo...', + successMessage: 'Detection complete', + errorMessage: 'Detection failed', + estimatedDurationMinutes: 1, + reportFile: '', + docsUrl: POSTHOG_DOCS_URL, }; } -const NOOP_SPINNER: SpinnerHandle = { - start: () => undefined, - stop: () => undefined, - message: () => undefined, -}; +/** What the UI saw before `runAgent`: no run lifecycle, spinner, outro, or setup logs below warn. */ +function reachesUi(event: AgentProgress): boolean { + switch (event.kind) { + case 'lifecycle': + case 'completion': + case 'spinner': + return false; + case 'log': + return event.level === 'warn' || event.level === 'error'; + default: + return true; + } +} -/** - * Drive the wizard's agent loop on HAIKU_MODEL to scan the repo and return a - * structured detection report. Reuses the same setup every program uses, so - * MCP, tools, and credentials are wired identically. - */ +/** Scan the repo with Haiku through `runAgent`; each attempt is a fresh run with its own deadline. */ export async function detectProjectsWithAgent( session: WizardSession, options: AgenticDetectOptions, @@ -359,12 +365,8 @@ export async function detectProjectsWithAgent( rerankIds, onEvent, } = options; - const { accessToken, host } = session.credentials; - const cwd = session.installDir; - const runOptions = sessionToWizardOptions(session); - // Built here rather than inherited: this scan runs before - // `bootstrapProgram`, so there's no `boot.wizardMetadata` yet. + // Built here: the scan runs before the program's own run tags exist. const wizardMetadata = { ...buildRunTags({ programId, @@ -374,108 +376,76 @@ export async function detectProjectsWithAgent( }), call_type: CallType.detection, }; + const config: RunConfig = { + programId, + run: detectionRunDefinition( + buildPrompt(session.installDir, targets, purpose, recommend), + ), + composed: true, + binding: AGENTIC_DETECTION_BINDING, + // Only the orchestrator reads it; the scan is linear. + switchboard: { + program: programId, + composed: true, + flags: {}, + flagPayloads: {}, + }, + skillsBaseUrl: getSkillsBaseUrl(), + wizardFlags: {}, + wizardFlagPayloads: {}, + wizardMetadata, + allowedTools: ['Read', 'Grep', 'Glob'], + // The scan's scans count toward the program run's report. + scanReport: 'defer', + }; + const input: RunInput = { + installDir: session.installDir, + credentials: session.credentials, + project: null, + apiUser: null, + // No benchmark pipeline and no AIO capture: the scan never had either. + flags: { + ci: session.ci, + signup: session.signup, + debug: session.debug, + e2eAsk: false, + localMcp: false, + captureAio: false, + benchmark: false, + yaraReport: session.yaraReport, + }, + host: { projectId: session.projectId, apiKey: session.apiKey }, + }; + const reduceUi = createUiReducer(getUI()); + const forward = (event: AgentProgress): void => { + if (event.kind === 'activity') onEvent?.(event.line); + if (reachesUi(event)) reduceUi(event); + }; - const prompt = buildPrompt(cwd, targets, purpose, recommend); for (let attempt = 0; attempt < 2; attempt++) { const timeoutMs = attempt === 0 ? AGENTIC_DETECTION_FIRST_ATTEMPT_TIMEOUT_MS : AGENTIC_DETECTION_RETRY_TIMEOUT_MS; - const agent = await initializeAgent( - { - emit: createUiReducer(getUI()), - workingDirectory: cwd, - posthogMcpUrl: host.mcpUrl, - posthogApiKey: accessToken, - host, - detectPackageManager: detectNodePackageManagers, - skillsBaseUrl: getSkillsBaseUrl(), - programId, - integrationLabel: 'agentic-detect', - wizardMetadata, - allowedTools: ['Read', 'Grep', 'Glob'], - modelOverride: HAIKU_MODEL, - }, - runOptions, - ); - - // Keeps only the transcript tail — the report JSON is the last output. - const MAX_TRANSCRIPT_CHARS = 256 * 1024; - const collected: string[] = []; - let collectedChars = 0; - const collect = (text: string): void => { - collected.push(text); - collectedChars += text.length; - while (collectedChars > MAX_TRANSCRIPT_CHARS && collected.length > 1) { - collectedChars -= collected.shift()!.length; - } - }; - let resultText = ''; + const deadline = AbortSignal.timeout(timeoutMs); + const result = await runAgent(config, input, { + signal: deadline, + onProgress: forward, + }); - const middleware = { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - onMessage: (message: any): void => { - if (message?.type === 'assistant') { - for (const block of message.message?.content ?? []) { - if (block?.type === 'text' && typeof block.text === 'string') { - collect(block.text); - const line = block.text.trim(); - if (line && onEvent) { - onEvent(line.length > 100 ? `${line.slice(0, 100)}…` : line); - } - } else if (block?.type === 'tool_use') { - onEvent?.(formatToolUse(block)); - } - } - } else if ( - message?.type === 'result' && - typeof message.result === 'string' - ) { - resultText = message.result; - } - }, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - finalize: (_resultMessage: any, _durationMs: number): unknown => - undefined, - }; - - const result = await executeAgent( - agent, - prompt, - runOptions, - NOOP_SPINNER, - { - spinnerMessage: 'Scanning the repo...', - successMessage: 'Detection complete', - errorMessage: 'Detection failed', - requestRemark: false, - timeoutMs, - }, - middleware, - ); - - if ( - result.kind === 'failure' && - result.classification === AgentErrorType.AGENTIC_DETECTION_TIMEOUT - ) { + if (result.outcome === RunOutcome.Aborted && deadline.aborted) { if (attempt === 0) { onEvent?.('Project scan timed out; retrying...'); continue; } throw new AgenticDetectionTimeoutError(attempt + 1, timeoutMs); } - if (result.kind !== 'success') { - if (result.kind === 'decided_failure') { - throw result.failure.error ?? new Error(result.failure.message); - } - throw ( - result.error ?? - new Error(result.message || `Agent error: ${result.classification}`) - ); + if (result.outcome !== RunOutcome.Success) { + throw result.failure.error ?? new Error(result.failure.message); } // Transcript first, final message last — its verdicts win path conflicts. - const output = `${collected.join('\n')}\n${resultText}`; + const output = result.snapshot.transcriptTail ?? ''; const derived = deriveReportJson(output); if (derived !== null) { return coerceAgenticReport( diff --git a/src/programs/posthog-integration/ai-sdk-stamp.ts b/src/programs/posthog-integration/ai-sdk-stamp.ts new file mode 100644 index 000000000..a3945a6a8 --- /dev/null +++ b/src/programs/posthog-integration/ai-sdk-stamp.ts @@ -0,0 +1,38 @@ +/** The AI SDK org stamp. Session-free, so runProgram can load it without the session. */ + +import type { ApiUser } from '@shared/api'; +import { DiscoveredFeature } from '@shared/discovered-feature'; +import { analytics } from '@utils/analytics'; +import { AI_SOURCE_KINDS } from '@programs/warehouse-sources/registry'; +import type { DetectedSource } from '@programs/warehouse-sources/types'; + +/** What the org stamp reads, with no session. */ +export type AiSdkStampEvidence = { + apiUser: Pick | null; + discoveredFeatures: readonly DiscoveredFeature[]; + warehouseSources: readonly DetectedSource[]; + /** Scan consent was granted, so local detection results may be reported. */ + mayReportScanResults: boolean; +}; + +function hasAiSdkEvidence(evidence: AiSdkStampEvidence): boolean { + return ( + evidence.warehouseSources.some((s) => AI_SOURCE_KINDS.has(s.kind)) || + evidence.discoveredFeatures.includes(DiscoveredFeature.LLM) + ); +} + +/** + * Boolean only, on the org, never the list of kinds or any non-AI tool: a + * decline must not leak even the shape of what local detection saw. + */ +export function stampAiSdkDetected(evidence: AiSdkStampEvidence): void { + if (!evidence.mayReportScanResults) return; + const organizationId = evidence.apiUser?.organization?.id; + if (!organizationId) return; + if (!hasAiSdkEvidence(evidence)) return; + + analytics.groupIdentify('organization', organizationId, { + wizard_ai_sdk_detected: true, + }); +} diff --git a/src/programs/posthog-integration/detect.ts b/src/programs/posthog-integration/detect.ts index 09be0515d..bcfa60dc1 100644 --- a/src/programs/posthog-integration/detect.ts +++ b/src/programs/posthog-integration/detect.ts @@ -11,7 +11,6 @@ import type { ProgramReadyContext } from '@programs/program-step'; import { - DiscoveredFeature, mayReportScanResults, ScanConsent, type WizardSession, @@ -25,13 +24,18 @@ import { } from '@programs/detection/index'; import { analytics } from '@utils/analytics'; import { detectWarehouseSources } from '@programs/warehouse-sources/detect'; -import { AI_SOURCE_KINDS } from '@programs/warehouse-sources/registry'; -import type { DetectedSource } from '@programs/warehouse-sources/types'; import { DETECTED_WAREHOUSE_SOURCES_KEY, getDetectedWarehouseSources, } from '@programs/warehouse-source/detect'; import { findPackageJsons } from '@programs/shared/package-scanning'; +import { stampAiSdkDetected } from '@programs/posthog-integration/ai-sdk-stamp'; + +// Session-free, so runProgram can stamp without loading the session. +export { + stampAiSdkDetected, + type AiSdkStampEvidence, +} from '@programs/posthog-integration/ai-sdk-stamp'; export async function detectPostHogIntegration( ctx: ProgramReadyContext, @@ -144,33 +148,6 @@ function detectWarehouseSourcesForSuggestion( } } -function hasAiSdkEvidence( - session: WizardSession, - sources: DetectedSource[], -): boolean { - return ( - sources.some((s) => AI_SOURCE_KINDS.has(s.kind)) || - session.discoveredFeatures.includes(DiscoveredFeature.LLM) - ); -} - -/** - * Boolean only, on the org, never the list of kinds or any non-AI tool: a - * decline must not leak even the shape of what local detection saw. - */ -function stampAiSdkDetected( - session: WizardSession, - sources: DetectedSource[], -): void { - const organizationId = session.apiUser?.organization?.id; - if (!organizationId) return; - if (!hasAiSdkEvidence(session, sources)) return; - - analytics.groupIdentify('organization', organizationId, { - wizard_ai_sdk_detected: true, - }); -} - /** * Fires the org stamp once per session, right after `authenticate()` succeeds * — never from the consent path, since consent on the intro screen resolves @@ -194,9 +171,12 @@ export function maybeStampAiSdkDetected(session: WizardSession): void { // would latch this before consent exists and never stamp even once granted. if (session.aiSdkStampReported) return; session.aiSdkStampReported = true; - if (!mayReportScanResults(session)) return; - - stampAiSdkDetected(session, getDetectedWarehouseSources(session)); + stampAiSdkDetected({ + apiUser: session.apiUser, + discoveredFeatures: session.discoveredFeatures, + warehouseSources: getDetectedWarehouseSources(session), + mayReportScanResults: mayReportScanResults(session), + }); } /** diff --git a/src/programs/program-store.ts b/src/programs/program-store.ts index 67ffa533a..6cacfc8ca 100644 --- a/src/programs/program-store.ts +++ b/src/programs/program-store.ts @@ -1,4 +1,3 @@ -/* eslint-disable @typescript-eslint/no-unused-vars -- A shell: B3 fills in the bodies. */ import type { AgentProgress, ResolvedBinding, RunResult } from '@agent/types'; import type { ApiProject, ApiUser, Credentials } from '@shared/api'; @@ -46,50 +45,128 @@ export type AgentProgressAdapter = { finish(result: RunResult): void; }; +type RunEntry = { + runId: string; + result?: RunResult; +}; + +const MAX_DIAGNOSTICS = 10; + export class ProgramStore { + private readonly runs: RunEntry[] = []; + private readonly diagnostics: ProgramDiagnostic[] = []; + private readonly data: ProgramInvocationData; + private readonly onData?: (progress: ProgramDataProgress) => void; + constructor( options: { aiSdkStampReported?: boolean; onData?: (progress: ProgramDataProgress) => void; } = {}, ) { - throw new Error('ProgramStore: not implemented'); + this.onData = options.onData; + this.data = { + credentials: null, + apiProject: null, + apiUser: null, + detection: { frameworkContext: {} }, + binding: null, + aiSdkStampReported: options.aiSdkStampReported ?? false, + }; } readData(): ProgramInvocationData { - throw new Error('ProgramStore: not implemented'); + return structuredClone(this.data); } setAuthenticated( auth: Pick, ): void { - throw new Error('ProgramStore: not implemented'); + Object.assign(this.data, structuredClone(auth)); + this.emitData(); } setFrameworkContext(key: string, value: unknown): void { - throw new Error('ProgramStore: not implemented'); + this.data.detection.frameworkContext[key] = structuredClone(value); + this.emitData(); } setBinding(binding: ResolvedBinding): void { - throw new Error('ProgramStore: not implemented'); + this.data.binding = structuredClone(binding); + this.emitData(); } setAiSdkStampReported(): void { - throw new Error('ProgramStore: not implemented'); + if (this.data.aiSdkStampReported) return; + this.data.aiSdkStampReported = true; + this.emitData(); } beginRun( runId: string, observer?: (progress: ProgramRunProgress) => void, ): AgentProgressAdapter { - throw new Error('ProgramStore: not implemented'); + const run: RunEntry = { runId }; + this.runs.push(run); + + return { + onProgress: (event) => { + const source = { runId: run.runId, eventKind: event.kind }; + if (run.result) { + this.recordDiagnostic(source, 'progress after finish'); + return; + } + if (!observer) return; + this.deliver(source, () => + observer({ kind: 'run', runId, event: structuredClone(event) }), + ); + }, + finish: (result) => { + run.result = result; + }, + }; } settledRuns(): SettledProgramRun[] { - throw new Error('ProgramStore: not implemented'); + return this.runs.flatMap(({ runId, result }) => + result ? [{ runId, result }] : [], + ); } readDiagnostics(): ProgramDiagnostic[] { - throw new Error('ProgramStore: not implemented'); + return this.diagnostics.map((diagnostic) => ({ ...diagnostic })); + } + + private emitData(): void { + const onData = this.onData; + if (!onData) return; + this.deliver({ eventKind: 'data' }, () => + onData({ kind: 'program', data: this.readData() }), + ); + } + + /** Never waits for an observer; a throw or a rejection becomes a diagnostic. */ + private deliver(source: DiagnosticSource, send: () => unknown): void { + try { + const delivery = send(); + if ( + delivery && + typeof (delivery as PromiseLike).then === 'function' + ) { + void Promise.resolve(delivery).catch((error: unknown) => { + this.recordDiagnostic(source, error); + }); + } + } catch (error) { + this.recordDiagnostic(source, error); + } + } + + private recordDiagnostic(source: DiagnosticSource, error: unknown): void { + this.diagnostics.push({ + ...source, + message: error instanceof Error ? error.message : String(error), + }); + if (this.diagnostics.length > MAX_DIAGNOSTICS) this.diagnostics.shift(); } } diff --git a/src/programs/run-agent-legacy.ts b/src/programs/run-agent-legacy.ts index ee4134fd9..15efeb614 100644 --- a/src/programs/run-agent-legacy.ts +++ b/src/programs/run-agent-legacy.ts @@ -1,37 +1,25 @@ /** * The session-driven agent runner every existing caller uses. * - * `runProgramAgent(programConfig, session)` rebuilds today's behavior on top of the - * functional `runAgent(config, input, options)` in `@lib/agent/runner`: it - * runs the gates the TUI owns (health, settings, AI opt-in, post-auth steps), - * authenticates, resolves the program's binding, builds the agent's inputs - * from the session, maps every progress event back onto `getUI()` one call - * per event, answers the agent's questions through `getUI()`, and applies the - * result — `wizardAbort` with the outcome's terminal status for a decided - * failure, the terminal analytics event for a finished top-level run. + * `runProgramAgent(programConfig, session)` runs the gates the TUI owns + * (health, settings), then calls `runProgram` as its caller, backed by the + * session and `getUI()`: credentials come from `authenticate`, the AI + * opt-in and post-auth gates park on the UI, every progress event maps back + * onto `getUI()`, and the invocation's data projects back onto the session. + * It applies the result — `wizardAbort` with the outcome's terminal status for + * a decided failure, the terminal analytics event for a finished top-level run. * * This is the only file that knows about `getUI()`, the session and - * `wizardAbort` on the agent's behalf. Programs replace it in Release B. + * `wizardAbort` on the agent's behalf. + * + * ⚠️ Temporary adapter. It is removed later in the refactor. */ -import type { WizardSession } from '@lib/wizard-session'; +import { mayReportScanResults, type WizardSession } from '@lib/wizard-session'; import { analytics } from '@utils/analytics'; -import { getUI } from '@ui'; +import { getUI, type WizardUI } from '@ui'; import { createUiReducer, uiInteraction } from '@ui/agent-progress'; -import { - buildRunTags, - flushScanReport, - resolveBinding, - runAgent, - RunOutcome, - TASK_OUTCOMES_KEY, -} from '@agent'; -import type { - ProgramBinding, - RunConfig, - RunInput, - SwitchboardCtx, -} from '@agent/types'; +import { flushScanReport, RunOutcome, TASK_OUTCOMES_KEY } from '@agent'; import type { ProgramRun } from './program-run'; import { backupAndFixClaudeSettings, @@ -50,21 +38,17 @@ import { enableDebugLogs, logToFile, initLogFile } from '@utils/debug'; import { registerCleanup, wizardAbort } from '@utils/wizard-abort'; import { ErrorCodes } from '@shared/errors'; import { isNonInteractiveEnvironment } from '@utils/environment'; -import { - getSkillsBaseUrl, - Sequence, - WIZARD_ORCHESTRATOR_FLAG_KEY, - WIZARD_SELF_DRIVING_USE_PI_HARNESS_FLAG_KEY, - type Integration, -} from '@shared/constants'; +import { Sequence, type Integration } from '@shared/constants'; import { FRAMEWORK_REGISTRY } from '@programs/frameworks/registry'; import { postAuthGateSteps, type ProgramConfig } from './program-step'; -import { authenticate, refreshAccessTokenIfNeeded } from './authenticate'; -import { maybeStampAiSdkDetected } from './posthog-integration/detect'; +import { authenticate } from './authenticate'; import { removeAuditLedger, startAuditLedgerWatcher, } from './audit/ledger-watcher'; +import { getDetectedWarehouseSources } from './warehouse-source/detect'; +import { runProgram, type WizardFlagSnapshot } from './run-program'; +import type { ProgramInvocationData } from './program-store'; /** * Resolve a ProgramConfig's agent run definition and execute the pipeline. @@ -99,17 +83,19 @@ export async function runProgramAgent( ? await programConfig.run(session) : programConfig.run; - await runProgram(session, runDef, programConfig, options.composed ?? false); + await runSessionProgram( + session, + runDef, + programConfig, + options.composed ?? false, + ); } finally { releaseLedger(); } } -/** - * Gates → authenticate → flags → binding → the functional run → apply result. - * Every step happens in the order it did inside the agent's bootstrap. - */ -async function runProgram( +/** Gates → runProgram on the session's behalf → apply result. */ +async function runSessionProgram( session: WizardSession, run: ProgramRun, programConfig: ProgramConfig, @@ -134,171 +120,138 @@ async function runProgram( // 3. Settings conflicts await runSettingsGate(session); - analytics.wizardCapture('agent started', { - integration: run.integrationLabel, - program_id: programConfig.id, - skill_id: run.skillId ?? null, - }); - - // 4. Authenticate — idempotent within a run (see authenticate()). A second - // agent run in the same invocation (self-driving's integration phase) reuses - // the first login; it does not launch another OAuth. authenticate() also - // identifies the user and sets analytics groups. - await authenticate(session, programConfig.id); - maybeStampAiSdkDetected(session); - - // 4.5. AI opt-in enforcement. Parks here while AiOptInRequiredScreen is - // up if the org hasn't approved third-party AI — BEFORE the skill - // install and agent start, so no source leaves the machine. The screen - // alone is cosmetic; this await is the actual gate. Resolves - // immediately when the program declared requiresAi: false or in CI. - logToFile('[agent-runner] checking AI opt-in gate'); - await getUI().waitForAiOptIn(); - logToFile('[agent-runner] AI opt-in gate cleared'); - - // Park for any interactive step the user must complete AFTER authenticating - // but BEFORE the agent runs — e.g. the source-maps project picker, which - // needs credentials to scan and writes its choice to frameworkContext that - // the run prompt reads. Generic: await every gated step between auth and run. - for (const step of postAuthGateSteps(programConfig.steps)) { - logToFile(`[agent-runner] awaiting post-auth gate: ${step.id}`); - await getUI().waitForGate(step.id); - logToFile(`[agent-runner] post-auth gate cleared: ${step.id}`); - } - - // Feature flags. Both arms need these, and the routing decision reads them. - const wizardFlags = await analytics.getAllFlagsForWizard(); - const wizardFlagPayloads = analytics.getWizardFlagPayloads(); - - // Gateway trace tags for this run; the binding below stamps its axes on. - const wizardMetadata = buildRunTags({ - programId: programConfig.id, - integration: run.integrationLabel, - runId: analytics.runId, - build: analytics.build, - skillId: run.skillId, - }); - - // The agent can't swap tokens mid-run, so freshness is measured after every - // park above, right before the agent mints. - await refreshAccessTokenIfNeeded(session); - - // Credentials (incl. the resolved host family and its MCP url) live on - // `session.credentials`; narrow once at this boundary — `authenticate` above - // set them — so downstream readers get a non-null type without asserting. - const credentials = session.credentials!; - - // Resolve which sequence and harness will run a program (CLI → PostHog flag → - // per-program binding → default), tag both axes onto analytics, and hand the - // binding to the agent for dispatch. - const switchboard: SwitchboardCtx = { - program: programConfig.id, - composed, - flags: wizardFlags, - flagPayloads: wizardFlagPayloads, - cliHarness: session.harness, - cliSequence: session.sequence, - cliModel: session.model, - }; - const binding = resolveBinding(switchboard); - analytics.setTag('sequence', binding.sequence); - analytics.setTag('harness', binding.harness); - wizardMetadata.SEQUENCE = binding.sequence; - wizardMetadata.HARNESS = binding.harness; - captureSwitchboardDecision(switchboard, binding); - const ui = getUI(); - - // Cleanup coverage for the abort/cancel path: `wizardAbort` runs the - // registered cleanups, and the agent's own `finally` covers completion. - // flushScanReport is idempotent, so the overlap is a harmless no-op. - registerCleanup(() => { - const report = flushScanReport({ yaraReport: session.yaraReport }); - if (report) ui.log.info(report); - }); - - // Linear settings restoration fires on entry to the outro screen, so it - // is registered before the run can reach that screen. Same owner, same - // timing as before; the abort path still restores through the cleanup - // `backupAndFixClaudeSettings` registered. - if (binding.sequence === Sequence.linear) { - ui.onEnterScreen('outro', () => restoreClaudeSettings(session.installDir)); - } + const reduceUi = createUiReducer(ui); + const projectData = projectProgramData(ui, session); + + // runProgram turns a throwing capability into a failed run; the CLI roots expect the throw. + let capabilityFailure: { error: unknown } | undefined; + const keepFailure = (work: Promise): Promise => + work.catch((error: unknown) => { + capabilityFailure ??= { error }; + throw error; + }); const framework = session.integration ?? session.skillId ?? undefined; - const config: RunConfig = { - programId: programConfig.id, - run, - composed, - binding, - switchboard, - skillsBaseUrl: getSkillsBaseUrl(), - wizardFlags, - wizardFlagPayloads, - wizardMetadata, - allowedTools: programConfig.allowedTools, - disallowedTools: programConfig.disallowedTools, - agentFlow: programConfig.agentFlow, - excludedTaskTypes: programConfig.excludedTaskTypes, - seedTasks: programConfig.seedTasks - ? () => programConfig.seedTasks!(session) - : undefined, - hooks: { - postRun: run.postRun - ? (creds) => run.postRun!(session, creds) - : undefined, - buildOutroData: run.buildOutroData - ? (creds) => run.buildOutroData!(session, creds) ?? undefined + const result = await runProgram( + programConfig.id, + { + installDir: session.installDir, + run, + composed, + overrides: { + harness: session.harness, + sequence: session.sequence, + model: session.model, + }, + skillId: session.skillId ?? undefined, + integration: session.integration, + frameworkDocsUrl: framework + ? FRAMEWORK_REGISTRY[framework as Integration]?.metadata.docsUrl : undefined, - buildOutroNextSteps: run.buildOutroNextSteps - ? (creds, completed) => - run.buildOutroNextSteps!(session, creds, completed) + flags: { + ci: session.ci, + signup: session.signup, + debug: session.debug, + e2eAsk: session.e2eAsk, + localMcp: session.localMcp, + captureAio: session.captureAio, + benchmark: session.benchmark, + yaraReport: session.yaraReport, + }, + host: { + baseUrl: session.baseUrl, + region: session.region, + email: session.email, + projectId: session.projectId, + apiKey: session.apiKey, + }, + seedTasks: programConfig.seedTasks + ? () => programConfig.seedTasks!(session) : undefined, - recordTaskOutcomes: (outcomes) => { - session.frameworkContext[TASK_OUTCOMES_KEY] = outcomes; + hooks: { + postRun: run.postRun + ? (creds) => run.postRun!(session, creds) + : undefined, + buildOutroData: run.buildOutroData + ? (creds) => run.buildOutroData!(session, creds) ?? undefined + : undefined, + buildOutroNextSteps: run.buildOutroNextSteps + ? (creds, completed) => + run.buildOutroNextSteps!(session, creds, completed) + : undefined, + recordTaskOutcomes: (outcomes) => { + session.frameworkContext[TASK_OUTCOMES_KEY] = outcomes; + }, }, + program: { + requiresAi: programConfig.requiresAi, + agentFlow: programConfig.agentFlow, + allowedTools: programConfig.allowedTools, + disallowedTools: programConfig.disallowedTools, + excludedTaskTypes: programConfig.excludedTaskTypes, + postAuthGates: postAuthGateSteps(programConfig.steps).map( + (step) => step.id, + ), + }, + aiSdkStampReported: session.aiSdkStampReported, + discoveredFeatures: session.discoveredFeatures, + warehouseSources: getDetectedWarehouseSources(session), + mayReportScanResults: mayReportScanResults(session), }, - }; - const input: RunInput = { - installDir: session.installDir, - credentials, - project: session.apiProject, - apiUser: session.apiUser, - skillId: session.skillId ?? undefined, - integration: session.integration, - frameworkDocsUrl: framework - ? FRAMEWORK_REGISTRY[framework as Integration]?.metadata.docsUrl - : undefined, - flags: { - ci: session.ci, - signup: session.signup, - debug: session.debug, - e2eAsk: session.e2eAsk, - localMcp: session.localMcp, - captureAio: session.captureAio, - benchmark: session.benchmark, - yaraReport: session.yaraReport, - }, - host: { - baseUrl: session.baseUrl, - region: session.region, - email: session.email, - projectId: session.projectId, - apiKey: session.apiKey, + { + credentials: { + // Idempotent within a run: a second agent run in the same invocation + // (self-driving's integration phase) reuses the first login. + resolve: () => + keepFailure( + authenticate(session, programConfig.id).then(() => ({ + posthog: session.credentials!, + project: session.apiProject, + apiUser: session.apiUser, + })), + ), + }, + // The actual AI opt-in gate: it parks while AiOptInRequiredScreen is up, + // before the skill install and agent start, so no source leaves the machine. + awaitAiApproval: async () => { + logToFile('[agent-runner] checking AI opt-in gate'); + await ui.waitForAiOptIn(); + logToFile('[agent-runner] AI opt-in gate cleared'); + return true; + }, + // Each step the user settles between auth and run, such as the source-maps + // project picker, which writes its choice to frameworkContext for the prompt. + awaitPostAuthGates: async ({ gates }) => { + for (const gate of gates) { + logToFile(`[agent-runner] awaiting post-auth gate: ${gate}`); + await ui.waitForGate(gate); + logToFile(`[agent-runner] post-auth gate cleared: ${gate}`); + } + }, + featureFlags: () => keepFailure(loadWizardFlags()), + onProgress: (progress) => { + if (progress.kind === 'run') reduceUi(progress.event); + else projectData(progress.data); + }, + interaction: uiInteraction(ui), }, - }; - - const result = await runAgent(config, input, { - onProgress: createUiReducer(ui), - interaction: uiInteraction(ui), - }); + ); + // runProgram keeps a throwing progress handler or a late event as a diagnostic, so log it. + for (const diagnostic of result.diagnostics) { + const runLabel = 'runId' in diagnostic ? ` run=${diagnostic.runId}` : ''; + logToFile( + `[agent-runner] progress diagnostic (${diagnostic.eventKind}${runLabel}): ${diagnostic.message}`, + ); + } + if (capabilityFailure) throw capabilityFailure.error; - // The host owns process exits, terminal analytics and rethrowing crashes. + // The adapter owns process exits, terminal analytics and rethrowing crashes. if (result.outcome === RunOutcome.Crashed) { - throw result.failure.error; + throw result.failure?.error; } if (result.outcome !== RunOutcome.Success) { - if (result.failure.authErrorDetail) { + if (result.failure?.authErrorDetail) { ui.showAuthError(result.failure.authErrorDetail); } // The terminal status follows how the run ended, not whether an Error came back. @@ -307,7 +260,7 @@ async function runProgram( status: result.outcome === RunOutcome.Aborted ? 'cancelled' : 'error', }); } else if (!composed) { - // A composed sub-run leaves the terminal event to its host program's run. + // A composed sub-run leaves the terminal event to its parent program's run. // The run already succeeded: a failed flush is logged, never the outcome. try { await analytics.shutdown('success'); @@ -317,6 +270,57 @@ async function runProgram( } } +/** Mirror the invocation's data onto the session and the UI the TUI reads. */ +function projectProgramData( + ui: WizardUI, + session: WizardSession, +): (data: ProgramInvocationData) => void { + let bindingSeen = false; + return (data) => { + const current = session.credentials; + if ( + current && + data.credentials && + data.credentials.accessToken !== current.accessToken + ) { + // A refresh replaces only the token fields; the login keeps its host. + session.credentials = { + ...current, + accessToken: data.credentials.accessToken, + refreshToken: data.credentials.refreshToken, + expiresAt: data.credentials.expiresAt, + }; + ui.setAccessToken(session.credentials); + } + if (data.aiSdkStampReported) session.aiSdkStampReported = true; + if (!data.binding || bindingSeen) return; + bindingSeen = true; + + // Cleanup coverage for the abort/cancel path: `wizardAbort` runs the + // registered cleanups, and the agent's own `finally` covers completion. + // flushScanReport is idempotent, so the overlap is a harmless no-op. + registerCleanup(() => { + const report = flushScanReport({ yaraReport: session.yaraReport }); + if (report) ui.log.info(report); + }); + + // Linear settings restoration fires on entry to the outro screen, so it + // is registered before the run can reach that screen. The abort path + // still restores through the cleanup `backupAndFixClaudeSettings` + // registered. + if (data.binding.sequence === Sequence.linear) { + ui.onEnterScreen('outro', () => + restoreClaudeSettings(session.installDir), + ); + } + }; +} + +const loadWizardFlags = async (): Promise => ({ + flags: await analytics.getAllFlagsForWizard(), + payloads: analytics.getWizardFlagPayloads(), +}); + // ── Gates ───────────────────────────────────────────────────────────── async function runHealthGate( @@ -450,50 +454,3 @@ async function runSettingsGate(session: WizardSession): Promise { logToFile('[agent-runner] settings override resolved'); } } - -// ── Switchboard telemetry ───────────────────────────────────────────── - -/** - * One event + one log line per run: what entered the switchboard, which - * precedence rung decided each axis, and the final pick. - */ -function captureSwitchboardDecision( - ctx: SwitchboardCtx, - binding: ProgramBinding, -): void { - const trace = ctx.trace ?? {}; - // Unpinned orchestrator runs choose a model per task from the context-mill agent prompts; the orchestrator logs that map once the prompts load. - const perTaskModel = - binding.sequence === Sequence.orchestrator && trace.model === 'binding'; - const model = perTaskModel ? 'chosen-per-task' : binding.model; - const modelSource = perTaskModel ? 'agent-prompts' : trace.model; - analytics.wizardCapture('switchboard resolved', { - program: ctx.program, - flag_self_driving_use_pi_harness: - ctx.flags[WIZARD_SELF_DRIVING_USE_PI_HARNESS_FLAG_KEY], - flag_self_driving_pi_payload: JSON.stringify( - ctx.flagPayloads?.[WIZARD_SELF_DRIVING_USE_PI_HARNESS_FLAG_KEY] ?? null, - ), - flag_orchestrator: ctx.flags[WIZARD_ORCHESTRATOR_FLAG_KEY], - cli_harness: ctx.cliHarness, - cli_sequence: ctx.cliSequence, - cli_model: ctx.cliModel, - harness_source: trace.harness, - model_source: modelSource, - sequence_source: trace.sequence, - harness: binding.harness, - model, - thinking_level: binding.thinkingLevel, - sequence: binding.sequence, - }); - logToFile( - `[switchboard] decision: program=${ctx.program}` + - ` in(orchestrator=${ctx.flags[WIZARD_ORCHESTRATOR_FLAG_KEY] ?? '-'},` + - ` cli=${ctx.cliHarness ?? '-'}/${ctx.cliSequence ?? '-'}/${ - ctx.cliModel ?? '-' - })` + - ` → harness=${binding.harness} (${trace.harness ?? '?'})` + - ` model=${model} (${modelSource ?? '?'})` + - ` sequence=${binding.sequence} (${trace.sequence ?? '?'})`, - ); -} diff --git a/src/programs/run-program.ts b/src/programs/run-program.ts index 4c8c79280..6c36aa9cd 100644 --- a/src/programs/run-program.ts +++ b/src/programs/run-program.ts @@ -1,26 +1,43 @@ /** A caller-owned program invocation. No TUI store or session is required. */ -/* eslint-disable @typescript-eslint/no-unused-vars, @typescript-eslint/require-await -- A shell: B3 fills in the body. */ -import type { RunOutcome } from '@agent'; +import path from 'path'; +import { randomUUID } from 'crypto'; +import { buildRunTags, resolveBinding, runAgent, RunOutcome } from '@agent'; import type { AgentInteraction, AgentRunDefinition, + ProgramBinding, RunConfig, RunHooks, RunInput, RunResult, + SwitchboardCtx, } from '@agent/types'; -import type { Harness, Integration, Sequence } from '@shared/constants'; +import { + getSkillsBaseUrl, + Sequence, + WIZARD_ORCHESTRATOR_FLAG_KEY, + WIZARD_SELF_DRIVING_USE_PI_HARNESS_FLAG_KEY, + type Harness, + type Integration, +} from '@shared/constants'; +import { ErrorCodes } from '@shared/errors'; import type { DiscoveredFeature } from '@shared/discovered-feature'; +import { analytics, groupsFromUser } from '@utils/analytics'; +import { logToFile } from '@utils/debug'; import type { DetectedSource } from './warehouse-sources/types'; -import type { - CredentialsProvider, - ResolvedProgramCredentials, +import { configureOAuthSession, oauthCredentials } from '@shared/oauth-session'; +import { + rotateCredentials, + type CredentialsProvider, + type ResolvedProgramCredentials, } from './credentials'; -import type { - ProgramDiagnostic, - ProgramInvocationData, - ProgramProgress, - SettledProgramRun, +import { stampAiSdkDetected } from './posthog-integration/ai-sdk-stamp'; +import { + ProgramStore, + type ProgramDiagnostic, + type ProgramInvocationData, + type ProgramProgress, + type SettledProgramRun, } from './program-store'; /** Launch-time routing choices, such as the CLI's --harness, --sequence and --model. */ @@ -97,11 +114,284 @@ export interface ProgramRunOutcome { failure?: RunResult['failure']; // code and message on any non-success } +/** Handed to the caller's capabilities when it supplied no signal. */ +const NEVER_ABORTED = new AbortController().signal; + +const DEFAULT_FLAGS: RunInput['flags'] = { + ci: false, + signup: false, + debug: false, + e2eAsk: false, + localMcp: false, + captureAio: false, + benchmark: false, + yaraReport: false, +}; + +/** Fields that carry functions or class instances; everything else is data. */ +const KEPT_BY_REFERENCE = [ + 'credentials', + 'run', + 'program', + 'hooks', + 'seedTasks', +] as const satisfies readonly (keyof ProgramInput)[]; + +/** Copy the caller's input, so a later caller write cannot reach the run or its hooks. */ +function snapshotProgramInput(input: ProgramInput): ProgramInput { + const data: Partial = { ...input }; + for (const key of KEPT_BY_REFERENCE) delete data[key]; + return { ...input, ...structuredClone(data) }; +} + /** Run an existing program from explicit inputs, with invocation-owned state. */ export async function runProgram( programId: string, callerInput: ProgramInput, options: ProgramOptions = {}, ): Promise { - throw new Error('runProgram: not implemented'); + const input = snapshotProgramInput(callerInput); + const store = new ProgramStore({ + aiSdkStampReported: input.aiSdkStampReported, + onData: options.onProgress, + }); + const { installDir, run } = input; + const program = input.program ?? {}; + const artifacts: ProgramRunOutcome['artifacts'] = {}; + const runId = input.runId ?? randomUUID(); + const signal = options.signal ?? NEVER_ABORTED; + + const settle = ( + outcome: RunOutcome, + failure?: RunResult['failure'], + ): ProgramRunOutcome => ({ + programId, + outcome, + data: store.readData(), + settledRuns: store.settledRuns(), + diagnostics: store.readDiagnostics(), + artifacts, + ...(failure && { failure }), + }); + const fail = (message: string) => + settle(RunOutcome.Failed, { code: ErrorCodes.InternalUnhandled, message }); + const abort = (message: string) => + settle(RunOutcome.Aborted, { code: ErrorCodes.AgentAbort, message }); + const cancelled = () => abort('Run cancelled by the caller.'); + + if (signal.aborted) return cancelled(); + + analytics.wizardCapture('agent started', { + integration: run.integrationLabel, + program_id: programId, + skill_id: run.skillId ?? null, + }); + + /** Await a caller capability; a caller abort during the wait wins over its answer. */ + const park = async (work: Promise): Promise => { + const value = await work; + signal.throwIfAborted(); + return value; + }; + + let credentials = input.credentials; + const flags = { ...DEFAULT_FLAGS, ...input.flags }; + let flagSnapshot: WizardFlagSnapshot = { + flags: { ...input.wizardFlags }, + payloads: { ...input.wizardFlagPayloads }, + }; + // Everything before the agent starts: a rejection fails the run, unless the caller aborted. + try { + if (!credentials && options.credentials) { + credentials = await park( + options.credentials.resolve(programId, { signal }), + ); + } + if (!credentials) + return fail(`Credentials are required to run ${programId}.`); + store.setAuthenticated({ + credentials: credentials.posthog, + apiProject: credentials.project, + apiUser: credentials.apiUser, + }); + // Identify before flags are evaluated, so flags can target the user. + if (credentials.apiUser) analytics.identifyUser(credentials.apiUser); + analytics.setGroups( + groupsFromUser(credentials.apiUser, credentials.posthog.host.apiHost), + ); + if (!store.readData().aiSdkStampReported) { + store.setAiSdkStampReported(); + stampAiSdkDetected({ + apiUser: credentials.apiUser, + discoveredFeatures: input.discoveredFeatures ?? [], + warehouseSources: input.warehouseSources ?? [], + mayReportScanResults: input.mayReportScanResults ?? false, + }); + } + + if ( + program.requiresAi !== false && + !flags.ci && + !flags.signup && + credentials.apiUser?.organization?.is_ai_data_processing_approved !== true + ) { + if (!options.awaitAiApproval) { + return fail( + 'AI processing approval is required before this program can run.', + ); + } + const approved = await park( + options.awaitAiApproval({ programId, signal }), + ); + if (!approved) return abort('AI processing approval declined.'); + } + + const gates = program.postAuthGates ?? []; + if (gates.length > 0 && options.awaitPostAuthGates) { + await park(options.awaitPostAuthGates({ programId, gates, signal })); + } + + if (!input.wizardFlags && options.featureFlags) { + flagSnapshot = await park(options.featureFlags()); + } + + // Freshness is measured after every park above, right before the agent mints. + // Every rotation of this login, before or during the run, lands in data. + const login = credentials; + configureOAuthSession(login.posthog, { + rotate: (held) => rotateCredentials(held, input.host?.baseUrl), + onRefreshed: (refreshed) => + store.setAuthenticated({ + credentials: refreshed, + apiProject: login.project, + apiUser: login.apiUser, + }), + }); + // Not parked: a rotation spends the old refresh token, so the new one is kept before a cancel returns. + const posthog = (await oauthCredentials()) ?? login.posthog; + credentials = { ...login, posthog }; + signal.throwIfAborted(); + } catch (error) { + if (signal.aborted) return cancelled(); + return fail(error instanceof Error ? error.message : String(error)); + } + const wizardFlags = { ...flagSnapshot.flags }; + const wizardFlagPayloads = { ...flagSnapshot.payloads }; + + // Resolve which sequence and harness run the program (CLI → PostHog flag → + // per-program binding → default) and tag both axes onto analytics. + const switchboard: SwitchboardCtx = { + program: programId, + composed: input.composed ?? false, + flags: wizardFlags, + flagPayloads: wizardFlagPayloads, + cliHarness: input.overrides?.harness, + cliSequence: input.overrides?.sequence, + cliModel: input.overrides?.model, + }; + const binding = resolveBinding(switchboard); + analytics.setTag('sequence', binding.sequence); + analytics.setTag('harness', binding.harness); + captureSwitchboardDecision(switchboard, binding); + store.setBinding(binding); + + const wizardMetadata = { + ...buildRunTags({ + programId, + integration: run.integrationLabel, + runId: analytics.runId, + build: analytics.build, + skillId: run.skillId, + }), + SEQUENCE: binding.sequence, + HARNESS: binding.harness, + }; + artifacts.reportFile = path.resolve(installDir, run.reportFile); + const adapter = store.beginRun(runId, options.onProgress); + + const result = await runAgent( + { + programId, + run, + composed: input.composed ?? false, + binding, + switchboard, + skillsBaseUrl: getSkillsBaseUrl(), + wizardFlags, + wizardFlagPayloads, + wizardMetadata, + allowedTools: program.allowedTools, + disallowedTools: program.disallowedTools, + agentFlow: program.agentFlow, + excludedTaskTypes: program.excludedTaskTypes, + seedTasks: input.seedTasks, + hooks: input.hooks, + }, + { + installDir, + credentials: credentials.posthog, + project: credentials.project, + apiUser: credentials.apiUser, + skillId: input.skillId ?? run.skillId ?? run.integrationLabel, + integration: input.integration, + frameworkDocsUrl: input.frameworkDocsUrl, + flags, + host: { ...input.host }, + }, + { + interaction: options.interaction, + onProgress: (event) => adapter.onProgress(event), + signal: options.signal, + }, + ); + adapter.finish(result); + return settle( + result.outcome, + result.outcome === RunOutcome.Success ? undefined : result.failure, + ); +} + +/** + * One event + one log line per run: what entered the switchboard, which + * precedence rung decided each axis, and the final pick. + */ +function captureSwitchboardDecision( + ctx: SwitchboardCtx, + binding: ProgramBinding, +): void { + const trace = ctx.trace ?? {}; + // Unpinned orchestrator runs choose a model per task from the context-mill agent prompts; the orchestrator logs that map once the prompts load. + const perTaskModel = + binding.sequence === Sequence.orchestrator && trace.model === 'binding'; + const model = perTaskModel ? 'chosen-per-task' : binding.model; + const modelSource = perTaskModel ? 'agent-prompts' : trace.model; + analytics.wizardCapture('switchboard resolved', { + program: ctx.program, + flag_self_driving_use_pi_harness: + ctx.flags[WIZARD_SELF_DRIVING_USE_PI_HARNESS_FLAG_KEY], + flag_self_driving_pi_payload: JSON.stringify( + ctx.flagPayloads?.[WIZARD_SELF_DRIVING_USE_PI_HARNESS_FLAG_KEY] ?? null, + ), + flag_orchestrator: ctx.flags[WIZARD_ORCHESTRATOR_FLAG_KEY], + cli_harness: ctx.cliHarness, + cli_sequence: ctx.cliSequence, + cli_model: ctx.cliModel, + harness_source: trace.harness, + model_source: modelSource, + sequence_source: trace.sequence, + harness: binding.harness, + model, + thinking_level: binding.thinkingLevel, + sequence: binding.sequence, + }); + logToFile( + `[switchboard] decision: program=${ctx.program}` + + ` in(orchestrator=${ctx.flags[WIZARD_ORCHESTRATOR_FLAG_KEY] ?? '-'},` + + ` cli=${ctx.cliHarness ?? '-'}/${ctx.cliSequence ?? '-'}/${ + ctx.cliModel ?? '-' + })` + + ` → harness=${binding.harness} (${trace.harness ?? '?'})` + + ` model=${model} (${modelSource ?? '?'})` + + ` sequence=${binding.sequence} (${trace.sequence ?? '?'})`, + ); }