feat(worker): migrate agent runtime from Claude Agent SDK to pi harness

apps/worker/src/ai/claude-executor.ts

@@ -4,28 +4,80 @@
 // it under the terms of the GNU Affero General Public License version 3
 // as published by the Free Software Foundation.
 
-// Production Claude agent execution with retry, git checkpoints, and audit logging
+// Production agent execution on the pi harness, with git checkpoints and audit logging.
 
-import { type JsonSchemaOutputFormat, query } from '@anthropic-ai/claude-agent-sdk';
+import { createRequire } from 'node:module';
+import type { AgentMessage } from '@earendil-works/pi-agent-core';
+import {
+  type AgentSessionEvent,
+  createAgentSession,
+  DefaultResourceLoader,
+  getAgentDir,
+  ModelRegistry,
+  type ResourceLoader,
+  SessionManager,
+  SettingsManager,
+  type ToolDefinition,
+} from '@earendil-works/pi-coding-agent';
 import { fs, path } from 'zx';
 import type { AuditSession } from '../audit/index.js';
 import { deliverablesDir } from '../paths.js';
 import { isRetryableError, PentestError } from '../services/error-handling.js';
 import { AGENT_VALIDATORS } from '../session-manager.js';
 import type { ActivityLogger } from '../types/activity-logger.js';
-import { isSpendingCapBehavior } from '../utils/billing-detection.js';
+import { ErrorCode } from '../types/errors.js';
+import { isSpendingCapBehavior, matchesBillingTextPattern } from '../utils/billing-detection.js';
 import { formatTimestamp } from '../utils/formatting.js';
 import { Timer } from '../utils/metrics.js';
 import { createAuditLogger } from './audit-logger.js';
-import { dispatchMessage } from './message-handlers.js';
-import { type ModelTier, resolveModel, supportsAdaptiveThinking } from './models.js';
+import { type ModelTier, resolveModelSelection } from './models.js';
 import { detectExecutionContext, formatCompletionMessage, formatErrorOutput } from './output-formatters.js';
 import { createProgressManager } from './progress-manager.js';
+import { permissionConfigPath } from './settings-writer.js';
+import { createTaskTool, createTodoWriteTool } from './tools.js';
 
 declare global {
   var SHANNON_DISABLE_LOADER: boolean | undefined;
 }
 
+/** Built-in pi tools enabled for every agent (custom tool names are appended). */
+const BUILTIN_TOOLS = ['read', 'bash', 'edit', 'write', 'grep', 'find', 'ls'];
+
+const requireFromHere = createRequire(import.meta.url);
+let cachedExtensionDir: string | null | undefined;
+
+/** Resolve the installed @gotgenes/pi-permission-system package dir, or null. */
+function permissionExtensionDir(): string | null {
+  if (cachedExtensionDir !== undefined) return cachedExtensionDir;
+  try {
+    const entry = requireFromHere.resolve('@gotgenes/pi-permission-system');
+    cachedExtensionDir = path.dirname(path.dirname(entry));
+  } catch {
+    cachedExtensionDir = null;
+  }
+  return cachedExtensionDir;
+}
+
+/**
+ * Build a resource loader that loads the pi-permission-system extension — but only
+ * when a code_path deny config exists (written by settings-writer). Returns
+ * undefined otherwise, preserving default behavior (and zero overhead) for runs
+ * with no code_path avoids.
+ */
+async function buildPermissionResourceLoader(cwd: string, logger: ActivityLogger): Promise<ResourceLoader | undefined> {
+  if (!fs.existsSync(permissionConfigPath())) return undefined;
+  const extDir = permissionExtensionDir();
+  if (!extDir) {
+    logger.warn(
+      'code_path deny config present but @gotgenes/pi-permission-system not resolvable — skipping enforcement',
+    );
+    return undefined;
+  }
+  const loader = new DefaultResourceLoader({ cwd, agentDir: getAgentDir(), additionalExtensionPaths: [extDir] });
+  await loader.reload();
+  return loader;
+}
+
 export interface ClaudePromptResult {
   result?: string | null | undefined;
   success: boolean;
@@ -58,18 +110,8 @@ async function writeErrorLog(
     const errorLog = {
       timestamp: formatTimestamp(),
       agent: 'claude-executor',
-      error: {
-        name: err.constructor.name,
-        message: err.message,
-        code: err.code,
-        status: err.status,
-        stack: err.stack,
-      },
-      context: {
-        sourceDir,
-        prompt: `${fullPrompt.slice(0, 200)}...`,
-        retryable: isRetryableError(err),
-      },
+      error: { name: err.constructor.name, message: err.message, code: err.code, status: err.status, stack: err.stack },
+      context: { sourceDir, prompt: `${fullPrompt.slice(0, 200)}...`, retryable: isRetryableError(err) },
       duration,
     };
     const logPath = path.join(deliverablesDir(sourceDir), 'error.log');
@@ -86,34 +128,23 @@ export async function validateAgentOutput(
   logger: ActivityLogger,
 ): Promise<boolean> {
   logger.info(`Validating ${agentName} agent output`);
-
   try {
-    // Check if agent completed successfully (text result OR structured output)
     if (!result.success || (!result.result && result.structuredOutput === undefined)) {
       logger.error('Validation failed: Agent execution was unsuccessful');
       return false;
     }
-
-    // Get validator function for this agent
     const validator = agentName ? AGENT_VALIDATORS[agentName as keyof typeof AGENT_VALIDATORS] : undefined;
-
     if (!validator) {
       logger.warn(`No validator found for agent "${agentName}" - assuming success`);
-      logger.info('Validation passed: Unknown agent with successful result');
       return true;
     }
-
     logger.info(`Using validator for agent: ${agentName}`, { sourceDir });
-
-    // Apply validation function
     const validationResult = await validator(sourceDir, logger);
-
     if (validationResult) {
       logger.info('Validation passed: Required files/structure present');
     } else {
       logger.error('Validation failed: Missing required deliverable files');
     }
-
     return validationResult;
   } catch (error) {
     const errMsg = error instanceof Error ? error.message : String(error);
@@ -122,8 +153,40 @@ export async function validateAgentOutput(
   }
 }
 
-// Low-level SDK execution. Handles message streaming, progress, and audit logging.
-// Exported for Temporal activities to call single-attempt execution.
+/** Concatenate the text blocks of an assistant message (skips thinking + tool calls). */
+function extractAssistantText(message: AgentMessage): string {
+  if (message.role !== 'assistant') return '';
+  const blocks = message.content as Array<{ type: string; text?: string }>;
+  return blocks
+    .filter((c) => c.type === 'text')
+    .map((c) => c.text ?? '')
+    .join('\n');
+}
+
+/**
+ * Classify error-bearing text into a PentestError, mirroring the prior SDK error
+ * handling. Spending-cap / billing text is retryable (Temporal backs off and
+ * recovers when the cap resets); session limit is permanent.
+ */
+function classifyErrorText(content: string): PentestError | null {
+  if (!content) return null;
+  if (matchesBillingTextPattern(content)) {
+    return new PentestError(
+      `Billing limit reached: ${content.slice(0, 100)}`,
+      'billing',
+      true,
+      {},
+      ErrorCode.SPENDING_CAP_REACHED,
+    );
+  }
+  if (content.toLowerCase().includes('session limit reached')) {
+    return new PentestError('Session limit reached', 'billing', false);
+  }
+  return null;
+}
+
+// Low-level pi execution. Drives one agent session to completion with progress and
+// audit logging. Exported for Temporal activities to call single-attempt execution.
 export async function runClaudePrompt(
   prompt: string,
   sourceDir: string,
@@ -133,11 +196,10 @@ export async function runClaudePrompt(
   auditSession: AuditSession | null = null,
   logger: ActivityLogger,
   modelTier: ModelTier = 'medium',
-  outputFormat?: JsonSchemaOutputFormat,
+  callerTools?: ToolDefinition[],
   apiKey?: string,
   deliverablesSubdir?: string,
   providerConfig?: import('../types/config.js').ProviderConfig,
-  mcpServers?: Record<string, import('@anthropic-ai/claude-agent-sdk').McpServerConfig>,
 ): Promise<ClaudePromptResult> {
   // 1. Initialize timing and prompt
   const timer = new Timer(`agent-${description.toLowerCase().replace(/\s+/g, '-')}`);
@@ -151,132 +213,112 @@ export async function runClaudePrompt(
   );
   const auditLogger = createAuditLogger(auditSession);
 
-  logger.info(`Running Claude Code: ${description}...`);
+  logger.info(`Running pi agent: ${description}...`);
 
-  // 3. Build env vars to pass to SDK subprocesses
-  const sdkEnv: Record<string, string> = {
-    CLAUDE_CODE_MAX_OUTPUT_TOKENS: process.env.CLAUDE_CODE_MAX_OUTPUT_TOKENS || '64000',
-    PLAYWRIGHT_MCP_OUTPUT_DIR: deliverablesSubdir
-      ? path.join(sourceDir, path.dirname(deliverablesSubdir), '.playwright-cli')
-      : path.join(sourceDir, '.shannon', '.playwright-cli'),
-    // apiKey from ContainerConfig takes precedence over process.env
-    ...(apiKey && { ANTHROPIC_API_KEY: apiKey }),
-    // Deliverables subdir for save-deliverable CLI tool
-    ...(deliverablesSubdir && { SHANNON_DELIVERABLES_SUBDIR: deliverablesSubdir }),
-  };
+  // 3. Expose bash-invoked CLI tooling (playwright-cli, save-deliverable) to the
+  //    environment pi's bash tool inherits. These are constant per container, so
+  //    setting them on process.env is parallel-safe across this workflow's agents.
+  process.env.PLAYWRIGHT_MCP_OUTPUT_DIR = deliverablesSubdir
+    ? path.join(sourceDir, path.dirname(deliverablesSubdir), '.playwright-cli')
+    : path.join(sourceDir, '.shannon', '.playwright-cli');
+  if (deliverablesSubdir) process.env.SHANNON_DELIVERABLES_SUBDIR = deliverablesSubdir;
+  if (apiKey) process.env.ANTHROPIC_API_KEY = apiKey;
 
-  // 3a. Apply structured provider config directly to sdkEnv (no process.env mutation)
-  if (providerConfig) {
-    switch (providerConfig.providerType) {
-      case 'bedrock':
-        sdkEnv.CLAUDE_CODE_USE_BEDROCK = '1';
-        if (providerConfig.awsRegion) sdkEnv.AWS_REGION = providerConfig.awsRegion;
-        if (providerConfig.awsAccessKeyId) sdkEnv.AWS_ACCESS_KEY_ID = providerConfig.awsAccessKeyId;
-        if (providerConfig.awsSecretAccessKey) sdkEnv.AWS_SECRET_ACCESS_KEY = providerConfig.awsSecretAccessKey;
-        break;
-      case 'vertex':
-        sdkEnv.CLAUDE_CODE_USE_VERTEX = '1';
-        if (providerConfig.gcpRegion) sdkEnv.CLOUD_ML_REGION = providerConfig.gcpRegion;
-        if (providerConfig.gcpProjectId) sdkEnv.ANTHROPIC_VERTEX_PROJECT_ID = providerConfig.gcpProjectId;
-        if (providerConfig.gcpCredentialsPath)
-          sdkEnv.GOOGLE_APPLICATION_CREDENTIALS = providerConfig.gcpCredentialsPath;
-        break;
-      case 'litellm_router':
-        if (providerConfig.baseUrl) sdkEnv.ANTHROPIC_BASE_URL = providerConfig.baseUrl;
-        if (providerConfig.authToken) sdkEnv.ANTHROPIC_AUTH_TOKEN = providerConfig.authToken;
-        break;
-      default:
-        // 'anthropic_api' or unset — apiKey already handled above
-        if (providerConfig.apiKey && !apiKey) sdkEnv.ANTHROPIC_API_KEY = providerConfig.apiKey;
-        break;
-    }
-  }
-
-  // 3b. Passthrough env vars not already set by providerConfig or apiKey
-  const passthroughVars = [
-    ...(!sdkEnv.ANTHROPIC_API_KEY ? ['ANTHROPIC_API_KEY'] : []),
-    'CLAUDE_CODE_OAUTH_TOKEN',
-    ...(!sdkEnv.ANTHROPIC_BASE_URL ? ['ANTHROPIC_BASE_URL'] : []),
-    ...(!sdkEnv.ANTHROPIC_AUTH_TOKEN ? ['ANTHROPIC_AUTH_TOKEN'] : []),
-    ...(!sdkEnv.CLAUDE_CODE_USE_BEDROCK ? ['CLAUDE_CODE_USE_BEDROCK'] : []),
-    ...(!sdkEnv.AWS_REGION ? ['AWS_REGION'] : []),
-    'AWS_BEARER_TOKEN_BEDROCK',
-    ...(!sdkEnv.CLAUDE_CODE_USE_VERTEX ? ['CLAUDE_CODE_USE_VERTEX'] : []),
-    ...(!sdkEnv.CLOUD_ML_REGION ? ['CLOUD_ML_REGION'] : []),
-    ...(!sdkEnv.ANTHROPIC_VERTEX_PROJECT_ID ? ['ANTHROPIC_VERTEX_PROJECT_ID'] : []),
-    ...(!sdkEnv.GOOGLE_APPLICATION_CREDENTIALS ? ['GOOGLE_APPLICATION_CREDENTIALS'] : []),
-    'HOME',
-    'PATH',
-    'PLAYWRIGHT_MCP_EXECUTABLE_PATH',
+  // 4. Resolve model + auth, then assemble the tool set (universal task/todo tools
+  //    plus any caller-supplied collector/submit tools).
+  const selection = resolveModelSelection((auth) => ModelRegistry.create(auth), modelTier, apiKey, providerConfig);
+  // Load the code_path deny extension only when a deny config was written; the same
+  // loader is reused by child task sessions so they inherit the policy.
+  const resourceLoader = await buildPermissionResourceLoader(sourceDir, logger);
+  const customTools: ToolDefinition[] = [
+    createTaskTool({
+      model: selection.model,
+      thinkingLevel: selection.thinkingLevel,
+      authStorage: selection.authStorage,
+      cwd: sourceDir,
+      ...(resourceLoader && { resourceLoader }),
+    }),
+    createTodoWriteTool(auditLogger),
+    ...(callerTools ?? []),
   ];
-  for (const name of passthroughVars) {
-    const val = process.env[name];
-    if (val) {
-      sdkEnv[name] = val;
-    }
-  }
-
-  // 4. Configure SDK options
-  // Model override from providerConfig takes precedence over env-based resolveModel
-  const model = providerConfig?.modelOverrides?.[modelTier] ?? resolveModel(modelTier);
-  const adaptiveThinking = supportsAdaptiveThinking(model) && process.env.CLAUDE_ADAPTIVE_THINKING !== 'false';
-  const options = {
-    model,
-    maxTurns: 10_000,
-    cwd: sourceDir,
-    permissionMode: 'bypassPermissions' as const,
-    allowDangerouslySkipPermissions: true,
-    settingSources: ['user'] as ('user' | 'project' | 'local')[],
-    env: sdkEnv,
-    ...(adaptiveThinking && { thinking: { type: 'adaptive' as const } }),
-    ...(outputFormat && { outputFormat }),
-    ...(mcpServers && Object.keys(mcpServers).length > 0 && { mcpServers }),
-  };
-
-  if (!execContext.useCleanOutput) {
-    logger.info(`SDK Options: maxTurns=${options.maxTurns}, cwd=${sourceDir}, permissions=BYPASS`);
-  }
+  // pi's `tools` allowlist gates custom tools too — list every custom name.
+  const tools = [...BUILTIN_TOOLS, ...customTools.map((t) => t.name)];
 
   let turnCount = 0;
-  let result: string | null = null;
+  let pendingError: PentestError | null = null;
   let apiErrorDetected = false;
-  let totalCost = 0;
 
   progress.start();
 
   try {
-    // 6. Process the message stream
-    const messageLoopResult = await processMessageStream(
-      fullPrompt,
-      options,
-      { execContext, description, progress, auditLogger, logger },
-      timer,
-    );
+    const { session } = await createAgentSession({
+      cwd: sourceDir,
+      model: selection.model,
+      thinkingLevel: selection.thinkingLevel,
+      tools,
+      customTools,
+      authStorage: selection.authStorage,
+      sessionManager: SessionManager.inMemory(),
+      // Temporal owns retry; pi compaction stays on (no analog previously, guards
+      // against context overflow on long agent runs).
+      settingsManager: SettingsManager.inMemory({ retry: { enabled: false }, compaction: { enabled: true } }),
+      ...(resourceLoader && { resourceLoader }),
+    });
 
-    turnCount = messageLoopResult.turnCount;
-    result = messageLoopResult.result;
-    apiErrorDetected = messageLoopResult.apiErrorDetected;
-    totalCost = messageLoopResult.cost;
-    const model = messageLoopResult.model;
+    // 5. Map pi events to audit logging + progress + error capture.
+    session.subscribe((event: AgentSessionEvent) => {
+      switch (event.type) {
+        case 'turn_end': {
+          turnCount += 1;
+          const msg = event.message;
+          const text = extractAssistantText(msg);
+          if (text.trim()) {
+            void auditLogger.logLlmResponse(turnCount, text);
+            const billing = classifyErrorText(text);
+            if (billing) pendingError = billing;
+          }
+          if (msg.role === 'assistant' && msg.stopReason === 'error') {
+            apiErrorDetected = true;
+            pendingError =
+              pendingError ??
+              classifyErrorText(msg.errorMessage ?? '') ??
+              new PentestError(`Agent error: ${(msg.errorMessage ?? 'unknown').slice(0, 200)}`, 'unknown', true);
+          }
+          break;
+        }
+        case 'tool_execution_start':
+          void auditLogger.logToolStart(event.toolName, event.args);
+          break;
+        case 'tool_execution_end':
+          void auditLogger.logToolEnd(event.result);
+          break;
+        default:
+          break;
+      }
+    });
 
-    // === SPENDING CAP SAFEGUARD ===
-    // 7. Defense-in-depth: Detect spending cap that slipped through detectApiError().
-    // Uses consolidated billing detection from utils/billing-detection.ts
+    // 6. Run the agent to completion (resolves at agent_end).
+    await session.prompt(fullPrompt);
+    session.dispose();
+
+    // 7. Surface any error captured during the run.
+    if (pendingError) throw pendingError;
+
+    // 8. Read usage/cost and final text.
+    const stats = session.getSessionStats();
+    const totalCost = stats.cost;
+    const result = session.getLastAssistantText() ?? null;
+
+    // 9. Defense-in-depth: detect a spending cap that produced an empty/cheap run.
     if (isSpendingCapBehavior(turnCount, totalCost, result || '')) {
       throw new PentestError(
         `Spending cap likely reached (turns=${turnCount}, cost=$0): ${result?.slice(0, 100)}`,
         'billing',
-        true, // Retryable - Temporal will use 5-30 min backoff
+        true,
       );
     }
 
-    // 8. Finalize successful result
     const duration = timer.stop();
-
-    if (apiErrorDetected) {
-      logger.warn(`API Error detected in ${description} - will validate deliverables before failing`);
-    }
-
     progress.finish(formatCompletionMessage(execContext, description, turnCount, duration));
 
     return {
@@ -285,19 +327,14 @@ export async function runClaudePrompt(
       duration,
       turns: turnCount,
       cost: totalCost,
-      model,
+      model: selection.model.id,
       partialCost: totalCost,
       apiErrorDetected,
-      ...(messageLoopResult.structuredOutput !== undefined && {
-        structuredOutput: messageLoopResult.structuredOutput,
-      }),
     };
   } catch (error) {
-    // 9. Handle errors — log, write error file, return failure
+    // 10. Handle errors — log, write error file, return failure
     const duration = timer.stop();
-
     const err = error as Error & { code?: string; status?: number };
-
     await auditLogger.logError(err, duration, turnCount);
     progress.stop();
     outputLines(formatErrorOutput(err, execContext, description, duration, sourceDir, isRetryableError(err)));
@@ -309,96 +346,8 @@ export async function runClaudePrompt(
       prompt: `${fullPrompt.slice(0, 100)}...`,
       success: false,
       duration,
-      cost: totalCost,
+      cost: 0,
       retryable: isRetryableError(err),
     };
   }
 }
-
-interface MessageLoopResult {
-  turnCount: number;
-  result: string | null;
-  apiErrorDetected: boolean;
-  cost: number;
-  model?: string | undefined;
-  structuredOutput?: unknown;
-}
-
-interface MessageLoopDeps {
-  execContext: ReturnType<typeof detectExecutionContext>;
-  description: string;
-  progress: ReturnType<typeof createProgressManager>;
-  auditLogger: ReturnType<typeof createAuditLogger>;
-  logger: ActivityLogger;
-}
-
-async function processMessageStream(
-  fullPrompt: string,
-  options: NonNullable<Parameters<typeof query>[0]['options']>,
-  deps: MessageLoopDeps,
-  timer: Timer,
-): Promise<MessageLoopResult> {
-  const { execContext, description, progress, auditLogger, logger } = deps;
-  const HEARTBEAT_INTERVAL = 30000;
-
-  let turnCount = 0;
-  let result: string | null = null;
-  let apiErrorDetected = false;
-  let cost = 0;
-  let model: string | undefined;
-  let structuredOutput: unknown | undefined;
-  let lastHeartbeat = Date.now();
-
-  for await (const message of query({ prompt: fullPrompt, options })) {
-    // Heartbeat logging when loader is disabled
-    const now = Date.now();
-    if (global.SHANNON_DISABLE_LOADER && now - lastHeartbeat > HEARTBEAT_INTERVAL) {
-      logger.info(`[${Math.floor((now - timer.startTime) / 1000)}s] ${description} running... (Turn ${turnCount})`);
-      lastHeartbeat = now;
-    }
-
-    // Increment turn count for assistant messages
-    if (message.type === 'assistant') {
-      turnCount++;
-    }
-
-    const dispatchResult = await dispatchMessage(message as { type: string; subtype?: string }, turnCount, {
-      execContext,
-      description,
-      progress,
-      auditLogger,
-      logger,
-    });
-
-    if (dispatchResult.type === 'throw') {
-      throw dispatchResult.error;
-    }
-
-    if (dispatchResult.type === 'complete') {
-      result = dispatchResult.result;
-      cost = dispatchResult.cost;
-      if (dispatchResult.structuredOutput !== undefined) {
-        structuredOutput = dispatchResult.structuredOutput;
-      }
-      break;
-    }
-
-    if (dispatchResult.type === 'continue') {
-      if (dispatchResult.apiErrorDetected) {
-        apiErrorDetected = true;
-      }
-      if (dispatchResult.model) {
-        model = dispatchResult.model;
-      }
-    }
-  }
-
-  return {
-    turnCount,
-    result,
-    apiErrorDetected,
-    cost,
-    model,
-    ...(structuredOutput !== undefined && { structuredOutput }),
-  };
-}

apps/worker/src/ai/message-handlers.ts

@@ -1,408 +0,0 @@
-// Copyright (C) 2025 Keygraph, Inc.
-//
-// This program is free software: you can redistribute it and/or modify
-// it under the terms of the GNU Affero General Public License version 3
-// as published by the Free Software Foundation.
-
-import type { SDKAssistantMessageError } from '@anthropic-ai/claude-agent-sdk';
-import { PentestError } from '../services/error-handling.js';
-import type { ActivityLogger } from '../types/activity-logger.js';
-import { ErrorCode } from '../types/errors.js';
-import { matchesBillingTextPattern } from '../utils/billing-detection.js';
-import { formatTimestamp } from '../utils/formatting.js';
-import type { AuditLogger } from './audit-logger.js';
-import {
-  filterJsonToolCalls,
-  formatAssistantOutput,
-  formatResultOutput,
-  formatToolResultOutput,
-  formatToolUseOutput,
-} from './output-formatters.js';
-import type { ProgressManager } from './progress-manager.js';
-import type {
-  ApiErrorDetection,
-  AssistantMessage,
-  AssistantResult,
-  ContentBlock,
-  ExecutionContext,
-  ModelRefusalFallbackMessage,
-  ResultData,
-  ResultMessage,
-  SystemInitMessage,
-  ToolResultData,
-  ToolResultMessage,
-  ToolUseData,
-  ToolUseMessage,
-} from './types.js';
-
-// Handles both array and string content formats from SDK
-function extractMessageContent(message: AssistantMessage): string {
-  const messageContent = message.message;
-
-  if (Array.isArray(messageContent.content)) {
-    return messageContent.content
-      .filter((c: ContentBlock) => c.type !== 'thinking' && c.type !== 'redacted_thinking')
-      .map((c: ContentBlock) => c.text || JSON.stringify(c))
-      .join('\n');
-  }
-
-  return String(messageContent.content);
-}
-
-// Extracts only text content (no tool_use JSON) to avoid false positives in error detection
-function extractTextOnlyContent(message: AssistantMessage): string {
-  const messageContent = message.message;
-
-  if (Array.isArray(messageContent.content)) {
-    return messageContent.content
-      .filter((c: ContentBlock) => c.type === 'text' || c.text)
-      .map((c: ContentBlock) => c.text || '')
-      .join('\n');
-  }
-
-  return String(messageContent.content);
-}
-
-function detectApiError(content: string): ApiErrorDetection {
-  if (!content || typeof content !== 'string') {
-    return { detected: false };
-  }
-
-  const lowerContent = content.toLowerCase();
-
-  // === BILLING/SPENDING CAP ERRORS (Retryable with long backoff) ===
-  // When Claude Code hits its spending cap, it returns a short message like
-  // "Spending cap reached resets 8am" instead of throwing an error.
-  // These should retry with 5-30 min backoff so workflows can recover when cap resets.
-  if (matchesBillingTextPattern(content)) {
-    return {
-      detected: true,
-      shouldThrow: new PentestError(
-        `Billing limit reached: ${content.slice(0, 100)}`,
-        'billing',
-        true, // RETRYABLE - Temporal will use 5-30 min backoff
-        {},
-        ErrorCode.SPENDING_CAP_REACHED,
-      ),
-    };
-  }
-
-  // === SESSION LIMIT (Non-retryable) ===
-  // Different from spending cap - usually means something is fundamentally wrong
-  if (lowerContent.includes('session limit reached')) {
-    return {
-      detected: true,
-      shouldThrow: new PentestError('Session limit reached', 'billing', false),
-    };
-  }
-
-  // Non-fatal API errors - detected but continue
-  if (lowerContent.includes('api error') || lowerContent.includes('terminated')) {
-    return { detected: true };
-  }
-
-  return { detected: false };
-}
-
-// Maps SDK structured error types to our error handling.
-function handleStructuredError(errorType: SDKAssistantMessageError, content: string): ApiErrorDetection {
-  switch (errorType) {
-    case 'billing_error':
-      return {
-        detected: true,
-        shouldThrow: new PentestError(
-          `Billing error (structured): ${content.slice(0, 100)}`,
-          'billing',
-          true, // Retryable with backoff
-          {},
-          ErrorCode.INSUFFICIENT_CREDITS,
-        ),
-      };
-    case 'rate_limit':
-      return {
-        detected: true,
-        shouldThrow: new PentestError(
-          `Rate limit hit (structured): ${content.slice(0, 100)}`,
-          'network',
-          true, // Retryable with backoff
-          {},
-          ErrorCode.API_RATE_LIMITED,
-        ),
-      };
-    case 'authentication_failed':
-      return {
-        detected: true,
-        shouldThrow: new PentestError(
-          `Authentication failed: ${content.slice(0, 100)}`,
-          'config',
-          false, // Not retryable - needs API key fix
-        ),
-      };
-    case 'server_error':
-      return {
-        detected: true,
-        shouldThrow: new PentestError(
-          `Server error (structured): ${content.slice(0, 100)}`,
-          'network',
-          true, // Retryable
-        ),
-      };
-    case 'invalid_request':
-      return {
-        detected: true,
-        shouldThrow: new PentestError(
-          `Invalid request: ${content.slice(0, 100)}`,
-          'config',
-          false, // Not retryable - needs code fix
-        ),
-      };
-    case 'max_output_tokens':
-      return {
-        detected: true,
-        shouldThrow: new PentestError(
-          `Max output tokens reached: ${content.slice(0, 100)}`,
-          'billing',
-          true, // Retryable - may succeed with different content
-        ),
-      };
-    case 'overloaded':
-      return {
-        detected: true,
-        shouldThrow: new PentestError(
-          `Anthropic API overloaded (structured): ${content.slice(0, 100)}`,
-          'network',
-          true, // Retryable with backoff
-        ),
-      };
-    case 'model_not_found':
-      return {
-        detected: true,
-        shouldThrow: new PentestError(
-          `Model not found: ${content.slice(0, 100)}`,
-          'config',
-          false, // Not retryable - model ID is misconfigured
-        ),
-      };
-    case 'oauth_org_not_allowed':
-      return {
-        detected: true,
-        shouldThrow: new PentestError(
-          `Organization not allowed for this credential: ${content.slice(0, 100)}`,
-          'config',
-          false, // Not retryable - needs credential/org fix
-        ),
-      };
-    default:
-      return { detected: true };
-  }
-}
-
-function handleAssistantMessage(message: AssistantMessage, turnCount: number): AssistantResult {
-  const content = extractMessageContent(message);
-  const cleanedContent = filterJsonToolCalls(content);
-
-  // Prefer structured error field from SDK, fall back to text-sniffing
-  // Use text-only content for error detection to avoid false positives
-  // from tool_use JSON (e.g. security reports containing "usage limit")
-  let errorDetection: ApiErrorDetection;
-  if (message.error) {
-    errorDetection = handleStructuredError(message.error, content);
-  } else {
-    const textOnlyContent = extractTextOnlyContent(message);
-    errorDetection = detectApiError(textOnlyContent);
-  }
-
-  const result: AssistantResult = {
-    content,
-    cleanedContent,
-    apiErrorDetected: errorDetection.detected,
-    logData: {
-      turn: turnCount,
-      content,
-      timestamp: formatTimestamp(),
-    },
-  };
-
-  // Only add shouldThrow if it exists (exactOptionalPropertyTypes compliance)
-  if (errorDetection.shouldThrow) {
-    result.shouldThrow = errorDetection.shouldThrow;
-  }
-
-  return result;
-}
-
-// Final message of a query with cost/duration info
-function handleResultMessage(message: ResultMessage): ResultData {
-  const result: ResultData = {
-    result: message.result || null,
-    cost: message.total_cost_usd || 0,
-    duration_ms: message.duration_ms || 0,
-    permissionDenials: message.permission_denials?.length || 0,
-  };
-
-  // Only add subtype if it exists (exactOptionalPropertyTypes compliance)
-  if (message.subtype) {
-    result.subtype = message.subtype;
-  }
-
-  // Capture stop_reason for diagnostics (helps debug early stops, budget exceeded, etc.)
-  if (message.stop_reason !== undefined) {
-    result.stop_reason = message.stop_reason;
-    if (message.stop_reason && message.stop_reason !== 'end_turn') {
-      console.log(`    Stop reason: ${message.stop_reason}`);
-    }
-  }
-
-  if (message.structured_output !== undefined) {
-    result.structuredOutput = message.structured_output;
-  }
-
-  return result;
-}
-
-function handleToolUseMessage(message: ToolUseMessage): ToolUseData {
-  return {
-    toolName: message.name,
-    parameters: message.input || {},
-    timestamp: formatTimestamp(),
-  };
-}
-
-// Truncates long results for display (500 char limit), preserves full content for logging
-function handleToolResultMessage(message: ToolResultMessage): ToolResultData {
-  const content = message.content;
-  const contentStr = typeof content === 'string' ? content : JSON.stringify(content, null, 2);
-
-  const displayContent =
-    contentStr.length > 500
-      ? `${contentStr.slice(0, 500)}...\n[Result truncated - ${contentStr.length} total chars]`
-      : contentStr;
-
-  return {
-    content,
-    displayContent,
-    timestamp: formatTimestamp(),
-  };
-}
-
-function outputLines(lines: string[]): void {
-  for (const line of lines) {
-    console.log(line);
-  }
-}
-
-export type MessageDispatchAction =
-  | { type: 'continue'; apiErrorDetected?: boolean | undefined; model?: string | undefined }
-  | { type: 'complete'; result: string | null; cost: number; structuredOutput?: unknown }
-  | { type: 'throw'; error: Error };
-
-export interface MessageDispatchDeps {
-  execContext: ExecutionContext;
-  description: string;
-  progress: ProgressManager;
-  auditLogger: AuditLogger;
-  logger: ActivityLogger;
-}
-
-// Dispatches SDK messages to appropriate handlers and formatters
-export async function dispatchMessage(
-  message: { type: string; subtype?: string },
-  turnCount: number,
-  deps: MessageDispatchDeps,
-): Promise<MessageDispatchAction> {
-  const { execContext, description, progress, auditLogger, logger } = deps;
-
-  switch (message.type) {
-    case 'assistant': {
-      const assistantResult = handleAssistantMessage(message as AssistantMessage, turnCount);
-
-      if (assistantResult.shouldThrow) {
-        return { type: 'throw', error: assistantResult.shouldThrow };
-      }
-
-      if (assistantResult.cleanedContent.trim()) {
-        progress.stop();
-        outputLines(formatAssistantOutput(assistantResult.cleanedContent, execContext, turnCount, description));
-        progress.start();
-      }
-
-      await auditLogger.logLlmResponse(turnCount, assistantResult.content);
-
-      if (assistantResult.apiErrorDetected) {
-        logger.warn('API Error detected in assistant response');
-        return { type: 'continue', apiErrorDetected: true };
-      }
-
-      return { type: 'continue' };
-    }
-
-    case 'system': {
-      if (message.subtype === 'init') {
-        const initMsg = message as SystemInitMessage;
-        if (!execContext.useCleanOutput) {
-          logger.info(`Model: ${initMsg.model}, Permission: ${initMsg.permissionMode}`);
-        }
-        return { type: 'continue', model: initMsg.model };
-      }
-      if (message.subtype === 'model_refusal_fallback') {
-        const fallback = message as ModelRefusalFallbackMessage;
-        const category = fallback.api_refusal_category ?? 'policy';
-        await auditLogger.logNote(
-          'model-fallback',
-          `Model refused (${category}); fell back ${fallback.original_model} → ${fallback.fallback_model}`,
-        );
-        return { type: 'continue' };
-      }
-      return { type: 'continue' };
-    }
-
-    case 'user':
-    case 'tool_progress':
-    case 'tool_use_summary':
-    case 'auth_status':
-      return { type: 'continue' };
-
-    case 'tool_use': {
-      const toolData = handleToolUseMessage(message as unknown as ToolUseMessage);
-      outputLines(formatToolUseOutput(toolData.toolName, toolData.parameters));
-      await auditLogger.logToolStart(toolData.toolName, toolData.parameters);
-      return { type: 'continue' };
-    }
-
-    case 'tool_result': {
-      const toolResultData = handleToolResultMessage(message as unknown as ToolResultMessage);
-      outputLines(formatToolResultOutput(toolResultData.displayContent));
-      await auditLogger.logToolEnd(toolResultData.content);
-      return { type: 'continue' };
-    }
-
-    case 'result': {
-      const resultData = handleResultMessage(message as ResultMessage);
-      outputLines(formatResultOutput(resultData, !execContext.useCleanOutput));
-
-      if (resultData.subtype === 'error_max_structured_output_retries') {
-        return {
-          type: 'throw',
-          error: new PentestError(
-            'Structured output validation failed after max retries',
-            'validation',
-            true,
-            {},
-            ErrorCode.OUTPUT_VALIDATION_FAILED,
-          ),
-        };
-      }
-
-      return {
-        type: 'complete' as const,
-        result: resultData.result,
-        cost: resultData.cost,
-        ...(resultData.structuredOutput !== undefined && { structuredOutput: resultData.structuredOutput }),
-      };
-    }
-
-    default:
-      logger.info(`Unhandled message type: ${message.type}`);
-      return { type: 'continue' };
-  }
-}

apps/worker/src/ai/models.ts

@@ -5,17 +5,28 @@
 // as published by the Free Software Foundation.
 
 /**
- * Model tier definitions and resolution.
+ * Model tier definitions and resolution for the pi harness.
  *
  * Three tiers mapped to capability levels:
  * - "small"  (Haiku — summarization, structured extraction)
  * - "medium" (Sonnet — tool use, general analysis)
  * - "large"  (Opus — deep reasoning, complex analysis)
  *
- * Users override via ANTHROPIC_SMALL_MODEL / ANTHROPIC_MEDIUM_MODEL / ANTHROPIC_LARGE_MODEL,
- * which works across all providers (direct, Bedrock, Vertex).
+ * Users override per tier via ANTHROPIC_SMALL_MODEL / ANTHROPIC_MEDIUM_MODEL /
+ * ANTHROPIC_LARGE_MODEL, which works across all providers (Anthropic, Bedrock,
+ * Vertex, custom base URL).
+ *
+ * Resolution returns a pi `Model` object via `ModelRegistry.find`, plus the
+ * `thinkingLevel` and an `AuthStorage` primed with runtime credentials. Bedrock
+ * and Vertex authenticate from the process environment (the AWS_ and GOOGLE_ vars
+ * the CLI forwards), so they need no runtime API key.
  */
 
+import type { ThinkingLevel } from '@earendil-works/pi-agent-core';
+import type { Api, Model } from '@earendil-works/pi-ai';
+import { AuthStorage, type ModelRegistry } from '@earendil-works/pi-coding-agent';
+import type { ProviderConfig } from '../types/config.js';
+
 export type ModelTier = 'small' | 'medium' | 'large';
 
 const DEFAULT_MODELS: Readonly<Record<ModelTier, string>> = {
@@ -24,8 +35,24 @@ const DEFAULT_MODELS: Readonly<Record<ModelTier, string>> = {
   large: 'claude-opus-4-8',
 };
 
-/** Resolve a model tier to a concrete model ID. */
-export function resolveModel(tier: ModelTier = 'medium'): string {
+/** pi-ai provider id for a Shannon ProviderConfig.providerType. Default: anthropic. */
+function piProviderId(providerConfig?: ProviderConfig): string {
+  switch (providerConfig?.providerType) {
+    case 'bedrock':
+      return 'amazon-bedrock';
+    case 'vertex':
+      return 'google-vertex';
+    default:
+      // 'anthropic_api', 'custom_base_url', or unset all resolve to the anthropic
+      // provider; custom_base_url overrides baseUrl/auth below.
+      return 'anthropic';
+  }
+}
+
+/** Resolve a model tier to a concrete model ID (env override → providerConfig → default). */
+export function resolveModelId(tier: ModelTier = 'medium', providerConfig?: ProviderConfig): string {
+  const override = providerConfig?.modelOverrides?.[tier];
+  if (override) return override;
   switch (tier) {
     case 'small':
       return process.env.ANTHROPIC_SMALL_MODEL || DEFAULT_MODELS.small;
@@ -36,9 +63,75 @@ export function resolveModel(tier: ModelTier = 'medium'): string {
   }
 }
 
-/** Whether a model supports adaptive thinking. Opus 4.6, 4.7, and 4.8 only. */
-export function supportsAdaptiveThinking(model: string): boolean {
-  return /opus-4-[678]/.test(model);
+/**
+ * Resolve the thinking level for a run.
+ *
+ * The Claude Agent SDK enabled "adaptive" thinking only on capable models; pi uses
+ * explicit levels and clamps to model capability internally. We default to 'medium'
+ * and honour the existing CLAUDE_ADAPTIVE_THINKING=false kill switch (→ 'off'). An
+ * explicit CLAUDE_THINKING_LEVEL wins when set.
+ */
+export function resolveThinkingLevel(): ThinkingLevel {
+  if (process.env.CLAUDE_ADAPTIVE_THINKING === 'false') return 'off';
+  const explicit = process.env.CLAUDE_THINKING_LEVEL as ThinkingLevel | undefined;
+  if (explicit) return explicit;
+  return 'medium';
+}
+
+export interface ModelSelection {
+  model: Model<Api>;
+  thinkingLevel: ThinkingLevel;
+  authStorage: AuthStorage;
+  modelId: string;
+  providerId: string;
+}
+
+/**
+ * Build an AuthStorage primed with the right credential for the active provider,
+ * then resolve the tier's model from a fresh ModelRegistry.
+ *
+ * - Anthropic: runtime API key from ContainerConfig.apiKey → ProviderConfig.apiKey
+ *   → ANTHROPIC_API_KEY env. OAuth (CLAUDE_CODE_OAUTH_TOKEN) is read from env by pi.
+ * - Custom base URL (custom_base_url): the auth token is set as the anthropic
+ *   runtime key and the model's baseUrl is overridden.
+ * - Bedrock / Vertex: authenticate from the process environment (the AWS_ and
+ *   GOOGLE_ vars), no runtime key needed.
+ */
+export function resolveModelSelection(
+  registryFactory: (authStorage: AuthStorage) => ModelRegistry,
+  modelTier: ModelTier,
+  apiKey?: string,
+  providerConfig?: ProviderConfig,
+): ModelSelection {
+  const providerId = piProviderId(providerConfig);
+  const modelId = resolveModelId(modelTier, providerConfig);
+
+  const authStorage = AuthStorage.inMemory();
+
+  const anthropicKey = apiKey ?? providerConfig?.apiKey ?? process.env.ANTHROPIC_API_KEY;
+  if (providerId === 'anthropic') {
+    const token =
+      providerConfig?.providerType === 'custom_base_url' ? (providerConfig.authToken ?? anthropicKey) : anthropicKey;
+    if (token) authStorage.setRuntimeApiKey('anthropic', token);
+  }
+
+  const registry = registryFactory(authStorage);
+  const found = registry.find(providerId, modelId);
+  if (!found) {
+    throw new Error(`Model not found in pi registry: provider="${providerId}" model="${modelId}"`);
+  }
+
+  // Custom base URL: override the resolved model's endpoint.
+  const baseUrl = providerConfig?.providerType === 'custom_base_url' ? providerConfig.baseUrl : undefined;
+  const model: Model<Api> = baseUrl ? { ...found, baseUrl } : found;
+
+  return {
+    model,
+    thinkingLevel: resolveThinkingLevel(),
+    authStorage,
+    modelId,
+    providerId,
+  };
 }
 
 /**

apps/worker/src/ai/queue-schemas.ts

@@ -5,196 +5,114 @@
 // as published by the Free Software Foundation.
 
 /**
- * Zod schema definitions for vulnerability exploitation queue structured outputs.
+ * TypeBox schemas + submit-tool factory for vulnerability exploitation queues.
  *
- * Each vuln agent returns a structured JSON response matching its schema.
- * The SDK validates the output against the JSON Schema generated from these Zod definitions.
+ * pi has no JSON-schema output format, so each vuln agent's structured queue is
+ * captured via a `submit_exploitation_queue` custom tool whose parameters mirror
+ * the per-class schema below. The captured payload is written to
+ * `<class>_exploitation_queue.json` by the caller (agent-execution).
  */
 
-import type { JsonSchemaOutputFormat } from '@anthropic-ai/claude-agent-sdk';
-import { z } from 'zod';
+import { defineTool, type ToolDefinition } from '@earendil-works/pi-coding-agent';
+import { type Static, type TObject, Type } from 'typebox';
 import type { AgentName } from '../types/agents.js';
 
-// === Common Fields ===
-
 const ANALYSIS_NOTES_DESCRIPTION = 'Plain context for defenders (caveats, scope, what is at risk). Not attack steps.';
 
-function notesField(exploit: boolean) {
-  const f = z.string().optional();
-  return exploit ? f : f.describe(ANALYSIS_NOTES_DESCRIPTION);
-}
+const optStr = (description?: string) => Type.Optional(Type.String(description ? { description } : {}));
 
-function makeBase(exploit: boolean) {
-  return z.object({
-    ID: z.string(),
-    vulnerability_type: z.string(),
-    externally_exploitable: z.boolean(),
-    confidence: z.string(),
-    notes: notesField(exploit),
-  });
-}
-
-// === Per-Vuln-Type Schemas (used for type inference; notes description is mode-agnostic for types) ===
-
-const baseVulnerability = makeBase(true);
-
-const InjectionVulnerability = baseVulnerability.extend({
-  source: z.string().optional(),
-  combined_sources: z.string().optional(),
-  path: z.string().optional(),
-  sink_call: z.string().optional(),
-  slot_type: z.string().optional(),
-  sanitization_observed: z.string().optional(),
-  concat_occurrences: z.string().optional(),
-  verdict: z.string().optional(),
-  mismatch_reason: z.string().optional(),
-  witness_payload: z.string().optional(),
-});
-
-const XssVulnerability = baseVulnerability.extend({
-  source: z.string().optional(),
-  source_detail: z.string().optional(),
-  path: z.string().optional(),
-  sink_function: z.string().optional(),
-  render_context: z.string().optional(),
-  encoding_observed: z.string().optional(),
-  verdict: z.string().optional(),
-  mismatch_reason: z.string().optional(),
-  witness_payload: z.string().optional(),
-});
-
-const AuthVulnerability = baseVulnerability.extend({
-  source_endpoint: z.string().optional(),
-  vulnerable_code_location: z.string().optional(),
-  missing_defense: z.string().optional(),
-  exploitation_hypothesis: z.string().optional(),
-  suggested_exploit_technique: z.string().optional(),
-});
-
-const SsrfVulnerability = baseVulnerability.extend({
-  source_endpoint: z.string().optional(),
-  vulnerable_parameter: z.string().optional(),
-  vulnerable_code_location: z.string().optional(),
-  missing_defense: z.string().optional(),
-  exploitation_hypothesis: z.string().optional(),
-  suggested_exploit_technique: z.string().optional(),
-});
-
-const AuthzVulnerability = baseVulnerability.extend({
-  endpoint: z.string().optional(),
-  vulnerable_code_location: z.string().optional(),
-  role_context: z.string().optional(),
-  guard_evidence: z.string().optional(),
-  side_effect: z.string().optional(),
-  reason: z.string().optional(),
-  minimal_witness: z.string().optional(),
-});
-
-// === Inferred Entry Types (consumed by renderer) ===
-
-export type InjectionFinding = z.infer<typeof InjectionVulnerability>;
-export type XssFinding = z.infer<typeof XssVulnerability>;
-export type AuthFinding = z.infer<typeof AuthVulnerability>;
-export type SsrfFinding = z.infer<typeof SsrfVulnerability>;
-export type AuthzFinding = z.infer<typeof AuthzVulnerability>;
-
-// === Convert to JSON Schema for SDK ===
-
-// NOTE: The SDK's AJV validator expects draft-07. Zod defaults to draft-2020-12 which
-// causes the SDK to silently skip structured output.
-function toOutputFormat(zodSchema: z.ZodType): JsonSchemaOutputFormat {
-  return { type: 'json_schema', schema: z.toJSONSchema(zodSchema, { target: 'draft-07' }) as Record<string, unknown> };
-}
-
-// === Per-Mode Output Format Builders ===
-// Two maps cached at module load; the only per-mode difference is the
-// description on the `notes` field, which steers the LLM's writing.
-
-function buildOutputFormats(exploit: boolean): Partial<Record<AgentName, JsonSchemaOutputFormat>> {
-  const base = makeBase(exploit);
+/** Base fields shared by every queue entry. `notes` gains guidance in analysis mode. */
+function baseFields(exploit: boolean) {
   return {
-    'injection-vuln': toOutputFormat(
-      z.object({
-        vulnerabilities: z.array(
-          base.extend({
-            source: z.string().optional(),
-            combined_sources: z.string().optional(),
-            path: z.string().optional(),
-            sink_call: z.string().optional(),
-            slot_type: z.string().optional(),
-            sanitization_observed: z.string().optional(),
-            concat_occurrences: z.string().optional(),
-            verdict: z.string().optional(),
-            mismatch_reason: z.string().optional(),
-            witness_payload: z.string().optional(),
-          }),
-        ),
-      }),
-    ),
-    'xss-vuln': toOutputFormat(
-      z.object({
-        vulnerabilities: z.array(
-          base.extend({
-            source: z.string().optional(),
-            source_detail: z.string().optional(),
-            path: z.string().optional(),
-            sink_function: z.string().optional(),
-            render_context: z.string().optional(),
-            encoding_observed: z.string().optional(),
-            verdict: z.string().optional(),
-            mismatch_reason: z.string().optional(),
-            witness_payload: z.string().optional(),
-          }),
-        ),
-      }),
-    ),
-    'auth-vuln': toOutputFormat(
-      z.object({
-        vulnerabilities: z.array(
-          base.extend({
-            source_endpoint: z.string().optional(),
-            vulnerable_code_location: z.string().optional(),
-            missing_defense: z.string().optional(),
-            exploitation_hypothesis: z.string().optional(),
-            suggested_exploit_technique: z.string().optional(),
-          }),
-        ),
-      }),
-    ),
-    'ssrf-vuln': toOutputFormat(
-      z.object({
-        vulnerabilities: z.array(
-          base.extend({
-            source_endpoint: z.string().optional(),
-            vulnerable_parameter: z.string().optional(),
-            vulnerable_code_location: z.string().optional(),
-            missing_defense: z.string().optional(),
-            exploitation_hypothesis: z.string().optional(),
-            suggested_exploit_technique: z.string().optional(),
-          }),
-        ),
-      }),
-    ),
-    'authz-vuln': toOutputFormat(
-      z.object({
-        vulnerabilities: z.array(
-          base.extend({
-            endpoint: z.string().optional(),
-            vulnerable_code_location: z.string().optional(),
-            role_context: z.string().optional(),
-            guard_evidence: z.string().optional(),
-            side_effect: z.string().optional(),
-            reason: z.string().optional(),
-            minimal_witness: z.string().optional(),
-          }),
-        ),
-      }),
-    ),
+    ID: Type.String(),
+    vulnerability_type: Type.String(),
+    externally_exploitable: Type.Boolean(),
+    confidence: Type.String(),
+    notes: exploit ? optStr() : optStr(ANALYSIS_NOTES_DESCRIPTION),
   };
 }
 
-const OUTPUT_FORMATS_EXPLOIT = buildOutputFormats(true);
-const OUTPUT_FORMATS_ANALYSIS = buildOutputFormats(false);
+const injectionFields = {
+  source: optStr(),
+  combined_sources: optStr(),
+  path: optStr(),
+  sink_call: optStr(),
+  slot_type: optStr(),
+  sanitization_observed: optStr(),
+  concat_occurrences: optStr(),
+  verdict: optStr(),
+  mismatch_reason: optStr(),
+  witness_payload: optStr(),
+};
+
+const xssFields = {
+  source: optStr(),
+  source_detail: optStr(),
+  path: optStr(),
+  sink_function: optStr(),
+  render_context: optStr(),
+  encoding_observed: optStr(),
+  verdict: optStr(),
+  mismatch_reason: optStr(),
+  witness_payload: optStr(),
+};
+
+const authFields = {
+  source_endpoint: optStr(),
+  vulnerable_code_location: optStr(),
+  missing_defense: optStr(),
+  exploitation_hypothesis: optStr(),
+  suggested_exploit_technique: optStr(),
+};
+
+const ssrfFields = {
+  source_endpoint: optStr(),
+  vulnerable_parameter: optStr(),
+  vulnerable_code_location: optStr(),
+  missing_defense: optStr(),
+  exploitation_hypothesis: optStr(),
+  suggested_exploit_technique: optStr(),
+};
+
+const authzFields = {
+  endpoint: optStr(),
+  vulnerable_code_location: optStr(),
+  role_context: optStr(),
+  guard_evidence: optStr(),
+  side_effect: optStr(),
+  reason: optStr(),
+  minimal_witness: optStr(),
+};
+
+const PER_TYPE_FIELDS: Partial<Record<AgentName, Record<string, ReturnType<typeof optStr>>>> = {
+  'injection-vuln': injectionFields,
+  'xss-vuln': xssFields,
+  'auth-vuln': authFields,
+  'ssrf-vuln': ssrfFields,
+  'authz-vuln': authzFields,
+};
+
+/** Build the `{ vulnerabilities: [...] }` queue schema for an agent + mode. */
+function queueSchema(agentName: AgentName, exploit: boolean): TObject | null {
+  const extra = PER_TYPE_FIELDS[agentName];
+  if (!extra) return null;
+  return Type.Object({
+    vulnerabilities: Type.Array(Type.Object({ ...baseFields(exploit), ...extra })),
+  });
+}
+
+// === Inferred entry types (consumed by renderers) ===
+export type InjectionFinding = Static<ReturnType<typeof injectionEntry>>;
+export type XssFinding = Static<ReturnType<typeof xssEntry>>;
+export type AuthFinding = Static<ReturnType<typeof authEntry>>;
+export type SsrfFinding = Static<ReturnType<typeof ssrfEntry>>;
+export type AuthzFinding = Static<ReturnType<typeof authzEntry>>;
+
+const injectionEntry = () => Type.Object({ ...baseFields(true), ...injectionFields });
+const xssEntry = () => Type.Object({ ...baseFields(true), ...xssFields });
+const authEntry = () => Type.Object({ ...baseFields(true), ...authFields });
+const ssrfEntry = () => Type.Object({ ...baseFields(true), ...ssrfFields });
+const authzEntry = () => Type.Object({ ...baseFields(true), ...authzFields });
 
 const VULN_AGENT_QUEUE_FILENAMES: Partial<Record<AgentName, string>> = {
   'injection-vuln': 'injection_exploitation_queue.json',
@@ -204,12 +122,38 @@ const VULN_AGENT_QUEUE_FILENAMES: Partial<Record<AgentName, string>> = {
   'authz-vuln': 'authz_exploitation_queue.json',
 };
 
-/** Returns the structured output format for a vuln agent, or undefined for non-vuln agents. */
-export function getOutputFormat(agentName: AgentName, exploit = true): JsonSchemaOutputFormat | undefined {
-  return (exploit ? OUTPUT_FORMATS_EXPLOIT : OUTPUT_FORMATS_ANALYSIS)[agentName];
-}
-
 /** Returns the queue filename for a vuln agent, or undefined for non-vuln agents. */
 export function getQueueFilename(agentName: AgentName): string | undefined {
   return VULN_AGENT_QUEUE_FILENAMES[agentName];
 }
+
+export interface QueueSubmitTool {
+  tool: ToolDefinition;
+  getCaptured: () => unknown;
+}
+
+/**
+ * Build the `submit_exploitation_queue` tool for a vuln agent, or null for
+ * non-vuln agents. The agent calls it once with the full findings list; the
+ * captured payload is the structured queue.
+ */
+export function createQueueSubmitTool(agentName: AgentName, exploit: boolean): QueueSubmitTool | null {
+  const schema = queueSchema(agentName, exploit);
+  if (!schema) return null;
+  let captured: unknown;
+  const tool = defineTool({
+    name: 'submit_exploitation_queue',
+    label: 'Submit Exploitation Queue',
+    description:
+      'Submit the final structured list of analyzed vulnerabilities for this class. Call exactly once when ' +
+      'analysis is complete, with every finding included.',
+    promptSnippet: 'submit_exploitation_queue: record the final structured findings list (call once)',
+    parameters: schema,
+    execute: async (_toolCallId, params) => {
+      captured = params;
+      const count = (params as { vulnerabilities?: unknown[] }).vulnerabilities?.length ?? 0;
+      return { content: [{ type: 'text' as const, text: `Recorded ${count} findings.` }], details: {} };
+    },
+  });
+  return { tool, getCaptured: () => captured };
+}

apps/worker/src/ai/settings-writer.ts

@@ -5,37 +5,71 @@
 // as published by the Free Software Foundation.
 
 /**
- * Writes ~/.claude/settings.json with permissions.deny rules derived from
- * `code_path` avoid patterns. The SDK reads this via `settingSources: ['user']`;
- * deny rules fire even in `bypassPermissions` mode.
+ * Writes the @gotgenes/pi-permission-system global config from `code_path` avoid
+ * patterns. The executor loads the extension (see claude-executor) and pi enforces
+ * these path denies at the tool layer for every agent. Written to the global config
+ * dir under `agentDir` — the project-scoped path is gated behind project trust,
+ * which our headless runs do not grant; the global path is not.
  */
 
-import os from 'node:os';
+import { getAgentDir } from '@earendil-works/pi-coding-agent';
 import { fs, path } from 'zx';
 import type { DistributedConfig } from '../types/config.js';
 
-const FILE_TOOLS = ['Read', 'Edit'] as const;
-
-function denyEntriesFor(pattern: string): string[] {
-  const arg = `./${pattern.replace(/^[./]+/, '')}`;
-  return FILE_TOOLS.map((tool) => `${tool}(${arg})`);
+/** Absolute path to the pi-permission-system global config.json. */
+export function permissionConfigPath(): string {
+  return path.join(getAgentDir(), 'extensions', 'pi-permission-system', 'config.json');
 }
 
-export async function writeUserSettingsForCodePathAvoids(config: DistributedConfig | null): Promise<void> {
+/**
+ * Write (or remove) the pi-permission-system config derived from `code_path`
+ * avoid patterns.
+ *
+ * Each avoid maps to a cross-cutting `path` deny — the strongest surface, blocking
+ * the path across every tool and bash command, and not overridable by a per-tool
+ * allow. `"*": "allow"` keeps everything else permitted so the extension does not
+ * fall back to its default `ask` (which would block all access headlessly). When
+ * there are no avoids the config is removed, so the executor skips loading the
+ * extension entirely.
+ */
+export async function writeCodePathPermissionConfig(config: DistributedConfig | null): Promise<void> {
   const avoidPatterns = (config?.avoid ?? []).filter((r) => r.type === 'code_path').map((r) => r.value);
-  const settingsPath = path.join(os.homedir(), '.claude', 'settings.json');
+  const configPath = permissionConfigPath();
 
   if (avoidPatterns.length === 0) {
-    await fs.remove(settingsPath);
+    await fs.remove(configPath);
     return;
   }
 
-  const settings = {
-    permissions: {
-      deny: avoidPatterns.flatMap(denyEntriesFor),
+  // pi's matcher (wildcard-matcher.ts) has NO `**` globstar — it splits on each `*`
+  // and joins with `.*`, and a single `*` already matches any chars incl. `/`. Tool
+  // paths are compared as absolute (path-utils resolves them against cwd), so we
+  // collapse `**`→`*` and add a `*/`-prefixed variant that matches the path under
+  // any repo prefix. (A bare pattern never matches an absolute path.)
+  const pathDeny: Record<string, 'allow' | 'deny'> = { '*': 'allow' };
+  for (const pattern of avoidPatterns) {
+    const clean = pattern.replace(/^[./]+/, '').replace(/\*\*/g, '*');
+    // Deny the contents (under any repo prefix and as written)...
+    pathDeny[`*/${clean}`] = 'deny';
+    pathDeny[clean] = 'deny';
+    // ...and the folder path itself, so the directory entry is denied too — the
+    // contents patterns (…/*) require a trailing segment and wouldn't match it.
+    if (clean.endsWith('/*')) {
+      const folder = clean.slice(0, -2);
+      if (folder) {
+        pathDeny[`*/${folder}`] = 'deny';
+        pathDeny[folder] = 'deny';
+      }
+    }
+  }
+
+  const permissionConfig = {
+    permission: {
+      '*': 'allow',
+      path: pathDeny,
     },
   };
 
-  await fs.ensureDir(path.dirname(settingsPath));
-  await fs.writeJson(settingsPath, settings, { spaces: 2 });
+  await fs.ensureDir(path.dirname(configPath));
+  await fs.writeJson(configPath, permissionConfig, { spaces: 2 });
 }

apps/worker/src/ai/tools.ts

@@ -0,0 +1,137 @@
+// Copyright (C) 2025 Keygraph, Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License version 3
+// as published by the Free Software Foundation.
+
+/**
+ * Universal custom tools registered for every agent: `task` and `todo_write`.
+ *
+ * These replace the Claude Agent SDK built-ins that pi does not ship. `task`
+ * delegates a focused analysis to an in-process read-only child session (the
+ * Task sub-agent replacement); `todo_write` is a full-state-replace planning
+ * scratchpad mirrored to the workflow log.
+ */
+
+import type { ThinkingLevel } from '@earendil-works/pi-agent-core';
+import type { Api, Model } from '@earendil-works/pi-ai';
+import {
+  type AuthStorage,
+  createAgentSession,
+  defineTool,
+  type ResourceLoader,
+  SessionManager,
+  SettingsManager,
+  type ToolDefinition,
+} from '@earendil-works/pi-coding-agent';
+import { Type } from 'typebox';
+import type { AuditLogger } from './audit-logger.js';
+
+/** Read-only tool surface for delegated child sessions. No bash/edit/write. */
+const CHILD_TOOLS = ['read', 'grep', 'find', 'ls'];
+
+export interface TaskToolContext {
+  model: Model<Api>;
+  thinkingLevel: ThinkingLevel;
+  authStorage: AuthStorage;
+  cwd: string;
+  /** When set, child sessions inherit the code_path deny policy. */
+  resourceLoader?: ResourceLoader;
+}
+
+/**
+ * The `task` tool — delegate a focused, read-only analysis to a sub-agent.
+ *
+ * Spawns an in-process child `createAgentSession` scoped to read-only tools,
+ * drives it to completion, and returns its final text. Marked `parallel` so the
+ * model can fan out multiple delegations in one turn (the pre-recon/recon prompts
+ * rely on this). Children get no `task` tool of their own — delegation is one level.
+ */
+export function createTaskTool(ctx: TaskToolContext): ToolDefinition {
+  return defineTool({
+    name: 'task',
+    label: 'Task',
+    description:
+      'Delegate a focused source-code analysis question to a read-only sub-agent and return its findings. ' +
+      'Use for deep code reading, tracing, and attack-surface mapping. Issue multiple task calls in one ' +
+      'message to run analyses in parallel.',
+    promptSnippet: 'task: delegate read-only code analysis to a parallel sub-agent',
+    executionMode: 'parallel',
+    parameters: Type.Object({
+      description: Type.Optional(Type.String({ description: 'Short (3-5 word) label for the delegated analysis.' })),
+      prompt: Type.String({ description: 'The full analysis instruction for the sub-agent.' }),
+    }),
+    execute: async (_toolCallId, params) => {
+      const { session: child } = await createAgentSession({
+        cwd: ctx.cwd,
+        model: ctx.model,
+        thinkingLevel: ctx.thinkingLevel,
+        tools: CHILD_TOOLS,
+        authStorage: ctx.authStorage,
+        sessionManager: SessionManager.inMemory(),
+        settingsManager: SettingsManager.inMemory({
+          retry: { enabled: false },
+          compaction: { enabled: true },
+        }),
+        ...(ctx.resourceLoader && { resourceLoader: ctx.resourceLoader }),
+      });
+      try {
+        await child.prompt(params.prompt);
+        const text = child.getLastAssistantText() ?? '(sub-agent produced no output)';
+        return { content: [{ type: 'text' as const, text }], details: {} };
+      } finally {
+        child.dispose();
+      }
+    },
+  });
+}
+
+export interface TodoItem {
+  content: string;
+  status: 'pending' | 'in_progress' | 'completed';
+  activeForm: string;
+}
+
+/** Render a todo list as a compact checklist for the workflow log. */
+function renderTodos(todos: readonly TodoItem[]): string {
+  const mark = (s: TodoItem['status']): string => (s === 'completed' ? 'x' : s === 'in_progress' ? '~' : ' ');
+  return todos.map((t) => `[${mark(t.status)}] ${t.content}`).join('  ');
+}
+
+/**
+ * The `todo_write` tool — a full-state-replace planning scratchpad.
+ *
+ * Mirrors Claude Code's TodoWrite: each call carries the entire list and replaces
+ * stored state (no append/merge). No deliverable impact; every call is echoed to
+ * the workflow log so `shannon logs` shows the agent's live plan. State is per
+ * tool instance (one per agent execution).
+ */
+export function createTodoWriteTool(auditLogger: AuditLogger): ToolDefinition {
+  let current: TodoItem[] = [];
+  return defineTool({
+    name: 'todo_write',
+    label: 'Todo Write',
+    description:
+      'Record or update the task plan. Pass the COMPLETE todo list every call (full replace, not append). ' +
+      'Keep exactly one task in_progress; mark tasks completed as soon as they are done.',
+    promptSnippet: 'todo_write: track a plan as a checklist (pass the full list each call)',
+    parameters: Type.Object({
+      todos: Type.Array(
+        Type.Object({
+          content: Type.String({ description: 'Imperative task description, e.g. "Map SSRF sinks".' }),
+          status: Type.Union([Type.Literal('pending'), Type.Literal('in_progress'), Type.Literal('completed')]),
+          activeForm: Type.String({ description: 'Present-continuous form, e.g. "Mapping SSRF sinks".' }),
+        }),
+      ),
+    }),
+    execute: async (_toolCallId, params) => {
+      current = params.todos as TodoItem[];
+      const completed = current.filter((t) => t.status === 'completed').length;
+      await auditLogger.logNote('todo', renderTodos(current));
+      return {
+        content: [{ type: 'text' as const, text: `Todos updated (${current.length} items, ${completed} completed).` }],
+        details: {},
+      };
+    },
+  });
+}

apps/worker/src/ai/types.ts

@@ -4,9 +4,7 @@
 // it under the terms of the GNU Affero General Public License version 3
 // as published by the Free Software Foundation.
 
-// Type definitions for Claude executor message processing pipeline
-
-import type { SDKAssistantMessageError } from '@anthropic-ai/claude-agent-sdk';
+// Shared display/formatting types for the agent executor output layer.
 
 export interface ExecutionContext {
   isParallelExecution: boolean;
@@ -15,18 +13,6 @@ export interface ExecutionContext {
   agentKey: string;
 }
 
-export interface AssistantResult {
-  content: string;
-  cleanedContent: string;
-  apiErrorDetected: boolean;
-  shouldThrow?: Error;
-  logData: {
-    turn: number;
-    content: string;
-    timestamp: string;
-  };
-}
-
 export interface ResultData {
   result: string | null;
   cost: number;
@@ -36,77 +22,3 @@ export interface ResultData {
   permissionDenials: number;
   structuredOutput?: unknown;
 }
-
-export interface ToolUseData {
-  toolName: string;
-  parameters: Record<string, unknown>;
-  timestamp: string;
-}
-
-export interface ToolResultData {
-  content: unknown;
-  displayContent: string;
-  timestamp: string;
-}
-
-export interface ContentBlock {
-  type?: string;
-  text?: string;
-  thinking?: string;
-  data?: string;
-}
-
-export interface AssistantMessage {
-  type: 'assistant';
-  error?: SDKAssistantMessageError;
-  message: {
-    content: ContentBlock[] | string;
-  };
-}
-
-export interface ResultMessage {
-  type: 'result';
-  result?: string;
-  total_cost_usd?: number;
-  duration_ms?: number;
-  subtype?: string;
-  stop_reason?: string | null;
-  permission_denials?: unknown[];
-  structured_output?: unknown;
-}
-
-export interface ToolUseMessage {
-  type: 'tool_use';
-  name: string;
-  input?: Record<string, unknown>;
-}
-
-export interface ToolResultMessage {
-  type: 'tool_result';
-  content?: unknown;
-}
-
-export interface ApiErrorDetection {
-  detected: boolean;
-  shouldThrow?: Error;
-}
-
-export interface SystemInitMessage {
-  type: 'system';
-  subtype: 'init';
-  model?: string;
-  permissionMode?: string;
-}
-
-/** Emitted when a model refuses a request and the SDK falls back to another model (e.g. Fable 5 routing cybersecurity tasks to Opus 4.8). */
-export interface ModelRefusalFallbackMessage {
-  type: 'system';
-  subtype: 'model_refusal_fallback';
-  original_model: string;
-  fallback_model: string;
-  api_refusal_category?: string | null;
-}
-
-export interface UserMessage {
-  type: 'user';
-}