feat(ai-builder): Add per-stage model configuration for evaluations (no-changelog) (#24344)

Signed-off-by: Oleg Ivaniv <me@olegivaniv.com>

packages/@n8n/ai-workflow-builder.ee/evaluations/__tests__/cli.test.ts

@@ -14,6 +14,7 @@ import type { SimpleWorkflow } from '@/types/workflow';
 
 // Store mocks for dependencies
 const mockParseEvaluationArgs = jest.fn();
+const mockArgsToStageModels = jest.fn();
 const mockSetupTestEnvironment = jest.fn();
 const mockCreateAgent = jest.fn();
 const mockGenerateRunId = jest.fn();
@@ -30,6 +31,7 @@ const mockCreatePairwiseEvaluator = jest.fn();
 // Mock all external modules
 jest.mock('../cli/argument-parser', () => ({
 	parseEvaluationArgs: (): unknown => mockParseEvaluationArgs(),
+	argsToStageModels: (...args: unknown[]): unknown => mockArgsToStageModels(...args),
 	getDefaultDatasetName: (suite: unknown): unknown =>
 		suite === 'pairwise' ? 'notion-pairwise-workflows' : 'workflow-builder-canvas-prompts',
 	getDefaultExperimentName: (suite: unknown): unknown =>
@@ -98,9 +100,19 @@ function createMockArgs(overrides: Record<string, unknown> = {}) {
 
 /** Helper to create mock environment */
 function createMockEnvironment() {
+	const mockLlm = mock<BaseChatModel>();
 	return {
 		parsedNodeTypes: [] as INodeTypeDescription[],
-		llm: mock<BaseChatModel>(),
+		llms: {
+			default: mockLlm,
+			supervisor: mockLlm,
+			responder: mockLlm,
+			discovery: mockLlm,
+			builder: mockLlm,
+			configurator: mockLlm,
+			parameterUpdater: mockLlm,
+			judge: mockLlm,
+		},
 		lsClient: mock<Client>(),
 	};
 }
@@ -147,6 +159,7 @@ describe('CLI', () => {
 
 		// Setup default mocks
 		mockParseEvaluationArgs.mockReturnValue(createMockArgs());
+		mockArgsToStageModels.mockReturnValue({ default: 'claude-sonnet-4.5' });
 		mockSetupTestEnvironment.mockResolvedValue(createMockEnvironment());
 		mockCreateAgent.mockReturnValue(createMockAgentInstance());
 		mockGenerateRunId.mockReturnValue('test-run-id');

packages/@n8n/ai-workflow-builder.ee/evaluations/__tests__/runner-langsmith.test.ts

@@ -174,7 +174,8 @@ describe('Runner - LangSmith Mode', () => {
 			expect(isLangsmithTargetOutput(result)).toBe(true);
 			if (!isLangsmithTargetOutput(result)) throw new Error('Expected LangSmith target output');
 
-			expect(generateWorkflow).toHaveBeenCalledWith('Create a workflow');
+			// Callbacks are passed explicitly from the traceable wrapper (undefined in tests without traceable context)
+			expect(generateWorkflow).toHaveBeenCalledWith('Create a workflow', undefined);
 			expect(evaluator.evaluate).toHaveBeenCalledWith(
 				workflow,
 				expect.objectContaining({ prompt: 'Create a workflow' }),

packages/@n8n/ai-workflow-builder.ee/evaluations/cli/argument-parser.ts

@@ -1,9 +1,12 @@
 /* eslint-disable @typescript-eslint/naming-convention */
 import { z } from 'zod';
 
-import type { BuilderFeatureFlags } from '../../src/workflow-builder-agent.js';
+import { AVAILABLE_MODELS, DEFAULT_MODEL, type ModelId } from '@/llm-config';
+import type { BuilderFeatureFlags } from '@/workflow-builder-agent';
+
 import type { LangsmithExampleFilters } from '../harness/harness-types';
 import { DEFAULTS } from '../support/constants';
+import type { StageModels } from '../support/environment.js';
 
 export type EvaluationSuite = 'llm-judge' | 'pairwise' | 'programmatic' | 'similarity';
 export type EvaluationBackend = 'local' | 'langsmith';
@@ -32,10 +35,39 @@ export interface EvaluationArgs {
 	numJudges: number;
 
 	featureFlags?: BuilderFeatureFlags;
+
+	// Model configuration
+	/** Default model for all stages */
+	model: ModelId;
+	/** Model for LLM judge evaluation */
+	judgeModel?: ModelId;
+	/** Model for supervisor stage */
+	supervisorModel?: ModelId;
+	/** Model for responder stage */
+	responderModel?: ModelId;
+	/** Model for discovery stage */
+	discoveryModel?: ModelId;
+	/** Model for builder stage */
+	builderModel?: ModelId;
+	/** Model for configurator stage */
+	configuratorModel?: ModelId;
+	/** Model for parameter updater (within configurator) */
+	parameterUpdaterModel?: ModelId;
 }
 
 type CliValueKind = 'boolean' | 'string';
-type FlagGroup = 'input' | 'eval' | 'pairwise' | 'langsmith' | 'output' | 'feature' | 'advanced';
+type FlagGroup =
+	| 'input'
+	| 'eval'
+	| 'pairwise'
+	| 'langsmith'
+	| 'output'
+	| 'feature'
+	| 'model'
+	| 'advanced';
+
+// Model ID validation schema
+const modelIdSchema = z.enum(AVAILABLE_MODELS as [ModelId, ...ModelId[]]);
 
 const cliSchema = z
 	.object({
@@ -65,6 +97,16 @@ const cliSchema = z
 
 		langsmith: z.boolean().optional(),
 		templateExamples: z.boolean().default(false),
+
+		// Model configuration
+		model: modelIdSchema.default(DEFAULT_MODEL),
+		judgeModel: modelIdSchema.optional(),
+		supervisorModel: modelIdSchema.optional(),
+		responderModel: modelIdSchema.optional(),
+		discoveryModel: modelIdSchema.optional(),
+		builderModel: modelIdSchema.optional(),
+		configuratorModel: modelIdSchema.optional(),
+		parameterUpdaterModel: modelIdSchema.optional(),
 	})
 	.strict();
 
@@ -185,6 +227,56 @@ const FLAG_DEFS: Record<string, FlagDef> = {
 		desc: 'Enable template examples phase',
 	},
 
+	// Model configuration
+	'--model': {
+		key: 'model',
+		kind: 'string',
+		group: 'model',
+		desc: `Default model for all stages (default: ${DEFAULT_MODEL})`,
+	},
+	'--judge-model': {
+		key: 'judgeModel',
+		kind: 'string',
+		group: 'model',
+		desc: 'Model for LLM judge evaluation',
+	},
+	'--supervisor-model': {
+		key: 'supervisorModel',
+		kind: 'string',
+		group: 'model',
+		desc: 'Model for supervisor stage',
+	},
+	'--responder-model': {
+		key: 'responderModel',
+		kind: 'string',
+		group: 'model',
+		desc: 'Model for responder stage',
+	},
+	'--discovery-model': {
+		key: 'discoveryModel',
+		kind: 'string',
+		group: 'model',
+		desc: 'Model for discovery stage',
+	},
+	'--builder-model': {
+		key: 'builderModel',
+		kind: 'string',
+		group: 'model',
+		desc: 'Model for builder stage',
+	},
+	'--configurator-model': {
+		key: 'configuratorModel',
+		kind: 'string',
+		group: 'model',
+		desc: 'Model for configurator stage',
+	},
+	'--parameter-updater-model': {
+		key: 'parameterUpdaterModel',
+		kind: 'string',
+		group: 'model',
+		desc: 'Model for parameter updater',
+	},
+
 	// Advanced
 	'--judges': { key: 'numJudges', kind: 'string', group: 'advanced', desc: 'Number of LLM judges' },
 };
@@ -217,6 +309,7 @@ const GROUP_TITLES: Record<FlagGroup, string> = {
 	langsmith: 'LangSmith Options',
 	output: 'Output',
 	feature: 'Feature Flags',
+	model: 'Model Configuration',
 	advanced: 'Advanced',
 };
 
@@ -235,6 +328,7 @@ function formatHelp(): string {
 		'langsmith',
 		'output',
 		'feature',
+		'model',
 		'advanced',
 	];
 
@@ -431,6 +525,31 @@ export function parseEvaluationArgs(argv: string[] = process.argv.slice(2)): Eva
 		donts: parsed.donts,
 		numJudges: parsed.numJudges,
 		featureFlags,
+		// Model configuration
+		model: parsed.model,
+		judgeModel: parsed.judgeModel,
+		supervisorModel: parsed.supervisorModel,
+		responderModel: parsed.responderModel,
+		discoveryModel: parsed.discoveryModel,
+		builderModel: parsed.builderModel,
+		configuratorModel: parsed.configuratorModel,
+		parameterUpdaterModel: parsed.parameterUpdaterModel,
+	};
+}
+
+/**
+ * Converts EvaluationArgs to StageModels for use with environment setup.
+ */
+export function argsToStageModels(args: EvaluationArgs): StageModels {
+	return {
+		default: args.model,
+		supervisor: args.supervisorModel,
+		responder: args.responderModel,
+		discovery: args.discoveryModel,
+		builder: args.builderModel,
+		configurator: args.configuratorModel,
+		parameterUpdater: args.parameterUpdaterModel,
+		judge: args.judgeModel,
 	};
 }
 

packages/@n8n/ai-workflow-builder.ee/evaluations/cli/index.ts

@@ -5,7 +5,7 @@
  * Can be run directly or used as a reference for custom setups.
  */
 
-import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import type { Callbacks } from '@langchain/core/callbacks/manager';
 import type { INodeTypeDescription } from 'n8n-workflow';
 import pLimit from 'p-limit';
 
@@ -13,6 +13,7 @@ import type { SimpleWorkflow } from '@/types/workflow';
 import type { BuilderFeatureFlags } from '@/workflow-builder-agent';
 
 import {
+	argsToStageModels,
 	getDefaultDatasetName,
 	getDefaultExperimentName,
 	parseEvaluationArgs,
@@ -34,33 +35,29 @@ import {
 	loadDefaultTestCases,
 	getDefaultTestCaseIds,
 } from './csv-prompt-loader';
-import {
-	consumeGenerator,
-	getChatPayload,
-	getTracingCallbacks,
-} from '../harness/evaluation-helpers';
+import { consumeGenerator, getChatPayload } from '../harness/evaluation-helpers';
 import { createLogger } from '../harness/logger';
 import { generateRunId, isWorkflowStateValues } from '../langsmith/types';
 import { EVAL_TYPES, EVAL_USERS } from '../support/constants';
-import { setupTestEnvironment, createAgent } from '../support/environment';
+import { setupTestEnvironment, createAgent, type ResolvedStageLLMs } from '../support/environment';
 
 /**
  * Create a workflow generator function.
  * LangSmith tracing is handled via traceable() in the runner.
- * We bridge the trace context to LangChain via getTracingCallbacks().
+ * Callbacks are passed explicitly from the runner to ensure correct trace context
+ * under high concurrency (avoids AsyncLocalStorage race conditions).
  */
 function createWorkflowGenerator(
 	parsedNodeTypes: INodeTypeDescription[],
-	llm: BaseChatModel,
+	llms: ResolvedStageLLMs,
 	featureFlags?: BuilderFeatureFlags,
-): (prompt: string) => Promise<SimpleWorkflow> {
-	return async (prompt: string): Promise<SimpleWorkflow> => {
+): (prompt: string, callbacks?: Callbacks) => Promise<SimpleWorkflow> {
+	return async (prompt: string, callbacks?: Callbacks): Promise<SimpleWorkflow> => {
 		const runId = generateRunId();
-		const callbacks = await getTracingCallbacks();
 
 		const agent = createAgent({
 			parsedNodeTypes,
-			llm,
+			llms,
 			featureFlags,
 		});
 
@@ -149,30 +146,35 @@ export async function runV2Evaluation(): Promise<void> {
 		);
 	}
 
-	// Setup environment
+	// Setup environment with per-stage model configuration
 	const logger = createLogger(args.verbose);
 	const lifecycle = createConsoleLifecycle({ verbose: args.verbose, logger });
-	const env = await setupTestEnvironment(logger);
+	const stageModels = argsToStageModels(args);
+	const env = await setupTestEnvironment(stageModels, logger);
 
 	// Validate LangSmith client early if langsmith backend is requested
 	if (args.backend === 'langsmith' && !env.lsClient) {
 		throw new Error('LangSmith client not initialized - check LANGSMITH_API_KEY');
 	}
 
-	// Create workflow generator (tracing handled via traceable() in runner)
-	const generateWorkflow = createWorkflowGenerator(env.parsedNodeTypes, env.llm, args.featureFlags);
+	// Create workflow generator with per-stage LLMs
+	const generateWorkflow = createWorkflowGenerator(
+		env.parsedNodeTypes,
+		env.llms,
+		args.featureFlags,
+	);
 
-	// Create evaluators based on mode
+	// Create evaluators based on mode (using judge LLM for evaluation)
 	const evaluators: Array<Evaluator<EvaluationContext>> = [];
 
 	switch (args.suite) {
 		case 'llm-judge':
-			evaluators.push(createLLMJudgeEvaluator(env.llm, env.parsedNodeTypes));
+			evaluators.push(createLLMJudgeEvaluator(env.llms.judge, env.parsedNodeTypes));
 			evaluators.push(createProgrammaticEvaluator(env.parsedNodeTypes));
 			break;
 		case 'pairwise':
 			evaluators.push(
-				createPairwiseEvaluator(env.llm, {
+				createPairwiseEvaluator(env.llms.judge, {
 					numJudges: args.numJudges,
 				}),
 			);

packages/@n8n/ai-workflow-builder.ee/evaluations/harness/runner.ts

@@ -1,10 +1,11 @@
+import type { Callbacks } from '@langchain/core/callbacks/manager';
 import type { BaseMessage } from '@langchain/core/messages';
 import { evaluate } from 'langsmith/evaluation';
 import type { Run, Example } from 'langsmith/schemas';
 import { traceable } from 'langsmith/traceable';
 import pLimit from 'p-limit';
 
-import { runWithOptionalLimiter, withTimeout } from './evaluation-helpers';
+import { getTracingCallbacks, runWithOptionalLimiter, withTimeout } from './evaluation-helpers';
 import { toLangsmithEvaluationResult } from './feedback';
 import type {
 	Evaluator,
@@ -778,16 +779,22 @@ async function runLangsmith(config: LangsmithRunConfig): Promise<RunSummary> {
 
 	// Create traceable wrapper ONCE outside target function to avoid context leaking
 	// when running concurrent evaluations. Pass all parameters explicitly (no closures).
+	// IMPORTANT: Get callbacks INSIDE the traceable wrapper where AsyncLocalStorage context
+	// is correctly set, then pass them explicitly to genFn to avoid race conditions.
 	const traceableGenerateWorkflow = traceable(
 		async (args: {
 			prompt: string;
-			genFn: (prompt: string) => Promise<SimpleWorkflow>;
+			genFn: (prompt: string, callbacks?: Callbacks) => Promise<SimpleWorkflow>;
 			limiter?: LlmCallLimiter;
 			genTimeoutMs?: number;
 		}): Promise<SimpleWorkflow> => {
+			// Get callbacks inside traceable where context is correct
+			// Returns undefined if not in a traceable context (e.g., unit tests)
+			const callbacks = await getTracingCallbacks();
+
 			return await runWithOptionalLimiter(async () => {
 				return await withTimeout({
-					promise: args.genFn(args.prompt),
+					promise: args.genFn(args.prompt, callbacks),
 					timeoutMs: args.genTimeoutMs,
 					label: 'workflow_generation',
 				});

packages/@n8n/ai-workflow-builder.ee/evaluations/support/environment.ts

@@ -4,10 +4,11 @@ import { MemorySaver } from '@langchain/langgraph';
 import { Client } from 'langsmith/client';
 import type { INodeTypeDescription } from 'n8n-workflow';
 
+import { DEFAULT_MODEL, getApiKeyEnvVar, MODEL_FACTORIES, type ModelId } from '@/llm-config';
+import type { BuilderFeatureFlags } from '@/workflow-builder-agent';
+import { WorkflowBuilderAgent } from '@/workflow-builder-agent';
+
 import { loadNodesFromFile } from './load-nodes.js';
-import { anthropicClaudeSonnet45 } from '../../src/llm-config.js';
-import type { BuilderFeatureFlags } from '../../src/workflow-builder-agent.js';
-import { WorkflowBuilderAgent } from '../../src/workflow-builder-agent.js';
 import type { EvalLogger } from '../harness/logger.js';
 import {
 	createTraceFilters,
@@ -18,9 +19,52 @@ import {
 /** Maximum memory for trace queue (3GB) */
 const MAX_INGEST_MEMORY_BYTES = 3 * 1024 * 1024 * 1024;
 
+// ============================================================================
+// Stage Models Configuration
+// ============================================================================
+
+/**
+ * Configuration for per-stage model selection.
+ * All fields except 'default' are optional - unspecified stages use the default model.
+ */
+export interface StageModels {
+	/** Default model for all stages */
+	default: ModelId;
+	/** Model for supervisor stage (routing decisions) */
+	supervisor?: ModelId;
+	/** Model for responder stage (final user responses) */
+	responder?: ModelId;
+	/** Model for discovery stage (node discovery) */
+	discovery?: ModelId;
+	/** Model for builder stage (workflow structure) */
+	builder?: ModelId;
+	/** Model for configurator stage (node configuration) */
+	configurator?: ModelId;
+	/** Model for parameter updater (within configurator) */
+	parameterUpdater?: ModelId;
+	/** Model for LLM judge evaluation */
+	judge?: ModelId;
+}
+
+/**
+ * Resolved LLM instances for each stage.
+ * All fields are populated (using default model as fallback).
+ */
+export interface ResolvedStageLLMs {
+	default: BaseChatModel;
+	supervisor: BaseChatModel;
+	responder: BaseChatModel;
+	discovery: BaseChatModel;
+	builder: BaseChatModel;
+	configurator: BaseChatModel;
+	parameterUpdater: BaseChatModel;
+	judge: BaseChatModel;
+}
+
 export interface TestEnvironment {
 	parsedNodeTypes: INodeTypeDescription[];
-	llm: BaseChatModel;
+	/** Resolved LLM instances for each stage */
+	llms: ResolvedStageLLMs;
 	tracer?: LangChainTracer;
 	lsClient?: Client;
 	/** Trace filtering utilities (only present when minimal tracing is enabled) */
@@ -28,16 +72,48 @@ export interface TestEnvironment {
 }
 
 /**
- * Sets up the LLM with proper configuration
+ * Sets up an LLM with proper configuration
+ * @param modelId - Model identifier (defaults to DEFAULT_MODEL)
  * @returns Configured LLM instance
- * @throws Error if N8N_AI_ANTHROPIC_KEY environment variable is not set
+ * @throws Error if the required API key environment variable is not set
  */
-export async function setupLLM(): Promise<BaseChatModel> {
-	const apiKey = process.env.N8N_AI_ANTHROPIC_KEY;
+export async function setupLLM(modelId: ModelId = DEFAULT_MODEL): Promise<BaseChatModel> {
+	const envVar = getApiKeyEnvVar(modelId);
+	const apiKey = process.env[envVar];
 	if (!apiKey) {
-		throw new Error('N8N_AI_ANTHROPIC_KEY environment variable is required');
+		throw new Error(`${envVar} environment variable is required for model ${modelId}`);
 	}
-	return await anthropicClaudeSonnet45({ apiKey });
+	const factory = MODEL_FACTORIES[modelId];
+	return await factory({ apiKey });
+}
+
+/**
+ * Resolves all stage models to LLM instances.
+ * Unspecified stages fall back to the default model.
+ * @param stageModels - Per-stage model configuration
+ * @returns Resolved LLM instances for each stage
+ */
+export async function resolveStageModels(stageModels: StageModels): Promise<ResolvedStageLLMs> {
+	const defaultLLM = await setupLLM(stageModels.default);
+
+	// For stages without specific model, use default
+	// For parameter updater, fall back to configurator if not specified
+	const configuratorLLM = stageModels.configurator
+		? await setupLLM(stageModels.configurator)
+		: defaultLLM;
+
+	return {
+		default: defaultLLM,
+		supervisor: stageModels.supervisor ? await setupLLM(stageModels.supervisor) : defaultLLM,
+		responder: stageModels.responder ? await setupLLM(stageModels.responder) : defaultLLM,
+		discovery: stageModels.discovery ? await setupLLM(stageModels.discovery) : defaultLLM,
+		builder: stageModels.builder ? await setupLLM(stageModels.builder) : defaultLLM,
+		configurator: configuratorLLM,
+		parameterUpdater: stageModels.parameterUpdater
+			? await setupLLM(stageModels.parameterUpdater)
+			: configuratorLLM,
+		judge: stageModels.judge ? await setupLLM(stageModels.judge) : defaultLLM,
+	};
 }
 
 /**
@@ -97,24 +173,39 @@ export function createLangsmithClient(logger?: EvalLogger): LangsmithClientResul
 
 /**
  * Sets up the test environment with LLM, nodes, and tracing
+ * @param stageModels - Per-stage model configuration (optional, uses default model if not provided)
  * @param logger - Optional logger for trace filter output
  * @returns Test environment configuration
  */
-export async function setupTestEnvironment(logger?: EvalLogger): Promise<TestEnvironment> {
+export async function setupTestEnvironment(
+	stageModels?: StageModels,
+	logger?: EvalLogger,
+): Promise<TestEnvironment> {
 	const parsedNodeTypes = loadNodesFromFile();
-	const llm = await setupLLM();
+
+	// Use provided stage models or default configuration
+	const models: StageModels = stageModels ?? { default: DEFAULT_MODEL };
+	const llms = await resolveStageModels(models);
+
 	const lsClientResult = createLangsmithClient(logger);
 
 	const lsClient = lsClientResult?.client;
 	const traceFilters = lsClientResult?.traceFilters;
 	const tracer = lsClient ? createTracer(lsClient, 'workflow-builder-evaluation') : undefined;
 
-	return { parsedNodeTypes, llm, tracer, lsClient, traceFilters };
+	return {
+		parsedNodeTypes,
+		llms,
+		tracer,
+		lsClient,
+		traceFilters,
+	};
 }
 
 export interface CreateAgentOptions {
 	parsedNodeTypes: INodeTypeDescription[];
-	llm: BaseChatModel;
+	/** Per-stage LLMs resolved from model configuration */
+	llms: ResolvedStageLLMs;
 	tracer?: LangChainTracer;
 	featureFlags?: BuilderFeatureFlags;
 	experimentName?: string;
@@ -126,12 +217,18 @@ export interface CreateAgentOptions {
  * @returns Configured WorkflowBuilderAgent
  */
 export function createAgent(options: CreateAgentOptions): WorkflowBuilderAgent {
-	const { parsedNodeTypes, llm, tracer, featureFlags, experimentName } = options;
+	const { parsedNodeTypes, llms, tracer, featureFlags, experimentName } = options;
 
 	return new WorkflowBuilderAgent({
 		parsedNodeTypes,
-		llmSimpleTask: llm,
-		llmComplexTask: llm,
+		stageLLMs: {
+			supervisor: llms.supervisor,
+			responder: llms.responder,
+			discovery: llms.discovery,
+			builder: llms.builder,
+			configurator: llms.configurator,
+			parameterUpdater: llms.parameterUpdater,
+		},
 		checkpointer: new MemorySaver(),
 		tracer,
 		featureFlags,

packages/@n8n/ai-workflow-builder.ee/src/agents/responder.agent.ts

@@ -1,7 +1,8 @@
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
-import type { AIMessage, BaseMessage } from '@langchain/core/messages';
-import { HumanMessage } from '@langchain/core/messages';
+import type { BaseMessage } from '@langchain/core/messages';
+import { AIMessage, HumanMessage } from '@langchain/core/messages';
 import { ChatPromptTemplate } from '@langchain/core/prompts';
+import type { RunnableConfig } from '@langchain/core/runnables';
 
 import {
 	buildResponderPrompt,
@@ -12,6 +13,7 @@ import {
 
 import type { CoordinationLogEntry } from '../types/coordination';
 import type { DiscoveryContext } from '../types/discovery-types';
+import { isAIMessage } from '../types/langchain';
 import type { SimpleWorkflow } from '../types/workflow';
 import {
 	getErrorEntry,
@@ -150,8 +152,10 @@ export class ResponderAgent {
 
 	/**
 	 * Invoke the responder agent with the given context
+	 * @param context - Responder context with messages and workflow state
+	 * @param config - Optional RunnableConfig for tracing callbacks
 	 */
-	async invoke(context: ResponderContext): Promise<AIMessage> {
+	async invoke(context: ResponderContext, config?: RunnableConfig): Promise<AIMessage> {
 		const agent = systemPrompt.pipe(this.llm);
 
 		const contextMessage = this.buildContextMessage(context);
@@ -159,6 +163,12 @@ export class ResponderAgent {
 			? [...context.messages, contextMessage]
 			: context.messages;
 
-		return (await agent.invoke({ messages: messagesToSend })) as AIMessage;
+		const result = await agent.invoke({ messages: messagesToSend }, config);
+		if (!isAIMessage(result)) {
+			return new AIMessage({
+				content: 'I encountered an issue generating a response. Please try again.',
+			});
+		}
+		return result;
 	}
 }

packages/@n8n/ai-workflow-builder.ee/src/agents/supervisor.agent.ts

@@ -2,6 +2,7 @@ import type { BaseChatModel } from '@langchain/core/language_models/chat_models'
 import type { BaseMessage } from '@langchain/core/messages';
 import { HumanMessage } from '@langchain/core/messages';
 import { ChatPromptTemplate } from '@langchain/core/prompts';
+import type { RunnableConfig } from '@langchain/core/runnables';
 import { z } from 'zod';
 
 import { buildSupervisorPrompt } from '@/prompts/agents/supervisor.prompt';
@@ -104,8 +105,10 @@ export class SupervisorAgent {
 
 	/**
 	 * Invoke the supervisor to get routing decision
+	 * @param context - Supervisor context with messages and workflow state
+	 * @param config - Optional RunnableConfig for tracing callbacks
 	 */
-	async invoke(context: SupervisorContext): Promise<SupervisorRouting> {
+	async invoke(context: SupervisorContext, config?: RunnableConfig): Promise<SupervisorRouting> {
 		const agent = systemPrompt.pipe<SupervisorRouting>(
 			this.llm.withStructuredOutput(supervisorRoutingSchema, {
 				name: 'routing_decision',
@@ -117,6 +120,6 @@ export class SupervisorAgent {
 			? [...context.messages, contextMessage]
 			: context.messages;
 
-		return await agent.invoke({ messages: messagesToSend });
+		return await agent.invoke({ messages: messagesToSend }, config);
 	}
 }

packages/@n8n/ai-workflow-builder.ee/src/ai-workflow-builder-agent.service.ts

@@ -169,9 +169,15 @@ export class AiWorkflowBuilderService {
 
 		const agent = new WorkflowBuilderAgent({
 			parsedNodeTypes: this.parsedNodeTypes,
-			// We use Sonnet both for simple and complex tasks
-			llmSimpleTask: anthropicClaude,
-			llmComplexTask: anthropicClaude,
+			// Use the same model for all stages in production
+			stageLLMs: {
+				supervisor: anthropicClaude,
+				responder: anthropicClaude,
+				discovery: anthropicClaude,
+				builder: anthropicClaude,
+				configurator: anthropicClaude,
+				parameterUpdater: anthropicClaude,
+			},
 			logger: this.logger,
 			checkpointer: this.sessionManager.getCheckpointer(),
 			tracer: tracingClient

packages/@n8n/ai-workflow-builder.ee/src/chains/conversation-compact.ts

@@ -1,6 +1,7 @@
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import type { BaseMessage } from '@langchain/core/messages';
 import { AIMessage, HumanMessage } from '@langchain/core/messages';
+import type { RunnableConfig } from '@langchain/core/runnables';
 import z from 'zod';
 
 import { compactPromptTemplate } from '@/prompts/chains/compact.prompt';
@@ -9,6 +10,7 @@ export async function conversationCompactChain(
 	llm: BaseChatModel,
 	messages: BaseMessage[],
 	previousSummary: string = '',
+	config?: RunnableConfig,
 ) {
 	// Use structured output for consistent summary format
 	const CompactedSession = z.object({
@@ -44,7 +46,7 @@ export async function conversationCompactChain(
 		conversationText,
 	});
 
-	const structuredOutput = await modelWithStructure.invoke(compactPrompt);
+	const structuredOutput = await modelWithStructure.invoke(compactPrompt, config);
 
 	const formattedSummary = `## Previous Conversation Summary
 

packages/@n8n/ai-workflow-builder.ee/src/chains/workflow-name.ts

@@ -1,9 +1,14 @@
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import type { RunnableConfig } from '@langchain/core/runnables';
 import z from 'zod';
 
 import { workflowNamingPromptTemplate } from '@/prompts/chains/workflow-name.prompt';
 
-export async function workflowNameChain(llm: BaseChatModel, initialPrompt: string) {
+export async function workflowNameChain(
+	llm: BaseChatModel,
+	initialPrompt: string,
+	config?: RunnableConfig,
+) {
 	// Use structured output for the workflow name to ensure it meets the required format and length
 	const nameSchema = z.object({
 		name: z.string().min(10).max(128).describe('Name of the workflow based on the prompt'),
@@ -15,7 +20,8 @@ export async function workflowNameChain(llm: BaseChatModel, initialPrompt: strin
 		initialPrompt,
 	});
 
-	const structuredOutput = (await modelWithStructure.invoke(prompt)) as z.infer<typeof nameSchema>;
+	const rawOutput = await modelWithStructure.invoke(prompt, config);
+	const structuredOutput = nameSchema.parse(rawOutput);
 
 	return {
 		name: structuredOutput.name,

packages/@n8n/ai-workflow-builder.ee/src/llm-config.ts

@@ -1,33 +1,23 @@
 // Different LLMConfig type for this file - specific to LLM providers
+import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
+
 import { MAX_OUTPUT_TOKENS } from '@/constants';
 
 import { getProxyAgent } from './utils/http-proxy-agent';
 
-interface LLMProviderConfig {
+/**
+ * Configuration for LLM provider initialization.
+ */
+export interface LLMProviderConfig {
 	apiKey: string;
 	baseUrl?: string;
 	headers?: Record<string, string>;
 }
 
-export const o4mini = async (config: LLMProviderConfig) => {
+export const gpt52 = async (config: LLMProviderConfig) => {
 	const { ChatOpenAI } = await import('@langchain/openai');
 	return new ChatOpenAI({
-		model: 'o4-mini-2025-04-16',
-		apiKey: config.apiKey,
-		configuration: {
-			baseURL: config.baseUrl,
-			defaultHeaders: config.headers,
-			fetchOptions: {
-				dispatcher: getProxyAgent(config.baseUrl ?? 'https://api.openai.com/v1'),
-			},
-		},
-	});
-};
-
-export const gpt41mini = async (config: LLMProviderConfig) => {
-	const { ChatOpenAI } = await import('@langchain/openai');
-	return new ChatOpenAI({
-		model: 'gpt-4.1-mini-2025-04-14',
+		model: 'gpt-5.2-2025-12-11',
 		apiKey: config.apiKey,
 		temperature: 0,
 		maxTokens: -1,
@@ -41,27 +31,10 @@ export const gpt41mini = async (config: LLMProviderConfig) => {
 	});
 };
 
-export const gpt41 = async (config: LLMProviderConfig) => {
-	const { ChatOpenAI } = await import('@langchain/openai');
-	return new ChatOpenAI({
-		model: 'gpt-4.1-2025-04-14',
-		apiKey: config.apiKey,
-		temperature: 0.3,
-		maxTokens: -1,
-		configuration: {
-			baseURL: config.baseUrl,
-			defaultHeaders: config.headers,
-			fetchOptions: {
-				dispatcher: getProxyAgent(config.baseUrl ?? 'https://api.openai.com/v1'),
-			},
-		},
-	});
-};
-
 export const anthropicClaudeSonnet45 = async (config: LLMProviderConfig) => {
 	const { ChatAnthropic } = await import('@langchain/anthropic');
 	const model = new ChatAnthropic({
-		model: 'claude-sonnet-4-5',
+		model: 'claude-sonnet-4-5-20250929',
 		apiKey: config.apiKey,
 		temperature: 0,
 		maxTokens: MAX_OUTPUT_TOKENS,
@@ -101,3 +74,165 @@ export const anthropicHaiku45 = async (config: LLMProviderConfig) => {
 
 	return model;
 };
+
+export const anthropicClaudeOpus45 = async (config: LLMProviderConfig) => {
+	const { ChatAnthropic } = await import('@langchain/anthropic');
+	const model = new ChatAnthropic({
+		model: 'claude-opus-4-5-20251101',
+		apiKey: config.apiKey,
+		temperature: 0,
+		maxTokens: MAX_OUTPUT_TOKENS,
+		anthropicApiUrl: config.baseUrl,
+		clientOptions: {
+			defaultHeaders: config.headers,
+			fetchOptions: {
+				dispatcher: getProxyAgent(config.baseUrl),
+			},
+		},
+	});
+
+	// Remove Langchain default topP parameter since Opus 4.5 doesn't allow setting both temperature and topP
+	delete model.topP;
+
+	return model;
+};
+
+// ============================================================================
+// OpenRouter Models
+// ============================================================================
+
+const OPENROUTER_BASE_URL = 'https://openrouter.ai/api/v1';
+
+/**
+ * Creates an OpenRouter model factory for a given model name.
+ * Uses OpenAI-compatible API with OpenRouter base URL.
+ */
+function createOpenRouterModel(modelName: string) {
+	return async (config: LLMProviderConfig) => {
+		const { ChatOpenAI } = await import('@langchain/openai');
+		return new ChatOpenAI({
+			model: modelName,
+			apiKey: config.apiKey,
+			temperature: 0,
+			maxTokens: -1,
+			configuration: {
+				baseURL: OPENROUTER_BASE_URL,
+				defaultHeaders: {
+					...config.headers,
+					'HTTP-Referer': 'https://n8n.io',
+					'X-Title': 'n8n AI Workflow Builder',
+				},
+				fetchOptions: {
+					dispatcher: getProxyAgent(OPENROUTER_BASE_URL),
+				},
+			},
+		});
+	};
+}
+
+// OpenRouter model factories
+export const glm47 = createOpenRouterModel('thudm/glm-4-plus');
+export const gemini3Flash = createOpenRouterModel('google/gemini-3-flash-preview');
+export const deepseekV32 = createOpenRouterModel('deepseek/deepseek-chat-v3-0324');
+export const gemini3Pro = createOpenRouterModel('google/gemini-3-pro-preview');
+export const devstral = createOpenRouterModel('mistralai/devstral-small');
+
+// ============================================================================
+// Model Registry
+// ============================================================================
+
+/**
+ * IMPORTANT: Generation stages currently only support Anthropic models.
+ *
+ * Non-Anthropic models (OpenAI, OpenRouter) are available for evaluation/judging
+ * purposes only. Using them for generation stages (supervisor, discovery, builder,
+ * configurator, responder, parameterUpdater) will likely fail due to:
+ *
+ * 1. Prompt caching: Our prompts use Anthropic's cache_control for efficiency
+ * 2. Tool schemas: add_nodes and update_parameters tools use passthrough() schemas
+ *    which only Anthropic models handle correctly
+ *
+ * TODO: Add provider-agnostic prompt/tool support to enable non-Anthropic generation.
+ */
+
+/**
+ * Available model identifiers for the eval CLI.
+ * These can be used with --model, --judge-model, and per-stage model flags.
+ */
+export type ModelId =
+	// Native models
+	| 'claude-opus-4.5'
+	| 'claude-sonnet-4.5'
+	| 'claude-haiku-4.5'
+	| 'gpt-5.2'
+	// OpenRouter models
+	| 'glm-4.7'
+	| 'gemini-3-flash'
+	| 'deepseek-v3.2'
+	| 'gemini-3-pro'
+	| 'devstral';
+
+/**
+ * Model factory functions mapped by model ID.
+ */
+export const MODEL_FACTORIES: Record<
+	ModelId,
+	(config: LLMProviderConfig) => Promise<BaseChatModel>
+> = {
+	// Native models
+	'claude-opus-4.5': anthropicClaudeOpus45,
+	'claude-sonnet-4.5': anthropicClaudeSonnet45,
+	'claude-haiku-4.5': anthropicHaiku45,
+	'gpt-5.2': gpt52,
+	// OpenRouter models
+	'glm-4.7': glm47,
+	'gemini-3-flash': gemini3Flash,
+	'deepseek-v3.2': deepseekV32,
+	'gemini-3-pro': gemini3Pro,
+	devstral,
+};
+
+/** OpenRouter model IDs for API key resolution */
+const OPENROUTER_MODELS: ModelId[] = [
+	'glm-4.7',
+	'gemini-3-flash',
+	'deepseek-v3.2',
+	'gemini-3-pro',
+	'devstral',
+];
+
+/**
+ * Get the required API key environment variable for a model.
+ */
+export function getApiKeyEnvVar(modelId: ModelId): string {
+	if (OPENROUTER_MODELS.includes(modelId)) {
+		return 'OPENROUTER_API_KEY';
+	}
+	if (modelId.startsWith('gpt')) {
+		return 'N8N_AI_OPENAI_KEY';
+	}
+	return 'N8N_AI_ANTHROPIC_KEY';
+}
+
+/**
+ * List of available model IDs for CLI help text.
+ * Explicitly defined to avoid type casting Object.keys().
+ */
+export const AVAILABLE_MODELS: readonly ModelId[] = [
+	// Native models
+	'claude-opus-4.5',
+	'claude-sonnet-4.5',
+	'claude-haiku-4.5',
+	'gpt-5.2',
+	// OpenRouter models
+	'glm-4.7',
+	'gemini-3-flash',
+	'deepseek-v3.2',
+	'gemini-3-pro',
+	'devstral',
+] as const;
+
+/**
+ * Default model used when no model is specified.
+ */
+export const DEFAULT_MODEL: ModelId = 'claude-sonnet-4.5';

packages/@n8n/ai-workflow-builder.ee/src/multi-agent-workflow-subgraphs.ts

@@ -1,5 +1,5 @@
-import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { HumanMessage } from '@langchain/core/messages';
+import type { RunnableConfig } from '@langchain/core/runnables';
 import { StateGraph, END, START, type MemorySaver } from '@langchain/langgraph';
 import type { Logger } from '@n8n/backend-common';
 import type { INodeTypeDescription } from 'n8n-workflow';
@@ -29,7 +29,7 @@ import {
 	handleCreateWorkflowName,
 	handleDeleteMessages,
 } from './utils/state-modifier';
-import type { BuilderFeatureFlags } from './workflow-builder-agent';
+import type { BuilderFeatureFlags, StageLLMs } from './workflow-builder-agent';
 
 /**
  * Maps routing decisions to graph node names.
@@ -47,8 +47,8 @@ function routeToNode(next: string): string {
 
 export interface MultiAgentSubgraphConfig {
 	parsedNodeTypes: INodeTypeDescription[];
-	llmSimpleTask: BaseChatModel;
-	llmComplexTask: BaseChatModel;
+	/** Per-stage LLM configuration */
+	stageLLMs: StageLLMs;
 	logger?: Logger;
 	instanceUrl?: string;
 	checkpointer?: MemorySaver;
@@ -60,7 +60,8 @@ export interface MultiAgentSubgraphConfig {
 }
 
 /**
- * Creates a subgraph node handler with standardized error handling
+ * Creates a subgraph node handler with standardized error handling.
+ * Accepts RunnableConfig as second parameter to propagate callbacks for tracing.
  */
 function createSubgraphNodeHandler<
 	TSubgraph extends BaseSubgraph<unknown, Record<string, unknown>, Record<string, unknown>>,
@@ -71,10 +72,15 @@ function createSubgraphNodeHandler<
 	logger?: Logger,
 	recursionLimit?: number,
 ) {
-	return async (state: typeof ParentGraphState.State) => {
+	return async (state: typeof ParentGraphState.State, config?: RunnableConfig) => {
 		try {
 			const input = subgraph.transformInput(state);
-			const result = await compiledGraph.invoke(input, { recursionLimit });
+			// Merge parent config (callbacks, metadata) with recursionLimit
+			const invokeConfig: RunnableConfig = {
+				...config,
+				recursionLimit,
+			};
+			const result = await compiledGraph.invoke(input, invokeConfig);
 			const output = subgraph.transformOutput(result, state);
 
 			return output;
@@ -122,7 +128,7 @@ function createSubgraphNodeHandler<
 export function createMultiAgentWorkflowWithSubgraphs(config: MultiAgentSubgraphConfig) {
 	const {
 		parsedNodeTypes,
-		llmComplexTask,
+		stageLLMs,
 		logger,
 		instanceUrl,
 		checkpointer,
@@ -131,30 +137,31 @@ export function createMultiAgentWorkflowWithSubgraphs(config: MultiAgentSubgraph
 		onGenerationSuccess,
 	} = config;
 
-	const supervisorAgent = new SupervisorAgent({ llm: llmComplexTask });
-	const responderAgent = new ResponderAgent({ llm: llmComplexTask });
+	const supervisorAgent = new SupervisorAgent({ llm: stageLLMs.supervisor });
+	const responderAgent = new ResponderAgent({ llm: stageLLMs.responder });
 
 	// Create subgraph instances
 	const discoverySubgraph = new DiscoverySubgraph();
 	const builderSubgraph = new BuilderSubgraph();
 	const configuratorSubgraph = new ConfiguratorSubgraph();
 
-	// Compile subgraphs
+	// Compile subgraphs with per-stage LLMs
 	const compiledDiscovery = discoverySubgraph.create({
 		parsedNodeTypes,
-		llm: llmComplexTask,
+		llm: stageLLMs.discovery,
 		logger,
 		featureFlags,
 	});
 	const compiledBuilder = builderSubgraph.create({
 		parsedNodeTypes,
-		llm: llmComplexTask,
+		llm: stageLLMs.builder,
 		logger,
 		featureFlags,
 	});
 	const compiledConfigurator = configuratorSubgraph.create({
 		parsedNodeTypes,
-		llm: llmComplexTask,
+		llm: stageLLMs.configurator,
+		llmParameterUpdater: stageLLMs.parameterUpdater,
 		logger,
 		instanceUrl,
 		featureFlags,
@@ -164,27 +171,35 @@ export function createMultiAgentWorkflowWithSubgraphs(config: MultiAgentSubgraph
 	return (
 		new StateGraph(ParentGraphState)
 			// Add Supervisor Node (only used for initial routing)
-			.addNode('supervisor', async (state) => {
-				const routing = await supervisorAgent.invoke({
-					messages: state.messages,
-					workflowJSON: state.workflowJSON,
-					coordinationLog: state.coordinationLog,
-					previousSummary: state.previousSummary,
-				});
+			// Accepts config as second param to propagate callbacks for tracing
+			.addNode('supervisor', async (state, config) => {
+				const routing = await supervisorAgent.invoke(
+					{
+						messages: state.messages,
+						workflowJSON: state.workflowJSON,
+						coordinationLog: state.coordinationLog,
+						previousSummary: state.previousSummary,
+					},
+					config,
+				);
 
 				return {
 					nextPhase: routing.next,
 				};
 			})
 			// Add Responder Node (synthesizes final user-facing response)
-			.addNode('responder', async (state) => {
-				const response = await responderAgent.invoke({
-					messages: state.messages,
-					coordinationLog: state.coordinationLog,
-					discoveryContext: state.discoveryContext,
-					workflowJSON: state.workflowJSON,
-					previousSummary: state.previousSummary,
-				});
+			// Accepts config as second param to propagate callbacks for tracing
+			.addNode('responder', async (state, config) => {
+				const response = await responderAgent.invoke(
+					{
+						messages: state.messages,
+						coordinationLog: state.coordinationLog,
+						discoveryContext: state.discoveryContext,
+						workflowJSON: state.workflowJSON,
+						previousSummary: state.previousSummary,
+					},
+					config,
+				);
 
 				// Call success callback only when generation completed without errors
 				if (onGenerationSuccess && !hasErrorInLog(state.coordinationLog)) {
@@ -212,25 +227,27 @@ export function createMultiAgentWorkflowWithSubgraphs(config: MultiAgentSubgraph
 				nextPhase: determineStateAction(state, autoCompactThresholdTokens),
 			}))
 			.addNode('cleanup_dangling', (state) => handleCleanupDangling(state.messages, logger))
-			.addNode('compact_messages', async (state) => {
+			.addNode('compact_messages', async (state, config) => {
 				const isAutoCompact = state.messages[state.messages.length - 1]?.content !== '/compact';
 				return await handleCompactMessages(
 					state.messages,
 					state.previousSummary ?? '',
-					llmComplexTask,
+					stageLLMs.responder,
 					isAutoCompact,
+					config,
 				);
 			})
 			.addNode('delete_messages', (state) => handleDeleteMessages(state.messages))
 			.addNode('clear_error_state', (state) => handleClearErrorState(state.coordinationLog, logger))
 			.addNode(
 				'create_workflow_name',
-				async (state) =>
+				async (state, config) =>
 					await handleCreateWorkflowName(
 						state.messages,
 						state.workflowJSON,
-						llmComplexTask,
+						stageLLMs.responder,
 						logger,
+						config,
 					),
 			)
 			// Add Subgraph Nodes (using helper to reduce duplication)

packages/@n8n/ai-workflow-builder.ee/src/subgraphs/configurator.subgraph.ts

@@ -101,6 +101,8 @@ export const ConfiguratorSubgraphState = Annotation.Root({
 export interface ConfiguratorSubgraphConfig {
 	parsedNodeTypes: INodeTypeDescription[];
 	llm: BaseChatModel;
+	/** Separate LLM for parameter updater chain (defaults to llm if not provided) */
+	llmParameterUpdater?: BaseChatModel;
 	logger?: Logger;
 	instanceUrl?: string;
 	featureFlags?: BuilderFeatureFlags;
@@ -124,11 +126,14 @@ export class ConfiguratorSubgraph extends BaseSubgraph<
 		// Check if template examples are enabled
 		const includeExamples = config.featureFlags?.templateExamples === true;
 
+		// Use separate LLM for parameter updater if provided
+		const parameterUpdaterLLM = config.llmParameterUpdater ?? config.llm;
+
 		// Create base tools
 		const baseTools = [
 			createUpdateNodeParametersTool(
 				config.parsedNodeTypes,
-				config.llm, // Uses same LLM for parameter updater chain
+				parameterUpdaterLLM, // Uses separate LLM for parameter updater chain
 				config.logger,
 				config.instanceUrl,
 			),

packages/@n8n/ai-workflow-builder.ee/src/subgraphs/test/integration/multi-agent-error-handling.integration.test.ts

@@ -110,8 +110,14 @@ describe('Multi-Agent Error Handling - Integration Tests (AI-1812)', () => {
 			// Don't use checkpointer for this test - we're just testing error message
 			const graph = createMultiAgentWorkflowWithSubgraphs({
 				parsedNodeTypes,
-				llmSimpleTask: llm,
-				llmComplexTask: llm,
+				stageLLMs: {
+					supervisor: llm,
+					responder: llm,
+					discovery: llm,
+					builder: llm,
+					configurator: llm,
+					parameterUpdater: llm,
+				},
 				logger: mockLogger,
 			});
 

packages/@n8n/ai-workflow-builder.ee/src/test/workflow-builder-agent.test.ts

@@ -56,8 +56,7 @@ import {
 
 describe('WorkflowBuilderAgent', () => {
 	let agent: WorkflowBuilderAgent;
-	let mockLlmSimple: BaseChatModel;
-	let mockLlmComplex: BaseChatModel;
+	let mockLlm: BaseChatModel;
 	let mockLogger: Logger;
 	let mockCheckpointer: MemorySaver;
 	let parsedNodeTypes: INodeTypeDescription[];
@@ -68,18 +67,12 @@ describe('WorkflowBuilderAgent', () => {
 	>;
 
 	beforeEach(() => {
-		mockLlmSimple = mock<BaseChatModel>({
+		mockLlm = mock<BaseChatModel>({
 			_llmType: jest.fn().mockReturnValue('test-llm'),
 			bindTools: jest.fn().mockReturnThis(),
 			invoke: jest.fn(),
 		});
 
-		mockLlmComplex = mock<BaseChatModel>({
-			_llmType: jest.fn().mockReturnValue('test-llm-complex'),
-			bindTools: jest.fn().mockReturnThis(),
-			invoke: jest.fn(),
-		});
-
 		mockLogger = mock<Logger>({
 			debug: jest.fn(),
 			info: jest.fn(),
@@ -108,8 +101,14 @@ describe('WorkflowBuilderAgent', () => {
 
 		config = {
 			parsedNodeTypes,
-			llmSimpleTask: mockLlmSimple,
-			llmComplexTask: mockLlmComplex,
+			stageLLMs: {
+				supervisor: mockLlm,
+				responder: mockLlm,
+				discovery: mockLlm,
+				builder: mockLlm,
+				configurator: mockLlm,
+				parameterUpdater: mockLlm,
+			},
 			logger: mockLogger,
 			checkpointer: mockCheckpointer,
 		};
@@ -172,7 +171,7 @@ describe('WorkflowBuilderAgent', () => {
 			mockCreateStreamProcessor.mockReturnValue(mockAsyncGenerator);
 
 			// Mock the LLM to return a simple response
-			(mockLlmSimple.invoke as jest.Mock).mockResolvedValue({
+			(mockLlm.invoke as jest.Mock).mockResolvedValue({
 				content: 'Mocked response',
 				tool_calls: [],
 			});

packages/@n8n/ai-workflow-builder.ee/src/utils/state-modifier.ts

@@ -1,6 +1,7 @@
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import type { BaseMessage } from '@langchain/core/messages';
 import { HumanMessage, RemoveMessage } from '@langchain/core/messages';
+import type { RunnableConfig } from '@langchain/core/runnables';
 import type { Logger } from '@n8n/backend-common';
 
 import { cleanupDanglingToolCallMessages } from './cleanup-dangling-tool-call-messages';
@@ -142,12 +143,19 @@ export function handleCleanupDangling(
  *
  * For manual /compact: Removes all messages, routes to responder for acknowledgment.
  * For auto-compact: Removes old messages, preserves last user message to continue processing.
+ *
+ * @param messages - Conversation messages to compact
+ * @param previousSummary - Previous conversation summary
+ * @param llm - Language model for summarization
+ * @param isAutoCompact - Whether this is auto-compaction (preserve last message) or manual
+ * @param config - Optional RunnableConfig for tracing callbacks
  */
 export async function handleCompactMessages(
 	messages: BaseMessage[],
 	previousSummary: string,
 	llm: BaseChatModel,
 	isAutoCompact: boolean,
+	config?: RunnableConfig,
 ): Promise<{
 	previousSummary: string;
 	messages: BaseMessage[];
@@ -158,7 +166,7 @@ export async function handleCompactMessages(
 		throw new Error('Cannot compact messages: no HumanMessage found');
 	}
 
-	const compactedMessages = await conversationCompactChain(llm, messages, previousSummary);
+	const compactedMessages = await conversationCompactChain(llm, messages, previousSummary, config);
 
 	// For manual /compact: just remove messages, responder will generate acknowledgment
 	// For auto-compact: remove messages but preserve the last user message to continue processing
@@ -252,12 +260,19 @@ export function handleClearErrorState(
 
 /**
  * Generates a workflow name from the initial user message.
+ *
+ * @param messages - Conversation messages
+ * @param workflowJSON - Current workflow state
+ * @param llm - Language model for name generation
+ * @param logger - Optional logger
+ * @param config - Optional RunnableConfig for tracing callbacks
  */
 export async function handleCreateWorkflowName(
 	messages: BaseMessage[],
 	workflowJSON: SimpleWorkflow,
 	llm: BaseChatModel,
 	logger?: Logger,
+	config?: RunnableConfig,
 ): Promise<{ workflowJSON: SimpleWorkflow }> {
 	if (messages.length === 1 && messages[0] instanceof HumanMessage) {
 		const initialMessage = messages[0];
@@ -267,7 +282,7 @@ export async function handleCreateWorkflowName(
 		}
 
 		logger?.debug('Generating workflow name');
-		const { name } = await workflowNameChain(llm, initialMessage.content);
+		const { name } = await workflowNameChain(llm, initialMessage.content, config);
 
 		return {
 			workflowJSON: { ...workflowJSON, name },

packages/@n8n/ai-workflow-builder.ee/src/workflow-builder-agent.ts

@@ -37,10 +37,23 @@ export type TypedStateSnapshot = Omit<StateSnapshot, 'values'> & {
 	values: typeof WorkflowState.State;
 };
 
+/**
+ * Per-stage LLM configuration for the workflow builder.
+ * All stages must be configured with an LLM instance.
+ */
+export interface StageLLMs {
+	supervisor: BaseChatModel;
+	responder: BaseChatModel;
+	discovery: BaseChatModel;
+	builder: BaseChatModel;
+	configurator: BaseChatModel;
+	parameterUpdater: BaseChatModel;
+}
+
 export interface WorkflowBuilderAgentConfig {
 	parsedNodeTypes: INodeTypeDescription[];
-	llmSimpleTask: BaseChatModel;
-	llmComplexTask: BaseChatModel;
+	/** Per-stage LLM configuration */
+	stageLLMs: StageLLMs;
 	logger?: Logger;
 	checkpointer: MemorySaver;
 	tracer?: LangChainTracer;
@@ -80,8 +93,7 @@ export interface ChatPayload {
 export class WorkflowBuilderAgent {
 	private checkpointer: MemorySaver;
 	private parsedNodeTypes: INodeTypeDescription[];
-	private llmSimpleTask: BaseChatModel;
-	private llmComplexTask: BaseChatModel;
+	private stageLLMs: StageLLMs;
 	private logger?: Logger;
 	private tracer?: LangChainTracer;
 	private instanceUrl?: string;
@@ -90,8 +102,7 @@ export class WorkflowBuilderAgent {
 
 	constructor(config: WorkflowBuilderAgentConfig) {
 		this.parsedNodeTypes = config.parsedNodeTypes;
-		this.llmSimpleTask = config.llmSimpleTask;
-		this.llmComplexTask = config.llmComplexTask;
+		this.stageLLMs = config.stageLLMs;
 		this.logger = config.logger;
 		this.checkpointer = config.checkpointer;
 		this.tracer = config.tracer;
@@ -107,8 +118,7 @@ export class WorkflowBuilderAgent {
 	private createWorkflow(featureFlags?: BuilderFeatureFlags) {
 		return createMultiAgentWorkflowWithSubgraphs({
 			parsedNodeTypes: this.parsedNodeTypes,
-			llmSimpleTask: this.llmSimpleTask,
-			llmComplexTask: this.llmComplexTask,
+			stageLLMs: this.stageLLMs,
 			logger: this.logger,
 			instanceUrl: this.instanceUrl,
 			checkpointer: this.checkpointer,