Wire task-level model routing into the host workflow

src/cli-workflow.test.ts

@@ -210,6 +210,71 @@ describe('runMainWorkflow', () => {
     expect(logger.warn).not.toHaveBeenCalled();
   });
 
+  it('stages routing before config generation and waits for selection before agent startup', async () => {
+    const callOrder: string[] = [];
+    const routingState = {
+      root: '/tmp/awf-test-routing',
+      inputDir: '/tmp/awf-test-routing/input',
+      outputDir: '/tmp/awf-test-routing/output',
+      inputFile: '/tmp/awf-test-routing/input/conversation.json',
+      containerInputFile: '/run/awf-routing/input/conversation.json',
+      containerOutputDir: '/run/awf-routing/output',
+    };
+    const config: WrapperConfig = {
+      ...baseConfig,
+      enableApiProxy: true,
+      modelRouting: {
+        objective: { goal: 'cost', mode: 'balanced' },
+        task: { conversationFile: '/host/conversation.json' },
+      },
+    };
+    const dependencies = createOrderedWorkflowDependencies(callOrder, 0, {
+      prepareRouting: jest.fn().mockImplementation(async () => {
+        callOrder.push('prepareRouting');
+        return routingState;
+      }),
+      startContainers: jest.fn().mockImplementation(async (
+        _workDir,
+        _allowedDomains,
+        _proxyLogsDir,
+        _skipPull,
+        _onNetworkReady,
+        onInfrastructureReady,
+      ) => {
+        callOrder.push('startContainers');
+        await onInfrastructureReady?.();
+      }),
+      waitForRoutingSelection: jest.fn().mockImplementation(async () => {
+        callOrder.push('waitForRoutingSelection');
+      }),
+      verifyRoutingCompletion: jest.fn().mockImplementation(async () => {
+        callOrder.push('verifyRoutingCompletion');
+      }),
+      cleanupRouting: jest.fn().mockImplementation(async () => {
+        callOrder.push('cleanupRouting');
+      }),
+    });
+    const { logger, performCleanup } = createOrderedWorkflowOptions(callOrder);
+
+    const exitCode = await runMainWorkflow(config, dependencies, { logger, performCleanup });
+
+    expect(exitCode).toBe(0);
+    expect(callOrder).toEqual([
+      'prepareRouting',
+      'ensureFirewallNetwork',
+      'setupHostIptables',
+      'writeConfigs',
+      'startContainers',
+      'waitForRoutingSelection',
+      'runAgentCommand',
+      'performCleanup',
+      'verifyRoutingCompletion',
+      'cleanupRouting',
+    ]);
+    expect(dependencies.waitForRoutingSelection).toHaveBeenCalledWith(routingState);
+    expect(dependencies.verifyRoutingCompletion).toHaveBeenCalledWith(routingState);
+  });
+
   it('skips host network setup and iptables in network-isolation mode', async () => {
     const callOrder: string[] = [];
     const dependencies = {

src/cli-workflow.ts

@@ -7,6 +7,7 @@ import { buildInternalServiceHosts } from './services/internal-service-hosts';
 import { TOPOLOGY_NETWORK_NAME, getTopologyContainerIps, patchComposeWithTopologyHosts } from './topology';
 import { validateEnclavesConfig } from './enclave/preflight';
 import { isEnclaveAgentGithubRouteEnabled } from './types/enclave-options';
+import type { ModelRoutingBootstrapState } from './types';
 
 /**
  * Dependencies injected into the main workflow.
@@ -59,6 +60,14 @@ export interface WorkflowDependencies {
    * recovery, reconciliation, and the broker admission channel.
    */
   startEnclaveDynamicDelegation?: (config: WrapperConfig) => Promise<void>;
+  /** Stages the private routing conversation before any container exists. */
+  prepareRouting?: (config: WrapperConfig) => Promise<ModelRoutingBootstrapState | undefined>;
+  /** Waits for the proxy-owned routing selection before the primary agent starts. */
+  waitForRoutingSelection?: (state: ModelRoutingBootstrapState | undefined) => Promise<void>;
+  /** Verifies the proxy's end-of-run routing records after sidecar shutdown. */
+  verifyRoutingCompletion?: (state: ModelRoutingBootstrapState | undefined) => Promise<void>;
+  /** Removes private routing state after verification unless containers are kept. */
+  cleanupRouting?: (config: WrapperConfig) => Promise<void>;
 }
 
 interface WorkflowCallbacks {
@@ -115,6 +124,14 @@ export async function runMainWorkflow(
     logger.info('Staging enclave repository seeds...');
     await dependencies.prepareEnclaves(config);
   }
+  let routingState: ModelRoutingBootstrapState | undefined;
+  if (config.modelRouting) {
+    if (!dependencies.prepareRouting) {
+      throw new Error('Model routing is enabled but no staging implementation was provided to runMainWorkflow');
+    }
+    logger.info('Staging model routing conversation...');
+    routingState = await dependencies.prepareRouting(config);
+  }
 
   // Step 0: Setup host-level network and iptables
   //
@@ -217,7 +234,7 @@ export async function runMainWorkflow(
         }
       : undefined;
 
-  const onInfrastructureReady = config.enclaves?.enabled
+  const enclaveInfrastructureReady = config.enclaves?.enabled
     ? async () => {
         if (!dependencies.connectEnclaveGateway || !dependencies.assertEnclaveGatewayReady) {
           throw new Error('Enclaves require an exclusive MCP gateway readiness implementation');
@@ -252,6 +269,21 @@ export async function runMainWorkflow(
         }
       }
     : undefined;
+  const routingInfrastructureReady = config.modelRouting
+    ? async () => {
+        if (!dependencies.waitForRoutingSelection) {
+          throw new Error('Model routing is enabled but no selection wait implementation was provided');
+        }
+        logger.info('Waiting for model routing selection...');
+        await dependencies.waitForRoutingSelection(routingState);
+      }
+    : undefined;
+  const onInfrastructureReady = enclaveInfrastructureReady || routingInfrastructureReady
+    ? async () => {
+        if (enclaveInfrastructureReady) await enclaveInfrastructureReady();
+        if (routingInfrastructureReady) await routingInfrastructureReady();
+      }
+    : undefined;
 
   try {
     await dependencies.startContainers(
@@ -296,6 +328,16 @@ export async function runMainWorkflow(
 
   // Step 4: Cleanup (logs will be preserved automatically if they exist)
   await performCleanup();
+  if (config.modelRouting) {
+    if (!dependencies.verifyRoutingCompletion) {
+      throw new Error('Model routing is enabled but no completion verification implementation was provided');
+    }
+    try {
+      await dependencies.verifyRoutingCompletion(routingState);
+    } finally {
+      await dependencies.cleanupRouting?.(config);
+    }
+  }
 
   if (result.exitCode === 0) {
     logger.success('Command completed successfully');

src/commands/main-action.test.ts

@@ -46,6 +46,7 @@ import * as externalRuntimeResolver from '../external-runtime-backend-resolver';
 import { MAIN_ACTION_STUB_CONFIG, setupMainActionTestHarness } from './main-action.test-utils';
 import type { WrapperConfig } from '../types';
 import { CloudHypervisorUnsupportedHostError } from '../cloud-hypervisor/errors';
+import { RoutingFailureExitError } from '../routing/bootstrap';
 
 const {
   mkdirSync: mockMkdirSync,
@@ -863,6 +864,20 @@ describe('createMainAction', () => {
   });
 
   describe('fatal error cleanup after containers started', () => {
+    it('preserves the routing exit code through an infrastructure readiness error', async () => {
+      const routingFailure = new RoutingFailureExitError('Model routing selection timed out');
+      const readinessFailure = Object.assign(
+        new Error('Model routing selection timed out'),
+        { cause: routingFailure },
+      );
+      mockedCliWorkflow.runMainWorkflow.mockRejectedValueOnce(readinessFailure);
+
+      const action = createMainAction(getOptionValueSource);
+      await action(['echo hi'], {});
+
+      expect(processExitSpy).toHaveBeenCalledWith(78);
+    });
+
     it('stops containers during cleanup when workflow fails after startup callbacks', async () => {
       mockedCliWorkflow.runMainWorkflow.mockImplementation(
         async (_config, _deps, callbacks) => {

src/commands/main-action.ts

@@ -52,6 +52,13 @@ import {
   formatCloudHypervisorDockerFallbackWarning,
   isCloudHypervisorUnsupportedHostError,
 } from '../cloud-hypervisor/errors';
+import {
+  cleanupRoutingState,
+  RoutingFailureExitError,
+  stageRoutingConversation,
+  verifyRoutingCompletion,
+  waitForRoutingSelection,
+} from '../routing/bootstrap';
 
 const SENSITIVE_CONFIG_KEYS = new Set([
   'openaiApiKey',
@@ -67,6 +74,17 @@ const SENSITIVE_CONFIG_KEYS = new Set([
 
 const REFLECT_COMMAND = 'curl --fail --silent --show-error --noproxy "*" http://api-proxy:10000/reflect';
 
+function findRoutingFailure(error: unknown): RoutingFailureExitError | undefined {
+  const seen = new Set<unknown>();
+  let current = error;
+  while (current instanceof Error && !seen.has(current)) {
+    if (current instanceof RoutingFailureExitError) return current;
+    seen.add(current);
+    current = (current as Error & { cause?: unknown }).cause;
+  }
+  return undefined;
+}
+
 function redactConfigForLogging(config: WrapperConfig): Record<string, unknown> {
   const redactedConfig: Record<string, unknown> = {};
   for (const [key, value] of Object.entries(config)) {
@@ -407,6 +425,7 @@ export function createMainAction(getOptionValueSource: OptionSourceResolver) {
           : fastKillAgentContainer()
       ),
       performCleanup: (signal) => performCleanup(signal),
+      cleanupRouting: () => cleanupRoutingState(config),
     });
 
     if (externalRuntimeBackend) {
@@ -490,6 +509,10 @@ export function createMainAction(getOptionValueSource: OptionSourceResolver) {
         assertEnclaveGithubGatewayReady,
         prepareEnclaves,
         startEnclaveDynamicDelegation,
+        prepareRouting: async (routingConfig) => stageRoutingConversation(routingConfig),
+        waitForRoutingSelection,
+        verifyRoutingCompletion: async (routingState) => verifyRoutingCompletion(routingState),
+        cleanupRouting: async (routingConfig) => cleanupRoutingState(routingConfig),
       },
       {
         logger,
@@ -512,8 +535,10 @@ export function createMainAction(getOptionValueSource: OptionSourceResolver) {
       writeStartupFailureDiagnostic(config, error);
     }
     await performCleanup();
-    console.error(`Process exiting with code: 1`);
-    process.exit(1);
+    cleanupRoutingState(config);
+    const fatalExitCode = findRoutingFailure(error)?.exitCode ?? 1;
+    console.error(`Process exiting with code: ${fatalExitCode}`);
+    process.exit(fatalExitCode);
   }
   };
 }

src/commands/signal-handler.test.ts

@@ -16,24 +16,26 @@ describe('registerSignalHandlers', () => {
     containersStarted: boolean;
     keepContainers: boolean;
     fastKillRejects?: boolean;
-  }): Promise<{ fastKill: jest.Mock; performCleanup: jest.Mock }> {
+  }): Promise<{ fastKill: jest.Mock; performCleanup: jest.Mock; cleanupRouting: jest.Mock }> {
     const fastKill = fastKillRejects
       ? jest.fn().mockRejectedValue(new Error('kill failed'))
       : jest.fn().mockResolvedValue(undefined);
     const performCleanup = jest.fn().mockResolvedValue(undefined);
+    const cleanupRouting = jest.fn();
 
     const deps: SignalHandlerDependencies = {
       getContainersStarted: () => containersStarted,
       keepContainers,
       fastKillAgentContainer: fastKill,
       performCleanup,
+      cleanupRouting,
     };
 
     registerSignalHandlers(deps);
     harness.handlers[signal]();
     await flushPromises();
 
-    return { fastKill, performCleanup };
+    return { fastKill, performCleanup, cleanupRouting };
   }
 
   it('registers SIGINT and SIGTERM handlers', () => {
@@ -42,6 +44,7 @@ describe('registerSignalHandlers', () => {
       keepContainers: false,
       fastKillAgentContainer: jest.fn().mockResolvedValue(undefined),
       performCleanup: jest.fn().mockResolvedValue(undefined),
+      cleanupRouting: jest.fn(),
     };
 
     registerSignalHandlers(deps);
@@ -68,6 +71,34 @@ describe('registerSignalHandlers', () => {
     }
   );
 
+  it.each(['SIGINT', 'SIGTERM'] as const)(
+    'cleans private routing state on %s',
+    async signal => {
+      const { cleanupRouting } = await runSignalScenario({
+        signal,
+        containersStarted: true,
+        keepContainers: false,
+      });
+
+      expect(cleanupRouting).toHaveBeenCalledTimes(1);
+    },
+  );
+
+  it('exits with the signal status when routing cleanup fails', async () => {
+    registerSignalHandlers({
+      getContainersStarted: () => false,
+      keepContainers: false,
+      fastKillAgentContainer: jest.fn().mockResolvedValue(undefined),
+      performCleanup: jest.fn().mockResolvedValue(undefined),
+      cleanupRouting: jest.fn(() => { throw new Error('cleanup failed'); }),
+    });
+
+    harness.handlers.SIGTERM();
+    await flushPromises();
+
+    expect(harness.processExitSpy).toHaveBeenCalledWith(143);
+  });
+
   it('skips fast-kill on SIGINT when containers are not started', async () => {
     const { fastKill, performCleanup } = await runSignalScenario({
       signal: 'SIGINT',

src/commands/signal-handler.ts

@@ -10,6 +10,8 @@ interface SignalHandlerDependencies {
   fastKillAgentContainer: () => Promise<void>;
   /** Runs the full cleanup sequence (stop containers, remove host iptables rules, etc.). */
   performCleanup: (signal?: string) => Promise<void>;
+  /** Removes private state created for model routing. */
+  cleanupRouting: () => void;
 }
 
 /**
@@ -25,6 +27,7 @@ export function registerSignalHandlers({
   keepContainers,
   fastKillAgentContainer,
   performCleanup,
+  cleanupRouting,
 }: SignalHandlerDependencies): void {
   process.on('SIGINT', () => {
     (async () => {
@@ -34,6 +37,11 @@ export function registerSignalHandlers({
         }
         await performCleanup('SIGINT');
       } finally {
+        try {
+          cleanupRouting();
+        } catch {
+          // Cleanup failure must not change the signal exit status.
+        }
         console.error(`Process exiting with code: 130`);
         process.exit(130); // Standard exit code for SIGINT
       }
@@ -48,6 +56,11 @@ export function registerSignalHandlers({
         }
         await performCleanup('SIGTERM');
       } finally {
+        try {
+          cleanupRouting();
+        } catch {
+          // Cleanup failure must not change the signal exit status.
+        }
         console.error(`Process exiting with code: 143`);
         process.exit(143); // Standard exit code for SIGTERM
       }