CalvinBackup/shannon
1
0
Fork 0
mirror of https://github.com/KeygraphHQ/shannon.git synced 2026-07-23 13:01:00 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor(worker): converge shared core with shannon-oss (#388)

Browse Source 
* fix(worker): port keygraph shared-core correctness fixes

* refactor(worker): adopt collectors/ and ai/pi/ layout; add task budget cap and cancellation

* refactor(worker): drop inconsistent Collector "Server" suffix

* refactor(worker): drop unused providerConfig/apiKey seams, resolve credentials from env only

* refactor(worker): port oss code_path pattern expansion + external_directory allow

* fix(worker): preserve dotfile paths in code_path avoid patterns (.env no longer stripped to env)

* feat(worker): render Unprocessed Vulnerabilities section in exploit deliverable (align with oss)

* feat(worker): request set_blind_spots for all vuln classes (align auth/ssrf with production prompts)

* refactor(worker): adopt unified permissionSystem* naming and helper layout

* refactor(worker): inline blind_spots into vuln deliverable section array

* chore(worker): drop unused zod dependency (tree is typebox-native)

* fix(worker): normalize base32 TOTP secret to accept padding and whitespace

* refactor(worker): adopt shared toolResult helper and flatSchema naming in collectors

* refactor(worker): use undefined over null in queue-schema builders

* docs(worker): converge renderer/collector doc comments to current pi terminology

* refactor(worker): adopt schema.ts cleanInput/stringEnum helpers in collectors

* feat(worker): converge exploit-collector/renderer with vendored; capture and render overview for blocked findings

* refactor(worker): converge session-tools/pipeline/exploitation-checker with vendored

* refactor(worker): converge task-tool usage reporting with vendored onUsage callback

* refactor(worker): converge structured output onto a submitTool executor channel

* docs(worker): expand exploit-renderer docstring to match shannon-oss

* docs(worker): adopt richer vuln-renderer docstring from shannon-oss

* docs(worker): neutralize billing-detection wording for shannon-oss parity

* fix(worker): verify checkpoint hash in the deliverables clone being reset

* fix(worker): fail fast on malformed exploitation queue JSON

* fix(worker): honor retryable flag when classifying exploitation-queue check failures

* fix(worker): fail fast on corrupted session.json in run-scope validation

* feat(worker): propagate Temporal cancellation signal into agent and auth pi sessions

* fix(worker): mark exploit agent complete when exploitation is skipped so resume skips it

* prompts: drop scan description from executive report prompt

* refactor(worker): add createGenericSubmitTool for raw JSON-schema submit tools

* refactor(worker): gate playwright-cli skill to browser agents via skillsOverride (adopt shannon-oss mechanism)

* docs(worker): correct formatLogTime comment to UTC to match toISOString

* refactor(worker): converge queue-schemas with shannon-oss (guarded count, decl order)

* refactor(worker): converge task-tool with shannon-oss (byte-identical; modelRegistry optional)

* fix(worker): use replaceLiteral for all prompt value insertions to prevent $-mangling

* fix(worker): classify agent execution failures by error type instead of hardcoding validation

* fix(worker): cap auth-failure detail at 250 chars to match shannon-oss

* style(worker): apply biome formatting

* refactor(worker): remove per-session task delegation cap from task tool
This commit is contained in:
ezl-keygraph
2026-07-16 14:09:02 +05:30
committed by GitHub
parent b845936aa5
commit d0b0ec3378
47 changed files with 2007 additions and 1334 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/worker/package.json
-1
View File
@@ -32,7 +32,6 @@
"dotenv": "^16.4.5",
"js-yaml": "^4.1.0",
"typebox": "1.1.38",
"zod": "^4.3.6",
"zx": "^8.0.0"
},
"devDependencies": {
apps/worker/prompts/report-executive.txt
-1
View File
@@ -21,7 +21,6 @@ IMPORTANT: You are MODIFYING an existing file, not creating a new one.
<target>
URL: {{WEB_URL}}
{{DESCRIPTION}}
Filesystem:
- {{REPO_PATH}}/ (read only)
apps/worker/prompts/validate-authentication.txt
+1 -5
View File
@@ -27,11 +27,7 @@ After verification confirms login_success, save the authenticated browser sessio
Run this only when login_success is true. Skip it on failure.
</publish_session>
<report_result>
When the login attempt concludes, call the `submit_auth_result` tool to report the outcome.
</report_result>
<critical>
- Submit each field (username, password, captcha, TOTP) exactly once.
- Any rejection = auth error: call `submit_auth_result` with `login_success: false` and stop. Do not retry.
- Any rejection = auth error: return `login_success: false` and stop. Do not retry.
</critical>
apps/worker/prompts/vuln-auth.txt
+6 -5
View File
@@ -195,22 +195,23 @@ For each check you perform from the list above (Transport, Rate Limiting, Sessio
</methodology_and_domain_expertise>
<deliverable_tools>
After completing your `todo_write` tasks and synthesizing findings, emit your specialist deliverable via 3 one-shot tools. Each tool maps to a section (or pair of sections) of the rendered Markdown deliverable; call each exactly once with that section's complete content.
After completing your `todo_write` tasks and synthesizing findings, emit your specialist deliverable via 4 one-shot tools. Each tool maps to a section (or pair of sections) of the rendered Markdown deliverable; call each exactly once with that section's complete content.
**Tool catalog:**
- `set_findings_summary` — Section 1 (Executive Summary key outcome) and Section 2 (Dominant Vulnerability Patterns)
- `set_strategic_intelligence` — Section 3 (Strategic Intelligence for Exploitation, with auth-specific sub-fields: authentication method, session token details, password policy)
- `set_safe_vectors` — Section 4 (Secure by Design: Validated Components)
- `set_blind_spots` — Section 5 (analysis constraints and blind spots)
The harness injects each tool's complete description and per-field guidance into your tool catalog — refer to the tool catalog for what each parameter expects.
**Call semantics:** All 3 tools are one-shot — each may be called exactly once with the section's complete content. Duplicate calls return `"already called"` and are no-ops. There is no incremental/append mode; synthesize each section's full content in working memory before emitting.
**Call semantics:** All 4 tools are one-shot — each may be called exactly once with the section's complete content. Duplicate calls return `"already called"` and are no-ops. There is no incremental/append mode; synthesize each section's full content in working memory before emitting.
**Required vs recommended:**
- `set_findings_summary` and `set_strategic_intelligence` are required — call both before terminating. They produce the load-bearing content the downstream `exploit-auth` agent reads.
- `set_safe_vectors` is recommended. An empty array is acceptable on runs with no validated-secure components, but explicit emission is preferred over skipping.
- `set_safe_vectors` and `set_blind_spots` are recommended. Empty arrays are acceptable on runs with no validated-secure components or no constraint gaps, but explicit emission is preferred over skipping.
**Relationship to the exploitation queue:** The exploitation queue (`auth_exploitation_queue.json`) is produced by calling the `submit_exploitation_queue` tool when your analysis is complete. The 3 tools produce the analysis deliverable Markdown; the structured-output queue is separate and follows the `exploitation_queue_format` schema documented above.
**Relationship to the exploitation queue:** The exploitation queue (`auth_exploitation_queue.json`) is produced by calling the `submit_exploitation_queue` tool when your analysis is complete. The 4 tools produce the analysis deliverable Markdown; the structured-output queue is separate and follows the `exploitation_queue_format` schema documented above.
</deliverable_tools>
@@ -218,7 +219,7 @@ The harness injects each tool's complete description and per-field guidance into
**COMPLETION REQUIREMENTS (ALL must be satisfied):**
1. **Systematic Analysis:** ALL relevant API endpoints and user-facing features identified in the reconnaissance deliverable must be analyzed for AuthN/AuthZ flaws.
2. **Deliverable Emission:** Call the 3 tools (`set_findings_summary`, `set_strategic_intelligence`, `set_safe_vectors`) exactly once each with their respective section content. `set_findings_summary` and `set_strategic_intelligence` are required; `set_safe_vectors` is recommended (an empty array is acceptable but explicit emission is preferred).
2. **Deliverable Emission:** Call the 4 tools (`set_findings_summary`, `set_strategic_intelligence`, `set_safe_vectors`, `set_blind_spots`) exactly once each with their respective section content. `set_findings_summary` and `set_strategic_intelligence` are required; `set_safe_vectors` and `set_blind_spots` are recommended (empty arrays are acceptable but explicit emission is preferred).
**Note:** The exploitation queue is produced by calling the `submit_exploitation_queue` tool when your analysis is complete — separate from the tools above. The analysis deliverable Markdown is rendered by the harness after your session ends from the tool calls.
apps/worker/prompts/vuln-ssrf.txt
+6 -5
View File
@@ -244,22 +244,23 @@ For each check you perform from the list above, you must make a final **verdict*
</methodology_and_domain_expertise>
<deliverable_tools>
After completing your `todo_write` tasks and synthesizing findings, emit your specialist deliverable via 3 one-shot tools. Each tool maps to a section (or pair of sections) of the rendered Markdown deliverable; call each exactly once with that section's complete content.
After completing your `todo_write` tasks and synthesizing findings, emit your specialist deliverable via 4 one-shot tools. Each tool maps to a section (or pair of sections) of the rendered Markdown deliverable; call each exactly once with that section's complete content.
**Tool catalog:**
- `set_findings_summary` — Section 1 (Executive Summary key outcome) and Section 2 (Dominant Vulnerability Patterns)
- `set_strategic_intelligence` — Section 3 (Strategic Intelligence for Exploitation, with SSRF-specific sub-fields: HTTP client library, request architecture, internal services)
- `set_safe_vectors` — Section 4 (Secure by Design: Validated Components)
- `set_blind_spots` — Section 5 (analysis constraints and blind spots)
The harness injects each tool's complete description and per-field guidance into your tool catalog — refer to the tool catalog for what each parameter expects.
**Call semantics:** All 3 tools are one-shot — each may be called exactly once with the section's complete content. Duplicate calls return `"already called"` and are no-ops. There is no incremental/append mode; synthesize each section's full content in working memory before emitting.
**Call semantics:** All 4 tools are one-shot — each may be called exactly once with the section's complete content. Duplicate calls return `"already called"` and are no-ops. There is no incremental/append mode; synthesize each section's full content in working memory before emitting.
**Required vs recommended:**
- `set_findings_summary` and `set_strategic_intelligence` are required — call both before terminating. They produce the load-bearing content the downstream `exploit-ssrf` agent reads.
- `set_safe_vectors` is recommended. An empty array is acceptable on runs with no validated-secure components, but explicit emission is preferred over skipping.
- `set_safe_vectors` and `set_blind_spots` are recommended. Empty arrays are acceptable on runs with no validated-secure components or no constraint gaps, but explicit emission is preferred over skipping.
**Relationship to the exploitation queue:** The exploitation queue (`ssrf_exploitation_queue.json`) is produced by calling the `submit_exploitation_queue` tool when your analysis is complete. The 3 tools produce the analysis deliverable Markdown; the structured-output queue is separate and follows the `exploitation_queue_format` schema documented above.
**Relationship to the exploitation queue:** The exploitation queue (`ssrf_exploitation_queue.json`) is produced by calling the `submit_exploitation_queue` tool when your analysis is complete. The 4 tools produce the analysis deliverable Markdown; the structured-output queue is separate and follows the `exploitation_queue_format` schema documented above.
</deliverable_tools>
@@ -267,7 +268,7 @@ The harness injects each tool's complete description and per-field guidance into
**COMPLETION REQUIREMENTS (ALL must be satisfied):**
1. **Systematic Analysis:** ALL relevant API endpoints and request-making features identified in the reconnaissance deliverable must be analyzed for SSRF vulnerabilities.
2. **Deliverable Emission:** Call the 3 tools (`set_findings_summary`, `set_strategic_intelligence`, `set_safe_vectors`) exactly once each with their respective section content. `set_findings_summary` and `set_strategic_intelligence` are required; `set_safe_vectors` is recommended (an empty array is acceptable but explicit emission is preferred).
2. **Deliverable Emission:** Call the 4 tools (`set_findings_summary`, `set_strategic_intelligence`, `set_safe_vectors`, `set_blind_spots`) exactly once each with their respective section content. `set_findings_summary` and `set_strategic_intelligence` are required; `set_safe_vectors` and `set_blind_spots` are recommended (empty arrays are acceptable but explicit emission is preferred).
**Note:** The exploitation queue is produced by calling the `submit_exploitation_queue` tool when your analysis is complete — separate from the tools above. The analysis deliverable Markdown is rendered by the harness after your session ends from the tool calls.
apps/worker/src/ai/models.ts
+21 -41
View File
@@ -16,18 +16,16 @@
* ANTHROPIC_LARGE_MODEL, which works across all providers (Anthropic, Bedrock,
* custom base URL).
*
* The active provider is chosen from an injected `providerConfig` (the Pro consumer)
* or, in OSS, from the env-var contract the CLI forwards (`CLAUDE_CODE_USE_BEDROCK`,
* `ANTHROPIC_BASE_URL`+`ANTHROPIC_AUTH_TOKEN`, else direct Anthropic). Resolution
* returns a pi `Model` via `ModelRegistry.find`, the `thinkingLevel`, and an
* `AuthStorage` primed with the right credential. Bedrock authenticates from the
* AWS_ env vars via pi-ai.
* The active provider is chosen from the env-var contract the CLI forwards
* (`CLAUDE_CODE_USE_BEDROCK`, `ANTHROPIC_BASE_URL`+`ANTHROPIC_AUTH_TOKEN`, else
* direct Anthropic). Resolution returns a pi `Model` via `ModelRegistry.find`, the
* `thinkingLevel`, and an `AuthStorage` primed with the right credential. Bedrock
* authenticates from the AWS_ env vars via pi-ai.
*/
import type { ThinkingLevel } from '@earendil-works/pi-agent-core';
import type { Api, Model } from '@earendil-works/pi-ai';
import { AuthStorage, type ModelRegistry } from '@earendil-works/pi-coding-agent';
import type { ProviderConfig } from '../types/config.js';
export type ModelTier = 'small' | 'medium' | 'large';
@@ -47,34 +45,20 @@ export interface EffectiveProvider {
}
/**
* Determine the active provider + auth.
*
* An explicit `providerConfig` (injected by the Pro consumer) wins; otherwise we
* fall back to the OSS env-var contract the CLI forwards: `CLAUDE_CODE_USE_BEDROCK`
* → Bedrock; `ANTHROPIC_BASE_URL`+`ANTHROPIC_AUTH_TOKEN` → custom base URL; else
* direct Anthropic (`ANTHROPIC_API_KEY`, or `CLAUDE_CODE_OAUTH_TOKEN`). Bedrock
* authenticates from the AWS_ env vars via pi-ai, so it needs no anthropic token.
* Determine the active provider + auth from the env-var contract the CLI forwards:
* `CLAUDE_CODE_USE_BEDROCK` → Bedrock; `ANTHROPIC_BASE_URL`+`ANTHROPIC_AUTH_TOKEN`
* → custom base URL; else direct Anthropic (`ANTHROPIC_API_KEY`, or
* `CLAUDE_CODE_OAUTH_TOKEN`). Bedrock authenticates from the AWS_ env vars via
* pi-ai, so it needs no anthropic token.
*/
export function resolveEffectiveProvider(apiKey?: string, providerConfig?: ProviderConfig): EffectiveProvider {
const anthropicKey = apiKey ?? providerConfig?.apiKey ?? process.env.ANTHROPIC_API_KEY;
const type = providerConfig?.providerType;
// Bedrock — explicit providerConfig or the env flag.
if (type === 'bedrock' || (!type && process.env.CLAUDE_CODE_USE_BEDROCK === '1')) {
export function resolveEffectiveProvider(): EffectiveProvider {
// Bedrock — env flag.
if (process.env.CLAUDE_CODE_USE_BEDROCK === '1') {
return { providerId: 'amazon-bedrock' };
}
// Custom base URL — explicit providerConfig.
if (type === 'custom_base_url') {
const eff: EffectiveProvider = { providerId: 'anthropic' };
if (providerConfig?.baseUrl) eff.baseUrl = providerConfig.baseUrl;
const token = providerConfig?.authToken ?? anthropicKey;
if (token) eff.anthropicToken = token;
return eff;
}
// Custom base URL — OSS env contract (no providerConfig).
if (!type && process.env.ANTHROPIC_BASE_URL && process.env.ANTHROPIC_AUTH_TOKEN) {
// Custom base URL — env contract.
if (process.env.ANTHROPIC_BASE_URL && process.env.ANTHROPIC_AUTH_TOKEN) {
return {
providerId: 'anthropic',
baseUrl: process.env.ANTHROPIC_BASE_URL,
@@ -82,17 +66,15 @@ export function resolveEffectiveProvider(apiKey?: string, providerConfig?: Provi
};
}
// Direct Anthropic (API key, or — env only — OAuth token).
// Direct Anthropic (API key, or OAuth token).
const eff: EffectiveProvider = { providerId: 'anthropic' };
const token = anthropicKey ?? (type ? undefined : process.env.CLAUDE_CODE_OAUTH_TOKEN);
const token = process.env.ANTHROPIC_API_KEY ?? process.env.CLAUDE_CODE_OAUTH_TOKEN;
if (token) eff.anthropicToken = token;
return eff;
}
/** Resolve a model tier to a concrete model ID (env override → providerConfig → default). */
export function resolveModelId(tier: ModelTier = 'medium', providerConfig?: ProviderConfig): string {
const override = providerConfig?.modelOverrides?.[tier];
if (override) return override;
/** Resolve a model tier to a concrete model ID (env override → default). */
export function resolveModelId(tier: ModelTier = 'medium'): string {
switch (tier) {
case 'small':
return process.env.ANTHROPIC_SMALL_MODEL || DEFAULT_MODELS.small;
@@ -137,11 +119,9 @@ export interface ModelSelection {
export function resolveModelSelection(
registryFactory: (authStorage: AuthStorage) => ModelRegistry,
modelTier: ModelTier,
apiKey?: string,
providerConfig?: ProviderConfig,
): ModelSelection {
const eff = resolveEffectiveProvider(apiKey, providerConfig);
const modelId = resolveModelId(modelTier, providerConfig);
const eff = resolveEffectiveProvider();
const modelId = resolveModelId(modelTier);
const authStorage = AuthStorage.inMemory();
if (eff.providerId === 'anthropic' && eff.anthropicToken) {
apps/worker/src/ai/pi/permission-system.ts
+141
View File
@@ -0,0 +1,141 @@
// 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.
/**
* code_path "avoid" enforcement for the pi harness, delegated to the
* @gotgenes/pi-permission-system extension.
*
* Each `code_path` avoid is translated into the extension's cross-cutting `path`
* deny surface — the strongest gate, blocking file access (read/edit/write/grep/
* find/ls) AND recognized bash file commands (cat/grep/sed/…) on any matching path,
* across every tool and child `task` session, not overridable by a per-tool allow.
*
* `external_directory: allow` keeps the extension from gating the agent's legitimate
* access outside the working directory once it is loaded (the pentest agent shells
* out to tools/paths outside the mounted repo). When there are no avoids the config
* is removed so the executor skips loading the extension entirely.
*/
import fs from 'node:fs';
import { createRequire } from 'node:module';
import path from 'node:path';
import { getAgentDir } from '@earendil-works/pi-coding-agent';
import type { DistributedConfig } from '../../types/config.js';
const PERMISSION_EXTENSION_ID = 'pi-permission-system';
/**
* Translate one avoid value into the extension's flat-wildcard `path` patterns.
*
* The extension's `*` already spans path separators (no `**` globstar), and tool
* paths are compared as absolute. A plain directory value is expanded to cover the
* directory itself and everything under it, in both cwd-relative and prefixed
* (absolute) positions. Glob values fold `**`→`*`; a `dir/*` contents glob also
* denies the directory entry itself.
*/
export function toPathPatterns(value: string): string[] {
// Strip only leading path prefixes ("/", "./", "../"); preserve a dotfile's dot
// (so `.env` stays `.env`, not `env`).
const base = value.replace(/^(?:\.{0,2}\/)+/, '').replace(/\/+$/, '');
if (!base) return [];
if (base.includes('*') || base.includes('?')) {
// The extension's `*` already spans path separators, so fold `**` to `*`.
const flat = base.replace(/\*\*\//g, '*/').replace(/\*\*/g, '*');
const tail = flat.replace(/^(?:\*\/)+/, '');
const patterns = [flat, `*/${tail}`];
// Depth-agnostic catch-all only for a bare-name tail (so `**/*.env` hits a
// root-level `.env`); a structured tail would over-match sibling names.
if (!tail.includes('/')) {
patterns.push(tail.startsWith('*') ? tail : `*${tail}`);
}
// A `dir/*` contents glob should also deny the directory entry itself — the
// contents patterns require a trailing segment and wouldn't match the folder.
if (flat.endsWith('/*')) {
const folder = flat.slice(0, -2);
if (folder && !folder.includes('*')) {
patterns.push(folder, `*/${folder}`);
}
}
return [...new Set(patterns)];
}
return [base, `${base}/*`, `*/${base}`, `*/${base}/*`];
}
interface PermissionSystemConfig {
permission: {
'*': 'allow';
path: Record<string, 'allow' | 'deny'>;
external_directory: 'allow';
};
}
/** Build the extension config that denies every avoid pattern across all tools. */
export function buildPermissionConfig(patterns: readonly string[]): PermissionSystemConfig {
// Default allow first; deny entries are appended so they win (last match wins).
const pathRules: Record<string, 'allow' | 'deny'> = { '*': 'allow' };
for (const pattern of patterns) {
for (const expanded of toPathPatterns(pattern)) {
pathRules[expanded] = 'deny';
}
}
return {
permission: {
'*': 'allow',
path: pathRules,
external_directory: 'allow',
},
};
}
/** Path to the extension's global config under the agent directory. */
export function permissionSystemConfigPath(agentDir: string): string {
return path.join(agentDir, 'extensions', PERMISSION_EXTENSION_ID, 'config.json');
}
/** True when a pi-permission-system config has been written (avoid rules exist). */
export function permissionSystemConfigExists(agentDir: string): boolean {
return fs.existsSync(permissionSystemConfigPath(agentDir));
}
/**
* Sync the distributed config's `code_path` avoids into the extension's global
* config (`<agentDir>/extensions/pi-permission-system/config.json`). When there
* are no avoids the config is removed so the executor skips loading the extension.
*
* Global (not project) config is used deliberately: it loads synchronously at
* extension init without depending on a session_start/ctx, it keeps the config
* out of the scanned repo, and it is idempotent across the agents of one run.
*/
export function syncPermissionSystemConfig(config: DistributedConfig | null): void {
const configPath = permissionSystemConfigPath(getAgentDir());
const avoidRules = (config?.avoid ?? []).filter((r) => r.type === 'code_path');
if (avoidRules.length === 0) {
fs.rmSync(configPath, { force: true });
return;
}
// Single-repo (fixed mount): patterns are the raw avoid values.
const patterns = avoidRules.map((r) => r.value);
fs.mkdirSync(path.dirname(configPath), { recursive: true });
fs.writeFileSync(configPath, JSON.stringify(buildPermissionConfig(patterns), null, 2));
}
/**
* Absolute path to the installed @gotgenes/pi-permission-system package directory,
* suitable for `DefaultResourceLoader`'s `additionalExtensionPaths`. The loader
* reads the package's `pi.extensions` manifest and loads the extension itself.
*
* The package's `.` export points at its service module, so we resolve that and
* walk up to the package root. Throws if the package is not resolvable.
*/
export function permissionSystemPackageDir(): string {
const require = createRequire(import.meta.url);
const servicePath = require.resolve('@gotgenes/pi-permission-system');
return path.resolve(path.dirname(servicePath), '..');
}
apps/worker/src/ai/pi-executor.ts → apps/worker/src/ai/pi/pi-executor.ts
+76 -47
View File
@@ -6,7 +6,7 @@
// Production agent execution on the pi harness, with git checkpoints and audit logging.
import { createRequire } from 'node:module';
import os from 'node:os';
import type { AgentMessage } from '@earendil-works/pi-agent-core';
import {
type AgentSessionEvent,
@@ -17,30 +17,34 @@ import {
type ResourceLoader,
SessionManager,
SettingsManager,
type Skill,
type ToolDefinition,
} from '@earendil-works/pi-coding-agent';
import { fs, path } from 'zx';
import type { AuditSession } from '../audit/index.js';
import { BASH_TIMEOUT_EXTENSION_DIR, deliverablesDir, PLAYWRIGHT_SKILL_DIR } from '../paths.js';
import { isRetryableError, PentestError } from '../services/error-handling.js';
import { AGENT_VALIDATORS } from '../session-manager.js';
import type { ActivityLogger } from '../types/activity-logger.js';
import { ErrorCode } from '../types/errors.js';
import { isSpendingCapBehavior, matchesBillingTextPattern } from '../utils/billing-detection.js';
import { formatTimestamp } from '../utils/formatting.js';
import { Timer } from '../utils/metrics.js';
import { createAuditLogger } from './audit-logger.js';
import { type ModelTier, resolveModelSelection } from './models.js';
import type { AuditSession } from '../../audit/index.js';
import { BASH_TIMEOUT_EXTENSION_DIR, deliverablesDir } from '../../paths.js';
import { isRetryableError, PentestError } from '../../services/error-handling.js';
import { AGENT_VALIDATORS } from '../../session-manager.js';
import type { ActivityLogger } from '../../types/activity-logger.js';
import { ErrorCode } from '../../types/errors.js';
import { isSpendingCapBehavior, matchesBillingTextPattern } from '../../utils/billing-detection.js';
import { isBrowserAgent } from '../../utils/browser-agents.js';
import { formatTimestamp } from '../../utils/formatting.js';
import { Timer } from '../../utils/metrics.js';
import { createAuditLogger } from '../audit-logger.js';
import { type ModelTier, resolveModelSelection } from '../models.js';
import {
detectExecutionContext,
formatAssistantOutput,
formatCompletionMessage,
formatErrorOutput,
formatToolCall,
} from './output-formatters.js';
import { createProgressManager } from './progress-manager.js';
import { permissionConfigPath } from './settings-writer.js';
import { createGlobTool, createTaskTool, createTodoWriteTool } from './tools.js';
} from '../output-formatters.js';
import { createProgressManager } from '../progress-manager.js';
import type { CapturedSubmitTool } from '../submit-tool.js';
import { permissionSystemConfigExists, permissionSystemPackageDir } from './permission-system.js';
import { createGlobTool, createTodoWriteTool } from './session-tools.js';
import { createTaskTool } from './task-tool.js';
declare global {
var SHANNON_DISABLE_LOADER: boolean | undefined;
@@ -49,40 +53,53 @@ declare global {
/** Built-in pi tools enabled for every agent (custom tool names are appended). */
const BUILTIN_TOOLS = ['read', 'bash', 'edit', 'write', 'grep', 'find', 'ls'];
const requireFromHere = createRequire(import.meta.url);
let cachedExtensionDir: string | null | undefined;
/** Resolve the installed @gotgenes/pi-permission-system package dir, or null. */
function permissionExtensionDir(): string | null {
if (cachedExtensionDir !== undefined) return cachedExtensionDir;
try {
const entry = requireFromHere.resolve('@gotgenes/pi-permission-system');
cachedExtensionDir = path.dirname(path.dirname(entry));
} catch {
cachedExtensionDir = null;
}
return cachedExtensionDir;
/** Build the playwright-cli Skill object injected for browser-using agents. */
function buildPlaywrightSkill(): Skill {
const filePath =
process.env.PLAYWRIGHT_CLI_SKILL_PATH ?? path.join(os.homedir(), '.claude/skills/playwright-cli/SKILL.md');
const baseDir = path.dirname(filePath);
return {
name: 'playwright-cli',
description:
'Drive a real browser via the playwright-cli binary. Use for any task that navigates, clicks, ' +
'fills forms, takes screenshots, or reads live pages.',
filePath,
baseDir,
sourceInfo: { path: filePath, source: 'custom', scope: 'user', origin: 'top-level', baseDir },
disableModelInvocation: false,
};
}
async function buildResourceLoader(cwd: string, logger: ActivityLogger): Promise<ResourceLoader> {
async function buildResourceLoader(
cwd: string,
logger: ActivityLogger,
agentName: string | null,
): Promise<ResourceLoader> {
// Always enforce bounded bash timeouts so an unbounded command cannot hang the agent.
const additionalExtensionPaths: string[] = [BASH_TIMEOUT_EXTENSION_DIR];
if (fs.existsSync(permissionConfigPath())) {
const extDir = permissionExtensionDir();
if (extDir) {
additionalExtensionPaths.push(extDir);
} else {
if (permissionSystemConfigExists(getAgentDir())) {
try {
additionalExtensionPaths.push(permissionSystemPackageDir());
} catch {
logger.warn(
'code_path deny config present but @gotgenes/pi-permission-system not resolvable — skipping enforcement',
);
}
}
// Only browser-driving agents get the playwright-cli skill; the rest run with no skills.
const loader = new DefaultResourceLoader({
cwd,
agentDir: getAgentDir(),
additionalSkillPaths: [PLAYWRIGHT_SKILL_DIR],
...(additionalExtensionPaths.length > 0 && { additionalExtensionPaths }),
...(isBrowserAgent(agentName)
? {
skillsOverride: (base) => ({
skills: [buildPlaywrightSkill()],
diagnostics: base.diagnostics,
}),
}
: { noSkills: true }),
});
await loader.reload();
return loader;
@@ -202,18 +219,20 @@ export async function runPiPrompt(
sourceDir: string,
context: string = '',
description: string = 'Agent analysis',
_agentName: string | null = null,
agentName: string | null = null,
auditSession: AuditSession | null = null,
logger: ActivityLogger,
modelTier: ModelTier = 'medium',
callerTools?: ToolDefinition[],
apiKey?: string,
deliverablesSubdir?: string,
providerConfig?: import('../types/config.js').ProviderConfig,
cancellationSignal?: AbortSignal,
submitTool?: CapturedSubmitTool,
): Promise<PiPromptResult> {
// 1. Initialize timing and prompt
// 1. Initialize timing and prompt. A submit tool appends its directive so the
// instruction to call it lives with the tool, not in every prompt file.
const timer = new Timer(`agent-${description.toLowerCase().replace(/\s+/g, '-')}`);
const fullPrompt = context ? `${context}\n\n${prompt}` : prompt;
const basePrompt = context ? `${context}\n\n${prompt}` : prompt;
const fullPrompt = submitTool?.directive ? basePrompt + submitTool.directive : basePrompt;
// 2. Set up progress and audit infrastructure
const execContext = detectExecutionContext(description);
@@ -232,27 +251,32 @@ export async function runPiPrompt(
? path.join(sourceDir, path.dirname(deliverablesSubdir), '.playwright-cli')
: path.join(sourceDir, '.shannon', '.playwright-cli');
if (deliverablesSubdir) process.env.SHANNON_DELIVERABLES_SUBDIR = deliverablesSubdir;
if (apiKey) process.env.ANTHROPIC_API_KEY = apiKey;
// 4. Resolve model + auth, then assemble the tool set (universal task/todo tools
// plus any caller-supplied collector/submit tools).
const selection = resolveModelSelection((auth) => ModelRegistry.create(auth), modelTier, apiKey, providerConfig);
const resourceLoader = await buildResourceLoader(sourceDir, logger);
// Accumulates cost from in-process `task` child sessions so the parent's reported
const selection = resolveModelSelection((auth) => ModelRegistry.create(auth), modelTier);
const resourceLoader = await buildResourceLoader(sourceDir, logger, agentName);
// Accumulates usage from in-process `task` child sessions so the parent's reported
// cost includes sub-agent spend (their getSessionStats is separate from ours).
const childUsage = { cost: 0 };
const childUsage = { cost: 0, inputTokens: 0, outputTokens: 0 };
const customTools: ToolDefinition[] = [
createTaskTool({
model: selection.model,
thinkingLevel: selection.thinkingLevel,
authStorage: selection.authStorage,
cwd: sourceDir,
childUsage,
onUsage: (usage) => {
childUsage.cost += usage.cost;
childUsage.inputTokens += usage.inputTokens;
childUsage.outputTokens += usage.outputTokens;
},
resourceLoader,
...(cancellationSignal && { cancellationSignal }),
}),
createTodoWriteTool(auditLogger),
createGlobTool(sourceDir),
...(callerTools ?? []),
...(submitTool ? [submitTool.tool] : []),
];
// pi's `tools` allowlist gates custom tools too — list every custom name.
const tools = [...BUILTIN_TOOLS, ...customTools.map((t) => t.name)];
@@ -357,6 +381,10 @@ export async function runPiPrompt(
const duration = timer.stop();
progress.finish(formatCompletionMessage(execContext, description, turnCount, duration));
// Capture the submit tool's structured payload so callers read it off the
// result instead of holding a reference to the tool.
const structuredOutput = submitTool?.getCaptured();
return {
result,
success: true,
@@ -366,6 +394,7 @@ export async function runPiPrompt(
model: selection.model.id,
partialCost: totalCost,
apiErrorDetected,
...(structuredOutput !== undefined && { structuredOutput }),
};
} catch (error) {
// 10. Handle errors — log, write error file, return failure
apps/worker/src/ai/pi/session-tools.ts
+116
View File
@@ -0,0 +1,116 @@
// 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.
/**
* Per-session custom tools registered for every agent: `todo_write` and `glob`.
*
* These replace harness built-ins that pi does not ship. `todo_write` is a
* full-state-replace planning scratchpad mirrored to the workflow log; `glob` is
* fast-glob file matching (pi has no `Glob` built-in).
*/
import { defineTool, type ToolDefinition } from '@earendil-works/pi-coding-agent';
import { Type } from 'typebox';
import { fs, glob, path } from 'zx';
import type { AuditLogger } from '../audit-logger.js';
export interface TodoItem {
content: string;
status: 'pending' | 'in_progress' | 'completed';
activeForm: string;
}
function renderTodos(todos: readonly TodoItem[]): string {
const mark = (status: TodoItem['status']): string => {
if (status === 'completed') return 'x';
if (status === 'in_progress') return '~';
return ' ';
};
return todos.map((todo) => `[${mark(todo.status)}] ${todo.content}`).join(' ');
}
export function createTodoWriteTool(auditLogger: AuditLogger): ToolDefinition {
let current: TodoItem[] = [];
return defineTool({
name: 'todo_write',
label: 'Todo Write',
description:
'Use this tool to create and manage a structured task list for your current session. ' +
'Pass the complete todo list on every call; it replaces the stored list entirely. Each ' +
'todo has a status of pending, in_progress, or completed.',
promptSnippet: 'todo_write: create and manage a structured task list',
parameters: Type.Object({
todos: Type.Array(
Type.Object({
content: Type.String({ description: 'Imperative task description, e.g. "Map SSRF sinks".' }),
status: Type.Union([Type.Literal('pending'), Type.Literal('in_progress'), Type.Literal('completed')]),
activeForm: Type.String({ description: 'Present-continuous form, e.g. "Mapping SSRF sinks".' }),
}),
),
}),
async execute(_toolCallId, params) {
current = params.todos as TodoItem[];
const completed = current.filter((todo) => todo.status === 'completed').length;
await auditLogger.logNote('todo', renderTodos(current));
return {
content: [
{
type: 'text' as const,
text: `Todos updated (${current.length} items, ${completed} completed).`,
},
],
details: undefined,
};
},
});
}
export function createGlobTool(cwd: string): ToolDefinition {
return defineTool({
name: 'glob',
label: 'Glob',
description:
'Fast file pattern matching. Supports glob patterns like "**/*.ts" or "src/**/*.{js,ts}". ' +
'Returns matching file paths sorted by modification time, most recent first.',
promptSnippet: 'glob: find files by name pattern',
parameters: Type.Object({
pattern: Type.String({ description: 'The glob pattern to match files against.' }),
path: Type.Optional(Type.String({ description: 'Directory to search in. Omit for the repository root.' })),
}),
async execute(_toolCallId, params) {
const searchRoot = params.path ? path.resolve(cwd, params.path) : cwd;
const matches = await glob.globby(params.pattern, {
cwd: searchRoot,
absolute: true,
dot: true,
onlyFiles: true,
followSymbolicLinks: false,
});
if (matches.length === 0) {
return { content: [{ type: 'text' as const, text: 'No files found' }], details: undefined };
}
const withMtime = await Promise.all(
matches.map(async (file) => {
try {
return { file, mtime: (await fs.stat(file)).mtimeMs };
} catch {
return { file, mtime: 0 };
}
}),
);
withMtime.sort((a, b) => b.mtime - a.mtime);
return {
content: [{ type: 'text' as const, text: withMtime.map((match) => match.file).join('\n') }],
details: undefined,
};
},
});
}
apps/worker/src/ai/pi/task-tool.ts
+156
View File
@@ -0,0 +1,156 @@
// 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.
/**
* Generic `task` tool — pi.dev ships no built-in Task tool, so this supplies the
* Task-delegation surface Shannon's prompts require.
*
* Shannon's prompts mandate Task delegation (recon source tracer; the vuln
* agents delegate *every* code review; the exploit agents delegate automation),
* so this tool is required for parity, not optional. It spawns a nested pi
* session with the parent's resolved model object (never a tier string — that
* would route sub-agents through hardcoded IDs and leak billing), the parent's
* resource loader, and a fixed child tool surface.
*/
import type { ThinkingLevel } from '@earendil-works/pi-agent-core';
import { type AssistantMessage, type Model, Type } from '@earendil-works/pi-ai';
import {
type AuthStorage,
createAgentSession,
defineTool,
getAgentDir,
type ModelRegistry,
type ResourceLoader,
SessionManager,
SettingsManager,
type ToolDefinition,
} from '@earendil-works/pi-coding-agent';
export interface TaskToolContext {
cwd: string;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
model: Model<any>;
thinkingLevel?: ThinkingLevel;
authStorage: AuthStorage;
/** Explicit model registry for sub-session resolution. Omit to inherit the parent's default. */
modelRegistry?: ModelRegistry;
resourceLoader: ResourceLoader;
cancellationSignal?: AbortSignal | undefined;
/**
* Reports the cost/tokens of each spawned sub-session back to the caller.
* Sub-agents run in their own pi sessions that the parent has no reference to,
* so without this their spend (the bulk of a whitebox run, since Shannon
* prompts delegate the heavy work) is invisible to billing.
*/
onUsage?: (usage: { cost: number; inputTokens: number; outputTokens: number }) => void;
}
const CHILD_TOOLS = ['read', 'grep', 'find', 'ls', 'write', 'bash'];
function textResult(text: string) {
return { content: [{ type: 'text' as const, text }], details: undefined };
}
export function createTaskTool(config: TaskToolContext): ToolDefinition {
const taskTool: ToolDefinition = defineTool({
name: 'task',
label: 'Task',
description:
'Delegate a focused task to a sub-agent that runs independently with its own tools and returns ' +
'the result. Use this to break complex work into smaller, parallelizable sub-tasks.',
executionMode: 'parallel',
promptSnippet: 'task - Delegate a focused task to a sub-agent with read, grep, find, ls, write, and bash.',
promptGuidelines: [
'Use the task tool to delegate focused work: code review, reconnaissance, automation scripting, validation.',
'Pass all necessary context in the "prompt" parameter — the sub-agent cannot see your conversation history.',
'The sub-agent can use read, grep, find, ls, write, and bash, but cannot call task or custom collector tools.',
'You can launch multiple task tool calls in a single message to run sub-tasks in parallel.',
],
parameters: Type.Object({
prompt: Type.String({
description: 'The task for the sub-agent to perform. Include all necessary context.',
}),
description: Type.Optional(Type.String({ description: 'A short (3-5 word) description of the task.' })),
}),
async execute(_toolCallId, params) {
const agentDir = getAgentDir();
const { session: subSession } = await createAgentSession({
cwd: config.cwd,
agentDir,
resourceLoader: config.resourceLoader,
model: config.model,
...(config.thinkingLevel && { thinkingLevel: config.thinkingLevel }),
tools: CHILD_TOOLS,
authStorage: config.authStorage,
...(config.modelRegistry && { modelRegistry: config.modelRegistry }),
sessionManager: SessionManager.inMemory(config.cwd),
settingsManager: SettingsManager.inMemory({
retry: { enabled: false },
compaction: { enabled: true },
}),
});
const abortChildSession = (): void => {
void subSession.abort().catch(() => {
// Parent logger is not available inside the tool; dispose still tears
// down the session if abort itself rejects.
});
};
const onCancellation = (): void => abortChildSession();
if (config.cancellationSignal?.aborted) {
abortChildSession();
} else {
config.cancellationSignal?.addEventListener('abort', onCancellation, { once: true });
}
let resultText = '';
let subCost = 0;
let subInputTokens = 0;
let subOutputTokens = 0;
subSession.subscribe((event) => {
if (event.type === 'turn_end') {
const msg = event.message as AssistantMessage | undefined;
for (const block of msg?.content ?? []) {
if (block.type === 'text' && block.text) {
resultText += (resultText ? '\n' : '') + block.text;
}
}
if (msg?.usage?.cost?.total != null) subCost += msg.usage.cost.total;
subInputTokens += msg?.usage?.input ?? 0;
subOutputTokens += msg?.usage?.output ?? 0;
}
});
let swallowedError: string | undefined;
try {
try {
await subSession.prompt(params.prompt);
} catch (err) {
const errorMsg = err instanceof Error ? err.message : String(err);
resultText += `\n[Sub-agent error: ${errorMsg}]`;
}
swallowedError = subSession.state.errorMessage;
// Read stats before dispose; reconcile cost the same way the parent does.
const subStats = subSession.getSessionStats();
if (subStats.cost > subCost) subCost = subStats.cost;
config.onUsage?.({ cost: subCost, inputTokens: subInputTokens, outputTokens: subOutputTokens });
} finally {
config.cancellationSignal?.removeEventListener('abort', onCancellation);
subSession.dispose();
}
if (swallowedError && !resultText.includes(swallowedError)) {
resultText += `\n[Sub-agent error: ${swallowedError}]`;
}
return textResult(resultText || '[Sub-agent produced no output]');
},
});
return taskTool;
}
apps/worker/src/ai/queue-schemas.ts
+65 -56
View File
@@ -7,19 +7,21 @@
/**
* TypeBox schemas + submit-tool factory for vulnerability exploitation queues.
*
* pi has no JSON-schema output format, so each vuln agent's structured queue is
* captured via a `submit_exploitation_queue` custom tool whose parameters mirror
* the per-class schema below. The captured payload is written to
* `<class>_exploitation_queue.json` by the caller (agent-execution).
* pi captures each vuln agent's structured queue via a `submit_exploitation_queue`
* custom tool whose parameters mirror the per-class schema below. Entry types are
* derived from the same schemas and consumed by the findings renderer.
*/
import { defineTool, type ToolDefinition } from '@earendil-works/pi-coding-agent';
import { defineTool } from '@earendil-works/pi-coding-agent';
import { type Static, type TObject, Type } from 'typebox';
import type { AgentName } from '../types/agents.js';
import type { CapturedSubmitTool } from './submit-tool.js';
const ANALYSIS_NOTES_DESCRIPTION = 'Plain context for defenders (caveats, scope, what is at risk). Not attack steps.';
const optStr = (description?: string) => Type.Optional(Type.String(description ? { description } : {}));
function optStr(description?: string) {
return Type.Optional(Type.String(description === undefined ? {} : { description }));
}
/** Base fields shared by every queue entry. `notes` gains guidance in analysis mode. */
function baseFields(exploit: boolean) {
@@ -84,6 +86,20 @@ const authzFields = {
minimal_witness: optStr(),
};
// === Per-entry schemas (single vulnerability). Entry types derive from these. ===
const injectionEntry = () => Type.Object({ ...baseFields(true), ...injectionFields });
const xssEntry = () => Type.Object({ ...baseFields(true), ...xssFields });
const authEntry = () => Type.Object({ ...baseFields(true), ...authFields });
const ssrfEntry = () => Type.Object({ ...baseFields(true), ...ssrfFields });
const authzEntry = () => Type.Object({ ...baseFields(true), ...authzFields });
export type InjectionFinding = Static<ReturnType<typeof injectionEntry>>;
export type XssFinding = Static<ReturnType<typeof xssEntry>>;
export type AuthFinding = Static<ReturnType<typeof authEntry>>;
export type SsrfFinding = Static<ReturnType<typeof ssrfEntry>>;
export type AuthzFinding = Static<ReturnType<typeof authzEntry>>;
const PER_TYPE_FIELDS: Partial<Record<AgentName, Record<string, ReturnType<typeof optStr>>>> = {
'injection-vuln': injectionFields,
'xss-vuln': xssFields,
@@ -92,28 +108,6 @@ const PER_TYPE_FIELDS: Partial<Record<AgentName, Record<string, ReturnType<typeo
'authz-vuln': authzFields,
};
/** Build the `{ vulnerabilities: [...] }` queue schema for an agent + mode. */
function queueSchema(agentName: AgentName, exploit: boolean): TObject | null {
const extra = PER_TYPE_FIELDS[agentName];
if (!extra) return null;
return Type.Object({
vulnerabilities: Type.Array(Type.Object({ ...baseFields(exploit), ...extra })),
});
}
// === Inferred entry types (consumed by renderers) ===
export type InjectionFinding = Static<ReturnType<typeof injectionEntry>>;
export type XssFinding = Static<ReturnType<typeof xssEntry>>;
export type AuthFinding = Static<ReturnType<typeof authEntry>>;
export type SsrfFinding = Static<ReturnType<typeof ssrfEntry>>;
export type AuthzFinding = Static<ReturnType<typeof authzEntry>>;
const injectionEntry = () => Type.Object({ ...baseFields(true), ...injectionFields });
const xssEntry = () => Type.Object({ ...baseFields(true), ...xssFields });
const authEntry = () => Type.Object({ ...baseFields(true), ...authFields });
const ssrfEntry = () => Type.Object({ ...baseFields(true), ...ssrfFields });
const authzEntry = () => Type.Object({ ...baseFields(true), ...authzFields });
const VULN_AGENT_QUEUE_FILENAMES: Partial<Record<AgentName, string>> = {
'injection-vuln': 'injection_exploitation_queue.json',
'xss-vuln': 'xss_exploitation_queue.json',
@@ -122,38 +116,53 @@ const VULN_AGENT_QUEUE_FILENAMES: Partial<Record<AgentName, string>> = {
'authz-vuln': 'authz_exploitation_queue.json',
};
/** Build the TypeBox submit-tool parameters for a vuln agent, or undefined for non-vuln agents. */
function queueSchema(agentName: AgentName, exploit: boolean): TObject | undefined {
const extra = PER_TYPE_FIELDS[agentName];
if (!extra) return undefined;
return Type.Object({
vulnerabilities: Type.Array(Type.Object({ ...baseFields(exploit), ...extra })),
});
}
/** Returns the queue filename for a vuln agent, or undefined for non-vuln agents. */
export function getQueueFilename(agentName: AgentName): string | undefined {
return VULN_AGENT_QUEUE_FILENAMES[agentName];
}
export interface QueueSubmitTool {
tool: ToolDefinition;
getCaptured: () => unknown;
}
/**
* Build the `submit_exploitation_queue` tool for a vuln agent, or null for
* non-vuln agents. The agent calls it once with the full findings list; the
* captured payload is the structured queue.
*/
export function createQueueSubmitTool(agentName: AgentName, exploit: boolean): QueueSubmitTool | null {
/** Build the pi submit tool that captures the exploitation queue for vuln agents. */
export function createQueueSubmitTool(agentName: AgentName, exploit = true): CapturedSubmitTool | undefined {
const schema = queueSchema(agentName, exploit);
if (!schema) return null;
let captured: unknown;
const tool = defineTool({
name: 'submit_exploitation_queue',
label: 'Submit Exploitation Queue',
description:
'Submit the final structured list of analyzed vulnerabilities for this class. Call exactly once when ' +
'analysis is complete, with every finding included.',
promptSnippet: 'submit_exploitation_queue: record the final structured findings list (call once)',
parameters: schema,
execute: async (_toolCallId, params) => {
captured = params;
const count = (params as { vulnerabilities?: unknown[] }).vulnerabilities?.length ?? 0;
return { content: [{ type: 'text' as const, text: `Recorded ${count} findings.` }], details: {} };
},
});
return { tool, getCaptured: () => captured };
if (!schema) return undefined;
let captured: unknown | undefined;
return {
tool: defineTool({
name: 'submit_exploitation_queue',
label: 'Submit Exploitation Queue',
description:
'Submit the final structured list of analyzed vulnerabilities for this class. Call exactly once when analysis is complete.',
promptSnippet: 'submit_exploitation_queue: record the final structured findings list (call once)',
promptGuidelines: [
'You MUST call submit_exploitation_queue exactly once as your final action.',
'Include every analyzed finding in the vulnerabilities array.',
],
parameters: schema,
async execute(_toolCallId, params) {
captured = params;
const count = Array.isArray((params as { vulnerabilities?: unknown }).vulnerabilities)
? (params as { vulnerabilities: unknown[] }).vulnerabilities.length
: 0;
return {
content: [{ type: 'text' as const, text: `Recorded ${count} findings.` }],
details: params,
terminate: true,
};
},
}),
getCaptured: () => captured,
directive:
'\n\nYou MUST call the submit_exploitation_queue tool exactly once as your final action ' +
'to deliver your structured exploitation queue. Do not output JSON as text. Fill every required parameter.',
};
}
apps/worker/src/ai/settings-writer.ts
-75
View File
@@ -1,75 +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.
/**
* Writes the @gotgenes/pi-permission-system global config from `code_path` avoid
* patterns. The executor loads the extension (see pi-executor) and pi enforces
* these path denies at the tool layer for every agent. Written to the global config
* dir under `agentDir` — the project-scoped path is gated behind project trust,
* which our headless runs do not grant; the global path is not.
*/
import { getAgentDir } from '@earendil-works/pi-coding-agent';
import { fs, path } from 'zx';
import type { DistributedConfig } from '../types/config.js';
/** Absolute path to the pi-permission-system global config.json. */
export function permissionConfigPath(): string {
return path.join(getAgentDir(), 'extensions', 'pi-permission-system', 'config.json');
}
/**
* Write (or remove) the pi-permission-system config derived from `code_path`
* avoid patterns.
*
* Each avoid maps to a cross-cutting `path` deny — the strongest surface, blocking
* the path across every tool and bash command, and not overridable by a per-tool
* allow. `"*": "allow"` keeps everything else permitted so the extension does not
* fall back to its default `ask` (which would block all access headlessly). When
* there are no avoids the config is removed, so the executor skips loading the
* extension entirely.
*/
export async function writeCodePathPermissionConfig(config: DistributedConfig | null): Promise<void> {
const avoidPatterns = (config?.avoid ?? []).filter((r) => r.type === 'code_path').map((r) => r.value);
const configPath = permissionConfigPath();
if (avoidPatterns.length === 0) {
await fs.remove(configPath);
return;
}
// pi's matcher (wildcard-matcher.ts) has NO `**` globstar — it splits on each `*`
// and joins with `.*`, and a single `*` already matches any chars incl. `/`. Tool
// paths are compared as absolute (path-utils resolves them against cwd), so we
// collapse `**`→`*` and add a `*/`-prefixed variant that matches the path under
// any repo prefix. (A bare pattern never matches an absolute path.)
const pathDeny: Record<string, 'allow' | 'deny'> = { '*': 'allow' };
for (const pattern of avoidPatterns) {
const clean = pattern.replace(/^[./]+/, '').replace(/\*\*/g, '*');
// Deny the contents (under any repo prefix and as written)...
pathDeny[`*/${clean}`] = 'deny';
pathDeny[clean] = 'deny';
// ...and the folder path itself, so the directory entry is denied too — the
// contents patterns (…/*) require a trailing segment and wouldn't match it.
if (clean.endsWith('/*')) {
const folder = clean.slice(0, -2);
if (folder) {
pathDeny[`*/${folder}`] = 'deny';
pathDeny[folder] = 'deny';
}
}
}
const permissionConfig = {
permission: {
'*': 'allow',
path: pathDeny,
},
};
await fs.ensureDir(path.dirname(configPath));
await fs.writeJson(configPath, permissionConfig, { spaces: 2 });
}
apps/worker/src/ai/submit-tool.ts
+60
View File
@@ -0,0 +1,60 @@
// Copyright (C) 2025 Keygraph, Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License version 3
// as published by the Free Software Foundation.
import { defineTool, type ToolDefinition } from '@earendil-works/pi-coding-agent';
import { Type } from 'typebox';
/**
* A pi custom submit tool plus the captured payload it records.
*
* pi ships no JSON-schema output format, so an agent that must return structured
* data does so by calling a purpose-built TypeBox tool. This bundles that tool
* with its capture accessor and the directive that instructs the model to call
* it. The executor owns the wiring — it registers the tool, appends the
* directive to the prompt, and reads `getCaptured()` back as `structuredOutput`
* — so callers never assemble it by hand.
*/
export interface CapturedSubmitTool {
readonly tool: ToolDefinition;
readonly getCaptured: () => unknown | undefined;
readonly directive?: string;
}
/**
* Build a `submit_result` tool from a raw JSON Schema, for agents whose result
* shape is not one of the built-in per-agent schemas (e.g. an out-of-tree agent
* with its own verdict schema). pi validates the tool call against `schema`
* before `execute()` runs, so a captured payload is already schema-valid — no
* separate validation pass is needed.
*/
export function createGenericSubmitTool(schema: Record<string, unknown>): CapturedSubmitTool {
let captured: unknown | undefined;
return {
tool: defineTool({
name: 'submit_result',
label: 'Submit Result',
description: 'Return your final structured answer. Call exactly once as your last action.',
promptSnippet: 'submit_result: deliver your structured answer (call once)',
promptGuidelines: [
'You MUST call submit_result exactly once as your final action.',
'Fill every required parameter. Do not output JSON as text.',
],
parameters: Type.Unsafe(schema),
async execute(_toolCallId, params) {
captured = params;
return {
content: [{ type: 'text' as const, text: 'Result submitted.' }],
details: params,
terminate: true,
};
},
}),
getCaptured: () => captured,
directive:
'\n\nYou MUST call the submit_result tool exactly once as your final action ' +
'to deliver your structured answer. Do not output JSON as text. Fill every required parameter.',
};
}
apps/worker/src/ai/tools.ts
-205
View File
@@ -1,205 +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.
/**
* Universal custom tools registered for every agent: `task`, `todo_write`, and `glob`.
*
* These replace harness built-ins that pi does not ship. `task` delegates a focused
* sub-task to an in-process child session (the Task sub-agent replacement);
* `todo_write` is a full-state-replace planning scratchpad mirrored to the workflow
* log; `glob` is fast-glob file matching (pi has no `Glob` built-in).
*/
import type { ThinkingLevel } from '@earendil-works/pi-agent-core';
import type { Api, Model } from '@earendil-works/pi-ai';
import {
type AuthStorage,
createAgentSession,
defineTool,
type ResourceLoader,
SessionManager,
SettingsManager,
type ToolDefinition,
} from '@earendil-works/pi-coding-agent';
import { Type } from 'typebox';
import { fs, glob, path } from 'zx';
import type { AuditLogger } from './audit-logger.js';
/** Tool surface for child sessions: read/search plus `write`+`bash` to author and run scripts. */
const CHILD_TOOLS = ['read', 'grep', 'find', 'ls', 'write', 'bash'];
export interface TaskToolContext {
model: Model<Api>;
thinkingLevel: ThinkingLevel;
authStorage: AuthStorage;
cwd: string;
/** When set, child sessions inherit the code_path deny policy. */
resourceLoader?: ResourceLoader;
/**
* Mutable accumulator: each child (sub-agent) session's cost is added here so the
* parent executor can include sub-agent spend in its reported cost. Child sessions
* keep their own `getSessionStats`, separate from the parent's.
*/
childUsage?: { cost: number };
}
/**
* The `task` tool — launch a new agent to handle a multi-step task autonomously.
*
* Spawns an in-process child session, drives it to completion, and returns its
* final text. Marked `parallel` for one-turn fan-out. Children get no `task` of
* their own — delegation is one level.
*/
export function createTaskTool(ctx: TaskToolContext): ToolDefinition {
return defineTool({
name: 'task',
label: 'Task',
description:
'Launch a new agent to handle complex, multi-step tasks autonomously. The agent runs on its own and ' +
'its final report is returned to you as the tool result (it is not shown to the user). Each invocation ' +
'is stateless — you cannot send follow-up messages, so give a complete, detailed instruction in a single ' +
'prompt and specify exactly what information the agent should return. Launch multiple agents concurrently ' +
'by issuing multiple task calls in a single message.',
promptSnippet: 'task: launch a new agent to handle a multi-step task',
executionMode: 'parallel',
parameters: Type.Object({
description: Type.Optional(Type.String({ description: 'Short (3-5 word) label for the delegated sub-task.' })),
prompt: Type.String({ description: 'The full instruction for the sub-agent.' }),
}),
execute: async (_toolCallId, params) => {
const { session: child } = await createAgentSession({
cwd: ctx.cwd,
model: ctx.model,
thinkingLevel: ctx.thinkingLevel,
tools: CHILD_TOOLS,
authStorage: ctx.authStorage,
sessionManager: SessionManager.inMemory(),
settingsManager: SettingsManager.inMemory({
retry: { enabled: false },
compaction: { enabled: true },
}),
...(ctx.resourceLoader && { resourceLoader: ctx.resourceLoader }),
});
try {
await child.prompt(params.prompt);
const text = child.getLastAssistantText() ?? '(sub-agent produced no output)';
return { content: [{ type: 'text' as const, text }], details: {} };
} finally {
// Roll the child's cost up to the parent before disposing (best-effort, and
// captured in `finally` so a failed child's partial spend still counts).
if (ctx.childUsage) {
try {
ctx.childUsage.cost += child.getSessionStats().cost;
} catch {
// ignore — cost capture is best-effort
}
}
child.dispose();
}
},
});
}
export interface TodoItem {
content: string;
status: 'pending' | 'in_progress' | 'completed';
activeForm: string;
}
/** Render a todo list as a compact checklist for the workflow log. */
function renderTodos(todos: readonly TodoItem[]): string {
const mark = (s: TodoItem['status']): string => (s === 'completed' ? 'x' : s === 'in_progress' ? '~' : ' ');
return todos.map((t) => `[${mark(t.status)}] ${t.content}`).join(' ');
}
/**
* The `todo_write` tool — a full-state-replace planning scratchpad.
*
* Mirrors the TodoWrite tool: each call carries the entire list and replaces
* stored state (no append/merge). No deliverable impact; every call is echoed to
* the workflow log so `shannon logs` shows the agent's live plan. State is per
* tool instance (one per agent execution).
*/
export function createTodoWriteTool(auditLogger: AuditLogger): ToolDefinition {
let current: TodoItem[] = [];
return defineTool({
name: 'todo_write',
label: 'Todo Write',
description:
'Use this tool to create and manage a structured task list for your current session. This helps you ' +
'track progress and organize complex, multi-step work, and gives visibility into what you are doing. ' +
'Pass the COMPLETE todo list on every call — it replaces the stored list entirely (no append or merge). ' +
'Each todo has a status of pending, in_progress, or completed; keep exactly one task in_progress at a ' +
'time and mark a task completed as soon as it is finished.',
promptSnippet: 'todo_write: create and manage a structured task list',
parameters: Type.Object({
todos: Type.Array(
Type.Object({
content: Type.String({ description: 'Imperative task description, e.g. "Map SSRF sinks".' }),
status: Type.Union([Type.Literal('pending'), Type.Literal('in_progress'), Type.Literal('completed')]),
activeForm: Type.String({ description: 'Present-continuous form, e.g. "Mapping SSRF sinks".' }),
}),
),
}),
execute: async (_toolCallId, params) => {
current = params.todos as TodoItem[];
const completed = current.filter((t) => t.status === 'completed').length;
await auditLogger.logNote('todo', renderTodos(current));
return {
content: [{ type: 'text' as const, text: `Todos updated (${current.length} items, ${completed} completed).` }],
details: {},
};
},
});
}
/**
* The `glob` tool — fast file pattern matching (pi ships no `Glob` built-in).
*
* Backed by the same fast-glob engine that classifies code_path rules as `[GLOB]`
* (see utils/glob.ts `isGlobPattern`), so it enumerates exactly the patterns the
* routing tags as globs — including `**` and `{a,b}`, which pi's `find` would not
* match the same way. Returns absolute paths, most-recently-modified first.
*/
export function createGlobTool(cwd: string): ToolDefinition {
return defineTool({
name: 'glob',
label: 'Glob',
description:
'Fast file pattern matching. Supports glob patterns like "**/*.ts" or "src/**/*.{js,ts}". Returns ' +
'matching file paths sorted by modification time (most recent first), one per line, or "No files found".',
promptSnippet: 'glob: find files by name pattern',
parameters: Type.Object({
pattern: Type.String({ description: 'The glob pattern to match files against.' }),
path: Type.Optional(Type.String({ description: 'Directory to search in. Omit to search the repository root.' })),
}),
execute: async (_toolCallId, params) => {
const searchRoot = params.path ? path.resolve(cwd, params.path) : cwd;
const matches = await glob.globby(params.pattern, {
cwd: searchRoot,
absolute: true,
dot: true,
onlyFiles: true,
followSymbolicLinks: false,
});
if (matches.length === 0) {
return { content: [{ type: 'text' as const, text: 'No files found' }], details: {} };
}
// Sort by mtime (most recent first) to match the canonical Glob contract.
const withMtime = await Promise.all(
matches.map(async (file) => {
try {
return { file, mtime: (await fs.stat(file)).mtimeMs };
} catch {
return { file, mtime: 0 };
}
}),
);
withMtime.sort((a, b) => b.mtime - a.mtime);
return { content: [{ type: 'text' as const, text: withMtime.map((m) => m.file).join('\n') }], details: {} };
},
});
}
apps/worker/src/audit/audit-session.ts
+1 -1
View File
@@ -210,7 +210,7 @@ export class AuditSession {
/**
* Update session status
*/
async updateSessionStatus(status: 'in-progress' | 'completed' | 'failed' | 'cancelled'): Promise<void> {
async updateSessionStatus(status: 'in-progress' | 'completed' | 'failed' | 'cancelled' | 'partial'): Promise<void> {
await this.ensureInitialized();
const unlock = await sessionMutex.lock(this.sessionId);
apps/worker/src/audit/metrics-tracker.ts
+6 -6
View File
@@ -57,7 +57,7 @@ interface SessionData {
id: string;
webUrl: string;
repoPath?: string;
status: 'in-progress' | 'completed' | 'failed' | 'cancelled';
status: 'in-progress' | 'completed' | 'failed' | 'cancelled' | 'partial';
createdAt: string;
completedAt?: string;
originalWorkflowId?: string; // First workflow that created this workspace
@@ -214,9 +214,9 @@ export class MetricsTracker {
agent.checkpoint = result.checkpoint;
}
} else {
if (result.isFinalAttempt) {
agent.status = 'failed';
}
// A non-final failed attempt stays in-progress (Temporal will retry); only the
// terminal attempt (or an unqualified failure) marks the agent failed.
agent.status = result.isFinalAttempt === false ? 'in-progress' : 'failed';
}
// 7. Clear active timer
@@ -232,12 +232,12 @@ export class MetricsTracker {
/**
* Update session status
*/
async updateSessionStatus(status: 'in-progress' | 'completed' | 'failed' | 'cancelled'): Promise<void> {
async updateSessionStatus(status: 'in-progress' | 'completed' | 'failed' | 'cancelled' | 'partial'): Promise<void> {
if (!this.data) return;
this.data.session.status = status;
if (status === 'completed' || status === 'failed' || status === 'cancelled') {
if (status === 'completed' || status === 'failed' || status === 'cancelled' || status === 'partial') {
this.data.session.completedAt = formatTimestamp();
}
apps/worker/src/audit/workflow-logger.ts
+2 -2
View File
@@ -31,7 +31,7 @@ export interface AgentMetricsSummary {
}
export interface WorkflowSummary {
status: 'completed' | 'failed' | 'cancelled';
status: 'completed' | 'failed' | 'cancelled' | 'partial';
totalDurationMs: number;
totalCostUsd: number;
completedAgents: string[];
@@ -134,7 +134,7 @@ export class WorkflowLogger {
}
/**
* Format timestamp for log line (local time, human readable)
* Format timestamp for log line (UTC, human readable)
*/
private formatLogTime(): string {
const now = new Date();
apps/worker/src/mcp-server/exploit-collector.ts → apps/worker/src/collectors/exploit-collector.ts
+54 -62
View File
@@ -36,8 +36,9 @@
*/
import { defineTool, type ToolDefinition } from '@earendil-works/pi-coding-agent';
import { type Static, Type } from 'typebox';
import { type TSchema, Type } from 'typebox';
import { Value } from 'typebox/value';
import { stringEnum } from './schema.js';
// ============================================================================
// CLASS DISCRIMINATOR
@@ -84,6 +85,7 @@ export type BlockedExploit = {
vulnerability_id: string;
title: string;
vulnerable_location: string;
overview: string;
prerequisites?: string | null;
confidence: (typeof CONFIDENCE_VALUES)[number];
current_blocker: string;
@@ -101,7 +103,7 @@ export type AddExploitInput = ExploitedExploit | BlockedExploit;
// SCHEMA BUILDER
// ============================================================================
function buildSchemas(validIds: ReadonlySet<string>) {
export function buildSchemas(validIds: ReadonlySet<string>) {
const vulnerabilityIdField = Type.String({
minLength: 1,
description:
@@ -147,7 +149,7 @@ function buildSchemas(validIds: ReadonlySet<string>) {
}),
);
const statusField = Type.Union([Type.Literal('exploited'), Type.Literal('blocked')], {
const statusField = stringEnum(['exploited', 'blocked'], {
description:
'Verdict bucket. Set to "exploited" only after reaching Proof of Exploitation Level 3+ with ' +
'concrete impact evidence (extracted data, executed JavaScript, account takeover, internal ' +
@@ -157,11 +159,11 @@ function buildSchemas(validIds: ReadonlySet<string>) {
'route those to your workspace tracking file, not this tool.',
});
// Per-status fields. All optional at the flat parameter layer because a single
// Per-status fields. All optional at the flat shape layer because a single
// Type.Object cannot express a top-level discriminated union; the handler
// re-validates against the discriminated union below for true enforcement.
// re-validates against the strict union below for true enforcement.
const severityField = Type.Optional(
Type.Union([...SEVERITY_VALUES.map((v) => Type.Literal(v)), Type.Null()], {
Type.Union([stringEnum(SEVERITY_VALUES), Type.Null()], {
description:
'REQUIRED when status="exploited". Severity of the demonstrated impact. Critical = Level 4 ' +
'(admin credentials extracted, sensitive data dumped, system commands executed, full account ' +
@@ -205,7 +207,7 @@ function buildSchemas(validIds: ReadonlySet<string>) {
);
const confidenceField = Type.Optional(
Type.Union([...CONFIDENCE_VALUES.map((v) => Type.Literal(v)), Type.Null()], {
Type.Union([stringEnum(CONFIDENCE_VALUES), Type.Null()], {
description:
'REQUIRED when status="blocked". Confidence that this finding is a real vulnerability that ' +
'would be exploited if the external blocker were removed. High = code analysis strongly ' +
@@ -273,10 +275,10 @@ function buildSchemas(validIds: ReadonlySet<string>) {
}),
);
// The flat parameter schema passed to defineTool(). The harness uses this to
// build the agent's tool catalog. Per-status enforcement happens in the
// handler via the discriminated union below.
const flatShape = Type.Object({
// The flat shape passed to defineTool. pi uses this to build the agent's
// tool catalog. Per-status enforcement happens in the handler via the
// strict union below.
const flatSchema = Type.Object({
status: statusField,
vulnerability_id: vulnerabilityIdField,
title: titleField,
@@ -306,7 +308,7 @@ function buildSchemas(validIds: ReadonlySet<string>) {
vulnerable_location: vulnerableLocationField,
overview: overviewField,
prerequisites: prerequisitesField,
severity: Type.Union(SEVERITY_VALUES.map((v) => Type.Literal(v))),
severity: stringEnum(SEVERITY_VALUES),
impact: Type.String({ minLength: 1 }),
exploitation_steps: Type.Array(Type.String({ minLength: 1 }), { minItems: 1 }),
proof_of_impact: Type.String({ minLength: 1 }),
@@ -318,8 +320,9 @@ function buildSchemas(validIds: ReadonlySet<string>) {
vulnerability_id: vulnerabilityIdField,
title: titleField,
vulnerable_location: vulnerableLocationField,
overview: overviewField,
prerequisites: prerequisitesField,
confidence: Type.Union(CONFIDENCE_VALUES.map((v) => Type.Literal(v))),
confidence: stringEnum(CONFIDENCE_VALUES),
current_blocker: Type.String({ minLength: 1 }),
potential_impact: Type.String({ minLength: 1 }),
evidence_of_vulnerability: Type.String({ minLength: 1 }),
@@ -331,53 +334,45 @@ function buildSchemas(validIds: ReadonlySet<string>) {
const StrictSchema = Type.Union([ExploitedSchema, BlockedSchema]);
return { flatShape, StrictSchema };
return { flatSchema, StrictSchema };
}
type FlatInput = Static<ReturnType<typeof buildSchemas>['flatShape']>;
type StrictInput = Static<ReturnType<typeof buildSchemas>['StrictSchema']>;
// ============================================================================
// RESPONSE HELPERS
// ============================================================================
interface ToolResult {
content: Array<{ type: 'text'; text: string }>;
details: Record<string, unknown>;
isError?: boolean;
}
function createToolResult(response: { status: string; [key: string]: unknown }): ToolResult {
const isError = response.status === 'error';
function toolResult(payload: Record<string, unknown>) {
return {
content: [{ type: 'text' as const, text: JSON.stringify(response, null, 2) }],
details: {},
...(isError && { isError: true }),
content: [{ type: 'text' as const, text: JSON.stringify(payload, null, 2) }],
details: undefined,
};
}
function successResult(data: Record<string, unknown>): ToolResult {
return createToolResult({ status: 'success', ...data });
function successResult(data: Record<string, unknown>) {
return toolResult({ status: 'success', ...data });
}
function errorResult(message: string, errorType = 'ValidationError', retryable = true): ToolResult {
return createToolResult({ status: 'error', message, errorType, retryable });
function errorResult(message: string, errorType = 'ValidationError', retryable = true) {
return toolResult({ status: 'error', message, errorType, retryable });
}
function formatValueErrors(schema: ReturnType<typeof buildSchemas>['StrictSchema'], value: unknown): string {
return [...Value.Errors(schema, value)]
.map((issue) => {
const path = issue.instancePath.length > 0 ? issue.instancePath.replace(/^\//, '').replace(/\//g, '.') : '(root)';
return `- ${path}: ${issue.message}`;
})
.join('\n');
function formatValueErrors(schema: TSchema, value: unknown): string {
const issues: string[] = [];
for (const err of Value.Errors(schema, value)) {
const path =
err.instancePath && err.instancePath.length > 0
? err.instancePath.replace(/^\//, '').replace(/\//g, '.')
: '(root)';
issues.push(`- ${path}: ${err.message}`);
}
return issues.join('\n');
}
// ============================================================================
// TOOL FACTORY
// COLLECTOR FACTORY
// ============================================================================
export interface ExploitCollectorServer {
export interface ExploitCollector {
tools: ToolDefinition[];
getAll(): AddExploitInput[];
}
@@ -387,10 +382,10 @@ export interface CreateExploitCollectorOptions {
validIds: ReadonlySet<string>;
}
export function createExploitCollector(options: CreateExploitCollectorOptions): ExploitCollectorServer {
export function createExploitCollector(options: CreateExploitCollectorOptions): ExploitCollector {
const { vulnClass, validIds } = options;
const exploits: AddExploitInput[] = [];
const { flatShape, StrictSchema } = buildSchemas(validIds);
const { flatSchema, StrictSchema } = buildSchemas(validIds);
const addExploitTool = defineTool({
name: 'add_exploit',
@@ -405,21 +400,8 @@ export function createExploitCollector(options: CreateExploitCollectorOptions):
'IDs. FALSE POSITIVE findings do NOT use this tool — they go to your workspace tracking file. ' +
'After all queue vulnerabilities have been emitted, the host renderer assembles the ' +
'deliverable Markdown from your recorded calls.',
parameters: flatShape,
execute: async (_toolCallId, args): Promise<ToolResult> => {
const input = args as FlatInput;
// Strict queue-ID validation: reject hallucinated or typo'd IDs with the valid-ID list.
if (!validIds.has(input.vulnerability_id)) {
return errorResult(
`Vulnerability ID not in this run's queue. Valid IDs: ` +
`${formatValidIdsPreview(validIds)}. ` +
'Check the queue.json for the canonical ID — likely a typo or hallucinated ID.',
'ValidationError',
true,
);
}
parameters: flatSchema,
async execute(_toolCallId, input) {
// Re-validate against the strict discriminated union for per-status enforcement.
if (!Value.Check(StrictSchema, input)) {
return errorResult(
@@ -430,9 +412,19 @@ export function createExploitCollector(options: CreateExploitCollectorOptions):
true,
);
}
// Strip excess properties from the flat input so only the chosen status's
// fields survive (mirrors the prior discriminated-union parse).
const typed = Value.Clean(StrictSchema, structuredClone(input)) as StrictInput as AddExploitInput;
const typed = Value.Clean(StrictSchema, structuredClone(input)) as AddExploitInput;
// Reject IDs not in this run's queue (typo'd or hallucinated).
if (!validIds.has(typed.vulnerability_id)) {
return errorResult(
`Vulnerability ID "${typed.vulnerability_id}" not in this run's queue. Valid IDs: ` +
`${formatValidIdsPreview(validIds)}. ` +
'Check the queue.json for the canonical ID — likely a typo or hallucinated ID.',
'ValidationError',
true,
);
}
const existing = exploits.find((e) => e.vulnerability_id === typed.vulnerability_id);
if (existing) {
return errorResult(
@@ -448,7 +440,7 @@ export function createExploitCollector(options: CreateExploitCollectorOptions):
});
return {
tools: [addExploitTool] as ToolDefinition[],
tools: [addExploitTool],
getAll: (): AddExploitInput[] => [...exploits],
};
}
apps/worker/src/mcp-server/pre-recon-collector.ts → apps/worker/src/collectors/pre-recon-collector.ts
+152 -142
View File
@@ -21,6 +21,7 @@
import { defineTool, type ToolDefinition } from '@earendil-works/pi-coding-agent';
import { type Static, Type } from 'typebox';
import { cleanInput } from './schema.js';
// ============================================================================
// SHARED SCHEMA
@@ -64,100 +65,107 @@ export const ExecutiveSummaryInputSchema = Type.Object({
}),
});
const ArchitectureSchema = Type.Object({
framework_and_language: Type.String({
minLength: 1,
description: 'Framework and language details with their security implications.',
}),
architectural_pattern: Type.String({
minLength: 1,
description: 'Architectural pattern (monolith, microservices, hybrid) with trust boundary analysis.',
}),
critical_security_components: Type.String({
minLength: 1,
description: 'Critical security components with focus on auth, authz, and data protection.',
}),
});
const DataSecuritySchema = Type.Object({
database_security: Type.String({
minLength: 1,
description: 'Analyze encryption, access controls, and query safety in database interactions.',
}),
data_flow_security: Type.String({
minLength: 1,
description: 'Identify sensitive data paths and the protection mechanisms applied along them.',
}),
multi_tenant_isolation: Type.String({
minLength: 1,
description:
'Assess tenant separation effectiveness. If the application is single-tenant, state that ' +
'explicitly rather than leaving the field thin.',
}),
});
const AttackSurfaceSchema = Type.Object({
external_entry_points: Type.String({
minLength: 1,
description: 'Detailed analysis of each public interface that is network-accessible.',
}),
internal_service_communication: Type.String({
minLength: 1,
description:
'Trust relationships and security assumptions between network-reachable services. ' +
'If the application is a single service with no internal RPC fabric, state that.',
}),
input_validation_patterns: Type.String({
minLength: 1,
description: 'How user input is handled and validated in network-accessible endpoints.',
}),
background_processing: Type.String({
minLength: 1,
description:
'Async job security and privilege models for jobs triggered by network requests. ' +
'If no async/background processing exists, state that.',
}),
});
const InfrastructureSchema = Type.Object({
secrets_management: Type.String({ minLength: 1, description: 'How secrets are stored, rotated, and accessed.' }),
configuration_security: Type.String({
minLength: 1,
description:
'Environment separation and secret handling. Specifically search for infrastructure ' +
'configuration (e.g., Nginx, Kubernetes Ingress, CDN settings) that defines security ' +
'headers like Strict-Transport-Security (HSTS) and Cache-Control, and report what was found.',
}),
external_dependencies: Type.String({
minLength: 1,
description: 'Third-party services and their security implications.',
}),
monitoring_and_logging: Type.String({
minLength: 1,
description: 'Security event visibility — what is logged, where it goes, and who can see it.',
}),
});
export const ApplicationIntelligenceInputSchema = Type.Object({
architecture: Type.Object(ArchitectureSchema.properties, {
description:
'Architecture & Technology Stack — driven by the Architecture Scanner sub-agent. ' +
'Becomes Section 2 of the rendered deliverable.',
}),
data_security: Type.Object(DataSecuritySchema.properties, {
description:
'Data Security & Storage — driven by the Data Security Auditor sub-agent. ' +
'Becomes Section 4 of the rendered deliverable.',
}),
attack_surface: Type.Object(AttackSurfaceSchema.properties, {
description:
'Attack Surface Analysis — driven by Entry Point Mapper + Architecture Scanner sub-agents. ' +
'Only include entry points confirmed to be in-scope (network-reachable). ' +
'Becomes Section 5 of the rendered deliverable.',
}),
infrastructure: Type.Object(InfrastructureSchema.properties, {
description: 'Infrastructure & Operational Security. Becomes Section 6 of the rendered deliverable.',
}),
architecture: Type.Object(
{
framework_and_language: Type.String({
minLength: 1,
description: 'Framework and language details with their security implications.',
}),
architectural_pattern: Type.String({
minLength: 1,
description: 'Architectural pattern (monolith, microservices, hybrid) with trust boundary analysis.',
}),
critical_security_components: Type.String({
minLength: 1,
description: 'Critical security components with focus on auth, authz, and data protection.',
}),
},
{
description:
'Architecture & Technology Stack — driven by the Architecture Scanner sub-agent. ' +
'Becomes Section 2 of the rendered deliverable.',
},
),
data_security: Type.Object(
{
database_security: Type.String({
minLength: 1,
description: 'Analyze encryption, access controls, and query safety in database interactions.',
}),
data_flow_security: Type.String({
minLength: 1,
description: 'Identify sensitive data paths and the protection mechanisms applied along them.',
}),
multi_tenant_isolation: Type.String({
minLength: 1,
description:
'Assess tenant separation effectiveness. If the application is single-tenant, state that ' +
'explicitly rather than leaving the field thin.',
}),
},
{
description:
'Data Security & Storage — driven by the Data Security Auditor sub-agent. ' +
'Becomes Section 4 of the rendered deliverable.',
},
),
attack_surface: Type.Object(
{
external_entry_points: Type.String({
minLength: 1,
description: 'Detailed analysis of each public interface that is network-accessible.',
}),
internal_service_communication: Type.String({
minLength: 1,
description:
'Trust relationships and security assumptions between network-reachable services. ' +
'If the application is a single service with no internal RPC fabric, state that.',
}),
input_validation_patterns: Type.String({
minLength: 1,
description: 'How user input is handled and validated in network-accessible endpoints.',
}),
background_processing: Type.String({
minLength: 1,
description:
'Async job security and privilege models for jobs triggered by network requests. ' +
'If no async/background processing exists, state that.',
}),
},
{
description:
'Attack Surface Analysis — driven by Entry Point Mapper + Architecture Scanner sub-agents. ' +
'Only include entry points confirmed to be in-scope (network-reachable). ' +
'Becomes Section 5 of the rendered deliverable.',
},
),
infrastructure: Type.Object(
{
secrets_management: Type.String({
minLength: 1,
description: 'How secrets are stored, rotated, and accessed.',
}),
configuration_security: Type.String({
minLength: 1,
description:
'Environment separation and secret handling. Specifically search for infrastructure ' +
'configuration (e.g., Nginx, Kubernetes Ingress, CDN settings) that defines security ' +
'headers like Strict-Transport-Security (HSTS) and Cache-Control, and report what was found.',
}),
external_dependencies: Type.String({
minLength: 1,
description: 'Third-party services and their security implications.',
}),
monitoring_and_logging: Type.String({
minLength: 1,
description: 'Security event visibility — what is logged, where it goes, and who can see it.',
}),
},
{
description: 'Infrastructure & Operational Security. Becomes Section 6 of the rendered deliverable.',
},
),
});
export const AuthDeepDiveInputSchema = Type.Object({
@@ -173,7 +181,10 @@ export const AuthDeepDiveInputSchema = Type.Object({
'Session management and token security. Pinpoint the exact file and line(s) of code where ' +
'session cookie flags (HttpOnly, Secure, SameSite) are configured.',
}),
authz_model: Type.String({ minLength: 1, description: 'Authorization model and potential bypass scenarios.' }),
authz_model: Type.String({
minLength: 1,
description: 'Authorization model and potential bypass scenarios.',
}),
multi_tenancy: Type.String({
minLength: 1,
description: 'Multi-tenancy security implementation. If the application is single-tenant, state that explicitly.',
@@ -393,44 +404,45 @@ export type PreReconCallStatus = Readonly<Record<PreReconToolName, PreReconToolS
// RESPONSE HELPERS
// ============================================================================
interface ToolResult {
content: Array<{ type: 'text'; text: string }>;
details: Record<string, unknown>;
isError?: boolean;
function toolResult(payload: Record<string, unknown>) {
return {
content: [{ type: 'text' as const, text: JSON.stringify(payload, null, 2) }],
details: undefined,
};
}
function successResult(data: Record<string, unknown>): ToolResult {
const response = { status: 'success', ...data };
return { content: [{ type: 'text' as const, text: JSON.stringify(response, null, 2) }], details: {} };
function successResult(data: Record<string, unknown>) {
return toolResult({ status: 'success', ...data });
}
function errorResult(message: string, errorType = 'ValidationError', retryable = true): ToolResult {
const response = { status: 'error', message, errorType, retryable };
return { content: [{ type: 'text' as const, text: JSON.stringify(response, null, 2) }], details: {}, isError: true };
function errorResult(message: string, errorType = 'ValidationError', retryable = true) {
return toolResult({ status: 'error', message, errorType, retryable });
}
// ============================================================================
// TOOLS FACTORY
// COLLECTOR FACTORY
// ============================================================================
export interface PreReconCollectorServer {
interface PreReconState {
executive_summary?: ExecutiveSummaryInput;
application_intelligence?: ApplicationIntelligenceInput;
auth_deep_dive?: AuthDeepDiveInput;
codebase_indexing?: CodebaseIndexingInput;
critical_file_paths?: CriticalFilePathsInput;
xss_sinks?: XssSinksInput;
ssrf_sinks?: SsrfSinksInput;
}
export interface PreReconCollector {
tools: ToolDefinition[];
getAll(): PreReconData;
getCallStatus(): PreReconCallStatus;
}
export function createPreReconCollectorServer(): PreReconCollectorServer {
const state: {
executive_summary?: ExecutiveSummaryInput;
application_intelligence?: ApplicationIntelligenceInput;
auth_deep_dive?: AuthDeepDiveInput;
codebase_indexing?: CodebaseIndexingInput;
critical_file_paths?: CriticalFilePathsInput;
xss_sinks?: XssSinksInput;
ssrf_sinks?: SsrfSinksInput;
} = {};
export function createPreReconCollector(): PreReconCollector {
const state: PreReconState = {};
function alreadyCalled(toolName: PreReconToolName): ToolResult {
function alreadyCalled(toolName: PreReconToolName) {
return errorResult(
`${toolName} has already been called. Each set_* tool may only be called once per run.`,
'DuplicateError',
@@ -446,9 +458,9 @@ export function createPreReconCollectorServer(): PreReconCollectorServer {
'Call exactly once before terminating. Becomes Section 1 of the rendered deliverable. ' +
'Duplicate calls are rejected.',
parameters: ExecutiveSummaryInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.executive_summary) return alreadyCalled('set_executive_summary');
state.executive_summary = input;
state.executive_summary = cleanInput(ExecutiveSummaryInputSchema, input);
return successResult({ set: 'set_executive_summary' });
},
});
@@ -461,9 +473,9 @@ export function createPreReconCollectorServer(): PreReconCollectorServer {
'and infrastructure — in a single call. Call exactly once before terminating. ' +
'Becomes Sections 2, 4, 5, and 6 of the rendered deliverable. Duplicate calls are rejected.',
parameters: ApplicationIntelligenceInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.application_intelligence) return alreadyCalled('set_application_intelligence');
state.application_intelligence = input;
state.application_intelligence = cleanInput(ApplicationIntelligenceInputSchema, input);
return successResult({ set: 'set_application_intelligence' });
},
});
@@ -475,9 +487,9 @@ export function createPreReconCollectorServer(): PreReconCollectorServer {
'Record the authentication & authorization deep dive. Call exactly once before terminating. ' +
'Becomes Section 3 of the rendered deliverable. Duplicate calls are rejected.',
parameters: AuthDeepDiveInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.auth_deep_dive) return alreadyCalled('set_auth_deep_dive');
state.auth_deep_dive = input;
state.auth_deep_dive = cleanInput(AuthDeepDiveInputSchema, input);
return successResult({ set: 'set_auth_deep_dive' });
},
});
@@ -489,9 +501,9 @@ export function createPreReconCollectorServer(): PreReconCollectorServer {
'Record the overall codebase indexing narrative. Call exactly once before terminating. ' +
'Becomes Section 7 of the rendered deliverable. Duplicate calls are rejected.',
parameters: CodebaseIndexingInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.codebase_indexing) return alreadyCalled('set_codebase_indexing');
state.codebase_indexing = input;
state.codebase_indexing = cleanInput(CodebaseIndexingInputSchema, input);
return successResult({ set: 'set_codebase_indexing' });
},
});
@@ -504,9 +516,9 @@ export function createPreReconCollectorServer(): PreReconCollectorServer {
'before terminating. Becomes Section 8 of the rendered deliverable. The next agent uses this ' +
'as a starting point for manual review. Duplicate calls are rejected.',
parameters: CriticalFilePathsInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.critical_file_paths) return alreadyCalled('set_critical_file_paths');
state.critical_file_paths = input;
state.critical_file_paths = cleanInput(CriticalFilePathsInputSchema, input);
return successResult({ set: 'set_critical_file_paths' });
},
});
@@ -521,9 +533,9 @@ export function createPreReconCollectorServer(): PreReconCollectorServer {
"the vuln-xss agent's testing todos downstream. Becomes Section 9 of the rendered deliverable. " +
'Duplicate calls are rejected.',
parameters: XssSinksInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.xss_sinks) return alreadyCalled('set_xss_sinks');
state.xss_sinks = input;
state.xss_sinks = cleanInput(XssSinksInputSchema, input);
return successResult({ set: 'set_xss_sinks' });
},
});
@@ -538,23 +550,13 @@ export function createPreReconCollectorServer(): PreReconCollectorServer {
"the vuln-ssrf agent's testing todos downstream. Becomes Section 10 of the rendered deliverable. " +
'Duplicate calls are rejected.',
parameters: SsrfSinksInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.ssrf_sinks) return alreadyCalled('set_ssrf_sinks');
state.ssrf_sinks = input;
state.ssrf_sinks = cleanInput(SsrfSinksInputSchema, input);
return successResult({ set: 'set_ssrf_sinks' });
},
});
const tools: ToolDefinition[] = [
setExecutiveSummary,
setApplicationIntelligence,
setAuthDeepDive,
setCodebaseIndexing,
setCriticalFilePaths,
setXssSinks,
setSsrfSinks,
];
function statusOf<K extends PreReconToolName>(key: K): PreReconToolStatus {
const flagMap: Record<PreReconToolName, unknown> = {
set_executive_summary: state.executive_summary,
@@ -569,7 +571,15 @@ export function createPreReconCollectorServer(): PreReconCollectorServer {
}
return {
tools,
tools: [
setExecutiveSummary,
setApplicationIntelligence,
setAuthDeepDive,
setCodebaseIndexing,
setCriticalFilePaths,
setXssSinks,
setSsrfSinks,
],
getAll: (): PreReconData => ({
...(state.executive_summary && { executive_summary: state.executive_summary }),
...(state.application_intelligence && { application_intelligence: state.application_intelligence }),
apps/worker/src/mcp-server/recon-collector.ts → apps/worker/src/collectors/recon-collector.ts
+218 -224
View File
@@ -5,7 +5,7 @@
// as published by the Free Software Foundation.
/**
* Recon Collector MCP Server
* Recon Collector tools
*
* Exposes nine TypeBox-validated tools that feed the recon_deliverable.md
* renderer — eight one-shot `set_*` tools, one per deliverable section, plus a
@@ -22,6 +22,7 @@
import { defineTool, type ToolDefinition } from '@earendil-works/pi-coding-agent';
import { type Static, Type } from 'typebox';
import { type SinkRef, SinkRefSchema } from './pre-recon-collector.js';
import { cleanInput, stringEnum } from './schema.js';
// ============================================================================
// PER-TOOL INPUT SCHEMAS
@@ -52,122 +53,127 @@ export const TechnologyStackInputSchema = Type.Object({
}),
});
const SessionFlowSchema = Type.Object({
entry_points: Type.String({
minLength: 1,
description: 'Authentication entry points (e.g., /login, /register, /auth/sso).',
}),
mechanism: Type.String({
minLength: 1,
description:
'Describe the step-by-step authentication process: credential submission, token generation, ' +
'cookie setting, redirects, etc.',
}),
code_pointers: Type.String({
minLength: 1,
description:
'Pointers to the primary files and functions in the codebase that manage authentication and ' + 'session logic.',
}),
});
const RoleAssignmentSchema = Type.Object({
role_determination: Type.String({
minLength: 1,
description: 'How roles are assigned post-authentication — database lookup, JWT claims, external service, etc.',
}),
default_role: Type.String({ minLength: 1, description: 'What role new users get by default.' }),
role_upgrade_path: Type.String({
minLength: 1,
description:
'How users can gain higher privileges — admin approval, self-service, automatic, etc. ' +
'If no upgrade path exists, state that.',
}),
code_implementation: Type.String({
minLength: 1,
description: 'Where role assignment logic is implemented (file paths and functions).',
}),
});
const PrivilegeStorageSchema = Type.Object({
storage_location: Type.String({
minLength: 1,
description: 'Where user privileges are stored — JWT claims, session data, database, external service.',
}),
validation_points: Type.String({
minLength: 1,
description: 'Where role checks happen — middleware, decorators, inline checks.',
}),
cache_session_persistence: Type.String({
minLength: 1,
description: 'How long privileges are cached, and when they are refreshed.',
}),
code_pointers: Type.String({ minLength: 1, description: 'Files that handle privilege validation.' }),
});
const RoleSwitchingImpersonationSchema = Type.Object({
applicable: Type.Boolean({
description:
'False only if the application has no impersonation, sudo-mode, or role-switching features ' +
'at all. When false, the other fields in this object may be null.',
}),
impersonation_features: Type.Union([Type.String(), Type.Null()], {
description:
'Any ability for admins or higher-privilege users to impersonate other users. Pass null when ' +
'applicable is false.',
}),
role_switching: Type.Union([Type.String(), Type.Null()], {
description: 'Temporary privilege elevation mechanisms like "sudo mode". Pass null when applicable is false.',
}),
audit_trail: Type.Union([Type.String(), Type.Null()], {
description:
'Whether role switches or impersonation events are logged, and where. Pass null when applicable is false.',
}),
code_implementation: Type.Union([Type.String(), Type.Null()], {
description: 'Where these features are implemented (file paths and functions). Pass null when applicable is false.',
}),
});
export const AuthenticationInputSchema = Type.Object({
session_flow: Type.Object(SessionFlowSchema.properties, {
description:
'Authentication & Session Management Flow — overall entry points, mechanism, and code pointers. ' +
'Becomes Section 3 of the rendered deliverable.',
}),
role_assignment: Type.Object(RoleAssignmentSchema.properties, {
description: 'Role Assignment Process — how roles are determined post-authentication. ' + 'Becomes Section 3.1.',
}),
privilege_storage: Type.Object(PrivilegeStorageSchema.properties, {
description:
'Privilege Storage & Validation — where privileges live and where they are checked. ' + 'Becomes Section 3.2.',
}),
role_switching_impersonation: Type.Object(RoleSwitchingImpersonationSchema.properties, {
description:
'Role Switching & Impersonation — impersonation, sudo mode, audit trails. Becomes Section 3.3. ' +
'Set applicable=false if no such features exist; the other fields may be null in that case.',
}),
session_flow: Type.Object(
{
entry_points: Type.String({
minLength: 1,
description: 'Authentication entry points (e.g., /login, /register, /auth/sso).',
}),
mechanism: Type.String({
minLength: 1,
description:
'Describe the step-by-step authentication process: credential submission, token generation, ' +
'cookie setting, redirects, etc.',
}),
code_pointers: Type.String({
minLength: 1,
description:
'Pointers to the primary files and functions in the codebase that manage authentication and ' +
'session logic.',
}),
},
{
description:
'Authentication & Session Management Flow — overall entry points, mechanism, and code pointers. ' +
'Becomes Section 3 of the rendered deliverable.',
},
),
role_assignment: Type.Object(
{
role_determination: Type.String({
minLength: 1,
description: 'How roles are assigned post-authentication — database lookup, JWT claims, external service, etc.',
}),
default_role: Type.String({
minLength: 1,
description: 'What role new users get by default.',
}),
role_upgrade_path: Type.String({
minLength: 1,
description:
'How users can gain higher privileges — admin approval, self-service, automatic, etc. ' +
'If no upgrade path exists, state that.',
}),
code_implementation: Type.String({
minLength: 1,
description: 'Where role assignment logic is implemented (file paths and functions).',
}),
},
{
description: 'Role Assignment Process — how roles are determined post-authentication. Becomes Section 3.1.',
},
),
privilege_storage: Type.Object(
{
storage_location: Type.String({
minLength: 1,
description: 'Where user privileges are stored — JWT claims, session data, database, external service.',
}),
validation_points: Type.String({
minLength: 1,
description: 'Where role checks happen — middleware, decorators, inline checks.',
}),
cache_session_persistence: Type.String({
minLength: 1,
description: 'How long privileges are cached, and when they are refreshed.',
}),
code_pointers: Type.String({
minLength: 1,
description: 'Files that handle privilege validation.',
}),
},
{
description:
'Privilege Storage & Validation — where privileges live and where they are checked. ' + 'Becomes Section 3.2.',
},
),
role_switching_impersonation: Type.Object(
{
applicable: Type.Boolean({
description:
'False only if the application has no impersonation, sudo-mode, or role-switching features ' +
'at all. When false, the other fields in this object may be null.',
}),
impersonation_features: Type.Union([Type.String(), Type.Null()], {
description:
'Any ability for admins or higher-privilege users to impersonate other users. Pass null when ' +
'applicable is false.',
}),
role_switching: Type.Union([Type.String(), Type.Null()], {
description: 'Temporary privilege elevation mechanisms like "sudo mode". Pass null when applicable is false.',
}),
audit_trail: Type.Union([Type.String(), Type.Null()], {
description:
'Whether role switches or impersonation events are logged, and where. Pass null when applicable is false.',
}),
code_implementation: Type.Union([Type.String(), Type.Null()], {
description:
'Where these features are implemented (file paths and functions). Pass null when applicable is false.',
}),
},
{
description:
'Role Switching & Impersonation — impersonation, sudo mode, audit trails. Becomes Section 3.3. ' +
'Set applicable=false if no such features exist; the other fields may be null in that case.',
},
),
});
const HttpMethodSchema = Type.Union(
[
Type.Literal('GET'),
Type.Literal('POST'),
Type.Literal('PUT'),
Type.Literal('PATCH'),
Type.Literal('DELETE'),
Type.Literal('OPTIONS'),
Type.Literal('HEAD'),
Type.Literal('WS'),
],
{ description: 'HTTP method. Use WS for WebSocket upgrade endpoints.' },
);
const HTTP_METHOD_VALUES = ['GET', 'POST', 'PUT', 'PATCH', 'DELETE', 'OPTIONS', 'HEAD', 'WS'] as const;
const EndpointSchema = Type.Object({
method: HttpMethodSchema,
method: stringEnum(HTTP_METHOD_VALUES, {
description: 'HTTP method. Use WS for WebSocket upgrade endpoints.',
}),
path: Type.String({
minLength: 1,
description: 'Endpoint path with parameter placeholders, e.g. "/api/users/{user_id}".',
}),
required_role: Type.String({ minLength: 1, description: 'Minimum role needed (anon, user, admin, etc.).' }),
required_role: Type.String({
minLength: 1,
description: 'Minimum role needed (anon, user, admin, etc.).',
}),
object_id_parameters: Type.Array(Type.String(), {
description: 'Parameters that identify specific objects (user_id, order_id, etc.). Empty array if none.',
}),
@@ -177,7 +183,10 @@ const EndpointSchema = Type.Object({
'How access is controlled — middleware, decorator, inline check. ' +
'E.g. "Bearer Token + ownership check", "requireAuth() + requireAdmin()", "None".',
}),
description: Type.String({ minLength: 1, description: "Brief description of the endpoint's purpose." }),
description: Type.String({
minLength: 1,
description: "Brief description of the endpoint's purpose.",
}),
code_pointer: Type.String({
minLength: 1,
description: 'File path and (where possible) line number of the handler. E.g. "auth.controller.ts:45".',
@@ -215,55 +224,33 @@ export const InputVectorsInputSchema = Type.Object({
}),
});
const EntityTypeSchema = Type.Union([
Type.Literal('ExternAsset'),
Type.Literal('Service'),
Type.Literal('Identity'),
Type.Literal('DataStore'),
Type.Literal('AdminPlane'),
Type.Literal('ThirdParty'),
]);
const ENTITY_TYPE_VALUES = ['ExternAsset', 'Service', 'Identity', 'DataStore', 'AdminPlane', 'ThirdParty'] as const;
const EntityZoneSchema = Type.Union([
Type.Literal('Internet'),
Type.Literal('Edge'),
Type.Literal('App'),
Type.Literal('Data'),
Type.Literal('Admin'),
Type.Literal('BuildCI'),
Type.Literal('ThirdParty'),
]);
const ENTITY_ZONE_VALUES = ['Internet', 'Edge', 'App', 'Data', 'Admin', 'BuildCI', 'ThirdParty'] as const;
const DataLabelSchema = Type.Union([
Type.Literal('PII'),
Type.Literal('Tokens'),
Type.Literal('Payments'),
Type.Literal('Secrets'),
Type.Literal('Public'),
]);
const DATA_LABEL_VALUES = ['PII', 'Tokens', 'Payments', 'Secrets', 'Public'] as const;
const FlowChannelSchema = Type.Union([
Type.Literal('HTTP'),
Type.Literal('HTTPS'),
Type.Literal('TCP'),
Type.Literal('Message'),
Type.Literal('File'),
Type.Literal('Token'),
]);
const FLOW_CHANNEL_VALUES = ['HTTP', 'HTTPS', 'TCP', 'Message', 'File', 'Token'] as const;
const GuardCategorySchema = Type.Union([
Type.Literal('Auth'),
Type.Literal('Network'),
Type.Literal('Protocol'),
Type.Literal('Env'),
Type.Literal('RateLimit'),
Type.Literal('Authorization'),
Type.Literal('ObjectOwnership'),
]);
const GUARD_CATEGORY_VALUES = [
'Auth',
'Network',
'Protocol',
'Env',
'RateLimit',
'Authorization',
'ObjectOwnership',
] as const;
const EntityMetadataPairSchema = Type.Object({
key: Type.String({ minLength: 1, description: 'Metadata key (e.g., "Hosts", "Endpoints", "Engine", "Issuer").' }),
value: Type.String({ minLength: 1, description: 'Metadata value for this key.' }),
key: Type.String({
minLength: 1,
description: 'Metadata key (e.g., "Hosts", "Endpoints", "Engine", "Issuer").',
}),
value: Type.String({
minLength: 1,
description: 'Metadata value for this key.',
}),
});
const EntitySchema = Type.Object({
@@ -271,13 +258,13 @@ const EntitySchema = Type.Object({
minLength: 1,
description: 'Unique short name for the entity (e.g., "ExampleWebApp", "PostgreSQL-DB", "IdentityProvider").',
}),
type: Type.Union(EntityTypeSchema.anyOf, {
type: stringEnum(ENTITY_TYPE_VALUES, {
description:
'Entity type. ExternAsset = client-side asset; Service = backend service; Identity = identity ' +
'provider; DataStore = database / cache / object store; AdminPlane = admin/control surface; ' +
'ThirdParty = external integration.',
}),
zone: Type.Union(EntityZoneSchema.anyOf, {
zone: stringEnum(ENTITY_ZONE_VALUES, {
description:
'Trust zone. Internet = public; Edge = CDN/WAF/reverse-proxy tier; App = application/business logic; ' +
'Data = persistent storage; Admin = administrative surface; BuildCI = build/CI/CD infrastructure; ' +
@@ -287,7 +274,7 @@ const EntitySchema = Type.Object({
minLength: 1,
description: 'Short technology/framework description (e.g., "Node/Express", "Postgres 14", "AWS S3").',
}),
data: Type.Array(DataLabelSchema, {
data: Type.Array(stringEnum(DATA_LABEL_VALUES), {
description: 'Data labels handled by this entity. Empty array if the entity handles only Public data.',
}),
notes: Type.String({
@@ -302,12 +289,15 @@ const EntitySchema = Type.Object({
});
const FlowSchema = Type.Object({
from: Type.String({ minLength: 1, description: 'Source entity title — must match a title from the entities array.' }),
from: Type.String({
minLength: 1,
description: 'Source entity title — must match a title from the entities array.',
}),
to: Type.String({
minLength: 1,
description: 'Destination entity title — must match a title from the entities array.',
}),
channel: Type.Union(FlowChannelSchema.anyOf, { description: 'Transport channel for this flow.' }),
channel: stringEnum(FLOW_CHANNEL_VALUES, { description: 'Transport channel for this flow.' }),
path_port: Type.String({
minLength: 1,
description: 'Path and/or port for this flow. E.g. ":443 /api/users/me", ":5432", "queue: orders".',
@@ -317,7 +307,7 @@ const FlowSchema = Type.Object({
'Guard names that gate this flow. Each should match a name from the guards array. Empty array ' +
'means no guards apply (publicly accessible).',
}),
touches: Type.Array(DataLabelSchema, {
touches: Type.Array(stringEnum(DATA_LABEL_VALUES), {
description: 'Data labels this flow carries. Empty array if only Public data flows.',
}),
});
@@ -327,14 +317,17 @@ const GuardSchema = Type.Object({
minLength: 1,
description: 'Short guard identifier (e.g., "auth:user", "ownership:user", "vpc-only", "mtls").',
}),
category: Type.Union(GuardCategorySchema.anyOf, {
category: stringEnum(GUARD_CATEGORY_VALUES, {
description:
'Guard category. Auth = authentication identity; Authorization = role/scope check; ' +
'ObjectOwnership = ownership-based check; Network = network-level restriction; ' +
'Protocol = protocol-level requirement; Env = environment-bound restriction; ' +
'RateLimit = throttling.',
}),
statement: Type.String({ minLength: 1, description: 'One-sentence description of what this guard enforces.' }),
statement: Type.String({
minLength: 1,
description: 'One-sentence description of what this guard enforces.',
}),
});
export const NetworkMapInputSchema = Type.Object({
@@ -348,17 +341,25 @@ export const NetworkMapInputSchema = Type.Object({
'How entities communicate. Becomes Section 6.3. The from/to fields cross-reference entities ' +
'by title; the guards field cross-references guards by name.',
}),
guards: Type.Array(GuardSchema, { description: 'Catalog of guards referenced by flows. Becomes Section 6.4.' }),
guards: Type.Array(GuardSchema, {
description: 'Catalog of guards referenced by flows. Becomes Section 6.4.',
}),
});
const RoleSchema = Type.Object({
name: Type.String({ minLength: 1, description: 'Role name (e.g., "anon", "user", "admin", "team_admin").' }),
name: Type.String({
minLength: 1,
description: 'Role name (e.g., "anon", "user", "admin", "team_admin").',
}),
privilege_level: Type.Integer({
minimum: 0,
maximum: 10,
description: 'Privilege rank from 0 (lowest, anonymous) to 10 (highest, full admin).',
}),
scope_domain: Type.String({ minLength: 1, description: 'Scope of this role: Global, Org, Team, Project, etc.' }),
scope_domain: Type.String({
minLength: 1,
description: 'Scope of this role: Global, Org, Team, Project, etc.',
}),
code_implementation: Type.String({
minLength: 1,
description: 'Where this role is defined or checked (middleware, decorator, file:line, etc.).',
@@ -421,10 +422,8 @@ export const RoleArchitectureInputSchema = Type.Object({
const PRIORITY_VALUES = ['High', 'Medium', 'Low'] as const;
const PrioritySchema = Type.Union([Type.Literal('High'), Type.Literal('Medium'), Type.Literal('Low')]);
const HorizontalCandidateSchema = Type.Object({
priority: Type.Union(PrioritySchema.anyOf, {
priority: stringEnum(PRIORITY_VALUES, {
description: 'Priority: High, Medium, or Low, based on data sensitivity (title-case literals).',
}),
endpoint_pattern: Type.String({
@@ -458,7 +457,7 @@ const VerticalCandidateSchema = Type.Object({
minLength: 1,
description: 'What the endpoint does (e.g., "Administrative functions", "User management").',
}),
risk_level: Type.Union(PrioritySchema.anyOf, {
risk_level: stringEnum(PRIORITY_VALUES, {
description: 'Risk level: High, Medium, or Low (title-case literals).',
}),
});
@@ -605,24 +604,19 @@ export interface ReconCallStatus {
// RESPONSE HELPERS
// ============================================================================
interface ToolResult {
content: Array<{ type: 'text'; text: string }>;
details: Record<string, unknown>;
}
function createToolResult(response: { status: string; [key: string]: unknown }): ToolResult {
function toolResult(payload: Record<string, unknown>) {
return {
content: [{ type: 'text' as const, text: JSON.stringify(response, null, 2) }],
details: {},
content: [{ type: 'text' as const, text: JSON.stringify(payload, null, 2) }],
details: undefined,
};
}
function successResult(data: Record<string, unknown>): ToolResult {
return createToolResult({ status: 'success', ...data });
function successResult(data: Record<string, unknown>) {
return toolResult({ status: 'success', ...data });
}
function errorResult(message: string, errorType = 'ValidationError', retryable = true): ToolResult {
return createToolResult({ status: 'error', message, errorType, retryable });
function errorResult(message: string, errorType = 'ValidationError', retryable = true) {
return toolResult({ status: 'error', message, errorType, retryable });
}
function endpointKey(method: string, path: string): string {
@@ -630,32 +624,34 @@ function endpointKey(method: string, path: string): string {
}
// ============================================================================
// SERVER FACTORY
// COLLECTOR FACTORY
// ============================================================================
export interface ReconCollectorServer {
interface ReconState {
executive_summary?: ExecutiveSummaryInput;
technology_stack?: TechnologyStackInput;
authentication?: AuthenticationInput;
input_vectors?: InputVectorsInput;
network_map?: NetworkMapInput;
role_architecture?: RoleArchitectureInput;
authz_candidates?: AuthzCandidatesInput;
injection_sources?: InjectionSourcesInput;
}
export interface ReconCollector {
tools: ToolDefinition[];
getAll(): ReconData;
getCallStatus(): ReconCallStatus;
}
export function createReconCollectorServer(): ReconCollectorServer {
const state: {
executive_summary?: ExecutiveSummaryInput;
technology_stack?: TechnologyStackInput;
authentication?: AuthenticationInput;
input_vectors?: InputVectorsInput;
network_map?: NetworkMapInput;
role_architecture?: RoleArchitectureInput;
authz_candidates?: AuthzCandidatesInput;
injection_sources?: InjectionSourcesInput;
} = {};
export function createReconCollector(): ReconCollector {
const state: ReconState = {};
const endpoints: Endpoint[] = [];
const seenEndpointKeys = new Set<string>();
let addEndpointsCalls = 0;
function alreadyCalled(toolName: ReconOneShotToolName): ToolResult {
function alreadyCalled(toolName: ReconOneShotToolName) {
return errorResult(
`${toolName} has already been called. Each set_* tool may only be called once per run.`,
'DuplicateError',
@@ -671,9 +667,9 @@ export function createReconCollectorServer(): ReconCollectorServer {
'user-facing components. Call exactly once before terminating. Becomes Section 1 of the rendered ' +
'deliverable. Duplicate calls are rejected.',
parameters: ExecutiveSummaryInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.executive_summary) return alreadyCalled('set_executive_summary');
state.executive_summary = input;
state.executive_summary = cleanInput(ExecutiveSummaryInputSchema, input);
return successResult({ set: 'set_executive_summary' });
},
});
@@ -685,9 +681,9 @@ export function createReconCollectorServer(): ReconCollectorServer {
'Record the technology and service map: frontend, backend, and infrastructure. Call exactly once ' +
'before terminating. Becomes Section 2 of the rendered deliverable. Duplicate calls are rejected.',
parameters: TechnologyStackInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.technology_stack) return alreadyCalled('set_technology_stack');
state.technology_stack = input;
state.technology_stack = cleanInput(TechnologyStackInputSchema, input);
return successResult({ set: 'set_technology_stack' });
},
});
@@ -702,9 +698,9 @@ export function createReconCollectorServer(): ReconCollectorServer {
'role_switching_impersonation.applicable=false (with the other fields null) if no such features ' +
'exist. Duplicate calls are rejected.',
parameters: AuthenticationInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.authentication) return alreadyCalled('set_authentication');
state.authentication = input;
state.authentication = cleanInput(AuthenticationInputSchema, input);
return successResult({ set: 'set_authentication' });
},
});
@@ -720,7 +716,7 @@ export function createReconCollectorServer(): ReconCollectorServer {
'Section 4 of the rendered deliverable and drives vuln-authz / vuln-injection todos downstream. ' +
'The renderer sorts by (path, method) before rendering, so emission order does not affect output.',
parameters: AddEndpointsInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
addEndpointsCalls += 1;
const added: string[] = [];
const skipped: string[] = [];
@@ -731,7 +727,7 @@ export function createReconCollectorServer(): ReconCollectorServer {
continue;
}
seenEndpointKeys.add(key);
endpoints.push(ep);
endpoints.push(cleanInput(EndpointSchema, ep));
added.push(key);
}
return successResult({
@@ -751,9 +747,9 @@ export function createReconCollectorServer(): ReconCollectorServer {
'and cookie values. Call exactly once before terminating. Becomes Section 5 of the rendered ' +
'deliverable. Drives downstream vulnerability analysis. Duplicate calls are rejected.',
parameters: InputVectorsInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.input_vectors) return alreadyCalled('set_input_vectors');
state.input_vectors = input;
state.input_vectors = cleanInput(InputVectorsInputSchema, input);
return successResult({ set: 'set_input_vectors' });
},
});
@@ -767,9 +763,9 @@ export function createReconCollectorServer(): ReconCollectorServer {
'(Guards Directory) of the rendered deliverable. The renderer splits the entities array into ' +
'the 6.1 and 6.2 tables and sorts each array deterministically. Duplicate calls are rejected.',
parameters: NetworkMapInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.network_map) return alreadyCalled('set_network_map');
state.network_map = input;
state.network_map = cleanInput(NetworkMapInputSchema, input);
return successResult({ set: 'set_network_map' });
},
});
@@ -783,9 +779,9 @@ export function createReconCollectorServer(): ReconCollectorServer {
'7.3 (Role Entry Points), and 7.4 (Role-to-Code Mapping) of the rendered deliverable. The renderer ' +
'splits the roles array into the per-section tables. Duplicate calls are rejected.',
parameters: RoleArchitectureInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.role_architecture) return alreadyCalled('set_role_architecture');
state.role_architecture = input;
state.role_architecture = cleanInput(RoleArchitectureInputSchema, input);
return successResult({ set: 'set_role_architecture' });
},
});
@@ -800,9 +796,9 @@ export function createReconCollectorServer(): ReconCollectorServer {
'sub-arrays in horizontal → vertical → context order, which vuln-authz reads as its todo list. ' +
'Duplicate calls are rejected.',
parameters: AuthzCandidatesInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.authz_candidates) return alreadyCalled('set_authz_candidates');
state.authz_candidates = input;
state.authz_candidates = cleanInput(AuthzCandidatesInputSchema, input);
return successResult({ set: 'set_authz_candidates' });
},
});
@@ -817,25 +813,13 @@ export function createReconCollectorServer(): ReconCollectorServer {
'of this kind"). Becomes Section 9 of the rendered deliverable. Drives the vuln-injection agent\'s ' +
'todos downstream. Duplicate calls are rejected.',
parameters: InjectionSourcesInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.injection_sources) return alreadyCalled('set_injection_sources');
state.injection_sources = input;
state.injection_sources = cleanInput(InjectionSourcesInputSchema, input);
return successResult({ set: 'set_injection_sources' });
},
});
const tools: ToolDefinition[] = [
setExecutiveSummary,
setTechnologyStack,
setAuthentication,
addEndpoints,
setInputVectors,
setNetworkMap,
setRoleArchitecture,
setAuthzCandidates,
setInjectionSources,
];
function statusOf<K extends ReconOneShotToolName>(key: K): ReconToolStatus {
const flagMap: Record<ReconOneShotToolName, unknown> = {
set_executive_summary: state.executive_summary,
@@ -851,7 +835,17 @@ export function createReconCollectorServer(): ReconCollectorServer {
}
return {
tools,
tools: [
setExecutiveSummary,
setTechnologyStack,
setAuthentication,
addEndpoints,
setInputVectors,
setNetworkMap,
setRoleArchitecture,
setAuthzCandidates,
setInjectionSources,
],
getAll: (): ReconData => ({
...(state.executive_summary && { executive_summary: state.executive_summary }),
...(state.technology_stack && { technology_stack: state.technology_stack }),
apps/worker/src/collectors/schema.ts
+29
View File
@@ -0,0 +1,29 @@
// Copyright (C) 2025 Keygraph, Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License version 3
// as published by the Free Software Foundation.
import { type Static, type TSchema, type TSchemaOptions, Type } from 'typebox';
import { Value } from 'typebox/value';
/**
* String-literal enum schema whose `Static` resolves to the exact value union.
*
* Mapping `Type.Literal` over an array loses tuple typing (`Static` widens to
* `never`), so enums are authored as a JSON-Schema `{ type: 'string', enum }`
* via `Type.Unsafe` — the same shape the previous Zod `z.enum` schemas produced,
* and validated at runtime by pi's TypeBox checker.
*/
export function stringEnum<const T extends readonly string[]>(values: T, options: TSchemaOptions = {}) {
return Type.Unsafe<T[number]>({ ...options, type: 'string', enum: [...values] });
}
/**
* Strips keys not declared on `schema` before storage, matching the previous
* Zod bridge's `safeParse()` behavior (Zod objects silently strip unknown
* keys by default; TypeBox schemas accept them unless explicitly cleaned).
*/
export function cleanInput<T extends TSchema>(schema: T, input: Static<T>): Static<T> {
return Value.Clean(schema, structuredClone(input)) as Static<T>;
}
apps/worker/src/mcp-server/vuln-collector.ts → apps/worker/src/collectors/vuln-collector.ts
+31 -50
View File
@@ -26,7 +26,8 @@
*/
import { defineTool, type ToolDefinition } from '@earendil-works/pi-coding-agent';
import { type Static, Type } from 'typebox';
import { type Static, type TObject, Type } from 'typebox';
import { cleanInput } from './schema.js';
// ============================================================================
// CLASS DISCRIMINATOR
@@ -35,12 +36,6 @@ import { type Static, Type } from 'typebox';
export const VULN_CLASSES = ['injection', 'xss', 'auth', 'ssrf', 'authz'] as const;
export type VulnClass = (typeof VULN_CLASSES)[number];
// Classes whose deliverables carry a Section 5 (blind spots). The auth and ssrf
// analyses have no blind-spots section, so the set_blind_spots tool is withheld
// from those agents and the renderer omits the section. Single source of truth
// for both the tool registration and the rendering gate.
export const BLIND_SPOTS_CLASSES: ReadonlySet<VulnClass> = new Set<VulnClass>(['injection', 'xss', 'authz']);
// ============================================================================
// SHARED SCHEMAS — set_findings_summary, set_safe_vectors, set_blind_spots
// ============================================================================
@@ -277,13 +272,13 @@ const AuthzStrategicIntelSchema = Type.Object({
}),
});
const STRATEGIC_INTEL_SCHEMAS = {
export const STRATEGIC_INTEL_SCHEMAS: Record<VulnClass, TObject> = {
injection: InjectionStrategicIntelSchema,
xss: XssStrategicIntelSchema,
auth: AuthStrategicIntelSchema,
ssrf: SsrfStrategicIntelSchema,
authz: AuthzStrategicIntelSchema,
} as const;
};
// ============================================================================
// EXPORTED TYPES
@@ -335,48 +330,42 @@ export type VulnCallStatus = Readonly<Record<VulnToolName, VulnToolStatus>>;
// RESPONSE HELPERS
// ============================================================================
interface ToolResult {
[x: string]: unknown;
content: Array<{ type: 'text'; text: string }>;
details: Record<string, unknown>;
isError: boolean;
}
function createToolResult(response: { status: string; [key: string]: unknown }): ToolResult {
function toolResult(payload: Record<string, unknown>) {
return {
content: [{ type: 'text' as const, text: JSON.stringify(response, null, 2) }],
details: {},
isError: response.status === 'error',
content: [{ type: 'text' as const, text: JSON.stringify(payload, null, 2) }],
details: undefined,
};
}
function successResult(data: Record<string, unknown>): ToolResult {
return createToolResult({ status: 'success', ...data });
function successResult(data: Record<string, unknown>) {
return toolResult({ status: 'success', ...data });
}
function errorResult(message: string, errorType = 'ValidationError', retryable = true): ToolResult {
return createToolResult({ status: 'error', message, errorType, retryable });
function errorResult(message: string, errorType = 'ValidationError', retryable = true) {
return toolResult({ status: 'error', message, errorType, retryable });
}
// ============================================================================
// COLLECTOR FACTORY
// ============================================================================
export interface VulnCollectorServer {
interface VulnState {
findings_summary?: FindingsSummaryInput;
strategic_intelligence?: StrategicIntelligenceInput;
safe_vectors?: SafeVectorsInput;
blind_spots?: BlindSpotsInput;
}
export interface VulnCollector {
tools: ToolDefinition[];
getAll(): VulnCollectorData;
getCallStatus(): VulnCallStatus;
}
export function createVulnCollector(vulnClass: VulnClass): VulnCollectorServer {
const state: {
findings_summary?: FindingsSummaryInput;
strategic_intelligence?: StrategicIntelligenceInput;
safe_vectors?: SafeVectorsInput;
blind_spots?: BlindSpotsInput;
} = {};
export function createVulnCollector(vulnClass: VulnClass): VulnCollector {
const state: VulnState = {};
function alreadyCalled(toolName: VulnToolName): ToolResult {
function alreadyCalled(toolName: VulnToolName) {
return errorResult(
`${toolName} has already been called. Each tool may only be called once per run.`,
'DuplicateError',
@@ -395,9 +384,9 @@ export function createVulnCollector(vulnClass: VulnClass): VulnCollectorServer {
'patterns array is acceptable (renders as "No dominant patterns identified") but key_outcome ' +
'is always required.',
parameters: FindingsSummaryInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.findings_summary) return alreadyCalled('set_findings_summary');
state.findings_summary = input;
state.findings_summary = cleanInput(FindingsSummaryInputSchema, input);
return successResult({ set: 'set_findings_summary' });
},
});
@@ -413,9 +402,9 @@ export function createVulnCollector(vulnClass: VulnClass): VulnCollectorServer {
'Required. Duplicate calls return "already called" and are no-ops. Write "Not applicable" as ' +
'the field value when a sub-field does not apply to this run (rather than omitting).',
parameters: intelSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.strategic_intelligence) return alreadyCalled('set_strategic_intelligence');
state.strategic_intelligence = input as unknown as StrategicIntelligenceInput;
state.strategic_intelligence = cleanInput(intelSchema, input) as unknown as StrategicIntelligenceInput;
return successResult({ set: 'set_strategic_intelligence' });
},
});
@@ -431,9 +420,9 @@ export function createVulnCollector(vulnClass: VulnClass): VulnCollectorServer {
'(subject, location) before rendering, so emission order does not affect output. Duplicate ' +
'calls return "already called" and are no-ops.',
parameters: SafeVectorsInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.safe_vectors) return alreadyCalled('set_safe_vectors');
state.safe_vectors = input;
state.safe_vectors = cleanInput(SafeVectorsInputSchema, input);
return successResult({ set: 'set_safe_vectors', count: input.vectors.length });
},
});
@@ -448,21 +437,13 @@ export function createVulnCollector(vulnClass: VulnClass): VulnCollectorServer {
'either documented gaps or an explicit "no gaps" signal). Duplicate calls return "already ' +
'called" and are no-ops.',
parameters: BlindSpotsInputSchema,
execute: async (_toolCallId, input): Promise<ToolResult> => {
async execute(_toolCallId, input) {
if (state.blind_spots) return alreadyCalled('set_blind_spots');
state.blind_spots = input;
state.blind_spots = cleanInput(BlindSpotsInputSchema, input);
return successResult({ set: 'set_blind_spots', count: input.items.length });
},
});
// set_blind_spots is withheld from classes without a Section 5 (auth, ssrf).
const tools = [
setFindingsSummary,
setStrategicIntelligence,
setSafeVectors,
...(BLIND_SPOTS_CLASSES.has(vulnClass) ? [setBlindSpots] : []),
];
function statusOf<K extends VulnToolName>(key: K): VulnToolStatus {
const flagMap: Record<VulnToolName, unknown> = {
set_findings_summary: state.findings_summary,
@@ -474,7 +455,7 @@ export function createVulnCollector(vulnClass: VulnClass): VulnCollectorServer {
}
return {
tools: tools as ToolDefinition[],
tools: [setFindingsSummary, setStrategicIntelligence, setSafeVectors, setBlindSpots],
getAll: (): VulnCollectorData => ({
...(state.findings_summary && { findings_summary: state.findings_summary }),
...(state.strategic_intelligence && { strategic_intelligence: state.strategic_intelligence }),
apps/worker/src/paths.ts
-3
View File
@@ -1,7 +1,6 @@
/** Centralized path constants for the worker package */
import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
/** Worker package root (apps/worker/) resolved from compiled dist/ files */
@@ -10,8 +9,6 @@ const WORKER_ROOT = path.resolve(import.meta.dirname, '..');
export const PROMPTS_DIR = path.join(WORKER_ROOT, 'prompts');
export const CONFIGS_DIR = path.join(WORKER_ROOT, 'configs');
export const PLAYWRIGHT_SKILL_DIR = path.join(os.homedir(), '.claude', 'skills', 'playwright-cli');
/** Compiled pi extension dir that enforces bounded `bash` timeouts (resolved from dist/) */
export const BASH_TIMEOUT_EXTENSION_DIR = path.join(import.meta.dirname, 'ai', 'extensions', 'bash-timeout');
apps/worker/src/scripts/generate-totp.ts
+7 -4
View File
@@ -9,8 +9,7 @@
/**
* generate-totp CLI
*
* Generates 6-digit TOTP codes for authentication.
* Replaces the MCP generate_totp tool.
* Generates a TOTP code for the target's MFA.
* Based on RFC 6238 (TOTP) and RFC 4226 (HOTP).
*
* Usage:
@@ -129,8 +128,12 @@ function main(): void {
process.exit(1);
}
// Strip base32 padding ('=') and whitespace so grouped/padded secrets
// (e.g. "JBSW Y3DP" or "...PXP=") pass validation instead of being rejected.
const normalizedSecret = secret.replace(/[=\s]/g, '');
const base32Regex = /^[A-Z2-7]+$/i;
if (!base32Regex.test(secret)) {
if (!base32Regex.test(normalizedSecret)) {
console.log(
JSON.stringify({
status: 'error',
@@ -142,7 +145,7 @@ function main(): void {
}
try {
const totpCode = generateTOTP(secret);
const totpCode = generateTOTP(normalizedSecret);
const expiresIn = 30 - (Math.floor(Date.now() / 1000) % 30);
console.log(
apps/worker/src/services/agent-execution.ts
+127 -42
View File
@@ -23,7 +23,7 @@
*/
import { fs, path } from 'zx';
import { type PiPromptResult, runPiPrompt, validateAgentOutput } from '../ai/pi-executor.js';
import { type PiPromptResult, runPiPrompt, validateAgentOutput } from '../ai/pi/pi-executor.js';
import { createQueueSubmitTool, getQueueFilename } from '../ai/queue-schemas.js';
import type { AuditSession } from '../audit/index.js';
import { authStateFile } from '../audit/utils.js';
@@ -35,9 +35,10 @@ import { ErrorCode, type PentestErrorType } from '../types/errors.js';
import type { AgentMetrics } from '../types/metrics.js';
import { err, isErr, ok, type Result } from '../types/result.js';
import { isSpendingCapBehavior } from '../utils/billing-detection.js';
import { getAgentGitPaths } from './agent-git-paths.js';
import type { ConfigLoaderService } from './config-loader.js';
import { PentestError } from './error-handling.js';
import { commitGitSuccess, createGitCheckpoint, getGitCommitHash, rollbackGitWorkspace } from './git-manager.js';
import { commitGitSuccess, createGitCheckpoint, rollbackGitWorkspace, withGitRepoLock } from './git-manager.js';
import { loadPrompt } from './prompt-manager.js';
/**
@@ -52,12 +53,11 @@ export interface AgentExecutionInput {
configYAML?: string | undefined;
pipelineTestingMode?: boolean | undefined;
attemptNumber: number;
apiKey?: string | undefined;
promptDir?: string | undefined;
providerConfig?: import('../types/config.js').ProviderConfig | undefined;
customTools?: import('@earendil-works/pi-coding-agent').ToolDefinition[];
// Renders the deliverable to disk; invoked after validation, before the success commit.
writeDeliverable?: (deliverablesPath: string) => Promise<void>;
cancellationSignal?: AbortSignal | undefined;
}
interface FailAgentOpts {
@@ -71,6 +71,48 @@ interface FailAgentOpts {
context: Record<string, unknown>;
}
function errorCodeFromResult(result: PiPromptResult): ErrorCode {
if (result.errorType && Object.values(ErrorCode).includes(result.errorType as ErrorCode)) {
return result.errorType as ErrorCode;
}
return ErrorCode.AGENT_EXECUTION_FAILED;
}
function categoryForErrorCode(code: ErrorCode): PentestErrorType {
switch (code) {
case ErrorCode.SPENDING_CAP_REACHED:
case ErrorCode.INSUFFICIENT_CREDITS:
case ErrorCode.BILLING_ERROR:
case ErrorCode.API_RATE_LIMITED:
return 'billing';
case ErrorCode.GIT_CHECKPOINT_FAILED:
case ErrorCode.GIT_ROLLBACK_FAILED:
return 'filesystem';
case ErrorCode.PROMPT_LOAD_FAILED:
return 'prompt';
default:
return 'validation';
}
}
/** Wrap a failed git operation result into a PentestError attributed to the agent. */
function gitFailureForAgent(
agentName: AgentName,
operation: string,
error: Error | undefined,
code: ErrorCode = ErrorCode.GIT_CHECKPOINT_FAILED,
): PentestError {
const retryable = error instanceof PentestError ? error.retryable : true;
const message = error?.message ?? 'unknown git failure';
return new PentestError(
`Failed to ${operation} for ${agentName}: ${message}`,
'filesystem',
retryable,
{ agentName, originalError: message },
code,
);
}
/**
* Service for executing agents with full lifecycle management.
*
@@ -109,12 +151,12 @@ export class AgentExecutionService {
configYAML,
pipelineTestingMode = false,
attemptNumber,
apiKey,
promptDir,
providerConfig,
customTools,
writeDeliverable,
cancellationSignal,
} = input;
const gitPaths = getAgentGitPaths(agentName);
// 1. Load config (pre-parsed configData → raw YAML → file path)
const configResult = await this.configLoader.loadOptional(configPath, configData, configYAML);
@@ -148,9 +190,16 @@ export class AgentExecutionService {
);
}
// 3. Create git checkpoint before execution
// 3. Create git checkpoint before execution (scoped to this agent's paths)
try {
await createGitCheckpoint(deliverablesPath, agentName, attemptNumber, logger);
const checkpointResult = await createGitCheckpoint(deliverablesPath, agentName, attemptNumber, logger, gitPaths);
if (!checkpointResult.success) {
const code =
checkpointResult.error instanceof PentestError && checkpointResult.error.code
? checkpointResult.error.code
: ErrorCode.GIT_CHECKPOINT_FAILED;
return err(gitFailureForAgent(agentName, 'create git checkpoint', checkpointResult.error, code));
}
} catch (error) {
const errorMessage = error instanceof Error ? error.message : String(error);
return err(
@@ -170,7 +219,6 @@ export class AgentExecutionService {
// 5. Execute agent. Vuln agents get a submit tool that captures the structured
// exploitation queue (pi has no JSON-schema output format).
const submitTool = createQueueSubmitTool(agentName, distributedConfig?.exploit ?? true);
const callerTools = [...(customTools ?? []), ...(submitTool ? [submitTool.tool] : [])];
const result: PiPromptResult = await runPiPrompt(
prompt,
repoPath,
@@ -180,10 +228,10 @@ export class AgentExecutionService {
auditSession,
logger,
AGENTS[agentName].modelTier,
callerTools,
apiKey,
customTools,
path.relative(repoPath, deliverablesPath),
providerConfig,
cancellationSignal,
submitTool,
);
// 6. Spending cap check - defense-in-depth
@@ -205,55 +253,75 @@ export class AgentExecutionService {
// 7. Handle execution failure
if (!result.success) {
const errorCode = errorCodeFromResult(result);
return this.failAgent(agentName, deliverablesPath, auditSession, logger, {
attemptNumber,
result,
rollbackReason: 'execution failure',
errorMessage: result.error || 'Agent execution failed',
errorCode: ErrorCode.AGENT_EXECUTION_FAILED,
category: 'validation',
errorCode,
category: categoryForErrorCode(errorCode),
retryable: result.retryable ?? true,
context: { agentName, originalError: result.error },
});
}
// 8. Write structured output to disk (vuln agents only) from the submit-tool capture
const queueFilename = getQueueFilename(agentName);
if (submitTool && queueFilename) {
const captured = submitTool.getCaptured();
if (captured !== undefined) {
result.structuredOutput = captured; // carry for the validation gate below
// 8-11. Write structured output, validate, render, and commit under one repo lock so
// the write→validate→commit sequence is atomic against concurrent sibling agents.
let commitHash: string | undefined;
const finalizationError = await withGitRepoLock(async (): Promise<PentestError | null> => {
// 8. Write structured output to disk (vuln agents only) from the executor's capture
const queueFilename = getQueueFilename(agentName);
if (submitTool && queueFilename && result.structuredOutput !== undefined) {
await fs.ensureDir(deliverablesPath);
const queuePath = path.join(deliverablesPath, queueFilename);
await fs.writeFile(queuePath, JSON.stringify(captured, null, 2), 'utf8');
await fs.writeFile(queuePath, JSON.stringify(result.structuredOutput, null, 2), 'utf8');
logger.info(`Wrote structured output queue to ${queueFilename}`);
}
}
// 9. Validate output
const validationPassed = await validateAgentOutput(result, agentName, deliverablesPath, logger);
if (!validationPassed) {
// 9. Validate output
const validationPassed = await validateAgentOutput(result, agentName, deliverablesPath, logger);
if (!validationPassed) {
return new PentestError(
`Agent ${agentName} failed output validation`,
'validation',
true,
{ agentName, deliverableFilename: AGENTS[agentName].deliverableFilename },
ErrorCode.OUTPUT_VALIDATION_FAILED,
);
}
// 10. Render the deliverable to disk so the success commit below stages it
if (writeDeliverable) {
await writeDeliverable(deliverablesPath);
}
// 11. Success - commit deliverables (scoped) and capture the checkpoint hash
const commitResult = await commitGitSuccess(deliverablesPath, agentName, logger, gitPaths);
if (!commitResult.success) {
return gitFailureForAgent(agentName, 'commit successful results', commitResult.error);
}
commitHash = commitResult.commitHash;
return null;
});
if (finalizationError) {
const rollbackReason =
finalizationError.code === ErrorCode.OUTPUT_VALIDATION_FAILED
? 'validation failure'
: 'post-processing failure';
return this.failAgent(agentName, deliverablesPath, auditSession, logger, {
attemptNumber,
result,
rollbackReason: 'validation failure',
errorMessage: `Agent ${agentName} failed output validation`,
errorCode: ErrorCode.OUTPUT_VALIDATION_FAILED,
category: 'validation',
retryable: true,
context: { agentName, deliverableFilename: AGENTS[agentName].deliverableFilename },
rollbackReason,
errorMessage: finalizationError.message,
errorCode: finalizationError.code ?? ErrorCode.AGENT_EXECUTION_FAILED,
category: finalizationError.type,
retryable: finalizationError.retryable,
context: { agentName, ...finalizationError.context },
});
}
// 10. Render the deliverable to disk so the success commit below stages it
if (writeDeliverable) {
await writeDeliverable(deliverablesPath);
}
// 11. Success - commit deliverables, then capture checkpoint hash
await commitGitSuccess(deliverablesPath, agentName, logger);
const commitHash = await getGitCommitHash(deliverablesPath);
const endResult: AgentEndResult = {
attemptNumber,
duration_ms: result.duration,
@@ -274,7 +342,12 @@ export class AgentExecutionService {
logger: ActivityLogger,
opts: FailAgentOpts,
): Promise<Result<AgentEndResult, PentestError>> {
await rollbackGitWorkspace(deliverablesPath, opts.rollbackReason, logger);
const rollbackResult = await rollbackGitWorkspace(
deliverablesPath,
opts.rollbackReason,
logger,
getAgentGitPaths(agentName),
);
const endResult: AgentEndResult = {
attemptNumber: opts.attemptNumber,
@@ -286,7 +359,19 @@ export class AgentExecutionService {
};
await auditSession.endAgent(agentName, endResult);
return err(new PentestError(opts.errorMessage, opts.category, opts.retryable, opts.context, opts.errorCode));
const context = rollbackResult.success
? opts.context
: {
...opts.context,
rollbackFailed: true,
rollbackError: rollbackResult.error?.message ?? 'unknown rollback failure',
rollbackErrorCode:
rollbackResult.error instanceof PentestError
? (rollbackResult.error.code ?? ErrorCode.GIT_ROLLBACK_FAILED)
: ErrorCode.GIT_ROLLBACK_FAILED,
};
return err(new PentestError(opts.errorMessage, opts.category, opts.retryable, context, opts.errorCode));
}
/**
apps/worker/src/services/agent-git-paths.ts
+31
View File
@@ -0,0 +1,31 @@
// 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.
/**
* Per-agent git path scoping.
*
* Under parallel agent execution the deliverables git is shared, so each agent's
* checkpoint/commit/rollback must be limited to the files that agent actually
* writes. This resolves those paths from the agent's deliverable filename plus
* its structured exploitation queue (vuln agents only).
*/
import { getQueueFilename } from '../ai/queue-schemas.js';
import { AGENTS } from '../session-manager.js';
import type { AgentName } from '../types/agents.js';
/**
* Deliverable files an agent writes into the deliverables directory. Used to
* scope git operations so one agent never touches a sibling agent's output.
*/
export function getAgentGitPaths(agentName: AgentName): string[] {
const paths = [AGENTS[agentName].deliverableFilename];
const queueFilename = getQueueFilename(agentName);
if (queueFilename) {
paths.push(queueFilename);
}
return [...new Set(paths)];
}
apps/worker/src/services/error-handling.ts
+24 -1
View File
@@ -117,8 +117,10 @@ function classifyByErrorCode(code: ErrorCode, retryableFromError: boolean): { ty
case ErrorCode.PROMPT_LOAD_FAILED:
return { type: 'ConfigurationError', retryable: false };
// Git errors - non-retryable (indicates workspace corruption)
case ErrorCode.GIT_CHECKPOINT_FAILED:
return { type: 'GitError', retryable: retryableFromError };
// Rollback errors leave the workspace state untrusted.
case ErrorCode.GIT_ROLLBACK_FAILED:
return { type: 'GitError', retryable: false };
@@ -196,6 +198,27 @@ export function classifyErrorForTemporal(error: unknown): { type: string; retrya
return { type: 'PermissionError', retryable: false };
}
// Out of memory - deterministic resource exhaustion, retrying won't help
if (message.includes('out of memory')) {
return { type: 'OutOfMemoryError', retryable: false };
}
// Invalid prompt - malformed/rejected prompt content won't fix itself on retry
if (message.includes('invalid prompt')) {
return { type: 'InvalidPromptError', retryable: false };
}
// Session limit reached - distinct from billing/rate-limit; needs manual intervention
if (message.includes('session limit reached')) {
return { type: 'SessionLimitError', retryable: false };
}
// Overloaded - provider's own error-type token is authoritative regardless of the
// HTTP status it arrives under (seen in production under 400, not just 529)
if (message.includes('overloaded_error') || message.includes('overloaded')) {
return { type: 'OverloadedError', retryable: true };
}
// === OUTPUT VALIDATION ERRORS (Retryable) ===
// Agent didn't produce expected deliverables - retry may succeed
// IMPORTANT: Must come BEFORE generic 'validation' check below
apps/worker/src/services/exploit-renderer.ts
+39 -4
View File
@@ -11,16 +11,26 @@
* covers all 5 exploitation agents (injection, xss, auth, ssrf, authz). The
* per-class deltas are limited to title and ID prefix; every section, label,
* and sort rule is class-agnostic. Section headers and bolded field labels
* give downstream report-executive — which reads prose with bolded labels —
* a consistent structure to parse, with a single canonical label per field
* across all classes.
* mirror the prescribed-Markdown skeleton from the existing exploit-*.txt
* prompts so downstream report-executive — which reads prose with bolded
* labels — sees the same structure it sees today.
*
* Field-label drift across the 5 prompts ("Evidence of Vulnerability" vs
* "Why We Believe This Is Vulnerable"; "Attempted Exploitation" vs
* "What We Tried") is canonicalized here to a single label per field across
* all classes.
*
* Sort order is owned by the renderer:
* - Successfully Exploited: severity desc (critical → low), then ID asc.
* - Potential / Validation Blocked: confidence desc (high → low), then ID asc.
*
* ## Unprocessed Vulnerabilities surfaces queue IDs the collector did not see —
* the v1 stand-in for required-call enforcement. The activity passes idToType
* (queue ID → vulnerability_type, built from queue.json) so each entry renders
* as `- {ID} ({vulnerability_type})`. Omitted when every queue ID was emitted.
*/
import type { AddExploitInput, VulnClass } from '../mcp-server/exploit-collector.js';
import type { AddExploitInput, VulnClass } from '../collectors/exploit-collector.js';
// ============================================================================
// PER-CLASS CONSTANTS
@@ -121,6 +131,7 @@ function renderBlockedFinding(entry: BlockedEntry): string {
lines.push('');
lines.push('**Summary:**');
lines.push(`- **Vulnerable location:** ${entry.vulnerable_location}`);
lines.push(`- **Overview:** ${entry.overview}`);
lines.push(`- **Current Blocker:** ${entry.current_blocker}`);
lines.push(`- **Potential Impact:** ${entry.potential_impact}`);
lines.push(`- **Confidence:** ${entry.confidence.toUpperCase()}`);
@@ -171,6 +182,22 @@ function renderBlockedSection(entries: readonly BlockedEntry[]): string {
return [heading, '', blocks.join('\n\n')].join('\n');
}
function renderUnprocessedSection(missingIds: readonly string[], idToType: ReadonlyMap<string, string>): string {
const heading = '## Unprocessed Vulnerabilities';
const sortedIds = [...missingIds].sort((a, b) => a.localeCompare(b));
const lines = sortedIds.map((id) => {
const type = idToType.get(id);
return type ? `- ${id} (${type})` : `- ${id}`;
});
return [
heading,
'',
'The following queue vulnerabilities did not receive a definitive verdict during this run:',
'',
lines.join('\n'),
].join('\n');
}
// ============================================================================
// PUBLIC ENTRY POINT
// ============================================================================
@@ -190,7 +217,15 @@ export function renderExploitDeliverable(
const exploited = state.filter((e): e is ExploitedEntry => e.status === 'exploited');
const blocked = state.filter((e): e is BlockedEntry => e.status === 'blocked');
const emittedIds = new Set(state.map((e) => e.vulnerability_id));
const missingIds = [...idToType.keys()].filter((id) => !emittedIds.has(id));
const sections: string[] = [title, '', renderExploitedSection(exploited), '', renderBlockedSection(blocked)];
if (missingIds.length > 0) {
sections.push('');
sections.push(renderUnprocessedSection(missingIds, idToType));
}
return `${sections.join('\n').trimEnd()}\n`;
}
apps/worker/src/services/exploitation-checker.ts
+7 -16
View File
@@ -34,7 +34,7 @@ export class ExploitationCheckerService {
* @param repoPath - Path to the repository containing deliverables
* @param logger - ActivityLogger for structured logging
* @returns ExploitationDecision indicating whether to exploit
* @throws PentestError if validation fails and is retryable
* @throws PentestError if queue validation fails (retryable or not)
*/
async checkQueue(vulnType: VulnType, repoPath: string, logger: ActivityLogger): Promise<ExploitationDecision> {
const result = await validateQueueSafe(vulnType, repoPath);
@@ -47,21 +47,12 @@ export class ExploitationCheckerService {
return decision;
}
// Validation failed - check if we should retry or skip
// Validation failed. Throw in every case so the error reaches the workflow's
// per-class handler. Laundering a failure into a "no vulnerabilities" decision
// would render an un-assessed class as clean. Retryable vs non-retryable is
// decided downstream at the activity boundary.
const error = result.error;
if (error.retryable) {
// Re-throw retryable errors so caller can handle retry
logger.warn(`${vulnType}: ${error.message} (retryable)`);
throw error;
}
// Non-retryable error - skip exploitation gracefully
logger.warn(`${vulnType}: ${error.message}, skipping exploitation`);
return {
shouldExploit: false,
shouldRetry: false,
vulnerabilityCount: 0,
vulnType,
};
logger.warn(`${vulnType}: ${error.message}${error.retryable ? ' (retryable)' : ' (non-retryable, failing class)'}`);
throw error;
}
}
apps/worker/src/services/git-manager.ts
+177 -81
View File
@@ -4,6 +4,7 @@
// it under the terms of the GNU Affero General Public License version 3
// as published by the Free Software Foundation.
import { AsyncLocalStorage } from 'node:async_hooks';
import { $ } from 'zx';
import type { ActivityLogger } from '../types/activity-logger.js';
import { ErrorCode } from '../types/errors.js';
@@ -25,18 +26,72 @@ export async function isGitRepository(dir: string): Promise<boolean> {
interface GitOperationResult {
success: boolean;
hadChanges?: boolean;
changes?: string[];
commitHash?: string;
error?: Error;
}
/**
* Get list of changed files from git status --porcelain output
* Get list of changed files from git status --porcelain -z output.
* When paths is provided, the status query is scoped to those paths.
*/
async function getChangedFiles(sourceDir: string, operationDescription: string): Promise<string[]> {
const status = await executeGitCommandWithRetry(['git', 'status', '--porcelain'], sourceDir, operationDescription);
return status.stdout
.trim()
.split('\n')
.filter((line) => line.length > 0);
async function getChangedFiles(
sourceDir: string,
operationDescription: string,
paths?: readonly string[],
): Promise<string[]> {
const args = ['git', 'status', '--porcelain', '-z'];
if (paths && paths.length > 0) {
args.push('--', ...paths);
}
const status = await executeGitCommandWithRetry(args, sourceDir, operationDescription);
return parsePorcelainZ(status.stdout);
}
/**
* Parse `git status --porcelain -z` output.
*
* -z uses NUL separators and raw (unquoted) byte paths, sidestepping the
* fragile whitespace/quote handling of the default porcelain v1 format.
* Each entry is `XY<space>PATH\0`; renames/copies (X = 'R' or 'C') emit an
* additional `ORIG\0` token immediately after the entry, which we skip.
*/
export function parsePorcelainZ(raw: string): string[] {
if (raw.length === 0) {
return [];
}
const tokens = raw.split('\0');
const entries: string[] = [];
for (let i = 0; i < tokens.length; i++) {
const tok = tokens[i];
if (!tok || tok.length < 4) {
continue;
}
entries.push(tok);
const x = tok[0];
if (x === 'R' || x === 'C') {
i++;
}
}
return entries;
}
function changedPathFromStatus(entry: string): string {
return entry.slice(3);
}
async function stageChanges(sourceDir: string, description: string, paths?: readonly string[]): Promise<string[]> {
const changes = await getChangedFiles(sourceDir, description, paths);
if (paths && paths.length > 0) {
const changedPaths = [...new Set(changes.map(changedPathFromStatus).filter((p) => p.length > 0))];
if (changedPaths.length > 0) {
await executeGitCommandWithRetry(['git', 'add', '-A', '--', ...changedPaths], sourceDir, description);
}
return changes;
}
await executeGitCommandWithRetry(['git', 'add', '-A'], sourceDir, description);
return changes;
}
/**
@@ -102,6 +157,28 @@ class GitSemaphore {
const gitSemaphore = new GitSemaphore();
// Tracks whether the current async context already holds the repo lock, so a
// composite operation (e.g. status → add → commit) can call nested git helpers
// without re-acquiring the semaphore and deadlocking on itself.
const gitLockContext = new AsyncLocalStorage<boolean>();
/**
* Run an operation while holding the repo-wide git lock. Reentrant: a nested
* call inside an already-locked context runs immediately instead of blocking.
*/
export async function withGitRepoLock<T>(operation: () => Promise<T>): Promise<T> {
if (gitLockContext.getStore()) {
return operation();
}
await gitSemaphore.acquire();
try {
return await gitLockContext.run(true, operation);
} finally {
gitSemaphore.release();
}
}
const GIT_LOCK_ERROR_PATTERNS = [
'index.lock',
'unable to lock',
@@ -121,48 +198,49 @@ export async function executeGitCommandWithRetry(
description: string,
maxRetries: number = 5,
): Promise<{ stdout: string; stderr: string }> {
await gitSemaphore.acquire();
try {
for (let attempt = 1; attempt <= maxRetries; attempt++) {
try {
const [cmd, ...args] = commandArgs;
const result = await $`cd ${sourceDir} && ${cmd} ${args}`;
return result;
} catch (error) {
const errMsg = error instanceof Error ? error.message : String(error);
if (isGitLockError(errMsg) && attempt < maxRetries) {
const delay = 2 ** (attempt - 1) * 1000;
// executeGitCommandWithRetry is also called outside activity context
// (e.g., from resume logic), so we use console.warn as a fallback here
console.warn(
`Git lock conflict during ${description} (attempt ${attempt}/${maxRetries}). Retrying in ${delay}ms...`,
);
await new Promise((resolve) => setTimeout(resolve, delay));
continue;
}
throw error;
}
}
throw new PentestError(
`Git command failed after ${maxRetries} retries`,
'filesystem',
true, // Retryable - transient git lock issues
{ maxRetries, description },
ErrorCode.GIT_CHECKPOINT_FAILED,
);
} finally {
gitSemaphore.release();
if (!gitLockContext.getStore()) {
return withGitRepoLock(() => executeGitCommandWithRetry(commandArgs, sourceDir, description, maxRetries));
}
for (let attempt = 1; attempt <= maxRetries; attempt++) {
try {
const [cmd, ...args] = commandArgs;
const result = await $`cd ${sourceDir} && ${cmd} ${args}`;
return result;
} catch (error) {
const errMsg = error instanceof Error ? error.message : String(error);
if (isGitLockError(errMsg) && attempt < maxRetries) {
const delay = 2 ** (attempt - 1) * 1000;
// executeGitCommandWithRetry is also called outside activity context
// (e.g., from resume logic), so we use console.warn as a fallback here
console.warn(
`Git lock conflict during ${description} (attempt ${attempt}/${maxRetries}). Retrying in ${delay}ms...`,
);
await new Promise((resolve) => setTimeout(resolve, delay));
continue;
}
throw error;
}
}
throw new PentestError(
`Git command failed after ${maxRetries} retries`,
'filesystem',
true, // Retryable - transient git lock issues
{ maxRetries, description },
ErrorCode.GIT_CHECKPOINT_FAILED,
);
}
// Two-phase reset: hard reset (tracked files) + clean (untracked files)
// Two-phase reset: hard reset (tracked files) + clean (untracked files).
// When paths is provided, the untracked clean is scoped to those paths so a
// failing agent's rollback can't delete a concurrent sibling agent's scratch.
export async function rollbackGitWorkspace(
sourceDir: string,
reason: string = 'retry preparation',
logger: ActivityLogger,
paths?: readonly string[],
): Promise<GitOperationResult> {
// Skip git operations if not a git repository
if (!(await isGitRepository(sourceDir))) {
@@ -172,10 +250,13 @@ export async function rollbackGitWorkspace(
logger.info(`Rolling back workspace for ${reason}`);
try {
const changes = await getChangedFiles(sourceDir, 'status check for rollback');
await executeGitCommandWithRetry(['git', 'reset', '--hard', 'HEAD'], sourceDir, 'hard reset for rollback');
await executeGitCommandWithRetry(['git', 'clean', '-fd'], sourceDir, 'cleaning untracked files for rollback');
const changes = await withGitRepoLock(async () => {
const pendingChanges = await getChangedFiles(sourceDir, 'status check for rollback');
await executeGitCommandWithRetry(['git', 'reset', '--hard', 'HEAD'], sourceDir, 'hard reset for rollback');
const cleanArgs = paths && paths.length > 0 ? ['git', 'clean', '-fd', '--', ...paths] : ['git', 'clean', '-fd'];
await executeGitCommandWithRetry(cleanArgs, sourceDir, 'cleaning untracked files for rollback');
return pendingChanges;
});
logChangeSummary(
changes,
@@ -208,6 +289,7 @@ export async function createGitCheckpoint(
description: string,
attempt: number,
logger: ActivityLogger,
paths?: readonly string[],
): Promise<GitOperationResult> {
// Skip git operations if not a git repository
if (!(await isGitRepository(sourceDir))) {
@@ -217,33 +299,37 @@ export async function createGitCheckpoint(
logger.info(`Creating checkpoint for ${description} (attempt ${attempt})`);
try {
// 1. On retries, clean workspace to prevent pollution from previous attempt
if (attempt > 1) {
const cleanResult = await rollbackGitWorkspace(sourceDir, `${description} (retry cleanup)`, logger);
if (!cleanResult.success) {
logger.warn(`Workspace cleanup failed, continuing anyway: ${cleanResult.error?.message}`);
const result = await withGitRepoLock(async (): Promise<GitOperationResult> => {
// 1. On retries, clean workspace to prevent pollution from previous attempt
if (attempt > 1) {
const cleanResult = await rollbackGitWorkspace(sourceDir, `${description} (retry cleanup)`, logger, paths);
if (!cleanResult.success) {
return cleanResult;
}
}
// 2. Stage scoped changes and commit checkpoint
const changes = await stageChanges(sourceDir, 'staging changes', paths);
const hasChanges = changes.length > 0;
await executeGitCommandWithRetry(
['git', 'commit', '-m', `📍 Checkpoint: ${description} (attempt ${attempt})`, '--allow-empty'],
sourceDir,
'creating commit',
);
const commitHash = await getGitCommitHash(sourceDir);
return { success: true, hadChanges: hasChanges, changes, ...(commitHash && { commitHash }) };
});
if (result.success) {
if (result.hadChanges) {
logger.info('Checkpoint created with scoped changes staged');
} else {
logger.info('Empty checkpoint created (no scoped workspace changes)');
}
}
// 2. Detect existing changes
const changes = await getChangedFiles(sourceDir, 'status check');
const hasChanges = changes.length > 0;
// 3. Stage and commit checkpoint
await executeGitCommandWithRetry(['git', 'add', '-A'], sourceDir, 'staging changes');
await executeGitCommandWithRetry(
['git', 'commit', '-m', `📍 Checkpoint: ${description} (attempt ${attempt})`, '--allow-empty'],
sourceDir,
'creating commit',
);
// 4. Log result
if (hasChanges) {
logger.info('Checkpoint created with uncommitted changes staged');
} else {
logger.info('Empty checkpoint created (no workspace changes)');
}
return { success: true };
return result;
} catch (error) {
const result = toErrorResult(error);
logger.warn(`Checkpoint creation failed after retries: ${result.error?.message}`);
@@ -255,6 +341,7 @@ export async function commitGitSuccess(
sourceDir: string,
description: string,
logger: ActivityLogger,
paths?: readonly string[],
): Promise<GitOperationResult> {
// Skip git operations if not a git repository
if (!(await isGitRepository(sourceDir))) {
@@ -264,22 +351,31 @@ export async function commitGitSuccess(
logger.info(`Committing successful results for ${description}`);
try {
const changes = await getChangedFiles(sourceDir, 'status check for success commit');
const result = await withGitRepoLock(async (): Promise<GitOperationResult> => {
const changes = await stageChanges(sourceDir, 'staging changes for success commit', paths);
await executeGitCommandWithRetry(['git', 'add', '-A'], sourceDir, 'staging changes for success commit');
await executeGitCommandWithRetry(
['git', 'commit', '-m', `✅ ${description}: completed successfully`, '--allow-empty'],
sourceDir,
'creating success commit',
);
await executeGitCommandWithRetry(
['git', 'commit', '-m', `✅ ${description}: completed successfully`, '--allow-empty'],
sourceDir,
'creating success commit',
);
const commitHash = await getGitCommitHash(sourceDir);
return {
success: true,
hadChanges: changes.length > 0,
changes,
...(commitHash && { commitHash }),
};
});
logChangeSummary(
changes,
result.changes ?? [],
'Success commit created with {count} file changes:',
'Empty success commit created (agent made no file changes)',
logger,
);
return { success: true };
return result;
} catch (error) {
const result = toErrorResult(error);
logger.warn(`Success commit failed after retries: ${result.error?.message}`);
@@ -296,7 +392,7 @@ export async function getGitCommitHash(sourceDir: string): Promise<string | null
return null;
}
try {
const result = await $`cd ${sourceDir} && git rev-parse HEAD`;
const result = await executeGitCommandWithRetry(['git', 'rev-parse', 'HEAD'], sourceDir, 'read HEAD commit');
return result.stdout.trim();
} catch {
return null;
apps/worker/src/services/index.ts
+2 -2
View File
@@ -11,8 +11,8 @@
* Services are pure domain logic with no Temporal dependencies.
*/
export type { PiPromptResult } from '../ai/pi-executor.js';
export { runPiPrompt } from '../ai/pi-executor.js';
export type { PiPromptResult } from '../ai/pi/pi-executor.js';
export { runPiPrompt } from '../ai/pi/pi-executor.js';
export type { AgentExecutionInput } from './agent-execution.js';
export { AgentExecutionService } from './agent-execution.js';
export { ConfigLoaderService } from './config-loader.js';
apps/worker/src/services/pre-recon-renderer.ts
+8 -7
View File
@@ -7,14 +7,15 @@
/**
* Deterministic pre-recon collector → markdown renderer.
*
* Converts the typed payload bag harvested from the pre-recon-collector MCP
* server into the pre_recon_deliverable.md Markdown layout. No LLM in the
* Converts the typed payload bag harvested from the pre-recon-collector pi
* tool set into the Markdown layout that pre_recon_deliverable.md previously had
* when the agent emitted it directly via chunked Write/Edit. No LLM in the
* loop; section ordering, headings, and template are owned here.
*
* Any tool the agent skips becomes a `[Section X: not provided]` placeholder
* rather than an activity failure. Every section renderer accepts the
* corresponding collected field as possibly undefined and falls back to the
* placeholder helper when absent.
* Required-call enforcement is deferred for v1. Any tool the agent skips becomes
* a `[Section X: not provided]` placeholder rather than an activity failure.
* Every section renderer accepts the corresponding collected field as possibly
* undefined and falls back to the placeholder helper when absent.
*/
import type {
@@ -27,7 +28,7 @@ import type {
SinkRef,
SsrfSinksInput,
XssSinksInput,
} from '../mcp-server/pre-recon-collector.js';
} from '../collectors/pre-recon-collector.js';
// ============================================================================
// STATIC PROSE
apps/worker/src/services/preflight.ts
+5 -26
View File
@@ -353,29 +353,10 @@ async function probeCredentialsWithPi(
}
/** Validate credentials via a minimal pi session. */
async function validateCredentials(
logger: ActivityLogger,
apiKey?: string,
providerConfig?: import('../types/config.js').ProviderConfig,
): Promise<Result<void, PentestError>> {
// 0. If providerConfig is present, credentials are managed by the caller.
// The executor/provider layer owns providerConfig resolution — no env preflight needed.
if (providerConfig) {
logger.info(
`Provider config present (type: ${providerConfig.providerType || 'anthropic_api'}) — skipping env-based credential validation`,
);
return ok(undefined);
}
// 0b. If apiKey provided via config, set it in env for pi validation
// This avoids requiring process.env.ANTHROPIC_API_KEY when key is threaded via input
if (apiKey) {
process.env.ANTHROPIC_API_KEY = apiKey;
}
async function validateCredentials(logger: ActivityLogger): Promise<Result<void, PentestError>> {
// Resolve the active provider through the same precedence the executor uses, so
// preflight validates exactly the credentials the run will use (no drift).
const eff = resolveEffectiveProvider(apiKey);
const eff = resolveEffectiveProvider();
// 1. Bedrock mode — validate required AWS credentials are present (pi-ai owns the
// live AWS auth, so there is no cheap session probe here)
@@ -425,7 +406,7 @@ async function validateCredentials(
);
}
const usingApiKey = Boolean(apiKey ?? process.env.ANTHROPIC_API_KEY);
const usingApiKey = Boolean(process.env.ANTHROPIC_API_KEY);
const authType = usingApiKey ? 'API key' : 'OAuth token';
logger.info(`Validating ${authType} via pi...`);
const probe = await probeCredentialsWithPi(authType, eff.anthropicToken);
@@ -573,8 +554,6 @@ export async function runPreflightChecks(
configPath: string | undefined,
logger: ActivityLogger,
skipGitCheck?: boolean,
apiKey?: string,
providerConfig?: import('../types/config.js').ProviderConfig,
): Promise<Result<void, PentestError>> {
// 1. Repository check (free — filesystem only)
const repoResult = await validateRepo(repoPath, logger, skipGitCheck);
@@ -601,8 +580,8 @@ export async function runPreflightChecks(
}
}
// 4. Credential check (cheap — 1 pi round-trip, skipped when providerConfig present)
const credResult = await validateCredentials(logger, apiKey, providerConfig);
// 4. Credential check (cheap — 1 pi round-trip)
const credResult = await validateCredentials(logger);
if (!credResult.ok) {
return credResult;
}
apps/worker/src/services/prompt-manager.ts
+57 -34
View File
@@ -170,39 +170,46 @@ async function buildLoginInstructions(
if (authentication.credentials) {
if (authentication.credentials.username) {
userInstructions = userInstructions.replace(/\$username/g, authentication.credentials.username);
userInstructions = replaceLiteral(userInstructions, /\$username/g, authentication.credentials.username);
}
if (authentication.credentials.password) {
userInstructions = userInstructions.replace(/\$password/g, authentication.credentials.password);
userInstructions = replaceLiteral(userInstructions, /\$password/g, authentication.credentials.password);
}
if (authentication.credentials.totp_secret) {
userInstructions = userInstructions.replace(
userInstructions = replaceLiteral(
userInstructions,
/\$totp/g,
`generated TOTP code using secret "${authentication.credentials.totp_secret}"`,
);
}
if (authentication.credentials.email_login?.address) {
userInstructions = userInstructions.replace(/\$email_address/g, authentication.credentials.email_login.address);
userInstructions = replaceLiteral(
userInstructions,
/\$email_address/g,
authentication.credentials.email_login.address,
);
}
if (authentication.credentials.email_login?.password) {
userInstructions = userInstructions.replace(
userInstructions = replaceLiteral(
userInstructions,
/\$email_password/g,
authentication.credentials.email_login.password,
);
}
if (authentication.credentials.email_login?.totp_secret) {
userInstructions = userInstructions.replace(
userInstructions = replaceLiteral(
userInstructions,
/\$email_totp/g,
`generated TOTP code using secret "${authentication.credentials.email_login.totp_secret}"`,
);
}
}
loginInstructions = loginInstructions.replace(/{{user_instructions}}/g, userInstructions);
loginInstructions = replaceLiteral(loginInstructions, /{{user_instructions}}/g, userInstructions);
// 5. Replace TOTP secret placeholder if present in template
if (authentication.credentials?.totp_secret) {
loginInstructions = loginInstructions.replace(/{{totp_secret}}/g, authentication.credentials.totp_secret);
loginInstructions = replaceLiteral(loginInstructions, /{{totp_secret}}/g, authentication.credentials.totp_secret);
}
return loginInstructions;
@@ -242,11 +249,21 @@ async function processIncludes(content: string, baseDir: string): Promise<string
);
for (const replacement of replacements) {
content = content.replace(replacement.placeholder, replacement.content);
content = replaceLiteral(content, replacement.placeholder, replacement.content);
}
return content;
}
/**
* Replaces `pattern` with `replacement` treating the replacement as a literal
* string. Native `String.replace` interprets `$&`, `$1`, `$$` in the replacement
* as special patterns, which mangles credential and config values that legitimately
* contain `$`. The function form of `replace` bypasses that interpretation.
*/
function replaceLiteral(input: string, pattern: RegExp | string, replacement: string): string {
return input.replace(pattern, () => replacement);
}
function buildAuthContext(config: DistributedConfig | null): string {
if (!config?.authentication) {
return 'No authentication configured - unauthenticated testing only';
@@ -288,12 +305,18 @@ async function interpolateVariables(
});
}
let result = template
.replace(/{{WEB_URL}}/g, variables.webUrl)
.replace(/{{REPO_PATH}}/g, variables.repoPath)
.replace(/{{PLAYWRIGHT_SESSION}}/g, variables.PLAYWRIGHT_SESSION || 'agent1')
.replace(/{{AUTH_CONTEXT}}/g, buildAuthContext(config))
.replace(/{{DESCRIPTION}}/g, config?.description ? `Description: ${config.description}` : '');
// replaceLiteral is used for all value insertions so config values that
// contain `$&`/`$$`/`$1`/etc. aren't mangled as replacement patterns.
let result = template;
result = replaceLiteral(result, /{{WEB_URL}}/g, variables.webUrl);
result = replaceLiteral(result, /{{REPO_PATH}}/g, variables.repoPath);
result = replaceLiteral(result, /{{PLAYWRIGHT_SESSION}}/g, variables.PLAYWRIGHT_SESSION || 'agent1');
result = replaceLiteral(result, /{{AUTH_CONTEXT}}/g, buildAuthContext(config));
result = replaceLiteral(
result,
/{{DESCRIPTION}}/g,
config?.description ? `Description: ${config.description}` : '',
);
const avoidUrlRules = config?.avoid?.filter((r) => r.type !== 'code_path') ?? [];
const focusUrlRules = config?.focus?.filter((r) => r.type !== 'code_path') ?? [];
@@ -302,7 +325,8 @@ async function interpolateVariables(
} else {
const avoidStr = avoidUrlRules.length > 0 ? avoidUrlRules.map((r) => `- ${r.description}`).join('\n') : 'None';
const focusStr = focusUrlRules.length > 0 ? focusUrlRules.map((r) => `- ${r.description}`).join('\n') : 'None';
result = result.replace(/{{RULES_AVOID}}/g, avoidStr).replace(/{{RULES_FOCUS}}/g, focusStr);
result = replaceLiteral(result, /{{RULES_AVOID}}/g, avoidStr);
result = replaceLiteral(result, /{{RULES_FOCUS}}/g, focusStr);
}
const avoidCodeRules = (config?.avoid ?? []).filter((r) => r.type === 'code_path');
@@ -310,14 +334,13 @@ async function interpolateVariables(
if (avoidCodeRules.length === 0 && focusCodeRules.length === 0) {
result = result.replace(/<code_path_rules>[\s\S]*?<\/code_path_rules>\s*/g, '');
} else {
result = result
.replace(/{{CODE_RULES_AVOID}}/g, renderCodePathRules(config?.avoid ?? []))
.replace(/{{CODE_RULES_FOCUS}}/g, renderCodePathRules(config?.focus ?? []));
result = replaceLiteral(result, /{{CODE_RULES_AVOID}}/g, renderCodePathRules(config?.avoid ?? []));
result = replaceLiteral(result, /{{CODE_RULES_FOCUS}}/g, renderCodePathRules(config?.focus ?? []));
}
const roe = config?.rules_of_engagement?.trim() ?? '';
if (roe) {
result = result.replace(/{{RULES_OF_ENGAGEMENT}}/g, roe);
result = replaceLiteral(result, /{{RULES_OF_ENGAGEMENT}}/g, roe);
} else {
result = result.replace(/<rules_of_engagement>[\s\S]*?<\/rules_of_engagement>\s*/g, '');
}
@@ -325,35 +348,35 @@ async function interpolateVariables(
if (!config?.authentication) {
result = result.replace(/<shared_authenticated_session>[\s\S]*?<\/shared_authenticated_session>\s*/g, '');
} else {
result = result.replace(/{{AUTH_STATE_FILE}}/g, variables.AUTH_STATE_FILE);
result = replaceLiteral(result, /{{AUTH_STATE_FILE}}/g, variables.AUTH_STATE_FILE);
}
if (config?.authentication?.login_flow) {
const loginInstructions = await buildLoginInstructions(config.authentication, logger, promptsBaseDir);
result = result.replace(/{{LOGIN_INSTRUCTIONS}}/g, loginInstructions);
result = replaceLiteral(result, /{{LOGIN_INSTRUCTIONS}}/g, loginInstructions);
} else {
result = result.replace(/{{LOGIN_INSTRUCTIONS}}/g, '');
}
const vulnClasses = config?.vuln_classes ?? [];
result = result.replace(
result = replaceLiteral(
result,
/{{VULN_CLASSES_TESTED}}/g,
vulnClasses.length > 0 ? vulnClasses.join(', ') : 'injection, xss, auth, authz, ssrf',
);
result = result.replace(/{{VULN_SUMMARY_SUBSECTIONS}}/g, renderVulnSummarySubsections(vulnClasses));
result = replaceLiteral(result, /{{VULN_SUMMARY_SUBSECTIONS}}/g, renderVulnSummarySubsections(vulnClasses));
const exploitEnabled = config?.exploit ?? true;
result = result
.replace(/{{EXPLOITATION}}/g, exploitEnabled ? 'enabled' : 'disabled')
.replace(/{{REPORT_VULN_HEADING}}/g, exploitEnabled ? 'Exploitation Evidence' : 'Findings')
.replace(
/{{REPORT_VULN_SUBHEADING}}/g,
exploitEnabled ? 'Successfully Exploited Vulnerabilities' : 'Identified Vulnerabilities',
);
result = replaceLiteral(result, /{{EXPLOITATION}}/g, exploitEnabled ? 'enabled' : 'disabled');
result = replaceLiteral(result, /{{REPORT_VULN_HEADING}}/g, exploitEnabled ? 'Exploitation Evidence' : 'Findings');
result = replaceLiteral(
result,
/{{REPORT_VULN_SUBHEADING}}/g,
exploitEnabled ? 'Successfully Exploited Vulnerabilities' : 'Identified Vulnerabilities',
);
result = result
.replace(/{{REPORT_FILTERS_BLOCK}}/g, renderReportFiltersBlock(config?.report))
.replace(/{{REPORT_FILTER_RULES}}/g, renderReportFilterRules(config?.report));
result = replaceLiteral(result, /{{REPORT_FILTERS_BLOCK}}/g, renderReportFiltersBlock(config?.report));
result = replaceLiteral(result, /{{REPORT_FILTER_RULES}}/g, renderReportFilterRules(config?.report));
// Collapse runs of 3+ newlines (left behind by tag-strip and empty-fragment substitutions).
result = result.replace(/\n{3,}/g, '\n\n');
apps/worker/src/services/recon-renderer.ts
+9 -7
View File
@@ -7,13 +7,15 @@
/**
* Deterministic recon collector → markdown renderer.
*
* Converts the typed payload bag harvested from the recon-collector MCP server
* into the recon_deliverable.md Markdown layout. No LLM in the loop; section
* ordering, headings, sort, and the Section 0 boilerplate are owned here.
* Converts the typed payload bag harvested from the recon-collector tool
* into the Markdown layout that recon_deliverable.md previously had when the
* agent emitted it directly via chunked Write/Edit. No LLM in the loop;
* section ordering, headings, sort, and the Section 0 boilerplate are owned
* here.
*
* Any tool the agent skips becomes a `[Section X: not provided]` placeholder
* rather than an activity failure. Every section renderer accepts its input as
* optional.
* Required-call enforcement is deferred. Any tool the agent skips becomes a
* `[Section X: not provided]` placeholder rather than an activity failure.
* Every section renderer accepts its input as optional.
*/
import type {
@@ -36,7 +38,7 @@ import type {
SinkRef,
TechnologyStackInput,
VerticalCandidate,
} from '../mcp-server/recon-collector.js';
} from '../collectors/recon-collector.js';
type RoleSwitchingImpersonation = AuthenticationInput['role_switching_impersonation'];
type EntityZone = Entity['zone'];
apps/worker/src/services/validate-authentication.ts
+48 -36
View File
@@ -13,14 +13,15 @@
*/
import { readFile, rm } from 'node:fs/promises';
import { defineTool, type ToolDefinition } from '@earendil-works/pi-coding-agent';
import { defineTool } from '@earendil-works/pi-coding-agent';
import { Type } from 'typebox';
import { runPiPrompt } from '../ai/pi-executor.js';
import { runPiPrompt } from '../ai/pi/pi-executor.js';
import type { CapturedSubmitTool } from '../ai/submit-tool.js';
import type { AuditSession } from '../audit/index.js';
import { authStateFile } from '../audit/utils.js';
import type { ActivityLogger } from '../types/activity-logger.js';
import type { AgentEndResult } from '../types/audit.js';
import type { DistributedConfig, ProviderConfig } from '../types/config.js';
import type { DistributedConfig } from '../types/config.js';
import { ErrorCode } from '../types/errors.js';
import { err, ok, type Result } from '../types/result.js';
import { PentestError } from './error-handling.js';
@@ -40,30 +41,45 @@ interface AuthValidationVerdict {
}
/** Submit tool capturing the login verdict (pi has no JSON-schema output format). */
function createAuthSubmitTool(): { tool: ToolDefinition; getCaptured: () => AuthValidationVerdict | undefined } {
function createAuthSubmitTool(): CapturedSubmitTool {
let captured: AuthValidationVerdict | undefined;
const tool = defineTool({
name: 'submit_auth_result',
label: 'Submit Auth Result',
description: 'Report the login outcome. Call exactly once when the login attempt has concluded.',
parameters: Type.Object({
login_success: Type.Boolean(),
failure_point: Type.Optional(
Type.Union([Type.Literal('username_or_password'), Type.Literal('totp_secret'), Type.Literal('out_of_band')]),
),
failure_detail: Type.Optional(
Type.String({
description:
'Free-form 1-2 sentence diagnostic of what the page showed (error messages, page state) when login failed. Required when login_success is false. Mask any sensitive values.',
}),
),
return {
tool: defineTool({
name: 'submit_auth_result',
label: 'Submit Auth Result',
description: 'Report the login outcome. Call exactly once when the login attempt has concluded.',
promptSnippet: 'submit_auth_result: record the authentication validation verdict',
promptGuidelines: [
'You MUST call submit_auth_result exactly once as your final action.',
'Set login_success to true only after saving the authenticated browser session.',
],
parameters: Type.Object({
login_success: Type.Boolean(),
failure_point: Type.Optional(
Type.Union([Type.Literal('username_or_password'), Type.Literal('totp_secret'), Type.Literal('out_of_band')]),
),
failure_detail: Type.Optional(
Type.String({
maxLength: 250,
description:
'Free-form 1-2 sentence diagnostic of what the page showed (error messages, page state) when login failed. Required when login_success is false. Mask any sensitive values.',
}),
),
}),
execute: async (_toolCallId, params) => {
captured = params as AuthValidationVerdict;
return {
content: [{ type: 'text' as const, text: 'Auth result recorded.' }],
details: params,
terminate: true,
};
},
}),
execute: async (_toolCallId, params) => {
captured = params as AuthValidationVerdict;
return { content: [{ type: 'text' as const, text: 'Auth result recorded.' }], details: {} };
},
});
return { tool, getCaptured: () => captured };
getCaptured: () => captured,
directive:
'\n\nYou MUST call the submit_auth_result tool exactly once as your final action ' +
'to deliver the authentication verdict. Do not output JSON as text.',
};
}
const AGENT_NAME = 'validate-authentication';
@@ -75,11 +91,10 @@ export interface ValidateAuthInput {
readonly logger: ActivityLogger;
readonly auditSession: AuditSession;
readonly attemptNumber: number;
readonly apiKey?: string;
readonly providerConfig?: ProviderConfig;
readonly deliverablesSubdir?: string;
readonly promptDir?: string;
readonly pipelineTestingMode?: boolean;
readonly cancellationSignal?: AbortSignal;
}
export async function validateAuthentication(input: ValidateAuthInput): Promise<Result<void, PentestError>> {
@@ -90,11 +105,10 @@ export async function validateAuthentication(input: ValidateAuthInput): Promise<
logger,
auditSession,
attemptNumber,
apiKey,
providerConfig,
deliverablesSubdir,
promptDir,
pipelineTestingMode,
cancellationSignal,
} = input;
const authentication = distributedConfig.authentication;
@@ -122,7 +136,7 @@ export async function validateAuthentication(input: ValidateAuthInput): Promise<
await auditSession.startAgent(AGENT_NAME, prompt, attemptNumber);
const startTime = Date.now();
const submit = createAuthSubmitTool();
const submitTool = createAuthSubmitTool();
const result = await runPiPrompt(
prompt,
repoPath,
@@ -132,13 +146,11 @@ export async function validateAuthentication(input: ValidateAuthInput): Promise<
auditSession,
logger,
'medium',
[submit.tool],
apiKey,
undefined, // callerTools
deliverablesSubdir,
providerConfig,
cancellationSignal,
submitTool,
);
const verdict = submit.getCaptured();
if (verdict !== undefined) result.structuredOutput = verdict;
let classification = classifyResult(result, authentication);
@@ -219,7 +231,7 @@ function countStorageEntries(parsed: unknown, key: 'cookies' | 'origins'): numbe
}
function classifyResult(
result: import('../ai/pi-executor.js').PiPromptResult,
result: import('../ai/pi/pi-executor.js').PiPromptResult,
authentication: NonNullable<DistributedConfig['authentication']>,
): Result<void, PentestError> {
if (!result.success) {
apps/worker/src/services/vuln-renderer.ts
+14 -10
View File
@@ -10,12 +10,18 @@
* Single entry point renderVulnDeliverable(vulnClass, data) covers all 5
* vulnerability classes (injection, xss, auth, ssrf, authz). Per-class title,
* §3 sub-header set, §4 column shape, and §4 section heading are selected by
* branching on vulnClass.
* branching on vulnClass. The .md byte-stability concern is structural only —
* sub-headers, column ordering, sort order — narrative prose varies per run.
*
* Missing tools surface as placeholder sections, not activity failures.
* Required tools (set_findings_summary, set_strategic_intelligence) produce
* loud `[Section X: not provided]` placeholders; recommended tools
* (set_safe_vectors, set_blind_spots) produce quiet "None identified" prose.
* Required-call enforcement is deferred for v1. Missing tools surface as
* placeholder sections, not activity failures. Required tools (set_findings_summary,
* set_strategic_intelligence) produce loud `[Section X: not provided]`
* placeholders; recommended tools (set_safe_vectors, set_blind_spots) produce
* quiet "None identified" prose.
*
* The exploitation queue (`{class}_exploitation_queue.json`) is unrelated —
* it is written by the structured-output submit path in agent-execution.ts and
* is not touched here.
*/
import type {
@@ -25,8 +31,7 @@ import type {
StrategicIntelligenceInput,
VulnClass,
VulnCollectorData,
} from '../mcp-server/vuln-collector.js';
import { BLIND_SPOTS_CLASSES } from '../mcp-server/vuln-collector.js';
} from '../collectors/vuln-collector.js';
// ============================================================================
// PER-CLASS CONSTANTS
@@ -221,9 +226,8 @@ export function renderVulnDeliverable(vulnClass: VulnClass, data: VulnCollectorD
'',
renderSafeVectors(vulnClass, data.safe_vectors),
'',
renderBlindSpots(data.blind_spots),
'',
];
if (BLIND_SPOTS_CLASSES.has(vulnClass)) {
sections.push(renderBlindSpots(data.blind_spots), '');
}
return `${sections.join('\n').trimEnd()}\n`;
}
apps/worker/src/temporal/activities.ts
+74 -36
View File
@@ -18,14 +18,15 @@
import fs from 'node:fs/promises';
import path from 'node:path';
import { ApplicationFailure, Context, heartbeat } from '@temporalio/activity';
import { syncPermissionSystemConfig } from '../ai/pi/permission-system.js';
import { writePlaywrightStealthConfig } from '../ai/playwright-config-writer.js';
import { writeCodePathPermissionConfig } from '../ai/settings-writer.js';
import { AuditSession } from '../audit/index.js';
import type { ResumeAttempt } from '../audit/metrics-tracker.js';
import { authStateFile, generateAuditPath, generateSessionJsonPath, type SessionMetadata } from '../audit/utils.js';
import type { WorkflowSummary } from '../audit/workflow-logger.js';
import type { CheckpointContext } from '../interfaces/checkpoint-provider.js';
import { DEFAULT_DELIVERABLES_SUBDIR, deliverablesDir, resolveSessionJsonPath } from '../paths.js';
import { getAgentGitPaths } from '../services/agent-git-paths.js';
import { getContainer, getOrCreateContainer, removeContainer } from '../services/container.js';
import { classifyErrorForTemporal, PentestError } from '../services/error-handling.js';
import { ExploitationCheckerService } from '../services/exploitation-checker.js';
@@ -38,7 +39,7 @@ import { validateAuthentication } from '../services/validate-authentication.js';
import { AGENTS } from '../session-manager.js';
import type { AgentName } from '../types/agents.js';
import { ALL_AGENTS } from '../types/agents.js';
import type { ContainerConfig, ProviderConfig, VulnClass } from '../types/config.js';
import type { ContainerConfig, VulnClass } from '../types/config.js';
import { ErrorCode } from '../types/errors.js';
import { isErr } from '../types/result.js';
import { atomicWrite, fileExists, readJson } from '../utils/file-io.js';
@@ -58,7 +59,7 @@ const HEARTBEAT_INTERVAL_MS = 2000;
* Input for all agent activities.
*
* Config fields are optional with sensible defaults. When provided, they
* flow through to getOrCreateContainer() for path and credential configuration.
* flow through to getOrCreateContainer() for path configuration.
*/
export interface ActivityInput {
webUrl: string;
@@ -71,13 +72,11 @@ export interface ActivityInput {
// Config fields — serializable, read by getOrCreateContainer()
configYAML?: string;
apiKey?: string;
deliverablesSubdir?: string;
auditDir?: string;
promptDir?: string;
sastSarifPath?: string;
skipGitCheck?: boolean;
providerConfig?: ProviderConfig;
}
/**
@@ -119,9 +118,7 @@ function buildContainerConfig(input: ActivityInput): ContainerConfig {
return {
deliverablesSubdir: input.deliverablesSubdir ?? DEFAULT_DELIVERABLES_SUBDIR,
auditDir: input.auditDir ?? './workspaces',
...(input.apiKey !== undefined && { apiKey: input.apiKey }),
...(input.promptDir !== undefined && { promptDir: input.promptDir }),
...(input.providerConfig !== undefined && { providerConfig: input.providerConfig }),
};
}
@@ -189,12 +186,11 @@ async function runAgentActivity(
configPath,
pipelineTestingMode,
attemptNumber,
...(input.apiKey !== undefined && { apiKey: input.apiKey }),
...(input.providerConfig !== undefined && { providerConfig: input.providerConfig }),
...(input.promptDir !== undefined && { promptDir: input.promptDir }),
...(input.configYAML !== undefined && { configYAML: input.configYAML }),
...(customTools && { customTools }),
...(writeDeliverable && { writeDeliverable }),
cancellationSignal: Context.current().cancellationSignal,
},
auditSession,
logger,
@@ -254,10 +250,10 @@ async function runAgentActivity(
}
export async function runPreReconAgent(input: ActivityInput): Promise<AgentMetrics> {
const { createPreReconCollectorServer } = await import('../mcp-server/pre-recon-collector.js');
const { createPreReconCollector } = await import('../collectors/pre-recon-collector.js');
const { renderPreRecon } = await import('../services/pre-recon-renderer.js');
const collector = createPreReconCollectorServer();
const collector = createPreReconCollector();
const writeDeliverable = async (deliverablesPath: string): Promise<void> => {
const logger = createActivityLogger();
@@ -276,10 +272,10 @@ export async function runPreReconAgent(input: ActivityInput): Promise<AgentMetri
}
export async function runReconAgent(input: ActivityInput): Promise<AgentMetrics> {
const { createReconCollectorServer } = await import('../mcp-server/recon-collector.js');
const { createReconCollector } = await import('../collectors/recon-collector.js');
const { renderRecon } = await import('../services/recon-renderer.js');
const collector = createReconCollectorServer();
const collector = createReconCollector();
const writeDeliverable = async (deliverablesPath: string): Promise<void> => {
const logger = createActivityLogger();
@@ -302,7 +298,7 @@ async function runVulnAgentWithCollector(
vulnClass: 'injection' | 'xss' | 'auth' | 'ssrf' | 'authz',
input: ActivityInput,
): Promise<AgentMetrics> {
const { createVulnCollector } = await import('../mcp-server/vuln-collector.js');
const { createVulnCollector } = await import('../collectors/vuln-collector.js');
const { renderVulnDeliverable } = await import('../services/vuln-renderer.js');
const collector = createVulnCollector(vulnClass);
@@ -358,7 +354,19 @@ async function readExploitQueue(queuePath: string): Promise<{ validIds: Set<stri
if (!(await fileExists(queuePath))) {
return { validIds, idToType };
}
const doc = await readJson<ExploitQueueDocument>(queuePath);
let doc: ExploitQueueDocument;
try {
doc = await readJson<ExploitQueueDocument>(queuePath);
} catch (error) {
const rawMessage = error instanceof Error ? error.message : String(error);
const failure = ApplicationFailure.nonRetryable(
truncateErrorMessage(`Invalid exploitation queue ${queuePath}: ${rawMessage}`),
'InvalidExploitationQueueError',
[{ queuePath }],
);
truncateStackTrace(failure);
throw failure;
}
for (const entry of doc.vulnerabilities ?? []) {
if (!entry.ID) continue;
validIds.add(entry.ID);
@@ -372,7 +380,7 @@ async function runExploitAgentWithCollector(
vulnClass: 'injection' | 'xss' | 'auth' | 'ssrf' | 'authz',
input: ActivityInput,
): Promise<AgentMetrics> {
const { createExploitCollector } = await import('../mcp-server/exploit-collector.js');
const { createExploitCollector } = await import('../collectors/exploit-collector.js');
const { renderExploitDeliverable } = await import('../services/exploit-renderer.js');
const dir = deliverablesDir(input.repoPath, input.deliverablesSubdir);
@@ -453,15 +461,7 @@ export async function runPreflightValidation(input: ActivityInput): Promise<void
const logger = createActivityLogger();
logger.info('Running preflight validation...', { attempt: attemptNumber });
const result = await runPreflightChecks(
input.webUrl,
input.repoPath,
input.configPath,
logger,
input.skipGitCheck,
input.apiKey,
input.providerConfig,
);
const result = await runPreflightChecks(input.webUrl, input.repoPath, input.configPath, logger, input.skipGitCheck);
if (isErr(result)) {
const classified = classifyErrorForTemporal(result.error);
@@ -546,11 +546,10 @@ export async function runAuthenticationValidation(input: ActivityInput): Promise
logger,
auditSession,
attemptNumber,
...(input.apiKey !== undefined && { apiKey: input.apiKey }),
...(input.providerConfig !== undefined && { providerConfig: input.providerConfig }),
...(input.deliverablesSubdir !== undefined && { deliverablesSubdir: input.deliverablesSubdir }),
...(input.promptDir !== undefined && { promptDir: input.promptDir }),
...(input.pipelineTestingMode !== undefined && { pipelineTestingMode: input.pipelineTestingMode }),
cancellationSignal: Context.current().cancellationSignal,
});
if (isErr(result)) {
@@ -657,7 +656,7 @@ export async function syncCodePathDenyRules(input: ActivityInput): Promise<void>
const config = configResult.value;
const denyCount = (config?.avoid ?? []).filter((r) => r.type === 'code_path').length;
await writeCodePathPermissionConfig(config);
syncPermissionSystemConfig(config);
logger.info(
denyCount > 0
? `Synced ${denyCount} code_path deny rule(s) to the pi-permission-system config`
@@ -726,7 +725,29 @@ export async function checkExploitationQueue(input: ActivityInput, vulnType: Vul
// Pass deliverablesPath (not repoPath) — validators expect the deliverables directory
const delivPath = deliverablesDir(repoPath, input.deliverablesSubdir);
return checker.checkQueue(vulnType, delivPath, logger);
try {
return await checker.checkQueue(vulnType, delivPath, logger);
} catch (error) {
const classified = classifyErrorForTemporal(error);
const message = truncateErrorMessage(error instanceof Error ? error.message : String(error));
const details = [{ phase: 'check-exploitation-queue', vulnType }];
const queueValidationFailure = error instanceof PentestError && error.type === 'validation';
// A code-less PentestError (e.g. a filesystem read failure) is classified by
// string-matching, which can miss its retryable flag. Trust the flag directly so
// a non-retryable error never gets a Temporal retry.
const pentestNonRetryable = error instanceof PentestError && !error.retryable;
const failure =
queueValidationFailure || pentestNonRetryable || !classified.retryable
? ApplicationFailure.nonRetryable(
message,
queueValidationFailure ? 'InvalidExploitationQueueError' : classified.type,
details,
)
: ApplicationFailure.create({ message, type: classified.type, details });
truncateStackTrace(failure);
throw failure;
}
}
interface RunScope {
@@ -870,7 +891,16 @@ export async function persistOrValidateRunScope(
await auditSession.initialize(input.workflowId);
const sessionPath = generateSessionJsonPath(sessionMetadata);
const session = await readJson<SessionJson>(sessionPath);
let session: SessionJson;
try {
session = await readJson<SessionJson>(sessionPath);
} catch (error) {
const rawMessage = error instanceof Error ? error.message : String(error);
throw ApplicationFailure.nonRetryable(
`Corrupted session.json in workspace ${input.sessionId}: ${rawMessage}`,
'CorruptedSessionError',
);
}
if (session.session.scope) {
const recorded = session.session.scope;
@@ -933,11 +963,12 @@ export async function restoreGitCheckpoint(
const logger = createActivityLogger();
logger.info(`Restoring deliverables to ${checkpointHash}...`);
// Validate hash exists in this clone before attempting reset
// Validate the hash exists in the deliverables clone (the repo actually being
// reset below) before attempting reset.
try {
await executeGitCommandWithRetry(
['git', 'rev-parse', '--verify', checkpointHash],
repoPath,
['git', 'cat-file', '-e', `${checkpointHash}^{commit}`],
deliverablesPath,
'verify checkpoint hash exists',
);
} catch {
@@ -950,7 +981,13 @@ export async function restoreGitCheckpoint(
deliverablesPath,
'reset deliverables to checkpoint',
);
await executeGitCommandWithRetry(['git', 'clean', '-fd'], deliverablesPath, 'clean untracked deliverables');
// Scope the untracked clean so a completed agent's deliverables survive: exclude every
// completed agent's paths, cleaning only leftovers from the incomplete agents being re-run.
const incompleteSet = new Set<AgentName>(incompleteAgents);
const completedPaths = ALL_AGENTS.filter((name) => !incompleteSet.has(name)).flatMap(getAgentGitPaths);
const cleanArgs = ['git', 'clean', '-fd', ...completedPaths.flatMap((completedPath) => ['-e', completedPath])];
await executeGitCommandWithRetry(cleanArgs, deliverablesPath, 'clean untracked deliverables');
// Explicitly delete partial deliverables for incomplete agents
for (const agentName of incompleteAgents) {
@@ -1063,8 +1100,9 @@ export async function logWorkflowComplete(input: ActivityInput, summary: Workflo
await auditSession.logWorkflowComplete(cumulativeSummary);
// 6. Surface the final report at the run root. Done here (not in the report phase)
// so it also runs when a resume skips an already-complete report phase.
if (summary.status === 'completed') {
// so it also runs when a resume skips an already-complete report phase. A partial
// run still assembles a report (only some classes were not assessed), so surface it too.
if (summary.status === 'completed' || summary.status === 'partial') {
try {
await copyReportToRunRoot(
input.repoPath,
apps/worker/src/temporal/pipeline.ts
+1
View File
@@ -14,4 +14,5 @@ export type {
ResumeState,
VulnExploitPipelineResult,
} from './shared.js';
export { PipelineExecutionError } from './shared.js';
export { pentestPipeline } from './workflows.js';
apps/worker/src/temporal/shared.ts
+21 -5
View File
@@ -2,7 +2,7 @@ import { defineQuery } from '@temporalio/workflow';
export type { AgentMetrics } from '../types/metrics.js';
import type { DistributedConfig, PipelineConfig, ProviderConfig, VulnClass } from '../types/config.js';
import type { DistributedConfig, PipelineConfig, VulnClass } from '../types/config.js';
import type { ErrorCode } from '../types/errors.js';
import type { AgentMetrics } from '../types/metrics.js';
@@ -21,14 +21,12 @@ export interface PipelineInput {
// Config fields — serializable, flow through to ActivityInput → getOrCreateContainer()
configYAML?: string; // Raw YAML string (parsed in activity, not workflow — workflow sandbox can't use Node.js)
configData?: DistributedConfig; // Pre-parsed config (bypasses file loading)
apiKey?: string; // API key override (avoids process.env mutation)
deliverablesSubdir?: string; // Override deliverables path (default: '.shannon/deliverables')
auditDir?: string; // Override audit log directory (default: './workspaces')
promptDir?: string; // Override prompt template directory
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, custom base URL, etc.)
vulnClasses?: VulnClass[]; // omitted = all five
exploit?: boolean; // false skips the exploitation phase
}
@@ -49,10 +47,13 @@ export interface PipelineSummary {
}
export interface PipelineState {
status: 'running' | 'completed' | 'failed' | 'cancelled';
status: 'running' | 'completed' | 'failed' | 'cancelled' | 'partial';
currentPhase: string | null;
currentAgent: string | null;
completedAgents: string[];
// Vuln classes whose pipeline failed while at least one other succeeded. Drives the
// partial terminal status so a crashed class isn't reported as if it fully passed.
failedPipelines: { vulnType: VulnClass; error: string }[];
failedAgent: string | null;
error: string | null;
errorCode?: ErrorCode;
@@ -61,6 +62,21 @@ export interface PipelineState {
summary: PipelineSummary | null;
}
/**
* Thrown by pentestPipeline() when the run fails, carrying the fully-populated
* PipelineState (real agentMetrics, completedAgents, summary) so a consumer can
* report actual spend instead of synthesizing a zeroed failed state. `cause`
* preserves the original error for classification and Temporal failure reporting.
*/
export class PipelineExecutionError extends Error {
override name = 'PipelineExecutionError' as const;
readonly state: PipelineState;
constructor(message: string, state: PipelineState, options?: { cause?: unknown }) {
super(message, options);
this.state = state;
}
}
// Extended state returned by getProgress query (includes computed fields)
export interface PipelineProgress extends PipelineState {
workflowId: string;
@@ -69,7 +85,7 @@ export interface PipelineProgress extends PipelineState {
// Result from a single vuln→exploit pipeline
export interface VulnExploitPipelineResult {
vulnType: string;
vulnType: VulnClass;
vulnMetrics: AgentMetrics | null;
exploitMetrics: AgentMetrics | null;
exploitDecision: {
apps/worker/src/temporal/summary-mapper.ts
+4 -1
View File
@@ -19,7 +19,10 @@ import type { PipelineState } from './shared.js';
* safely imported into Temporal workflows. The caller must ensure
* state.summary is set before calling (via computeSummary).
*/
export function toWorkflowSummary(state: PipelineState, status: 'completed' | 'failed' | 'cancelled'): WorkflowSummary {
export function toWorkflowSummary(
state: PipelineState,
status: 'completed' | 'failed' | 'cancelled' | 'partial',
): WorkflowSummary {
// state.summary must be computed before calling this mapper
const summary = state.summary;
if (!summary) {
apps/worker/src/temporal/worker.ts
+6 -2
View File
@@ -299,8 +299,12 @@ async function loadOrchestrationConfig(configPath: string | undefined): Promise<
...(config.vuln_classes && config.vuln_classes.length > 0 && { vulnClasses: [...config.vuln_classes] }),
...(config.exploit !== undefined && { exploit: config.exploit === 'true' }),
};
} catch {
return { pipelineConfig: {} };
} catch (error) {
// A broken config must fail the run, not silently fall back to empty
// defaults that quietly change scope (vuln classes, exploit, retries).
const message = error instanceof Error ? error.message : String(error);
console.error(`Failed to parse config ${configPath}: ${message}`);
process.exit(1);
}
}
apps/worker/src/temporal/workflows.ts
+165 -56
View File
@@ -25,6 +25,7 @@
import {
ApplicationFailure,
CancellationScope,
isCancellation,
log,
proxyActivities,
@@ -39,6 +40,7 @@ import type { ActivityInput } from './activities.js';
import {
type AgentMetrics,
getProgress,
PipelineExecutionError,
type PipelineInput,
type PipelineProgress,
type PipelineState,
@@ -165,6 +167,15 @@ function computeSummary(state: PipelineState): PipelineSummary {
};
}
const MAX_PIPELINE_ERROR_MESSAGE_LENGTH = 2000;
function truncatePipelineErrorMessage(message: string): string {
if (message.length <= MAX_PIPELINE_ERROR_MESSAGE_LENGTH) {
return message;
}
return `${message.slice(0, MAX_PIPELINE_ERROR_MESSAGE_LENGTH - 20)}\n[truncated]`;
}
/**
* Core pipeline orchestration. Coordinates the pentest pipeline stages.
*
@@ -203,6 +214,7 @@ export async function pentestPipeline(input: PipelineInput): Promise<PipelineSta
currentPhase: null,
currentAgent: null,
completedAgents: [],
failedPipelines: [],
failedAgent: null,
error: null,
startTime: Date.now(),
@@ -237,13 +249,11 @@ export async function pentestPipeline(input: PipelineInput): Promise<PipelineSta
}),
// Config fields — flow through to getOrCreateContainer()
...(input.configYAML !== undefined && { configYAML: input.configYAML }),
...(input.apiKey !== undefined && { apiKey: input.apiKey }),
...(input.deliverablesSubdir !== undefined && { deliverablesSubdir: input.deliverablesSubdir }),
...(input.auditDir !== undefined && { auditDir: input.auditDir }),
...(input.promptDir !== undefined && { promptDir: input.promptDir }),
...(input.sastSarifPath !== undefined && { sastSarifPath: input.sastSarifPath }),
...(input.skipGitCheck !== undefined && { skipGitCheck: input.skipGitCheck }),
...(input.providerConfig !== undefined && { providerConfig: input.providerConfig }),
};
const selectedVulnClasses: readonly VulnClass[] =
@@ -373,24 +383,70 @@ export async function pentestPipeline(input: PipelineInput): Promise<PipelineSta
];
}
// Aggregate errors from settled pipeline promises.
// A rejected settle can be a genuine Temporal cancellation (runVulnExploitPipeline rethrows in
// its isCancellation branch). Cancellation must win over failed/partial classification — there is
// no report worth shipping once the user has cancelled, and rethrowing lets the workflow's outer
// isCancellation handler produce a real cancelled state instead of a hard failure.
function throwIfPipelineCancelled(results: PromiseSettledResult<VulnExploitPipelineResult>[]): void {
const cancelled = results.find(
(r): r is PromiseRejectedResult => r.status === 'rejected' && isCancellation(r.reason),
);
if (cancelled) {
throw cancelled.reason;
}
}
// Classify the settled pipeline results into clean / partial / fail-hard.
// Metrics and completedAgents are updated incrementally inside runVulnExploitPipeline
// so that getProgress queries reflect real-time status during execution.
function aggregatePipelineResults(results: PromiseSettledResult<VulnExploitPipelineResult>[]): void {
const failedPipelines: string[] = [];
function aggregatePipelineResults(
results: PromiseSettledResult<VulnExploitPipelineResult>[],
alreadyCompletedPipelineCount: number,
): void {
throwIfPipelineCancelled(results);
const failed: { vulnType: VulnClass; error: string }[] = [];
// A rejected settle is now unexpected (runVulnExploitPipeline catches and returns its error in
// the value). Without a value we cannot attribute the failure to a class, so we treat it as a
// hard failure rather than risk an under-qualified report.
const unattributable: string[] = [];
for (const result of results) {
if (result.status === 'rejected') {
const errorMsg = result.reason instanceof Error ? result.reason.message : String(result.reason);
failedPipelines.push(errorMsg);
if (result.status === 'fulfilled') {
if (result.value.error !== null) {
failed.push({ vulnType: result.value.vulnType, error: result.value.error });
}
} else {
const rawMessage = result.reason instanceof Error ? result.reason.message : String(result.reason);
unattributable.push(truncatePipelineErrorMessage(rawMessage));
}
}
if (failedPipelines.length > 0) {
log.warn(`${failedPipelines.length} pipeline(s) failed`, {
failures: failedPipelines,
});
const failedCount = failed.length + unattributable.length;
const totalPipelineCount = results.length + alreadyCompletedPipelineCount;
if (failedCount === 0) {
return;
}
// All run pipelines failed, or a failure we cannot attribute to a class → fail-hard. There is
// no report worth shipping, and we must never render an un-assessed class as if it passed.
if (failedCount === totalPipelineCount || unattributable.length > 0) {
const allErrors = [...failed.map((f) => `${f.vulnType}: ${f.error}`), ...unattributable];
const message = `${failedCount} vulnerability/exploitation pipeline(s) failed`;
state.status = 'failed';
state.failedAgent = 'pipelines';
state.error = `${message}: ${allErrors.join('; ')}`;
log.warn(message, { failures: allErrors });
throw ApplicationFailure.nonRetryable(state.error, 'PipelineFailedError', [{ failures: allErrors }]);
}
// Partial: at least one class succeeded and at least one failed. Record the failed classes and
// set the partial terminal status; do NOT throw — the successful pipelines still ship.
state.failedPipelines = failed;
state.status = 'partial';
log.warn(`${failed.length} of ${totalPipelineCount} pipeline(s) failed — continuing with partial results`, {
failures: failed.map((f) => `${f.vulnType}: ${f.error}`),
});
}
// Run thunks with a concurrency limit, returning PromiseSettledResult for each.
@@ -480,56 +536,87 @@ export async function pentestPipeline(input: PipelineInput): Promise<PipelineSta
const vulnAgentName = `${vulnType}-vuln`;
const exploitAgentName = `${vulnType}-exploit`;
// 1. Run vulnerability analysis (or skip if resumed)
let vulnMetrics: AgentMetrics | null = null;
if (!shouldSkip(vulnAgentName)) {
vulnMetrics = await runVulnAgent();
state.agentMetrics[vulnAgentName] = vulnMetrics;
state.completedAgents.push(vulnAgentName);
if (input.checkpointsEnabled) {
await a.saveCheckpoint(activityInput, vulnAgentName, 'vulnerability-analysis', state);
// A class failure must not reject the pipeline set — that would lose the class identity
// (results are completion-ordered) and force fail-hard. Catch here and return the error in
// the result's `error` field so aggregatePipelineResults can attribute it to `vulnType`.
try {
// 1. Run vulnerability analysis (or skip if resumed)
let vulnMetrics: AgentMetrics | null = null;
if (!shouldSkip(vulnAgentName)) {
vulnMetrics = await runVulnAgent();
state.agentMetrics[vulnAgentName] = vulnMetrics;
state.completedAgents.push(vulnAgentName);
if (input.checkpointsEnabled) {
await a.saveCheckpoint(activityInput, vulnAgentName, 'vulnerability-analysis', state);
}
} else {
log.info(`Skipping ${vulnAgentName} (already complete)`);
state.completedAgents.push(vulnAgentName);
}
} else {
log.info(`Skipping ${vulnAgentName} (already complete)`);
state.completedAgents.push(vulnAgentName);
}
// 1.5. Merge external findings from consumer provider into exploitation queue
await a.mergeFindingsIntoQueue(activityInput, vulnType);
// 1.5. Merge external findings from consumer provider into exploitation queue
await a.mergeFindingsIntoQueue(activityInput, vulnType);
// 2. Check exploitation queue for actionable findings
const decision = await a.checkExploitationQueue(activityInput, vulnType);
// 2. Check exploitation queue for actionable findings
const decision = await a.checkExploitationQueue(activityInput, vulnType);
// 3. Previously-completed exploits are preserved regardless of mode; new exploits gated by mode.
let exploitMetrics: AgentMetrics | null = null;
if (shouldSkip(exploitAgentName)) {
log.info(`Skipping ${exploitAgentName} (already complete)`);
state.completedAgents.push(exploitAgentName);
} else if (decision.shouldExploit && exploit) {
exploitMetrics = await runExploitAgent();
state.agentMetrics[exploitAgentName] = exploitMetrics;
state.completedAgents.push(exploitAgentName);
if (input.checkpointsEnabled) {
await a.saveCheckpoint(activityInput, exploitAgentName, 'exploitation', state);
// 3. Previously-completed exploits are preserved regardless of mode; new exploits gated by mode.
let exploitMetrics: AgentMetrics | null = null;
if (shouldSkip(exploitAgentName)) {
log.info(`Skipping ${exploitAgentName} (already complete)`);
state.completedAgents.push(exploitAgentName);
} else if (decision.shouldExploit && exploit) {
exploitMetrics = await runExploitAgent();
state.agentMetrics[exploitAgentName] = exploitMetrics;
state.completedAgents.push(exploitAgentName);
if (input.checkpointsEnabled) {
await a.saveCheckpoint(activityInput, exploitAgentName, 'exploitation', state);
}
} else {
// Exploitation did not run (exploit mode off, or no actionable findings) — still
// mark the agent complete so a resume does not treat it as unfinished work.
log.info(
`Marking ${exploitAgentName} complete (${decision.shouldExploit ? 'exploit mode disabled' : 'no actionable findings'})`,
);
state.completedAgents.push(exploitAgentName);
if (input.checkpointsEnabled) {
await a.saveCheckpoint(activityInput, exploitAgentName, 'exploitation', state);
}
}
}
return {
vulnType,
vulnMetrics,
exploitMetrics,
exploitDecision: {
shouldExploit: decision.shouldExploit,
vulnerabilityCount: decision.vulnerabilityCount,
},
error: null,
};
return {
vulnType,
vulnMetrics,
exploitMetrics,
exploitDecision: {
shouldExploit: decision.shouldExploit,
vulnerabilityCount: decision.vulnerabilityCount,
},
error: null,
};
} catch (error) {
// Let cancellation propagate to the workflow-level handler.
if (isCancellation(error)) {
throw error;
}
const rawMessage = error instanceof Error ? error.message : String(error);
const message = truncatePipelineErrorMessage(rawMessage);
log.warn(`Pipeline ${vulnType} failed`, { error: message });
return {
vulnType,
vulnMetrics: state.agentMetrics[vulnAgentName] ?? null,
exploitMetrics: state.agentMetrics[exploitAgentName] ?? null,
exploitDecision: null,
error: message,
};
}
}
const maxConcurrent = input.pipelineConfig?.max_concurrent_pipelines ?? 5;
const pipelineConfigs = buildPipelineConfigs();
const pipelineThunks: Array<() => Promise<VulnExploitPipelineResult>> = [];
let alreadyCompletedPipelineCount = 0;
for (const config of pipelineConfigs) {
// Excluded classes drop entirely; any prior deliverables stay on disk but don't count this run.
@@ -542,11 +629,12 @@ export async function pentestPipeline(input: PipelineInput): Promise<PipelineSta
} else {
log.info(`Skipping entire ${config.vulnType} pipeline (both agents complete)`);
state.completedAgents.push(config.vulnAgent, config.exploitAgent);
alreadyCompletedPipelineCount++;
}
}
const pipelineResults = await runWithConcurrencyLimit(pipelineThunks, maxConcurrent);
aggregatePipelineResults(pipelineResults);
aggregatePipelineResults(pipelineResults, alreadyCompletedPipelineCount);
state.currentPhase = 'exploitation';
state.currentAgent = null;
@@ -584,13 +672,16 @@ export async function pentestPipeline(input: PipelineInput): Promise<PipelineSta
await a.saveCheckpoint(activityInput, 'report-output', 'reporting', state);
}
state.status = 'completed';
// Preserve a partial verdict (set by aggregatePipelineResults) — a clean run is 'completed',
// a run where some classes were not assessed is 'partial'.
const terminalStatus: 'completed' | 'partial' = state.failedPipelines.length > 0 ? 'partial' : 'completed';
state.status = terminalStatus;
state.currentPhase = null;
state.currentAgent = null;
state.summary = computeSummary(state);
// Log workflow completion summary
await a.logWorkflowComplete(activityInput, toWorkflowSummary(state, 'completed'));
await a.logWorkflowComplete(activityInput, toWorkflowSummary(state, terminalStatus));
return state;
} catch (error) {
@@ -599,7 +690,17 @@ export async function pentestPipeline(input: PipelineInput): Promise<PipelineSta
state.status = 'cancelled';
state.error = `Cancelled during phase: ${state.currentPhase ?? 'unknown'}`;
state.summary = computeSummary(state);
await a.logWorkflowComplete(activityInput, toWorkflowSummary(state, 'cancelled'));
// Finalization runs I/O activities; shield them from the cancellation so the
// cancelled state is still logged rather than aborted mid-write.
await CancellationScope.nonCancellable(async () => {
try {
await a.logWorkflowComplete(activityInput, toWorkflowSummary(state, 'cancelled'));
} catch (completionError) {
log.warn('Failed to finalize cancelled workflow', {
error: completionError instanceof Error ? completionError.message : String(completionError),
});
}
});
return state;
}
@@ -613,9 +714,17 @@ export async function pentestPipeline(input: PipelineInput): Promise<PipelineSta
state.summary = computeSummary(state);
// Log workflow failure summary
await a.logWorkflowComplete(activityInput, toWorkflowSummary(state, 'failed'));
try {
await a.logWorkflowComplete(activityInput, toWorkflowSummary(state, 'failed'));
} catch (completionError) {
log.warn('Failed to finalize failed workflow', {
error: completionError instanceof Error ? completionError.message : String(completionError),
});
}
throw error;
// Carry the populated state so a consumer can report real spend instead of a zeroed
// failed state. The original error rides as `cause` for classification/reporting.
throw new PipelineExecutionError(state.error ?? 'Pipeline failed', state, { cause: error });
}
}
apps/worker/src/types/config.ts
+2 -25
View File
@@ -91,40 +91,17 @@ export interface DistributedConfig {
rules_of_engagement: string;
}
/**
* LLM provider configuration for multi-provider support.
*
* Resolved by the pi model/provider layer at execution time. Recognized
* providerType values: 'bedrock', 'custom_base_url', 'anthropic_api'.
* When omitted or 'anthropic_api', falls back to apiKey + ANTHROPIC_API_KEY.
*/
export interface ProviderConfig {
readonly providerType?: string;
readonly apiKey?: string;
readonly awsRegion?: string;
readonly awsAccessKeyId?: string;
readonly awsSecretAccessKey?: string;
readonly baseUrl?: string;
readonly authToken?: string;
readonly modelOverrides?: Record<string, string>;
readonly supportsStructuredOutput?: boolean;
}
/**
* Runtime configuration for the DI container.
*
* Abstracts path conventions and credential threading so consumers
* can override OSS defaults without modifying source files.
* Abstracts path conventions so consumers can override OSS defaults
* without modifying source files.
*/
export interface ContainerConfig {
/** Subdirectory for deliverables relative to repoPath. Default: '.shannon/deliverables' */
readonly deliverablesSubdir: string;
/** Directory for audit logs. Default: './workspaces' */
readonly auditDir: string;
/** API key override — when set, executor reads from config instead of process.env */
readonly apiKey?: string;
/** Prompt directory override — when set, prompt manager loads from this path */
readonly promptDir?: string;
/** LLM provider configuration for the pi executor */
readonly providerConfig?: ProviderConfig;
}
apps/worker/src/utils/billing-detection.ts
+5 -5
View File
@@ -9,7 +9,7 @@
*
* Anthropic's spending cap behavior is inconsistent:
* - Sometimes a proper provider error (billing_error)
* - Sometimes the agent responds with text about the cap
* - Sometimes the model responds with text about the cap
* - Sometimes partial billing before cutoff
*
* This module provides defense-in-depth detection with shared pattern lists
@@ -17,8 +17,8 @@
*/
/**
* Text patterns for provider/harness output sniffing (what the agent says).
* Used by the pi event stream and the behavioral heuristic.
* Text patterns for model-output sniffing (what the model says).
* Used by the pi executor and the behavioral heuristic.
*/
export const BILLING_TEXT_PATTERNS = [
'spending cap',
@@ -48,7 +48,7 @@ export const BILLING_API_PATTERNS = [
/**
* Checks if text matches any billing text pattern.
* Used for sniffing agent output content for spending cap messages.
* Used for sniffing model output content for spending cap messages.
*/
export function matchesBillingTextPattern(text: string): boolean {
const lowerText = text.toLowerCase();
@@ -67,7 +67,7 @@ export function matchesBillingApiPattern(message: string): boolean {
/**
* Behavioral heuristic for detecting spending cap.
*
* When the agent hits a spending cap, it often returns a short message
* When the model hits a spending cap, it often returns a short message
* with $0 cost. Legitimate agent work NEVER costs $0 with only 1-2 turns.
*
* This combines three signals:
apps/worker/src/utils/browser-agents.ts
+33
View File
@@ -0,0 +1,33 @@
// 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.
/**
* Agents that drive a live browser via the playwright-cli skill. These get the
* skill registered through pi's `skillsOverride`.
*
* `validate-authentication` is not an AgentName in the main pipeline graph but
* runs the same executor, so it is included by string.
*/
export const BROWSER_AGENTS: ReadonlySet<string> = new Set([
'recon',
'injection-vuln',
'xss-vuln',
'auth-vuln',
'ssrf-vuln',
'authz-vuln',
'injection-exploit',
'xss-exploit',
'auth-exploit',
'ssrf-exploit',
'authz-exploit',
'validate-authentication',
'verify-exploit',
]);
/** Whether the given agent uses the browser (and therefore needs the playwright-cli skill under pi). */
export function isBrowserAgent(agentName: string | null | undefined): boolean {
return agentName != null && BROWSER_AGENTS.has(agentName);
}
pnpm-lock.yaml
Generated
-3
View File
@@ -81,9 +81,6 @@ importers:
typebox:
specifier: 1.1.38
version: 1.1.38
zod:
specifier: ^4.3.6
version: 4.3.6
zx:
specifier: ^8.0.0
version: 8.8.5