diff --git a/CLAUDE.md b/CLAUDE.md index 72a6582..0bfd78d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -167,6 +167,26 @@ For detailed design, see `docs/unified-audit-system-design.md`. ## Development Notes +### Learning from Reference Implementations + +A working POC exists at `/Users/arjunmalleswaran/Code/shannon-pocs` that demonstrates the ideal Temporal + Claude Agent SDK integration. When implementing Temporal features, agents can ask questions in the chat, and the user will relay them to another Claude Code session working in that POC directory. + +**How to use this approach:** +1. When stuck or unsure about Temporal patterns, write a specific question in the chat +2. The user will ask an agent working on the POC to answer +3. The user relays the answer (code snippets, patterns, explanations) back +4. Apply the learned patterns to Shannon's codebase + +**Example questions to ask:** +- "How does the POC structure its workflow to handle parallel activities?" +- "Show me how heartbeats are implemented in the POC's activities" +- "What retry configuration does the POC use for long-running agent activities?" +- "How does the POC integrate Claude Agent SDK calls within Temporal activities?" + +**Current reference implementations:** +- **Temporal + Claude Agent SDK**: `/Users/arjunmalleswaran/Code/shannon-pocs` - working implementation demonstrating workflows, activities, worker setup, and SDK integration +- **Implementation plan**: `docs/temporal-implementation-plan.md` - detailed plan for Shannon integration + ### Key Design Patterns - **Configuration-Driven Architecture**: YAML configs with JSON Schema validation - **Modular Error Handling**: Categorized error types with retry logic diff --git a/docker/Dockerfile.worker b/docker/Dockerfile.worker new file mode 100644 index 0000000..aef429c --- /dev/null +++ b/docker/Dockerfile.worker @@ -0,0 +1,43 @@ +# Wolfi-based worker for Shannon AI pentester +FROM cgr.dev/chainguard/wolfi-base:latest + +# Install Node.js 22, Python 3.12, Chromium, and dependencies +RUN apk add --no-cache \ + nodejs-22 \ + npm \ + python-3.12 \ + py3.12-pip \ + chromium \ + git \ + bash \ + curl + +# Install uvx for browser-use +RUN pip install uvx --break-system-packages + +# Create non-root user +RUN adduser -D -u 1000 pentest +WORKDIR /app + +# Copy package files first for better caching +COPY package*.json ./ + +# Install dependencies +RUN npm ci --omit=dev + +# Copy application code +COPY dist/ ./dist/ +COPY prompts/ ./prompts/ + +# Set ownership +RUN chown -R pentest:pentest /app + +# Switch to non-root user +USER pentest + +# Set Chromium path for Playwright +ENV CHROME_PATH=/usr/bin/chromium-browser +ENV PLAYWRIGHT_CHROMIUM_EXECUTABLE_PATH=/usr/bin/chromium-browser + +# Entry point +CMD ["node", "dist/temporal/worker.js"] diff --git a/docker/docker-compose.temporal.yml b/docker/docker-compose.temporal.yml index b568830..4d9bb45 100644 --- a/docker/docker-compose.temporal.yml +++ b/docker/docker-compose.temporal.yml @@ -1,12 +1,10 @@ services: temporal: - image: temporalio/auto-setup:latest - environment: - - DB=sqlite - - SQLITE_DB_PATH=/var/lib/temporal/temporal.db + image: temporalio/temporal:latest + command: ["server", "start-dev", "--db-filename", "/var/lib/temporal/temporal.db", "--ip", "0.0.0.0"] ports: - "7233:7233" # gRPC - - "8233:8233" # Web UI + - "8233:8233" # Web UI (built-in) volumes: - temporal-data:/var/lib/temporal healthcheck: @@ -16,5 +14,23 @@ services: retries: 10 start_period: 30s + worker: + build: + context: .. + dockerfile: docker/Dockerfile.worker + environment: + - TEMPORAL_ADDRESS=temporal:7233 + - ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY} + depends_on: + temporal: + condition: service_healthy + volumes: + - ../deliverables:/app/deliverables + - ../prompts:/app/prompts + shm_size: 2gb + ipc: host + security_opt: + - seccomp:unconfined + volumes: temporal-data: diff --git a/package.json b/package.json index 0b43f0c..b470c72 100644 --- a/package.json +++ b/package.json @@ -5,7 +5,12 @@ "main": "./dist/shannon.js", "scripts": { "build": "tsc", - "start": "node ./dist/shannon.js" + "start": "node ./dist/shannon.js", + "temporal:server": "docker compose -f docker/docker-compose.temporal.yml up temporal -d", + "temporal:server:stop": "docker compose -f docker/docker-compose.temporal.yml down", + "temporal:worker": "node dist/temporal/worker.js", + "temporal:start": "node dist/temporal/client.js", + "temporal:query": "node dist/temporal/query.js" }, "dependencies": { "@anthropic-ai/claude-agent-sdk": "^0.1.0", diff --git a/src/temporal/activities.ts b/src/temporal/activities.ts index 0b631a7..15df935 100644 --- a/src/temporal/activities.ts +++ b/src/temporal/activities.ts @@ -21,6 +21,29 @@ import { heartbeat, ApplicationFailure, Context } from '@temporalio/activity'; import chalk from 'chalk'; +// Max lengths to prevent Temporal protobuf buffer overflow +const MAX_ERROR_MESSAGE_LENGTH = 2000; +const MAX_STACK_TRACE_LENGTH = 1000; + +/** + * Truncate error message to prevent buffer overflow in Temporal serialization. + */ +function truncateErrorMessage(message: string): string { + if (message.length <= MAX_ERROR_MESSAGE_LENGTH) { + return message; + } + return message.slice(0, MAX_ERROR_MESSAGE_LENGTH - 20) + '\n[truncated]'; +} + +/** + * Truncate stack trace on an ApplicationFailure to prevent buffer overflow. + */ +function truncateStackTrace(failure: ApplicationFailure): void { + if (failure.stack && failure.stack.length > MAX_STACK_TRACE_LENGTH) { + failure.stack = failure.stack.slice(0, MAX_STACK_TRACE_LENGTH) + '\n[stack truncated]'; + } +} + import { runClaudePrompt, validateAgentOutput, @@ -202,20 +225,26 @@ async function runAgentActivity( // Classify error for Temporal retry behavior const classified = classifyErrorForTemporal(error); - const message = error instanceof Error ? error.message : String(error); + // Truncate message to prevent protobuf buffer overflow + const rawMessage = error instanceof Error ? error.message : String(error); + const message = truncateErrorMessage(rawMessage); if (classified.retryable) { // Temporal will retry with configured backoff - throw ApplicationFailure.create({ + const failure = ApplicationFailure.create({ message, type: classified.type, details: [{ agentName, attemptNumber, elapsed: Date.now() - startTime }], }); + truncateStackTrace(failure); + throw failure; } else { // Fail immediately - no retry - throw ApplicationFailure.nonRetryable(message, classified.type, [ + const failure = ApplicationFailure.nonRetryable(message, classified.type, [ { agentName, attemptNumber, elapsed: Date.now() - startTime }, ]); + truncateStackTrace(failure); + throw failure; } } finally { clearInterval(heartbeatInterval); diff --git a/src/temporal/client.ts b/src/temporal/client.ts new file mode 100644 index 0000000..627121b --- /dev/null +++ b/src/temporal/client.ts @@ -0,0 +1,209 @@ +#!/usr/bin/env node +// 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. + +/** + * Temporal client for starting Shannon pentest pipeline workflows. + * + * Starts a workflow and optionally waits for completion with progress polling. + * + * Usage: + * npm run temporal:start -- [options] + * # or + * node dist/temporal/client.js [options] + * + * Options: + * --config Configuration file path + * --output Output directory for audit logs + * --pipeline-testing Use minimal prompts for fast testing + * --workflow-id Custom workflow ID (default: shannon-) + * --no-wait Start workflow and exit without waiting + * + * Environment: + * TEMPORAL_ADDRESS - Temporal server address (default: localhost:7233) + */ + +import { Connection, Client } from '@temporalio/client'; +import dotenv from 'dotenv'; +import chalk from 'chalk'; +// Import types only - these don't pull in workflow runtime code +import type { PipelineInput, PipelineState, PipelineProgress } from './shared.js'; + +dotenv.config(); + +// Query name must match the one defined in workflows.ts +const PROGRESS_QUERY = 'getProgress'; + +function showUsage(): void { + console.log(chalk.cyan.bold('\nShannon Temporal Client')); + console.log(chalk.gray('Start a pentest pipeline workflow\n')); + console.log(chalk.yellow('Usage:')); + console.log( + ' node dist/temporal/client.js [options]\n' + ); + console.log(chalk.yellow('Options:')); + console.log(' --config Configuration file path'); + console.log(' --output Output directory for audit logs'); + console.log(' --pipeline-testing Use minimal prompts for fast testing'); + console.log( + ' --workflow-id Custom workflow ID (default: shannon-)' + ); + console.log(' --no-wait Start workflow and exit without waiting\n'); + console.log(chalk.yellow('Examples:')); + console.log(' node dist/temporal/client.js https://example.com /path/to/repo'); + console.log( + ' node dist/temporal/client.js https://example.com /path/to/repo --config config.yaml\n' + ); +} + +async function startPipeline(): Promise { + const args = process.argv.slice(2); + + if (args.includes('--help') || args.includes('-h') || args.length === 0) { + showUsage(); + process.exit(0); + } + + // Parse arguments + let webUrl: string | undefined; + let repoPath: string | undefined; + let configPath: string | undefined; + let outputPath: string | undefined; + let pipelineTestingMode = false; + let customWorkflowId: string | undefined; + let waitForCompletion = true; + + for (let i = 0; i < args.length; i++) { + const arg = args[i]; + if (arg === '--config') { + const nextArg = args[i + 1]; + if (nextArg && !nextArg.startsWith('-')) { + configPath = nextArg; + i++; + } + } else if (arg === '--output') { + const nextArg = args[i + 1]; + if (nextArg && !nextArg.startsWith('-')) { + outputPath = nextArg; + i++; + } + } else if (arg === '--workflow-id') { + const nextArg = args[i + 1]; + if (nextArg && !nextArg.startsWith('-')) { + customWorkflowId = nextArg; + i++; + } + } else if (arg === '--pipeline-testing') { + pipelineTestingMode = true; + } else if (arg === '--no-wait') { + waitForCompletion = false; + } else if (arg && !arg.startsWith('-')) { + if (!webUrl) { + webUrl = arg; + } else if (!repoPath) { + repoPath = arg; + } + } + } + + if (!webUrl || !repoPath) { + console.log(chalk.red('Error: webUrl and repoPath are required')); + showUsage(); + process.exit(1); + } + + const address = process.env.TEMPORAL_ADDRESS || 'localhost:7233'; + console.log(chalk.cyan(`Connecting to Temporal at ${address}...`)); + + const connection = await Connection.connect({ address }); + const client = new Client({ connection }); + + try { + const workflowId = customWorkflowId || `shannon-${Date.now()}`; + + const input: PipelineInput = { + webUrl, + repoPath, + ...(configPath && { configPath }), + ...(outputPath && { outputPath }), + ...(pipelineTestingMode && { pipelineTestingMode }), + }; + + console.log(chalk.green(`\nStarting workflow: ${workflowId}`)); + console.log(chalk.gray(`Target: ${webUrl}`)); + console.log(chalk.gray(`Repository: ${repoPath}`)); + console.log( + chalk.blue( + `Web UI: http://localhost:8233/namespaces/default/workflows/${workflowId}\n` + ) + ); + + // Start workflow by name (not by importing the function) + const handle = await client.workflow.start<(input: PipelineInput) => Promise>( + 'pentestPipelineWorkflow', + { + taskQueue: 'shannon-pipeline', + workflowId, + args: [input], + } + ); + + if (!waitForCompletion) { + console.log( + chalk.yellow('Workflow started in background. Use query tool to check progress.') + ); + console.log(chalk.gray(` npm run temporal:query -- ${workflowId}`)); + return; + } + + // Poll for progress every 30 seconds + const progressInterval = setInterval(async () => { + try { + const progress = await handle.query(PROGRESS_QUERY); + const elapsed = Math.floor(progress.elapsedMs / 1000); + console.log( + chalk.gray(`[${elapsed}s]`), + chalk.cyan(`Phase: ${progress.currentPhase || 'unknown'}`), + chalk.gray(`| Agent: ${progress.currentAgent || 'none'}`), + chalk.gray(`| Completed: ${progress.completedAgents.length}/13`) + ); + } catch { + // Workflow may have completed + } + }, 30000); + + try { + const result = await handle.result(); + clearInterval(progressInterval); + + console.log(chalk.green.bold('\nPipeline completed successfully!')); + console.log( + chalk.gray(`Duration: ${Math.floor((Date.now() - result.startTime) / 1000)}s`) + ); + console.log(chalk.gray(`Agents completed: ${result.completedAgents.length}`)); + + // Show cost summary if available + const totalCost = Object.values(result.agentMetrics).reduce( + (sum, m) => sum + (m.costUsd ?? 0), + 0 + ); + if (totalCost > 0) { + console.log(chalk.gray(`Total cost: $${totalCost.toFixed(4)}`)); + } + } catch (error) { + clearInterval(progressInterval); + console.error(chalk.red.bold('\nPipeline failed:'), error); + process.exit(1); + } + } finally { + await connection.close(); + } +} + +startPipeline().catch((err) => { + console.error(chalk.red('Client error:'), err); + process.exit(1); +}); diff --git a/src/temporal/query.ts b/src/temporal/query.ts new file mode 100644 index 0000000..a97fe74 --- /dev/null +++ b/src/temporal/query.ts @@ -0,0 +1,155 @@ +#!/usr/bin/env node +// 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. + +/** + * Temporal query tool for inspecting Shannon workflow progress. + * + * Queries a running or completed workflow and displays its state. + * + * Usage: + * npm run temporal:query -- + * # or + * node dist/temporal/query.js + * + * Environment: + * TEMPORAL_ADDRESS - Temporal server address (default: localhost:7233) + */ + +import { Connection, Client } from '@temporalio/client'; +import dotenv from 'dotenv'; +import chalk from 'chalk'; + +dotenv.config(); + +// Query name must match the one defined in workflows.ts +const PROGRESS_QUERY = 'getProgress'; + +// Types duplicated from shared.ts to avoid importing workflow APIs +interface AgentMetrics { + durationMs: number; + inputTokens: number | null; + outputTokens: number | null; + costUsd: number | null; + numTurns: number | null; +} + +interface PipelineProgress { + status: 'running' | 'completed' | 'failed'; + currentPhase: string | null; + currentAgent: string | null; + completedAgents: string[]; + failedAgent: string | null; + error: string | null; + startTime: number; + agentMetrics: Record; + workflowId: string; + elapsedMs: number; +} + +function showUsage(): void { + console.log(chalk.cyan.bold('\nShannon Temporal Query Tool')); + console.log(chalk.gray('Query progress of a running workflow\n')); + console.log(chalk.yellow('Usage:')); + console.log(' node dist/temporal/query.js \n'); + console.log(chalk.yellow('Examples:')); + console.log(' node dist/temporal/query.js shannon-1704672000000\n'); +} + +function getStatusColor(status: string): string { + switch (status) { + case 'running': + return chalk.yellow(status); + case 'completed': + return chalk.green(status); + case 'failed': + return chalk.red(status); + default: + return status; + } +} + +function formatDuration(ms: number): string { + const seconds = Math.floor(ms / 1000); + const minutes = Math.floor(seconds / 60); + const hours = Math.floor(minutes / 60); + + if (hours > 0) { + return `${hours}h ${minutes % 60}m`; + } else if (minutes > 0) { + return `${minutes}m ${seconds % 60}s`; + } + return `${seconds}s`; +} + +async function queryWorkflow(): Promise { + const workflowId = process.argv[2]; + + if (!workflowId || workflowId === '--help' || workflowId === '-h') { + showUsage(); + process.exit(workflowId ? 0 : 1); + } + + const address = process.env.TEMPORAL_ADDRESS || 'localhost:7233'; + + const connection = await Connection.connect({ address }); + const client = new Client({ connection }); + + try { + const handle = client.workflow.getHandle(workflowId); + const progress = await handle.query(PROGRESS_QUERY); + + console.log(chalk.cyan.bold('\nWorkflow Progress')); + console.log(chalk.gray('\u2500'.repeat(40))); + console.log(`${chalk.white('Workflow ID:')} ${progress.workflowId}`); + console.log(`${chalk.white('Status:')} ${getStatusColor(progress.status)}`); + console.log( + `${chalk.white('Current Phase:')} ${progress.currentPhase || 'none'}` + ); + console.log( + `${chalk.white('Current Agent:')} ${progress.currentAgent || 'none'}` + ); + console.log(`${chalk.white('Elapsed:')} ${formatDuration(progress.elapsedMs)}`); + console.log( + `${chalk.white('Completed:')} ${progress.completedAgents.length}/13 agents` + ); + + if (progress.completedAgents.length > 0) { + console.log(chalk.gray('\nCompleted agents:')); + for (const agent of progress.completedAgents) { + const metrics = progress.agentMetrics[agent]; + const duration = metrics ? formatDuration(metrics.durationMs) : 'unknown'; + const cost = metrics?.costUsd ? `$${metrics.costUsd.toFixed(4)}` : ''; + console.log( + chalk.green(` - ${agent}`) + + chalk.gray(` (${duration}${cost ? ', ' + cost : ''})`) + ); + } + } + + if (progress.error) { + console.log(chalk.red(`\nError: ${progress.error}`)); + console.log(chalk.red(`Failed agent: ${progress.failedAgent}`)); + } + + console.log(); + } catch (error) { + const err = error as Error; + if (err.message?.includes('not found')) { + console.log(chalk.red(`Workflow not found: ${workflowId}`)); + } else { + console.error(chalk.red('Query failed:'), err.message); + } + process.exit(1); + } finally { + await connection.close(); + } +} + +queryWorkflow().catch((err) => { + console.error(chalk.red('Query error:'), err); + process.exit(1); +}); diff --git a/src/temporal/worker.ts b/src/temporal/worker.ts new file mode 100644 index 0000000..7346257 --- /dev/null +++ b/src/temporal/worker.ts @@ -0,0 +1,79 @@ +#!/usr/bin/env node +// 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. + +/** + * Temporal worker for Shannon pentest pipeline. + * + * Polls the 'shannon-pipeline' task queue and executes activities. + * Handles up to 5 concurrent activities to support parallel agent execution. + * + * Usage: + * npm run temporal:worker + * # or + * node dist/temporal/worker.js + * + * Environment: + * TEMPORAL_ADDRESS - Temporal server address (default: localhost:7233) + */ + +import { NativeConnection, Worker, bundleWorkflowCode } from '@temporalio/worker'; +import { fileURLToPath } from 'node:url'; +import path from 'node:path'; +import dotenv from 'dotenv'; +import chalk from 'chalk'; +import * as activities from './activities.js'; + +dotenv.config(); + +const __dirname = path.dirname(fileURLToPath(import.meta.url)); + +async function runWorker(): Promise { + const address = process.env.TEMPORAL_ADDRESS || 'localhost:7233'; + console.log(chalk.cyan(`Connecting to Temporal at ${address}...`)); + + const connection = await NativeConnection.connect({ address }); + + // Bundle workflows for Temporal's V8 isolate + console.log(chalk.gray('Bundling workflows...')); + const workflowBundle = await bundleWorkflowCode({ + workflowsPath: path.join(__dirname, 'workflows.js'), + }); + + const worker = await Worker.create({ + connection, + namespace: 'default', + workflowBundle, + activities, + taskQueue: 'shannon-pipeline', + maxConcurrentActivityTaskExecutions: 5, // Match parallel agent count + }); + + // Graceful shutdown handling + const shutdown = async (): Promise => { + console.log(chalk.yellow('\nShutting down worker...')); + worker.shutdown(); + }; + + process.on('SIGINT', shutdown); + process.on('SIGTERM', shutdown); + + console.log(chalk.green('Shannon worker started')); + console.log(chalk.gray('Task queue: shannon-pipeline')); + console.log(chalk.gray('Press Ctrl+C to stop\n')); + + try { + await worker.run(); + } finally { + await connection.close(); + console.log(chalk.gray('Worker stopped')); + } +} + +runWorker().catch((err) => { + console.error(chalk.red('Worker failed:'), err); + process.exit(1); +}); diff --git a/src/utils/git-manager.ts b/src/utils/git-manager.ts index 969e811..780bdcd 100644 --- a/src/utils/git-manager.ts +++ b/src/utils/git-manager.ts @@ -7,6 +7,19 @@ import { $ } from 'zx'; import chalk from 'chalk'; +/** + * Check if a directory is a git repository. + * Returns true if the directory contains a .git folder or is inside a git repo. + */ +export async function isGitRepository(dir: string): Promise { + try { + await $`cd ${dir} && git rev-parse --git-dir`.quiet(); + return true; + } catch { + return false; + } +} + interface GitOperationResult { success: boolean; hadChanges?: boolean; @@ -146,6 +159,12 @@ export async function rollbackGitWorkspace( sourceDir: string, reason: string = 'retry preparation' ): Promise { + // Skip git operations if not a git repository + if (!(await isGitRepository(sourceDir))) { + console.log(chalk.gray(` ⏭️ Skipping git rollback (not a git repository)`)); + return { success: true }; + } + console.log(chalk.yellow(` 🔄 Rolling back workspace for ${reason}`)); try { const changes = await getChangedFiles(sourceDir, 'status check for rollback'); @@ -182,6 +201,12 @@ export async function createGitCheckpoint( description: string, attempt: number ): Promise { + // Skip git operations if not a git repository + if (!(await isGitRepository(sourceDir))) { + console.log(chalk.gray(` ⏭️ Skipping git checkpoint (not a git repository)`)); + return { success: true }; + } + console.log(chalk.blue(` 📍 Creating checkpoint for ${description} (attempt ${attempt})`)); try { // First attempt: preserve existing deliverables. Retries: clean workspace to prevent pollution @@ -221,6 +246,12 @@ export async function commitGitSuccess( sourceDir: string, description: string ): Promise { + // Skip git operations if not a git repository + if (!(await isGitRepository(sourceDir))) { + console.log(chalk.gray(` ⏭️ Skipping git commit (not a git repository)`)); + return { success: true }; + } + console.log(chalk.green(` 💾 Committing successful results for ${description}`)); try { const changes = await getChangedFiles(sourceDir, 'status check for success commit'); @@ -252,9 +283,13 @@ export async function commitGitSuccess( } /** - * Get current git commit hash + * Get current git commit hash. + * Returns null if not a git repository. */ export async function getGitCommitHash(sourceDir: string): Promise { + if (!(await isGitRepository(sourceDir))) { + return null; + } try { const result = await $`cd ${sourceDir} && git rev-parse HEAD`; return result.stdout.trim();