feat: add Temporal worker, client, and query tools (phase 5)

- Add worker.ts with workflow bundling and graceful shutdown
- Add client.ts CLI to start pipelines with progress polling
- Add query.ts CLI to inspect running workflow state
- Fix buffer overflow by truncating error messages and stack traces
- Skip git operations gracefully on non-git repositories
- Add kill.sh/start.sh dev scripts and Dockerfile.worker
This commit is contained in:
ajmallesh
2026-01-12 14:27:59 -08:00
parent eb7eced23f
commit 05f8e2382c
9 changed files with 601 additions and 10 deletions
+20
View File
@@ -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
+43
View File
@@ -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"]
+21 -5
View File
@@ -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:
+6 -1
View File
@@ -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",
+32 -3
View File
@@ -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);
+209
View File
@@ -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 -- <webUrl> <repoPath> [options]
* # or
* node dist/temporal/client.js <webUrl> <repoPath> [options]
*
* Options:
* --config <path> Configuration file path
* --output <path> Output directory for audit logs
* --pipeline-testing Use minimal prompts for fast testing
* --workflow-id <id> Custom workflow ID (default: shannon-<timestamp>)
* --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 <webUrl> <repoPath> [options]\n'
);
console.log(chalk.yellow('Options:'));
console.log(' --config <path> Configuration file path');
console.log(' --output <path> Output directory for audit logs');
console.log(' --pipeline-testing Use minimal prompts for fast testing');
console.log(
' --workflow-id <id> Custom workflow ID (default: shannon-<timestamp>)'
);
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<void> {
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<PipelineState>>(
'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<PipelineProgress>(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);
});
+155
View File
@@ -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 -- <workflowId>
* # or
* node dist/temporal/query.js <workflowId>
*
* 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<string, AgentMetrics>;
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 <workflowId>\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<void> {
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<PipelineProgress>(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);
});
+79
View File
@@ -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<void> {
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<void> => {
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);
});
+36 -1
View File
@@ -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<boolean> {
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<GitOperationResult> {
// 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<GitOperationResult> {
// 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<GitOperationResult> {
// 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<string | null> {
if (!(await isGitRepository(sourceDir))) {
return null;
}
try {
const result = await $`cd ${sourceDir} && git rev-parse HEAD`;
return result.stdout.trim();