From c7be324083f9edbf016a525f498278ec670cfaa1 Mon Sep 17 00:00:00 2001 From: Chris Farhood Date: Thu, 23 Apr 2026 13:32:23 -0400 Subject: [PATCH] backport: provider extensions and drop claude-code-router mode Cherry-pick of KeygraphHQ/shannon#295 (581c208). Upstream changes: removes router mode from CLI/worker, adds provider extensions, new report-output-provider and checkpoint-provider interfaces, refactored workflow orchestration. Conflicts resolved: kept our README.md, CLAUDE.md, and deleted compose files. Co-Authored-By: Claude Opus 4.6 (1M context) --- .env.example | 26 +- .github/ISSUE_TEMPLATE/bug_report.yml | 2 - apps/cli/infra/router-config.json | 31 -- apps/cli/src/commands/setup.ts | 49 +- apps/cli/src/commands/start.ts | 68 ++- apps/cli/src/config/resolver.ts | 21 +- apps/cli/src/config/writer.ts | 1 - apps/cli/src/docker.ts | 471 ++++++++---------- apps/cli/src/env.ts | 14 +- apps/cli/src/index.ts | 7 - apps/worker/src/ai/claude-executor.ts | 5 +- apps/worker/src/ai/message-handlers.ts | 7 +- apps/worker/src/ai/router-utils.ts | 27 - .../src/interfaces/checkpoint-provider.ts | 50 +- .../src/interfaces/findings-provider.ts | 2 +- apps/worker/src/interfaces/index.ts | 4 +- .../src/interfaces/report-output-provider.ts | 22 + apps/worker/src/services/container.ts | 33 +- apps/worker/src/services/index.ts | 4 +- apps/worker/src/services/preflight.ts | 4 +- apps/worker/src/services/reporting.ts | 13 +- apps/worker/src/session-manager.ts | 1 - apps/worker/src/temporal/activities.ts | 90 +++- apps/worker/src/temporal/shared.ts | 4 +- apps/worker/src/temporal/workflows.ts | 39 +- apps/worker/src/types/config.ts | 1 - apps/worker/src/utils/billing-detection.ts | 1 - 27 files changed, 458 insertions(+), 539 deletions(-) delete mode 100644 apps/cli/infra/router-config.json delete mode 100644 apps/worker/src/ai/router-utils.ts create mode 100644 apps/worker/src/interfaces/report-output-provider.ts diff --git a/.env.example b/.env.example index e843035..a7262e9 100644 --- a/.env.example +++ b/.env.example @@ -5,7 +5,7 @@ CLAUDE_CODE_MAX_OUTPUT_TOKENS=64000 # ============================================================================= -# OPTION 1: Direct Anthropic (default, no router) +# OPTION 1: Direct Anthropic # ============================================================================= ANTHROPIC_API_KEY=your-api-key-here @@ -19,20 +19,6 @@ ANTHROPIC_API_KEY=your-api-key-here # ANTHROPIC_BASE_URL=https://your-proxy.example.com # ANTHROPIC_AUTH_TOKEN=your-auth-token # Auth token for the custom endpoint -# ============================================================================= -# OPTION 3: Router Mode (use alternative providers) -# ============================================================================= -# Enable router mode by running: ./shannon start ... ROUTER=true -# Then configure ONE of the providers below: - -# --- OpenAI --- -# OPENAI_API_KEY=sk-your-openai-key -# ROUTER_DEFAULT=openai,gpt-5.2 - -# --- OpenRouter (access Gemini 3 models via single API) --- -# OPENROUTER_API_KEY=sk-or-your-openrouter-key -# ROUTER_DEFAULT=openrouter,google/gemini-3-flash-preview - # ============================================================================= # Model Tier Overrides (Anthropic API / OAuth / Custom Base URL / Bedrock) # ============================================================================= @@ -43,7 +29,7 @@ ANTHROPIC_API_KEY=your-api-key-here # ANTHROPIC_LARGE_MODEL=... # Large tier (default: claude-opus-4-6) # ============================================================================= -# OPTION 4: AWS Bedrock +# OPTION 3: AWS Bedrock # ============================================================================= # https://aws.amazon.com/blogs/machine-learning/accelerate-ai-development-with-amazon-bedrock-api-keys/ # Requires the model tier overrides above to be set with Bedrock-specific model IDs. @@ -57,7 +43,7 @@ ANTHROPIC_API_KEY=your-api-key-here # AWS_BEARER_TOKEN_BEDROCK=your-bearer-token # ============================================================================= -# OPTION 5: Google Vertex AI +# OPTION 4: Google Vertex AI # ============================================================================= # https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-partner-models # Requires a GCP service account with roles/aiplatform.user. @@ -72,9 +58,3 @@ ANTHROPIC_API_KEY=your-api-key-here # CLOUD_ML_REGION=us-east5 # ANTHROPIC_VERTEX_PROJECT_ID=your-gcp-project-id # GOOGLE_APPLICATION_CREDENTIALS=./credentials/google-sa-key.json - -# ============================================================================= -# Available Models -# ============================================================================= -# OpenAI: gpt-5.2, gpt-5-mini -# OpenRouter: google/gemini-3-flash-preview diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index acac933..fe7437e 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -120,8 +120,6 @@ body: - "Custom base URL (proxy/gateway)" - "AWS Bedrock" - "Google Vertex AI" - - "Router - OpenAI" - - "Router - OpenRouter" validations: required: true diff --git a/apps/cli/infra/router-config.json b/apps/cli/infra/router-config.json deleted file mode 100644 index 9cc0422..0000000 --- a/apps/cli/infra/router-config.json +++ /dev/null @@ -1,31 +0,0 @@ -{ - "HOST": "0.0.0.0", - "APIKEY": "shannon-router-key", - "LOG": true, - "LOG_LEVEL": "info", - "NON_INTERACTIVE_MODE": true, - "API_TIMEOUT_MS": 600000, - "Providers": [ - { - "name": "openai", - "api_base_url": "https://api.openai.com/v1/chat/completions", - "api_key": "$OPENAI_API_KEY", - "models": ["gpt-5.2", "gpt-5-mini"], - "transformer": { - "use": [["maxcompletiontokens", { "max_completion_tokens": 16384 }]] - } - }, - { - "name": "openrouter", - "api_base_url": "https://openrouter.ai/api/v1/chat/completions", - "api_key": "$OPENROUTER_API_KEY", - "models": ["google/gemini-3-flash-preview"], - "transformer": { - "use": ["openrouter"] - } - } - ], - "Router": { - "default": "$ROUTER_DEFAULT" - } -} diff --git a/apps/cli/src/commands/setup.ts b/apps/cli/src/commands/setup.ts index 34cdbe0..42e7e77 100644 --- a/apps/cli/src/commands/setup.ts +++ b/apps/cli/src/commands/setup.ts @@ -13,7 +13,7 @@ import { type ShannonConfig, saveConfig } from '../config/writer.js'; const SHANNON_HOME = path.join(os.homedir(), '.shannon'); -type Provider = 'anthropic' | 'custom_base_url' | 'bedrock' | 'vertex' | 'router'; +type Provider = 'anthropic' | 'custom_base_url' | 'bedrock' | 'vertex'; export async function setup(): Promise { p.intro('Shannon Setup'); @@ -26,7 +26,6 @@ export async function setup(): Promise { { value: 'custom_base_url' as const, label: 'Custom Base URL', hint: 'proxies, gateways' }, { value: 'bedrock' as const, label: 'Claude via AWS Bedrock' }, { value: 'vertex' as const, label: 'Claude via Google Vertex AI' }, - { value: 'router' as const, label: 'Router', hint: 'experimental' }, ], }); if (p.isCancel(provider)) return cancelAndExit(); @@ -51,8 +50,6 @@ async function setupProvider(provider: Provider): Promise { return setupBedrock(); case 'vertex': return setupVertex(); - case 'router': - return setupRouter(); } } @@ -282,50 +279,6 @@ async function setupVertex(): Promise { }; } -async function setupRouter(): Promise { - const routerProvider = await p.select({ - message: 'Router provider', - options: [ - { value: 'openai' as const, label: 'OpenAI' }, - { value: 'openrouter' as const, label: 'OpenRouter' }, - ], - }); - if (p.isCancel(routerProvider)) return cancelAndExit(); - - const apiKey = await promptSecret( - routerProvider === 'openai' ? 'Enter your OpenAI API key' : 'Enter your OpenRouter API key', - ); - - let defaultModel: string; - if (routerProvider === 'openai') { - const model = await p.select({ - message: 'Default model', - options: [ - { value: 'gpt-5.2' as const, label: 'GPT-5.2' }, - { value: 'gpt-5-mini' as const, label: 'GPT-5 Mini' }, - ], - }); - if (p.isCancel(model)) return cancelAndExit(); - defaultModel = `openai,${model}`; - } else { - const model = await p.select({ - message: 'Default model', - options: [{ value: 'google/gemini-3-flash-preview' as const, label: 'Google Gemini 3 Flash Preview' }], - }); - if (p.isCancel(model)) return cancelAndExit(); - defaultModel = `openrouter,${model}`; - } - - const router: ShannonConfig['router'] = { default: defaultModel }; - if (routerProvider === 'openai') { - router.openai_key = apiKey; - } else { - router.openrouter_key = apiKey; - } - - return { router }; -} - // === Helpers === async function promptSecret(message: string): Promise { diff --git a/apps/cli/src/commands/start.ts b/apps/cli/src/commands/start.ts index 5f9404d..d8f989e 100644 --- a/apps/cli/src/commands/start.ts +++ b/apps/cli/src/commands/start.ts @@ -5,12 +5,11 @@ * and npx mode (Docker Hub pull, ~/.shannon/). */ +import { execFileSync } from 'node:child_process'; import fs from 'node:fs'; -import os from 'node:os'; import path from 'node:path'; -import { getOrchestrator } from '../backend.js'; -import { randomSuffix } from '../docker.js'; -import { buildEnvFlags, isRouterConfigured, loadEnv, validateCredentials } from '../env.js'; +import { ensureImage, ensureInfra, randomSuffix, spawnWorker } from '../docker.js'; +import { buildEnvFlags, loadEnv, validateCredentials } from '../env.js'; import { getCredentialsPath, getWorkspacesDir, initHome } from '../home.js'; import { isLocal } from '../mode.js'; import { resolveConfig, resolveRepo } from '../paths.js'; @@ -23,7 +22,6 @@ export interface StartArgs { workspace?: string; output?: string; pipelineTesting: boolean; - router: boolean; version: string; } @@ -32,13 +30,12 @@ export async function start(args: StartArgs): Promise { initHome(); loadEnv(); - // 2. Validate credentials and auto-detect router mode + // 2. Validate credentials const creds = validateCredentials(); if (!creds.valid) { console.error(`ERROR: ${creds.error}`); process.exit(1); } - const useRouter = args.router || isRouterConfigured(); // 3. Resolve paths const repo = resolveRepo(args.repo); @@ -49,27 +46,20 @@ export async function start(args: StartArgs): Promise { fs.mkdirSync(workspacesDir, { recursive: true }); fs.chmodSync(workspacesDir, 0o777); - // 5. Handle router env - if (useRouter) { - process.env.ANTHROPIC_BASE_URL = 'http://shannon-router:3456'; - process.env.ANTHROPIC_AUTH_TOKEN = 'shannon-router-key'; - } + // 5. Ensure image (auto-build in dev, pull in npx) and start infra + ensureImage(args.version); + await ensureInfra(); - // 6. Ensure image and start infra via orchestrator - const orchestrator = await getOrchestrator(); - orchestrator.ensureImage(args.version); - await orchestrator.ensureInfra(useRouter); - - // 7. Generate unique task queue and container name + // 6. Generate unique task queue and container name const suffix = randomSuffix(); const taskQueue = `shannon-${suffix}`; const containerName = `shannon-worker-${suffix}`; - // 8. Generate workspace name if not provided + // 7. Generate workspace name if not provided const workspace = args.workspace ?? `${new URL(args.url).hostname.replace(/[^a-zA-Z0-9-]/g, '-')}_shannon-${Date.now()}`; - // 9. Create writable overlay directories (mounted over :ro repo paths inside container) + // 8. Create writable overlay directories (mounted over :ro repo paths inside container) // Workspace dir must be 0o777 so the container user (UID 1001) can create audit subdirs const workspacePath = path.join(workspacesDir, workspace); fs.mkdirSync(workspacePath, { recursive: true }); @@ -80,12 +70,10 @@ export async function start(args: StartArgs): Promise { fs.chmodSync(dirPath, 0o777); } - // 10. Pre-create overlay mount points (Linux :ro mounts can't auto-create them) - if (os.platform() === 'linux') { - const shannonDir = path.join(repo.hostPath, '.shannon'); - for (const dir of ['deliverables', 'scratchpad', '.playwright-cli']) { - fs.mkdirSync(path.join(shannonDir, dir), { recursive: true }); - } + // 9. Pre-create overlay mount points (:ro mounts can't auto-create them) + const shannonDir = path.join(repo.hostPath, '.shannon'); + for (const dir of ['deliverables', 'scratchpad', '.playwright-cli']) { + fs.mkdirSync(path.join(shannonDir, dir), { recursive: true }); } const credentialsPath = getCredentialsPath(); @@ -95,20 +83,20 @@ export async function start(args: StartArgs): Promise { process.env.GOOGLE_APPLICATION_CREDENTIALS = '/app/credentials/google-sa-key.json'; } - // 11. Resolve output directory + // 10. Resolve output directory const outputDir = args.output ? path.resolve(args.output) : undefined; if (outputDir) { fs.mkdirSync(outputDir, { recursive: true }); } - // 12. Resolve prompts directory (local mode only) + // 11. Resolve prompts directory (local mode only) const promptsDir = isLocal() ? path.resolve('apps/worker/prompts') : undefined; - // 13. Display splash screen + // 12. Display splash screen displaySplash(isLocal() ? undefined : args.version); - // 14. Spawn worker via orchestrator - const handle = orchestrator.spawnWorker({ + // 13. Spawn worker container + const proc = spawnWorker({ version: args.version, url: args.url, repo, @@ -124,8 +112,8 @@ export async function start(args: StartArgs): Promise { ...(args.pipelineTesting && { pipelineTesting: true }), }); - // 15. Wait for workflow to register, then display info - handle.onError((err) => { + // 14. Wait for workflow to register, then display info + proc.on('error', (err) => { console.error(`Failed to start worker: ${err.message}`); process.exit(1); }); @@ -173,7 +161,7 @@ export async function start(args: StartArgs): Promise { // Clear waiting line and show info process.stdout.write('\r\x1b[K'); - printInfo(args, useRouter, workspace, workflowId, repo.hostPath, workspacesDir); + printInfo(args, workspace, workflowId, repo.hostPath, workspacesDir); return; } } catch { @@ -182,14 +170,18 @@ export async function start(args: StartArgs): Promise { process.stdout.write('.'); }, 2000); - // Stop the worker only if it hasn't started yet + // Stop the worker container only if it hasn't started yet let cleaned = false; const cleanup = (): void => { if (cleaned || started) return; cleaned = true; clearInterval(pollInterval); console.log(`\nStopping worker ${containerName}...`); - handle.kill(); + try { + execFileSync('docker', ['stop', containerName], { stdio: 'pipe' }); + } catch { + // Container may have already exited + } }; process.on('SIGINT', () => { @@ -205,7 +197,6 @@ export async function start(args: StartArgs): Promise { function printInfo( args: StartArgs, - routerActive: boolean, workspace: string, workflowId: string, repoPath: string, @@ -223,9 +214,6 @@ function printInfo( if (args.pipelineTesting) { console.log(' Mode: Pipeline Testing'); } - if (routerActive) { - console.log(' Router: Enabled'); - } console.log(''); console.log(' Monitor:'); if (workflowId) { diff --git a/apps/cli/src/config/resolver.ts b/apps/cli/src/config/resolver.ts index 039c362..98f0fad 100644 --- a/apps/cli/src/config/resolver.ts +++ b/apps/cli/src/config/resolver.ts @@ -44,11 +44,6 @@ const CONFIG_MAP: readonly ConfigMapping[] = [ { env: 'ANTHROPIC_BASE_URL', toml: 'custom_base_url.base_url', type: 'string' }, { env: 'ANTHROPIC_AUTH_TOKEN', toml: 'custom_base_url.auth_token', type: 'string' }, - // Router - { env: 'ROUTER_DEFAULT', toml: 'router.default', type: 'string' }, - { env: 'OPENAI_API_KEY', toml: 'router.openai_key', type: 'string' }, - { env: 'OPENROUTER_API_KEY', toml: 'router.openrouter_key', type: 'string' }, - // Model tiers { env: 'ANTHROPIC_SMALL_MODEL', toml: 'models.small', type: 'string' }, { env: 'ANTHROPIC_MEDIUM_MODEL', toml: 'models.medium', type: 'string' }, @@ -165,20 +160,6 @@ function validateProviderFields(config: TOMLConfig, provider: string, errors: st validateModelTiers(config, 'vertex', errors); break; } - - case 'router': { - if (!keys.includes('default')) { - errors.push('[router] missing required key: default'); - } - if (!keys.includes('openai_key') && !keys.includes('openrouter_key')) { - errors.push('[router] requires either openai_key or openrouter_key'); - } - const models = config.models as Record | undefined; - if (models && typeof models === 'object' && Object.keys(models).length > 0) { - errors.push('[models] is not supported with [router]'); - } - break; - } } } @@ -242,7 +223,7 @@ function validateConfig(config: TOMLConfig): string[] { } // 4. Only one provider section allowed (ignore empty sections) - const PROVIDER_SECTIONS = ['anthropic', 'custom_base_url', 'bedrock', 'vertex', 'router'] as const; + const PROVIDER_SECTIONS = ['anthropic', 'custom_base_url', 'bedrock', 'vertex'] as const; const present = PROVIDER_SECTIONS.filter((s) => { const section = config[s]; return section && typeof section === 'object' && Object.keys(section).length > 0; diff --git a/apps/cli/src/config/writer.ts b/apps/cli/src/config/writer.ts index 58ee7e9..82aa63b 100644 --- a/apps/cli/src/config/writer.ts +++ b/apps/cli/src/config/writer.ts @@ -13,7 +13,6 @@ export interface ShannonConfig { custom_base_url?: { base_url?: string; auth_token?: string }; bedrock?: { use?: boolean; region?: string; token?: string }; vertex?: { use?: boolean; region?: string; project_id?: string; key_path?: string }; - router?: { default?: string; openai_key?: string; openrouter_key?: string }; models?: { small?: string; medium?: string; large?: string }; } diff --git a/apps/cli/src/docker.ts b/apps/cli/src/docker.ts index 18eb270..84c3cc4 100644 --- a/apps/cli/src/docker.ts +++ b/apps/cli/src/docker.ts @@ -12,22 +12,27 @@ import path from 'node:path'; import { setTimeout as sleep } from 'node:timers/promises'; import { fileURLToPath } from 'node:url'; import { getMode } from './mode.js'; -import type { Orchestrator, WorkerHandle, WorkerOptions } from './orchestrator.js'; - -export type { WorkerOptions }; const __dirname = path.dirname(fileURLToPath(import.meta.url)); const NPX_IMAGE_REPO = 'keygraph/shannon'; const DEV_IMAGE = 'shannon-worker'; +export function getWorkerImage(version: string): string { + return getMode() === 'local' ? DEV_IMAGE : `${NPX_IMAGE_REPO}:${version}`; +} + +function getComposeFile(): string { + return getMode() === 'local' + ? path.resolve('docker-compose.yml') + : path.resolve(__dirname, '..', 'infra', 'compose.yml'); +} + /** Generate an 8-char random hex suffix for container/queue names. */ export function randomSuffix(): string { return crypto.randomBytes(4).toString('hex'); } -// === Internal Helpers === - /** Run a command silently, return true if it succeeds. */ function runQuiet(cmd: string, args: string[]): boolean { try { @@ -47,16 +52,83 @@ function runOutput(cmd: string, args: string[]): string { } } -function getComposeFile(): string { - return getMode() === 'local' - ? path.resolve('docker-compose.yml') - : path.resolve(__dirname, '..', 'infra', 'compose.yml'); +/** + * Check if Temporal is running and healthy. + */ +export function isTemporalReady(): boolean { + const output = runOutput('docker', [ + 'exec', + 'shannon-temporal', + 'temporal', + 'operator', + 'cluster', + 'health', + '--address', + 'localhost:7233', + ]); + return output.includes('SERVING'); } -/** Check if the router container is running and healthy. */ -function isRouterReady(): boolean { - const status = runOutput('docker', ['inspect', '--format', '{{.State.Health.Status}}', 'shannon-router']); - return status === 'healthy'; +/** + * Ensure Temporal is running via compose. + */ +export async function ensureInfra(): Promise { + if (isTemporalReady()) { + return; + } + + const composeFile = getComposeFile(); + console.log('Starting Shannon infrastructure...'); + execFileSync('docker', ['compose', '-f', composeFile, 'up', '-d'], { stdio: 'inherit' }); + + console.log('Waiting for Temporal to be ready...'); + for (let i = 0; i < 30; i++) { + if (isTemporalReady()) { + console.log('Temporal is ready!'); + return; + } + await sleep(2000); + } + console.error('Timeout waiting for Temporal'); + process.exit(1); +} + +/** + * Build the worker image locally (local mode only). + */ +export function buildImage(noCache: boolean): void { + console.log(`Building ${DEV_IMAGE}...`); + const args = ['build']; + if (noCache) args.push('--no-cache'); + args.push('-t', DEV_IMAGE, '.'); + execFileSync('docker', args, { stdio: 'inherit' }); + console.log(`Build complete: ${DEV_IMAGE}`); +} + +/** + * Ensure the worker image is available. + * Local mode: auto-builds if missing. NPX mode: pulls from Docker Hub. + */ +export function ensureImage(version: string): void { + const image = getWorkerImage(version); + const exists = runQuiet('docker', ['image', 'inspect', image]); + if (exists) return; + + if (getMode() === 'local') { + console.log('Worker image not found, building...'); + buildImage(false); + } else { + console.log(`Pulling ${image}...`); + try { + execFileSync('docker', ['pull', image], { stdio: 'inherit' }); + } catch { + console.error(`\nERROR: Failed to pull ${image}`); + console.error('The image may not be available for your platform yet.'); + console.error('Check https://hub.docker.com/r/keygraph/shannon for available tags.'); + process.exit(1); + } + pruneOldImages(version); + } } /** @@ -73,259 +145,140 @@ function addHostFlag(): string[] { return []; } -/** Remove old keygraph/shannon images that don't match the current version. */ +export interface WorkerOptions { + version: string; + url: string; + repo: { hostPath: string; containerPath: string }; + workspacesDir: string; + taskQueue: string; + containerName: string; + envFlags: string[]; + config?: { hostPath: string; containerPath: string }; + credentials?: string; + promptsDir?: string; + outputDir?: string; + workspace: string; + pipelineTesting?: boolean; +} + +/** + * Spawn the worker container in detached mode and return the process. + */ +export function spawnWorker(opts: WorkerOptions): ChildProcess { + const args = ['run', '-d', '--rm', '--name', opts.containerName, '--network', 'shannon-net']; + + // Add host flag for Linux + args.push(...addHostFlag()); + + // UID remapping for Linux bind mounts + if (os.platform() === 'linux' && process.getuid && process.getgid) { + args.push('-e', `SHANNON_HOST_UID=${process.getuid()}`, '-e', `SHANNON_HOST_GID=${process.getgid()}`); + } + + // Volume mounts + args.push('-v', `${opts.workspacesDir}:/app/workspaces`); + args.push('-v', `${opts.repo.hostPath}:${opts.repo.containerPath}:ro`); + + // Writable overlays: shadow .shannon/ inside the :ro repo with workspace-backed dirs + const workspacePath = path.join(opts.workspacesDir, opts.workspace); + args.push('-v', `${path.join(workspacePath, 'deliverables')}:${opts.repo.containerPath}/.shannon/deliverables`); + args.push('-v', `${path.join(workspacePath, 'scratchpad')}:${opts.repo.containerPath}/.shannon/scratchpad`); + args.push('-v', `${path.join(workspacePath, '.playwright-cli')}:${opts.repo.containerPath}/.shannon/.playwright-cli`); + + // Local mode: mount prompts for live editing + if (opts.promptsDir) { + args.push('-v', `${opts.promptsDir}:/app/apps/worker/prompts:ro`); + } + + if (opts.config) { + args.push('-v', `${opts.config.hostPath}:${opts.config.containerPath}:ro`); + } + + // Output directory for deliverables copy + if (opts.outputDir) { + args.push('-v', `${opts.outputDir}:/app/output`); + } + + // Mount credentials file to fixed container path + if (opts.credentials) { + args.push('-v', `${opts.credentials}:/app/credentials/google-sa-key.json:ro`); + } + + // Environment + args.push(...opts.envFlags); + + // Container settings + args.push('--shm-size', '2gb', '--security-opt', 'seccomp=unconfined'); + + // Image + args.push(getWorkerImage(opts.version)); + + // Worker command + args.push('node', 'apps/worker/dist/temporal/worker.js', opts.url, opts.repo.containerPath); + args.push('--task-queue', opts.taskQueue); + if (opts.config) { + args.push('--config', opts.config.containerPath); + } + if (opts.outputDir) { + args.push('--output', '/app/output'); + } + args.push('--workspace', opts.workspace); + if (opts.pipelineTesting) { + args.push('--pipeline-testing'); + } + + // Prevent MSYS/Git Bash from converting Unix paths (e.g. /repos/my-repo) to Windows paths + return spawn('docker', args, { + stdio: 'pipe', + ...(os.platform() === 'win32' && { env: { ...process.env, MSYS_NO_PATHCONV: '1' } }), + }); +} + +/** + * Stop all running shannon-worker-* containers. + */ +export function stopWorkers(): void { + const workers = runOutput('docker', ['ps', '-q', '--filter', 'name=shannon-worker-']); + if (!workers) return; + + const ids = workers.split('\n').filter(Boolean); + console.log('Stopping worker containers...'); + execFileSync('docker', ['stop', ...ids], { stdio: 'inherit' }); +} + +/** + * Tear down the compose stack. + */ +export function stopInfra(clean: boolean): void { + const composeFile = getComposeFile(); + const args = ['compose', '-f', composeFile, 'down']; + if (clean) args.push('-v'); + execFileSync('docker', args, { stdio: 'inherit' }); +} + +/** + * Remove old keygraph/shannon images that don't match the current version. + */ function pruneOldImages(currentVersion: string): void { const output = runOutput('docker', ['images', NPX_IMAGE_REPO, '--format', '{{.Tag}}']); if (!output) return; - const stale = output.split('\n').filter((tag) => tag && tag !== currentVersion); + const currentTag = currentVersion; + const stale = output.split('\n').filter((tag) => tag && tag !== currentTag); for (const tag of stale) { runQuiet('docker', ['rmi', `${NPX_IMAGE_REPO}:${tag}`]); } } -// === DockerOrchestrator === - -/** Docker-based orchestration backend. */ -export class DockerOrchestrator implements Orchestrator { - getWorkerImage(version: string): string { - return getMode() === 'local' ? DEV_IMAGE : `${NPX_IMAGE_REPO}:${version}`; - } - - isTemporalReady(): boolean { - const output = runOutput('docker', [ - 'exec', - 'shannon-temporal', - 'temporal', - 'operator', - 'cluster', - 'health', - '--address', - 'localhost:7233', - ]); - return output.includes('SERVING'); - } - - async ensureInfra(useRouter: boolean): Promise { - const temporalReady = this.isTemporalReady(); - const routerNeeded = useRouter && !isRouterReady(); - - if (temporalReady && !routerNeeded) { - return; - } - - const composeFile = getComposeFile(); - const composeArgs = ['compose', '-f', composeFile]; - if (useRouter) composeArgs.push('--profile', 'router'); - composeArgs.push('up', '-d'); - - if (temporalReady && routerNeeded) { - console.log('Starting router...'); - } else { - console.log('Starting Shannon infrastructure...'); - } - execFileSync('docker', composeArgs, { stdio: 'inherit' }); - - // Wait for Temporal if it wasn't already running - if (!temporalReady) { - console.log('Waiting for Temporal to be ready...'); - for (let i = 0; i < 30; i++) { - if (this.isTemporalReady()) { - console.log('Temporal is ready!'); - break; - } - if (i === 29) { - console.error('Timeout waiting for Temporal'); - process.exit(1); - } - await sleep(2000); - } - } - - // Wait for router if needed - if (routerNeeded) { - console.log('Waiting for router to be ready...'); - for (let i = 0; i < 15; i++) { - if (isRouterReady()) { - console.log('Router is ready!'); - return; - } - await sleep(2000); - } - console.error('Timeout waiting for router'); - process.exit(1); - } - } - - ensureImage(version: string): void { - const image = this.getWorkerImage(version); - const exists = runQuiet('docker', ['image', 'inspect', image]); - if (exists) return; - - if (getMode() === 'local') { - console.log('Worker image not found, building...'); - this.buildImage(false); - } else { - console.log(`Pulling ${image}...`); - try { - execFileSync('docker', ['pull', image], { stdio: 'inherit' }); - } catch { - console.error(`\nERROR: Failed to pull ${image}`); - console.error('The image may not be available for your platform yet.'); - console.error('Check https://hub.docker.com/r/keygraph/shannon for available tags.'); - process.exit(1); - } - pruneOldImages(version); - } - } - - spawnWorker(opts: WorkerOptions): WorkerHandle { - const args = ['run', '-d', '--rm', '--name', opts.containerName, '--network', 'shannon-net']; - - // Add host flag for Linux - args.push(...addHostFlag()); - - // UID remapping for Linux bind mounts - if (os.platform() === 'linux' && process.getuid && process.getgid) { - args.push('-e', `SHANNON_HOST_UID=${process.getuid()}`, '-e', `SHANNON_HOST_GID=${process.getgid()}`); - } - - // Volume mounts - args.push('-v', `${opts.workspacesDir}:/app/workspaces`); - args.push('-v', `${opts.repo.hostPath}:${opts.repo.containerPath}:ro`); - - // Writable overlays: shadow .shannon/ inside the :ro repo with workspace-backed dirs - const workspacePath = path.join(opts.workspacesDir, opts.workspace); - args.push('-v', `${path.join(workspacePath, 'deliverables')}:${opts.repo.containerPath}/.shannon/deliverables`); - args.push('-v', `${path.join(workspacePath, 'scratchpad')}:${opts.repo.containerPath}/.shannon/scratchpad`); - args.push( - '-v', - `${path.join(workspacePath, '.playwright-cli')}:${opts.repo.containerPath}/.shannon/.playwright-cli`, - ); - - // Local mode: mount prompts for live editing - if (opts.promptsDir) { - args.push('-v', `${opts.promptsDir}:/app/apps/worker/prompts:ro`); - } - - if (opts.config) { - args.push('-v', `${opts.config.hostPath}:${opts.config.containerPath}:ro`); - } - - // Output directory for deliverables copy - if (opts.outputDir) { - args.push('-v', `${opts.outputDir}:/app/output`); - } - - // Mount credentials file to fixed container path - if (opts.credentials) { - args.push('-v', `${opts.credentials}:/app/credentials/google-sa-key.json:ro`); - } - - // Environment - args.push(...opts.envFlags); - - // Container settings - args.push('--shm-size', '2gb', '--security-opt', 'seccomp=unconfined'); - - // Image - args.push(this.getWorkerImage(opts.version)); - - // Worker command - args.push('node', 'apps/worker/dist/temporal/worker.js', opts.url, opts.repo.containerPath); - args.push('--task-queue', opts.taskQueue); - if (opts.config) { - args.push('--config', opts.config.containerPath); - } - if (opts.outputDir) { - args.push('--output', '/app/output'); - } - args.push('--workspace', opts.workspace); - if (opts.pipelineTesting) { - args.push('--pipeline-testing'); - } - - // Prevent MSYS/Git Bash from converting Unix paths (e.g. /repos/my-repo) to Windows paths - const proc = spawn('docker', args, { - stdio: 'pipe', - ...(os.platform() === 'win32' && { env: { ...process.env, MSYS_NO_PATHCONV: '1' } }), - }); - - return new DockerWorkerHandle(proc, opts.containerName); - } - - stopWorkers(): void { - const workers = runOutput('docker', ['ps', '-q', '--filter', 'name=shannon-worker-']); - if (!workers) return; - - const ids = workers.split('\n').filter(Boolean); - console.log('Stopping worker containers...'); - execFileSync('docker', ['stop', ...ids], { stdio: 'inherit' }); - } - - stopInfra(clean: boolean): void { - const composeFile = getComposeFile(); - const args = ['compose', '-f', composeFile, '--profile', 'router', 'down']; - if (clean) args.push('-v'); - execFileSync('docker', args, { stdio: 'inherit' }); - } - - listRunningWorkers(): string { - return runOutput('docker', [ - 'ps', - '--filter', - 'name=shannon-worker-', - '--format', - 'table {{.Names}}\t{{.Status}}\t{{.RunningFor}}', - ]); - } - - runEphemeral(image: string, args: string[], mounts: string[]): void { - const dockerArgs = ['run', '--rm']; - for (const mount of mounts) { - dockerArgs.push('-v', mount); - } - dockerArgs.push(image, ...args); - execFileSync('docker', dockerArgs, { - stdio: 'inherit', - ...(os.platform() === 'win32' && { env: { ...process.env, MSYS_NO_PATHCONV: '1' } }), - }); - } - - /** Build the worker image locally (local mode only). */ - buildImage(noCache: boolean): void { - console.log(`Building ${DEV_IMAGE}...`); - const args = ['build']; - if (noCache) args.push('--no-cache'); - args.push('-t', DEV_IMAGE, '.'); - execFileSync('docker', args, { stdio: 'inherit' }); - console.log(`Build complete: ${DEV_IMAGE}`); - } -} - -/** WorkerHandle wrapping a Docker container's ChildProcess. */ -class DockerWorkerHandle implements WorkerHandle { - constructor( - private readonly proc: ChildProcess, - private readonly containerName: string, - ) {} - - onError(cb: (err: Error) => void): void { - this.proc.on('error', cb); - } - - kill(): void { - try { - execFileSync('docker', ['stop', this.containerName], { stdio: 'pipe' }); - } catch { - // Container may have already exited - } - } -} - -// === Backward-compatible exports === - -// NOTE: Used by commands/build.ts which doesn't go through the orchestrator -export function buildImage(noCache: boolean): void { - new DockerOrchestrator().buildImage(noCache); +/** + * List running worker containers. + */ +export function listRunningWorkers(): string { + return runOutput('docker', [ + 'ps', + '--filter', + 'name=shannon-worker-', + '--format', + 'table {{.Names}}\t{{.Status}}\t{{.RunningFor}}', + ]); } diff --git a/apps/cli/src/env.ts b/apps/cli/src/env.ts index d06c1c8..183a60b 100644 --- a/apps/cli/src/env.ts +++ b/apps/cli/src/env.ts @@ -14,7 +14,6 @@ export const FORWARD_VARS = [ 'ANTHROPIC_API_KEY', 'ANTHROPIC_BASE_URL', 'ANTHROPIC_AUTH_TOKEN', - 'ROUTER_DEFAULT', 'CLAUDE_CODE_OAUTH_TOKEN', 'CLAUDE_CODE_USE_BEDROCK', 'AWS_REGION', @@ -27,8 +26,6 @@ export const FORWARD_VARS = [ 'ANTHROPIC_MEDIUM_MODEL', 'ANTHROPIC_LARGE_MODEL', 'CLAUDE_CODE_MAX_OUTPUT_TOKENS', - 'OPENAI_API_KEY', - 'OPENROUTER_API_KEY', ] as const; /** @@ -81,12 +78,7 @@ export function buildEnvRecord(): Record { interface CredentialValidation { valid: boolean; error?: string; - mode: 'api-key' | 'oauth' | 'custom-base-url' | 'bedrock' | 'vertex' | 'router'; -} - -/** Check if router credentials are present in the environment. */ -export function isRouterConfigured(): boolean { - return !!(process.env.ROUTER_DEFAULT && (process.env.OPENAI_API_KEY || process.env.OPENROUTER_API_KEY)); + mode: 'api-key' | 'oauth' | 'custom-base-url' | 'bedrock' | 'vertex'; } /** Check if a custom Anthropic-compatible base URL is configured. */ @@ -102,7 +94,6 @@ function detectProviders(): string[] { if (isCustomBaseUrlConfigured()) providers.push('Custom Base URL'); if (process.env.CLAUDE_CODE_USE_BEDROCK === '1') providers.push('AWS Bedrock'); if (process.env.CLAUDE_CODE_USE_VERTEX === '1') providers.push('Google Vertex'); - if (isRouterConfigured()) providers.push('Router'); return providers; } @@ -168,9 +159,6 @@ export function validateCredentials(): CredentialValidation { } return { valid: true, mode: 'vertex' }; } - if (isRouterConfigured()) { - return { valid: true, mode: 'router' }; - } const hint = getMode() === 'local' diff --git a/apps/cli/src/index.ts b/apps/cli/src/index.ts index 8b1a79a..8bbfef1 100644 --- a/apps/cli/src/index.ts +++ b/apps/cli/src/index.ts @@ -70,7 +70,6 @@ Options for 'start': -o, --output Copy deliverables to this directory after run -w, --workspace Named workspace (auto-resumes if exists) --pipeline-testing Use minimal prompts for fast testing - --router Route requests through claude-code-router Examples: ${prefix} start -u https://example.com -r ${mode === 'local' ? 'my-repo' : './my-repo'} @@ -95,7 +94,6 @@ interface ParsedStartArgs { workspace?: string; output?: string; pipelineTesting: boolean; - router: boolean; } function parseStartArgs(argv: string[]): ParsedStartArgs { @@ -105,7 +103,6 @@ function parseStartArgs(argv: string[]): ParsedStartArgs { let workspace: string | undefined; let output: string | undefined; let pipelineTesting = false; - let router = false; for (let i = 0; i < argv.length; i++) { const arg = argv[i]; @@ -150,9 +147,6 @@ function parseStartArgs(argv: string[]): ParsedStartArgs { case '--pipeline-testing': pipelineTesting = true; break; - case '--router': - router = true; - break; default: console.error(`Unknown option: ${arg}`); console.error(`Run "${getMode() === 'local' ? './shannon' : 'npx @keygraph/shannon'} help" for usage`); @@ -170,7 +164,6 @@ function parseStartArgs(argv: string[]): ParsedStartArgs { url, repo, pipelineTesting, - router, ...(config && { config }), ...(workspace && { workspace }), ...(output && { output }), diff --git a/apps/worker/src/ai/claude-executor.ts b/apps/worker/src/ai/claude-executor.ts index 32f4b49..b98202f 100644 --- a/apps/worker/src/ai/claude-executor.ts +++ b/apps/worker/src/ai/claude-executor.ts @@ -21,7 +21,6 @@ import { dispatchMessage } from './message-handlers.js'; import { type ModelTier, resolveModel } from './models.js'; import { detectExecutionContext, formatCompletionMessage, formatErrorOutput } from './output-formatters.js'; import { createProgressManager } from './progress-manager.js'; -import { getActualModelName } from './router-utils.js'; declare global { var SHANNON_DISABLE_LOADER: boolean | undefined; @@ -184,7 +183,6 @@ export async function runClaudePrompt( case 'litellm_router': if (providerConfig.baseUrl) sdkEnv.ANTHROPIC_BASE_URL = providerConfig.baseUrl; if (providerConfig.authToken) sdkEnv.ANTHROPIC_AUTH_TOKEN = providerConfig.authToken; - if (providerConfig.routerDefault) sdkEnv.ROUTER_DEFAULT = providerConfig.routerDefault; break; default: // 'anthropic_api' or unset — apiKey already handled above @@ -385,9 +383,8 @@ async function processMessageStream( if (dispatchResult.apiErrorDetected) { apiErrorDetected = true; } - // Capture model from SystemInitMessage, but override with router model if applicable if (dispatchResult.model) { - model = getActualModelName(dispatchResult.model); + model = dispatchResult.model; } } } diff --git a/apps/worker/src/ai/message-handlers.ts b/apps/worker/src/ai/message-handlers.ts index 148df45..81768a5 100644 --- a/apps/worker/src/ai/message-handlers.ts +++ b/apps/worker/src/ai/message-handlers.ts @@ -19,7 +19,6 @@ import { formatToolUseOutput, } from './output-formatters.js'; import type { ProgressManager } from './progress-manager.js'; -import { getActualModelName } from './router-utils.js'; import type { ApiErrorDetection, AssistantMessage, @@ -309,12 +308,10 @@ export async function dispatchMessage( case 'system': { if (message.subtype === 'init') { const initMsg = message as SystemInitMessage; - const actualModel = getActualModelName(initMsg.model); if (!execContext.useCleanOutput) { - logger.info(`Model: ${actualModel}, Permission: ${initMsg.permissionMode}`); + logger.info(`Model: ${initMsg.model}, Permission: ${initMsg.permissionMode}`); } - // Return actual model for tracking in audit logs - return { type: 'continue', model: actualModel }; + return { type: 'continue', model: initMsg.model }; } return { type: 'continue' }; } diff --git a/apps/worker/src/ai/router-utils.ts b/apps/worker/src/ai/router-utils.ts deleted file mode 100644 index b41185c..0000000 --- a/apps/worker/src/ai/router-utils.ts +++ /dev/null @@ -1,27 +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. - -/** - * Get the actual model name being used. - * When using claude-code-router, the SDK reports its configured model (claude-sonnet) - * but the actual model is determined by ROUTER_DEFAULT env var. - */ -export function getActualModelName(sdkReportedModel?: string): string | undefined { - const routerBaseUrl = process.env.ANTHROPIC_BASE_URL; - const routerDefault = process.env.ROUTER_DEFAULT; - - // If router mode is active and ROUTER_DEFAULT is set, use that - if (routerBaseUrl && routerDefault) { - // ROUTER_DEFAULT format: "provider,model" (e.g., "gemini,gemini-2.5-pro") - const parts = routerDefault.split(','); - if (parts.length >= 2) { - return parts.slice(1).join(','); // Handle model names with commas - } - } - - // Fall back to SDK-reported model - return sdkReportedModel; -} diff --git a/apps/worker/src/interfaces/checkpoint-provider.ts b/apps/worker/src/interfaces/checkpoint-provider.ts index 158bcb1..a066043 100644 --- a/apps/worker/src/interfaces/checkpoint-provider.ts +++ b/apps/worker/src/interfaces/checkpoint-provider.ts @@ -1,21 +1,59 @@ /** * CheckpointProvider — injectable interface for external state persistence. * - * Called after each agent completes to allow external progress tracking. - * During the concurrent vulnerability-exploitation phase, 5 pipelines run - * in parallel — onAgentComplete fires per-agent for granular progress. + * Called before and after each agent to support skip-guard (resume) and + * post-agent artifact persistence. During the concurrent vulnerability-exploitation + * phase, 5 pipelines run in parallel — methods fire per-agent for granular control. * - * Default: no-op. + * Default: no-op (skip nothing, persist nothing). */ -import type { PipelineState } from '../temporal/shared.js'; +import type { AgentMetrics, PipelineState } from '../temporal/shared.js'; + +/** Result of a pre-agent skip check. */ +export interface SkipDecision { + readonly skip: boolean; + readonly metrics?: AgentMetrics; // Required when skip=true +} + +/** File-system context passed after agent completion for artifact persistence. */ +export interface CheckpointContext { + readonly repoPath: string; + readonly sessionId: string; + readonly deliverablesSubdir: string; + readonly outputPath?: string; +} export interface CheckpointProvider { - onAgentComplete(agentName: string, phase: string, state: PipelineState): Promise; + /** + * Called before an agent activity executes. + * Return { skip: true, metrics } to skip the agent (e.g., output files already exist). + * Return { skip: false } to run normally. + */ + shouldSkipAgent( + agentName: string, + repoPath: string, + deliverablesSubdir: string, + ): Promise; + + /** + * Called after an agent activity succeeds. + * Receives pipeline state and optional file context for artifact persistence. + */ + onAgentComplete( + agentName: string, + phase: string, + state: PipelineState, + context?: CheckpointContext, + ): Promise; } /** Default no-op implementation — no external checkpointing. */ export class NoOpCheckpointProvider implements CheckpointProvider { + async shouldSkipAgent(): Promise { + return { skip: false }; + } + async onAgentComplete(): Promise { // No-op } diff --git a/apps/worker/src/interfaces/findings-provider.ts b/apps/worker/src/interfaces/findings-provider.ts index e111839..5770b0b 100644 --- a/apps/worker/src/interfaces/findings-provider.ts +++ b/apps/worker/src/interfaces/findings-provider.ts @@ -1,7 +1,7 @@ /** * FindingsProvider — injectable interface for external findings integration. * - * Allows external security data (SAST, SCA, secrets, etc.) to be merged + * Allows external security data from consumer-supplied sources to be merged * into the exploitation pipeline between vulnerability analysis and exploitation. * * Default: no-op returning { mergedCount: 0 }. diff --git a/apps/worker/src/interfaces/index.ts b/apps/worker/src/interfaces/index.ts index c73e3bc..fac478f 100644 --- a/apps/worker/src/interfaces/index.ts +++ b/apps/worker/src/interfaces/index.ts @@ -5,7 +5,9 @@ * Consumers can provide alternate implementations via the DI container. */ -export type { CheckpointProvider } from './checkpoint-provider.js'; +export type { CheckpointProvider, CheckpointContext, SkipDecision } from './checkpoint-provider.js'; export { NoOpCheckpointProvider } from './checkpoint-provider.js'; export type { FindingsProvider } from './findings-provider.js'; export { NoOpFindingsProvider } from './findings-provider.js'; +export type { ReportOutputProvider } from './report-output-provider.js'; +export { NoOpReportOutputProvider } from './report-output-provider.js'; diff --git a/apps/worker/src/interfaces/report-output-provider.ts b/apps/worker/src/interfaces/report-output-provider.ts new file mode 100644 index 0000000..05a7631 --- /dev/null +++ b/apps/worker/src/interfaces/report-output-provider.ts @@ -0,0 +1,22 @@ +/** + * ReportOutputProvider — injectable interface for emitting an optional + * additional artifact alongside the assembled markdown report. + * + * Runs after the report agent has finalized + * `comprehensive_security_assessment_report.md`. Consumers can override to + * produce derived outputs; the default no-op produces nothing. + */ + +import type { ActivityInput } from '../temporal/activities.js'; +import type { ActivityLogger } from '../types/activity-logger.js'; + +export interface ReportOutputProvider { + generate(input: ActivityInput, logger: ActivityLogger): Promise<{ outputPath?: string }>; +} + +/** Default no-op implementation — no additional output produced. */ +export class NoOpReportOutputProvider implements ReportOutputProvider { + async generate(): Promise<{ outputPath?: string }> { + return {}; + } +} diff --git a/apps/worker/src/services/container.ts b/apps/worker/src/services/container.ts index 05265a7..a9586b3 100644 --- a/apps/worker/src/services/container.ts +++ b/apps/worker/src/services/container.ts @@ -22,6 +22,8 @@ import type { CheckpointProvider } from '../interfaces/checkpoint-provider.js'; import { NoOpCheckpointProvider } from '../interfaces/checkpoint-provider.js'; import type { FindingsProvider } from '../interfaces/findings-provider.js'; import { NoOpFindingsProvider } from '../interfaces/findings-provider.js'; +import type { ReportOutputProvider } from '../interfaces/report-output-provider.js'; +import { NoOpReportOutputProvider } from '../interfaces/report-output-provider.js'; import type { ContainerConfig } from '../types/config.js'; import { AgentExecutionService } from './agent-execution.js'; import { ConfigLoaderService } from './config-loader.js'; @@ -40,6 +42,7 @@ export interface ContainerDependencies { readonly config: ContainerConfig; readonly findingsProvider?: FindingsProvider; readonly checkpointProvider?: CheckpointProvider; + readonly reportOutputProvider?: ReportOutputProvider; } /** @@ -59,6 +62,7 @@ export class Container { readonly exploitationChecker: ExploitationCheckerService; readonly findingsProvider: FindingsProvider; readonly checkpointProvider: CheckpointProvider; + readonly reportOutputProvider: ReportOutputProvider; constructor(deps: ContainerDependencies) { this.sessionMetadata = deps.sessionMetadata; @@ -72,6 +76,7 @@ export class Container { // Wire providers with default no-ops when not provided this.findingsProvider = deps.findingsProvider ?? new NoOpFindingsProvider(); this.checkpointProvider = deps.checkpointProvider ?? new NoOpCheckpointProvider(); + this.reportOutputProvider = deps.reportOutputProvider ?? new NoOpReportOutputProvider(); } } @@ -87,6 +92,32 @@ const DEFAULT_CONFIG: ContainerConfig = { auditDir: './workspaces', }; +/** + * Factory function for creating containers. + * + * Default: creates a plain Container with NoOp providers. Consumers can call + * setContainerFactory() at worker startup to inject custom provider + * implementations into every container. + */ +type ContainerFactory = ( + workflowId: string, + sessionMetadata: SessionMetadata, + config: ContainerConfig, +) => Container; + +let containerFactory: ContainerFactory = (_workflowId, sessionMetadata, config) => + new Container({ sessionMetadata, config }); + +/** + * Override the default container factory. + * + * Call once at worker startup to inject providers into all containers + * created during the worker's lifetime. + */ +export function setContainerFactory(factory: ContainerFactory): void { + containerFactory = factory; +} + /** * Get or create a Container for a workflow. * @@ -106,7 +137,7 @@ export function getOrCreateContainer( let container = containers.get(workflowId); if (!container) { - container = new Container({ sessionMetadata, config }); + container = containerFactory(workflowId, sessionMetadata, config); containers.set(workflowId, container); } diff --git a/apps/worker/src/services/index.ts b/apps/worker/src/services/index.ts index 2412f96..b864d27 100644 --- a/apps/worker/src/services/index.ts +++ b/apps/worker/src/services/index.ts @@ -16,7 +16,9 @@ export { AgentExecutionService } from './agent-execution.js'; export { ConfigLoaderService } from './config-loader.js'; export type { ContainerDependencies } from './container.js'; -export { Container, getContainer, getOrCreateContainer, removeContainer } from './container.js'; +export { Container, getContainer, getOrCreateContainer, removeContainer, setContainerFactory } from './container.js'; export { ExploitationCheckerService } from './exploitation-checker.js'; export { loadPrompt } from './prompt-manager.js'; export { assembleFinalReport, injectModelIntoReport } from './reporting.js'; +export type { ClaudePromptResult } from '../ai/claude-executor.js'; +export { runClaudePrompt } from '../ai/claude-executor.js'; diff --git a/apps/worker/src/services/preflight.ts b/apps/worker/src/services/preflight.ts index 1ce0944..ba4b234 100644 --- a/apps/worker/src/services/preflight.ts +++ b/apps/worker/src/services/preflight.ts @@ -14,7 +14,7 @@ * Checks run sequentially, cheapest first: * 1. Repository path exists and contains .git * 2. Config file parses and validates (if provided) - * 3. Credentials validate via Claude Agent SDK query (API key, OAuth, Bedrock, Vertex AI, or router mode) + * 3. Credentials validate via Claude Agent SDK query (API key, OAuth, Bedrock, or Vertex AI) * 4. Target URL is reachable from the container (DNS + HTTP) */ @@ -473,7 +473,7 @@ async function validateTargetUrl(targetUrl: string, logger: ActivityLogger): Pro * * 1. Repository path exists and contains .git * 2. Config file parses and validates (if configPath provided) - * 3. Credentials validate (API key, OAuth, or router mode) + * 3. Credentials validate (API key, OAuth, Bedrock, or Vertex AI) * 4. Target URL is reachable from the container * * Returns on first failure. diff --git a/apps/worker/src/services/reporting.ts b/apps/worker/src/services/reporting.ts index a9c294a..1cb11a6 100644 --- a/apps/worker/src/services/reporting.ts +++ b/apps/worker/src/services/reporting.ts @@ -17,7 +17,11 @@ interface DeliverableFile { } // Pure function: Assemble final report from specialist deliverables -export async function assembleFinalReport(sourceDir: string, logger: ActivityLogger): Promise { +export async function assembleFinalReport( + sourceDir: string, + deliverablesSubdir: string | undefined, + logger: ActivityLogger, +): Promise { const deliverableFiles: DeliverableFile[] = [ { name: 'Injection', path: 'injection_exploitation_evidence.md', required: false }, { name: 'XSS', path: 'xss_exploitation_evidence.md', required: false }, @@ -29,7 +33,7 @@ export async function assembleFinalReport(sourceDir: string, logger: ActivityLog const sections: string[] = []; for (const file of deliverableFiles) { - const filePath = path.join(deliverablesDir(sourceDir), file.path); + const filePath = path.join(deliverablesDir(sourceDir, deliverablesSubdir), file.path); try { if (await fs.pathExists(filePath)) { const content = await fs.readFile(filePath, 'utf8'); @@ -56,7 +60,7 @@ export async function assembleFinalReport(sourceDir: string, logger: ActivityLog } const finalContent = sections.join('\n\n'); - const outputDir = deliverablesDir(sourceDir); + const outputDir = deliverablesDir(sourceDir, deliverablesSubdir); const finalReportPath = path.join(outputDir, 'comprehensive_security_assessment_report.md'); try { @@ -82,6 +86,7 @@ export async function assembleFinalReport(sourceDir: string, logger: ActivityLog */ export async function injectModelIntoReport( repoPath: string, + deliverablesSubdir: string | undefined, outputPath: string, logger: ActivityLogger, ): Promise { @@ -118,7 +123,7 @@ export async function injectModelIntoReport( logger.info(`Injecting model info into report: ${modelStr}`); // 3. Read the final report - const reportPath = path.join(deliverablesDir(repoPath), 'comprehensive_security_assessment_report.md'); + const reportPath = path.join(deliverablesDir(repoPath, deliverablesSubdir), 'comprehensive_security_assessment_report.md'); if (!(await fs.pathExists(reportPath))) { logger.warn('Final report not found, skipping model injection'); diff --git a/apps/worker/src/session-manager.ts b/apps/worker/src/session-manager.ts index 9b808a1..d7fb981 100644 --- a/apps/worker/src/session-manager.ts +++ b/apps/worker/src/session-manager.ts @@ -103,7 +103,6 @@ export const AGENTS: Readonly> = Object.freez prerequisites: ['injection-exploit', 'xss-exploit', 'auth-exploit', 'ssrf-exploit', 'authz-exploit'], promptTemplate: 'report-executive', deliverableFilename: 'comprehensive_security_assessment_report.md', - modelTier: 'small', }, }); diff --git a/apps/worker/src/temporal/activities.ts b/apps/worker/src/temporal/activities.ts index 18e229b..de6063f 100644 --- a/apps/worker/src/temporal/activities.ts +++ b/apps/worker/src/temporal/activities.ts @@ -22,7 +22,8 @@ import { AuditSession } from '../audit/index.js'; import type { ResumeAttempt } from '../audit/metrics-tracker.js'; import type { SessionMetadata } from '../audit/utils.js'; import type { WorkflowSummary } from '../audit/workflow-logger.js'; -import { DEFAULT_DELIVERABLES_SUBDIR, deliverablesDir } from '../paths.js'; +import type { ContainerConfig, ProviderConfig } from '../types/config.js'; +import type { CheckpointContext } from '../interfaces/checkpoint-provider.js'; import { getContainer, getOrCreateContainer, removeContainer } from '../services/container.js'; import { classifyErrorForTemporal, PentestError } from '../services/error-handling.js'; import { ExploitationCheckerService } from '../services/exploitation-checker.js'; @@ -33,9 +34,9 @@ import { assembleFinalReport, injectModelIntoReport } from '../services/reportin import { AGENTS } from '../session-manager.js'; import type { AgentName } from '../types/agents.js'; import { ALL_AGENTS } from '../types/agents.js'; -import type { ContainerConfig, ProviderConfig } from '../types/config.js'; import { ErrorCode } from '../types/errors.js'; import { isErr } from '../types/result.js'; +import { DEFAULT_DELIVERABLES_SUBDIR, deliverablesDir } from '../paths.js'; import { fileExists, readJson } from '../utils/file-io.js'; import { createActivityLogger } from './activity-logger.js'; import type { AgentMetrics, PipelineState, ResumeState } from './shared.js'; @@ -131,6 +132,20 @@ function buildContainerConfig(input: ActivityInput): ContainerConfig { */ async function runAgentActivity(agentName: AgentName, input: ActivityInput): Promise { const { repoPath, configPath, pipelineTestingMode = false, workflowId, webUrl } = input; + + // Skip guard: the checkpoint provider decides whether to run the agent. + // The default NoOp provider always returns { skip: false }. + const skipContainer = getContainer(workflowId) ?? + getOrCreateContainer(workflowId, buildSessionMetadata(input), buildContainerConfig(input)); + const decision = await skipContainer.checkpointProvider.shouldSkipAgent( + agentName, + repoPath, + input.deliverablesSubdir ?? DEFAULT_DELIVERABLES_SUBDIR, + ); + if (decision.skip && decision.metrics) { + return decision.metrics; + } + const startTime = Date.now(); const attemptNumber = Context.current().info.attempt; @@ -288,7 +303,7 @@ export async function runReportAgent(input: ActivityInput): Promise { * Assemble the final report by concatenating exploitation evidence files. */ export async function assembleReportActivity(input: ActivityInput): Promise { - const { repoPath } = input; + const { repoPath, deliverablesSubdir } = input; const logger = createActivityLogger(); logger.info('Assembling deliverables from specialist agents...'); try { - await assembleFinalReport(repoPath, logger); + await assembleFinalReport(repoPath, deliverablesSubdir, logger); } catch (error) { const err = error as Error; logger.warn(`Error assembling final report: ${err.message}`); @@ -401,11 +408,11 @@ export async function assembleReportActivity(input: ActivityInput): Promise { - const { repoPath, sessionId, outputPath } = input; + const { repoPath, sessionId, outputPath, deliverablesSubdir } = input; const logger = createActivityLogger(); const effectiveOutputPath = outputPath ? path.join(outputPath, sessionId) : path.join('./workspaces', sessionId); try { - await injectModelIntoReport(repoPath, effectiveOutputPath, logger); + await injectModelIntoReport(repoPath, deliverablesSubdir, effectiveOutputPath, logger); } catch (error) { const err = error as Error; logger.warn(`Error injecting model into report: ${err.message}`); @@ -593,6 +600,18 @@ export async function restoreGitCheckpoint( const logger = createActivityLogger(); logger.info(`Restoring deliverables to ${checkpointHash}...`); + // Validate hash exists in this clone before attempting reset + try { + await executeGitCommandWithRetry( + ['git', 'rev-parse', '--verify', checkpointHash], + repoPath, + 'verify checkpoint hash exists' + ); + } catch { + logger.info(`Checkpoint hash not found in clone, skipping git reset: ${checkpointHash}`); + return; + } + await executeGitCommandWithRetry( ['git', 'reset', '--hard', checkpointHash], deliverablesPath, @@ -744,5 +763,42 @@ export async function saveCheckpoint( ): Promise { const container = getContainer(input.workflowId); if (!container?.checkpointProvider) return; - return container.checkpointProvider.onAgentComplete(agentName, phase, state); + + const context: CheckpointContext = { + repoPath: input.repoPath, + sessionId: input.sessionId, + deliverablesSubdir: input.deliverablesSubdir ?? DEFAULT_DELIVERABLES_SUBDIR, + ...(input.outputPath !== undefined && { outputPath: input.outputPath }), + }; + + return container.checkpointProvider.onAgentComplete(agentName, phase, state, context); +} + +/** + * Generate an optional additional output alongside the assembled markdown report. + * + * Delegates to the ReportOutputProvider registered in the DI container. + * Default: no-op. Consumers can override this activity at the worker level + * to emit derived outputs from the final report. + */ +export async function generateReportOutputActivity(input: ActivityInput): Promise { + const container = getContainer(input.workflowId); + if (!container?.reportOutputProvider) return; + + const logger = createActivityLogger(); + + // Resolve promptDir against the worker root so providers are cwd-independent. + const resolvedInput: ActivityInput = { + ...input, + ...(input.promptDir !== undefined && { + promptDir: path.isAbsolute(input.promptDir) + ? input.promptDir + : path.resolve(process.env.SHANNON_WORKER_ROOT ?? process.cwd(), input.promptDir), + }), + }; + + const result = await container.reportOutputProvider.generate(resolvedInput, logger); + if (result.outputPath) { + logger.info(`Report output written to ${result.outputPath}`); + } } diff --git a/apps/worker/src/temporal/shared.ts b/apps/worker/src/temporal/shared.ts index 583fce3..32519cf 100644 --- a/apps/worker/src/temporal/shared.ts +++ b/apps/worker/src/temporal/shared.ts @@ -25,10 +25,10 @@ export interface PipelineInput { deliverablesSubdir?: string; // Override deliverables path (default: '.shannon/deliverables') auditDir?: string; // Override audit log directory (default: './workspaces') promptDir?: string; // Override prompt template directory - sastSarifPath?: string; // Path to SARIF file (gates SAST-enhanced mode) + sastSarifPath?: string; // Optional path for consumer-supplied findings input checkpointsEnabled?: boolean; // Enable checkpoint activities (default: false) skipGitCheck?: boolean; // Skip .git directory validation in preflight (e.g. when .git is removed after clone) - providerConfig?: ProviderConfig; // LLM provider configuration (Bedrock, Vertex, LiteLLM, etc.) + providerConfig?: ProviderConfig; // LLM provider configuration (Bedrock, Vertex, etc.) } export interface ResumeState { diff --git a/apps/worker/src/temporal/workflows.ts b/apps/worker/src/temporal/workflows.ts index 8ae27c8..cb6516b 100644 --- a/apps/worker/src/temporal/workflows.ts +++ b/apps/worker/src/temporal/workflows.ts @@ -332,30 +332,14 @@ export async function pentestPipeline(input: PipelineInput): Promise[]): void { const failedPipelines: string[] = []; for (const result of results) { - if (result.status === 'fulfilled') { - const { vulnType, vulnMetrics, exploitMetrics } = result.value; - - const vulnAgentName = `${vulnType}-vuln`; - if (vulnMetrics) { - state.agentMetrics[vulnAgentName] = vulnMetrics; - state.completedAgents.push(vulnAgentName); - } else if (shouldSkip(vulnAgentName)) { - state.completedAgents.push(vulnAgentName); - } - - const exploitAgentName = `${vulnType}-exploit`; - if (exploitMetrics) { - state.agentMetrics[exploitAgentName] = exploitMetrics; - state.completedAgents.push(exploitAgentName); - } else if (shouldSkip(exploitAgentName)) { - state.completedAgents.push(exploitAgentName); - } - } else { + if (result.status === 'rejected') { const errorMsg = result.reason instanceof Error ? result.reason.message : String(result.reason); failedPipelines.push(errorMsg); } @@ -442,14 +426,17 @@ export async function pentestPipeline(input: PipelineInput): Promise; readonly supportsStructuredOutput?: boolean; } diff --git a/apps/worker/src/utils/billing-detection.ts b/apps/worker/src/utils/billing-detection.ts index 6f25c72..dd019c2 100644 --- a/apps/worker/src/utils/billing-detection.ts +++ b/apps/worker/src/utils/billing-detection.ts @@ -26,7 +26,6 @@ export const BILLING_TEXT_PATTERNS = [ 'cap reached', 'budget exceeded', 'usage limit', - 'resets', ] as const; /**