From d0b0ec3378ef159871ce521dbb6043a7b97889e3 Mon Sep 17 00:00:00 2001 From: ezl-keygraph Date: Thu, 16 Jul 2026 14:09:02 +0530 Subject: [PATCH] refactor(worker): converge shared core with shannon-oss (#388) * 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 --- apps/worker/package.json | 1 - apps/worker/prompts/report-executive.txt | 1 - .../prompts/validate-authentication.txt | 6 +- apps/worker/prompts/vuln-auth.txt | 11 +- apps/worker/prompts/vuln-ssrf.txt | 11 +- apps/worker/src/ai/models.ts | 62 +-- apps/worker/src/ai/pi/permission-system.ts | 141 ++++++ apps/worker/src/ai/{ => pi}/pi-executor.ts | 123 +++-- apps/worker/src/ai/pi/session-tools.ts | 116 +++++ apps/worker/src/ai/pi/task-tool.ts | 156 +++++++ apps/worker/src/ai/queue-schemas.ts | 121 ++--- apps/worker/src/ai/settings-writer.ts | 75 --- apps/worker/src/ai/submit-tool.ts | 60 +++ apps/worker/src/ai/tools.ts | 205 -------- apps/worker/src/audit/audit-session.ts | 2 +- apps/worker/src/audit/metrics-tracker.ts | 12 +- apps/worker/src/audit/workflow-logger.ts | 4 +- .../exploit-collector.ts | 116 +++-- .../pre-recon-collector.ts | 294 ++++++------ .../recon-collector.ts | 442 +++++++++--------- apps/worker/src/collectors/schema.ts | 29 ++ .../vuln-collector.ts | 81 ++-- apps/worker/src/paths.ts | 3 - apps/worker/src/scripts/generate-totp.ts | 11 +- apps/worker/src/services/agent-execution.ts | 169 +++++-- apps/worker/src/services/agent-git-paths.ts | 31 ++ apps/worker/src/services/error-handling.ts | 25 +- apps/worker/src/services/exploit-renderer.ts | 43 +- .../src/services/exploitation-checker.ts | 23 +- apps/worker/src/services/git-manager.ts | 258 ++++++---- apps/worker/src/services/index.ts | 4 +- .../worker/src/services/pre-recon-renderer.ts | 15 +- apps/worker/src/services/preflight.ts | 31 +- apps/worker/src/services/prompt-manager.ts | 91 ++-- apps/worker/src/services/recon-renderer.ts | 16 +- .../src/services/validate-authentication.ts | 84 ++-- apps/worker/src/services/vuln-renderer.ts | 24 +- apps/worker/src/temporal/activities.ts | 110 +++-- apps/worker/src/temporal/pipeline.ts | 1 + apps/worker/src/temporal/shared.ts | 26 +- apps/worker/src/temporal/summary-mapper.ts | 5 +- apps/worker/src/temporal/worker.ts | 8 +- apps/worker/src/temporal/workflows.ts | 221 ++++++--- apps/worker/src/types/config.ts | 27 +- apps/worker/src/utils/billing-detection.ts | 10 +- apps/worker/src/utils/browser-agents.ts | 33 ++ pnpm-lock.yaml | 3 - 47 files changed, 2007 insertions(+), 1334 deletions(-) create mode 100644 apps/worker/src/ai/pi/permission-system.ts rename apps/worker/src/ai/{ => pi}/pi-executor.ts (76%) create mode 100644 apps/worker/src/ai/pi/session-tools.ts create mode 100644 apps/worker/src/ai/pi/task-tool.ts delete mode 100644 apps/worker/src/ai/settings-writer.ts create mode 100644 apps/worker/src/ai/submit-tool.ts delete mode 100644 apps/worker/src/ai/tools.ts rename apps/worker/src/{mcp-server => collectors}/exploit-collector.ts (85%) rename apps/worker/src/{mcp-server => collectors}/pre-recon-collector.ts (77%) rename apps/worker/src/{mcp-server => collectors}/recon-collector.ts (75%) create mode 100644 apps/worker/src/collectors/schema.ts rename apps/worker/src/{mcp-server => collectors}/vuln-collector.ts (89%) create mode 100644 apps/worker/src/services/agent-git-paths.ts create mode 100644 apps/worker/src/utils/browser-agents.ts diff --git a/apps/worker/package.json b/apps/worker/package.json index 63b61f9..d843178 100644 --- a/apps/worker/package.json +++ b/apps/worker/package.json @@ -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": { diff --git a/apps/worker/prompts/report-executive.txt b/apps/worker/prompts/report-executive.txt index d06cc7d..bcb3be1 100644 --- a/apps/worker/prompts/report-executive.txt +++ b/apps/worker/prompts/report-executive.txt @@ -21,7 +21,6 @@ IMPORTANT: You are MODIFYING an existing file, not creating a new one. URL: {{WEB_URL}} -{{DESCRIPTION}} Filesystem: - {{REPO_PATH}}/ (read only) diff --git a/apps/worker/prompts/validate-authentication.txt b/apps/worker/prompts/validate-authentication.txt index 4e5ba66..9dd9aa9 100644 --- a/apps/worker/prompts/validate-authentication.txt +++ b/apps/worker/prompts/validate-authentication.txt @@ -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. - -When the login attempt concludes, call the `submit_auth_result` tool to report the outcome. - - - 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. diff --git a/apps/worker/prompts/vuln-auth.txt b/apps/worker/prompts/vuln-auth.txt index fedb60c..7acaea3 100644 --- a/apps/worker/prompts/vuln-auth.txt +++ b/apps/worker/prompts/vuln-auth.txt @@ -195,22 +195,23 @@ For each check you perform from the list above (Transport, Rate Limiting, Sessio -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. @@ -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. diff --git a/apps/worker/prompts/vuln-ssrf.txt b/apps/worker/prompts/vuln-ssrf.txt index f3ddb3c..5063cec 100644 --- a/apps/worker/prompts/vuln-ssrf.txt +++ b/apps/worker/prompts/vuln-ssrf.txt @@ -244,22 +244,23 @@ For each check you perform from the list above, you must make a final **verdict* -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. @@ -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. diff --git a/apps/worker/src/ai/models.ts b/apps/worker/src/ai/models.ts index 97716d9..0bbadbb 100644 --- a/apps/worker/src/ai/models.ts +++ b/apps/worker/src/ai/models.ts @@ -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) { diff --git a/apps/worker/src/ai/pi/permission-system.ts b/apps/worker/src/ai/pi/permission-system.ts new file mode 100644 index 0000000..a1febac --- /dev/null +++ b/apps/worker/src/ai/pi/permission-system.ts @@ -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; + 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 = { '*': '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 (`/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), '..'); +} diff --git a/apps/worker/src/ai/pi-executor.ts b/apps/worker/src/ai/pi/pi-executor.ts similarity index 76% rename from apps/worker/src/ai/pi-executor.ts rename to apps/worker/src/ai/pi/pi-executor.ts index a3dfa2c..acc7ded 100644 --- a/apps/worker/src/ai/pi-executor.ts +++ b/apps/worker/src/ai/pi/pi-executor.ts @@ -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 { +async function buildResourceLoader( + cwd: string, + logger: ActivityLogger, + agentName: string | null, +): Promise { // 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 { - // 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 diff --git a/apps/worker/src/ai/pi/session-tools.ts b/apps/worker/src/ai/pi/session-tools.ts new file mode 100644 index 0000000..9f77c7f --- /dev/null +++ b/apps/worker/src/ai/pi/session-tools.ts @@ -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, + }; + }, + }); +} diff --git a/apps/worker/src/ai/pi/task-tool.ts b/apps/worker/src/ai/pi/task-tool.ts new file mode 100644 index 0000000..c0865be --- /dev/null +++ b/apps/worker/src/ai/pi/task-tool.ts @@ -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; + 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; +} diff --git a/apps/worker/src/ai/queue-schemas.ts b/apps/worker/src/ai/queue-schemas.ts index a26b34d..3dc968c 100644 --- a/apps/worker/src/ai/queue-schemas.ts +++ b/apps/worker/src/ai/queue-schemas.ts @@ -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 - * `_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>; +export type XssFinding = Static>; +export type AuthFinding = Static>; +export type SsrfFinding = Static>; +export type AuthzFinding = Static>; + const PER_TYPE_FIELDS: Partial>>> = { 'injection-vuln': injectionFields, 'xss-vuln': xssFields, @@ -92,28 +108,6 @@ const PER_TYPE_FIELDS: Partial>; -export type XssFinding = Static>; -export type AuthFinding = Static>; -export type SsrfFinding = Static>; -export type AuthzFinding = Static>; - -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> = { 'injection-vuln': 'injection_exploitation_queue.json', 'xss-vuln': 'xss_exploitation_queue.json', @@ -122,38 +116,53 @@ const VULN_AGENT_QUEUE_FILENAMES: Partial> = { '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.', + }; } diff --git a/apps/worker/src/ai/settings-writer.ts b/apps/worker/src/ai/settings-writer.ts deleted file mode 100644 index 46c126d..0000000 --- a/apps/worker/src/ai/settings-writer.ts +++ /dev/null @@ -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 { - 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 = { '*': '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 }); -} diff --git a/apps/worker/src/ai/submit-tool.ts b/apps/worker/src/ai/submit-tool.ts new file mode 100644 index 0000000..e461302 --- /dev/null +++ b/apps/worker/src/ai/submit-tool.ts @@ -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): 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.', + }; +} diff --git a/apps/worker/src/ai/tools.ts b/apps/worker/src/ai/tools.ts deleted file mode 100644 index 585ffd7..0000000 --- a/apps/worker/src/ai/tools.ts +++ /dev/null @@ -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; - 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: {} }; - }, - }); -} diff --git a/apps/worker/src/audit/audit-session.ts b/apps/worker/src/audit/audit-session.ts index 9d364fe..39dc203 100644 --- a/apps/worker/src/audit/audit-session.ts +++ b/apps/worker/src/audit/audit-session.ts @@ -210,7 +210,7 @@ export class AuditSession { /** * Update session status */ - async updateSessionStatus(status: 'in-progress' | 'completed' | 'failed' | 'cancelled'): Promise { + async updateSessionStatus(status: 'in-progress' | 'completed' | 'failed' | 'cancelled' | 'partial'): Promise { await this.ensureInitialized(); const unlock = await sessionMutex.lock(this.sessionId); diff --git a/apps/worker/src/audit/metrics-tracker.ts b/apps/worker/src/audit/metrics-tracker.ts index 9bad57c..914c8d1 100644 --- a/apps/worker/src/audit/metrics-tracker.ts +++ b/apps/worker/src/audit/metrics-tracker.ts @@ -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 { + async updateSessionStatus(status: 'in-progress' | 'completed' | 'failed' | 'cancelled' | 'partial'): Promise { 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(); } diff --git a/apps/worker/src/audit/workflow-logger.ts b/apps/worker/src/audit/workflow-logger.ts index 9662b14..5bdf7bb 100644 --- a/apps/worker/src/audit/workflow-logger.ts +++ b/apps/worker/src/audit/workflow-logger.ts @@ -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(); diff --git a/apps/worker/src/mcp-server/exploit-collector.ts b/apps/worker/src/collectors/exploit-collector.ts similarity index 85% rename from apps/worker/src/mcp-server/exploit-collector.ts rename to apps/worker/src/collectors/exploit-collector.ts index fe05f8b..6e95787 100644 --- a/apps/worker/src/mcp-server/exploit-collector.ts +++ b/apps/worker/src/collectors/exploit-collector.ts @@ -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) { +export function buildSchemas(validIds: ReadonlySet) { const vulnerabilityIdField = Type.String({ minLength: 1, description: @@ -147,7 +149,7 @@ function buildSchemas(validIds: ReadonlySet) { }), ); - 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) { '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) { ); 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) { }), ); - // 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) { 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) { 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) { const StrictSchema = Type.Union([ExploitedSchema, BlockedSchema]); - return { flatShape, StrictSchema }; + return { flatSchema, StrictSchema }; } -type FlatInput = Static['flatShape']>; -type StrictInput = Static['StrictSchema']>; - // ============================================================================ // RESPONSE HELPERS // ============================================================================ -interface ToolResult { - content: Array<{ type: 'text'; text: string }>; - details: Record; - isError?: boolean; -} - -function createToolResult(response: { status: string; [key: string]: unknown }): ToolResult { - const isError = response.status === 'error'; +function toolResult(payload: Record) { 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): ToolResult { - return createToolResult({ status: 'success', ...data }); +function successResult(data: Record) { + 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['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; } -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 => { - 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], }; } diff --git a/apps/worker/src/mcp-server/pre-recon-collector.ts b/apps/worker/src/collectors/pre-recon-collector.ts similarity index 77% rename from apps/worker/src/mcp-server/pre-recon-collector.ts rename to apps/worker/src/collectors/pre-recon-collector.ts index 38b2dfb..2c0ee91 100644 --- a/apps/worker/src/mcp-server/pre-recon-collector.ts +++ b/apps/worker/src/collectors/pre-recon-collector.ts @@ -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; - details: Record; - isError?: boolean; +function toolResult(payload: Record) { + return { + content: [{ type: 'text' as const, text: JSON.stringify(payload, null, 2) }], + details: undefined, + }; } -function successResult(data: Record): ToolResult { - const response = { status: 'success', ...data }; - return { content: [{ type: 'text' as const, text: JSON.stringify(response, null, 2) }], details: {} }; +function successResult(data: Record) { + 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 => { + 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 => { + 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 => { + 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 => { + 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 => { + 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 => { + 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 => { + 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(key: K): PreReconToolStatus { const flagMap: Record = { 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 }), diff --git a/apps/worker/src/mcp-server/recon-collector.ts b/apps/worker/src/collectors/recon-collector.ts similarity index 75% rename from apps/worker/src/mcp-server/recon-collector.ts rename to apps/worker/src/collectors/recon-collector.ts index 0662673..f48a338 100644 --- a/apps/worker/src/mcp-server/recon-collector.ts +++ b/apps/worker/src/collectors/recon-collector.ts @@ -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; -} - -function createToolResult(response: { status: string; [key: string]: unknown }): ToolResult { +function toolResult(payload: Record) { 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): ToolResult { - return createToolResult({ status: 'success', ...data }); +function successResult(data: Record) { + 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(); 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 => { + 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 => { + 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 => { + 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 => { + 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 => { + 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 => { + 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 => { + 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 => { + 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 => { + 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(key: K): ReconToolStatus { const flagMap: Record = { 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 }), diff --git a/apps/worker/src/collectors/schema.ts b/apps/worker/src/collectors/schema.ts new file mode 100644 index 0000000..8b9f240 --- /dev/null +++ b/apps/worker/src/collectors/schema.ts @@ -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(values: T, options: TSchemaOptions = {}) { + return Type.Unsafe({ ...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(schema: T, input: Static): Static { + return Value.Clean(schema, structuredClone(input)) as Static; +} diff --git a/apps/worker/src/mcp-server/vuln-collector.ts b/apps/worker/src/collectors/vuln-collector.ts similarity index 89% rename from apps/worker/src/mcp-server/vuln-collector.ts rename to apps/worker/src/collectors/vuln-collector.ts index ee96f9f..deb08da 100644 --- a/apps/worker/src/mcp-server/vuln-collector.ts +++ b/apps/worker/src/collectors/vuln-collector.ts @@ -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 = new Set(['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 = { injection: InjectionStrategicIntelSchema, xss: XssStrategicIntelSchema, auth: AuthStrategicIntelSchema, ssrf: SsrfStrategicIntelSchema, authz: AuthzStrategicIntelSchema, -} as const; +}; // ============================================================================ // EXPORTED TYPES @@ -335,48 +330,42 @@ export type VulnCallStatus = Readonly>; // RESPONSE HELPERS // ============================================================================ -interface ToolResult { - [x: string]: unknown; - content: Array<{ type: 'text'; text: string }>; - details: Record; - isError: boolean; -} - -function createToolResult(response: { status: string; [key: string]: unknown }): ToolResult { +function toolResult(payload: Record) { 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): ToolResult { - return createToolResult({ status: 'success', ...data }); +function successResult(data: Record) { + 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 => { + 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 => { + 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 => { + 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 => { + 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(key: K): VulnToolStatus { const flagMap: Record = { 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 }), diff --git a/apps/worker/src/paths.ts b/apps/worker/src/paths.ts index 03126d2..daab067 100644 --- a/apps/worker/src/paths.ts +++ b/apps/worker/src/paths.ts @@ -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'); diff --git a/apps/worker/src/scripts/generate-totp.ts b/apps/worker/src/scripts/generate-totp.ts index 09aa74a..7214061 100644 --- a/apps/worker/src/scripts/generate-totp.ts +++ b/apps/worker/src/scripts/generate-totp.ts @@ -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( diff --git a/apps/worker/src/services/agent-execution.ts b/apps/worker/src/services/agent-execution.ts index c70f10f..2e8385e 100644 --- a/apps/worker/src/services/agent-execution.ts +++ b/apps/worker/src/services/agent-execution.ts @@ -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; + cancellationSignal?: AbortSignal | undefined; } interface FailAgentOpts { @@ -71,6 +71,48 @@ interface FailAgentOpts { context: Record; } +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 => { + // 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> { - 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)); } /** diff --git a/apps/worker/src/services/agent-git-paths.ts b/apps/worker/src/services/agent-git-paths.ts new file mode 100644 index 0000000..e3966de --- /dev/null +++ b/apps/worker/src/services/agent-git-paths.ts @@ -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)]; +} diff --git a/apps/worker/src/services/error-handling.ts b/apps/worker/src/services/error-handling.ts index e6e1b0d..39a8c7f 100644 --- a/apps/worker/src/services/error-handling.ts +++ b/apps/worker/src/services/error-handling.ts @@ -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 diff --git a/apps/worker/src/services/exploit-renderer.ts b/apps/worker/src/services/exploit-renderer.ts index 4b352ca..6c87bd5 100644 --- a/apps/worker/src/services/exploit-renderer.ts +++ b/apps/worker/src/services/exploit-renderer.ts @@ -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 { + 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`; } diff --git a/apps/worker/src/services/exploitation-checker.ts b/apps/worker/src/services/exploitation-checker.ts index 3786243..79238cc 100644 --- a/apps/worker/src/services/exploitation-checker.ts +++ b/apps/worker/src/services/exploitation-checker.ts @@ -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 { 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; } } diff --git a/apps/worker/src/services/git-manager.ts b/apps/worker/src/services/git-manager.ts index de8ebd1..0b47cf5 100644 --- a/apps/worker/src/services/git-manager.ts +++ b/apps/worker/src/services/git-manager.ts @@ -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 { 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 { - 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 { + 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 `XYPATH\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 { + 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(); + +/** + * 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(operation: () => Promise): Promise { + 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 { // 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 { // 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 => { + // 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 { // 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 => { + 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> { - // 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> { // 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> { // 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; } diff --git a/apps/worker/src/services/prompt-manager.ts b/apps/worker/src/services/prompt-manager.ts index 68d1b5a..3d8d45e 100644 --- a/apps/worker/src/services/prompt-manager.ts +++ b/apps/worker/src/services/prompt-manager.ts @@ -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 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(/[\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(/[\s\S]*?<\/rules_of_engagement>\s*/g, ''); } @@ -325,35 +348,35 @@ async function interpolateVariables( if (!config?.authentication) { result = result.replace(/[\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'); diff --git a/apps/worker/src/services/recon-renderer.ts b/apps/worker/src/services/recon-renderer.ts index 8d37135..9b00be9 100644 --- a/apps/worker/src/services/recon-renderer.ts +++ b/apps/worker/src/services/recon-renderer.ts @@ -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']; diff --git a/apps/worker/src/services/validate-authentication.ts b/apps/worker/src/services/validate-authentication.ts index 1f3e870..9732f26 100644 --- a/apps/worker/src/services/validate-authentication.ts +++ b/apps/worker/src/services/validate-authentication.ts @@ -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> { @@ -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, ): Result { if (!result.success) { diff --git a/apps/worker/src/services/vuln-renderer.ts b/apps/worker/src/services/vuln-renderer.ts index 0f8eae5..70daa15 100644 --- a/apps/worker/src/services/vuln-renderer.ts +++ b/apps/worker/src/services/vuln-renderer.ts @@ -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`; } diff --git a/apps/worker/src/temporal/activities.ts b/apps/worker/src/temporal/activities.ts index f1628f3..182af19 100644 --- a/apps/worker/src/temporal/activities.ts +++ b/apps/worker/src/temporal/activities.ts @@ -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 { - 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 => { const logger = createActivityLogger(); @@ -276,10 +272,10 @@ export async function runPreReconAgent(input: ActivityInput): Promise { - 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 => { const logger = createActivityLogger(); @@ -302,7 +298,7 @@ async function runVulnAgentWithCollector( vulnClass: 'injection' | 'xss' | 'auth' | 'ssrf' | 'authz', input: ActivityInput, ): Promise { - 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(queuePath); + let doc: ExploitQueueDocument; + try { + doc = await readJson(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 { - 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 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(sessionPath); + let session: SessionJson; + try { + session = await readJson(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(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, diff --git a/apps/worker/src/temporal/pipeline.ts b/apps/worker/src/temporal/pipeline.ts index 911ffc3..5e17560 100644 --- a/apps/worker/src/temporal/pipeline.ts +++ b/apps/worker/src/temporal/pipeline.ts @@ -14,4 +14,5 @@ export type { ResumeState, VulnExploitPipelineResult, } from './shared.js'; +export { PipelineExecutionError } from './shared.js'; export { pentestPipeline } from './workflows.js'; diff --git a/apps/worker/src/temporal/shared.ts b/apps/worker/src/temporal/shared.ts index 9ce9cf9..0712a40 100644 --- a/apps/worker/src/temporal/shared.ts +++ b/apps/worker/src/temporal/shared.ts @@ -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: { diff --git a/apps/worker/src/temporal/summary-mapper.ts b/apps/worker/src/temporal/summary-mapper.ts index d8a92c8..3979898 100644 --- a/apps/worker/src/temporal/summary-mapper.ts +++ b/apps/worker/src/temporal/summary-mapper.ts @@ -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) { diff --git a/apps/worker/src/temporal/worker.ts b/apps/worker/src/temporal/worker.ts index 1eadb0a..c752125 100644 --- a/apps/worker/src/temporal/worker.ts +++ b/apps/worker/src/temporal/worker.ts @@ -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); } } diff --git a/apps/worker/src/temporal/workflows.ts b/apps/worker/src/temporal/workflows.ts index 53039ff..da3a4b3 100644 --- a/apps/worker/src/temporal/workflows.ts +++ b/apps/worker/src/temporal/workflows.ts @@ -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[]): 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[]): void { - const failedPipelines: string[] = []; + function aggregatePipelineResults( + results: PromiseSettledResult[], + 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 Promise> = []; + 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 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 { + 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; - 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; } diff --git a/apps/worker/src/utils/billing-detection.ts b/apps/worker/src/utils/billing-detection.ts index 7d41354..5509292 100644 --- a/apps/worker/src/utils/billing-detection.ts +++ b/apps/worker/src/utils/billing-detection.ts @@ -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: diff --git a/apps/worker/src/utils/browser-agents.ts b/apps/worker/src/utils/browser-agents.ts new file mode 100644 index 0000000..aa18544 --- /dev/null +++ b/apps/worker/src/utils/browser-agents.ts @@ -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 = 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); +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 7e91e17..e67f124 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -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