Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions src/lib/agent/__tests__/variant-gating.test.ts
Original file line number Diff line number Diff line change
@@ -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' });
});
});
12 changes: 12 additions & 0 deletions src/lib/agent/agent-interface.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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<string, string> = {},
): 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
Expand Down
110 changes: 92 additions & 18 deletions src/lib/agent/agent-runner.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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';
Expand Down Expand Up @@ -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
Expand All @@ -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<string, string>;
wizardMetadata: Record<string, string>;
}

// ── Helpers ──────────────────────────────────────────────────────────

/**
Expand Down Expand Up @@ -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<void> {
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<BootstrapResult> {
// 1. Init logging + debug
initLogFile();
session.skillId = config.skillId ?? config.integrationLabel;
Expand Down Expand Up @@ -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');
Comment on lines +329 to +330

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we still need this? isn't everything 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<void> {
const {
skillsBaseUrl,
projectApiKey,
host,
accessToken,
projectId,
cloudRegion,
mcpUrl,
wizardFlags,
wizardMetadata,
} = boot;

// 5. Skill install (if skillId provided)
let skillPath: string | undefined;
if (config.skillId) {
Expand All @@ -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);
Expand Down
2 changes: 2 additions & 0 deletions src/lib/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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). */
Expand Down
Loading