refactor: consolidate file layout and break circular dependencies
- Move error-handling, git-manager, prompt-manager, queue-validation, and reporting into src/services/ - Delete src/constants.ts — relocate AGENT_VALIDATORS and MCP_AGENT_MAPPING into session-manager.ts alongside agent definitions - Delete src/utils/output-formatter.ts — absorb filterJsonToolCalls and getAgentPrefix into ai/output-formatters.ts - Extract ActivityLogger interface into src/types/activity-logger.ts to break temporal/ → services circular dependency - Consolidate VulnType, ExploitationDecision into types/agents.ts and SessionMetadata into types/audit.ts - Remove dead timingResults/costResults globals from utils/metrics.ts and all consumers
This commit is contained in:
ajmallesh
@@ -1,319 +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.
|
||||
|
||||
import { $ } from 'zx';
|
||||
import { PentestError } from '../error-handling.js';
|
||||
import { ErrorCode } from '../types/errors.js';
|
||||
import type { ActivityLogger } from '../temporal/activity-logger.js';
|
||||
|
||||
/**
|
||||
* Check if a directory is a git repository.
|
||||
* Returns true if the directory contains a .git folder or is inside a git repo.
|
||||
*/
|
||||
export async function isGitRepository(dir: string): Promise<boolean> {
|
||||
try {
|
||||
await $`cd ${dir} && git rev-parse --git-dir`.quiet();
|
||||
return true;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
interface GitOperationResult {
|
||||
success: boolean;
|
||||
hadChanges?: boolean;
|
||||
error?: Error;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get list of changed files from git status --porcelain output
|
||||
*/
|
||||
async function getChangedFiles(
|
||||
sourceDir: string,
|
||||
operationDescription: string
|
||||
): Promise<string[]> {
|
||||
const status = await executeGitCommandWithRetry(
|
||||
['git', 'status', '--porcelain'],
|
||||
sourceDir,
|
||||
operationDescription
|
||||
);
|
||||
return status.stdout
|
||||
.trim()
|
||||
.split('\n')
|
||||
.filter((line) => line.length > 0);
|
||||
}
|
||||
|
||||
/**
|
||||
* Log a summary of changed files with truncation for long lists
|
||||
*/
|
||||
function logChangeSummary(
|
||||
changes: string[],
|
||||
messageWithChanges: string,
|
||||
messageWithoutChanges: string,
|
||||
logger: ActivityLogger,
|
||||
level: 'info' | 'warn' = 'info',
|
||||
maxToShow: number = 5
|
||||
): void {
|
||||
if (changes.length > 0) {
|
||||
const msg = messageWithChanges.replace('{count}', String(changes.length));
|
||||
const fileList = changes.slice(0, maxToShow).map((c) => ` ${c}`).join(', ');
|
||||
const suffix = changes.length > maxToShow
|
||||
? ` ... and ${changes.length - maxToShow} more files`
|
||||
: '';
|
||||
logger[level](`${msg} ${fileList}${suffix}`);
|
||||
} else {
|
||||
logger[level](messageWithoutChanges);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert unknown error to GitOperationResult
|
||||
*/
|
||||
function toErrorResult(error: unknown): GitOperationResult {
|
||||
const errMsg = error instanceof Error ? error.message : String(error);
|
||||
return {
|
||||
success: false,
|
||||
error: error instanceof Error ? error : new Error(errMsg),
|
||||
};
|
||||
}
|
||||
|
||||
// Serializes git operations to prevent index.lock conflicts during parallel agent execution
|
||||
class GitSemaphore {
|
||||
private queue: Array<() => void> = [];
|
||||
private running: boolean = false;
|
||||
|
||||
async acquire(): Promise<void> {
|
||||
return new Promise((resolve) => {
|
||||
this.queue.push(resolve);
|
||||
this.process();
|
||||
});
|
||||
}
|
||||
|
||||
release(): void {
|
||||
this.running = false;
|
||||
this.process();
|
||||
}
|
||||
|
||||
private process(): void {
|
||||
if (!this.running && this.queue.length > 0) {
|
||||
this.running = true;
|
||||
const resolve = this.queue.shift();
|
||||
resolve!();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const gitSemaphore = new GitSemaphore();
|
||||
|
||||
const GIT_LOCK_ERROR_PATTERNS = [
|
||||
'index.lock',
|
||||
'unable to lock',
|
||||
'Another git process',
|
||||
'fatal: Unable to create',
|
||||
'fatal: index file',
|
||||
];
|
||||
|
||||
function isGitLockError(errorMessage: string): boolean {
|
||||
return GIT_LOCK_ERROR_PATTERNS.some((pattern) => errorMessage.includes(pattern));
|
||||
}
|
||||
|
||||
// Retries git commands on lock conflicts with exponential backoff
|
||||
export async function executeGitCommandWithRetry(
|
||||
commandArgs: string[],
|
||||
sourceDir: string,
|
||||
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 = Math.pow(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();
|
||||
}
|
||||
}
|
||||
|
||||
// Two-phase reset: hard reset (tracked files) + clean (untracked files)
|
||||
export async function rollbackGitWorkspace(
|
||||
sourceDir: string,
|
||||
reason: string = 'retry preparation',
|
||||
logger: ActivityLogger
|
||||
): Promise<GitOperationResult> {
|
||||
// Skip git operations if not a git repository
|
||||
if (!(await isGitRepository(sourceDir))) {
|
||||
logger.info('Skipping git rollback (not a git repository)');
|
||||
return { success: true };
|
||||
}
|
||||
|
||||
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'
|
||||
);
|
||||
|
||||
logChangeSummary(
|
||||
changes,
|
||||
'Rollback completed - removed {count} contaminated changes:',
|
||||
'Rollback completed - no changes to remove',
|
||||
logger,
|
||||
'info',
|
||||
3
|
||||
);
|
||||
return { success: true };
|
||||
} catch (error) {
|
||||
const errMsg = error instanceof Error ? error.message : String(error);
|
||||
logger.error(`Rollback failed after retries: ${errMsg}`);
|
||||
return {
|
||||
success: false,
|
||||
error: new PentestError(
|
||||
`Git rollback failed: ${errMsg}`,
|
||||
'filesystem',
|
||||
false, // Non-retryable - rollback is best-effort cleanup
|
||||
{ sourceDir, reason },
|
||||
ErrorCode.GIT_ROLLBACK_FAILED
|
||||
),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Creates checkpoint before each attempt. First attempt preserves workspace; retries clean it.
|
||||
export async function createGitCheckpoint(
|
||||
sourceDir: string,
|
||||
description: string,
|
||||
attempt: number,
|
||||
logger: ActivityLogger
|
||||
): Promise<GitOperationResult> {
|
||||
// Skip git operations if not a git repository
|
||||
if (!(await isGitRepository(sourceDir))) {
|
||||
logger.info('Skipping git checkpoint (not a git repository)');
|
||||
return { success: true };
|
||||
}
|
||||
|
||||
logger.info(`Creating checkpoint for ${description} (attempt ${attempt})`);
|
||||
try {
|
||||
// First attempt: preserve existing deliverables. Retries: clean workspace to prevent pollution
|
||||
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 changes = await getChangedFiles(sourceDir, 'status check');
|
||||
const hasChanges = changes.length > 0;
|
||||
|
||||
await executeGitCommandWithRetry(['git', 'add', '-A'], sourceDir, 'staging changes');
|
||||
await executeGitCommandWithRetry(
|
||||
['git', 'commit', '-m', `📍 Checkpoint: ${description} (attempt ${attempt})`, '--allow-empty'],
|
||||
sourceDir,
|
||||
'creating commit'
|
||||
);
|
||||
|
||||
if (hasChanges) {
|
||||
logger.info('Checkpoint created with uncommitted changes staged');
|
||||
} else {
|
||||
logger.info('Empty checkpoint created (no workspace changes)');
|
||||
}
|
||||
return { success: true };
|
||||
} catch (error) {
|
||||
const result = toErrorResult(error);
|
||||
logger.warn(`Checkpoint creation failed after retries: ${result.error?.message}`);
|
||||
return result;
|
||||
}
|
||||
}
|
||||
|
||||
export async function commitGitSuccess(
|
||||
sourceDir: string,
|
||||
description: string,
|
||||
logger: ActivityLogger
|
||||
): Promise<GitOperationResult> {
|
||||
// Skip git operations if not a git repository
|
||||
if (!(await isGitRepository(sourceDir))) {
|
||||
logger.info('Skipping git commit (not a git repository)');
|
||||
return { success: true };
|
||||
}
|
||||
|
||||
logger.info(`Committing successful results for ${description}`);
|
||||
try {
|
||||
const changes = await getChangedFiles(sourceDir, 'status check for success commit');
|
||||
|
||||
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'
|
||||
);
|
||||
|
||||
logChangeSummary(
|
||||
changes,
|
||||
'Success commit created with {count} file changes:',
|
||||
'Empty success commit created (agent made no file changes)',
|
||||
logger
|
||||
);
|
||||
return { success: true };
|
||||
} catch (error) {
|
||||
const result = toErrorResult(error);
|
||||
logger.warn(`Success commit failed after retries: ${result.error?.message}`);
|
||||
return result;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get current git commit hash.
|
||||
* Returns null if not a git repository.
|
||||
*/
|
||||
export async function getGitCommitHash(sourceDir: string): Promise<string | null> {
|
||||
if (!(await isGitRepository(sourceDir))) {
|
||||
return null;
|
||||
}
|
||||
try {
|
||||
const result = await $`cd ${sourceDir} && git rev-parse HEAD`;
|
||||
return result.stdout.trim();
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
@@ -26,33 +26,3 @@ export class Timer {
|
||||
return end - this.startTime;
|
||||
}
|
||||
}
|
||||
|
||||
interface TimingResultsAgents {
|
||||
[key: string]: number;
|
||||
}
|
||||
|
||||
interface TimingResults {
|
||||
total: Timer | null;
|
||||
agents: TimingResultsAgents;
|
||||
}
|
||||
|
||||
interface CostResultsAgents {
|
||||
[key: string]: number;
|
||||
}
|
||||
|
||||
interface CostResults {
|
||||
agents: CostResultsAgents;
|
||||
total: number;
|
||||
}
|
||||
|
||||
// Global timing and cost tracker
|
||||
export const timingResults: TimingResults = {
|
||||
total: null,
|
||||
agents: {},
|
||||
};
|
||||
|
||||
export const costResults: CostResults = {
|
||||
agents: {},
|
||||
total: 0,
|
||||
};
|
||||
|
||||
|
||||
@@ -1,264 +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.
|
||||
|
||||
import { AGENTS } from '../session-manager.js';
|
||||
|
||||
interface ToolCallInput {
|
||||
url?: string;
|
||||
element?: string;
|
||||
key?: string;
|
||||
fields?: unknown[];
|
||||
text?: string;
|
||||
action?: string;
|
||||
description?: string;
|
||||
todos?: Array<{
|
||||
status: string;
|
||||
content: string;
|
||||
}>;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
interface ToolCall {
|
||||
name: string;
|
||||
input?: ToolCallInput;
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract domain from URL for display
|
||||
*/
|
||||
function extractDomain(url: string): string {
|
||||
try {
|
||||
const urlObj = new URL(url);
|
||||
return urlObj.hostname || url.slice(0, 30);
|
||||
} catch {
|
||||
return url.slice(0, 30);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Summarize TodoWrite updates into clean progress indicators
|
||||
*/
|
||||
function summarizeTodoUpdate(input: ToolCallInput | undefined): string | null {
|
||||
if (!input?.todos || !Array.isArray(input.todos)) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const todos = input.todos;
|
||||
const completed = todos.filter((t) => t.status === 'completed');
|
||||
const inProgress = todos.filter((t) => t.status === 'in_progress');
|
||||
|
||||
// Show recently completed tasks
|
||||
if (completed.length > 0) {
|
||||
const recent = completed[completed.length - 1]!;
|
||||
return `✅ ${recent.content}`;
|
||||
}
|
||||
|
||||
// Show current in-progress task
|
||||
if (inProgress.length > 0) {
|
||||
const current = inProgress[0]!;
|
||||
return `🔄 ${current.content}`;
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get agent prefix for parallel execution
|
||||
*/
|
||||
export function getAgentPrefix(description: string): string {
|
||||
// Map agent names to their prefixes
|
||||
const agentPrefixes: Record<string, string> = {
|
||||
'injection-vuln': '[Injection]',
|
||||
'xss-vuln': '[XSS]',
|
||||
'auth-vuln': '[Auth]',
|
||||
'authz-vuln': '[Authz]',
|
||||
'ssrf-vuln': '[SSRF]',
|
||||
'injection-exploit': '[Injection]',
|
||||
'xss-exploit': '[XSS]',
|
||||
'auth-exploit': '[Auth]',
|
||||
'authz-exploit': '[Authz]',
|
||||
'ssrf-exploit': '[SSRF]',
|
||||
};
|
||||
|
||||
// First try to match by agent name directly
|
||||
for (const [agentName, prefix] of Object.entries(agentPrefixes)) {
|
||||
const agent = AGENTS[agentName as keyof typeof AGENTS];
|
||||
if (agent && description.includes(agent.displayName)) {
|
||||
return prefix;
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to partial matches for backwards compatibility
|
||||
if (description.includes('injection')) return '[Injection]';
|
||||
if (description.includes('xss')) return '[XSS]';
|
||||
if (description.includes('authz')) return '[Authz]'; // Check authz before auth
|
||||
if (description.includes('auth')) return '[Auth]';
|
||||
if (description.includes('ssrf')) return '[SSRF]';
|
||||
|
||||
return '[Agent]';
|
||||
}
|
||||
|
||||
/**
|
||||
* Format browser tool calls into clean progress indicators
|
||||
*/
|
||||
function formatBrowserAction(toolCall: ToolCall): string {
|
||||
const toolName = toolCall.name;
|
||||
const input = toolCall.input || {};
|
||||
|
||||
// Core Browser Operations
|
||||
if (toolName === 'mcp__playwright__browser_navigate') {
|
||||
const url = input.url || '';
|
||||
const domain = extractDomain(url);
|
||||
return `🌐 Navigating to ${domain}`;
|
||||
}
|
||||
|
||||
if (toolName === 'mcp__playwright__browser_navigate_back') {
|
||||
return `⬅️ Going back`;
|
||||
}
|
||||
|
||||
// Page Interaction
|
||||
if (toolName === 'mcp__playwright__browser_click') {
|
||||
const element = input.element || 'element';
|
||||
return `🖱️ Clicking ${element.slice(0, 25)}`;
|
||||
}
|
||||
|
||||
if (toolName === 'mcp__playwright__browser_hover') {
|
||||
const element = input.element || 'element';
|
||||
return `👆 Hovering over ${element.slice(0, 20)}`;
|
||||
}
|
||||
|
||||
if (toolName === 'mcp__playwright__browser_type') {
|
||||
const element = input.element || 'field';
|
||||
return `⌨️ Typing in ${element.slice(0, 20)}`;
|
||||
}
|
||||
|
||||
if (toolName === 'mcp__playwright__browser_press_key') {
|
||||
const key = input.key || 'key';
|
||||
return `⌨️ Pressing ${key}`;
|
||||
}
|
||||
|
||||
// Form Handling
|
||||
if (toolName === 'mcp__playwright__browser_fill_form') {
|
||||
const fieldCount = input.fields?.length || 0;
|
||||
return `📝 Filling ${fieldCount} form fields`;
|
||||
}
|
||||
|
||||
if (toolName === 'mcp__playwright__browser_select_option') {
|
||||
return `📋 Selecting dropdown option`;
|
||||
}
|
||||
|
||||
if (toolName === 'mcp__playwright__browser_file_upload') {
|
||||
return `📁 Uploading file`;
|
||||
}
|
||||
|
||||
// Page Analysis
|
||||
if (toolName === 'mcp__playwright__browser_snapshot') {
|
||||
return `📸 Taking page snapshot`;
|
||||
}
|
||||
|
||||
if (toolName === 'mcp__playwright__browser_take_screenshot') {
|
||||
return `📸 Taking screenshot`;
|
||||
}
|
||||
|
||||
if (toolName === 'mcp__playwright__browser_evaluate') {
|
||||
return `🔍 Running JavaScript analysis`;
|
||||
}
|
||||
|
||||
// Waiting & Monitoring
|
||||
if (toolName === 'mcp__playwright__browser_wait_for') {
|
||||
if (input.text) {
|
||||
return `⏳ Waiting for "${input.text.slice(0, 20)}"`;
|
||||
}
|
||||
return `⏳ Waiting for page response`;
|
||||
}
|
||||
|
||||
if (toolName === 'mcp__playwright__browser_console_messages') {
|
||||
return `📜 Checking console logs`;
|
||||
}
|
||||
|
||||
if (toolName === 'mcp__playwright__browser_network_requests') {
|
||||
return `🌐 Analyzing network traffic`;
|
||||
}
|
||||
|
||||
// Tab Management
|
||||
if (toolName === 'mcp__playwright__browser_tabs') {
|
||||
const action = input.action || 'managing';
|
||||
return `🗂️ ${action} browser tab`;
|
||||
}
|
||||
|
||||
// Dialog Handling
|
||||
if (toolName === 'mcp__playwright__browser_handle_dialog') {
|
||||
return `💬 Handling browser dialog`;
|
||||
}
|
||||
|
||||
// Fallback for any missed tools
|
||||
const actionType = toolName.split('_').pop();
|
||||
return `🌐 Browser: ${actionType}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Filter out JSON tool calls from content, with special handling for Task calls
|
||||
*/
|
||||
export function filterJsonToolCalls(content: string | null | undefined): string {
|
||||
if (!content || typeof content !== 'string') {
|
||||
return content || '';
|
||||
}
|
||||
|
||||
const lines = content.split('\n');
|
||||
const processedLines: string[] = [];
|
||||
|
||||
for (const line of lines) {
|
||||
const trimmed = line.trim();
|
||||
|
||||
// Skip empty lines
|
||||
if (trimmed === '') {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check if this is a JSON tool call
|
||||
if (trimmed.startsWith('{"type":"tool_use"')) {
|
||||
try {
|
||||
const toolCall = JSON.parse(trimmed) as ToolCall;
|
||||
|
||||
// Special handling for Task tool calls
|
||||
if (toolCall.name === 'Task') {
|
||||
const description = toolCall.input?.description || 'analysis agent';
|
||||
processedLines.push(`🚀 Launching ${description}`);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Special handling for TodoWrite tool calls
|
||||
if (toolCall.name === 'TodoWrite') {
|
||||
const summary = summarizeTodoUpdate(toolCall.input);
|
||||
if (summary) {
|
||||
processedLines.push(summary);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Special handling for browser tool calls
|
||||
if (toolCall.name.startsWith('mcp__playwright__browser_')) {
|
||||
const browserAction = formatBrowserAction(toolCall);
|
||||
if (browserAction) {
|
||||
processedLines.push(browserAction);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Hide all other tool calls (Read, Write, Grep, etc.)
|
||||
continue;
|
||||
} catch {
|
||||
// If JSON parsing fails, treat as regular text
|
||||
processedLines.push(line);
|
||||
}
|
||||
} else {
|
||||
// Keep non-JSON lines (assistant text)
|
||||
processedLines.push(line);
|
||||
}
|
||||
}
|
||||
|
||||
return processedLines.join('\n');
|
||||
}
|
||||
Reference in New Issue
Block a user