From 70cd200ddd6e73400e54e73125c6e5498a0a10fa Mon Sep 17 00:00:00 2001 From: "Vincent (Wen Yu) Ge" Date: Mon, 6 Apr 2026 10:57:27 -0400 Subject: [PATCH 1/2] Use a queue and break up query --- src/lib/__tests__/agent-interface.test.ts | 61 +++++ src/lib/__tests__/agent-runner.test.ts | 19 ++ src/lib/__tests__/workflow-queue.test.ts | 93 ++++++++ src/lib/agent-interface.ts | 107 +++++++-- src/lib/agent-runner.ts | 262 ++++++++++++++++++---- src/lib/workflow-queue.ts | 75 +++++++ src/lib/yara-hooks.ts | 16 +- 7 files changed, 566 insertions(+), 67 deletions(-) create mode 100644 src/lib/__tests__/agent-runner.test.ts create mode 100644 src/lib/__tests__/workflow-queue.test.ts create mode 100644 src/lib/workflow-queue.ts diff --git a/src/lib/__tests__/agent-interface.test.ts b/src/lib/__tests__/agent-interface.test.ts index e1ba0d83..4fd05ad1 100644 --- a/src/lib/__tests__/agent-interface.test.ts +++ b/src/lib/__tests__/agent-interface.test.ts @@ -279,6 +279,61 @@ describe('runAgent', () => { // ui.log.error should NOT have been called (errors suppressed for user) expect(mockUIInstance.log.error).not.toHaveBeenCalled(); }); + + it('passes resume to the SDK and returns the active session id', async () => { + function* mockGeneratorWithSessionId() { + yield { + type: 'system', + subtype: 'init', + model: 'claude-opus-4-5-20251101', + tools: [], + mcp_servers: [], + session_id: 'session-123', + }; + + yield { + type: 'assistant', + session_id: 'session-123', + message: { + content: [{ type: 'text', text: '[STATUS] Working through queue' }], + }, + }; + + yield { + type: 'result', + subtype: 'success', + is_error: false, + session_id: 'session-123', + result: 'Queued step complete', + }; + } + + mockQuery.mockReturnValue(mockGeneratorWithSessionId()); + + const result = await runAgent( + defaultAgentConfig, + 'queued prompt', + defaultOptions, + mockSpinner as unknown as SpinnerHandle, + { + successMessage: 'Queued success', + resumeSessionId: 'session-123', + requestRemark: false, + captureOutputText: true, + captureSessionId: true, + }, + ); + + expect(mockQuery).toHaveBeenCalledWith( + expect.objectContaining({ + options: expect.objectContaining({ + resume: 'session-123', + }), + }), + ); + expect(result.sessionId).toBe('session-123'); + expect(result.outputText).toContain('Queued step complete'); + }); }); }); @@ -382,4 +437,10 @@ describe('createStopHook', () => { expect(first).toHaveProperty('decision', 'block'); expect((first as { reason: string }).reason).toContain('WIZARD-REMARK'); }); + + it('can skip the remark collection phase for intermediate queued steps', () => { + const hook = createStopHook([], [], { requestRemark: false }); + + expect(hook(hookInput)).toEqual({}); + }); }); diff --git a/src/lib/__tests__/agent-runner.test.ts b/src/lib/__tests__/agent-runner.test.ts new file mode 100644 index 00000000..bba9dbdf --- /dev/null +++ b/src/lib/__tests__/agent-runner.test.ts @@ -0,0 +1,19 @@ +import { extractInstalledSkillId } from '../agent-runner'; + +describe('extractInstalledSkillId', () => { + it('extracts the installed skill id from bootstrap output', () => { + const output = ` +Bootstrap complete. +[WIZARD-SKILL-ID] integration-nextjs-app-router +Waiting for next queued step. +`; + + expect(extractInstalledSkillId(output)).toBe( + 'integration-nextjs-app-router', + ); + }); + + it('returns null when the marker is missing', () => { + expect(extractInstalledSkillId('No marker present')).toBeNull(); + }); +}); diff --git a/src/lib/__tests__/workflow-queue.test.ts b/src/lib/__tests__/workflow-queue.test.ts new file mode 100644 index 00000000..ab0e62f9 --- /dev/null +++ b/src/lib/__tests__/workflow-queue.test.ts @@ -0,0 +1,93 @@ +import { + WizardWorkflowQueue, + createInitialWizardWorkflowQueue, + type WorkflowStepSeed, +} from '../workflow-queue'; + +const BASIC_INTEGRATION_STEPS: WorkflowStepSeed[] = [ + { stepId: '1.0-begin', referenceFilename: 'basic-integration-1.0-begin.md' }, + { stepId: '1.1-edit', referenceFilename: 'basic-integration-1.1-edit.md' }, + { + stepId: '1.2-revise', + referenceFilename: 'basic-integration-1.2-revise.md', + }, + { + stepId: '1.3-conclude', + referenceFilename: 'basic-integration-1.3-conclude.md', + }, +]; + +describe('WizardWorkflowQueue', () => { + it('seeds a queue from workflow steps in the expected order', () => { + const queue = createInitialWizardWorkflowQueue(BASIC_INTEGRATION_STEPS); + + expect(queue.toArray()).toEqual([ + { id: 'bootstrap', kind: 'bootstrap' }, + { + id: 'workflow:1.0-begin', + kind: 'workflow', + referenceFilename: 'basic-integration-1.0-begin.md', + }, + { + id: 'workflow:1.1-edit', + kind: 'workflow', + referenceFilename: 'basic-integration-1.1-edit.md', + }, + { + id: 'workflow:1.2-revise', + kind: 'workflow', + referenceFilename: 'basic-integration-1.2-revise.md', + }, + { + id: 'workflow:1.3-conclude', + kind: 'workflow', + referenceFilename: 'basic-integration-1.3-conclude.md', + }, + { id: 'env-vars', kind: 'env-vars' }, + ]); + }); + + it('builds a queue from arbitrary steps, not just basic-integration', () => { + const customSteps: WorkflowStepSeed[] = [ + { stepId: 'setup', referenceFilename: 'feature-flags-setup.md' }, + { stepId: 'verify', referenceFilename: 'feature-flags-verify.md' }, + ]; + const queue = createInitialWizardWorkflowQueue(customSteps); + + expect(queue.toArray()).toEqual([ + { id: 'bootstrap', kind: 'bootstrap' }, + { + id: 'workflow:setup', + kind: 'workflow', + referenceFilename: 'feature-flags-setup.md', + }, + { + id: 'workflow:verify', + kind: 'workflow', + referenceFilename: 'feature-flags-verify.md', + }, + { id: 'env-vars', kind: 'env-vars' }, + ]); + }); + + it('supports enqueue and dequeue operations', () => { + const queue = new WizardWorkflowQueue(); + + queue.enqueue({ id: 'bootstrap', kind: 'bootstrap' }); + queue.enqueue({ + id: 'workflow:1.0-begin', + kind: 'workflow', + referenceFilename: 'basic-integration-1.0-begin.md', + }); + + expect(queue.peek()).toEqual({ id: 'bootstrap', kind: 'bootstrap' }); + expect(queue.dequeue()).toEqual({ id: 'bootstrap', kind: 'bootstrap' }); + expect(queue).toHaveLength(1); + expect(queue.dequeue()).toEqual({ + id: 'workflow:1.0-begin', + kind: 'workflow', + referenceFilename: 'basic-integration-1.0-begin.md', + }); + expect(queue).toHaveLength(0); + }); +}); diff --git a/src/lib/agent-interface.ts b/src/lib/agent-interface.ts index 1dca0cce..84d1249f 100644 --- a/src/lib/agent-interface.ts +++ b/src/lib/agent-interface.ts @@ -71,6 +71,10 @@ export const AgentSignals = { WIZARD_REMARK: '[WIZARD-REMARK]', /** Signal prefix for benchmark logging */ BENCHMARK: '[BENCHMARK]', + /** Signal emitted by YARA scanner on critical security violation */ + YARA_CRITICAL: '[YARA CRITICAL]', + /** Signal emitted by YARA scanner on scanner error */ + YARA_SCANNER_ERROR: '[YARA] Scanner error', } as const; export type AgentSignal = (typeof AgentSignals)[keyof typeof AgentSignals]; @@ -295,9 +299,13 @@ export type StopHookResult = export function createStopHook( featureQueue: readonly AdditionalFeature[], collectedText?: string[], + options?: { + requestRemark?: boolean; + }, ): (input: { stop_hook_active: boolean }) => StopHookResult { let featureIndex = 0; let remarkRequested = false; + const requestRemark = options?.requestRemark ?? true; return (input: { stop_hook_active: boolean }): StopHookResult => { logToFile('Stop hook triggered', { @@ -326,7 +334,7 @@ export function createStopHook( } // Phase 2: collect remark (once) - if (!remarkRequested) { + if (requestRemark && !remarkRequested) { remarkRequested = true; logToFile('Stop hook: requesting reflection'); return { @@ -692,8 +700,8 @@ function checkYaraViolation( spinner: SpinnerHandle, ): { error: AgentErrorType } | null { if ( - outputText.includes('[YARA CRITICAL]') || - outputText.includes('[YARA] Scanner error') + outputText.includes(AgentSignals.YARA_CRITICAL) || + outputText.includes(AgentSignals.YARA_SCANNER_ERROR) ) { logToFile('Agent error: YARA_VIOLATION'); spinner.stop('Security violation detected'); @@ -719,16 +727,31 @@ export async function runAgent( successMessage?: string; errorMessage?: string; additionalFeatureQueue?: readonly AdditionalFeature[]; + resumeSessionId?: string; + requestRemark?: boolean; + captureOutputText?: boolean; + captureSessionId?: boolean; + finalizeMiddleware?: boolean; }, middleware?: { onMessage(message: any): void; finalize(resultMessage: any, totalDurationMs: number): any; }, -): Promise<{ error?: AgentErrorType; message?: string }> { +): Promise<{ + error?: AgentErrorType; + message?: string; + sessionId?: string; + outputText?: string; +}> { const { spinnerMessage = 'Customizing your PostHog setup...', successMessage = 'PostHog integration complete', errorMessage = 'Integration failed', + resumeSessionId, + requestRemark = true, + captureOutputText = false, + captureSessionId = false, + finalizeMiddleware = true, } = config ?? {}; const { query } = await getSDKModule(); @@ -742,6 +765,7 @@ export async function runAgent( const startTime = Date.now(); const collectedText: string[] = []; + let sessionId: string | undefined; // Track if we received a successful result (before any cleanup errors) let receivedSuccessResult = false; let loggedInitialContext = false; @@ -760,7 +784,7 @@ export async function runAgent( const createPromptStream = async function* () { yield { type: 'user', - session_id: '', + session_id: resumeSessionId ?? '', message: { role: 'user', content: prompt }, parent_tool_use_id: null, }; @@ -770,7 +794,12 @@ export async function runAgent( // Helper to handle successful completion (used in normal path and race condition recovery) const completeWithSuccess = ( suppressedError?: Error, - ): { error?: AgentErrorType; message?: string } => { + ): { + error?: AgentErrorType; + message?: string; + sessionId?: string; + outputText?: string; + } => { const durationMs = Date.now() - startTime; const durationSeconds = Math.round(durationMs / 1000); @@ -804,13 +833,18 @@ export async function runAgent( duration_ms: durationMs, duration_seconds: durationSeconds, }); - try { - middleware?.finalize(lastResultMessage, durationMs); - } catch (e) { - logToFile(`${AgentSignals.BENCHMARK} Middleware finalize error:`, e); + if (finalizeMiddleware) { + try { + middleware?.finalize(lastResultMessage, durationMs); + } catch (e) { + logToFile(`${AgentSignals.BENCHMARK} Middleware finalize error:`, e); + } } spinner.stop(successMessage); - return {}; + return { + ...(captureSessionId && sessionId ? { sessionId } : {}), + ...(captureOutputText ? { outputText } : {}), + }; }; // Event plan file watcher — cleaned up in finally block @@ -840,6 +874,7 @@ export async function runAgent( const response = query({ prompt: createPromptStream(), options: { + resume: resumeSessionId, model: agentConfig.model, cwd: agentConfig.workingDirectory, permissionMode: 'acceptEdits', @@ -928,6 +963,7 @@ export async function runAgent( createStopHook( config?.additionalFeatureQueue ?? [], collectedText, + { requestRemark }, ), ], timeout: 30, @@ -979,6 +1015,15 @@ export async function runAgent( // Process the async generator for await (const message of response) { + if ( + message && + typeof message === 'object' && + 'session_id' in message && + typeof message.session_id === 'string' + ) { + sessionId = message.session_id; + } + // Log initial context size on the first assistant response so we can // detect sudden shifts in starting context (e.g. MCP schema bloat). if (!loggedInitialContext && message.type === 'assistant') { @@ -1060,13 +1105,21 @@ export async function runAgent( if (outputText.includes(AgentSignals.ERROR_MCP_MISSING)) { logToFile('Agent error: MCP_MISSING'); spinner.stop('Agent could not access PostHog MCP'); - return { error: AgentErrorType.MCP_MISSING }; + return { + error: AgentErrorType.MCP_MISSING, + sessionId, + outputText, + }; } if (outputText.includes(AgentSignals.ERROR_RESOURCE_MISSING)) { logToFile('Agent error: RESOURCE_MISSING'); spinner.stop('Agent could not access setup resource'); - return { error: AgentErrorType.RESOURCE_MISSING }; + return { + error: AgentErrorType.RESOURCE_MISSING, + sessionId, + outputText, + }; } // Check for API errors (rate limits, etc.) @@ -1079,13 +1132,23 @@ export async function runAgent( if (outputText.includes('API Error: 429')) { logToFile('Agent error: RATE_LIMIT'); spinner.stop('Rate limit exceeded'); - return { error: AgentErrorType.RATE_LIMIT, message: apiErrorMessage }; + return { + error: AgentErrorType.RATE_LIMIT, + message: apiErrorMessage, + sessionId, + outputText, + }; } if (outputText.includes('API Error:')) { logToFile('Agent error: API_ERROR'); spinner.stop('API error occurred'); - return { error: AgentErrorType.API_ERROR, message: apiErrorMessage }; + return { + error: AgentErrorType.API_ERROR, + message: apiErrorMessage, + sessionId, + outputText, + }; } return completeWithSuccess(); @@ -1117,13 +1180,23 @@ export async function runAgent( if (outputText.includes('API Error: 429')) { logToFile('Agent error (caught): RATE_LIMIT'); spinner.stop('Rate limit exceeded'); - return { error: AgentErrorType.RATE_LIMIT, message: apiErrorMessage }; + return { + error: AgentErrorType.RATE_LIMIT, + message: apiErrorMessage, + sessionId, + outputText, + }; } if (outputText.includes('API Error:')) { logToFile('Agent error (caught): API_ERROR'); spinner.stop('API error occurred'); - return { error: AgentErrorType.API_ERROR, message: apiErrorMessage }; + return { + error: AgentErrorType.API_ERROR, + message: apiErrorMessage, + sessionId, + outputText, + }; } // No API error found, re-throw the original exception diff --git a/src/lib/agent-runner.ts b/src/lib/agent-runner.ts index d8686ae7..a2ff2828 100644 --- a/src/lib/agent-runner.ts +++ b/src/lib/agent-runner.ts @@ -1,6 +1,5 @@ import { DEFAULT_PACKAGE_INSTALLATION, - SPINNER_MESSAGE, type FrameworkConfig, } from './framework-config'; import { type WizardSession, OutroKind } from './wizard-session'; @@ -39,6 +38,13 @@ import { registerCleanup, } from '../utils/wizard-abort'; import { formatScanReport, writeScanReport } from './yara-hooks'; +import { + createInitialWizardWorkflowQueue, + type WizardWorkflowQueueItem, + type WorkflowStepSeed, +} from './workflow-queue'; + +const WIZARD_SKILL_ID_SIGNAL = '[WIZARD-SKILL-ID]'; /** * Build a WizardOptions bag from a WizardSession (for code that still expects WizardOptions). @@ -210,17 +216,13 @@ export async function runAgentWizard( analytics.setTag(key, value); }); - const integrationPrompt = buildIntegrationPrompt( - config, - { - frameworkVersion: frameworkVersion || 'latest', - typescript: typeScriptDetected, - projectApiKey, - host, - projectId, - }, - frameworkContext, - ); + const promptContext = { + frameworkVersion: frameworkVersion || 'latest', + typescript: typeScriptDetected, + projectApiKey, + host, + projectId, + }; // Initialize and run agent const spinner = getUI().spinner(); @@ -272,20 +274,78 @@ export async function runAgentWizard( ? createBenchmarkPipeline(spinner, sessionToOptions(session)) : undefined; - const agentResult = await runAgent( - agent, - integrationPrompt, - sessionToOptions(session), - spinner, + // TODO: S5 — seed from context-mill workflow manifest instead of this static list + const workflowSteps: WorkflowStepSeed[] = [ { - estimatedDurationMinutes: config.ui.estimatedDurationMinutes, - spinnerMessage: SPINNER_MESSAGE, - successMessage: config.ui.successMessage, - errorMessage: 'Integration failed', - additionalFeatureQueue: session.additionalFeatureQueue, + stepId: '1.0-begin', + referenceFilename: 'basic-integration-1.0-begin.md', }, - middleware, - ); + { stepId: '1.1-edit', referenceFilename: 'basic-integration-1.1-edit.md' }, + { + stepId: '1.2-revise', + referenceFilename: 'basic-integration-1.2-revise.md', + }, + { + stepId: '1.3-conclude', + referenceFilename: 'basic-integration-1.3-conclude.md', + }, + ]; + const queue = createInitialWizardWorkflowQueue(workflowSteps); + let queuedSessionId: string | undefined; + let installedSkillId: string | undefined; + let agentResult: Awaited> = {}; + + while (queue.length > 0) { + const queueItem = queue.dequeue()!; + const prompt = buildQueuedPrompt( + queueItem, + config, + promptContext, + frameworkContext, + installedSkillId, + ); + + agentResult = await runAgent( + agent, + prompt, + sessionToOptions(session), + spinner, + { + estimatedDurationMinutes: config.ui.estimatedDurationMinutes, + spinnerMessage: getQueueSpinnerMessage(queueItem), + successMessage: getQueueSuccessMessage(queueItem, config), + errorMessage: `Integration failed during ${queueItem.id}`, + additionalFeatureQueue: + queueItem.id === 'env-vars' ? session.additionalFeatureQueue : [], + resumeSessionId: queuedSessionId, + requestRemark: queueItem.id === 'env-vars', + captureOutputText: queueItem.kind === 'bootstrap', + captureSessionId: true, + finalizeMiddleware: queue.length === 0, + }, + middleware, + ); + + queuedSessionId = agentResult.sessionId ?? queuedSessionId; + + if (queueItem.kind === 'bootstrap') { + installedSkillId = + extractInstalledSkillId(agentResult.outputText ?? '') ?? undefined; + if (!installedSkillId) { + await wizardAbort({ + message: + 'The wizard could not determine which integration skill was installed during bootstrap.', + error: new WizardError( + 'Bootstrap step did not emit installed skill id', + ), + }); + } + } + + if (agentResult.error) { + break; + } + } // Handle error cases detected in agent output if (agentResult.error === AgentErrorType.MCP_MISSING) { @@ -397,7 +457,37 @@ export async function runAgentWizard( /** * Build the integration prompt for the agent. */ -function buildIntegrationPrompt( +function buildQueuedPrompt( + queueItem: WizardWorkflowQueueItem, + config: FrameworkConfig, + context: { + frameworkVersion: string; + typescript: boolean; + projectApiKey: string; + host: string; + projectId: number; + }, + frameworkContext: Record, + installedSkillId?: string, +): string { + if (queueItem.kind === 'bootstrap') { + return buildBootstrapPrompt(config, context, frameworkContext); + } + + if (queueItem.kind === 'workflow') { + if (!installedSkillId) { + throw new Error('Workflow step requires installed skill id'); + } + return buildWorkflowStepPrompt( + queueItem.referenceFilename, + installedSkillId, + ); + } + + return buildEnvVarPrompt(config, context); +} + +function buildProjectContextBlock( config: FrameworkConfig, context: { frameworkVersion: string; @@ -417,11 +507,7 @@ function buildIntegrationPrompt( ? '\n' + additionalLines.map((line) => `- ${line}`).join('\n') : ''; - return `You have access to the PostHog MCP server which provides skills to integrate PostHog into this ${ - config.metadata.name - } project. - -Project context: + return `Project context: - PostHog Project ID: ${context.projectId} - Framework: ${config.metadata.name} ${context.frameworkVersion} - TypeScript: ${context.typescript ? 'Yes' : 'No'} @@ -430,9 +516,25 @@ Project context: - Project type: ${config.prompts.projectTypeDetection} - Package installation: ${ config.prompts.packageInstallation ?? DEFAULT_PACKAGE_INSTALLATION - }${additionalContext} + }${additionalContext}`; +} -Instructions (follow these steps IN ORDER - do not skip or reorder): +function buildBootstrapPrompt( + config: FrameworkConfig, + context: { + frameworkVersion: string; + typescript: boolean; + projectApiKey: string; + host: string; + projectId: number; + }, + frameworkContext: Record, +): string { + return `You have access to the PostHog MCP server which provides skills to integrate PostHog into this ${ + config.metadata.name + } project. + +${buildProjectContextBlock(config, context, frameworkContext)} STEP 1: Call load_skill_menu (from the wizard-tools MCP server) to see available skills. If the tool fails, emit: ${ @@ -449,17 +551,95 @@ STEP 2: Call install_skill (from the wizard-tools MCP server) with the chosen sk STEP 3: Load the installed skill's SKILL.md file to understand what references are available. -STEP 4: Follow the skill's workflow files in sequence. Look for numbered workflow files in the references (e.g., files with patterns like "1.0-", "1.1-", "1.2-"). Start with the first one and proceed through each step until completion. Each workflow file will tell you what to do and which file comes next. Never directly write PostHog tokens directly to code files; always use environment variables. +STEP 4: When preparation is complete, emit exactly one line in this format: +${WIZARD_SKILL_ID_SIGNAL} -STEP 5: Set up environment variables for PostHog using the wizard-tools MCP server (this runs locally — secret values never leave the machine): - - Use check_env_keys to see which keys already exist in the project's .env file (e.g. .env.local or .env). - - Use set_env_values to create or update the PostHog public token and host, using the appropriate environment variable naming convention for ${ - config.metadata.name - }, which you'll find in example code. The tool will also ensure .gitignore coverage. Don't assume the presence of keys means the value is up to date. Write the correct value each time. - - Reference these environment variables in the code files you create instead of hardcoding the public token and host. +Important: +- Do NOT execute any of the workflow reference files yet. +- Do NOT set up environment variables yet. +- Stop after preparation is complete. +- Use the detect_package_manager tool (from the wizard-tools MCP server) to determine which package manager the project uses. Do not manually search for lockfiles or config files. Always install packages as a background task. Don't await completion; proceed with other work immediately after starting the installation. You must read a file immediately before attempting to write it, even if you have previously read it; failure to do so will cause a tool failure. -Important: Use the detect_package_manager tool (from the wizard-tools MCP server) to determine which package manager the project uses. Do not manually search for lockfiles or config files. Always install packages as a background task. Don't await completion; proceed with other work immediately after starting the installation. You must read a file immediately before attempting to write it, even if you have previously read it; failure to do so will cause a tool failure. +`; +} +function buildWorkflowStepPrompt( + referenceFilename: string, + installedSkillId: string, +): string { + return `Continue the existing conversation. -`; +Read and follow this workflow reference: +\`.claude/skills/${installedSkillId}/references/${referenceFilename}\` + +Important: +- Complete only this workflow step. +- Do NOT continue to any other workflow file. +- Do NOT set up environment variables yet. +- Stop when this step is complete.`; +} + +function buildEnvVarPrompt( + config: FrameworkConfig, + context: { + frameworkVersion: string; + typescript: boolean; + projectApiKey: string; + host: string; + projectId: number; + }, +): string { + return `Continue the existing conversation. + +Execute the final queued environment-variable setup step for this ${ + config.metadata.name + } project. + +${buildProjectContextBlock(config, context, {})} + +Set up environment variables for PostHog using the wizard-tools MCP server (this runs locally — secret values never leave the machine): +- Use check_env_keys to see which keys already exist in the project's .env file (e.g. .env.local or .env). +- Use set_env_values to create or update the PostHog public token and host, using the appropriate environment variable naming convention for ${ + config.metadata.name + }, which you'll find in example code. The tool will also ensure .gitignore coverage. Don't assume the presence of keys means the value is up to date. Write the correct value each time. +- Reference these environment variables in the code files you create instead of hardcoding the public token and host. + +Stop after the environment-variable setup step is complete.`; +} + +function getQueueSpinnerMessage(queueItem: WizardWorkflowQueueItem): string { + switch (queueItem.kind) { + case 'bootstrap': + return 'Preparing integration...'; + case 'workflow': + return `Running step ${queueItem.id.replace('workflow:', '')}...`; + case 'env-vars': + return 'Finalizing environment variables...'; + } +} + +function getQueueSuccessMessage( + queueItem: WizardWorkflowQueueItem, + config: FrameworkConfig, +): string { + switch (queueItem.kind) { + case 'bootstrap': + return 'Integration prepared'; + case 'workflow': + return `Step ${queueItem.id.replace('workflow:', '')} complete`; + case 'env-vars': + return config.ui.successMessage; + } +} + +export function extractInstalledSkillId(outputText: string): string | null { + const match = outputText.match( + new RegExp( + `${WIZARD_SKILL_ID_SIGNAL.replace( + /[.*+?^${}()|[\]\\]/g, + '\\$&', + )}\\s+([A-Za-z0-9._-]+)`, + ), + ); + return match?.[1] ?? null; } diff --git a/src/lib/workflow-queue.ts b/src/lib/workflow-queue.ts new file mode 100644 index 00000000..6b283742 --- /dev/null +++ b/src/lib/workflow-queue.ts @@ -0,0 +1,75 @@ +export type WizardWorkflowQueueItem = + | { + id: 'bootstrap'; + kind: 'bootstrap'; + } + | { + id: string; + kind: 'workflow'; + referenceFilename: string; + } + | { + id: 'env-vars'; + kind: 'env-vars'; + }; + +export class WizardWorkflowQueue { + private items: WizardWorkflowQueueItem[]; + + constructor(items: WizardWorkflowQueueItem[] = []) { + this.items = [...items]; + } + + enqueue(item: WizardWorkflowQueueItem): void { + this.items.push(item); + } + + dequeue(): WizardWorkflowQueueItem | undefined { + return this.items.shift(); + } + + peek(): WizardWorkflowQueueItem | undefined { + return this.items[0]; + } + + toArray(): WizardWorkflowQueueItem[] { + return [...this.items]; + } + + get length(): number { + return this.items.length; + } +} + +/** + * Describes a workflow step that can be seeded into the queue. + * Eventually this comes from a context-mill manifest; for now it's + * passed in by the caller so the queue itself stays generic. + */ +export interface WorkflowStepSeed { + /** Unique id for the step, e.g. "1.0-begin" */ + stepId: string; + /** Filename inside the skill's references/ dir, e.g. "basic-integration-1.0-begin.md" */ + referenceFilename: string; +} + +/** + * Build the initial queue from an ordered list of workflow steps. + * The queue is always: bootstrap → workflow steps → env-vars. + */ +export function createInitialWizardWorkflowQueue( + steps: WorkflowStepSeed[], +): WizardWorkflowQueue { + const items: WizardWorkflowQueueItem[] = [ + { id: 'bootstrap', kind: 'bootstrap' }, + ...steps.map( + (step): WizardWorkflowQueueItem => ({ + id: `workflow:${step.stepId}`, + kind: 'workflow', + referenceFilename: step.referenceFilename, + }), + ), + { id: 'env-vars', kind: 'env-vars' }, + ]; + return new WizardWorkflowQueue(items); +} diff --git a/src/lib/yara-hooks.ts b/src/lib/yara-hooks.ts index 1621e10e..61a4b4a7 100644 --- a/src/lib/yara-hooks.ts +++ b/src/lib/yara-hooks.ts @@ -19,6 +19,7 @@ import type { YaraMatch, ScanResult } from './yara-scanner'; import { logToFile } from '../utils/debug'; import { analytics } from '../utils/analytics'; import { isSkillInstallCommand } from './skill-install'; +import { AgentSignals } from './agent-interface'; // ─── Types ─────────────────────────────────────────────────────── // Using loose types to avoid tight coupling to SDK version. @@ -226,7 +227,7 @@ export function createPreToolUseYaraHooks(): HookCallbackMatcher[] { // Fail closed: block the command if scanning fails return Promise.resolve({ decision: 'block', - reason: '[YARA] Scanner error — command blocked as a precaution.', + reason: `${AgentSignals.YARA_SCANNER_ERROR} — command blocked as a precaution.`, }); } }, @@ -296,8 +297,7 @@ export function createPostToolUseYaraHooks(): HookCallbackMatcher[] { return Promise.resolve({ hookSpecificOutput: { hookEventName: 'PostToolUse', - additionalContext: - '[YARA] Scanner error — you MUST revert this change as a precaution.', + additionalContext: `${AgentSignals.YARA_SCANNER_ERROR} — you MUST revert this change as a precaution.`, }, }); } @@ -342,7 +342,7 @@ export function createPostToolUseYaraHooks(): HookCallbackMatcher[] { // Prompt injection: abort the session — context is poisoned return Promise.resolve({ stopReason: - `[YARA CRITICAL] ${match.rule.name}: Prompt injection detected in file content. ` + + `${AgentSignals.YARA_CRITICAL} ${match.rule.name}: Prompt injection detected in file content. ` + `Agent context is potentially poisoned. Session terminated for safety.`, }); } @@ -365,8 +365,7 @@ export function createPostToolUseYaraHooks(): HookCallbackMatcher[] { logToFile('[YARA] PostToolUse Read/Grep hook error:', error); // Fail closed: terminate session if scanning fails on read content return Promise.resolve({ - stopReason: - '[YARA] Scanner error while scanning read content — session terminated as a precaution.', + stopReason: `${AgentSignals.YARA_SCANNER_ERROR} while scanning read content — session terminated as a precaution.`, }); } }, @@ -419,15 +418,14 @@ export function createPostToolUseYaraHooks(): HookCallbackMatcher[] { return { stopReason: - `[YARA CRITICAL] Poisoned skill detected in ${skillDir}: ${match.rule.name}. ` + + `${AgentSignals.YARA_CRITICAL} Poisoned skill detected in ${skillDir}: ${match.rule.name}. ` + `The downloaded skill contains potential prompt injection. Session terminated for safety.`, }; } catch (error) { logToFile('[YARA] PostToolUse skill install hook error:', error); // Fail closed: terminate if skill scanning fails return { - stopReason: - '[YARA] Scanner error while scanning skill files — session terminated as a precaution.', + stopReason: `${AgentSignals.YARA_SCANNER_ERROR} while scanning skill files — session terminated as a precaution.`, }; } }, From cb21681ff3e997d5da71f38b0394c5eae703cf8a Mon Sep 17 00:00:00 2001 From: "Vincent (Wen Yu) Ge" Date: Mon, 6 Apr 2026 12:14:28 -0400 Subject: [PATCH 2/2] build a queue --- src/lib/__tests__/workflow-queue.test.ts | 89 ++++++++++++ src/lib/agent-runner.ts | 168 +++++++++++++---------- src/lib/workflow-queue.ts | 49 ++++++- 3 files changed, 232 insertions(+), 74 deletions(-) diff --git a/src/lib/__tests__/workflow-queue.test.ts b/src/lib/__tests__/workflow-queue.test.ts index ab0e62f9..d2f1d1c1 100644 --- a/src/lib/__tests__/workflow-queue.test.ts +++ b/src/lib/__tests__/workflow-queue.test.ts @@ -1,6 +1,8 @@ import { WizardWorkflowQueue, createInitialWizardWorkflowQueue, + createPostBootstrapQueue, + parseWorkflowStepsFromSkillMd, type WorkflowStepSeed, } from '../workflow-queue'; @@ -70,6 +72,22 @@ describe('WizardWorkflowQueue', () => { ]); }); + it('createPostBootstrapQueue omits bootstrap', () => { + const queue = createPostBootstrapQueue(BASIC_INTEGRATION_STEPS); + const items = queue.toArray(); + + expect(items[0]).toEqual({ + id: 'workflow:1.0-begin', + kind: 'workflow', + referenceFilename: 'basic-integration-1.0-begin.md', + }); + expect(items[items.length - 1]).toEqual({ + id: 'env-vars', + kind: 'env-vars', + }); + expect(items.find((i) => i.id === 'bootstrap')).toBeUndefined(); + }); + it('supports enqueue and dequeue operations', () => { const queue = new WizardWorkflowQueue(); @@ -91,3 +109,74 @@ describe('WizardWorkflowQueue', () => { expect(queue).toHaveLength(0); }); }); + +describe('parseWorkflowStepsFromSkillMd', () => { + it('parses workflow steps from SKILL.md frontmatter', () => { + const skillMd = `--- +name: integration-nextjs-app-router +description: PostHog integration for Next.js App Router applications +metadata: + author: PostHog + version: dev +workflow: + - step_id: 1.0-begin + reference: basic-integration-1.0-begin.md + title: PostHog Setup - Begin + next: + - basic-integration-1.1-edit.md + - step_id: 1.1-edit + reference: basic-integration-1.1-edit.md + title: PostHog Setup - Edit + next: + - basic-integration-1.2-revise.md + - step_id: 1.2-revise + reference: basic-integration-1.2-revise.md + title: PostHog Setup - Revise + next: + - basic-integration-1.3-conclude.md + - step_id: 1.3-conclude + reference: basic-integration-1.3-conclude.md + title: PostHog Setup - Conclusion + next: [] +--- + +# PostHog integration for Next.js App Router +`; + + const steps = parseWorkflowStepsFromSkillMd(skillMd); + + expect(steps).toEqual([ + { + stepId: '1.0-begin', + referenceFilename: 'basic-integration-1.0-begin.md', + }, + { + stepId: '1.1-edit', + referenceFilename: 'basic-integration-1.1-edit.md', + }, + { + stepId: '1.2-revise', + referenceFilename: 'basic-integration-1.2-revise.md', + }, + { + stepId: '1.3-conclude', + referenceFilename: 'basic-integration-1.3-conclude.md', + }, + ]); + }); + + it('returns empty array when no frontmatter', () => { + expect(parseWorkflowStepsFromSkillMd('# No frontmatter')).toEqual([]); + }); + + it('returns empty array when no workflow key', () => { + const skillMd = `--- +name: feature-flags-nextjs +description: docs only +--- + +# Feature flags +`; + expect(parseWorkflowStepsFromSkillMd(skillMd)).toEqual([]); + }); +}); diff --git a/src/lib/agent-runner.ts b/src/lib/agent-runner.ts index a2ff2828..4f0cb847 100644 --- a/src/lib/agent-runner.ts +++ b/src/lib/agent-runner.ts @@ -1,3 +1,5 @@ +import fs from 'fs'; +import path from 'path'; import { DEFAULT_PACKAGE_INSTALLATION, type FrameworkConfig, @@ -39,9 +41,9 @@ import { } from '../utils/wizard-abort'; import { formatScanReport, writeScanReport } from './yara-hooks'; import { - createInitialWizardWorkflowQueue, + createPostBootstrapQueue, + parseWorkflowStepsFromSkillMd, type WizardWorkflowQueueItem, - type WorkflowStepSeed, } from './workflow-queue'; const WIZARD_SKILL_ID_SIGNAL = '[WIZARD-SKILL-ID]'; @@ -274,76 +276,106 @@ export async function runAgentWizard( ? createBenchmarkPipeline(spinner, sessionToOptions(session)) : undefined; - // TODO: S5 — seed from context-mill workflow manifest instead of this static list - const workflowSteps: WorkflowStepSeed[] = [ - { - stepId: '1.0-begin', - referenceFilename: 'basic-integration-1.0-begin.md', - }, - { stepId: '1.1-edit', referenceFilename: 'basic-integration-1.1-edit.md' }, - { - stepId: '1.2-revise', - referenceFilename: 'basic-integration-1.2-revise.md', - }, + // ── Step 1: Bootstrap — install the skill and get its ID ── + + let agentResult = await runAgent( + agent, + buildBootstrapPrompt(config, promptContext, frameworkContext), + sessionToOptions(session), + spinner, { - stepId: '1.3-conclude', - referenceFilename: 'basic-integration-1.3-conclude.md', + estimatedDurationMinutes: config.ui.estimatedDurationMinutes, + spinnerMessage: 'Preparing integration...', + successMessage: 'Integration prepared', + errorMessage: 'Integration failed during bootstrap', + additionalFeatureQueue: [], + requestRemark: false, + captureOutputText: true, + captureSessionId: true, + finalizeMiddleware: false, }, - ]; - const queue = createInitialWizardWorkflowQueue(workflowSteps); - let queuedSessionId: string | undefined; - let installedSkillId: string | undefined; - let agentResult: Awaited> = {}; - - while (queue.length > 0) { - const queueItem = queue.dequeue()!; - const prompt = buildQueuedPrompt( - queueItem, - config, - promptContext, - frameworkContext, + middleware, + ); + + const queuedSessionId = agentResult.sessionId; + const installedSkillId = + extractInstalledSkillId(agentResult.outputText ?? '') ?? undefined; + + if (!installedSkillId) { + await wizardAbort({ + message: + 'The wizard could not determine which integration skill was installed during bootstrap.', + error: new WizardError('Bootstrap step did not emit installed skill id'), + }); + } + + // ── Step 2: Read SKILL.md and seed the queue from its frontmatter ── + + if (!agentResult.error && installedSkillId) { + const skillMdPath = path.join( + session.installDir, + '.claude', + 'skills', installedSkillId, + 'SKILL.md', ); + const skillMdContent = fs.readFileSync(skillMdPath, 'utf-8'); + const workflowSteps = parseWorkflowStepsFromSkillMd(skillMdContent); - agentResult = await runAgent( - agent, - prompt, - sessionToOptions(session), - spinner, - { - estimatedDurationMinutes: config.ui.estimatedDurationMinutes, - spinnerMessage: getQueueSpinnerMessage(queueItem), - successMessage: getQueueSuccessMessage(queueItem, config), - errorMessage: `Integration failed during ${queueItem.id}`, - additionalFeatureQueue: - queueItem.id === 'env-vars' ? session.additionalFeatureQueue : [], - resumeSessionId: queuedSessionId, - requestRemark: queueItem.id === 'env-vars', - captureOutputText: queueItem.kind === 'bootstrap', - captureSessionId: true, - finalizeMiddleware: queue.length === 0, - }, - middleware, + if (workflowSteps.length === 0) { + logToFile( + '[agent-runner] No workflow steps found in SKILL.md frontmatter, aborting', + ); + await wizardAbort({ + message: + 'The installed skill does not contain workflow steps in its metadata.', + error: new WizardError('No workflow steps in SKILL.md frontmatter'), + }); + } + + logToFile( + `[agent-runner] Seeded queue from SKILL.md: ${workflowSteps + .map((s) => s.stepId) + .join(', ')}`, ); - queuedSessionId = agentResult.sessionId ?? queuedSessionId; + // ── Step 3: Execute workflow steps + env-vars from the queue ── - if (queueItem.kind === 'bootstrap') { - installedSkillId = - extractInstalledSkillId(agentResult.outputText ?? '') ?? undefined; - if (!installedSkillId) { - await wizardAbort({ - message: - 'The wizard could not determine which integration skill was installed during bootstrap.', - error: new WizardError( - 'Bootstrap step did not emit installed skill id', - ), - }); - } - } + const queue = createPostBootstrapQueue(workflowSteps); + + while (queue.length > 0) { + const queueItem = queue.dequeue()!; + const prompt = buildQueuedPrompt( + queueItem, + config, + promptContext, + installedSkillId, + ); + + agentResult = await runAgent( + agent, + prompt, + sessionToOptions(session), + spinner, + { + estimatedDurationMinutes: config.ui.estimatedDurationMinutes, + spinnerMessage: getQueueSpinnerMessage(queueItem), + successMessage: getQueueSuccessMessage(queueItem, config), + errorMessage: `Integration failed during ${queueItem.id}`, + additionalFeatureQueue: + queueItem.id === 'env-vars' ? session.additionalFeatureQueue : [], + resumeSessionId: queuedSessionId, + requestRemark: queueItem.id === 'env-vars', + captureOutputText: false, + captureSessionId: false, + finalizeMiddleware: queue.length === 0, + }, + middleware, + ); - if (agentResult.error) { - break; + if (agentResult.error) { + break; + } } } @@ -467,17 +499,9 @@ function buildQueuedPrompt( host: string; projectId: number; }, - frameworkContext: Record, - installedSkillId?: string, + installedSkillId: string, ): string { - if (queueItem.kind === 'bootstrap') { - return buildBootstrapPrompt(config, context, frameworkContext); - } - if (queueItem.kind === 'workflow') { - if (!installedSkillId) { - throw new Error('Workflow step requires installed skill id'); - } return buildWorkflowStepPrompt( queueItem.referenceFilename, installedSkillId, diff --git a/src/lib/workflow-queue.ts b/src/lib/workflow-queue.ts index 6b283742..c8ab586e 100644 --- a/src/lib/workflow-queue.ts +++ b/src/lib/workflow-queue.ts @@ -43,8 +43,7 @@ export class WizardWorkflowQueue { /** * Describes a workflow step that can be seeded into the queue. - * Eventually this comes from a context-mill manifest; for now it's - * passed in by the caller so the queue itself stays generic. + * Parsed from SKILL.md frontmatter's `workflow` array. */ export interface WorkflowStepSeed { /** Unique id for the step, e.g. "1.0-begin" */ @@ -53,6 +52,32 @@ export interface WorkflowStepSeed { referenceFilename: string; } +/** + * Parse workflow steps from SKILL.md content. + * + * Extracts `step_id` and `reference` from the YAML frontmatter's + * `workflow` array. Uses simple regex — no YAML library needed + * since we control the output format in skill-generator. + */ +export function parseWorkflowStepsFromSkillMd( + skillMdContent: string, +): WorkflowStepSeed[] { + const fmMatch = skillMdContent.match(/^---\n([\s\S]*?)\n---/); + if (!fmMatch) return []; + const frontmatter = fmMatch[1]; + + const steps: WorkflowStepSeed[] = []; + const entryRegex = /step_id:\s*(.+)\n\s*reference:\s*(.+)/g; + let match; + while ((match = entryRegex.exec(frontmatter)) !== null) { + steps.push({ + stepId: match[1].trim(), + referenceFilename: match[2].trim(), + }); + } + return steps; +} + /** * Build the initial queue from an ordered list of workflow steps. * The queue is always: bootstrap → workflow steps → env-vars. @@ -73,3 +98,23 @@ export function createInitialWizardWorkflowQueue( ]; return new WizardWorkflowQueue(items); } + +/** + * Build a queue with only workflow steps + env-vars (no bootstrap). + * Used after bootstrap has already run and SKILL.md has been parsed. + */ +export function createPostBootstrapQueue( + steps: WorkflowStepSeed[], +): WizardWorkflowQueue { + const items: WizardWorkflowQueueItem[] = [ + ...steps.map( + (step): WizardWorkflowQueueItem => ({ + id: `workflow:${step.stepId}`, + kind: 'workflow', + referenceFilename: step.referenceFilename, + }), + ), + { id: 'env-vars', kind: 'env-vars' }, + ]; + return new WizardWorkflowQueue(items); +}