diff --git a/src/lib/agent/__tests__/variant-gating.test.ts b/src/lib/agent/__tests__/variant-gating.test.ts new file mode 100644 index 00000000..699bd096 --- /dev/null +++ b/src/lib/agent/__tests__/variant-gating.test.ts @@ -0,0 +1,36 @@ +import { + buildWizardMetadata, + isOrchestratorEnabled, +} from '@lib/agent/agent-interface'; + +describe('isOrchestratorEnabled', () => { + it('is true only when the wizard-orchestrator flag is true', () => { + expect(isOrchestratorEnabled({ 'wizard-orchestrator': 'true' })).toBe(true); + }); + + it('is false when the flag is false, another flag, or absent', () => { + expect(isOrchestratorEnabled({ 'wizard-orchestrator': 'false' })).toBe( + false, + ); + expect(isOrchestratorEnabled({ 'wizard-variant': 'orchestrator' })).toBe( + false, + ); + expect(isOrchestratorEnabled({})).toBe(false); + expect(isOrchestratorEnabled()).toBe(false); + }); +}); + +describe('buildWizardMetadata', () => { + it('selects a known variant header from the flag', () => { + expect(buildWizardMetadata({ 'wizard-variant': 'subagents' })).toEqual({ + VARIANT: 'subagents', + }); + }); + + it('falls back to the base variant for unknown or missing flags', () => { + expect(buildWizardMetadata({ 'wizard-variant': 'nope' })).toEqual({ + VARIANT: 'base', + }); + expect(buildWizardMetadata({})).toEqual({ VARIANT: 'base' }); + }); +}); diff --git a/src/lib/agent/agent-interface.ts b/src/lib/agent/agent-interface.ts index 8c3be73e..35c45d03 100644 --- a/src/lib/agent/agent-interface.ts +++ b/src/lib/agent/agent-interface.ts @@ -14,6 +14,7 @@ import { POSTHOG_PROPERTY_HEADER_PREFIX, WIZARD_VARIANT_FLAG_KEY, WIZARD_VARIANTS, + WIZARD_ORCHESTRATOR_FLAG_KEY, WIZARD_USER_AGENT, } from '@lib/constants'; import { @@ -363,6 +364,17 @@ export function buildWizardMetadata( return { ...variant }; } +/** + * Whether this run uses the experimental task-queue orchestrator. Gated by the + * boolean `wizard-orchestrator` feature flag, targeted to the user in the wizard's + * analytics project. + */ +export function isOrchestratorEnabled( + flags: Record = {}, +): boolean { + return flags[WIZARD_ORCHESTRATOR_FLAG_KEY] === 'true'; +} + /** * Build env for the SDK subprocess: process.env plus ANTHROPIC_CUSTOM_HEADERS, which always * includes `x-posthog-use-bedrock-fallback: true` so the LLM gateway falls back to Bedrock on diff --git a/src/lib/agent/agent-runner.ts b/src/lib/agent/agent-runner.ts index 78bb01e9..a6169460 100644 --- a/src/lib/agent/agent-runner.ts +++ b/src/lib/agent/agent-runner.ts @@ -9,9 +9,11 @@ * - What MCP servers and package manager detector to use * - What happens after the agent completes * - * The pipeline itself is fixed: - * init → health check → settings → OAuth → [skill install] → - * agent init → prompt → run → errors → [postRun] → outro + * The pipeline runs a shared bootstrap (logging, health check, settings, OAuth, + * flags, MCP url), then forks. The `orchestrator` variant routes to the + * experimental task-queue runner. Every other variant runs the fixed linear + * pipeline: + * [skill install] → agent init → prompt → run → errors → [postRun] → outro */ import { @@ -51,7 +53,7 @@ import { getSkillsBaseUrl } from '@lib/constants'; import { runtimeEnv } from '@env'; import { installSkillById, type InstallSkillResult } from '@lib/wizard-tools'; import { createWizardAskBridge } from '@lib/wizard-ask-bridge'; -import type { WizardRunOptions } from '@utils/types'; +import type { WizardRunOptions, CloudRegion } from '@utils/types'; import type { ProgramConfig } from '@lib/programs/program-step'; import { assemblePrompt, type PromptContext } from './agent-prompt'; @@ -106,7 +108,7 @@ export interface ProgramRun { buildOutroData?: ( session: WizardSession, credentials: Credentials, - cloudRegion: import('@utils/types').CloudRegion | undefined, + cloudRegion: CloudRegion | undefined, ) => WizardSession['outroData']; /** * Per-run cap on `wizard_ask` invocations. Defaults to 10. The 4th call @@ -115,6 +117,23 @@ export interface ProgramRun { maxQuestions?: number; } +/** + * Result of the shared bootstrap, consumed by both the linear and the + * orchestrator arm. Credentials, role, and user are already applied to the + * session by `bootstrapProgram`; this carries the values both arms still need. + */ +export interface BootstrapResult { + skillsBaseUrl: string; + projectApiKey: Credentials['projectApiKey']; + host: Credentials['host']; + accessToken: Credentials['accessToken']; + projectId: Credentials['projectId']; + cloudRegion: CloudRegion; + mcpUrl: string; + wizardFlags: Record; + wizardMetadata: Record; +} + // ── Helpers ────────────────────────────────────────────────────────── /** @@ -170,16 +189,31 @@ export async function runAgent( /** * Run a program's agent pipeline. * - * This is the single execution path for all programs — both skill-based - * (revenue analytics) and framework-based (core integration). The - * `ProgramRun` controls what varies between them; `programConfig` carries - * the program-level static metadata (tool allow/disallow lists, etc.). + * Runs the shared bootstrap, then forks on the `wizard-variant` flag. The + * `orchestrator` variant routes to the experimental task-queue runner; every + * other variant runs the linear pipeline. */ export async function runProgram( session: WizardSession, config: ProgramRun, programConfig: ProgramConfig, ): Promise { + const boot = await bootstrapProgram(session, config, programConfig); + + return runLinearProgram(session, config, programConfig, boot); +} + +/** + * Shared setup for both arms: logging, health check, settings conflicts, OAuth + * and credentials, then the feature flags, variant metadata, and MCP url. Sets + * `session.credentials`, role, and user as a side effect. Returns the values the + * arms still need. + */ +async function bootstrapProgram( + session: WizardSession, + config: ProgramRun, + programConfig: ProgramConfig, +): Promise { // 1. Init logging + debug initLogFile(); session.skillId = config.skillId ?? config.integrationLabel; @@ -283,6 +317,55 @@ export async function runProgram( analytics.setGroups(groupsFromUser(user, host)); + // Feature flags, variant metadata, and MCP url. Both arms need these, and the + // fork decision reads the flags. + const wizardFlags = await analytics.getAllFlagsForWizard(); + const wizardMetadata = buildWizardMetadata(wizardFlags); + + const mcpUrl = session.localMcp + ? 'http://localhost:8787/mcp' + : runtimeEnv('MCP_URL') || + (cloudRegion === 'eu' + ? 'https://mcp-eu.posthog.com/mcp' + : 'https://mcp.posthog.com/mcp'); + + return { + skillsBaseUrl, + projectApiKey, + host, + accessToken, + projectId, + cloudRegion, + mcpUrl, + wizardFlags, + wizardMetadata, + }; +} + +/** + * The linear pipeline. Single execution path for all non-orchestrator programs, + * both skill-based (revenue analytics) and framework-based (core integration). + * The `ProgramRun` controls what varies between them; `programConfig` carries the + * program-level static metadata (tool allow/disallow lists, etc.). + */ +async function runLinearProgram( + session: WizardSession, + config: ProgramRun, + programConfig: ProgramConfig, + boot: BootstrapResult, +): Promise { + const { + skillsBaseUrl, + projectApiKey, + host, + accessToken, + projectId, + cloudRegion, + mcpUrl, + wizardFlags, + wizardMetadata, + } = boot; + // 5. Skill install (if skillId provided) let skillPath: string | undefined; if (config.skillId) { @@ -302,15 +385,6 @@ export async function runProgram( // 6. Initialize agent const spinner = getUI().spinner(); - const wizardFlags = await analytics.getAllFlagsForWizard(); - const wizardMetadata = buildWizardMetadata(wizardFlags); - - const mcpUrl = session.localMcp - ? 'http://localhost:8787/mcp' - : runtimeEnv('MCP_URL') || - (cloudRegion === 'eu' - ? 'https://mcp-eu.posthog.com/mcp' - : 'https://mcp.posthog.com/mcp'); const restoreSettings = () => restoreClaudeSettings(session.installDir); getUI().onEnterScreen('outro', restoreSettings); diff --git a/src/lib/constants.ts b/src/lib/constants.ts index bb7f3580..9f7357c3 100644 --- a/src/lib/constants.ts +++ b/src/lib/constants.ts @@ -141,6 +141,8 @@ export const WIZARD_INTERACTION_EVENT_NAME = 'wizard interaction'; export const WIZARD_REMARK_EVENT_NAME = 'wizard remark'; /** Feature flag key whose value selects a variant from WIZARD_VARIANTS. */ export const WIZARD_VARIANT_FLAG_KEY = 'wizard-variant'; +/** Boolean feature flag that routes a run to the experimental orchestrator runner. */ +export const WIZARD_ORCHESTRATOR_FLAG_KEY = 'wizard-orchestrator'; /** Feature flag key that gates the intro-screen "Tools" menu. */ export const WIZARD_TOOLS_MENU_FLAG_KEY = 'wizard-tools-menu'; /** Variant key -> metadata for wizard run (VARIANT flag selects which entry to use). */