From 438d8ec138c3b0edad55441ddbd13969aae39e48 Mon Sep 17 00:00:00 2001 From: j3ssie Date: Tue, 10 Feb 2026 08:44:48 +0700 Subject: [PATCH] feat: implement agent executor with tool calling, sub-agents, and comprehensive test suite - Add AgentExecutor implementing LLM-based agentic loop with tool calling, max iterations, and stop conditions - Introduce agent preset tools (bash, file_exists, http_get, run_module, etc.) with extensible registry pattern - Add sub-agent spawning capability via spawn_agent tool call with recursive depth limits and validation - Implement ToolExecutor for custom tool execution with template rendering and error handling - Add agent session persistence and memory management with sliding window configuration - Create comprehensive E2E test suite covering 15+ agent workflow scenarios (minimal, custom tools, planning, multi-goal, structured output, tracing hooks, file tools, orchestration, Python tools, sub-agents, nested sub-agents, and validation) - Add agent-and-llm test data directory with 17 YAML workflow fixtures - Update integration tests to include agent workflow directories - Add AgentTool and AgentConfig types with validation for duplicate names and unknown presets - Implement LLM streaming test utilities - Update documentation (CLAUDE.md, HACKING.md, README.md) with agent features and CLI examples --- CLAUDE.md | 48 +- HACKING.md | 199 +- README.md | 5 +- internal/broker/redis_event.go | 2 +- internal/core/agent_tool_presets.go | 396 ++++ internal/core/agent_tool_presets_test.go | 230 ++ internal/core/agent_types.go | 92 + internal/core/agent_types_test.go | 346 +++ internal/core/llm_types.go | 23 + internal/core/preferences_test.go | 10 +- internal/core/step.go | 73 +- internal/core/step_test.go | 60 + internal/core/types.go | 1 + internal/database/agent_session.go | 60 + internal/database/database.go | 1 + internal/database/jsonl.go | 2 +- internal/distributed/master.go | 2 +- internal/executor/agent_executor.go | 1220 +++++++++++ internal/executor/agent_executor_test.go | 1862 +++++++++++++++++ internal/executor/dispatcher.go | 277 ++- internal/executor/executor.go | 62 +- internal/executor/llm_executor.go | 277 ++- internal/executor/llm_streaming_test.go | 988 +++++++++ internal/executor/tool_executor.go | 306 +++ internal/executor/tool_executor_test.go | 284 +++ internal/functions/constants.go | 61 +- internal/functions/db_functions.go | 2 +- internal/functions/file_functions.go | 103 + internal/functions/file_functions_test.go | 219 ++ internal/functions/goja_runtime.go | 5 + internal/functions/url_functions_test.go | 10 +- internal/functions/util_functions.go | 145 ++ internal/functions/util_functions_test.go | 148 ++ internal/linter/rules.go | 190 ++ internal/parser/parser.go | 78 + internal/terminal/printer.go | 4 +- internal/terminal/symbols.go | 5 + pkg/cli/root.go | 15 +- test/e2e/agent_test.go | 709 +++++++ test/e2e/e2e_test.go | 6 + test/integration/workflow_test.go | 32 +- .../sample-jsonl-output/semgrep-data.json | 1 + .../test-agent-custom-tools.yaml | 33 + .../agent-and-llm/test-agent-exports.yaml | 30 + .../agent-and-llm/test-agent-file-tools.yaml | 31 + .../agent-and-llm/test-agent-minimal.yaml | 19 + .../agent-and-llm/test-agent-multi-goal.yaml | 28 + .../test-agent-orchestration.yaml | 29 + .../agent-and-llm/test-agent-planning.yaml | 28 + .../test-agent-python-tools.yaml | 29 + .../agent-and-llm/test-agent-structured.yaml | 21 + .../test-agent-sub-agents-nested.yaml | 57 + ...test-agent-sub-agents-validation-fail.yaml | 28 + .../agent-and-llm/test-agent-sub-agents.yaml | 36 + .../agent-and-llm/test-agent-tracing.yaml | 22 + .../test-agent-unknown-preset.yaml | 21 + .../test-agent-validation-fail.yaml | 22 + .../workflows/agent-and-llm/test-agent.yaml | 96 + .../{ => agent-and-llm}/test-llm.yaml | 0 59 files changed, 8983 insertions(+), 106 deletions(-) create mode 100644 internal/core/agent_tool_presets.go create mode 100644 internal/core/agent_tool_presets_test.go create mode 100644 internal/core/agent_types.go create mode 100644 internal/core/agent_types_test.go create mode 100644 internal/core/step_test.go create mode 100644 internal/database/agent_session.go create mode 100644 internal/executor/agent_executor.go create mode 100644 internal/executor/agent_executor_test.go create mode 100644 internal/executor/llm_streaming_test.go create mode 100644 internal/executor/tool_executor.go create mode 100644 internal/executor/tool_executor_test.go create mode 100644 test/e2e/agent_test.go create mode 100644 test/testdata/sample-jsonl-output/semgrep-data.json create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-custom-tools.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-exports.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-file-tools.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-minimal.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-multi-goal.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-orchestration.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-planning.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-python-tools.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-structured.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-sub-agents-nested.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-sub-agents-validation-fail.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-sub-agents.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-tracing.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-unknown-preset.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent-validation-fail.yaml create mode 100644 test/testdata/workflows/agent-and-llm/test-agent.yaml rename test/testdata/workflows/{ => agent-and-llm}/test-llm.yaml (100%) diff --git a/CLAUDE.md b/CLAUDE.md index 575eb22..2896eff 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -51,7 +51,7 @@ CLI/API (pkg/cli, pkg/server) ↓ Executor (internal/executor) - coordinates workflow execution ↓ -StepDispatcher - routes to: BashExecutor, FunctionExecutor, ForeachExecutor, ParallelExecutor, RemoteBashExecutor, HTTPExecutor, LLMExecutor +StepDispatcher - routes to: BashExecutor, FunctionExecutor, ForeachExecutor, ParallelExecutor, RemoteBashExecutor, HTTPExecutor, LLMExecutor, AgentExecutor ↓ Runner (internal/runner) - executes commands via: HostRunner, DockerRunner, SSHRunner ``` @@ -79,7 +79,7 @@ Runner (internal/runner) - executes commands via: HostRunner, DockerRunner, SSHR ```go WorkflowKind: "module" | "flow" // module = single unit, flow = orchestrates modules -StepType: "bash" | "function" | "parallel-steps" | "foreach" | "remote-bash" | "http" | "llm" +StepType: "bash" | "function" | "parallel-steps" | "foreach" | "remote-bash" | "http" | "llm" | "agent" RunnerType: "host" | "docker" | "ssh" TriggerType: "cron" | "event" | "watch" | "manual" ``` @@ -110,7 +110,7 @@ Use `goto: _end` to terminate workflow. - `{{Variable}}` - standard template variables (Target, Output, threads, etc.) - `[[variable]]` - foreach loop variables (to avoid conflicts) -- Functions evaluated via Goja JS runtime: `file_exists()`, `file_length()`, `trim()`, etc. +- Functions evaluated via Goja JS runtime: `file_exists()`, `file_length()`, `trim()`, `exec_python()`, etc. ### Platform Variables @@ -121,6 +121,46 @@ Built-in variables for environment detection: - `{{PlatformInKubernetes}}` - "true" if running in Kubernetes pod - `{{PlatformCloudProvider}}` - Cloud provider (aws, gcp, azure, local) +### Agent Step Type + +The `agent` step type provides an agentic LLM execution loop with tool calling, sub-agent orchestration, and memory management. + +Key YAML fields: +- `query` / `queries` - Task prompt (single or multi-goal) +- `agent_tools` - List of preset or custom tools available to the agent +- `max_iterations` - Maximum tool-calling loop iterations (required, > 0) +- `system_prompt` - System prompt for the agent +- `sub_agents` - Inline sub-agents spawnable via `spawn_agent` tool call +- `memory` - Sliding window config (`max_messages`, `summarize_on_truncate`, `persist_path`, `resume_path`) +- `models` - Preferred models tried in order before falling back to default +- `output_schema` - JSON schema enforced on final output +- `plan_prompt` - Optional planning stage prompt run before the main loop +- `stop_condition` - JS expression evaluated after each iteration +- `on_tool_start` / `on_tool_end` - JS hook expressions for tool call tracing +- `parallel_tool_calls` - Enable/disable parallel tool execution (default: true) + +Preset tools: `bash`, `read_file`, `read_lines`, `file_exists`, `file_length`, `append_file`, `save_content`, `glob`, `grep_string`, `grep_regex`, `http_get`, `http_request`, `jq`, `exec_python`, `exec_python_file`, `run_module`, `run_flow` + +Available exports: `agent_content`, `agent_history`, `agent_iterations`, `agent_total_tokens`, `agent_prompt_tokens`, `agent_completion_tokens`, `agent_tool_results`, `agent_plan`, `agent_goal_results` + +```yaml +steps: + - name: analyze-target + type: agent + query: "Enumerate subdomains of {{Target}} and summarize findings." + system_prompt: "You are a security reconnaissance agent." + max_iterations: 10 + agent_tools: + - preset: bash + - preset: read_file + - preset: save_content + memory: + max_messages: 30 + persist_path: "{{Output}}/agent/conversation.json" + exports: + findings: "{{agent_content}}" +``` + ## CLI Commands ```bash @@ -208,6 +248,8 @@ REST API documentation with curl examples is in `docs/api/`. Key endpoint catego **New Utility Function**: Add Go implementation in `internal/functions/`, register in `goja_runtime.go` +**New Agent Preset Tool**: Add to `PresetToolRegistry` in `internal/core/agent_tool_presets.go`, add case in `buildPresetCallExpr()` in `internal/executor/agent_executor.go` + ## Architecture Notes - **Executor**: Fresh instances created per target/request - no global singleton diff --git a/HACKING.md b/HACKING.md index 6995b84..6b56fea 100644 --- a/HACKING.md +++ b/HACKING.md @@ -8,6 +8,7 @@ This document describes the technical architecture and development practices for - [Architecture Overview](#architecture-overview) - [Core Components](#core-components) - [Workflow Engine](#workflow-engine) +- [Agent Step Type](#agent-step-type) - [Execution Pipeline](#execution-pipeline) - [Runner System](#runner-system) - [Authentication Middleware](#authentication-middleware) @@ -128,7 +129,7 @@ type Workflow struct { type Step struct { Name string - Type StepType // bash, function, foreach, parallel-steps, remote-bash, http, llm + Type StepType // bash, function, foreach, parallel-steps, remote-bash, http, llm, agent PreCondition string // Skip condition Command string // For bash/remote-bash Commands []string // Multiple commands @@ -140,6 +141,23 @@ type Step struct { ParallelSteps []Step // For parallel-steps type StepRunner RunnerType // For remote-bash: docker or ssh StepRunnerConfig *StepRunnerConfig // Runner config for remote-bash + + // Agent step fields + Query string // Task prompt for the agent + Queries []string // Multiple queries (multi-goal mode) + SystemPrompt string // System prompt for the agent + AgentTools []AgentToolDef // Preset or custom tools + MaxIterations int // Max tool-calling loop iterations + Models []string // Preferred models (tried in order) + SubAgents []SubAgentDef // Inline sub-agents spawnable via spawn_agent + MaxAgentDepth int // Max nesting depth for sub-agents (default: 3) + Memory *AgentMemoryConfig // Sliding window, summarization, persistence + OutputSchema string // JSON schema for structured output + StopCondition string // JS expression evaluated after each iteration + PlanPrompt string // Planning stage prompt + OnToolStart string // JS hook before each tool call + OnToolEnd string // JS hook after each tool call + Exports map[string]string OnSuccess []Action OnError []Action @@ -208,8 +226,175 @@ steps: The `_end` special value terminates workflow execution from the current step. +## Agent Step Type + +The `agent` step type implements an agentic LLM execution loop. It sends a query to the LLM with available tools, executes tool calls returned by the LLM, feeds results back, and repeats until the LLM responds without tool calls or `max_iterations` is reached. + +### Execution Flow + +``` +1. Planning stage (optional) ──▶ LLM generates a plan from plan_prompt +2. Initialize conversation ──▶ system_prompt + query + plan (if any) +3. Main agent loop: + a. Send conversation to LLM (with tools) + b. If no tool_calls → done + c. Execute tool calls (parallel or sequential) + d. Append tool results to conversation + e. Evaluate stop_condition (if defined) + f. Apply memory window (if configured) + g. Repeat until max_iterations +4. Structured output (optional) ──▶ Final LLM call with output_schema +5. Persist conversation (if memory.persist_path set) +``` + +### YAML Structure + +```yaml +steps: + - name: my-agent + type: agent + query: "Analyze {{Target}} and report findings." + system_prompt: "You are a security analyst." + max_iterations: 10 + agent_tools: + - preset: bash + - preset: read_file + - preset: http_get + - name: custom_tool + description: "My custom tool" + parameters: + type: object + properties: + input: + type: string + required: [input] + handler: 'process(args.input)' + models: + - gpt-4o + - claude-sonnet-4-20250514 + memory: + max_messages: 30 + summarize_on_truncate: true + persist_path: "{{Output}}/agent/conversation.json" + resume_path: "{{Output}}/agent/conversation.json" + stop_condition: 'contains(agent_content, "DONE")' + output_schema: '{"type":"object","properties":{"summary":{"type":"string"}}}' + plan_prompt: "Create a plan for analyzing the target." + on_tool_start: 'log_info("Tool: " + tool_name)' + on_tool_end: 'log_info("Result: " + result)' + parallel_tool_calls: true + exports: + findings: "{{agent_content}}" +``` + +### Preset Tools + +All preset tools are defined in `PresetToolRegistry` (`internal/core/agent_tool_presets.go`): + +| Preset | Description | Parameters | +|--------|-------------|------------| +| `bash` | Execute a shell command | `command` | +| `read_file` | Read file contents | `path` | +| `read_lines` | Read file as array of lines | `path` | +| `file_exists` | Check if a file exists | `path` | +| `file_length` | Count non-empty lines in a file | `path` | +| `append_file` | Append content from source to dest | `dest`, `content` | +| `save_content` | Write string content to a file | `content`, `path` | +| `glob` | Find files matching a glob pattern | `pattern` | +| `grep_string` | Search file for lines containing a string | `source`, `str` | +| `grep_regex` | Search file for lines matching a regex | `source`, `pattern` | +| `http_get` | Make an HTTP GET request | `url` | +| `http_request` | Make an HTTP request with method/headers/body | `url`, `method`, `body`?, `headers`? | +| `jq` | Query JSON data using jq syntax | `json_data`, `expression` | +| `exec_python` | Run inline Python code | `code` | +| `exec_python_file` | Run a Python file | `path` | +| `run_module` | Run an osmedeus module | `module`, `target`, `params`? | +| `run_flow` | Run an osmedeus flow | `flow`, `target`, `params`? | + +### Custom Tool Definition + +Custom tools use a JS handler expression. The parsed arguments are available as the `args` object: + +```yaml +agent_tools: + - name: check_domain + description: "Validate if a string is a valid domain" + parameters: + type: object + properties: + domain: + type: string + required: [domain] + handler: 'contains(args.domain, ".")' +``` + +### Sub-Agent Orchestration + +Agents can delegate to sub-agents via the `spawn_agent` tool (automatically added when `sub_agents` is defined): + +```yaml +steps: + - name: orchestrator + type: agent + query: "Analyze {{Target}} by coordinating specialists" + system_prompt: "You are an orchestrator. Delegate tasks to sub-agents." + max_iterations: 10 + max_agent_depth: 3 + agent_tools: + - preset: bash + sub_agents: + - name: recon_agent + description: "Specialized agent for reconnaissance" + system_prompt: "You are a recon specialist" + max_iterations: 5 + agent_tools: + - preset: bash + - preset: http_get + - name: vuln_scanner + description: "Specialized agent for vulnerability scanning" + max_iterations: 5 + agent_tools: + - preset: bash + - preset: read_file +``` + +Sub-agents are implemented via `SubAgentToolExecutor` in `internal/executor/tool_executor.go`. Child token counts are merged into the parent via `agentState.MergeTokens()`. + +### Memory Management + +- **Sliding window**: `max_messages` limits conversation history (system message always kept) +- **Summarization**: `summarize_on_truncate: true` uses LLM to summarize dropped messages +- **Persistence**: `persist_path` saves conversation JSON after completion +- **Resume**: `resume_path` loads a prior conversation on start + +### Available Exports + +| Export | Description | +|--------|-------------| +| `agent_content` | Final text output from the agent | +| `agent_history` | Full conversation history as JSON | +| `agent_iterations` | Number of iterations completed | +| `agent_total_tokens` | Total tokens used (including sub-agents) | +| `agent_prompt_tokens` | Prompt tokens used | +| `agent_completion_tokens` | Completion tokens used | +| `agent_tool_results` | All tool call results as JSON | +| `agent_plan` | Plan generated by planning stage (if used) | +| `agent_goal_results` | Results per query in multi-goal mode (JSON) | + +### Tool Hooks + +Hook expressions receive these variables: +- `tool_name` - Name of the tool being called +- `tool_args` - JSON string of tool arguments +- `result` - Tool result (empty in `on_tool_start`) +- `duration` - Execution time in ms (0 in `on_tool_start`) +- `iteration` - Current agent iteration number +- `error` - Error string (empty if no error) + ### Execution Context + + ```go // internal/core/context.go @@ -282,9 +467,9 @@ Lookup order: ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │ BashExecutor │ │FunctionExec │ │ForeachExec │ └──────────────┘ └──────────────┘ └──────────────┘ - ┌──────────────┐ ┌──────────────┐ - │ HTTPExecutor │ │ LLMExecutor │ - └──────────────┘ └──────────────┘ + ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ + │ HTTPExecutor │ │ LLMExecutor │ │AgentExecutor │ + └──────────────┘ └──────────────┘ └──────────────┘ │ │ │ └────────────────────────────┼────────────────────────────┘ ▼ @@ -356,6 +541,7 @@ Built-in executors registered at startup: - `RemoteBashExecutor` - handles `remote-bash` steps - `HTTPExecutor` - handles `http` steps - `LLMExecutor` - handles `llm` steps +- `AgentExecutor` - handles `agent` steps (agentic LLM loop with tool calling) ## Run Control Plane @@ -604,6 +790,8 @@ _ = vm.Set("my_new_function", vf.myNewFunction) const FnMyNewFunction = "my_new_function" ``` +Notable utility functions include `exec_python(code)` and `exec_python_file(path)` for running Python code, and `run_module(module, target, params)` / `run_flow(flow, target, params)` for launching osmedeus workflows as subprocesses. + ### Output and Control Functions These functions provide output and execution control within workflows: @@ -1167,7 +1355,8 @@ test/e2e/ # E2E CLI tests ├── worker_test.go # Worker command tests ├── distributed_test.go # Distributed scan e2e tests ├── ssh_test.go # SSH runner e2e tests (module & step level) -└── api_test.go # API endpoint e2e tests (all routes) +├── api_test.go # API endpoint e2e tests (all routes) +└── agent_test.go # Agent step e2e tests ``` ### Running Tests diff --git a/README.md b/README.md index 20db4be..42c40d8 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,8 @@ Built for both beginners and experts, it delivers powerful, composable automatio - **Distributed Execution** - Scale with Redis-based master-worker pattern for parallel scanning - **Notifications** - Telegram bot and webhook integrations - **Cloud Storage** - S3-compatible storage for artifact management -- **LLM Integration** - AI-powered workflow steps with chat completions and embeddings +- **LLM Integration** - AI-powered workflow steps with chat completions, embeddings, and agentic tool-calling loops +- **Agent Step Type** - Agentic LLM execution with tool calling, sub-agents, and memory management See [Documentation Page](https://docs.osmedeus.org/) for more details. @@ -113,7 +114,7 @@ For more CLI usage and example commands, refer to the [CLI Reference](https://do │ │ CONFIG ──▶ PARSER ──▶ EXECUTOR ──▶ STEP DISPATCHER ──▶ RUNNER │ │ │ │ │ │ │ │ │ Step Executors: bash | function | parallel | foreach | remote-bash │ │ -│ │ http | llm │ │ +│ │ http | llm | agent │ │ │ │ │ │ │ │ │ Runners: HostRunner | DockerRunner | SSHRunner │ │ │ └─────────────────────────────────────────────────────────────────────┘ │ diff --git a/internal/broker/redis_event.go b/internal/broker/redis_event.go index 11063db..877cc5c 100644 --- a/internal/broker/redis_event.go +++ b/internal/broker/redis_event.go @@ -5,8 +5,8 @@ package broker import ( "context" - "github.com/j3ssie/osmedeus/v5/internal/json" "fmt" + "github.com/j3ssie/osmedeus/v5/internal/json" "sync" "time" diff --git a/internal/core/agent_tool_presets.go b/internal/core/agent_tool_presets.go new file mode 100644 index 0000000..30035f3 --- /dev/null +++ b/internal/core/agent_tool_presets.go @@ -0,0 +1,396 @@ +package core + +import ( + "fmt" + "strings" +) + +// PresetToolDef holds a preset tool's metadata for auto-generating LLMTool schemas +type PresetToolDef struct { + Description string + Parameters map[string]interface{} +} + +// PresetToolRegistry maps preset names to their tool definitions. +// These are used to auto-generate OpenAI-compatible function schemas +// when an agent step references a tool by preset name. +var PresetToolRegistry = map[string]PresetToolDef{ + "bash": { + Description: "Execute a shell command and return its output", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "command": map[string]interface{}{ + "type": "string", + "description": "The shell command to execute", + }, + }, + "required": []string{"command"}, + }, + }, + "read_file": { + Description: "Read the contents of a file", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "path": map[string]interface{}{ + "type": "string", + "description": "Path to the file to read", + }, + }, + "required": []string{"path"}, + }, + }, + "read_lines": { + Description: "Read a file and return its contents as an array of lines", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "path": map[string]interface{}{ + "type": "string", + "description": "Path to the file to read", + }, + }, + "required": []string{"path"}, + }, + }, + "file_exists": { + Description: "Check if a file exists at the given path", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "path": map[string]interface{}{ + "type": "string", + "description": "Path to check", + }, + }, + "required": []string{"path"}, + }, + }, + "file_length": { + Description: "Count the number of non-empty lines in a file", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "path": map[string]interface{}{ + "type": "string", + "description": "Path to the file", + }, + }, + "required": []string{"path"}, + }, + }, + "append_file": { + Description: "Append content from source file to destination file", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "dest": map[string]interface{}{ + "type": "string", + "description": "Destination file path", + }, + "content": map[string]interface{}{ + "type": "string", + "description": "Source file path to append from", + }, + }, + "required": []string{"dest", "content"}, + }, + }, + "save_content": { + Description: "Write string content to a file (overwrites if exists)", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "content": map[string]interface{}{ + "type": "string", + "description": "Content to write", + }, + "path": map[string]interface{}{ + "type": "string", + "description": "File path to write to", + }, + }, + "required": []string{"content", "path"}, + }, + }, + "glob": { + Description: "Find files matching a glob pattern", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "pattern": map[string]interface{}{ + "type": "string", + "description": "Glob pattern (e.g., '*.txt', '/path/**/*.json')", + }, + }, + "required": []string{"pattern"}, + }, + }, + "grep_string": { + Description: "Search a file for lines containing a string", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "source": map[string]interface{}{ + "type": "string", + "description": "File path to search in", + }, + "str": map[string]interface{}{ + "type": "string", + "description": "String to search for", + }, + }, + "required": []string{"source", "str"}, + }, + }, + "grep_regex": { + Description: "Search a file for lines matching a regex pattern", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "source": map[string]interface{}{ + "type": "string", + "description": "File path to search in", + }, + "pattern": map[string]interface{}{ + "type": "string", + "description": "Regex pattern to match", + }, + }, + "required": []string{"source", "pattern"}, + }, + }, + "http_get": { + Description: "Make an HTTP GET request and return the response", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "url": map[string]interface{}{ + "type": "string", + "description": "URL to send GET request to", + }, + }, + "required": []string{"url"}, + }, + }, + "http_request": { + Description: "Make an HTTP request with specified method, headers, and body", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "method": map[string]interface{}{ + "type": "string", + "description": "HTTP method (GET, POST, PUT, DELETE, etc.)", + }, + "url": map[string]interface{}{ + "type": "string", + "description": "URL to send request to", + }, + "body": map[string]interface{}{ + "type": "string", + "description": "Request body (for POST/PUT)", + }, + "headers": map[string]interface{}{ + "type": "string", + "description": "Headers as JSON string", + }, + }, + "required": []string{"url", "method"}, + }, + }, + "jq": { + Description: "Query JSON data using jq expression syntax", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "json_data": map[string]interface{}{ + "type": "string", + "description": "JSON string to query", + }, + "expression": map[string]interface{}{ + "type": "string", + "description": "jq expression (e.g., '.name', '.items[].id')", + }, + }, + "required": []string{"json_data", "expression"}, + }, + }, + "exec_python": { + Description: "Run inline Python code and return stdout (prefers python3)", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "code": map[string]interface{}{ + "type": "string", + "description": "Python code to execute", + }, + }, + "required": []string{"code"}, + }, + }, + "exec_python_file": { + Description: "Run a Python file and return stdout (prefers python3)", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "path": map[string]interface{}{ + "type": "string", + "description": "Path to the Python file to execute", + }, + }, + "required": []string{"path"}, + }, + }, + "run_module": { + Description: "Run an osmedeus module as a subprocess", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "module": map[string]interface{}{ + "type": "string", + "description": "Module name to run", + }, + "target": map[string]interface{}{ + "type": "string", + "description": "Target to scan", + }, + "params": map[string]interface{}{ + "type": "string", + "description": "Optional comma-separated key=value params (e.g., 'threads=10,deep=true')", + }, + }, + "required": []string{"module", "target"}, + }, + }, + "run_flow": { + Description: "Run an osmedeus flow as a subprocess", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "flow": map[string]interface{}{ + "type": "string", + "description": "Flow name to run", + }, + "target": map[string]interface{}{ + "type": "string", + "description": "Target to scan", + }, + "params": map[string]interface{}{ + "type": "string", + "description": "Optional comma-separated key=value params (e.g., 'threads=10,deep=true')", + }, + }, + "required": []string{"flow", "target"}, + }, + }, +} + +// GetPresetTool returns the LLMTool schema for a preset tool name. +// Returns the tool and true if found, zero value and false otherwise. +func GetPresetTool(presetName string) (LLMTool, bool) { + preset, ok := PresetToolRegistry[presetName] + if !ok { + return LLMTool{}, false + } + + return LLMTool{ + Type: "function", + Function: LLMToolFunction{ + Name: presetName, + Description: preset.Description, + Parameters: preset.Parameters, + }, + }, true +} + +// ResolveAgentTools converts a list of AgentToolDef into OpenAI-compatible LLMTool schemas. +// Preset tools are looked up from PresetToolRegistry; custom tools are converted directly. +func ResolveAgentTools(defs []AgentToolDef) ([]LLMTool, error) { + tools := make([]LLMTool, 0, len(defs)) + + for _, def := range defs { + if def.IsPreset() { + tool, ok := GetPresetTool(def.Preset) + if !ok { + return nil, fmt.Errorf("unknown preset tool: %s", def.Preset) + } + tools = append(tools, tool) + } else { + // Custom tool + if def.Name == "" { + return nil, fmt.Errorf("custom agent tool requires 'name' field") + } + tool := LLMTool{ + Type: "function", + Function: LLMToolFunction{ + Name: def.Name, + Description: def.Description, + Parameters: def.Parameters, + }, + } + tools = append(tools, tool) + } + } + + return tools, nil +} + +// SpawnAgentToolName is the tool name used for sub-agent spawning. +const SpawnAgentToolName = "spawn_agent" + +// BuildSpawnAgentTool dynamically generates the spawn_agent tool schema +// with an enum of available sub-agent names and their descriptions. +func BuildSpawnAgentTool(subAgents []SubAgentDef) LLMTool { + // Build enum of sub-agent names + names := make([]interface{}, 0, len(subAgents)) + var descParts []string + for _, sa := range subAgents { + names = append(names, sa.Name) + desc := sa.Name + if sa.Description != "" { + desc += ": " + sa.Description + } + descParts = append(descParts, desc) + } + + description := "Spawn a sub-agent to handle a specialized task. Available agents: " + strings.Join(descParts, "; ") + + return LLMTool{ + Type: "function", + Function: LLMToolFunction{ + Name: SpawnAgentToolName, + Description: description, + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "agent": map[string]interface{}{ + "type": "string", + "description": "Name of the sub-agent to spawn", + "enum": names, + }, + "query": map[string]interface{}{ + "type": "string", + "description": "The task or query to delegate to the sub-agent", + }, + }, + "required": []string{"agent", "query"}, + }, + }, + } +} + +// ResolveAgentToolsWithSubAgents resolves agent tools and appends the spawn_agent +// tool if sub-agents are defined. +func ResolveAgentToolsWithSubAgents(defs []AgentToolDef, subAgents []SubAgentDef) ([]LLMTool, error) { + tools, err := ResolveAgentTools(defs) + if err != nil { + return nil, err + } + + if len(subAgents) > 0 { + tools = append(tools, BuildSpawnAgentTool(subAgents)) + } + + return tools, nil +} diff --git a/internal/core/agent_tool_presets_test.go b/internal/core/agent_tool_presets_test.go new file mode 100644 index 0000000..e77a635 --- /dev/null +++ b/internal/core/agent_tool_presets_test.go @@ -0,0 +1,230 @@ +package core + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGetPresetTool(t *testing.T) { + t.Run("known preset", func(t *testing.T) { + tool, ok := GetPresetTool("bash") + assert.True(t, ok) + assert.Equal(t, "function", tool.Type) + assert.Equal(t, "bash", tool.Function.Name) + assert.NotEmpty(t, tool.Function.Description) + assert.NotNil(t, tool.Function.Parameters) + }) + + t.Run("all presets have valid schemas", func(t *testing.T) { + for name := range PresetToolRegistry { + tool, ok := GetPresetTool(name) + assert.True(t, ok, "preset %s should exist", name) + assert.Equal(t, "function", tool.Type) + assert.Equal(t, name, tool.Function.Name) + assert.NotEmpty(t, tool.Function.Description, "preset %s should have description", name) + assert.NotNil(t, tool.Function.Parameters, "preset %s should have parameters", name) + + // Verify parameters have 'type' and 'properties' + params := tool.Function.Parameters + assert.Equal(t, "object", params["type"], "preset %s parameters should be object type", name) + props, ok := params["properties"] + assert.True(t, ok, "preset %s should have properties", name) + assert.NotNil(t, props, "preset %s properties should not be nil", name) + } + }) + + t.Run("unknown preset", func(t *testing.T) { + _, ok := GetPresetTool("nonexistent") + assert.False(t, ok) + }) +} + +func TestGetPresetTool_SpecificPresets(t *testing.T) { + presets := []string{ + "bash", "read_file", "read_lines", "file_exists", "file_length", + "append_file", "save_content", "glob", "grep_string", "grep_regex", + "http_get", "http_request", "jq", + "exec_python", "exec_python_file", "run_module", "run_flow", + } + + for _, name := range presets { + t.Run(name, func(t *testing.T) { + tool, ok := GetPresetTool(name) + require.True(t, ok, "preset %s should be registered", name) + assert.Equal(t, name, tool.Function.Name) + + // Verify required fields are present + params := tool.Function.Parameters + required, hasRequired := params["required"] + if hasRequired { + reqSlice, ok := required.([]string) + assert.True(t, ok, "required should be []string for %s", name) + assert.NotEmpty(t, reqSlice, "required should not be empty for %s", name) + } + }) + } +} + +func TestResolveAgentTools(t *testing.T) { + t.Run("preset tools only", func(t *testing.T) { + defs := []AgentToolDef{ + {Preset: "bash"}, + {Preset: "read_file"}, + {Preset: "file_exists"}, + } + + tools, err := ResolveAgentTools(defs) + require.NoError(t, err) + assert.Len(t, tools, 3) + assert.Equal(t, "bash", tools[0].Function.Name) + assert.Equal(t, "read_file", tools[1].Function.Name) + assert.Equal(t, "file_exists", tools[2].Function.Name) + }) + + t.Run("custom tool", func(t *testing.T) { + defs := []AgentToolDef{ + { + Name: "my_tool", + Description: "My custom tool", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "input": map[string]interface{}{ + "type": "string", + }, + }, + }, + Handler: "exec_cmd(args.input)", + }, + } + + tools, err := ResolveAgentTools(defs) + require.NoError(t, err) + assert.Len(t, tools, 1) + assert.Equal(t, "function", tools[0].Type) + assert.Equal(t, "my_tool", tools[0].Function.Name) + assert.Equal(t, "My custom tool", tools[0].Function.Description) + }) + + t.Run("mixed preset and custom", func(t *testing.T) { + defs := []AgentToolDef{ + {Preset: "bash"}, + { + Name: "custom", + Description: "Custom tool", + Parameters: map[string]interface{}{"type": "object"}, + }, + {Preset: "read_file"}, + } + + tools, err := ResolveAgentTools(defs) + require.NoError(t, err) + assert.Len(t, tools, 3) + assert.Equal(t, "bash", tools[0].Function.Name) + assert.Equal(t, "custom", tools[1].Function.Name) + assert.Equal(t, "read_file", tools[2].Function.Name) + }) + + t.Run("unknown preset fails", func(t *testing.T) { + defs := []AgentToolDef{ + {Preset: "nonexistent_tool"}, + } + + _, err := ResolveAgentTools(defs) + assert.Error(t, err) + assert.Contains(t, err.Error(), "unknown preset tool") + }) + + t.Run("custom tool without name fails", func(t *testing.T) { + defs := []AgentToolDef{ + {Description: "no name tool"}, + } + + _, err := ResolveAgentTools(defs) + assert.Error(t, err) + assert.Contains(t, err.Error(), "requires 'name' field") + }) + + t.Run("empty list", func(t *testing.T) { + tools, err := ResolveAgentTools([]AgentToolDef{}) + require.NoError(t, err) + assert.Empty(t, tools) + }) +} + +func TestBuildSpawnAgentTool(t *testing.T) { + subAgents := []SubAgentDef{ + {Name: "recon_agent", Description: "Recon specialist"}, + {Name: "vuln_scanner", Description: "Vulnerability scanner"}, + } + + tool := BuildSpawnAgentTool(subAgents) + + assert.Equal(t, "function", tool.Type) + assert.Equal(t, SpawnAgentToolName, tool.Function.Name) + assert.Contains(t, tool.Function.Description, "recon_agent") + assert.Contains(t, tool.Function.Description, "vuln_scanner") + assert.Contains(t, tool.Function.Description, "Recon specialist") + + // Verify parameters schema + params := tool.Function.Parameters + assert.Equal(t, "object", params["type"]) + + props, ok := params["properties"].(map[string]interface{}) + require.True(t, ok) + + // Verify agent param has enum + agentParam, ok := props["agent"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "string", agentParam["type"]) + enumVals, ok := agentParam["enum"].([]interface{}) + require.True(t, ok) + assert.Len(t, enumVals, 2) + assert.Equal(t, "recon_agent", enumVals[0]) + assert.Equal(t, "vuln_scanner", enumVals[1]) + + // Verify query param + queryParam, ok := props["query"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "string", queryParam["type"]) + + // Verify required + required, ok := params["required"].([]string) + require.True(t, ok) + assert.Contains(t, required, "agent") + assert.Contains(t, required, "query") +} + +func TestResolveAgentToolsWithSubAgents(t *testing.T) { + t.Run("no sub-agents — same as ResolveAgentTools", func(t *testing.T) { + defs := []AgentToolDef{{Preset: "bash"}} + tools, err := ResolveAgentToolsWithSubAgents(defs, nil) + require.NoError(t, err) + assert.Len(t, tools, 1) + assert.Equal(t, "bash", tools[0].Function.Name) + }) + + t.Run("with sub-agents — spawn_agent appended", func(t *testing.T) { + defs := []AgentToolDef{{Preset: "bash"}, {Preset: "read_file"}} + subAgents := []SubAgentDef{ + {Name: "recon", Description: "Recon agent"}, + } + + tools, err := ResolveAgentToolsWithSubAgents(defs, subAgents) + require.NoError(t, err) + assert.Len(t, tools, 3) // bash + read_file + spawn_agent + + assert.Equal(t, "bash", tools[0].Function.Name) + assert.Equal(t, "read_file", tools[1].Function.Name) + assert.Equal(t, SpawnAgentToolName, tools[2].Function.Name) + }) + + t.Run("error propagated from ResolveAgentTools", func(t *testing.T) { + defs := []AgentToolDef{{Preset: "nonexistent"}} + _, err := ResolveAgentToolsWithSubAgents(defs, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "unknown preset tool") + }) +} diff --git a/internal/core/agent_types.go b/internal/core/agent_types.go new file mode 100644 index 0000000..2bdf70c --- /dev/null +++ b/internal/core/agent_types.go @@ -0,0 +1,92 @@ +package core + +// AgentToolDef defines a tool available to the agent. +// Supports two styles: +// - Preset: references a built-in osmedeus function by name (schema auto-generated) +// - Custom: explicit name, description, parameters schema, and handler expression +type AgentToolDef struct { + // Preset tool — name of an osmedeus function (schema auto-generated from PresetToolRegistry) + Preset string `yaml:"preset,omitempty" json:"preset,omitempty"` + + // Custom tool fields + Name string `yaml:"name,omitempty" json:"name,omitempty"` + Description string `yaml:"description,omitempty" json:"description,omitempty"` + Parameters map[string]interface{} `yaml:"parameters,omitempty" json:"parameters,omitempty"` + + // Handler is a JS expression executed when the tool is called. + // The parsed arguments are available as the `args` object. + Handler string `yaml:"handler,omitempty" json:"handler,omitempty"` +} + +// IsPreset returns true if this is a preset tool definition +func (t *AgentToolDef) IsPreset() bool { + return t.Preset != "" +} + +// AgentMemoryConfig configures conversation memory for agent steps +type AgentMemoryConfig struct { + // MaxMessages is the sliding window size for conversation messages. + // When exceeded, oldest non-system messages are dropped. + // 0 = unlimited (keep all messages in context). + MaxMessages int `yaml:"max_messages,omitempty" json:"max_messages,omitempty"` + + // SummarizeOnTruncate enables LLM-based summarization of dropped messages + // instead of silently discarding them when the sliding window is exceeded. + SummarizeOnTruncate bool `yaml:"summarize_on_truncate,omitempty" json:"summarize_on_truncate,omitempty"` + + // PersistPath is the file path to save the conversation JSON after completion. + PersistPath string `yaml:"persist_path,omitempty" json:"persist_path,omitempty"` + + // ResumePath is the file path to load a prior conversation from on start. + ResumePath string `yaml:"resume_path,omitempty" json:"resume_path,omitempty"` +} + +// DefaultMaxAgentDepth is the maximum nesting depth for sub-agent spawning. +const DefaultMaxAgentDepth = 3 + +// SubAgentDef defines an inline sub-agent that can be spawned by a parent agent. +type SubAgentDef struct { + Name string `yaml:"name" json:"name"` + Description string `yaml:"description,omitempty" json:"description,omitempty"` + SystemPrompt string `yaml:"system_prompt,omitempty" json:"system_prompt,omitempty"` + AgentTools []AgentToolDef `yaml:"agent_tools,omitempty" json:"agent_tools,omitempty"` + MaxIterations int `yaml:"max_iterations,omitempty" json:"max_iterations,omitempty"` + Models []string `yaml:"models,omitempty" json:"models,omitempty"` + LLMConfig *LLMStepConfig `yaml:"llm_config,omitempty" json:"llm_config,omitempty"` + OutputSchema string `yaml:"output_schema,omitempty" json:"output_schema,omitempty"` + Memory *AgentMemoryConfig `yaml:"memory,omitempty" json:"memory,omitempty"` + StopCondition string `yaml:"stop_condition,omitempty" json:"stop_condition,omitempty"` + SubAgents []SubAgentDef `yaml:"sub_agents,omitempty" json:"sub_agents,omitempty"` // Recursive nesting + OnToolStart string `yaml:"on_tool_start,omitempty" json:"on_tool_start,omitempty"` + OnToolEnd string `yaml:"on_tool_end,omitempty" json:"on_tool_end,omitempty"` +} + +// DeepCopy creates an independent deep copy of the SubAgentDef. +func (s *SubAgentDef) DeepCopy() SubAgentDef { + cp := *s + + if len(s.AgentTools) > 0 { + cp.AgentTools = make([]AgentToolDef, len(s.AgentTools)) + copy(cp.AgentTools, s.AgentTools) + } + if len(s.Models) > 0 { + cp.Models = make([]string, len(s.Models)) + copy(cp.Models, s.Models) + } + if s.LLMConfig != nil { + cfg := *s.LLMConfig + cp.LLMConfig = &cfg + } + if s.Memory != nil { + mem := *s.Memory + cp.Memory = &mem + } + if len(s.SubAgents) > 0 { + cp.SubAgents = make([]SubAgentDef, len(s.SubAgents)) + for i, sa := range s.SubAgents { + cp.SubAgents[i] = sa.DeepCopy() + } + } + + return cp +} diff --git a/internal/core/agent_types_test.go b/internal/core/agent_types_test.go new file mode 100644 index 0000000..dc0e841 --- /dev/null +++ b/internal/core/agent_types_test.go @@ -0,0 +1,346 @@ +package core + +import ( + "testing" + + "github.com/goccy/go-yaml" + "github.com/stretchr/testify/assert" +) + +func TestAgentToolDef_IsPreset(t *testing.T) { + t.Run("preset tool", func(t *testing.T) { + tool := AgentToolDef{Preset: "bash"} + assert.True(t, tool.IsPreset()) + }) + + t.Run("custom tool", func(t *testing.T) { + tool := AgentToolDef{Name: "my_tool", Description: "A custom tool"} + assert.False(t, tool.IsPreset()) + }) + + t.Run("empty tool", func(t *testing.T) { + tool := AgentToolDef{} + assert.False(t, tool.IsPreset()) + }) +} + +func TestAgentToolDef_YAMLUnmarshal(t *testing.T) { + t.Run("preset tool", func(t *testing.T) { + yamlData := `preset: bash` + var tool AgentToolDef + err := yaml.Unmarshal([]byte(yamlData), &tool) + assert.NoError(t, err) + assert.Equal(t, "bash", tool.Preset) + assert.True(t, tool.IsPreset()) + }) + + t.Run("custom tool with handler", func(t *testing.T) { + yamlData := ` +name: nuclei_scan +description: "Run nuclei scanner" +parameters: + type: object + properties: + url: + type: string + description: "Target URL" + required: + - url +handler: 'exec_cmd("nuclei -u " + args.url)' +` + var tool AgentToolDef + err := yaml.Unmarshal([]byte(yamlData), &tool) + assert.NoError(t, err) + assert.Equal(t, "nuclei_scan", tool.Name) + assert.Equal(t, "Run nuclei scanner", tool.Description) + assert.NotNil(t, tool.Parameters) + assert.Equal(t, `exec_cmd("nuclei -u " + args.url)`, tool.Handler) + assert.False(t, tool.IsPreset()) + }) +} + +func TestAgentMemoryConfig_YAMLUnmarshal(t *testing.T) { + t.Run("full memory config", func(t *testing.T) { + yamlData := ` +max_messages: 50 +persist_path: "/tmp/agent/conv.json" +resume_path: "/tmp/agent/prev.json" +` + var mem AgentMemoryConfig + err := yaml.Unmarshal([]byte(yamlData), &mem) + assert.NoError(t, err) + assert.Equal(t, 50, mem.MaxMessages) + assert.Equal(t, "/tmp/agent/conv.json", mem.PersistPath) + assert.Equal(t, "/tmp/agent/prev.json", mem.ResumePath) + }) + + t.Run("empty memory config", func(t *testing.T) { + yamlData := `{}` + var mem AgentMemoryConfig + err := yaml.Unmarshal([]byte(yamlData), &mem) + assert.NoError(t, err) + assert.Equal(t, 0, mem.MaxMessages) + assert.Empty(t, mem.PersistPath) + assert.Empty(t, mem.ResumePath) + }) +} + +func TestStep_IsAgentStep(t *testing.T) { + t.Run("agent step", func(t *testing.T) { + step := Step{Type: StepTypeAgent} + assert.True(t, step.IsAgentStep()) + }) + + t.Run("non-agent step", func(t *testing.T) { + step := Step{Type: StepTypeBash} + assert.False(t, step.IsAgentStep()) + }) +} + +func TestStep_AgentFields_YAMLUnmarshal(t *testing.T) { + yamlData := ` +name: test-agent +type: agent +query: "Analyze {{Target}}" +system_prompt: "You are a helpful assistant." +max_iterations: 10 +stop_condition: 'contains(agent_content, "DONE")' +parallel_tool_calls: false +agent_tools: + - preset: bash + - preset: read_file + - name: custom_tool + description: "A custom tool" + parameters: + type: object + properties: + input: + type: string + handler: 'trim(args.input)' +memory: + max_messages: 50 + persist_path: "/tmp/conv.json" +llm_config: + model: "gpt-4o" + temperature: 0.2 +exports: + result: "{{agent_content}}" +` + var step Step + err := yaml.Unmarshal([]byte(yamlData), &step) + assert.NoError(t, err) + + assert.Equal(t, "test-agent", step.Name) + assert.Equal(t, StepTypeAgent, step.Type) + assert.Equal(t, "Analyze {{Target}}", step.Query) + assert.Equal(t, "You are a helpful assistant.", step.SystemPrompt) + assert.Equal(t, 10, step.MaxIterations) + assert.Equal(t, `contains(agent_content, "DONE")`, step.StopCondition) + assert.NotNil(t, step.ParallelToolCalls) + assert.False(t, *step.ParallelToolCalls) + + // Agent tools + assert.Len(t, step.AgentTools, 3) + assert.Equal(t, "bash", step.AgentTools[0].Preset) + assert.Equal(t, "read_file", step.AgentTools[1].Preset) + assert.Equal(t, "custom_tool", step.AgentTools[2].Name) + assert.Equal(t, `trim(args.input)`, step.AgentTools[2].Handler) + + // Memory + assert.NotNil(t, step.Memory) + assert.Equal(t, 50, step.Memory.MaxMessages) + assert.Equal(t, "/tmp/conv.json", step.Memory.PersistPath) + + // LLM Config + assert.NotNil(t, step.LLMConfig) + assert.Equal(t, "gpt-4o", step.LLMConfig.Model) + + // Exports + assert.Equal(t, "{{agent_content}}", step.Exports["result"]) +} + +func TestSubAgentDef_DeepCopy(t *testing.T) { + t.Run("deep copy independence", func(t *testing.T) { + original := SubAgentDef{ + Name: "recon", + Description: "Recon agent", + SystemPrompt: "You are a recon specialist", + AgentTools: []AgentToolDef{ + {Preset: "bash"}, + {Preset: "http_get"}, + }, + Models: []string{"gpt-4o", "claude-3"}, + MaxIterations: 5, + LLMConfig: &LLMStepConfig{Model: "gpt-4o"}, + Memory: &AgentMemoryConfig{MaxMessages: 20, PersistPath: "/tmp/conv.json"}, + StopCondition: `contains(agent_content, "DONE")`, + SubAgents: []SubAgentDef{ + { + Name: "nested", + AgentTools: []AgentToolDef{{Preset: "read_file"}}, + }, + }, + OnToolStart: `log_info("start")`, + OnToolEnd: `log_info("end")`, + } + + cp := original.DeepCopy() + + // Verify values match + assert.Equal(t, "recon", cp.Name) + assert.Equal(t, "Recon agent", cp.Description) + assert.Equal(t, "You are a recon specialist", cp.SystemPrompt) + assert.Len(t, cp.AgentTools, 2) + assert.Len(t, cp.Models, 2) + assert.Equal(t, "gpt-4o", cp.LLMConfig.Model) + assert.Equal(t, 20, cp.Memory.MaxMessages) + assert.Len(t, cp.SubAgents, 1) + assert.Equal(t, "nested", cp.SubAgents[0].Name) + + // Modify original, verify copy is independent + original.AgentTools[0].Preset = "modified" + assert.Equal(t, "bash", cp.AgentTools[0].Preset) + + original.Models[0] = "modified" + assert.Equal(t, "gpt-4o", cp.Models[0]) + + original.LLMConfig.Model = "modified" + assert.Equal(t, "gpt-4o", cp.LLMConfig.Model) + + original.Memory.MaxMessages = 999 + assert.Equal(t, 20, cp.Memory.MaxMessages) + + original.SubAgents[0].Name = "modified" + assert.Equal(t, "nested", cp.SubAgents[0].Name) + }) + + t.Run("nil pointers", func(t *testing.T) { + original := SubAgentDef{ + Name: "minimal", + } + cp := original.DeepCopy() + assert.Equal(t, "minimal", cp.Name) + assert.Nil(t, cp.LLMConfig) + assert.Nil(t, cp.Memory) + assert.Nil(t, cp.SubAgents) + }) +} + +func TestSubAgentDef_YAMLUnmarshal(t *testing.T) { + yamlData := ` +name: recon_agent +description: "Specialized agent for recon" +system_prompt: "You are a recon specialist" +max_iterations: 5 +agent_tools: + - preset: bash + - preset: http_get +sub_agents: + - name: port_scanner + description: "Scans ports" + agent_tools: + - preset: bash +` + var sa SubAgentDef + err := yaml.Unmarshal([]byte(yamlData), &sa) + assert.NoError(t, err) + assert.Equal(t, "recon_agent", sa.Name) + assert.Equal(t, "Specialized agent for recon", sa.Description) + assert.Equal(t, "You are a recon specialist", sa.SystemPrompt) + assert.Equal(t, 5, sa.MaxIterations) + assert.Len(t, sa.AgentTools, 2) + assert.Len(t, sa.SubAgents, 1) + assert.Equal(t, "port_scanner", sa.SubAgents[0].Name) +} + +func TestStep_SubAgents_YAMLUnmarshal(t *testing.T) { + yamlData := ` +name: orchestrator +type: agent +query: "Coordinate analysis" +system_prompt: "You are an orchestrator" +max_iterations: 10 +max_agent_depth: 2 +agent_tools: + - preset: bash +sub_agents: + - name: recon_agent + description: "Recon specialist" + system_prompt: "You do recon" + max_iterations: 5 + agent_tools: + - preset: bash + - preset: http_get + - name: vuln_scanner + description: "Vulnerability scanner" + agent_tools: + - preset: bash +` + var step Step + err := yaml.Unmarshal([]byte(yamlData), &step) + assert.NoError(t, err) + assert.Equal(t, "orchestrator", step.Name) + assert.Equal(t, 2, step.MaxAgentDepth) + assert.Len(t, step.SubAgents, 2) + assert.Equal(t, "recon_agent", step.SubAgents[0].Name) + assert.Equal(t, "Recon specialist", step.SubAgents[0].Description) + assert.Len(t, step.SubAgents[0].AgentTools, 2) + assert.Equal(t, "vuln_scanner", step.SubAgents[1].Name) +} + +func TestStep_Clone_AgentFields(t *testing.T) { + ptc := false + step := Step{ + Name: "agent-step", + Type: StepTypeAgent, + Query: "test query", + SystemPrompt: "test prompt", + MaxIterations: 10, + StopCondition: "test condition", + ParallelToolCalls: &ptc, + AgentTools: []AgentToolDef{ + {Preset: "bash"}, + {Name: "custom", Handler: "handler()"}, + }, + Memory: &AgentMemoryConfig{ + MaxMessages: 50, + PersistPath: "/tmp/conv.json", + }, + SubAgents: []SubAgentDef{ + {Name: "sub1", AgentTools: []AgentToolDef{{Preset: "bash"}}}, + }, + MaxAgentDepth: 2, + } + + cloned := step.Clone() + + // Verify values are copied + assert.Equal(t, "agent-step", cloned.Name) + assert.Equal(t, "test query", cloned.Query) + assert.Equal(t, "test prompt", cloned.SystemPrompt) + assert.Equal(t, 10, cloned.MaxIterations) + assert.Equal(t, "test condition", cloned.StopCondition) + assert.NotNil(t, cloned.ParallelToolCalls) + assert.False(t, *cloned.ParallelToolCalls) + + // Agent tools are deep copied + assert.Len(t, cloned.AgentTools, 2) + assert.Equal(t, "bash", cloned.AgentTools[0].Preset) + + // Modify original, verify clone is independent + step.AgentTools[0].Preset = "modified" + assert.Equal(t, "bash", cloned.AgentTools[0].Preset) + + // Memory is deep copied + assert.NotNil(t, cloned.Memory) + assert.Equal(t, 50, cloned.Memory.MaxMessages) + step.Memory.MaxMessages = 100 + assert.Equal(t, 50, cloned.Memory.MaxMessages) + + // SubAgents are deep copied + assert.Len(t, cloned.SubAgents, 1) + assert.Equal(t, "sub1", cloned.SubAgents[0].Name) + step.SubAgents[0].Name = "modified" + assert.Equal(t, "sub1", cloned.SubAgents[0].Name) + assert.Equal(t, 2, cloned.MaxAgentDepth) +} diff --git a/internal/core/llm_types.go b/internal/core/llm_types.go index 196c534..618bdb7 100644 --- a/internal/core/llm_types.go +++ b/internal/core/llm_types.go @@ -1,5 +1,28 @@ package core +import ( + "fmt" + + "github.com/j3ssie/osmedeus/v5/internal/json" +) + +// ParseOutputSchema converts a JSON string into an LLMResponseFormat suitable +// for the OpenAI response_format parameter. The schemaJSON should be a valid +// JSON schema object, e.g. '{"type":"object","properties":{...}}'. +func ParseOutputSchema(schemaJSON string) (*LLMResponseFormat, error) { + var schema map[string]interface{} + if err := json.Unmarshal([]byte(schemaJSON), &schema); err != nil { + return nil, fmt.Errorf("invalid output_schema JSON: %w", err) + } + return &LLMResponseFormat{ + Type: "json_schema", + JSONSchema: map[string]interface{}{ + "name": "output_schema", + "schema": schema, + }, + }, nil +} + // LLMMessageRole represents the role of a message sender type LLMMessageRole string diff --git a/internal/core/preferences_test.go b/internal/core/preferences_test.go index 31dc9b5..22c3c9d 100644 --- a/internal/core/preferences_test.go +++ b/internal/core/preferences_test.go @@ -63,11 +63,11 @@ func TestPreferences_GetEmptyTarget(t *testing.T) { func TestPreferences_UnmarshalYAML_EmptyTarget(t *testing.T) { tests := []struct { - name string - input string - wantNil bool - wantVal bool - wantErr bool + name string + input string + wantNil bool + wantVal bool + wantErr bool }{ { name: "empty_target true", diff --git a/internal/core/step.go b/internal/core/step.go index 94becf7..531f89d 100644 --- a/internal/core/step.go +++ b/internal/core/step.go @@ -232,11 +232,45 @@ type Step struct { EmbeddingInput []string `yaml:"embedding_input,omitempty"` ExtraLLMParams map[string]interface{} `yaml:"extra_llm_parameters,omitempty"` + // Agent step fields + Query string `yaml:"query,omitempty"` + Queries []string `yaml:"queries,omitempty"` // Multiple queries executed sequentially (multi-goal) + SystemPrompt string `yaml:"system_prompt,omitempty"` + AgentTools []AgentToolDef `yaml:"agent_tools,omitempty"` + ParallelToolCalls *bool `yaml:"parallel_tool_calls,omitempty"` + MaxIterations int `yaml:"max_iterations,omitempty"` + StopCondition string `yaml:"stop_condition,omitempty"` + Memory *AgentMemoryConfig `yaml:"memory,omitempty"` + + // Agent planning stage + PlanPrompt string `yaml:"plan_prompt,omitempty"` + PlanMaxTokens *int `yaml:"plan_max_tokens,omitempty"` + + // Agent model preferences (tried in order before falling back to default) + Models []string `yaml:"models,omitempty"` + + // Agent structured output (enforced on final iteration) + // JSON string, e.g. '{"type":"object","properties":{"key":{"type":"string"}}}' + OutputSchema string `yaml:"output_schema,omitempty"` + + // Agent tool tracing hooks (JS expressions) + OnToolStart string `yaml:"on_tool_start,omitempty"` // Evaluated before each tool call + OnToolEnd string `yaml:"on_tool_end,omitempty"` // Evaluated after each tool call + + // Sub-agents that can be spawned by this agent via tool calls + SubAgents []SubAgentDef `yaml:"sub_agents,omitempty"` + // Maximum nesting depth for sub-agent spawning (default: 3) + MaxAgentDepth int `yaml:"max_agent_depth,omitempty"` + + // Streaming (applies to both llm and agent steps) + Stream *bool `yaml:"stream,omitempty"` // Enable streaming output (overrides llm_config.stream and global config) + // Common fields - Exports map[string]string `yaml:"exports"` - OnSuccess []Action `yaml:"on_success"` - OnError []Action `yaml:"on_error"` - Decision *DecisionConfig `yaml:"decision,omitempty"` + SuppressDetails bool `yaml:"suppress_details"` // Hide command/function details from output + Exports map[string]string `yaml:"exports"` + OnSuccess []Action `yaml:"on_success"` + OnError []Action `yaml:"on_error"` + Decision *DecisionConfig `yaml:"decision,omitempty"` } // DecisionCase represents a single case in switch-style decision @@ -310,6 +344,11 @@ func (s *Step) IsLLMStep() bool { return s.Type == StepTypeLLM } +// IsAgentStep returns true if this is an agent step +func (s *Step) IsAgentStep() bool { + return s.Type == StepTypeAgent +} + // GetStepRunner returns the step runner type, defaulting to host/local func (s *Step) GetStepRunner() RunnerType { if s.StepRunner == "" { @@ -441,5 +480,31 @@ func (s *Step) Clone() *Step { } } + // Deep copy Agent fields + if len(s.AgentTools) > 0 { + cloned.AgentTools = make([]AgentToolDef, len(s.AgentTools)) + copy(cloned.AgentTools, s.AgentTools) + } + if s.Memory != nil { + mem := *s.Memory + cloned.Memory = &mem + } + if len(s.Queries) > 0 { + cloned.Queries = make([]string, len(s.Queries)) + copy(cloned.Queries, s.Queries) + } + if len(s.Models) > 0 { + cloned.Models = make([]string, len(s.Models)) + copy(cloned.Models, s.Models) + } + + // Deep copy SubAgents + if len(s.SubAgents) > 0 { + cloned.SubAgents = make([]SubAgentDef, len(s.SubAgents)) + for i, sa := range s.SubAgents { + cloned.SubAgents[i] = sa.DeepCopy() + } + } + return &cloned } diff --git a/internal/core/step_test.go b/internal/core/step_test.go new file mode 100644 index 0000000..ad89480 --- /dev/null +++ b/internal/core/step_test.go @@ -0,0 +1,60 @@ +package core + +import ( + "testing" + + "github.com/goccy/go-yaml" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestStepSuppressDetails(t *testing.T) { + tests := []struct { + name string + yaml string + expected bool + }{ + { + name: "suppress_details true", + yaml: ` +name: test-step +type: function +suppress_details: true +functions: + - exec_cmd("echo hello") +`, + expected: true, + }, + { + name: "suppress_details false", + yaml: ` +name: test-step +type: function +suppress_details: false +functions: + - exec_cmd("echo hello") +`, + expected: false, + }, + { + name: "suppress_details omitted defaults to false", + yaml: ` +name: test-step +type: function +functions: + - exec_cmd("echo hello") +`, + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var step Step + err := yaml.Unmarshal([]byte(tt.yaml), &step) + require.NoError(t, err) + assert.Equal(t, tt.expected, step.SuppressDetails) + assert.Equal(t, "test-step", step.Name) + }) + } +} diff --git a/internal/core/types.go b/internal/core/types.go index 86211d0..4912513 100644 --- a/internal/core/types.go +++ b/internal/core/types.go @@ -24,6 +24,7 @@ const ( StepTypeRemoteBash StepType = "remote-bash" StepTypeHTTP StepType = "http" StepTypeLLM StepType = "llm" + StepTypeAgent StepType = "agent" ) // TriggerType represents trigger types diff --git a/internal/database/agent_session.go b/internal/database/agent_session.go new file mode 100644 index 0000000..884c1ae --- /dev/null +++ b/internal/database/agent_session.go @@ -0,0 +1,60 @@ +package database + +import ( + "context" + "time" + + "github.com/uptrace/bun" +) + +// AgentSession persists agent execution sessions for querying, debugging, and resuming +type AgentSession struct { + bun.BaseModel `bun:"table:agent_sessions,alias:as"` + + ID int64 `bun:"id,pk,autoincrement"` + RunID int64 `bun:"run_id"` + StepName string `bun:"step_name,notnull"` + Query string `bun:"query"` + PlanContent string `bun:"plan_content"` + FinalContent string `bun:"final_content"` + Iterations int `bun:"iterations"` + TotalTokens int `bun:"total_tokens"` + PromptTokens int `bun:"prompt_tokens"` + CompletionTokens int `bun:"completion_tokens"` + ToolCallsJSON string `bun:"tool_calls_json"` + ConversationJSON string `bun:"conversation_json"` + Status string `bun:"status,notnull"` + DurationMs int64 `bun:"duration_ms"` + CreatedAt time.Time `bun:"created_at,notnull,default:current_timestamp"` +} + +// CreateAgentSession inserts a new agent session record +func CreateAgentSession(ctx context.Context, session *AgentSession) error { + if db == nil { + return nil // No database configured, skip persistence + } + _, err := db.NewInsert().Model(session).Exec(ctx) + return err +} + +// GetAgentSessionsByRun returns all agent sessions for a given run +func GetAgentSessionsByRun(ctx context.Context, runID int64) ([]AgentSession, error) { + var sessions []AgentSession + err := db.NewSelect(). + Model(&sessions). + Where("run_id = ?", runID). + Order("created_at ASC"). + Scan(ctx) + return sessions, err +} + +// GetAgentSessionsByStep returns all agent sessions for a given step name +func GetAgentSessionsByStep(ctx context.Context, stepName string) ([]AgentSession, error) { + var sessions []AgentSession + err := db.NewSelect(). + Model(&sessions). + Where("step_name = ?", stepName). + Order("created_at DESC"). + Scan(ctx) + return sessions, err +} diff --git a/internal/database/database.go b/internal/database/database.go index 4fb4c56..9ead51f 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -159,6 +159,7 @@ func Migrate(ctx context.Context) error { (*Vulnerability)(nil), (*AssetDiffSnapshot)(nil), (*VulnDiffSnapshot)(nil), + (*AgentSession)(nil), } for _, model := range models { diff --git a/internal/database/jsonl.go b/internal/database/jsonl.go index 996ea6e..a991a10 100644 --- a/internal/database/jsonl.go +++ b/internal/database/jsonl.go @@ -3,8 +3,8 @@ package database import ( "bufio" "context" - "github.com/j3ssie/osmedeus/v5/internal/json" "fmt" + "github.com/j3ssie/osmedeus/v5/internal/json" "io" "net" "os" diff --git a/internal/distributed/master.go b/internal/distributed/master.go index f0c0931..80ab419 100644 --- a/internal/distributed/master.go +++ b/internal/distributed/master.go @@ -2,8 +2,8 @@ package distributed import ( "context" - "github.com/j3ssie/osmedeus/v5/internal/json" "fmt" + "github.com/j3ssie/osmedeus/v5/internal/json" "os" "sync" "time" diff --git a/internal/executor/agent_executor.go b/internal/executor/agent_executor.go new file mode 100644 index 0000000..8dc65d8 --- /dev/null +++ b/internal/executor/agent_executor.go @@ -0,0 +1,1220 @@ +package executor + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/j3ssie/osmedeus/v5/internal/config" + "github.com/j3ssie/osmedeus/v5/internal/core" + "github.com/j3ssie/osmedeus/v5/internal/functions" + "github.com/j3ssie/osmedeus/v5/internal/json" + "github.com/j3ssie/osmedeus/v5/internal/logger" + "github.com/j3ssie/osmedeus/v5/internal/template" + "go.uber.org/zap" +) + +// AgentExecutor implements the agentic loop for type: agent steps. +// It queries the LLM, handles tool calls by executing them via the ToolExecutorRegistry, +// feeds results back, and repeats until the LLM responds without tool calls or +// the max_iterations limit is reached. +type AgentExecutor struct { + templateEngine template.TemplateEngine + functionRegistry *functions.Registry + config *config.Config + silent bool + currentDepth int // 0 = top-level + maxDepth int // default 3 +} + +// NewAgentExecutor creates a new agent executor +func NewAgentExecutor(engine template.TemplateEngine, funcRegistry *functions.Registry) *AgentExecutor { + return &AgentExecutor{ + templateEngine: engine, + functionRegistry: funcRegistry, + } +} + +// Name returns the executor name +func (e *AgentExecutor) Name() string { + return "agent" +} + +// StepTypes returns the step types this executor handles +func (e *AgentExecutor) StepTypes() []core.StepType { + return []core.StepType{core.StepTypeAgent} +} + +// SetConfig sets the application config +func (e *AgentExecutor) SetConfig(cfg *config.Config) { + e.config = cfg +} + +// SetSilent enables or disables silent mode +func (e *AgentExecutor) SetSilent(s bool) { + e.silent = s +} + +// SetDepthContext sets the current nesting depth and maximum allowed depth. +func (e *AgentExecutor) SetDepthContext(depth, maxDepth int) { + e.currentDepth = depth + e.maxDepth = maxDepth +} + +// agentState tracks the agent's runtime state during execution +type agentState struct { + messages []ChatMessage + totalTokens int + promptTokens int + completionTokens int + iteration int + toolResults []map[string]interface{} + finalContent string + planContent string // populated by planning stage + goalResults []map[string]interface{} // results from each goal in multi-goal mode + toolRegistry *ToolExecutorRegistry // pluggable tool dispatch + tokenMu sync.Mutex // protects token fields for concurrent sub-agent merging +} + +// MergeTokens safely adds child agent token counts into this state. +// Thread-safe for concurrent sub-agent spawning. +func (s *agentState) MergeTokens(total, prompt, completion int) { + s.tokenMu.Lock() + defer s.tokenMu.Unlock() + s.totalTokens += total + s.promptTokens += prompt + s.completionTokens += completion +} + +// Execute runs the agent loop +func (e *AgentExecutor) Execute(ctx context.Context, step *core.Step, execCtx *core.ExecutionContext) (*core.StepResult, error) { + log := logger.Get() + result := &core.StepResult{ + StepName: step.Name, + Status: core.StepStatusRunning, + StartTime: time.Now(), + Exports: make(map[string]interface{}), + } + + // Validate config + if e.config == nil { + return e.fail(result, fmt.Errorf("agent executor config not set")) + } + + // Validate required fields: need query OR queries (not both) + if step.Query == "" && len(step.Queries) == 0 { + return e.fail(result, fmt.Errorf("agent step '%s' requires 'query' or 'queries' field", step.Name)) + } + if step.Query != "" && len(step.Queries) > 0 { + return e.fail(result, fmt.Errorf("agent step '%s' cannot have both 'query' and 'queries'", step.Name)) + } + if step.MaxIterations <= 0 { + return e.fail(result, fmt.Errorf("agent step '%s' requires 'max_iterations' > 0", step.Name)) + } + if len(step.AgentTools) == 0 { + return e.fail(result, fmt.Errorf("agent step '%s' requires 'agent_tools' field", step.Name)) + } + + // Parse output schema if specified + var outputSchema *core.LLMResponseFormat + if step.OutputSchema != "" { + var err error + outputSchema, err = core.ParseOutputSchema(step.OutputSchema) + if err != nil { + return e.fail(result, fmt.Errorf("agent step '%s': %w", step.Name, err)) + } + } + + // Resolve agent tools to OpenAI-compatible schemas (with spawn_agent if sub-agents present) + tools, err := core.ResolveAgentToolsWithSubAgents(step.AgentTools, step.SubAgents) + if err != nil { + return e.fail(result, fmt.Errorf("agent step '%s': %w", step.Name, err)) + } + + // Get merged LLM config + llmConfig := e.getMergedConfig(step) + + // Compute effective max depth + effectiveMaxDepth := step.MaxAgentDepth + if effectiveMaxDepth <= 0 { + effectiveMaxDepth = core.DefaultMaxAgentDepth + } + // If spawned as child, inherit parent's max depth + if e.maxDepth > 0 { + effectiveMaxDepth = e.maxDepth + } + + // Initialize state before building tool registry (needed as parentState arg) + state := &agentState{} + + // Build ToolExecutorRegistry with sub-agent support if needed + var toolRegistry *ToolExecutorRegistry + if len(step.SubAgents) > 0 { + toolRegistry = BuildToolRegistryWithSubAgents( + step.AgentTools, e.functionRegistry, + e.templateEngine, e.config, e.silent, + e.currentDepth, effectiveMaxDepth, + state, step.SubAgents, + ) + } else { + toolRegistry = BuildToolRegistry(step.AgentTools, e.functionRegistry) + } + state.toolRegistry = toolRegistry + + // Load resumed conversation if specified + if step.Memory != nil && step.Memory.ResumePath != "" { + if err := e.loadConversation(state, step.Memory.ResumePath); err != nil { + log.Warn("Failed to load conversation for resume, starting fresh", + zap.String("path", step.Memory.ResumePath), + zap.Error(err), + ) + } + } + + // Determine queries to execute (Group 3: multi-goal) + queries := []string{step.Query} + if len(step.Queries) > 0 { + queries = step.Queries + } + + // Run planning stage if configured (Group 2) + if step.PlanPrompt != "" { + planContent, err := e.executePlanningStage(ctx, state, step, llmConfig) + if err != nil { + log.Warn("Planning stage failed, continuing without plan", + zap.String("step", step.Name), + zap.Error(err), + ) + } else { + state.planContent = planContent + } + } + + // Execute each query (single or multi-goal) + for goalIdx, query := range queries { + if query == "" { + continue + } + + // Build initial messages for this goal + e.initMessages(state, query, step, llmConfig, goalIdx) + + log.Debug("Starting agent loop", + zap.String("step", step.Name), + zap.Int("goal", goalIdx+1), + zap.Int("total_goals", len(queries)), + zap.Int("max_iterations", step.MaxIterations), + zap.Int("tools", len(tools)), + ) + + // Main agent loop + for state.iteration = 1; state.iteration <= step.MaxIterations; state.iteration++ { + log.Debug("Agent iteration", + zap.String("step", step.Name), + zap.Int("iteration", state.iteration), + zap.Int("messages", len(state.messages)), + ) + + // Build LLM request, potentially with structured output on final iteration + var responseFormat *core.LLMResponseFormat + if outputSchema != nil && state.iteration == step.MaxIterations { + responseFormat = outputSchema + } + + // Call LLM with optional model fallback (Group 5) + response, err := e.callLLMWithFallback(ctx, state, tools, llmConfig, step.Models, responseFormat) + if err != nil { + return e.fail(result, fmt.Errorf("agent step '%s' iteration %d: %w", step.Name, state.iteration, err)) + } + + if response == nil || len(response.Choices) == 0 { + return e.fail(result, fmt.Errorf("agent step '%s': empty response from LLM", step.Name)) + } + + // Track tokens + state.totalTokens += response.Usage.TotalTokens + state.promptTokens += response.Usage.PromptTokens + state.completionTokens += response.Usage.CompletionTokens + + choice := response.Choices[0] + + // Append assistant message to conversation + state.messages = append(state.messages, choice.Message) + + // Extract content + if content, ok := choice.Message.Content.(string); ok { + state.finalContent = content + } + + // Check if we're done (no tool calls) + if len(choice.Message.ToolCalls) == 0 { + // If we have OutputSchema and this isn't the max iteration, + // try to enforce structured output on the next call + if outputSchema != nil && state.iteration < step.MaxIterations { + log.Debug("Agent completed (no tool calls), structured output available", + zap.String("step", step.Name), + zap.Int("iterations", state.iteration), + ) + } else { + log.Debug("Agent completed (no tool calls)", + zap.String("step", step.Name), + zap.Int("iterations", state.iteration), + ) + } + break + } + + // Execute tool calls with tracing hooks (Group 7) + toolMessages, err := e.executeToolCalls(ctx, choice.Message.ToolCalls, state, step, execCtx) + if err != nil { + return e.fail(result, fmt.Errorf("agent step '%s' iteration %d tool execution: %w", step.Name, state.iteration, err)) + } + + // Append tool results to conversation + state.messages = append(state.messages, toolMessages...) + + // Evaluate stop condition if defined + if step.StopCondition != "" { + vars := execCtx.GetVariables() + vars["agent_content"] = state.finalContent + vars["iteration"] = state.iteration + shouldStop, err := e.functionRegistry.EvaluateCondition(step.StopCondition, vars) + if err != nil { + log.Warn("Stop condition evaluation failed", + zap.String("step", step.Name), + zap.Error(err), + ) + } else if shouldStop { + log.Debug("Agent stopped by stop_condition", + zap.String("step", step.Name), + zap.Int("iteration", state.iteration), + ) + break + } + } + + // Apply sliding window if configured (with optional compression — Group 4) + if step.Memory != nil && step.Memory.MaxMessages > 0 { + if step.Memory.SummarizeOnTruncate { + e.applyMessageWindowWithSummary(ctx, state, step.Memory.MaxMessages, llmConfig) + } else { + e.applyMessageWindow(state, step.Memory.MaxMessages) + } + } + } + + // Record goal result (Group 3) + goalResult := map[string]interface{}{ + "query": query, + "content": state.finalContent, + } + state.goalResults = append(state.goalResults, goalResult) + } + + // If OutputSchema is set and we haven't gotten structured output yet, + // make a final structured output request (Group 6) + if outputSchema != nil { + structuredContent := e.requestStructuredOutput(ctx, state, outputSchema, llmConfig) + if structuredContent != "" { + state.finalContent = structuredContent + } + } + + // Print final output (skip if streaming — tokens were already printed in real-time) + if !e.silent && !llmConfig.Stream && state.finalContent != "" { + printLLMOutput(state.finalContent) + } + + // Persist conversation if configured + if step.Memory != nil && step.Memory.PersistPath != "" { + if err := e.persistConversation(state, step.Memory.PersistPath); err != nil { + log.Warn("Failed to persist agent conversation", + zap.String("path", step.Memory.PersistPath), + zap.Error(err), + ) + } + } + + // Set exports + historyJSON, err := json.Marshal(state.messages) + if err != nil { + log.Warn("Failed to marshal agent history", zap.Error(err)) + historyJSON = []byte("[]") + } + toolResultsJSON, err := json.Marshal(state.toolResults) + if err != nil { + log.Warn("Failed to marshal agent tool results", zap.Error(err)) + toolResultsJSON = []byte("[]") + } + + result.Exports["agent_content"] = state.finalContent + result.Exports["agent_history"] = string(historyJSON) + // Cap iteration count: the for-loop post-increments past MaxIterations on natural exit + iterations := state.iteration + if iterations > step.MaxIterations { + iterations = step.MaxIterations + } + result.Exports["agent_iterations"] = iterations + result.Exports["agent_total_tokens"] = state.totalTokens + result.Exports["agent_prompt_tokens"] = state.promptTokens + result.Exports["agent_completion_tokens"] = state.completionTokens + result.Exports["agent_tool_results"] = string(toolResultsJSON) + + // Planning stage export (Group 2) + if state.planContent != "" { + result.Exports["agent_plan"] = state.planContent + } + + // Multi-goal results export (Group 3) + if len(state.goalResults) > 1 { + goalResultsJSON, err := json.Marshal(state.goalResults) + if err != nil { + log.Warn("Failed to marshal goal results", zap.Error(err)) + goalResultsJSON = []byte("[]") + } + result.Exports["agent_goal_results"] = string(goalResultsJSON) + } + + result.Output = state.finalContent + result.Status = core.StepStatusSuccess + result.EndTime = time.Now() + result.Duration = result.EndTime.Sub(result.StartTime) + + return result, nil +} + +// executePlanningStage runs the planning phase before the main agent loop (Group 2) +func (e *AgentExecutor) executePlanningStage( + ctx context.Context, + state *agentState, + step *core.Step, + llmConfig *MergedLLMConfig, +) (string, error) { + log := logger.Get() + log.Debug("Executing planning stage", + zap.String("step", step.Name), + ) + + // Build planning messages + var planMessages []ChatMessage + + // System prompt if available + systemPrompt := step.SystemPrompt + if systemPrompt == "" && llmConfig.SystemPrompt != "" { + systemPrompt = llmConfig.SystemPrompt + } + if systemPrompt != "" { + planMessages = append(planMessages, ChatMessage{ + Role: string(core.LLMRoleSystem), + Content: systemPrompt, + }) + } + + // Plan prompt as user message + planMessages = append(planMessages, ChatMessage{ + Role: string(core.LLMRoleUser), + Content: step.PlanPrompt, + }) + + // Build planning request (no tools — just text generation) + planConfig := *llmConfig + if step.PlanMaxTokens != nil { + planConfig.MaxTokens = *step.PlanMaxTokens + } + + planState := &agentState{messages: planMessages} + response, err := e.callLLM(ctx, planState, nil, &planConfig) + if err != nil { + return "", fmt.Errorf("planning request failed: %w", err) + } + + if response == nil || len(response.Choices) == 0 { + return "", fmt.Errorf("empty response from planning request") + } + + // Track tokens from planning + state.totalTokens += response.Usage.TotalTokens + state.promptTokens += response.Usage.PromptTokens + state.completionTokens += response.Usage.CompletionTokens + + planContent := "" + if content, ok := response.Choices[0].Message.Content.(string); ok { + planContent = content + } + + log.Debug("Planning stage complete", + zap.String("step", step.Name), + zap.Int("plan_length", len(planContent)), + ) + + return planContent, nil +} + +// initMessages builds the initial conversation messages for a goal +func (e *AgentExecutor) initMessages(state *agentState, query string, step *core.Step, llmConfig *MergedLLMConfig, goalIdx int) { + // For the first goal, build from scratch or resume + if goalIdx == 0 { + // If we loaded resumed messages, just append the new query + if len(state.messages) > 0 { + state.messages = append(state.messages, ChatMessage{ + Role: string(core.LLMRoleUser), + Content: query, + }) + return + } + + // Start fresh conversation + systemPrompt := step.SystemPrompt + if systemPrompt == "" && llmConfig.SystemPrompt != "" { + systemPrompt = llmConfig.SystemPrompt + } + if systemPrompt != "" { + state.messages = append(state.messages, ChatMessage{ + Role: string(core.LLMRoleSystem), + Content: systemPrompt, + }) + } + + // Prepend plan if available (Group 2) + if state.planContent != "" { + state.messages = append(state.messages, ChatMessage{ + Role: string(core.LLMRoleAssistant), + Content: "Here is my plan:\n\n" + state.planContent, + }) + } + + // User query + state.messages = append(state.messages, ChatMessage{ + Role: string(core.LLMRoleUser), + Content: query, + }) + return + } + + // For subsequent goals in multi-goal mode, append new user message + state.messages = append(state.messages, ChatMessage{ + Role: string(core.LLMRoleUser), + Content: query, + }) +} + +// callLLMWithFallback calls the LLM with optional per-agent model fallback (Group 5) +func (e *AgentExecutor) callLLMWithFallback( + ctx context.Context, + state *agentState, + tools []core.LLMTool, + llmConfig *MergedLLMConfig, + models []string, + responseFormat *core.LLMResponseFormat, +) (*ChatCompletionResponse, error) { + // If step specifies preferred models, try each in order + if len(models) > 0 { + var lastErr error + for _, model := range models { + modelConfig := *llmConfig + modelConfig.Model = model + if responseFormat != nil { + modelConfig.ResponseFormat = responseFormat + } + response, err := e.callLLM(ctx, state, tools, &modelConfig) + if err == nil && (response == nil || response.Error == nil) { + return response, nil + } + lastErr = err + logger.Get().Debug("Model fallback: trying next model", + zap.String("failed_model", model), + zap.Error(err), + ) + } + // Fall through to default provider rotation + logger.Get().Warn("All specified models failed, falling back to default", + zap.Error(lastErr), + ) + } + + // Default path: use standard provider rotation + if responseFormat != nil { + configWithFormat := *llmConfig + configWithFormat.ResponseFormat = responseFormat + return e.callLLM(ctx, state, tools, &configWithFormat) + } + return e.callLLM(ctx, state, tools, llmConfig) +} + +// callLLM sends the current conversation to the LLM and returns the response +func (e *AgentExecutor) callLLM( + ctx context.Context, + state *agentState, + tools []core.LLMTool, + llmConfig *MergedLLMConfig, +) (*ChatCompletionResponse, error) { + log := logger.Get() + + request := &ChatCompletionRequest{ + Model: llmConfig.Model, + Messages: state.messages, + MaxTokens: llmConfig.MaxTokens, + Temperature: llmConfig.Temperature, + TopP: llmConfig.TopP, + TopK: llmConfig.TopK, + Tools: tools, + Stream: llmConfig.Stream, + ResponseFormat: llmConfig.ResponseFormat, + } + + // Provider rotation with retries + var response *ChatCompletionResponse + var lastErr error + + maxRetries := llmConfig.MaxRetries + if maxRetries <= 0 { + maxRetries = 3 + } + providerCount := e.config.LLM.GetProviderCount() + if providerCount == 0 { + return nil, fmt.Errorf("no LLM providers configured") + } + + totalAttempts := maxRetries * providerCount + + for attempt := 0; attempt < totalAttempts; attempt++ { + provider := e.config.LLM.GetCurrentProvider() + if provider == nil { + lastErr = fmt.Errorf("no LLM providers available") + break + } + + if llmConfig.Model == "" { + request.Model = provider.Model + } + + log.Debug("Agent LLM request", + zap.String("provider", provider.Provider), + zap.String("model", request.Model), + zap.Int("attempt", attempt+1), + ) + + response, lastErr = e.sendChatRequest(ctx, provider, request, llmConfig) + + if lastErr == nil && response.Error == nil { + break + } + + if isProviderError(lastErr) || isRateLimitError(response) { + e.config.LLM.RotateProvider() + } + + if attempt < totalAttempts-1 { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(time.Duration(attempt+1) * 500 * time.Millisecond): + } + } + } + + if lastErr != nil { + return nil, lastErr + } + + if response != nil && response.Error != nil { + return nil, fmt.Errorf("LLM API error: %s (%s)", response.Error.Message, response.Error.Type) + } + + return response, nil +} + +// sendChatRequest sends an HTTP request to the LLM provider (delegates to LLMExecutor's logic) +func (e *AgentExecutor) sendChatRequest( + ctx context.Context, + provider *config.LLMProvider, + request *ChatCompletionRequest, + llmConfig *MergedLLMConfig, +) (*ChatCompletionResponse, error) { + llmExec := &LLMExecutor{config: e.config} + return llmExec.sendChatRequest(ctx, provider, request, llmConfig) +} + +// executeToolCalls executes tool calls from the LLM response, with tracing hooks (Group 7) +func (e *AgentExecutor) executeToolCalls( + ctx context.Context, + toolCalls []core.LLMToolCall, + state *agentState, + step *core.Step, + execCtx *core.ExecutionContext, +) ([]ChatMessage, error) { + log := logger.Get() + + parallelToolCalls := true + if step.ParallelToolCalls != nil { + parallelToolCalls = *step.ParallelToolCalls + } + + if parallelToolCalls && len(toolCalls) > 1 { + return e.executeToolCallsParallel(ctx, toolCalls, state, step, execCtx, log) + } + + return e.executeToolCallsSequential(ctx, toolCalls, state, step, execCtx, log) +} + +// executeToolCallsSequential executes tool calls one at a time +func (e *AgentExecutor) executeToolCallsSequential( + ctx context.Context, + toolCalls []core.LLMToolCall, + state *agentState, + step *core.Step, + execCtx *core.ExecutionContext, + log *zap.Logger, +) ([]ChatMessage, error) { + messages := make([]ChatMessage, 0, len(toolCalls)) + + for _, tc := range toolCalls { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + // Execute on_tool_start hook (Group 7) + e.executeToolHook(step.OnToolStart, tc.Function.Name, tc.Function.Arguments, "", 0, state.iteration, nil, execCtx) + + startTime := time.Now() + result, err := e.executeSingleToolCall(ctx, tc, state, execCtx, log) + duration := time.Since(startTime) + if err != nil { + result = fmt.Sprintf("Error executing tool '%s': %s", tc.Function.Name, err.Error()) + } + + // Execute on_tool_end hook (Group 7) + e.executeToolHook(step.OnToolEnd, tc.Function.Name, tc.Function.Arguments, result, duration.Milliseconds(), state.iteration, err, execCtx) + + log.Debug("Tool call completed", + zap.String("tool", tc.Function.Name), + zap.String("call_id", tc.ID), + zap.Duration("duration", duration), + ) + + state.toolResults = append(state.toolResults, map[string]interface{}{ + "tool_call_id": tc.ID, + "tool_name": tc.Function.Name, + "result": result, + }) + + messages = append(messages, ChatMessage{ + Role: string(core.LLMRoleTool), + Content: result, + ToolCallID: tc.ID, + }) + } + + return messages, nil +} + +// executeToolCallsParallel executes tool calls concurrently +func (e *AgentExecutor) executeToolCallsParallel( + ctx context.Context, + toolCalls []core.LLMToolCall, + state *agentState, + step *core.Step, + execCtx *core.ExecutionContext, + log *zap.Logger, +) ([]ChatMessage, error) { + type toolResult struct { + index int + message ChatMessage + } + + results := make(chan toolResult, len(toolCalls)) + var wg sync.WaitGroup + + for i, tc := range toolCalls { + wg.Add(1) + go func(idx int, call core.LLMToolCall) { + defer wg.Done() + + select { + case <-ctx.Done(): + return + default: + } + + // Execute on_tool_start hook (Group 7) + e.executeToolHook(step.OnToolStart, call.Function.Name, call.Function.Arguments, "", 0, state.iteration, nil, execCtx) + + startTime := time.Now() + result, err := e.executeSingleToolCall(ctx, call, state, execCtx, log) + duration := time.Since(startTime) + if err != nil { + result = fmt.Sprintf("Error executing tool '%s': %s", call.Function.Name, err.Error()) + } + + // Execute on_tool_end hook (Group 7) + e.executeToolHook(step.OnToolEnd, call.Function.Name, call.Function.Arguments, result, duration.Milliseconds(), state.iteration, err, execCtx) + + log.Debug("Tool call completed", + zap.String("tool", call.Function.Name), + zap.String("call_id", call.ID), + zap.Duration("duration", duration), + ) + + results <- toolResult{ + index: idx, + message: ChatMessage{ + Role: string(core.LLMRoleTool), + Content: result, + ToolCallID: call.ID, + }, + } + }(i, tc) + } + + wg.Wait() + close(results) + + // Collect and order results, populate toolResults + ordered := make([]ChatMessage, len(toolCalls)) + for r := range results { + ordered[r.index] = r.message + } + for i, msg := range ordered { + state.toolResults = append(state.toolResults, map[string]interface{}{ + "tool_call_id": toolCalls[i].ID, + "tool_name": toolCalls[i].Function.Name, + "result": msg.Content, + }) + } + + return ordered, nil +} + +// executeSingleToolCall executes a single tool call via the ToolExecutorRegistry +func (e *AgentExecutor) executeSingleToolCall( + ctx context.Context, + tc core.LLMToolCall, + state *agentState, + execCtx *core.ExecutionContext, + log *zap.Logger, +) (string, error) { + // Use ToolExecutorRegistry if available (Group 1) + if state.toolRegistry != nil { + return executeToolCallViaRegistry(ctx, tc, state.toolRegistry, execCtx, log) + } + + // Fallback to legacy dispatch (should not happen in normal flow) + return e.executeSingleToolCallLegacy(tc, execCtx, log) +} + +// executeSingleToolCallLegacy is the legacy dispatch path (kept for safety) +func (e *AgentExecutor) executeSingleToolCallLegacy( + tc core.LLMToolCall, + execCtx *core.ExecutionContext, + log *zap.Logger, +) (string, error) { + funcName := tc.Function.Name + argsJSON := tc.Function.Arguments + + log.Debug("Executing tool call (legacy)", + zap.String("tool", funcName), + zap.String("args", argsJSON), + ) + + var args map[string]interface{} + if argsJSON != "" { + if err := json.Unmarshal([]byte(argsJSON), &args); err != nil { + return "", fmt.Errorf("failed to parse tool arguments for %s: %w", funcName, err) + } + } + if args == nil { + args = make(map[string]interface{}) + } + + return e.executePresetTool(funcName, args, execCtx) +} + +// executePresetTool runs a preset tool by building a function call expression +func (e *AgentExecutor) executePresetTool(funcName string, args map[string]interface{}, execCtx *core.ExecutionContext) (string, error) { + expr := buildPresetCallExpr(funcName, args) + vars := execCtx.GetVariables() + result, err := e.functionRegistry.Execute(expr, vars) + if err != nil { + return "", fmt.Errorf("preset tool '%s' failed: %w", funcName, err) + } + return formatToolResult(result), nil +} + +// executeToolHook evaluates a JS hook expression with tool call context (Group 7) +func (e *AgentExecutor) executeToolHook( + hook string, + toolName string, + toolArgs string, + result string, + durationMs int64, + iteration int, + toolErr error, + execCtx *core.ExecutionContext, +) { + if hook == "" || e.functionRegistry == nil { + return + } + + log := logger.Get() + vars := execCtx.GetVariables() + vars["tool_name"] = toolName + vars["tool_args"] = toolArgs + vars["result"] = result + vars["duration"] = durationMs + vars["iteration"] = iteration + if toolErr != nil { + vars["error"] = toolErr.Error() + } else { + vars["error"] = "" + } + + if _, err := e.functionRegistry.Execute(hook, vars); err != nil { + log.Warn("Tool hook execution failed (non-blocking)", + zap.String("hook", hook), + zap.Error(err), + ) + } +} + +// requestStructuredOutput makes a final LLM call to enforce structured output (Group 6) +func (e *AgentExecutor) requestStructuredOutput( + ctx context.Context, + state *agentState, + schema *core.LLMResponseFormat, + llmConfig *MergedLLMConfig, +) string { + log := logger.Get() + + // Only request if we have a schema and the last response might not be structured + if schema == nil { + return "" + } + + // Check if current content is already valid JSON matching the schema + var check interface{} + if json.Unmarshal([]byte(state.finalContent), &check) == nil { + // Already valid JSON, likely structured + return "" + } + + // Add instruction to produce structured output + state.messages = append(state.messages, ChatMessage{ + Role: string(core.LLMRoleUser), + Content: "Please provide your final answer in the structured JSON format specified.", + }) + + configWithSchema := *llmConfig + configWithSchema.ResponseFormat = schema + + response, err := e.callLLM(ctx, state, nil, &configWithSchema) + if err != nil { + log.Warn("Structured output request failed", + zap.Error(err), + ) + return "" + } + + if response == nil || len(response.Choices) == 0 { + return "" + } + + // Track tokens + state.totalTokens += response.Usage.TotalTokens + state.promptTokens += response.Usage.PromptTokens + state.completionTokens += response.Usage.CompletionTokens + + if content, ok := response.Choices[0].Message.Content.(string); ok { + return content + } + return "" +} + +// buildPresetCallExpr builds a JS function call expression from a preset tool name and arguments +func buildPresetCallExpr(funcName string, args map[string]interface{}) string { + switch funcName { + case "bash": + return fmt.Sprintf("bash(%s)", jsQuote(getStringArg(args, "command"))) + case "read_file": + return fmt.Sprintf("read_file(%s)", jsQuote(getStringArg(args, "path"))) + case "read_lines": + return fmt.Sprintf("read_lines(%s)", jsQuote(getStringArg(args, "path"))) + case "file_exists": + return fmt.Sprintf("file_exists(%s)", jsQuote(getStringArg(args, "path"))) + case "file_length": + return fmt.Sprintf("file_length(%s)", jsQuote(getStringArg(args, "path"))) + case "append_file": + return fmt.Sprintf("append_file(%s, %s)", jsQuote(getStringArg(args, "dest")), jsQuote(getStringArg(args, "content"))) + case "save_content": + return fmt.Sprintf("save_content(%s, %s)", jsQuote(getStringArg(args, "content")), jsQuote(getStringArg(args, "path"))) + case "glob": + return fmt.Sprintf("glob(%s)", jsQuote(getStringArg(args, "pattern"))) + case "grep_string": + return fmt.Sprintf("grep_string(%s, %s)", jsQuote(getStringArg(args, "source")), jsQuote(getStringArg(args, "str"))) + case "grep_regex": + return fmt.Sprintf("grep_regex(%s, %s)", jsQuote(getStringArg(args, "source")), jsQuote(getStringArg(args, "pattern"))) + case "http_get": + return fmt.Sprintf("http_get(%s)", jsQuote(getStringArg(args, "url"))) + case "http_request": + return fmt.Sprintf("http_request(%s, %s, %s, %s)", + jsQuote(getStringArg(args, "url")), + jsQuote(getStringArg(args, "method")), + jsQuote(getStringArg(args, "headers")), + jsQuote(getStringArg(args, "body")), + ) + case "jq": + return fmt.Sprintf("jq(%s, %s)", jsQuote(getStringArg(args, "json_data")), jsQuote(getStringArg(args, "expression"))) + case "exec_python": + return fmt.Sprintf("exec_python(%s)", jsQuote(getStringArg(args, "code"))) + case "exec_python_file": + return fmt.Sprintf("exec_python_file(%s)", jsQuote(getStringArg(args, "path"))) + case "run_module": + return fmt.Sprintf("run_module(%s, %s, %s)", jsQuote(getStringArg(args, "module")), jsQuote(getStringArg(args, "target")), jsQuote(getStringArg(args, "params"))) + case "run_flow": + return fmt.Sprintf("run_flow(%s, %s, %s)", jsQuote(getStringArg(args, "flow")), jsQuote(getStringArg(args, "target")), jsQuote(getStringArg(args, "params"))) + default: + var argStrs []string + for _, v := range args { + argStrs = append(argStrs, jsQuote(fmt.Sprintf("%v", v))) + } + return fmt.Sprintf("%s(%s)", funcName, strings.Join(argStrs, ", ")) + } +} + +// getStringArg safely extracts a string argument from the args map +func getStringArg(args map[string]interface{}, key string) string { + if v, ok := args[key]; ok { + if s, ok := v.(string); ok { + return s + } + return fmt.Sprintf("%v", v) + } + return "" +} + +// jsQuote returns a JavaScript string literal, escaping special characters +func jsQuote(s string) string { + s = strings.ReplaceAll(s, `\`, `\\`) + s = strings.ReplaceAll(s, `"`, `\"`) + s = strings.ReplaceAll(s, "\n", `\n`) + s = strings.ReplaceAll(s, "\r", `\r`) + s = strings.ReplaceAll(s, "\t", `\t`) + return `"` + s + `"` +} + +// formatToolResult converts a tool execution result to a string for the LLM +func formatToolResult(result interface{}) string { + if result == nil { + return "" + } + + switch v := result.(type) { + case string: + return v + case bool: + if v { + return "true" + } + return "false" + case int, int64, float64: + return fmt.Sprintf("%v", v) + default: + jsonBytes, err := json.Marshal(v) + if err != nil { + return fmt.Sprintf("%v", v) + } + return string(jsonBytes) + } +} + +// applyMessageWindow trims the conversation to max_messages, keeping the system message +func (e *AgentExecutor) applyMessageWindow(state *agentState, maxMessages int) { + if maxMessages <= 0 || len(state.messages) <= maxMessages { + return + } + + var systemMsg *ChatMessage + startIdx := 0 + if len(state.messages) > 0 && state.messages[0].Role == string(core.LLMRoleSystem) { + systemMsg = &state.messages[0] + startIdx = 1 + } + + nonSystemMsgs := state.messages[startIdx:] + keepCount := maxMessages + if systemMsg != nil { + keepCount-- + } + + if len(nonSystemMsgs) > keepCount { + nonSystemMsgs = nonSystemMsgs[len(nonSystemMsgs)-keepCount:] + } + + if systemMsg != nil { + state.messages = make([]ChatMessage, 0, keepCount+1) + state.messages = append(state.messages, *systemMsg) + state.messages = append(state.messages, nonSystemMsgs...) + } else { + state.messages = nonSystemMsgs + } +} + +// applyMessageWindowWithSummary trims with LLM-based summarization of dropped messages (Group 4) +func (e *AgentExecutor) applyMessageWindowWithSummary( + ctx context.Context, + state *agentState, + maxMessages int, + llmConfig *MergedLLMConfig, +) { + if maxMessages <= 0 || len(state.messages) <= maxMessages { + return + } + + log := logger.Get() + + var systemMsg *ChatMessage + startIdx := 0 + if len(state.messages) > 0 && state.messages[0].Role == string(core.LLMRoleSystem) { + systemMsg = &state.messages[0] + startIdx = 1 + } + + nonSystemMsgs := state.messages[startIdx:] + keepCount := maxMessages + if systemMsg != nil { + keepCount-- // Account for system message + keepCount-- // Account for summary message we'll insert + } + if keepCount < 1 { + keepCount = 1 + } + + if len(nonSystemMsgs) <= keepCount { + return + } + + // Messages to be dropped + dropCount := len(nonSystemMsgs) - keepCount + droppedMsgs := nonSystemMsgs[:dropCount] + keptMsgs := nonSystemMsgs[dropCount:] + + // Build summary of dropped messages + var summaryParts []string + for _, msg := range droppedMsgs { + if content, ok := msg.Content.(string); ok && content != "" { + role := msg.Role + // Truncate long messages + if len(content) > 200 { + content = content[:200] + "..." + } + summaryParts = append(summaryParts, fmt.Sprintf("[%s]: %s", role, content)) + } + } + + if len(summaryParts) == 0 { + // No meaningful content to summarize, just truncate + e.applyMessageWindow(state, maxMessages) + return + } + + // Ask LLM to summarize the dropped context + summaryPrompt := "Summarize the following conversation context concisely, preserving key information:\n\n" + strings.Join(summaryParts, "\n") + summaryMessages := []ChatMessage{ + {Role: string(core.LLMRoleUser), Content: summaryPrompt}, + } + + summaryState := &agentState{messages: summaryMessages} + summaryConfig := *llmConfig + summaryConfig.MaxTokens = 300 + + response, err := e.callLLM(ctx, summaryState, nil, &summaryConfig) + if err != nil { + log.Warn("Conversation summarization failed, falling back to simple truncation", + zap.Error(err), + ) + e.applyMessageWindow(state, maxMessages) + return + } + + summaryContent := "" + if response != nil && len(response.Choices) > 0 { + if content, ok := response.Choices[0].Message.Content.(string); ok { + summaryContent = content + } + // Track summarization tokens + state.totalTokens += response.Usage.TotalTokens + state.promptTokens += response.Usage.PromptTokens + state.completionTokens += response.Usage.CompletionTokens + } + + if summaryContent == "" { + e.applyMessageWindow(state, maxMessages) + return + } + + // Rebuild messages with summary + state.messages = make([]ChatMessage, 0, keepCount+2) + if systemMsg != nil { + state.messages = append(state.messages, *systemMsg) + } + state.messages = append(state.messages, ChatMessage{ + Role: string(core.LLMRoleSystem), + Content: "[Summary of earlier conversation]\n" + summaryContent, + }) + state.messages = append(state.messages, keptMsgs...) +} + +// loadConversation loads a prior conversation from a JSON file +func (e *AgentExecutor) loadConversation(state *agentState, path string) error { + data, err := os.ReadFile(path) + if err != nil { + return fmt.Errorf("failed to read conversation file: %w", err) + } + + var messages []ChatMessage + if err := json.Unmarshal(data, &messages); err != nil { + return fmt.Errorf("failed to parse conversation file: %w", err) + } + + state.messages = messages + return nil +} + +// persistConversation saves the conversation to a JSON file +func (e *AgentExecutor) persistConversation(state *agentState, path string) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o755); err != nil { + return fmt.Errorf("failed to create directory: %w", err) + } + + data, err := json.MarshalIndent(state.messages, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal conversation: %w", err) + } + + if err := os.WriteFile(path, data, 0o644); err != nil { + return fmt.Errorf("failed to write conversation file: %w", err) + } + + return nil +} + +// getMergedConfig merges global and step-level LLM configuration +func (e *AgentExecutor) getMergedConfig(step *core.Step) *MergedLLMConfig { + llmExec := &LLMExecutor{config: e.config} + return llmExec.getMergedConfig(step) +} + +// fail is a helper that sets the result to failed state +func (e *AgentExecutor) fail(result *core.StepResult, err error) (*core.StepResult, error) { + result.Status = core.StepStatusFailed + result.Error = err + result.EndTime = time.Now() + result.Duration = result.EndTime.Sub(result.StartTime) + return result, err +} diff --git a/internal/executor/agent_executor_test.go b/internal/executor/agent_executor_test.go new file mode 100644 index 0000000..38bf672 --- /dev/null +++ b/internal/executor/agent_executor_test.go @@ -0,0 +1,1862 @@ +package executor + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + + "github.com/j3ssie/osmedeus/v5/internal/config" + "github.com/j3ssie/osmedeus/v5/internal/core" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// mockLLMResponse creates a ChatCompletionResponse with text content and no tool calls +func mockLLMResponse(content string) ChatCompletionResponse { + return ChatCompletionResponse{ + ID: "test-id", + Model: "test-model", + Choices: []ChatChoice{ + { + Index: 0, + Message: ChatMessage{ + Role: "assistant", + Content: content, + }, + FinishReason: "stop", + }, + }, + Usage: ChatUsage{ + PromptTokens: 10, + CompletionTokens: 20, + TotalTokens: 30, + }, + } +} + +// mockLLMToolCallResponse creates a response with tool calls +func mockLLMToolCallResponse(toolCalls []core.LLMToolCall) ChatCompletionResponse { + return ChatCompletionResponse{ + ID: "test-id", + Model: "test-model", + Choices: []ChatChoice{ + { + Index: 0, + Message: ChatMessage{ + Role: "assistant", + Content: "", + ToolCalls: toolCalls, + }, + FinishReason: "tool_calls", + }, + }, + Usage: ChatUsage{ + PromptTokens: 10, + CompletionTokens: 20, + TotalTokens: 30, + }, + } +} + +// newMockLLMServer creates a mock OpenAI-compatible API server. +// handler is called for each request and should write the response. +func newMockLLMServer(handler http.HandlerFunc) *httptest.Server { + return httptest.NewServer(handler) +} + +// newMockConfig creates a config with a mock LLM provider pointing to the given URL +func newMockConfig(t *testing.T, serverURL string) *config.Config { + t.Helper() + baseDir := t.TempDir() + return &config.Config{ + BaseFolder: baseDir, + WorkspacesPath: filepath.Join(baseDir, "workspaces"), + LLM: config.LLMConfig{ + LLMProviders: []config.LLMProvider{ + { + Provider: "mock", + BaseURL: serverURL, + AuthToken: "test-token", + Model: "test-model", + }, + }, + MaxTokens: 1000, + MaxRetries: 1, + Timeout: "30s", + }, + } +} + +// ============================================================================ +// Agent Executor Unit Tests +// ============================================================================ + +func TestAgentExecutor_Name(t *testing.T) { + executor := NewAgentExecutor(nil, nil) + assert.Equal(t, "agent", executor.Name()) +} + +func TestAgentExecutor_StepTypes(t *testing.T) { + executor := NewAgentExecutor(nil, nil) + types := executor.StepTypes() + assert.Len(t, types, 1) + assert.Equal(t, core.StepTypeAgent, types[0]) +} + +func TestAgentExecutor_ValidationErrors(t *testing.T) { + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + t.Run("missing config", func(t *testing.T) { + executor := NewAgentExecutor(nil, nil) + step := &core.Step{Name: "test", Type: core.StepTypeAgent} + + result, err := executor.Execute(ctx, step, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "config not set") + assert.Equal(t, core.StepStatusFailed, result.Status) + }) + + t.Run("missing query", func(t *testing.T) { + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {}) + defer server.Close() + cfg := newMockConfig(t, server.URL) + + executor := NewAgentExecutor(nil, nil) + executor.SetConfig(cfg) + + step := &core.Step{ + Name: "test", + Type: core.StepTypeAgent, + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "requires 'query'") + assert.Equal(t, core.StepStatusFailed, result.Status) + }) + + t.Run("missing max_iterations", func(t *testing.T) { + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {}) + defer server.Close() + cfg := newMockConfig(t, server.URL) + + executor := NewAgentExecutor(nil, nil) + executor.SetConfig(cfg) + + step := &core.Step{ + Name: "test", + Type: core.StepTypeAgent, + Query: "test query", + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "requires 'max_iterations'") + assert.Equal(t, core.StepStatusFailed, result.Status) + }) + + t.Run("missing agent_tools", func(t *testing.T) { + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {}) + defer server.Close() + cfg := newMockConfig(t, server.URL) + + executor := NewAgentExecutor(nil, nil) + executor.SetConfig(cfg) + + step := &core.Step{ + Name: "test", + Type: core.StepTypeAgent, + Query: "test query", + MaxIterations: 5, + } + + result, err := executor.Execute(ctx, step, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "requires 'agent_tools'") + assert.Equal(t, core.StepStatusFailed, result.Status) + }) + + t.Run("unknown preset tool", func(t *testing.T) { + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {}) + defer server.Close() + cfg := newMockConfig(t, server.URL) + + executor := NewAgentExecutor(nil, nil) + executor.SetConfig(cfg) + + step := &core.Step{ + Name: "test", + Type: core.StepTypeAgent, + Query: "test query", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "nonexistent_tool"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "unknown preset tool") + assert.Equal(t, core.StepStatusFailed, result.Status) + }) +} + +func TestAgentExecutor_SimpleCompletion(t *testing.T) { + // Mock server that responds without tool calls (agent should complete in 1 iteration) + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + resp := mockLLMResponse("The analysis is complete.") + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "simple-agent", + Type: core.StepTypeAgent, + Query: "Analyze the target", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, "The analysis is complete.", result.Output) + assert.Equal(t, "The analysis is complete.", result.Exports["agent_content"]) + assert.Equal(t, 1, result.Exports["agent_iterations"]) + assert.Equal(t, 30, result.Exports["agent_total_tokens"]) + assert.NotEmpty(t, result.Exports["agent_history"]) +} + +func TestAgentExecutor_ToolCallLoop(t *testing.T) { + // Mock server that: + // 1st call: returns a tool call for exec_cmd + // 2nd call: returns final text response + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + if count == 1 { + // First call: return tool call + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_1", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "bash", + Arguments: `{"command": "echo hello"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + } else { + // Second call: return final response + resp := mockLLMResponse("Command output was: hello") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "tool-call-agent", + Type: core.StepTypeAgent, + Query: "Run echo hello", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, "Command output was: hello", result.Output) + assert.Equal(t, 2, result.Exports["agent_iterations"]) + // 2 calls * 30 tokens each + assert.Equal(t, 60, result.Exports["agent_total_tokens"]) +} + +func TestAgentExecutor_MaxIterationsLimit(t *testing.T) { + // Mock server that always returns tool calls (never finishes) + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_1", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "bash", + Arguments: `{"command": "echo loop"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "loop-agent", + Type: core.StepTypeAgent, + Query: "Keep running", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, 3, result.Exports["agent_iterations"]) +} + +func TestAgentExecutor_StopCondition(t *testing.T) { + // Mock server: first call returns "checking...", second returns "DONE" + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + if count == 1 { + // Return tool call first + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_1", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "bash", + Arguments: `{"command": "echo check"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + } else { + // Return text with DONE keyword + resp := mockLLMResponse("Analysis DONE successfully") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "stop-agent", + Type: core.StepTypeAgent, + Query: "Analyze and say DONE when finished", + MaxIterations: 10, + StopCondition: `contains(agent_content, "DONE")`, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Contains(t, result.Output, "DONE") + // Should have stopped after 2 iterations (not 10) + assert.Equal(t, 2, result.Exports["agent_iterations"]) +} + +func TestAgentExecutor_SystemPrompt(t *testing.T) { + // Verify that system prompt is included in the request + var receivedMessages []ChatMessage + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + receivedMessages = req.Messages + + w.Header().Set("Content-Type", "application/json") + resp := mockLLMResponse("Response") + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "system-prompt-agent", + Type: core.StepTypeAgent, + SystemPrompt: "You are a security analyst.", + Query: "Analyze target", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + _, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + + // Verify messages contain system prompt and user query + require.Len(t, receivedMessages, 2) + assert.Equal(t, "system", receivedMessages[0].Role) + assert.Equal(t, "You are a security analyst.", receivedMessages[0].Content) + assert.Equal(t, "user", receivedMessages[1].Role) + assert.Equal(t, "Analyze target", receivedMessages[1].Content) +} + +func TestAgentExecutor_ToolsInRequest(t *testing.T) { + // Verify that tools are included in the LLM request + var receivedTools []core.LLMTool + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + receivedTools = req.Tools + + w.Header().Set("Content-Type", "application/json") + resp := mockLLMResponse("Done") + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "tools-check", + Type: core.StepTypeAgent, + Query: "Test", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{ + {Preset: "bash"}, + {Preset: "read_file"}, + { + Name: "custom", + Description: "Custom tool", + Parameters: map[string]interface{}{"type": "object"}, + }, + }, + } + + _, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + + // Verify 3 tools were sent to the LLM + assert.Len(t, receivedTools, 3) + assert.Equal(t, "bash", receivedTools[0].Function.Name) + assert.Equal(t, "read_file", receivedTools[1].Function.Name) + assert.Equal(t, "custom", receivedTools[2].Function.Name) +} + +func TestAgentExecutor_MemoryPersist(t *testing.T) { + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + resp := mockLLMResponse("Persisted response") + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + persistPath := filepath.Join(t.TempDir(), "agent", "conversation.json") + + step := &core.Step{ + Name: "persist-agent", + Type: core.StepTypeAgent, + Query: "Test persist", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + Memory: &core.AgentMemoryConfig{ + PersistPath: persistPath, + }, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + + // Verify conversation was persisted + data, err := os.ReadFile(persistPath) + require.NoError(t, err) + + var messages []ChatMessage + err = json.Unmarshal(data, &messages) + require.NoError(t, err) + assert.GreaterOrEqual(t, len(messages), 2) // at least user + assistant +} + +func TestAgentExecutor_MemoryResume(t *testing.T) { + // Create a conversation file to resume from + tmpDir := t.TempDir() + resumePath := filepath.Join(tmpDir, "resume.json") + + priorMessages := []ChatMessage{ + {Role: "system", Content: "You are a helper."}, + {Role: "user", Content: "Previous question"}, + {Role: "assistant", Content: "Previous answer"}, + } + data, _ := json.MarshalIndent(priorMessages, "", " ") + os.WriteFile(resumePath, data, 0o644) + + var receivedMessages []ChatMessage + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + receivedMessages = req.Messages + + w.Header().Set("Content-Type", "application/json") + resp := mockLLMResponse("Resumed response") + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "resume-agent", + Type: core.StepTypeAgent, + Query: "Follow-up question", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + Memory: &core.AgentMemoryConfig{ + ResumePath: resumePath, + }, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + + // Verify resumed messages + new query + // Should be: system + previous user + previous assistant + new user query + assert.Len(t, receivedMessages, 4) + assert.Equal(t, "system", receivedMessages[0].Role) + assert.Equal(t, "user", receivedMessages[1].Role) + assert.Equal(t, "Previous question", receivedMessages[1].Content) + assert.Equal(t, "assistant", receivedMessages[2].Role) + assert.Equal(t, "user", receivedMessages[3].Role) + assert.Equal(t, "Follow-up question", receivedMessages[3].Content) +} + +func TestAgentExecutor_CustomToolHandler(t *testing.T) { + // First call returns a tool call for "custom_tool", second returns final answer + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + if count == 1 { + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_1", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "greet", + Arguments: `{"name": "World"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + } else { + resp := mockLLMResponse("Greeting sent") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "custom-handler-agent", + Type: core.StepTypeAgent, + Query: "Greet World", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{ + { + Name: "greet", + Description: "Greet someone", + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "name": map[string]interface{}{"type": "string"}, + }, + }, + Handler: `"Hello, " + args.name + "!"`, + }, + }, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, "Greeting sent", result.Output) +} + +func TestAgentExecutor_ContextCancellation(t *testing.T) { + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + // This should not be reached because context is already cancelled + w.Header().Set("Content-Type", "application/json") + resp := mockLLMResponse("Should not reach") + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // Cancel immediately + + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "cancel-agent", + Type: core.StepTypeAgent, + Query: "This should be cancelled", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + assert.Error(t, err) + assert.Equal(t, core.StepStatusFailed, result.Status) +} + +func TestAgentExecutor_MultipleToolCalls(t *testing.T) { + // Mock server: first call returns 2 tool calls, second returns final answer + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + if count == 1 { + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_1", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "bash", + Arguments: `{"command": "echo first"}`, + }, + }, + { + ID: "call_2", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "bash", + Arguments: `{"command": "echo second"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + } else { + resp := mockLLMResponse("Both commands executed") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "multi-tool-agent", + Type: core.StepTypeAgent, + Query: "Run two commands", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, "Both commands executed", result.Output) + assert.Equal(t, 2, result.Exports["agent_iterations"]) +} + +// ============================================================================ +// Helper Function Tests +// ============================================================================ + +func TestBuildPresetCallExpr(t *testing.T) { + tests := []struct { + name string + funcName string + args map[string]interface{} + expected string + }{ + { + name: "bash", + funcName: "bash", + args: map[string]interface{}{"command": "echo hello"}, + expected: `bash("echo hello")`, + }, + { + name: "read_file", + funcName: "read_file", + args: map[string]interface{}{"path": "/tmp/test.txt"}, + expected: `read_file("/tmp/test.txt")`, + }, + { + name: "save_content", + funcName: "save_content", + args: map[string]interface{}{"content": "hello", "path": "/tmp/out.txt"}, + expected: `save_content("hello", "/tmp/out.txt")`, + }, + { + name: "grep_string", + funcName: "grep_string", + args: map[string]interface{}{"source": "/tmp/file.txt", "str": "pattern"}, + expected: `grep_string("/tmp/file.txt", "pattern")`, + }, + { + name: "jq", + funcName: "jq", + args: map[string]interface{}{"json_data": `{"key":"value"}`, "expression": ".key"}, + expected: `jq("{\"key\":\"value\"}", ".key")`, + }, + { + name: "exec_python", + funcName: "exec_python", + args: map[string]interface{}{"code": "print('hello')"}, + expected: `exec_python("print('hello')")`, + }, + { + name: "exec_python_file", + funcName: "exec_python_file", + args: map[string]interface{}{"path": "/tmp/script.py"}, + expected: `exec_python_file("/tmp/script.py")`, + }, + { + name: "run_module", + funcName: "run_module", + args: map[string]interface{}{"module": "subdomain", "target": "example.com", "params": "threads=10"}, + expected: `run_module("subdomain", "example.com", "threads=10")`, + }, + { + name: "run_flow", + funcName: "run_flow", + args: map[string]interface{}{"flow": "general", "target": "example.com", "params": ""}, + expected: `run_flow("general", "example.com", "")`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := buildPresetCallExpr(tt.funcName, tt.args) + assert.Equal(t, tt.expected, result) + }) + } +} + +func TestJsQuote(t *testing.T) { + tests := []struct { + input string + expected string + }{ + {`hello`, `"hello"`}, + {`say "hi"`, `"say \"hi\""`}, + {"line1\nline2", `"line1\nline2"`}, + {`path\to\file`, `"path\\to\\file"`}, + {"tab\there", `"tab\there"`}, + } + + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + assert.Equal(t, tt.expected, jsQuote(tt.input)) + }) + } +} + +func TestFormatToolResult(t *testing.T) { + assert.Equal(t, "hello", formatToolResult("hello")) + assert.Equal(t, "true", formatToolResult(true)) + assert.Equal(t, "false", formatToolResult(false)) + assert.Equal(t, "42", formatToolResult(42)) + assert.Equal(t, "3.14", formatToolResult(3.14)) + assert.Equal(t, "", formatToolResult(nil)) + + // Complex type gets JSON serialized + result := formatToolResult(map[string]string{"key": "value"}) + assert.Contains(t, result, "key") + assert.Contains(t, result, "value") +} + +func TestGetStringArg(t *testing.T) { + args := map[string]interface{}{ + "str_arg": "hello", + "int_arg": 42, + "nil_arg": nil, + } + + assert.Equal(t, "hello", getStringArg(args, "str_arg")) + assert.Equal(t, "42", getStringArg(args, "int_arg")) + assert.Equal(t, "", getStringArg(args, "missing")) +} + +// ============================================================================ +// Integration Test: Full Agent Loop via Executor +// ============================================================================ + +func TestExecutor_AgentStep_FullIntegration(t *testing.T) { + // Mock LLM server that simulates a tool call loop + var callCount int32 + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + if count == 1 { + // First: ask to run a command + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_1", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "bash", + Arguments: `{"command": "echo integration_test"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + } else { + resp := mockLLMResponse("Integration test passed") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := testConfig(t) + cfg.LLM = config.LLMConfig{ + LLMProviders: []config.LLMProvider{ + { + Provider: "mock", + BaseURL: server.URL, + Model: "test-model", + }, + }, + MaxTokens: 1000, + MaxRetries: 1, + Timeout: "30s", + } + + module := &core.Workflow{ + Name: "test-agent-integration", + Kind: core.KindModule, + Steps: []core.Step{ + { + Name: "agent-step", + Type: core.StepTypeAgent, + Query: "Run echo integration_test", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{ + {Preset: "bash"}, + }, + Exports: map[string]string{ + "result": "{{agent_content}}", + }, + }, + { + Name: "verify-step", + Type: core.StepTypeBash, + Command: fmt.Sprintf("echo 'agent result: {{result}}'"), + }, + }, + } + + executor := NewExecutor() + executor.SetDryRun(false) + executor.SetSpinner(false) + + ctx := context.Background() + result, err := executor.ExecuteModule(ctx, module, map[string]string{ + "target": "test", + }, cfg) + + require.NoError(t, err) + assert.Equal(t, core.RunStatusCompleted, result.Status) + assert.Len(t, result.Steps, 2) + assert.Equal(t, core.StepStatusSuccess, result.Steps[0].Status) + assert.Equal(t, "Integration test passed", result.Steps[0].Output) + assert.Equal(t, core.StepStatusSuccess, result.Steps[1].Status) +} + +// ============================================================================ +// Group 2: Planning Stage Tests +// ============================================================================ + +func TestAgentExecutor_PlanningStage(t *testing.T) { + // Mock server: 1st call = planning response (no tools), 2nd call = final response + var callCount int32 + var receivedMessages [][]ChatMessage + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + receivedMessages = append(receivedMessages, req.Messages) + + if count == 1 { + // Planning phase: return plan text + resp := mockLLMResponse("Step 1: Scan ports\nStep 2: Check services\nStep 3: Report findings") + json.NewEncoder(w).Encode(resp) + } else { + // Execution phase: return final response + resp := mockLLMResponse("Plan executed successfully") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "plan-agent", + Type: core.StepTypeAgent, + PlanPrompt: "Create a plan for scanning test.com", + Query: "Execute the plan", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, "Plan executed successfully", result.Output) + + // Verify agent_plan export is populated + planContent, ok := result.Exports["agent_plan"] + assert.True(t, ok) + assert.Contains(t, planContent, "Step 1") + + // Verify planning request had the plan prompt + require.Len(t, receivedMessages, 2) + planMsgs := receivedMessages[0] + assert.Equal(t, "user", planMsgs[len(planMsgs)-1].Role) + assert.Equal(t, "Create a plan for scanning test.com", planMsgs[len(planMsgs)-1].Content) + + // Verify execution request includes the plan as assistant message + execMsgs := receivedMessages[1] + foundPlan := false + for _, msg := range execMsgs { + if msg.Role == "assistant" { + if content, ok := msg.Content.(string); ok && strings.Contains(content, "Step 1") { + foundPlan = true + } + } + } + assert.True(t, foundPlan, "execution messages should include the plan as assistant message") + + // Verify tokens are tracked from both planning and execution + assert.Equal(t, 60, result.Exports["agent_total_tokens"]) // 30 from planning + 30 from execution +} + +func TestAgentExecutor_PlanningStage_WithMaxTokens(t *testing.T) { + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + + // Return plan or final based on whether tools are present + if len(req.Tools) == 0 { + // Planning request — verify max_tokens + assert.Equal(t, 500, req.MaxTokens) + resp := mockLLMResponse("Short plan") + json.NewEncoder(w).Encode(resp) + } else { + resp := mockLLMResponse("Done") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + planMaxTokens := 500 + step := &core.Step{ + Name: "plan-tokens-agent", + Type: core.StepTypeAgent, + PlanPrompt: "Create a plan", + PlanMaxTokens: &planMaxTokens, + Query: "Execute", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) +} + +// ============================================================================ +// Group 3: Multi-Goal Execution Tests +// ============================================================================ + +func TestAgentExecutor_MultiGoal(t *testing.T) { + // Mock server: responds to each query with a unique response + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + resp := mockLLMResponse(fmt.Sprintf("Goal %d complete", count)) + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "multi-goal-agent", + Type: core.StepTypeAgent, + Queries: []string{"List files", "Summarize findings", "Report"}, + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + + assert.Equal(t, core.StepStatusSuccess, result.Status) + // Final content should be from the last goal + assert.Equal(t, "Goal 3 complete", result.Output) + + // Verify agent_goal_results export + goalResultsJSON, ok := result.Exports["agent_goal_results"] + assert.True(t, ok, "agent_goal_results should be exported for multi-goal") + assert.Contains(t, goalResultsJSON, "List files") + assert.Contains(t, goalResultsJSON, "Summarize findings") + assert.Contains(t, goalResultsJSON, "Report") + + // Tokens should be accumulated from all goals + assert.Equal(t, 90, result.Exports["agent_total_tokens"]) // 30 * 3 goals +} + +func TestAgentExecutor_QueryAndQueriesMutualExclusion(t *testing.T) { + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {}) + defer server.Close() + cfg := newMockConfig(t, server.URL) + + executor := NewAgentExecutor(nil, nil) + executor.SetConfig(cfg) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "both-query", + Type: core.StepTypeAgent, + Query: "single query", + Queries: []string{"query 1", "query 2"}, + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "cannot have both") + assert.Equal(t, core.StepStatusFailed, result.Status) +} + +func TestAgentExecutor_MultiGoal_SharedConversation(t *testing.T) { + // Verify subsequent goals see messages from previous goals + var receivedMessages [][]ChatMessage + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + receivedMessages = append(receivedMessages, req.Messages) + + resp := mockLLMResponse("response") + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "shared-conv", + Type: core.StepTypeAgent, + Queries: []string{"First query", "Second query"}, + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + + // Second goal should have more messages (accumulated from first goal) + require.Len(t, receivedMessages, 2) + assert.True(t, len(receivedMessages[1]) > len(receivedMessages[0]), + "second goal should see accumulated messages from first goal") +} + +// ============================================================================ +// Group 5: Model Fallback Tests +// ============================================================================ + +func TestAgentExecutor_ModelFallback(t *testing.T) { + var receivedModels []string + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + receivedModels = append(receivedModels, req.Model) + + // Fail for first model, succeed for second + if req.Model == "bad-model" { + w.WriteHeader(http.StatusInternalServerError) + json.NewEncoder(w).Encode(map[string]interface{}{ + "error": map[string]interface{}{ + "message": "model not found", + "type": "invalid_request_error", + }, + }) + return + } + + resp := mockLLMResponse("Fallback succeeded") + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "fallback-agent", + Type: core.StepTypeAgent, + Query: "Test fallback", + Models: []string{"bad-model", "test-model"}, + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, "Fallback succeeded", result.Output) + + // Verify both models were attempted + assert.Contains(t, receivedModels, "bad-model") + assert.Contains(t, receivedModels, "test-model") +} + +// ============================================================================ +// Group 6: Structured Output Tests +// ============================================================================ + +func TestAgentExecutor_StructuredOutput(t *testing.T) { + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + + if count == 1 { + // First call: agent completes with unstructured text + resp := mockLLMResponse("Found some results") + json.NewEncoder(w).Encode(resp) + } else { + // Second call: structured output request — verify response_format is set + assert.NotNil(t, req.ResponseFormat, "structured output call should have response_format") + resp := mockLLMResponse(`{"subdomains":[{"name":"sub.test.com","status":"active"}]}`) + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "structured-agent", + Type: core.StepTypeAgent, + Query: "Find subdomains", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + OutputSchema: `{"type":"object","properties":{"subdomains":{"type":"array"}}}`, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + // Final content should be the structured JSON + assert.Contains(t, result.Output, "subdomains") + assert.Contains(t, result.Output, "sub.test.com") +} + +func TestAgentExecutor_StructuredOutput_AlreadyJSON(t *testing.T) { + // If the agent already returns valid JSON, no extra call should be made + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + resp := mockLLMResponse(`{"results":["a","b"]}`) + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "already-json-agent", + Type: core.StepTypeAgent, + Query: "Return JSON", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + OutputSchema: `{"type":"object"}`, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + // Should only have 1 LLM call (no extra structured output request) + assert.Equal(t, 30, result.Exports["agent_total_tokens"]) +} + +// ============================================================================ +// Group 7: Tool Tracing Hooks Tests +// ============================================================================ + +func TestAgentExecutor_TracingHooks(t *testing.T) { + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + if count == 1 { + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_1", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "bash", + Arguments: `{"command": "echo traced"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + } else { + resp := mockLLMResponse("Traced complete") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + // Use hooks that call log_info (which is a registered function) + step := &core.Step{ + Name: "traced-agent", + Type: core.StepTypeAgent, + Query: "Run traced command", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + OnToolStart: `log_info("HOOK_START: " + tool_name)`, + OnToolEnd: `log_info("HOOK_END: " + tool_name + " dur=" + duration)`, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, "Traced complete", result.Output) +} + +func TestAgentExecutor_TracingHooks_EmptyHooks(t *testing.T) { + // Verify hooks are no-ops when empty strings + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + resp := mockLLMResponse("No hooks") + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "no-hooks-agent", + Type: core.StepTypeAgent, + Query: "Test without hooks", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + OnToolStart: "", + OnToolEnd: "", + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) +} + +// ============================================================================ +// Sub-Agent Orchestration Tests +// ============================================================================ + +func TestAgentExecutor_SubAgentSpawn(t *testing.T) { + // Mock server: + // 1st call (parent): returns spawn_agent tool call + // 2nd call (child): returns child final response + // 3rd call (parent): returns parent final response using child result + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + switch count { + case 1: + // Parent: spawn recon_agent + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_spawn", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "spawn_agent", + Arguments: `{"agent":"recon_agent","query":"scan ports on target"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + case 2: + // Child: return final result + resp := mockLLMResponse("Found open ports: 80, 443, 8080") + json.NewEncoder(w).Encode(resp) + case 3: + // Parent: final response using child result + resp := mockLLMResponse("Recon complete. Open ports: 80, 443, 8080") + json.NewEncoder(w).Encode(resp) + default: + resp := mockLLMResponse("unexpected call") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "orchestrator", + Type: core.StepTypeAgent, + Query: "Analyze target by coordinating specialists", + SystemPrompt: "You are an orchestrator. Delegate tasks to sub-agents.", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + SubAgents: []core.SubAgentDef{ + { + Name: "recon_agent", + Description: "Specialized agent for reconnaissance", + SystemPrompt: "You are a recon specialist", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + }, + }, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Contains(t, result.Output, "Recon complete") + // 3 LLM calls × 30 tokens each = 90 + assert.Equal(t, 90, result.Exports["agent_total_tokens"]) +} + +func TestAgentExecutor_SubAgentDepthLimit(t *testing.T) { + // Mock server: parent calls spawn_agent which tries to exceed depth + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + if count == 1 { + // Parent: spawn sub-agent + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_spawn", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "spawn_agent", + Arguments: `{"agent":"child","query":"do work"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + } else { + resp := mockLLMResponse("Done with depth error") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "depth-limit-parent", + Type: core.StepTypeAgent, + Query: "Delegate work", + MaxIterations: 5, + MaxAgentDepth: 1, // Only 1 level allowed, but parent is at 0 so child at 1 is OK + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + SubAgents: []core.SubAgentDef{ + { + Name: "child", + Description: "Child agent", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + SubAgents: []core.SubAgentDef{ + { + Name: "grandchild", + Description: "Grandchild agent", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + }, + }, + }, + }, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + // Parent should succeed even though the child depth limit was hit + // The error is returned as a tool result to the parent LLM + assert.Equal(t, core.StepStatusSuccess, result.Status) +} + +func TestAgentExecutor_SubAgentFailure(t *testing.T) { + // Mock server: child agent fails (e.g., config not available or bad response) + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + if count == 1 { + // Parent: spawn sub-agent + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_spawn", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "spawn_agent", + Arguments: `{"agent":"failing_agent","query":"do work"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + } else if count == 2 { + // Child: return server error + w.WriteHeader(http.StatusInternalServerError) + json.NewEncoder(w).Encode(map[string]interface{}{ + "error": map[string]interface{}{ + "message": "internal server error", + "type": "server_error", + }, + }) + } else { + // Parent: handle error gracefully + resp := mockLLMResponse("Sub-agent failed but I handled it") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "failure-parent", + Type: core.StepTypeAgent, + Query: "Delegate work", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + SubAgents: []core.SubAgentDef{ + { + Name: "failing_agent", + Description: "Agent that will fail", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + }, + }, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + // Parent should succeed — child failure is returned as tool result string + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Contains(t, result.Output, "handled") +} + +func TestAgentExecutor_SubAgentRecursive(t *testing.T) { + // Test 3-level nesting: parent → child → grandchild + // Mock server handles 5 calls: + // 1. Parent spawns child + // 2. Child spawns grandchild + // 3. Grandchild returns result + // 4. Child returns using grandchild result + // 5. Parent returns using child result + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + switch count { + case 1: + // Parent: spawn child + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_1", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "spawn_agent", + Arguments: `{"agent":"child","query":"intermediate task"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + case 2: + // Child: spawn grandchild + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: "call_2", + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "spawn_agent", + Arguments: `{"agent":"grandchild","query":"leaf task"}`, + }, + }, + }) + json.NewEncoder(w).Encode(resp) + case 3: + // Grandchild: return result + resp := mockLLMResponse("Grandchild result: data collected") + json.NewEncoder(w).Encode(resp) + case 4: + // Child: return using grandchild data + resp := mockLLMResponse("Child processed grandchild data") + json.NewEncoder(w).Encode(resp) + case 5: + // Parent: final response + resp := mockLLMResponse("Parent summarized: all levels complete") + json.NewEncoder(w).Encode(resp) + default: + resp := mockLLMResponse("unexpected") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "recursive-parent", + Type: core.StepTypeAgent, + Query: "Coordinate 3-level task", + MaxIterations: 5, + MaxAgentDepth: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + SubAgents: []core.SubAgentDef{ + { + Name: "child", + Description: "Intermediate agent", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + SubAgents: []core.SubAgentDef{ + { + Name: "grandchild", + Description: "Leaf agent", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + }, + }, + }, + }, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Contains(t, result.Output, "all levels complete") + // 5 LLM calls × 30 tokens = 150 + assert.Equal(t, 150, result.Exports["agent_total_tokens"]) +} + +func TestAgentState_MergeTokens(t *testing.T) { + state := &agentState{ + totalTokens: 100, + promptTokens: 60, + completionTokens: 40, + } + + state.MergeTokens(30, 20, 10) + assert.Equal(t, 130, state.totalTokens) + assert.Equal(t, 80, state.promptTokens) + assert.Equal(t, 50, state.completionTokens) +} + +// ============================================================================ +// Group 4: Conversation Compression Tests +// ============================================================================ + +func TestAgentExecutor_MessageWindowWithSummary(t *testing.T) { + // Mock server: returns tool calls for several iterations, then final answer. + // Also handles summarization request. + var callCount int32 + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + count := atomic.AddInt32(&callCount, 1) + w.Header().Set("Content-Type", "application/json") + + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + + // Check if this is a summarization request (no tools) + if len(req.Tools) == 0 && len(req.Messages) > 0 { + lastMsg := req.Messages[len(req.Messages)-1] + if content, ok := lastMsg.Content.(string); ok && strings.Contains(content, "Summarize") { + resp := mockLLMResponse("Summary of earlier context") + json.NewEncoder(w).Encode(resp) + return + } + } + + if count <= 3 { + // Return tool calls to build up conversation + resp := mockLLMToolCallResponse([]core.LLMToolCall{ + { + ID: fmt.Sprintf("call_%d", count), + Type: "function", + Function: core.LLMToolCallFunction{ + Name: "bash", + Arguments: fmt.Sprintf(`{"command": "echo iter_%d"}`, count), + }, + }, + }) + json.NewEncoder(w).Encode(resp) + } else { + resp := mockLLMResponse("Compression test complete") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "compress-agent", + Type: core.StepTypeAgent, + Query: "Run multiple commands", + MaxIterations: 10, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + Memory: &core.AgentMemoryConfig{ + MaxMessages: 5, + SummarizeOnTruncate: true, + }, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, "Compression test complete", result.Output) +} diff --git a/internal/executor/dispatcher.go b/internal/executor/dispatcher.go index 789f1e1..a30fdff 100644 --- a/internal/executor/dispatcher.go +++ b/internal/executor/dispatcher.go @@ -27,8 +27,9 @@ type StepDispatcher struct { runner runner.Runner enableBatch bool // Enable batch template rendering // Keep direct references to executors that need special configuration - bashExecutor *BashExecutor - llmExecutor *LLMExecutor + bashExecutor *BashExecutor + llmExecutor *LLMExecutor + agentExecutor *AgentExecutor } // SetDryRun enables or disables dry-run mode for the dispatcher @@ -39,6 +40,7 @@ func (d *StepDispatcher) SetDryRun(dryRun bool) { // SetSilent enables or disables silent mode for executors that support it func (d *StepDispatcher) SetSilent(silent bool) { d.llmExecutor.SetSilent(silent) + d.agentExecutor.SetSilent(silent) } // SetRunner sets the runner for command execution @@ -101,6 +103,7 @@ func NewStepDispatcherWithConfig(cfg StepDispatcherConfig) *StepDispatcher { // Create executors d.bashExecutor = NewBashExecutor(engine) d.llmExecutor = NewLLMExecutor(engine) + d.agentExecutor = NewAgentExecutor(engine, d.functionRegistry) // Register all built-in plugins d.registry.Register(d.bashExecutor) @@ -110,6 +113,7 @@ func NewStepDispatcherWithConfig(cfg StepDispatcherConfig) *StepDispatcher { d.registry.Register(NewRemoteBashExecutor(engine)) d.registry.Register(NewHTTPExecutor(engine)) d.registry.Register(d.llmExecutor) + d.registry.Register(d.agentExecutor) return d } @@ -122,6 +126,7 @@ func (d *StepDispatcher) RegisterPlugin(plugin StepExecutorPlugin) { // SetConfig passes config to executors that need it func (d *StepDispatcher) SetConfig(cfg *config.Config) { d.llmExecutor.SetConfig(cfg) + d.agentExecutor.SetConfig(cfg) } // Dispatch dispatches a step to the appropriate executor @@ -280,6 +285,39 @@ func collectRenderRequests(step *core.Step) []template.RenderRequest { add(fmt.Sprintf("Headers[%s]", k), v) } + // Agent step fields + add("Query", step.Query) + add("SystemPrompt", step.SystemPrompt) + add("StopCondition", step.StopCondition) + add("PlanPrompt", step.PlanPrompt) + add("OnToolStart", step.OnToolStart) + add("OnToolEnd", step.OnToolEnd) + for i, q := range step.Queries { + add(fmt.Sprintf("Queries[%d]", i), q) + } + if step.Memory != nil { + add("Memory.PersistPath", step.Memory.PersistPath) + add("Memory.ResumePath", step.Memory.ResumePath) + } + for i, tool := range step.AgentTools { + add(fmt.Sprintf("AgentTools[%d].Handler", i), tool.Handler) + } + + // Sub-agent fields (only top-level; nested sub-agents rendered on spawn) + for i, sa := range step.SubAgents { + add(fmt.Sprintf("SubAgents[%d].SystemPrompt", i), sa.SystemPrompt) + add(fmt.Sprintf("SubAgents[%d].StopCondition", i), sa.StopCondition) + add(fmt.Sprintf("SubAgents[%d].OnToolStart", i), sa.OnToolStart) + add(fmt.Sprintf("SubAgents[%d].OnToolEnd", i), sa.OnToolEnd) + if sa.Memory != nil { + add(fmt.Sprintf("SubAgents[%d].Memory.PersistPath", i), sa.Memory.PersistPath) + add(fmt.Sprintf("SubAgents[%d].Memory.ResumePath", i), sa.Memory.ResumePath) + } + for j, tool := range sa.AgentTools { + add(fmt.Sprintf("SubAgents[%d].AgentTools[%d].Handler", i, j), tool.Handler) + } + } + // RunnerConfig fields if step.StepRunnerConfig != nil && step.StepRunnerConfig.RunnerConfig != nil { cfg := step.StepRunnerConfig.RunnerConfig @@ -391,6 +429,90 @@ func (d *StepDispatcher) renderStepBatch(step *core.Step, vars map[string]any) ( if v := get("HostOutputFile"); v != "" { rendered.HostOutputFile = v } + // Agent step fields + if v := get("Query"); v != "" { + rendered.Query = v + } + if v := get("SystemPrompt"); v != "" { + rendered.SystemPrompt = v + } + if v := get("StopCondition"); v != "" { + rendered.StopCondition = v + } + if v := get("PlanPrompt"); v != "" { + rendered.PlanPrompt = v + } + if v := get("OnToolStart"); v != "" { + rendered.OnToolStart = v + } + if v := get("OnToolEnd"); v != "" { + rendered.OnToolEnd = v + } + // Render Queries slice + if len(step.Queries) > 0 { + rendered.Queries = make([]string, len(step.Queries)) + for i := range step.Queries { + rendered.Queries[i] = get(fmt.Sprintf("Queries[%d]", i)) + } + } + if step.Memory != nil { + mem := *step.Memory + if v := get("Memory.PersistPath"); v != "" { + mem.PersistPath = v + } + if v := get("Memory.ResumePath"); v != "" { + mem.ResumePath = v + } + rendered.Memory = &mem + } + // Apply results to agent tool handlers + if len(step.AgentTools) > 0 { + renderedTools := make([]core.AgentToolDef, len(step.AgentTools)) + copy(renderedTools, step.AgentTools) + for i := range renderedTools { + if v := get(fmt.Sprintf("AgentTools[%d].Handler", i)); v != "" { + renderedTools[i].Handler = v + } + } + rendered.AgentTools = renderedTools + } + // Apply results to sub-agent fields + if len(step.SubAgents) > 0 { + renderedSAs := make([]core.SubAgentDef, len(step.SubAgents)) + for i, sa := range step.SubAgents { + renderedSAs[i] = sa.DeepCopy() + if v := get(fmt.Sprintf("SubAgents[%d].SystemPrompt", i)); v != "" { + renderedSAs[i].SystemPrompt = v + } + if v := get(fmt.Sprintf("SubAgents[%d].StopCondition", i)); v != "" { + renderedSAs[i].StopCondition = v + } + if v := get(fmt.Sprintf("SubAgents[%d].OnToolStart", i)); v != "" { + renderedSAs[i].OnToolStart = v + } + if v := get(fmt.Sprintf("SubAgents[%d].OnToolEnd", i)); v != "" { + renderedSAs[i].OnToolEnd = v + } + if sa.Memory != nil { + if renderedSAs[i].Memory == nil { + mem := *sa.Memory + renderedSAs[i].Memory = &mem + } + if v := get(fmt.Sprintf("SubAgents[%d].Memory.PersistPath", i)); v != "" { + renderedSAs[i].Memory.PersistPath = v + } + if v := get(fmt.Sprintf("SubAgents[%d].Memory.ResumePath", i)); v != "" { + renderedSAs[i].Memory.ResumePath = v + } + } + for j := range sa.AgentTools { + if v := get(fmt.Sprintf("SubAgents[%d].AgentTools[%d].Handler", i, j)); v != "" { + renderedSAs[i].AgentTools[j].Handler = v + } + } + } + rendered.SubAgents = renderedSAs + } // Apply results to slice fields if len(step.Commands) > 0 { @@ -815,6 +937,157 @@ func (d *StepDispatcher) renderStepSequential(step *core.Step, vars map[string]a rendered.HostOutputFile = hostFile } + // Render agent step fields + if step.Query != "" { + q, err := d.templateEngine.Render(step.Query, vars) + if err != nil { + return nil, fmt.Errorf("error rendering query: %w", err) + } + rendered.Query = q + } + if len(step.Queries) > 0 { + qs, err := d.templateEngine.RenderSlice(step.Queries, vars) + if err != nil { + return nil, fmt.Errorf("error rendering queries: %w", err) + } + rendered.Queries = qs + } + if step.SystemPrompt != "" { + sp, err := d.templateEngine.Render(step.SystemPrompt, vars) + if err != nil { + return nil, fmt.Errorf("error rendering system_prompt: %w", err) + } + rendered.SystemPrompt = sp + } + if step.StopCondition != "" { + sc, err := d.templateEngine.Render(step.StopCondition, vars) + if err != nil { + return nil, fmt.Errorf("error rendering stop_condition: %w", err) + } + rendered.StopCondition = sc + } + if step.PlanPrompt != "" { + pp, err := d.templateEngine.Render(step.PlanPrompt, vars) + if err != nil { + return nil, fmt.Errorf("error rendering plan_prompt: %w", err) + } + rendered.PlanPrompt = pp + } + if step.OnToolStart != "" { + ots, err := d.templateEngine.Render(step.OnToolStart, vars) + if err != nil { + return nil, fmt.Errorf("error rendering on_tool_start: %w", err) + } + rendered.OnToolStart = ots + } + if step.OnToolEnd != "" { + ote, err := d.templateEngine.Render(step.OnToolEnd, vars) + if err != nil { + return nil, fmt.Errorf("error rendering on_tool_end: %w", err) + } + rendered.OnToolEnd = ote + } + if step.Memory != nil { + mem := *step.Memory + if mem.PersistPath != "" { + pp, err := d.templateEngine.Render(mem.PersistPath, vars) + if err != nil { + return nil, fmt.Errorf("error rendering memory.persist_path: %w", err) + } + mem.PersistPath = pp + } + if mem.ResumePath != "" { + rp, err := d.templateEngine.Render(mem.ResumePath, vars) + if err != nil { + return nil, fmt.Errorf("error rendering memory.resume_path: %w", err) + } + mem.ResumePath = rp + } + rendered.Memory = &mem + } + + // Render agent tool handlers + if len(step.AgentTools) > 0 { + renderedTools := make([]core.AgentToolDef, len(step.AgentTools)) + copy(renderedTools, step.AgentTools) + for i, tool := range renderedTools { + if tool.Handler != "" { + h, err := d.templateEngine.Render(tool.Handler, vars) + if err != nil { + return nil, fmt.Errorf("error rendering agent_tools[%d].handler: %w", i, err) + } + renderedTools[i].Handler = h + } + } + rendered.AgentTools = renderedTools + } + + // Render sub-agent fields (only top-level; nested sub-agents rendered on spawn) + if len(step.SubAgents) > 0 { + renderedSAs := make([]core.SubAgentDef, len(step.SubAgents)) + for i, sa := range step.SubAgents { + renderedSAs[i] = sa.DeepCopy() + if sa.SystemPrompt != "" { + sp, err := d.templateEngine.Render(sa.SystemPrompt, vars) + if err != nil { + return nil, fmt.Errorf("error rendering sub_agents[%d].system_prompt: %w", i, err) + } + renderedSAs[i].SystemPrompt = sp + } + if sa.StopCondition != "" { + sc, err := d.templateEngine.Render(sa.StopCondition, vars) + if err != nil { + return nil, fmt.Errorf("error rendering sub_agents[%d].stop_condition: %w", i, err) + } + renderedSAs[i].StopCondition = sc + } + if sa.OnToolStart != "" { + ots, err := d.templateEngine.Render(sa.OnToolStart, vars) + if err != nil { + return nil, fmt.Errorf("error rendering sub_agents[%d].on_tool_start: %w", i, err) + } + renderedSAs[i].OnToolStart = ots + } + if sa.OnToolEnd != "" { + ote, err := d.templateEngine.Render(sa.OnToolEnd, vars) + if err != nil { + return nil, fmt.Errorf("error rendering sub_agents[%d].on_tool_end: %w", i, err) + } + renderedSAs[i].OnToolEnd = ote + } + if sa.Memory != nil { + if renderedSAs[i].Memory == nil { + mem := *sa.Memory + renderedSAs[i].Memory = &mem + } + if sa.Memory.PersistPath != "" { + pp, err := d.templateEngine.Render(sa.Memory.PersistPath, vars) + if err != nil { + return nil, fmt.Errorf("error rendering sub_agents[%d].memory.persist_path: %w", i, err) + } + renderedSAs[i].Memory.PersistPath = pp + } + if sa.Memory.ResumePath != "" { + rp, err := d.templateEngine.Render(sa.Memory.ResumePath, vars) + if err != nil { + return nil, fmt.Errorf("error rendering sub_agents[%d].memory.resume_path: %w", i, err) + } + renderedSAs[i].Memory.ResumePath = rp + } + } + for j, tool := range sa.AgentTools { + if tool.Handler != "" { + h, err := d.templateEngine.Render(tool.Handler, vars) + if err != nil { + return nil, fmt.Errorf("error rendering sub_agents[%d].agent_tools[%d].handler: %w", i, j, err) + } + renderedSAs[i].AgentTools[j].Handler = h + } + } + } + rendered.SubAgents = renderedSAs + } + // Render LLM step fields if len(step.Messages) > 0 { if err := d.renderLLMMessages(&rendered, vars); err != nil { diff --git a/internal/executor/executor.go b/internal/executor/executor.go index a6eacd2..29137d6 100644 --- a/internal/executor/executor.go +++ b/internal/executor/executor.go @@ -2136,36 +2136,38 @@ func (e *Executor) executeStep(ctx context.Context, step *core.Step, execCtx *co fmt.Printf(" Pre-condition: %s %s\n", terminal.Gray(renderedCond), terminal.Yellow("(skipped in dry-run)")) } - if step.Command != "" { - // Render the command for display - rendered, _ := e.templateEngine.Render(step.Command, execCtx.GetVariables()) - fmt.Printf(" Would execute: %s\n", terminal.Gray(rendered)) - } - if len(step.Commands) > 0 { - fmt.Printf(" Would execute in %s:\n", step.Type) - for _, cmd := range step.Commands { - rendered, _ := e.templateEngine.Render(cmd, execCtx.GetVariables()) - fmt.Printf(" %s %s\n", terminal.SymbolBullet, terminal.Gray(rendered)) + if !step.SuppressDetails { + if step.Command != "" { + // Render the command for display + rendered, _ := e.templateEngine.Render(step.Command, execCtx.GetVariables()) + fmt.Printf(" Would execute: %s\n", terminal.Gray(rendered)) + } + if len(step.Commands) > 0 { + fmt.Printf(" Would execute in %s:\n", step.Type) + for _, cmd := range step.Commands { + rendered, _ := e.templateEngine.Render(cmd, execCtx.GetVariables()) + fmt.Printf(" %s %s\n", terminal.SymbolBullet, terminal.Gray(rendered)) + } } - } - // Display function(s) for function steps - if step.Function != "" { - rendered, _ := e.templateEngine.Render(step.Function, execCtx.GetVariables()) - fmt.Printf(" Would execute: %s\n", terminal.Gray(rendered)) - } - if len(step.Functions) > 0 { - fmt.Printf(" Would execute functions:\n") - for _, fn := range step.Functions { - rendered, _ := e.templateEngine.Render(fn, execCtx.GetVariables()) - fmt.Printf(" %s %s\n", terminal.SymbolBullet, terminal.Gray(rendered)) + // Display function(s) for function steps + if step.Function != "" { + rendered, _ := e.templateEngine.Render(step.Function, execCtx.GetVariables()) + fmt.Printf(" Would execute: %s\n", terminal.Gray(rendered)) } - } - if len(step.ParallelFunctions) > 0 { - fmt.Printf(" Would execute in parallel:\n") - for _, fn := range step.ParallelFunctions { - rendered, _ := e.templateEngine.Render(fn, execCtx.GetVariables()) - fmt.Printf(" %s %s\n", terminal.SymbolBullet, terminal.Gray(rendered)) + if len(step.Functions) > 0 { + fmt.Printf(" Would execute functions:\n") + for _, fn := range step.Functions { + rendered, _ := e.templateEngine.Render(fn, execCtx.GetVariables()) + fmt.Printf(" %s %s\n", terminal.SymbolBullet, terminal.Gray(rendered)) + } + } + if len(step.ParallelFunctions) > 0 { + fmt.Printf(" Would execute in parallel:\n") + for _, fn := range step.ParallelFunctions { + rendered, _ := e.templateEngine.Render(fn, execCtx.GetVariables()) + fmt.Printf(" %s %s\n", terminal.SymbolBullet, terminal.Gray(rendered)) + } } } @@ -2206,6 +2208,12 @@ func (e *Executor) executeStep(ctx context.Context, step *core.Step, execCtx *co stepCommandColored, _ = e.templateEngine.Render(stepCommandColored, execCtx.GetVariables()) } + // Suppress command details if requested + if step.SuppressDetails { + stepCommand = "" + stepCommandColored = "" + } + // Show step start (skip when progress bar is active) if e.progressBar == nil { e.printer.StepStartWithCommand(step.Name, stepSymbol, stepCommandColored, cmdPrefix) diff --git a/internal/executor/llm_executor.go b/internal/executor/llm_executor.go index 0ebf7fe..86a93f4 100644 --- a/internal/executor/llm_executor.go +++ b/internal/executor/llm_executor.go @@ -1,10 +1,11 @@ package executor import ( + "bufio" "bytes" "context" - "github.com/j3ssie/osmedeus/v5/internal/json" "fmt" + "github.com/j3ssie/osmedeus/v5/internal/json" "io" "net/http" "strings" @@ -134,6 +135,45 @@ type ChatError struct { Code string `json:"code"` } +// ChatCompletionStreamChunk is a single SSE event in a streaming response +type ChatCompletionStreamChunk struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []StreamChunkChoice `json:"choices"` + Usage *ChatUsage `json:"usage,omitempty"` + Error *ChatError `json:"error,omitempty"` +} + +// StreamChunkChoice is a single choice in a streaming chunk +type StreamChunkChoice struct { + Index int `json:"index"` + Delta ChatDelta `json:"delta"` + FinishReason string `json:"finish_reason,omitempty"` +} + +// ChatDelta contains incremental content from streaming +type ChatDelta struct { + Role string `json:"role,omitempty"` + Content string `json:"content,omitempty"` + ToolCalls []StreamToolCall `json:"tool_calls,omitempty"` +} + +// StreamToolCall is a partial tool call from streaming (index-based accumulation) +type StreamToolCall struct { + Index int `json:"index"` + ID string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + Function StreamToolCallFunction `json:"function,omitempty"` +} + +// StreamToolCallFunction contains partial function data from streaming +type StreamToolCallFunction struct { + Name string `json:"name,omitempty"` + Arguments string `json:"arguments,omitempty"` +} + // EmbeddingRequest represents a request for embeddings type EmbeddingRequest struct { Model string `json:"model"` @@ -316,8 +356,8 @@ func (e *LLMExecutor) executeChatCompletion( return result, err } - // Process response and exports - e.processChatResponse(result, step.Name, response) + // Process response and exports (skip print if streamed — output was already printed token-by-token) + e.processChatResponse(result, step.Name, response, llmConfig.Stream) result.Status = core.StepStatusSuccess result.EndTime = time.Now() @@ -496,13 +536,22 @@ func (e *LLMExecutor) buildChatRequest(step *core.Step, llmConfig *MergedLLMConf return request, nil } -// sendChatRequest sends an HTTP request to the LLM provider +// sendChatRequest sends an HTTP request to the LLM provider. +// When request.Stream is true, it delegates to sendChatRequestStreaming for SSE handling. func (e *LLMExecutor) sendChatRequest( ctx context.Context, provider *config.LLMProvider, request *ChatCompletionRequest, llmConfig *MergedLLMConfig, ) (*ChatCompletionResponse, error) { + if request.Stream { + return e.sendChatRequestStreaming(ctx, provider, request, llmConfig, func(token string) { + if !e.silent { + fmt.Print(token) + } + }) + } + // Marshal request to JSON body, err := json.Marshal(request) if err != nil { @@ -564,6 +613,215 @@ func (e *LLMExecutor) sendChatRequest( return &response, nil } +// sendChatRequestStreaming handles SSE streaming responses from the LLM provider. +// It reads the stream line-by-line, accumulates content and tool calls, and calls +// onToken for each text chunk for real-time display. +func (e *LLMExecutor) sendChatRequestStreaming( + ctx context.Context, + provider *config.LLMProvider, + request *ChatCompletionRequest, + llmConfig *MergedLLMConfig, + onToken func(token string), +) (*ChatCompletionResponse, error) { + log := logger.Get() + + // Marshal request to JSON + body, err := json.Marshal(request) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + // Create HTTP request + req, err := http.NewRequestWithContext(ctx, "POST", provider.BaseURL, bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + // Set headers + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + + if provider.AuthToken != "" { + req.Header.Set("Authorization", "Bearer "+provider.AuthToken) + } + + // Add custom headers + for key, value := range llmConfig.CustomHeaders { + req.Header.Set(key, value) + } + + // Set timeout + timeout, err := time.ParseDuration(llmConfig.Timeout) + if err != nil { + timeout = 120 * time.Second + } + client := &http.Client{Timeout: timeout} + + // Execute request + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + // Check for HTTP errors before reading the stream + if resp.StatusCode >= 400 { + respBody, _ := io.ReadAll(resp.Body) + var errResp ChatCompletionResponse + if json.Unmarshal(respBody, &errResp) == nil && errResp.Error != nil { + return &errResp, fmt.Errorf("HTTP %d: %s", resp.StatusCode, errResp.Error.Message) + } + return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(respBody)) + } + + // Parse SSE stream + var contentBuilder strings.Builder + var role string + var finishReason string + var responseID, responseModel string + var usage ChatUsage + // Accumulate tool calls by index + toolCallMap := make(map[int]*core.LLMToolCall) + + scanner := bufio.NewScanner(resp.Body) + // Increase buffer size for potentially large SSE lines + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + + for scanner.Scan() { + line := scanner.Text() + + // SSE format: lines starting with "data: " + if !strings.HasPrefix(line, "data: ") { + continue + } + + data := strings.TrimPrefix(line, "data: ") + + // Stream terminator + if data == "[DONE]" { + break + } + + var chunk ChatCompletionStreamChunk + if err := json.Unmarshal([]byte(data), &chunk); err != nil { + log.Debug("Failed to parse streaming chunk, skipping", + zap.String("data", data), + zap.Error(err), + ) + continue + } + + // Check for error in chunk + if chunk.Error != nil { + return &ChatCompletionResponse{Error: chunk.Error}, + fmt.Errorf("streaming error: %s", chunk.Error.Message) + } + + // Capture metadata from first chunk + if responseID == "" && chunk.ID != "" { + responseID = chunk.ID + } + if responseModel == "" && chunk.Model != "" { + responseModel = chunk.Model + } + + // Capture usage from final chunk + if chunk.Usage != nil { + usage = *chunk.Usage + } + + // Process choices + for _, choice := range chunk.Choices { + // Capture role from first delta + if choice.Delta.Role != "" { + role = choice.Delta.Role + } + + // Accumulate content + if choice.Delta.Content != "" { + contentBuilder.WriteString(choice.Delta.Content) + if onToken != nil { + onToken(choice.Delta.Content) + } + } + + // Accumulate tool calls (index-based) + for _, tc := range choice.Delta.ToolCalls { + existing, ok := toolCallMap[tc.Index] + if !ok { + existing = &core.LLMToolCall{ + Type: "function", + } + toolCallMap[tc.Index] = existing + } + if tc.ID != "" { + existing.ID = tc.ID + } + if tc.Type != "" { + existing.Type = tc.Type + } + if tc.Function.Name != "" { + existing.Function.Name = tc.Function.Name + } + if tc.Function.Arguments != "" { + existing.Function.Arguments += tc.Function.Arguments + } + } + + // Capture finish reason + if choice.FinishReason != "" { + finishReason = choice.FinishReason + } + } + } + + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("error reading stream: %w", err) + } + + // Print trailing newline after streaming output + if onToken != nil && contentBuilder.Len() > 0 { + onToken("\n") + } + + // Build tool calls slice from map (ordered by index) + var toolCalls []core.LLMToolCall + if len(toolCallMap) > 0 { + toolCalls = make([]core.LLMToolCall, len(toolCallMap)) + for idx, tc := range toolCallMap { + if idx < len(toolCalls) { + toolCalls[idx] = *tc + } + } + } + + // Assemble final response + if role == "" { + role = "assistant" + } + + var msgContent interface{} = contentBuilder.String() + + response := &ChatCompletionResponse{ + ID: responseID, + Model: responseModel, + Choices: []ChatChoice{ + { + Index: 0, + Message: ChatMessage{ + Role: role, + Content: msgContent, + ToolCalls: toolCalls, + }, + FinishReason: finishReason, + }, + }, + Usage: usage, + } + + return response, nil +} + // sendEmbeddingRequest sends an embedding request to the LLM provider func (e *LLMExecutor) sendEmbeddingRequest( ctx context.Context, @@ -661,7 +919,7 @@ func printLLMOutput(content string) { } // processChatResponse exports the LLM response to step result -func (e *LLMExecutor) processChatResponse(result *core.StepResult, stepName string, response *ChatCompletionResponse) { +func (e *LLMExecutor) processChatResponse(result *core.StepResult, stepName string, response *ChatCompletionResponse, streamed bool) { exportKey := sanitizeStepName(stepName) + "_llm_resp" // Build comprehensive export structure @@ -691,8 +949,8 @@ func (e *LLMExecutor) processChatResponse(result *core.StepResult, stepName stri // Set output to content for display if content, ok := choice.Message.Content.(string); ok { result.Output = content - // Print LLM output with symbol prefix and markdown formatting (skip in silent mode) - if !e.silent { + // Print LLM output with markdown formatting (skip if silent or already streamed) + if !e.silent && !streamed { printLLMOutput(content) } } @@ -837,6 +1095,11 @@ func (e *LLMExecutor) getMergedConfig(step *core.Step) *MergedLLMConfig { } } + // Top-level step.Stream overrides everything (highest precedence) + if step.Stream != nil { + merged.Stream = *step.Stream + } + return merged } diff --git a/internal/executor/llm_streaming_test.go b/internal/executor/llm_streaming_test.go new file mode 100644 index 0000000..a707798 --- /dev/null +++ b/internal/executor/llm_streaming_test.go @@ -0,0 +1,988 @@ +package executor + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/j3ssie/osmedeus/v5/internal/config" + "github.com/j3ssie/osmedeus/v5/internal/core" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeSSEChunk writes a single SSE data line +func writeSSEChunk(w http.ResponseWriter, data string) { + fmt.Fprintf(w, "data: %s\n\n", data) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } +} + +// makeStreamChunk creates a JSON string for a streaming chunk with content +func makeStreamChunk(id, model, content string) string { + chunk := ChatCompletionStreamChunk{ + ID: id, + Model: model, + Choices: []StreamChunkChoice{ + { + Index: 0, + Delta: ChatDelta{ + Content: content, + }, + }, + }, + } + b, _ := json.Marshal(chunk) + return string(b) +} + +// makeStreamChunkWithRole creates a streaming chunk with role +func makeStreamChunkWithRole(id, model, role string) string { + chunk := ChatCompletionStreamChunk{ + ID: id, + Model: model, + Choices: []StreamChunkChoice{ + { + Index: 0, + Delta: ChatDelta{ + Role: role, + }, + }, + }, + } + b, _ := json.Marshal(chunk) + return string(b) +} + +// makeStreamChunkDone creates a final chunk with finish_reason and optional usage +func makeStreamChunkDone(id, model string, usage *ChatUsage) string { + chunk := ChatCompletionStreamChunk{ + ID: id, + Model: model, + Choices: []StreamChunkChoice{ + { + Index: 0, + Delta: ChatDelta{}, + FinishReason: "stop", + }, + }, + Usage: usage, + } + b, _ := json.Marshal(chunk) + return string(b) +} + +// makeStreamToolCallChunk creates a streaming chunk with a tool call delta +func makeStreamToolCallChunk(id, model string, tcIndex int, tcID, tcType, funcName, funcArgs string) string { + tc := StreamToolCall{ + Index: tcIndex, + } + if tcID != "" { + tc.ID = tcID + } + if tcType != "" { + tc.Type = tcType + } + if funcName != "" { + tc.Function.Name = funcName + } + if funcArgs != "" { + tc.Function.Arguments = funcArgs + } + + chunk := ChatCompletionStreamChunk{ + ID: id, + Model: model, + Choices: []StreamChunkChoice{ + { + Index: 0, + Delta: ChatDelta{ + ToolCalls: []StreamToolCall{tc}, + }, + }, + }, + } + b, _ := json.Marshal(chunk) + return string(b) +} + +func newStreamingMockServer(handler http.HandlerFunc) *httptest.Server { + return httptest.NewServer(handler) +} + +// ============================================================================ +// Streaming Tests +// ============================================================================ + +func TestStreamingSSEParsing(t *testing.T) { + // Mock server returns SSE events with content tokens + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + + writeSSEChunk(w, makeStreamChunkWithRole("chat-1", "test-model", "assistant")) + writeSSEChunk(w, makeStreamChunk("chat-1", "test-model", "Hello")) + writeSSEChunk(w, makeStreamChunk("chat-1", "test-model", " world")) + writeSSEChunk(w, makeStreamChunk("chat-1", "test-model", "!")) + writeSSEChunk(w, makeStreamChunkDone("chat-1", "test-model", &ChatUsage{ + PromptTokens: 10, + CompletionTokens: 5, + TotalTokens: 15, + })) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{ + config: &config.Config{}, + silent: true, + } + + ctx := context.Background() + provider := &config.LLMProvider{ + BaseURL: server.URL, + AuthToken: "test-token", + Model: "test-model", + } + request := &ChatCompletionRequest{ + Model: "test-model", + Messages: []ChatMessage{{Role: "user", Content: "Hi"}}, + Stream: true, + } + llmConfig := &MergedLLMConfig{Timeout: "30s"} + + var tokens []string + response, err := executor.sendChatRequestStreaming(ctx, provider, request, llmConfig, func(token string) { + tokens = append(tokens, token) + }) + + require.NoError(t, err) + require.NotNil(t, response) + + // Verify accumulated content + require.Len(t, response.Choices, 1) + content, ok := response.Choices[0].Message.Content.(string) + require.True(t, ok) + assert.Equal(t, "Hello world!", content) + + // Verify metadata + assert.Equal(t, "chat-1", response.ID) + assert.Equal(t, "test-model", response.Model) + assert.Equal(t, "assistant", response.Choices[0].Message.Role) + assert.Equal(t, "stop", response.Choices[0].FinishReason) + + // Verify usage + assert.Equal(t, 10, response.Usage.PromptTokens) + assert.Equal(t, 5, response.Usage.CompletionTokens) + assert.Equal(t, 15, response.Usage.TotalTokens) +} + +func TestStreamingTokenCallback(t *testing.T) { + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + writeSSEChunk(w, makeStreamChunk("c1", "m1", "token1")) + writeSSEChunk(w, makeStreamChunk("c1", "m1", "token2")) + writeSSEChunk(w, makeStreamChunk("c1", "m1", "token3")) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", nil)) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + var tokens []string + response, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + func(token string) { tokens = append(tokens, token) }, + ) + + require.NoError(t, err) + require.NotNil(t, response) + + // Verify callback was called for each content delta + trailing newline + assert.Equal(t, []string{"token1", "token2", "token3", "\n"}, tokens) +} + +func TestStreamingToolCallAccumulation(t *testing.T) { + // Simulate streaming tool calls: name arrives in one chunk, arguments split across chunks + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + // First tool call: ID and name + writeSSEChunk(w, makeStreamToolCallChunk("c1", "m1", 0, "call_1", "function", "bash", "")) + // Arguments arrive in parts + writeSSEChunk(w, makeStreamToolCallChunk("c1", "m1", 0, "", "", "", `{"comma`)) + writeSSEChunk(w, makeStreamToolCallChunk("c1", "m1", 0, "", "", "", `nd": "echo hi"}`)) + + // Second tool call in same response + writeSSEChunk(w, makeStreamToolCallChunk("c1", "m1", 1, "call_2", "function", "read_file", "")) + writeSSEChunk(w, makeStreamToolCallChunk("c1", "m1", 1, "", "", "", `{"path": "/tmp/test"}`)) + + // Finish + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", nil)) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + response, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + require.NoError(t, err) + require.NotNil(t, response) + require.Len(t, response.Choices, 1) + + toolCalls := response.Choices[0].Message.ToolCalls + require.Len(t, toolCalls, 2) + + // First tool call + assert.Equal(t, "call_1", toolCalls[0].ID) + assert.Equal(t, "function", toolCalls[0].Type) + assert.Equal(t, "bash", toolCalls[0].Function.Name) + assert.Equal(t, `{"command": "echo hi"}`, toolCalls[0].Function.Arguments) + + // Second tool call + assert.Equal(t, "call_2", toolCalls[1].ID) + assert.Equal(t, "read_file", toolCalls[1].Function.Name) + assert.Equal(t, `{"path": "/tmp/test"}`, toolCalls[1].Function.Arguments) +} + +func TestStreamingDoneSignal(t *testing.T) { + // Verify data: [DONE] terminates the stream cleanly + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + writeSSEChunk(w, makeStreamChunk("c1", "m1", "before done")) + writeSSEChunk(w, "[DONE]") + // Anything after [DONE] should be ignored + writeSSEChunk(w, makeStreamChunk("c1", "m1", "SHOULD NOT APPEAR")) + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + response, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + require.NoError(t, err) + require.NotNil(t, response) + + content, ok := response.Choices[0].Message.Content.(string) + require.True(t, ok) + assert.Equal(t, "before done", content) + assert.NotContains(t, content, "SHOULD NOT APPEAR") +} + +func TestStreamingFallbackOnError(t *testing.T) { + // Server returns HTTP 500 — should error gracefully + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + json.NewEncoder(w).Encode(map[string]interface{}{ + "error": map[string]interface{}{ + "message": "server overloaded", + "type": "server_error", + "code": "500", + }, + }) + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + _, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + require.Error(t, err) + assert.Contains(t, err.Error(), "500") + assert.Contains(t, err.Error(), "server overloaded") +} + +func TestStreamingErrorInChunk(t *testing.T) { + // Server streams normally then sends an error chunk + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + writeSSEChunk(w, makeStreamChunk("c1", "m1", "partial ")) + + // Error chunk + errorChunk := ChatCompletionStreamChunk{ + Error: &ChatError{ + Message: "context length exceeded", + Type: "invalid_request_error", + Code: "context_length_exceeded", + }, + } + b, _ := json.Marshal(errorChunk) + writeSSEChunk(w, string(b)) + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + _, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + require.Error(t, err) + assert.Contains(t, err.Error(), "context length exceeded") +} + +func TestStreamingMalformedSSE(t *testing.T) { + // Server sends some malformed data — should skip and continue + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + writeSSEChunk(w, makeStreamChunk("c1", "m1", "good")) + writeSSEChunk(w, "{invalid json") + writeSSEChunk(w, makeStreamChunk("c1", "m1", " data")) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", nil)) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + response, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + require.NoError(t, err) + require.NotNil(t, response) + + content, ok := response.Choices[0].Message.Content.(string) + require.True(t, ok) + assert.Equal(t, "good data", content) +} + +func TestAgentExecutor_StreamFlagInRequest(t *testing.T) { + // Verify that Stream is passed through to the LLM request when llmConfig.Stream is true + var receivedStream bool + + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + receivedStream = req.Stream + + // Respond with non-streaming format since the mock doesn't do SSE + // (the test captures what was sent, not how it responds) + w.Header().Set("Content-Type", "text/event-stream") + writeSSEChunk(w, makeStreamChunkWithRole("c1", "m1", "assistant")) + writeSSEChunk(w, makeStreamChunk("c1", "m1", "streamed response")) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", &ChatUsage{ + PromptTokens: 5, CompletionTokens: 3, TotalTokens: 8, + })) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + cfg.LLM.Stream = true + + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "stream-agent", + Type: core.StepTypeAgent, + Query: "Test streaming", + MaxIterations: 3, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.True(t, receivedStream, "Stream should be true in the LLM request") + assert.Equal(t, "streamed response", result.Output) +} + +func TestStepLevelStreamOverride(t *testing.T) { + // Verify that step.Stream overrides global config + t.Run("step stream true overrides global false", func(t *testing.T) { + cfg := &config.Config{ + LLM: config.LLMConfig{ + Stream: false, + MaxTokens: 100, + MaxRetries: 1, + Timeout: "30s", + }, + } + executor := &LLMExecutor{config: cfg} + streamTrue := true + step := &core.Step{ + Name: "test", + Stream: &streamTrue, + } + + merged := executor.getMergedConfig(step) + assert.True(t, merged.Stream, "step.Stream=true should override global false") + }) + + t.Run("step stream false overrides global true", func(t *testing.T) { + cfg := &config.Config{ + LLM: config.LLMConfig{ + Stream: true, + MaxTokens: 100, + MaxRetries: 1, + Timeout: "30s", + }, + } + executor := &LLMExecutor{config: cfg} + streamFalse := false + step := &core.Step{ + Name: "test", + Stream: &streamFalse, + } + + merged := executor.getMergedConfig(step) + assert.False(t, merged.Stream, "step.Stream=false should override global true") + }) + + t.Run("step stream nil inherits global", func(t *testing.T) { + cfg := &config.Config{ + LLM: config.LLMConfig{ + Stream: true, + MaxTokens: 100, + MaxRetries: 1, + Timeout: "30s", + }, + } + executor := &LLMExecutor{config: cfg} + step := &core.Step{ + Name: "test", + } + + merged := executor.getMergedConfig(step) + assert.True(t, merged.Stream, "nil step.Stream should inherit global true") + }) + + t.Run("step stream overrides llm_config stream", func(t *testing.T) { + cfg := &config.Config{ + LLM: config.LLMConfig{ + Stream: false, + MaxTokens: 100, + MaxRetries: 1, + Timeout: "30s", + }, + } + executor := &LLMExecutor{config: cfg} + llmConfigStream := false + stepStream := true + step := &core.Step{ + Name: "test", + LLMConfig: &core.LLMStepConfig{ + Stream: &llmConfigStream, + }, + Stream: &stepStream, + } + + merged := executor.getMergedConfig(step) + assert.True(t, merged.Stream, "step.Stream should override llm_config.Stream") + }) +} + +func TestStreamingDispatchFromSendChatRequest(t *testing.T) { + // Verify that sendChatRequest dispatches to streaming when request.Stream is true + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + writeSSEChunk(w, makeStreamChunkWithRole("c1", "m1", "assistant")) + writeSSEChunk(w, makeStreamChunk("c1", "m1", "streamed via dispatch")) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", &ChatUsage{ + PromptTokens: 5, CompletionTokens: 3, TotalTokens: 8, + })) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{ + config: &config.Config{}, + silent: true, + } + + ctx := context.Background() + provider := &config.LLMProvider{BaseURL: server.URL, Model: "m1"} + request := &ChatCompletionRequest{ + Model: "m1", + Messages: []ChatMessage{{Role: "user", Content: "test"}}, + Stream: true, + } + llmConfig := &MergedLLMConfig{Timeout: "30s", Stream: true} + + response, err := executor.sendChatRequest(ctx, provider, request, llmConfig) + require.NoError(t, err) + require.NotNil(t, response) + + content, ok := response.Choices[0].Message.Content.(string) + require.True(t, ok) + assert.Equal(t, "streamed via dispatch", content) + assert.Equal(t, 8, response.Usage.TotalTokens) +} + +func TestStreamingEmptyContent(t *testing.T) { + // Stream that has no content tokens (e.g., only tool calls or empty response) + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + writeSSEChunk(w, makeStreamChunkWithRole("c1", "m1", "assistant")) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", nil)) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + var tokens []string + response, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + func(token string) { tokens = append(tokens, token) }, + ) + + require.NoError(t, err) + require.NotNil(t, response) + + content, ok := response.Choices[0].Message.Content.(string) + require.True(t, ok) + assert.Equal(t, "", content) + + // No content tokens, so no callback should have been called (no trailing newline either) + assert.Empty(t, tokens) +} + +func TestStreamingSSEWithComments(t *testing.T) { + // SSE spec allows comment lines starting with ":" — these should be ignored + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + // SSE comment (keep-alive) + fmt.Fprintf(w, ": this is a comment\n\n") + writeSSEChunk(w, makeStreamChunk("c1", "m1", "with comments")) + fmt.Fprintf(w, ": another comment\n\n") + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", nil)) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + response, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + require.NoError(t, err) + content, ok := response.Choices[0].Message.Content.(string) + require.True(t, ok) + assert.Equal(t, "with comments", content) +} + +func TestStreamingLLMStepIntegration(t *testing.T) { + // Test the full LLM step path with streaming via sendChatRequest dispatch + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + + if req.Stream { + w.Header().Set("Content-Type", "text/event-stream") + writeSSEChunk(w, makeStreamChunkWithRole("c1", "m1", "assistant")) + writeSSEChunk(w, makeStreamChunk("c1", "m1", "Streaming ")) + writeSSEChunk(w, makeStreamChunk("c1", "m1", "response")) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", &ChatUsage{ + PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15, + })) + writeSSEChunk(w, "[DONE]") + } else { + // Fallback non-streaming response + w.Header().Set("Content-Type", "application/json") + resp := mockLLMResponse("Non-streaming response") + json.NewEncoder(w).Encode(resp) + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + cfg.LLM.Stream = true + + dispatcher := NewStepDispatcher() + llmExec := NewLLMExecutor(dispatcher.GetTemplateEngine()) + llmExec.SetConfig(cfg) + llmExec.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "streaming-llm", + Type: core.StepTypeLLM, + Messages: []core.LLMMessage{ + {Role: core.LLMRoleUser, Content: "Hello"}, + }, + } + + result, err := llmExec.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, "Streaming response", result.Output) +} + +func TestStreamingAuthHeader(t *testing.T) { + // Verify auth headers are sent in streaming requests + var receivedAuth string + + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + receivedAuth = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "text/event-stream") + writeSSEChunk(w, makeStreamChunk("c1", "m1", "ok")) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", nil)) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + _, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, AuthToken: "secret-key", Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + require.NoError(t, err) + assert.Equal(t, "Bearer secret-key", receivedAuth) +} + +func TestStreamingCustomHeaders(t *testing.T) { + // Verify custom headers are sent in streaming requests + var receivedHeaders http.Header + + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + receivedHeaders = r.Header + w.Header().Set("Content-Type", "text/event-stream") + writeSSEChunk(w, makeStreamChunk("c1", "m1", "ok")) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", nil)) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + _, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{ + Timeout: "30s", + CustomHeaders: map[string]string{ + "X-Custom": "value", + }, + }, + nil, + ) + + require.NoError(t, err) + assert.Equal(t, "value", receivedHeaders.Get("X-Custom")) + assert.Equal(t, "text/event-stream", receivedHeaders.Get("Accept")) +} + +func TestStreamingNonStreamingFallback(t *testing.T) { + // When Stream is false, sendChatRequest should NOT use streaming + server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + + assert.False(t, req.Stream, "stream should be false in request") + + w.Header().Set("Content-Type", "application/json") + resp := mockLLMResponse("non-streaming") + json.NewEncoder(w).Encode(resp) + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + response, err := executor.sendChatRequest( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: false}, + &MergedLLMConfig{Timeout: "30s"}, + ) + + require.NoError(t, err) + require.NotNil(t, response) + + content, ok := response.Choices[0].Message.Content.(string) + require.True(t, ok) + assert.Equal(t, "non-streaming", content) +} + +func TestStreamingContentWithSpecialChars(t *testing.T) { + // Test streaming with content containing special characters, newlines, etc. + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + writeSSEChunk(w, makeStreamChunk("c1", "m1", "Hello\n")) + writeSSEChunk(w, makeStreamChunk("c1", "m1", "- bullet point\n")) + writeSSEChunk(w, makeStreamChunk("c1", "m1", `{"json": "value"}`)) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", nil)) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + response, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + require.NoError(t, err) + content, ok := response.Choices[0].Message.Content.(string) + require.True(t, ok) + expected := "Hello\n- bullet point\n" + `{"json": "value"}` + assert.Equal(t, expected, content) +} + +func TestStreamingAcceptHeader(t *testing.T) { + // Verify that streaming requests include Accept: text/event-stream + var receivedAccept string + + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + receivedAccept = r.Header.Get("Accept") + w.Header().Set("Content-Type", "text/event-stream") + writeSSEChunk(w, makeStreamChunk("c1", "m1", "ok")) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", nil)) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + _, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + require.NoError(t, err) + assert.Equal(t, "text/event-stream", receivedAccept) +} + +func TestStreamingEmptyDataLines(t *testing.T) { + // SSE with blank lines and event: fields (should be ignored, only data: parsed) + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + // Various non-data lines + fmt.Fprintf(w, "event: message\n") + fmt.Fprintf(w, "id: 1\n") + writeSSEChunk(w, makeStreamChunk("c1", "m1", "content")) + fmt.Fprintf(w, "\n") // blank line + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", nil)) + writeSSEChunk(w, "[DONE]") + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + response, err := executor.sendChatRequestStreaming( + context.Background(), + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + require.NoError(t, err) + content, ok := response.Choices[0].Message.Content.(string) + require.True(t, ok) + assert.Equal(t, "content", content) +} + +func TestStreamingLLMExecProcess_SkipsPrint(t *testing.T) { + // Verify that processChatResponse with streamed=true does not call printLLMOutput + // (We can't easily assert fmt.Print wasn't called, but we verify the Output is set correctly) + + executor := &LLMExecutor{config: &config.Config{}, silent: false} + result := &core.StepResult{ + Exports: make(map[string]interface{}), + } + + response := &ChatCompletionResponse{ + ID: "test", + Model: "m1", + Choices: []ChatChoice{ + { + Message: ChatMessage{ + Role: "assistant", + Content: "test content", + }, + FinishReason: "stop", + }, + }, + } + + // With streamed=true, should not print (we just verify it doesn't panic) + executor.processChatResponse(result, "test-step", response, true) + assert.Equal(t, "test content", result.Output) + + // Exports should still be populated + _, ok := result.Exports["test_step_llm_resp"] + assert.True(t, ok, "exports should be populated even when streamed") + _, ok = result.Exports["test_step_content"] + assert.True(t, ok, "content export should be populated") +} + +func TestStreamingAgentWithToolCalls(t *testing.T) { + // Test agent executor with streaming enabled, doing tool call + final response + var callCount int32 + + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var req ChatCompletionRequest + json.Unmarshal(body, &req) + + callCount++ + w.Header().Set("Content-Type", "text/event-stream") + + if callCount == 1 { + // First call: stream tool call + writeSSEChunk(w, makeStreamChunkWithRole("c1", "m1", "assistant")) + writeSSEChunk(w, makeStreamToolCallChunk("c1", "m1", 0, "call_1", "function", "bash", `{"command": "echo test"}`)) + finishChunk := ChatCompletionStreamChunk{ + ID: "c1", + Model: "m1", + Choices: []StreamChunkChoice{ + {Index: 0, Delta: ChatDelta{}, FinishReason: "tool_calls"}, + }, + Usage: &ChatUsage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15}, + } + b, _ := json.Marshal(finishChunk) + writeSSEChunk(w, string(b)) + writeSSEChunk(w, "[DONE]") + } else { + // Second call: stream final response + writeSSEChunk(w, makeStreamChunkWithRole("c1", "m1", "assistant")) + writeSSEChunk(w, makeStreamChunk("c1", "m1", "Tool executed ")) + writeSSEChunk(w, makeStreamChunk("c1", "m1", "successfully")) + writeSSEChunk(w, makeStreamChunkDone("c1", "m1", &ChatUsage{ + PromptTokens: 15, CompletionTokens: 5, TotalTokens: 20, + })) + writeSSEChunk(w, "[DONE]") + } + }) + defer server.Close() + + cfg := newMockConfig(t, server.URL) + cfg.LLM.Stream = true + + dispatcher := NewStepDispatcher() + executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry()) + executor.SetConfig(cfg) + executor.SetSilent(true) + + ctx := context.Background() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + step := &core.Step{ + Name: "streaming-tool-agent", + Type: core.StepTypeAgent, + Query: "Run echo test", + MaxIterations: 5, + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + + result, err := executor.Execute(ctx, step, execCtx) + require.NoError(t, err) + assert.Equal(t, core.StepStatusSuccess, result.Status) + assert.Equal(t, "Tool executed successfully", result.Output) + assert.Equal(t, 2, result.Exports["agent_iterations"]) + + // Verify tokens accumulated from both calls + totalTokens, ok := result.Exports["agent_total_tokens"].(int) + require.True(t, ok) + assert.Equal(t, 35, totalTokens) // 15 + 20 +} + +func TestStreamingRespectsContext(t *testing.T) { + // Verify that context cancellation is respected during streaming + var _ = strings.NewReader // ensure strings import is used + + server := newStreamingMockServer(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + // Send one chunk then block (simulating slow stream) + writeSSEChunk(w, makeStreamChunk("c1", "m1", "partial")) + // The server will naturally end when the connection is closed + }) + defer server.Close() + + executor := &LLMExecutor{config: &config.Config{}, silent: true} + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // Cancel immediately + + _, err := executor.sendChatRequestStreaming( + ctx, + &config.LLMProvider{BaseURL: server.URL, Model: "m1"}, + &ChatCompletionRequest{Model: "m1", Messages: []ChatMessage{{Role: "user", Content: "test"}}, Stream: true}, + &MergedLLMConfig{Timeout: "30s"}, + nil, + ) + + // Should error due to cancelled context + assert.Error(t, err) +} diff --git a/internal/executor/tool_executor.go b/internal/executor/tool_executor.go new file mode 100644 index 0000000..f89d77e --- /dev/null +++ b/internal/executor/tool_executor.go @@ -0,0 +1,306 @@ +package executor + +import ( + "context" + "fmt" + + "github.com/j3ssie/osmedeus/v5/internal/config" + "github.com/j3ssie/osmedeus/v5/internal/core" + "github.com/j3ssie/osmedeus/v5/internal/functions" + "github.com/j3ssie/osmedeus/v5/internal/json" + "github.com/j3ssie/osmedeus/v5/internal/logger" + "github.com/j3ssie/osmedeus/v5/internal/template" + "go.uber.org/zap" +) + +// ToolExecutor defines the interface for executing agent tool calls. +// Each implementation handles a specific tool (preset or custom). +type ToolExecutor interface { + // Name returns the tool name (matches the function name in LLM tool calls) + Name() string + // Execute runs the tool with the given arguments and returns the result string + Execute(ctx context.Context, args map[string]interface{}, execCtx *core.ExecutionContext) (string, error) +} + +// ToolExecutorRegistry manages a collection of ToolExecutor instances +type ToolExecutorRegistry struct { + executors map[string]ToolExecutor +} + +// NewToolExecutorRegistry creates an empty registry +func NewToolExecutorRegistry() *ToolExecutorRegistry { + return &ToolExecutorRegistry{ + executors: make(map[string]ToolExecutor), + } +} + +// Register adds a ToolExecutor to the registry +func (r *ToolExecutorRegistry) Register(te ToolExecutor) { + r.executors[te.Name()] = te +} + +// Get returns the ToolExecutor for the given name +func (r *ToolExecutorRegistry) Get(name string) (ToolExecutor, bool) { + te, ok := r.executors[name] + return te, ok +} + +// Execute dispatches a tool call to the appropriate executor +func (r *ToolExecutorRegistry) Execute(ctx context.Context, name string, args map[string]interface{}, execCtx *core.ExecutionContext) (string, error) { + te, ok := r.executors[name] + if !ok { + return "", fmt.Errorf("unknown tool: %s", name) + } + return te.Execute(ctx, args, execCtx) +} + +// PresetToolExecutor wraps a preset tool and executes it via the function registry +type PresetToolExecutor struct { + name string + registry *functions.Registry +} + +// NewPresetToolExecutor creates a preset tool executor +func NewPresetToolExecutor(name string, registry *functions.Registry) *PresetToolExecutor { + return &PresetToolExecutor{name: name, registry: registry} +} + +// Name returns the tool name +func (e *PresetToolExecutor) Name() string { + return e.name +} + +// Execute runs the preset tool by building a function call expression +func (e *PresetToolExecutor) Execute(_ context.Context, args map[string]interface{}, execCtx *core.ExecutionContext) (string, error) { + expr := buildPresetCallExpr(e.name, args) + vars := execCtx.GetVariables() + result, err := e.registry.Execute(expr, vars) + if err != nil { + return "", fmt.Errorf("preset tool '%s' failed: %w", e.name, err) + } + return formatToolResult(result), nil +} + +// CustomToolExecutor wraps a custom tool with a JS handler expression +type CustomToolExecutor struct { + name string + handler string + registry *functions.Registry +} + +// NewCustomToolExecutor creates a custom tool executor +func NewCustomToolExecutor(name, handler string, registry *functions.Registry) *CustomToolExecutor { + return &CustomToolExecutor{name: name, handler: handler, registry: registry} +} + +// Name returns the tool name +func (e *CustomToolExecutor) Name() string { + return e.name +} + +// Execute runs the custom tool handler JS expression +func (e *CustomToolExecutor) Execute(_ context.Context, args map[string]interface{}, execCtx *core.ExecutionContext) (string, error) { + vars := execCtx.GetVariables() + vars["args"] = args + result, err := e.registry.Execute(e.handler, vars) + if err != nil { + return "", fmt.Errorf("custom handler '%s' failed: %w", e.name, err) + } + return formatToolResult(result), nil +} + +// BuildToolRegistry constructs a ToolExecutorRegistry from the resolved agent tool definitions +func BuildToolRegistry(toolDefs []core.AgentToolDef, funcRegistry *functions.Registry) *ToolExecutorRegistry { + log := logger.Get() + reg := NewToolExecutorRegistry() + + for _, def := range toolDefs { + if def.IsPreset() { + reg.Register(NewPresetToolExecutor(def.Preset, funcRegistry)) + } else if def.Handler != "" { + reg.Register(NewCustomToolExecutor(def.Name, def.Handler, funcRegistry)) + } else { + log.Debug("Tool definition has no handler, registering as preset fallback", + zap.String("tool", def.Name), + ) + reg.Register(NewPresetToolExecutor(def.Name, funcRegistry)) + } + } + + return reg +} + +// SubAgentToolExecutor handles spawn_agent tool calls by creating +// a child AgentExecutor and running the specified sub-agent. +type SubAgentToolExecutor struct { + subAgents map[string]core.SubAgentDef + templateEngine template.TemplateEngine + funcRegistry *functions.Registry + config *config.Config + silent bool + currentDepth int + maxDepth int + parentState *agentState // for token merging +} + +// Name returns the tool name +func (e *SubAgentToolExecutor) Name() string { + return core.SpawnAgentToolName +} + +// Execute spawns a child agent and returns its output as the tool result. +func (e *SubAgentToolExecutor) Execute(ctx context.Context, args map[string]interface{}, execCtx *core.ExecutionContext) (string, error) { + log := logger.Get() + + agentName, _ := args["agent"].(string) + query, _ := args["query"].(string) + + if agentName == "" { + return "", fmt.Errorf("spawn_agent requires 'agent' parameter") + } + if query == "" { + return "", fmt.Errorf("spawn_agent requires 'query' parameter") + } + + // Look up sub-agent definition + saDef, ok := e.subAgents[agentName] + if !ok { + return "", fmt.Errorf("unknown sub-agent: %s", agentName) + } + + // Check depth limit + childDepth := e.currentDepth + 1 + if childDepth > e.maxDepth { + return "", fmt.Errorf("sub-agent depth limit exceeded (depth %d > max %d)", childDepth, e.maxDepth) + } + + log.Debug("Spawning sub-agent", + zap.String("agent", agentName), + zap.Int("depth", childDepth), + zap.Int("max_depth", e.maxDepth), + ) + + // Build synthetic step from SubAgentDef + syntheticStep := buildSyntheticStep(&saDef, query) + + // Create child AgentExecutor + childExec := NewAgentExecutor(e.templateEngine, e.funcRegistry) + childExec.SetConfig(e.config) + childExec.SetSilent(e.silent) + childExec.SetDepthContext(childDepth, e.maxDepth) + + // Execute child agent + result, err := childExec.Execute(ctx, syntheticStep, execCtx) + if err != nil { + // Return error as string tool result (don't crash parent) + errMsg := fmt.Sprintf("Sub-agent '%s' failed: %s", agentName, err.Error()) + log.Warn("Sub-agent execution failed", + zap.String("agent", agentName), + zap.Error(err), + ) + return errMsg, nil + } + + // Merge child tokens into parent + if e.parentState != nil { + childTotal, _ := result.Exports["agent_total_tokens"].(int) + childPrompt, _ := result.Exports["agent_prompt_tokens"].(int) + childCompletion, _ := result.Exports["agent_completion_tokens"].(int) + e.parentState.MergeTokens(childTotal, childPrompt, childCompletion) + } + + return result.Output, nil +} + +// buildSyntheticStep creates a Step from a SubAgentDef and query. +func buildSyntheticStep(sa *core.SubAgentDef, query string) *core.Step { + maxIterations := sa.MaxIterations + if maxIterations <= 0 { + maxIterations = 10 + } + + step := &core.Step{ + Name: "sub-agent-" + sa.Name, + Type: core.StepTypeAgent, + Query: query, + SystemPrompt: sa.SystemPrompt, + AgentTools: sa.AgentTools, + MaxIterations: maxIterations, + Models: sa.Models, + LLMConfig: sa.LLMConfig, + OutputSchema: sa.OutputSchema, + Memory: sa.Memory, + StopCondition: sa.StopCondition, + SubAgents: sa.SubAgents, + OnToolStart: sa.OnToolStart, + OnToolEnd: sa.OnToolEnd, + } + + return step +} + +// BuildToolRegistryWithSubAgents constructs a ToolExecutorRegistry that includes +// both standard tools and the spawn_agent tool for sub-agent delegation. +func BuildToolRegistryWithSubAgents( + toolDefs []core.AgentToolDef, + funcRegistry *functions.Registry, + engine template.TemplateEngine, + cfg *config.Config, + silent bool, + currentDepth int, + maxDepth int, + parentState *agentState, + subAgents []core.SubAgentDef, +) *ToolExecutorRegistry { + reg := BuildToolRegistry(toolDefs, funcRegistry) + + if len(subAgents) > 0 { + // Build name → def map + saMap := make(map[string]core.SubAgentDef, len(subAgents)) + for _, sa := range subAgents { + saMap[sa.Name] = sa + } + + reg.Register(&SubAgentToolExecutor{ + subAgents: saMap, + templateEngine: engine, + funcRegistry: funcRegistry, + config: cfg, + silent: silent, + currentDepth: currentDepth, + maxDepth: maxDepth, + parentState: parentState, + }) + } + + return reg +} + +// executeToolCallViaRegistry parses tool call arguments and dispatches to the registry +func executeToolCallViaRegistry( + ctx context.Context, + tc core.LLMToolCall, + toolRegistry *ToolExecutorRegistry, + execCtx *core.ExecutionContext, + log *zap.Logger, +) (string, error) { + funcName := tc.Function.Name + argsJSON := tc.Function.Arguments + + log.Debug("Executing tool call via registry", + zap.String("tool", funcName), + zap.String("args", argsJSON), + ) + + var args map[string]interface{} + if argsJSON != "" { + if err := json.Unmarshal([]byte(argsJSON), &args); err != nil { + return "", fmt.Errorf("failed to parse tool arguments for %s: %w", funcName, err) + } + } + if args == nil { + args = make(map[string]interface{}) + } + + return toolRegistry.Execute(ctx, funcName, args, execCtx) +} diff --git a/internal/executor/tool_executor_test.go b/internal/executor/tool_executor_test.go new file mode 100644 index 0000000..0b32788 --- /dev/null +++ b/internal/executor/tool_executor_test.go @@ -0,0 +1,284 @@ +package executor + +import ( + "context" + "testing" + + "github.com/j3ssie/osmedeus/v5/internal/core" + "github.com/j3ssie/osmedeus/v5/internal/functions" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestToolExecutorRegistry_RegisterAndGet(t *testing.T) { + reg := NewToolExecutorRegistry() + + preset := NewPresetToolExecutor("bash", nil) + reg.Register(preset) + + got, ok := reg.Get("bash") + assert.True(t, ok) + assert.Equal(t, "bash", got.Name()) + + _, ok = reg.Get("nonexistent") + assert.False(t, ok) +} + +func TestToolExecutorRegistry_ExecuteUnknown(t *testing.T) { + reg := NewToolExecutorRegistry() + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + _, err := reg.Execute(context.Background(), "unknown_tool", nil, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "unknown tool") +} + +func TestPresetToolExecutor_Name(t *testing.T) { + preset := NewPresetToolExecutor("bash", nil) + assert.Equal(t, "bash", preset.Name()) +} + +func TestPresetToolExecutor_Execute(t *testing.T) { + funcRegistry := functions.NewRegistry() + preset := NewPresetToolExecutor("bash", funcRegistry) + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + args := map[string]interface{}{"command": "echo hello"} + result, err := preset.Execute(context.Background(), args, execCtx) + require.NoError(t, err) + assert.Contains(t, result, "hello") +} + +func TestCustomToolExecutor_Name(t *testing.T) { + custom := NewCustomToolExecutor("greet", `"Hello " + args.name`, nil) + assert.Equal(t, "greet", custom.Name()) +} + +func TestCustomToolExecutor_Execute(t *testing.T) { + funcRegistry := functions.NewRegistry() + custom := NewCustomToolExecutor("greet", `"Hello " + args.name`, funcRegistry) + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + args := map[string]interface{}{"name": "World"} + result, err := custom.Execute(context.Background(), args, execCtx) + require.NoError(t, err) + assert.Equal(t, "Hello World", result) +} + +func TestBuildToolRegistry_PresetOnly(t *testing.T) { + funcRegistry := functions.NewRegistry() + toolDefs := []core.AgentToolDef{ + {Preset: "bash"}, + {Preset: "read_file"}, + } + + reg := BuildToolRegistry(toolDefs, funcRegistry) + + _, ok := reg.Get("bash") + assert.True(t, ok) + + _, ok = reg.Get("read_file") + assert.True(t, ok) + + _, ok = reg.Get("nonexistent") + assert.False(t, ok) +} + +func TestBuildToolRegistry_CustomOnly(t *testing.T) { + funcRegistry := functions.NewRegistry() + toolDefs := []core.AgentToolDef{ + { + Name: "greet", + Description: "Greet someone", + Handler: `"Hello " + args.name`, + }, + } + + reg := BuildToolRegistry(toolDefs, funcRegistry) + + te, ok := reg.Get("greet") + assert.True(t, ok) + assert.Equal(t, "greet", te.Name()) +} + +func TestBuildToolRegistry_Mixed(t *testing.T) { + funcRegistry := functions.NewRegistry() + toolDefs := []core.AgentToolDef{ + {Preset: "bash"}, + { + Name: "custom_tool", + Description: "A custom tool", + Handler: `"result"`, + }, + } + + reg := BuildToolRegistry(toolDefs, funcRegistry) + + _, ok := reg.Get("bash") + assert.True(t, ok) + + _, ok = reg.Get("custom_tool") + assert.True(t, ok) +} + +func TestSubAgentToolExecutor_Name(t *testing.T) { + exec := &SubAgentToolExecutor{} + assert.Equal(t, core.SpawnAgentToolName, exec.Name()) +} + +func TestSubAgentToolExecutor_Execute_UnknownAgent(t *testing.T) { + exec := &SubAgentToolExecutor{ + subAgents: map[string]core.SubAgentDef{ + "recon": {Name: "recon"}, + }, + } + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + _, err := exec.Execute(context.Background(), map[string]interface{}{ + "agent": "nonexistent", + "query": "do something", + }, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "unknown sub-agent") +} + +func TestSubAgentToolExecutor_Execute_MissingQuery(t *testing.T) { + exec := &SubAgentToolExecutor{ + subAgents: map[string]core.SubAgentDef{ + "recon": {Name: "recon"}, + }, + } + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + _, err := exec.Execute(context.Background(), map[string]interface{}{ + "agent": "recon", + }, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "requires 'query'") +} + +func TestSubAgentToolExecutor_Execute_MissingAgent(t *testing.T) { + exec := &SubAgentToolExecutor{ + subAgents: map[string]core.SubAgentDef{}, + } + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + _, err := exec.Execute(context.Background(), map[string]interface{}{ + "query": "do something", + }, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "requires 'agent'") +} + +func TestSubAgentToolExecutor_Execute_DepthLimit(t *testing.T) { + exec := &SubAgentToolExecutor{ + subAgents: map[string]core.SubAgentDef{ + "recon": {Name: "recon", AgentTools: []core.AgentToolDef{{Preset: "bash"}}}, + }, + currentDepth: 3, + maxDepth: 3, + } + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + _, err := exec.Execute(context.Background(), map[string]interface{}{ + "agent": "recon", + "query": "scan ports", + }, execCtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "depth limit exceeded") +} + +func TestBuildToolRegistryWithSubAgents(t *testing.T) { + funcRegistry := functions.NewRegistry() + toolDefs := []core.AgentToolDef{{Preset: "bash"}} + subAgents := []core.SubAgentDef{ + {Name: "recon", Description: "Recon agent", AgentTools: []core.AgentToolDef{{Preset: "bash"}}}, + } + + reg := BuildToolRegistryWithSubAgents( + toolDefs, funcRegistry, nil, nil, true, + 0, 3, nil, subAgents, + ) + + // Should have bash and spawn_agent + _, ok := reg.Get("bash") + assert.True(t, ok) + + _, ok = reg.Get(core.SpawnAgentToolName) + assert.True(t, ok) +} + +func TestBuildToolRegistryWithSubAgents_NoSubAgents(t *testing.T) { + funcRegistry := functions.NewRegistry() + toolDefs := []core.AgentToolDef{{Preset: "bash"}} + + reg := BuildToolRegistryWithSubAgents( + toolDefs, funcRegistry, nil, nil, true, + 0, 3, nil, nil, + ) + + // Should have bash only (no spawn_agent) + _, ok := reg.Get("bash") + assert.True(t, ok) + + _, ok = reg.Get(core.SpawnAgentToolName) + assert.False(t, ok) +} + +func TestBuildSyntheticStep(t *testing.T) { + sa := &core.SubAgentDef{ + Name: "recon", + SystemPrompt: "You are a recon specialist", + AgentTools: []core.AgentToolDef{{Preset: "bash"}, {Preset: "http_get"}}, + MaxIterations: 5, + Models: []string{"gpt-4o"}, + StopCondition: `contains(agent_content, "DONE")`, + SubAgents: []core.SubAgentDef{ + {Name: "nested", AgentTools: []core.AgentToolDef{{Preset: "bash"}}}, + }, + } + + step := buildSyntheticStep(sa, "scan ports on target") + assert.Equal(t, "sub-agent-recon", step.Name) + assert.Equal(t, core.StepTypeAgent, step.Type) + assert.Equal(t, "scan ports on target", step.Query) + assert.Equal(t, "You are a recon specialist", step.SystemPrompt) + assert.Len(t, step.AgentTools, 2) + assert.Equal(t, 5, step.MaxIterations) + assert.Equal(t, []string{"gpt-4o"}, step.Models) + assert.Len(t, step.SubAgents, 1) +} + +func TestBuildSyntheticStep_DefaultMaxIterations(t *testing.T) { + sa := &core.SubAgentDef{ + Name: "minimal", + AgentTools: []core.AgentToolDef{{Preset: "bash"}}, + } + step := buildSyntheticStep(sa, "test") + assert.Equal(t, 10, step.MaxIterations) +} + +func TestToolExecutorRegistry_FullFlow(t *testing.T) { + funcRegistry := functions.NewRegistry() + toolDefs := []core.AgentToolDef{ + {Preset: "bash"}, + { + Name: "echo_tool", + Description: "Echo something", + Handler: `"echoed: " + args.text`, + }, + } + + reg := BuildToolRegistry(toolDefs, funcRegistry) + execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com") + + // Test preset execution + result, err := reg.Execute(context.Background(), "bash", map[string]interface{}{"command": "echo test"}, execCtx) + require.NoError(t, err) + assert.Contains(t, result, "test") + + // Test custom execution + result, err = reg.Execute(context.Background(), "echo_tool", map[string]interface{}{"text": "hello"}, execCtx) + require.NoError(t, err) + assert.Equal(t, "echoed: hello", result) +} diff --git a/internal/functions/constants.go b/internal/functions/constants.go index f7dc357..38efb4d 100644 --- a/internal/functions/constants.go +++ b/internal/functions/constants.go @@ -26,6 +26,7 @@ const ( FnGrepString = "grep_string" // grep_string(source, str) -> string FnGrepRegex = "grep_regex" // grep_regex(source, pattern) -> string FnRemoveBlankLines = "remove_blank_lines" // remove_blank_lines(path) -> bool (in-place) + FnChunkFile = "chunk_file" // chunk_file(input, lines_per_chunk, output) -> bool ) // String Functions - String manipulation operations @@ -71,17 +72,21 @@ const ( // Utility Functions - General utility operations const ( - FnLen = "len" // len(val) -> int - FnIsEmpty = "is_empty" // is_empty(val) -> bool - FnIsNotEmpty = "is_not_empty" // is_not_empty(val) -> bool - FnPrintf = "printf" // printf(message) -> void (print message to stdout) - FnCatFile = "cat_file" // cat_file(path) -> void (print file content to stdout) - FnExit = "exit" // exit(code) -> void (exit scan with code) - FnExecCmd = "exec_cmd" // exec_cmd(command) -> string (alias for bash) - FnBash = "bash" - FnSleep = "sleep" // sleep(seconds) -> void (pause for n seconds) - FnCommandExists = "command_exists" // command_exists(command) -> bool (check if command exists in PATH) - FnPickValid = "pick_valid" // pick_valid(v1, v2, ..., v10) -> any (first valid value) + FnLen = "len" // len(val) -> int + FnIsEmpty = "is_empty" // is_empty(val) -> bool + FnIsNotEmpty = "is_not_empty" // is_not_empty(val) -> bool + FnPrintf = "printf" // printf(message) -> void (print message to stdout) + FnCatFile = "cat_file" // cat_file(path) -> void (print file content to stdout) + FnExit = "exit" // exit(code) -> void (exit scan with code) + FnExecCmd = "exec_cmd" // exec_cmd(command) -> string (alias for bash) + FnBash = "bash" + FnSleep = "sleep" // sleep(seconds) -> void (pause for n seconds) + FnCommandExists = "command_exists" // command_exists(command) -> bool (check if command exists in PATH) + FnPickValid = "pick_valid" // pick_valid(v1, v2, ..., v10) -> any (first valid value) + FnRunModule = "run_module" // run_module(module, target, params?) -> string (run osmedeus module) + FnRunFlow = "run_flow" // run_flow(flow, target, params?) -> string (run osmedeus flow) + FnExecPython = "exec_python" // exec_python(code) -> string (run inline Python, prefer python3) + FnExecPythonFile = "exec_python_file" // exec_python_file(path) -> string (run Python file, prefer python3) ) // Logging Functions - Log messages with level prefixes @@ -167,18 +172,18 @@ const ( // Unix Command Wrappers - Wrappers around common Unix commands const ( - FnSortUnix = "sort_unix" // sort_unix(inputFile, outputFile?) -> bool (LC_ALL=C sort -u) - FnWgetUnix = "wget_unix" // wget_unix(url, outputPath?) -> bool - FnWget = "wget" // wget(url, outputPath) -> bool (pure Go, segmented download) - FnGitClone = "git_clone" // git_clone(repo, dest?) -> bool - FnGitCloneSubfolder = "git_clone_subfolder" // git_clone_subfolder(git_url, subfolder, dest) -> bool - FnZipUnix = "zip_unix" // zip_unix(source, dest) -> bool (zip -r dest source) - FnUnzipUnix = "unzip_unix" // unzip_unix(source, dest?) -> bool (unzip source -d dest) - FnTarUnix = "tar_unix" // tar_unix(source, dest) -> bool (tar -czf dest source) - FnUntarUnix = "untar_unix" // untar_unix(source, dest?) -> bool (tar -xzf source -C dest) - FnDiffUnix = "diff_unix" // diff_unix(file1, file2, output?) -> string - FnSedStringReplace = "sed_string_replace" // sed_string_replace(sed_syntax, source, dest) -> bool - FnSedRegexReplace = "sed_regex_replace" // sed_regex_replace(sed_syntax, source, dest) -> bool + FnSortUnix = "sort_unix" // sort_unix(inputFile, outputFile?) -> bool (LC_ALL=C sort -u) + FnWgetUnix = "wget_unix" // wget_unix(url, outputPath?) -> bool + FnWget = "wget" // wget(url, outputPath) -> bool (pure Go, segmented download) + FnGitClone = "git_clone" // git_clone(repo, dest?) -> bool + FnGitCloneSubfolder = "git_clone_subfolder" // git_clone_subfolder(git_url, subfolder, dest) -> bool + FnZipUnix = "zip_unix" // zip_unix(source, dest) -> bool (zip -r dest source) + FnUnzipUnix = "unzip_unix" // unzip_unix(source, dest?) -> bool (unzip source -d dest) + FnTarUnix = "tar_unix" // tar_unix(source, dest) -> bool (tar -czf dest source) + FnUntarUnix = "untar_unix" // untar_unix(source, dest?) -> bool (tar -xzf source -C dest) + FnDiffUnix = "diff_unix" // diff_unix(file1, file2, output?) -> string + FnSedStringReplace = "sed_string_replace" // sed_string_replace(sed_syntax, source, dest) -> bool + FnSedRegexReplace = "sed_regex_replace" // sed_regex_replace(sed_syntax, source, dest) -> bool ) // Installer Functions - Download and install packages @@ -325,6 +330,7 @@ func AllFunctions() []string { FnGrepString, FnGrepRegex, FnRemoveBlankLines, + FnChunkFile, // String Functions FnTrim, @@ -373,6 +379,10 @@ func AllFunctions() []string { FnSleep, FnCommandExists, FnPickValid, + FnRunModule, + FnRunFlow, + FnExecPython, + FnExecPythonFile, // Logging Functions FnLogDebug, @@ -649,6 +659,7 @@ func FunctionRegistry() map[string][]FunctionInfo { {FnGrepString, "grep_string(source, str)", "Return lines containing string", "string", "grep_string('{{Output}}/in.txt', 'admin')"}, {FnGrepRegex, "grep_regex(source, pattern)", "Return lines matching regex", "string", "grep_regex('{{Output}}/in.txt', '.*api.*')"}, {FnRemoveBlankLines, "remove_blank_lines(path)", "Remove blank lines from file in-place", "bool", "remove_blank_lines('{{Output}}/urls.txt')"}, + {FnChunkFile, "chunk_file(input, lines_per_chunk, output)", "Split file into chunks and write manifest of chunk paths", "bool", "chunk_file('{{Output}}/urls.txt', 100, '{{Output}}/url_chunks.txt')"}, }, CategoryString: { {FnTrim, "trim(str)", "Trim whitespace", "string", "trim(' hello ')"}, @@ -689,6 +700,10 @@ func FunctionRegistry() map[string][]FunctionInfo { {FnSleep, "sleep(seconds)", "Pause for n seconds", "void", "sleep(5)"}, {FnCommandExists, "command_exists(command)", "Check if command exists in PATH", "bool", "command_exists('nmap')"}, {FnPickValid, "pick_valid(v1, v2, ..., v10)", "Return first valid value from up to 10 arguments", "any", "pick_valid('', '', 'hello', 'world')"}, + {FnRunModule, "run_module(module, target, params?)", "Run osmedeus module as subprocess, optional comma-separated key=value params", "string", "run_module('subdomain', 'example.com', 'threads=10,deep=true')"}, + {FnRunFlow, "run_flow(flow, target, params?)", "Run osmedeus flow as subprocess, optional comma-separated key=value params", "string", "run_flow('general', 'example.com')"}, + {FnExecPython, "exec_python(code)", "Run inline Python code via python3 -c (falls back to python)", "string", "exec_python('print(2+2)')"}, + {FnExecPythonFile, "exec_python_file(path)", "Run a Python file via python3 (falls back to python)", "string", "exec_python_file('/tmp/script.py')"}, }, CategoryLogging: { {FnLogDebug, "log_debug(message)", "Log debug message with [DEBUG] prefix", "void", "log_debug('Processing target')"}, diff --git a/internal/functions/db_functions.go b/internal/functions/db_functions.go index 7a5f7be..9ce3f92 100644 --- a/internal/functions/db_functions.go +++ b/internal/functions/db_functions.go @@ -3,8 +3,8 @@ package functions import ( "bufio" "context" - "github.com/j3ssie/osmedeus/v5/internal/json" "fmt" + "github.com/j3ssie/osmedeus/v5/internal/json" "os" "path/filepath" "strconv" diff --git a/internal/functions/file_functions.go b/internal/functions/file_functions.go index 31eec1b..d4da726 100644 --- a/internal/functions/file_functions.go +++ b/internal/functions/file_functions.go @@ -3,6 +3,7 @@ package functions import ( "archive/zip" "bufio" + "fmt" "io" "os" "path/filepath" @@ -460,6 +461,108 @@ func (vf *vmFunc) removeBlankLines(call goja.FunctionCall) goja.Value { return vf.vm.ToValue(true) } +// chunkFile splits an input file into chunks of N lines each, writing chunk paths to an output manifest. +// Blank lines are skipped. Chunk files are named {base}_part_{N}{ext} in the same directory as input. +// Usage: chunk_file(input, lines_per_chunk, output) -> bool +func (vf *vmFunc) chunkFile(call goja.FunctionCall) goja.Value { + input := call.Argument(0).String() + chunkSize := int(call.Argument(1).ToInteger()) + output := call.Argument(2).String() + logger.Get().Debug("Calling "+terminal.HiGreen("chunk_file"), + zap.String("input", input), zap.Int("chunkSize", chunkSize), zap.String("output", output)) + + if input == "undefined" || input == "" || output == "undefined" || output == "" { + logger.Get().Warn("chunk_file: input and output paths are required") + return vf.vm.ToValue(false) + } + if chunkSize <= 0 { + logger.Get().Warn("chunk_file: lines_per_chunk must be > 0", zap.Int("chunkSize", chunkSize)) + return vf.vm.ToValue(false) + } + + // Read input file, skipping blank lines + file, err := os.Open(input) + if err != nil { + logger.Get().Warn("chunk_file: failed to open input file", zap.String("input", input), zap.Error(err)) + return vf.vm.ToValue(false) + } + + var lines []string + scanner := bufio.NewScanner(file) + for scanner.Scan() { + line := scanner.Text() + if strings.TrimSpace(line) != "" { + lines = append(lines, line) + } + } + if err := scanner.Err(); err != nil { + _ = file.Close() + logger.Get().Warn("chunk_file: failed to read input file", zap.String("input", input), zap.Error(err)) + return vf.vm.ToValue(false) + } + _ = file.Close() + + if len(lines) == 0 { + logger.Get().Debug(terminal.HiGreen("chunk_file") + ": input file has no non-blank lines") + // Write empty manifest + if err := os.MkdirAll(filepath.Dir(output), 0755); err != nil { + return vf.vm.ToValue(false) + } + if err := os.WriteFile(output, []byte(""), 0644); err != nil { + return vf.vm.ToValue(false) + } + return vf.vm.ToValue(true) + } + + // Compute chunk file naming from input path + dir := filepath.Dir(input) + ext := filepath.Ext(input) + base := strings.TrimSuffix(filepath.Base(input), ext) + + var chunkPaths []string + chunkIdx := 0 + + for i := 0; i < len(lines); i += chunkSize { + end := i + chunkSize + if end > len(lines) { + end = len(lines) + } + chunk := lines[i:end] + + chunkName := fmt.Sprintf("%s_part_%d%s", base, chunkIdx, ext) + chunkPath := filepath.Join(dir, chunkName) + + content := strings.Join(chunk, "\n") + "\n" + if err := os.WriteFile(chunkPath, []byte(content), 0644); err != nil { + logger.Get().Warn("chunk_file: failed to write chunk", + zap.String("chunkPath", chunkPath), zap.Error(err)) + return vf.vm.ToValue(false) + } + + chunkPaths = append(chunkPaths, chunkPath) + chunkIdx++ + } + + // Write manifest (one chunk path per line) + if err := os.MkdirAll(filepath.Dir(output), 0755); err != nil { + logger.Get().Warn("chunk_file: failed to create output directory", zap.Error(err)) + return vf.vm.ToValue(false) + } + + manifest := strings.Join(chunkPaths, "\n") + "\n" + if err := os.WriteFile(output, []byte(manifest), 0644); err != nil { + logger.Get().Warn("chunk_file: failed to write manifest", zap.String("output", output), zap.Error(err)) + return vf.vm.ToValue(false) + } + + logger.Get().Debug(terminal.HiGreen("chunk_file")+" result", + zap.String("input", input), + zap.Int("totalLines", len(lines)), + zap.Int("chunks", len(chunkPaths)), + zap.String("output", output)) + return vf.vm.ToValue(true) +} + // zipDir creates a zip archive from a directory using Go's archive/zip // Usage: zip_dir(source, dest) -> bool func (vf *vmFunc) zipDir(call goja.FunctionCall) goja.Value { diff --git a/internal/functions/file_functions_test.go b/internal/functions/file_functions_test.go index 5d8f728..419d489 100644 --- a/internal/functions/file_functions_test.go +++ b/internal/functions/file_functions_test.go @@ -3,6 +3,7 @@ package functions import ( "os" "path/filepath" + "strings" "testing" "github.com/stretchr/testify/assert" @@ -101,3 +102,221 @@ func TestRemoveBlankLines(t *testing.T) { assert.Equal(t, " line with leading spaces\n\tmiddle with tab\nlast line\n", string(data)) }) } + +func TestChunkFile(t *testing.T) { + runtime := NewOttoRuntime() + + t.Run("splits file into equal chunks", func(t *testing.T) { + tmpDir := t.TempDir() + input := filepath.Join(tmpDir, "urls.txt") + output := filepath.Join(tmpDir, "chunks.txt") + + // Write 6 lines + content := "line1\nline2\nline3\nline4\nline5\nline6\n" + require.NoError(t, os.WriteFile(input, []byte(content), 0644)) + + result, err := runtime.Execute(`chunk_file("`+input+`", 2, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, true, result) + + // Verify manifest + manifest, err := os.ReadFile(output) + require.NoError(t, err) + lines := strings.Split(strings.TrimSpace(string(manifest)), "\n") + assert.Len(t, lines, 3, "should have 3 chunks for 6 lines with chunk_size=2") + + // Verify chunk 0 + chunk0, err := os.ReadFile(lines[0]) + require.NoError(t, err) + assert.Equal(t, "line1\nline2\n", string(chunk0)) + + // Verify chunk 1 + chunk1, err := os.ReadFile(lines[1]) + require.NoError(t, err) + assert.Equal(t, "line3\nline4\n", string(chunk1)) + + // Verify chunk 2 + chunk2, err := os.ReadFile(lines[2]) + require.NoError(t, err) + assert.Equal(t, "line5\nline6\n", string(chunk2)) + }) + + t.Run("handles uneven split (last chunk smaller)", func(t *testing.T) { + tmpDir := t.TempDir() + input := filepath.Join(tmpDir, "data.txt") + output := filepath.Join(tmpDir, "chunks.txt") + + // Write 5 lines + content := "a\nb\nc\nd\ne\n" + require.NoError(t, os.WriteFile(input, []byte(content), 0644)) + + result, err := runtime.Execute(`chunk_file("`+input+`", 3, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, true, result) + + manifest, err := os.ReadFile(output) + require.NoError(t, err) + lines := strings.Split(strings.TrimSpace(string(manifest)), "\n") + assert.Len(t, lines, 2, "should have 2 chunks for 5 lines with chunk_size=3") + + // First chunk has 3 lines + chunk0, err := os.ReadFile(lines[0]) + require.NoError(t, err) + assert.Equal(t, "a\nb\nc\n", string(chunk0)) + + // Second chunk has 2 lines + chunk1, err := os.ReadFile(lines[1]) + require.NoError(t, err) + assert.Equal(t, "d\ne\n", string(chunk1)) + }) + + t.Run("skips blank lines", func(t *testing.T) { + tmpDir := t.TempDir() + input := filepath.Join(tmpDir, "mixed.txt") + output := filepath.Join(tmpDir, "chunks.txt") + + content := "line1\n\nline2\n \nline3\n\n" + require.NoError(t, os.WriteFile(input, []byte(content), 0644)) + + result, err := runtime.Execute(`chunk_file("`+input+`", 2, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, true, result) + + manifest, err := os.ReadFile(output) + require.NoError(t, err) + lines := strings.Split(strings.TrimSpace(string(manifest)), "\n") + assert.Len(t, lines, 2, "should have 2 chunks for 3 non-blank lines with chunk_size=2") + + chunk0, err := os.ReadFile(lines[0]) + require.NoError(t, err) + assert.Equal(t, "line1\nline2\n", string(chunk0)) + + chunk1, err := os.ReadFile(lines[1]) + require.NoError(t, err) + assert.Equal(t, "line3\n", string(chunk1)) + }) + + t.Run("preserves extension in chunk names", func(t *testing.T) { + tmpDir := t.TempDir() + input := filepath.Join(tmpDir, "urls.txt") + output := filepath.Join(tmpDir, "chunks.txt") + + content := "url1\nurl2\nurl3\n" + require.NoError(t, os.WriteFile(input, []byte(content), 0644)) + + result, err := runtime.Execute(`chunk_file("`+input+`", 2, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, true, result) + + manifest, err := os.ReadFile(output) + require.NoError(t, err) + lines := strings.Split(strings.TrimSpace(string(manifest)), "\n") + + assert.Contains(t, lines[0], "urls_part_0.txt") + assert.Contains(t, lines[1], "urls_part_1.txt") + }) + + t.Run("empty file produces empty manifest", func(t *testing.T) { + tmpDir := t.TempDir() + input := filepath.Join(tmpDir, "empty.txt") + output := filepath.Join(tmpDir, "chunks.txt") + + require.NoError(t, os.WriteFile(input, []byte(""), 0644)) + + result, err := runtime.Execute(`chunk_file("`+input+`", 10, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, true, result) + + manifest, err := os.ReadFile(output) + require.NoError(t, err) + assert.Equal(t, "", string(manifest)) + }) + + t.Run("file with only blank lines produces empty manifest", func(t *testing.T) { + tmpDir := t.TempDir() + input := filepath.Join(tmpDir, "blanks.txt") + output := filepath.Join(tmpDir, "chunks.txt") + + require.NoError(t, os.WriteFile(input, []byte("\n\n \n\n"), 0644)) + + result, err := runtime.Execute(`chunk_file("`+input+`", 5, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, true, result) + + manifest, err := os.ReadFile(output) + require.NoError(t, err) + assert.Equal(t, "", string(manifest)) + }) + + t.Run("chunk size larger than file creates single chunk", func(t *testing.T) { + tmpDir := t.TempDir() + input := filepath.Join(tmpDir, "small.txt") + output := filepath.Join(tmpDir, "chunks.txt") + + content := "a\nb\nc\n" + require.NoError(t, os.WriteFile(input, []byte(content), 0644)) + + result, err := runtime.Execute(`chunk_file("`+input+`", 100, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, true, result) + + manifest, err := os.ReadFile(output) + require.NoError(t, err) + lines := strings.Split(strings.TrimSpace(string(manifest)), "\n") + assert.Len(t, lines, 1, "should have 1 chunk when chunk_size > total lines") + + chunk, err := os.ReadFile(lines[0]) + require.NoError(t, err) + assert.Equal(t, "a\nb\nc\n", string(chunk)) + }) + + t.Run("empty input path returns false", func(t *testing.T) { + tmpDir := t.TempDir() + output := filepath.Join(tmpDir, "chunks.txt") + + result, err := runtime.Execute(`chunk_file("", 10, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, false, result) + }) + + t.Run("empty output path returns false", func(t *testing.T) { + tmpDir := t.TempDir() + input := filepath.Join(tmpDir, "input.txt") + require.NoError(t, os.WriteFile(input, []byte("line\n"), 0644)) + + result, err := runtime.Execute(`chunk_file("`+input+`", 10, "")`, nil) + require.NoError(t, err) + assert.Equal(t, false, result) + }) + + t.Run("chunk size zero returns false", func(t *testing.T) { + tmpDir := t.TempDir() + input := filepath.Join(tmpDir, "input.txt") + output := filepath.Join(tmpDir, "chunks.txt") + require.NoError(t, os.WriteFile(input, []byte("line\n"), 0644)) + + result, err := runtime.Execute(`chunk_file("`+input+`", 0, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, false, result) + }) + + t.Run("negative chunk size returns false", func(t *testing.T) { + tmpDir := t.TempDir() + input := filepath.Join(tmpDir, "input.txt") + output := filepath.Join(tmpDir, "chunks.txt") + require.NoError(t, os.WriteFile(input, []byte("line\n"), 0644)) + + result, err := runtime.Execute(`chunk_file("`+input+`", -5, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, false, result) + }) + + t.Run("non-existent input file returns false", func(t *testing.T) { + tmpDir := t.TempDir() + output := filepath.Join(tmpDir, "chunks.txt") + + result, err := runtime.Execute(`chunk_file("/nonexistent/input.txt", 10, "`+output+`")`, nil) + require.NoError(t, err) + assert.Equal(t, false, result) + }) +} diff --git a/internal/functions/goja_runtime.go b/internal/functions/goja_runtime.go index 4fe8a24..e68e697 100644 --- a/internal/functions/goja_runtime.go +++ b/internal/functions/goja_runtime.go @@ -67,6 +67,7 @@ func (r *GojaRuntime) registerFunctionsOnVM(vm *goja.Runtime) { _ = vm.Set(FnGrepString, vf.grepString) _ = vm.Set(FnGrepRegex, vf.grepRegex) _ = vm.Set(FnRemoveBlankLines, vf.removeBlankLines) + _ = vm.Set(FnChunkFile, vf.chunkFile) // String functions _ = vm.Set(FnTrim, vf.trim) @@ -115,6 +116,10 @@ func (r *GojaRuntime) registerFunctionsOnVM(vm *goja.Runtime) { _ = vm.Set(FnSleep, vf.sleep) _ = vm.Set(FnCommandExists, vf.commandExists) _ = vm.Set(FnPickValid, vf.pickValid) + _ = vm.Set(FnRunModule, vf.runModule) + _ = vm.Set(FnRunFlow, vf.runFlow) + _ = vm.Set(FnExecPython, vf.execPython) + _ = vm.Set(FnExecPythonFile, vf.execPythonFile) // Logging functions _ = vm.Set(FnLogDebug, vf.logDebug) diff --git a/internal/functions/url_functions_test.go b/internal/functions/url_functions_test.go index d3b7006..f36d561 100644 --- a/internal/functions/url_functions_test.go +++ b/internal/functions/url_functions_test.go @@ -689,11 +689,11 @@ func TestParseURL_WithRegistry(t *testing.T) { func TestExtractDomainParts(t *testing.T) { tests := []struct { - name string - domain string - expectedSubdomain string - expectedRoot string - expectedTLD string + name string + domain string + expectedSubdomain string + expectedRoot string + expectedTLD string }{ { name: "simple domain", diff --git a/internal/functions/util_functions.go b/internal/functions/util_functions.go index 95cfb1a..343fc53 100644 --- a/internal/functions/util_functions.go +++ b/internal/functions/util_functions.go @@ -736,6 +736,64 @@ func (vf *vmFunc) bash(call goja.FunctionCall) goja.Value { return vf.execCmd(call) } +// findPythonBin returns "python3" if available, otherwise "python". +func findPythonBin() string { + if _, err := exec.LookPath("python3"); err == nil { + return "python3" + } + return "python" +} + +// execPython runs inline Python code via `python3 -c ''`. +// Usage: exec_python(code) -> string +func (vf *vmFunc) execPython(call goja.FunctionCall) goja.Value { + code := call.Argument(0).String() + logger.Get().Debug("Calling "+terminal.HiGreen("execPython"), zap.Int("codeLength", len(code))) + + if code == "undefined" || code == "" { + logger.Get().Warn("execPython: empty code provided") + return vf.vm.ToValue("") + } + + pythonBin := findPythonBin() + // @NOTE: This is intentional - exec_python() is a utility function exposed to workflow + // definitions for executing Python code. Input comes from trusted workflow YAML files. + cmd := exec.Command(pythonBin, "-c", code) + output, err := cmd.Output() + if err != nil { + logger.Get().Warn("execPython: command failed", zap.String("python", pythonBin), zap.Error(err)) + return vf.vm.ToValue("") + } + + logger.Get().Debug(terminal.HiGreen("execPython")+" result", zap.Int("outputLength", len(output))) + return vf.vm.ToValue(strings.TrimSpace(string(output))) +} + +// execPythonFile runs a Python file via `python3 `. +// Usage: exec_python_file(path) -> string +func (vf *vmFunc) execPythonFile(call goja.FunctionCall) goja.Value { + path := call.Argument(0).String() + logger.Get().Debug("Calling "+terminal.HiGreen("execPythonFile"), zap.String("path", path)) + + if path == "undefined" || path == "" { + logger.Get().Warn("execPythonFile: empty path provided") + return vf.vm.ToValue("") + } + + pythonBin := findPythonBin() + // @NOTE: This is intentional - exec_python_file() is a utility function exposed to workflow + // definitions for executing Python files. Input comes from trusted workflow YAML files. + cmd := exec.Command(pythonBin, path) + output, err := cmd.Output() + if err != nil { + logger.Get().Warn("execPythonFile: command failed", zap.String("python", pythonBin), zap.String("path", path), zap.Error(err)) + return vf.vm.ToValue("") + } + + logger.Get().Debug(terminal.HiGreen("execPythonFile")+" result", zap.String("path", path), zap.Int("outputLength", len(output))) + return vf.vm.ToValue(strings.TrimSpace(string(output))) +} + // commandExists checks if a command is available in PATH // Usage: commandExists(command) -> bool func (vf *vmFunc) commandExists(call goja.FunctionCall) goja.Value { @@ -1015,6 +1073,93 @@ func (vf *vmFunc) moveFile(call goja.FunctionCall) goja.Value { return vf.vm.ToValue(true) } +// parseParamsToFlags parses a comma-separated "key=value" string into -p flags. +// E.g. "threads=10,deep=true" -> ["-p", "threads=10", "-p", "deep=true"] +func parseParamsToFlags(params string) []string { + params = strings.TrimSpace(params) + if params == "" || params == "undefined" { + return nil + } + var flags []string + for _, pair := range strings.Split(params, ",") { + pair = strings.TrimSpace(pair) + if pair == "" || !strings.Contains(pair, "=") { + continue + } + flags = append(flags, "-p", pair) + } + return flags +} + +// runOsmedeus executes the current osmedeus binary with the given flag (-m or -f), name, target, and optional params. +// Returns the combined stdout+stderr output as a string. +func (vf *vmFunc) runOsmedeus(flag, name, target, params, funcName string) goja.Value { + // Find the current executable + exePath, err := os.Executable() + if err != nil { + logger.Get().Warn(funcName+": failed to find executable", zap.Error(err)) + return vf.vm.ToValue("") + } + + // Build command arguments + args := []string{"run", flag, name, "-t", target} + args = append(args, parseParamsToFlags(params)...) + + logger.Get().Debug("Calling "+terminal.HiGreen(funcName), + zap.String("exe", exePath), + zap.Strings("args", args)) + + // @NOTE: This is intentional - run_module/run_flow are utility functions exposed to workflow + // definitions for launching sub-scans. Input comes from trusted workflow YAML files. + cmd := exec.Command(exePath, args...) + output, err := cmd.CombinedOutput() + if err != nil { + logger.Get().Warn(funcName+": command failed", + zap.String("name", name), + zap.String("target", target), + zap.Error(err), + zap.String("output", string(output))) + } + + logger.Get().Debug(terminal.HiGreen(funcName)+" result", + zap.String("name", name), + zap.String("target", target), + zap.Int("outputLength", len(output))) + return vf.vm.ToValue(strings.TrimSpace(string(output))) +} + +// runModule runs an osmedeus module as a subprocess +// Usage: run_module(module, target, params?) -> string +// params is optional comma-separated key=value pairs: "threads=10,deep=true" +func (vf *vmFunc) runModule(call goja.FunctionCall) goja.Value { + module := call.Argument(0).String() + target := call.Argument(1).String() + params := call.Argument(2).String() + + if module == "undefined" || module == "" || target == "undefined" || target == "" { + logger.Get().Warn("run_module: module and target are required") + return vf.vm.ToValue("") + } + + return vf.runOsmedeus("-m", module, target, params, "run_module") +} + +// runFlow runs an osmedeus flow as a subprocess +// Usage: run_flow(flow, target, params?) -> string +// params is optional comma-separated key=value pairs: "threads=10,deep=true" +func (vf *vmFunc) runFlow(call goja.FunctionCall) goja.Value { + flow := call.Argument(0).String() + target := call.Argument(1).String() + params := call.Argument(2).String() + + if flow == "undefined" || flow == "" || target == "undefined" || target == "" { + logger.Get().Warn("run_flow: flow and target are required") + return vf.vm.ToValue("") + } + + return vf.runOsmedeus("-f", flow, target, params, "run_flow") +} + // copyFileBuffered copies a file using buffered I/O (memory-efficient for large files) func copyFileBuffered(source, dest string, mode os.FileMode) error { srcFile, err := os.Open(source) diff --git a/internal/functions/util_functions_test.go b/internal/functions/util_functions_test.go index 0909029..56a89d8 100644 --- a/internal/functions/util_functions_test.go +++ b/internal/functions/util_functions_test.go @@ -642,6 +642,154 @@ func TestPickValid(t *testing.T) { }) } +func TestParseParamsToFlags(t *testing.T) { + t.Run("empty string returns nil", func(t *testing.T) { + result := parseParamsToFlags("") + assert.Nil(t, result) + }) + + t.Run("undefined returns nil", func(t *testing.T) { + result := parseParamsToFlags("undefined") + assert.Nil(t, result) + }) + + t.Run("single param", func(t *testing.T) { + result := parseParamsToFlags("threads=10") + assert.Equal(t, []string{"-p", "threads=10"}, result) + }) + + t.Run("multiple params", func(t *testing.T) { + result := parseParamsToFlags("threads=10,deep=true") + assert.Equal(t, []string{"-p", "threads=10", "-p", "deep=true"}, result) + }) + + t.Run("params with spaces", func(t *testing.T) { + result := parseParamsToFlags(" threads=10 , deep=true ") + assert.Equal(t, []string{"-p", "threads=10", "-p", "deep=true"}, result) + }) + + t.Run("skips entries without equals sign", func(t *testing.T) { + result := parseParamsToFlags("threads=10,invalid,deep=true") + assert.Equal(t, []string{"-p", "threads=10", "-p", "deep=true"}, result) + }) + + t.Run("skips empty entries from trailing comma", func(t *testing.T) { + result := parseParamsToFlags("threads=10,") + assert.Equal(t, []string{"-p", "threads=10"}, result) + }) +} + +func TestRunModule(t *testing.T) { + runtime := NewOttoRuntime() + + t.Run("empty module returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`run_module("", "example.com")`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) + + t.Run("empty target returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`run_module("subdomain", "")`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) + + t.Run("undefined module returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`run_module(undefined, "example.com")`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) + + t.Run("undefined target returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`run_module("subdomain", undefined)`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) +} + +func TestRunFlow(t *testing.T) { + runtime := NewOttoRuntime() + + t.Run("empty flow returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`run_flow("", "example.com")`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) + + t.Run("empty target returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`run_flow("general", "")`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) + + t.Run("undefined flow returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`run_flow(undefined, "example.com")`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) + + t.Run("undefined target returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`run_flow("general", undefined)`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) +} + +func TestExecPython(t *testing.T) { + runtime := NewOttoRuntime() + + t.Run("simple print", func(t *testing.T) { + result, err := runtime.Execute(`exec_python("print('hello')")`, nil) + require.NoError(t, err) + assert.Equal(t, "hello", result) + }) + + t.Run("multiline code", func(t *testing.T) { + result, err := runtime.Execute(`exec_python("x = 2 + 3\nprint(x)")`, nil) + require.NoError(t, err) + assert.Equal(t, "5", result) + }) + + t.Run("empty code returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`exec_python("")`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) + + t.Run("invalid code returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`exec_python("import sys; sys.exit(1)")`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) +} + +func TestExecPythonFile(t *testing.T) { + runtime := NewOttoRuntime() + + t.Run("run temp python file", func(t *testing.T) { + tmpDir := t.TempDir() + pyFile := filepath.Join(tmpDir, "test.py") + err := os.WriteFile(pyFile, []byte("print('from file')"), 0644) + require.NoError(t, err) + + result, err := runtime.Execute(`exec_python_file("`+pyFile+`")`, nil) + require.NoError(t, err) + assert.Equal(t, "from file", result) + }) + + t.Run("empty path returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`exec_python_file("")`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) + + t.Run("nonexistent file returns empty string", func(t *testing.T) { + result, err := runtime.Execute(`exec_python_file("/tmp/nonexistent_py_file_12345.py")`, nil) + require.NoError(t, err) + assert.Equal(t, "", result) + }) +} + func TestMoveFile(t *testing.T) { runtime := NewOttoRuntime() diff --git a/internal/linter/rules.go b/internal/linter/rules.go index ea4e4eb..240e8ce 100644 --- a/internal/linter/rules.go +++ b/internal/linter/rules.go @@ -367,6 +367,76 @@ func (r *EmptyStepRule) Check(wast *WorkflowAST) []LintIssue { if len(step.Messages) == 0 && len(step.EmbeddingInput) == 0 { empty = true } + case core.StepTypeAgent: + if (step.Query == "" && len(step.Queries) == 0) || len(step.AgentTools) == 0 { + empty = true + } + // Agent-specific warnings + if step.Query != "" && len(step.Queries) > 0 { + line, col := wast.FindStepFieldPosition(i, "queries") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityWarning, + Message: fmt.Sprintf("Step '%s' has both 'query' and 'queries' — only one should be used", step.Name), + Suggestion: "Remove either 'query' or 'queries'", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].queries", i), + }) + } + if step.PlanPrompt != "" && step.Query == "" && len(step.Queries) == 0 { + line, col := wast.FindStepFieldPosition(i, "plan_prompt") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityWarning, + Message: fmt.Sprintf("Step '%s' has plan_prompt but no query/queries", step.Name), + Suggestion: "Add a 'query' or 'queries' field for the agent to execute after planning", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].plan_prompt", i), + }) + } + if step.MaxIterations > 50 { + line, col := wast.FindStepFieldPosition(i, "max_iterations") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityInfo, + Message: fmt.Sprintf("Step '%s' has max_iterations=%d which is suspiciously high", step.Name, step.MaxIterations), + Suggestion: "Consider lowering max_iterations to avoid excessive LLM calls", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].max_iterations", i), + }) + } + for j, tool := range step.AgentTools { + if tool.IsPreset() { + if _, ok := core.GetPresetTool(tool.Preset); !ok { + line, col := wast.FindStepFieldPosition(i, "agent_tools") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityWarning, + Message: fmt.Sprintf("Step '%s' references unknown preset tool '%s'", step.Name, tool.Preset), + Suggestion: "Check the preset tool name against the preset tool registry", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].agent_tools[%d].preset", i, j), + }) + } + } else if tool.Handler != "" && tool.Description == "" { + line, col := wast.FindStepFieldPosition(i, "agent_tools") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityWarning, + Message: fmt.Sprintf("Step '%s' has custom tool '%s' with handler but no description", step.Name, tool.Name), + Suggestion: "Add a description so the LLM knows when to use this tool", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].agent_tools[%d].description", i, j), + }) + } + } + // Sub-agent validation + issues = append(issues, validateSubAgents(step, i, wast, r)...) } if empty { @@ -727,6 +797,126 @@ func getStepNames(steps []core.Step) []string { return names } +// validateSubAgents checks sub-agent definitions for common issues +func validateSubAgents(step core.Step, stepIdx int, wast *WorkflowAST, r *EmptyStepRule) []LintIssue { + var issues []LintIssue + if len(step.SubAgents) == 0 { + return issues + } + + seenNames := make(map[string]bool) + for j, sa := range step.SubAgents { + // Error: sub-agent missing name + if sa.Name == "" { + line, col := wast.FindStepFieldPosition(stepIdx, "sub_agents") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityWarning, + Message: fmt.Sprintf("Step '%s' sub-agent at index %d is missing required 'name' field", step.Name, j), + Suggestion: "Add a 'name' field to the sub-agent definition", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].sub_agents[%d].name", stepIdx, j), + }) + } + + // Warning: sub-agent missing description + if sa.Name != "" && sa.Description == "" { + line, col := wast.FindStepFieldPosition(stepIdx, "sub_agents") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityInfo, + Message: fmt.Sprintf("Step '%s' sub-agent '%s' has no description", step.Name, sa.Name), + Suggestion: "Add a description so the LLM knows when to delegate to this agent", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].sub_agents[%d].description", stepIdx, j), + }) + } + + // Error: duplicate sub-agent names + if sa.Name != "" { + if seenNames[sa.Name] { + line, col := wast.FindStepFieldPosition(stepIdx, "sub_agents") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityError, + Message: fmt.Sprintf("Step '%s' has duplicate sub-agent name '%s'", step.Name, sa.Name), + Suggestion: "Use unique names for each sub-agent", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].sub_agents[%d].name", stepIdx, j), + }) + } + seenNames[sa.Name] = true + } + + // Warning: sub-agent with no agent_tools + if len(sa.AgentTools) == 0 { + line, col := wast.FindStepFieldPosition(stepIdx, "sub_agents") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityInfo, + Message: fmt.Sprintf("Step '%s' sub-agent '%s' has no agent_tools", step.Name, sa.Name), + Suggestion: "Add agent_tools so the sub-agent can perform actions", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].sub_agents[%d].agent_tools", stepIdx, j), + }) + } + + // Warning: unknown preset tool in sub-agent + for k, tool := range sa.AgentTools { + if tool.IsPreset() { + if _, ok := core.GetPresetTool(tool.Preset); !ok { + line, col := wast.FindStepFieldPosition(stepIdx, "sub_agents") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityWarning, + Message: fmt.Sprintf("Step '%s' sub-agent '%s' references unknown preset tool '%s'", step.Name, sa.Name, tool.Preset), + Suggestion: "Check the preset tool name against the preset tool registry", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].sub_agents[%d].agent_tools[%d].preset", stepIdx, j, k), + }) + } + } + } + } + + // Info: deep nesting warning + maxDepth := countSubAgentDepth(step.SubAgents) + if maxDepth > core.DefaultMaxAgentDepth { + line, col := wast.FindStepFieldPosition(stepIdx, "sub_agents") + issues = append(issues, LintIssue{ + Rule: r.Name(), + Severity: SeverityInfo, + Message: fmt.Sprintf("Step '%s' has sub-agent nesting depth of %d (default limit is %d)", step.Name, maxDepth, core.DefaultMaxAgentDepth), + Suggestion: "Consider increasing max_agent_depth or reducing nesting", + Line: line, + Column: col, + Field: fmt.Sprintf("steps[%d].sub_agents", stepIdx), + }) + } + + return issues +} + +// countSubAgentDepth returns the maximum nesting depth of sub-agents +func countSubAgentDepth(subAgents []core.SubAgentDef) int { + if len(subAgents) == 0 { + return 0 + } + maxChild := 0 + for _, sa := range subAgents { + childDepth := countSubAgentDepth(sa.SubAgents) + if childDepth > maxChild { + maxChild = childDepth + } + } + return 1 + maxChild +} + // GetDefaultRules returns all built-in linting rules func GetDefaultRules() []LinterRule { return []LinterRule{ diff --git a/internal/parser/parser.go b/internal/parser/parser.go index 1afd091..77a4c8d 100644 --- a/internal/parser/parser.go +++ b/internal/parser/parser.go @@ -362,6 +362,84 @@ func (p *Parser) validateStep(step *core.Step, index int) error { } } } + case core.StepTypeAgent: + // Validate agent step has query or queries (not both) + if step.Query == "" && len(step.Queries) == 0 { + return &ValidationError{ + Field: fmt.Sprintf("steps[%d].query", index), + Message: "agent step must have 'query' or 'queries'", + } + } + if step.Query != "" && len(step.Queries) > 0 { + return &ValidationError{ + Field: fmt.Sprintf("steps[%d]", index), + Message: "agent step cannot have both 'query' and 'queries'", + } + } + // Validate queries are not empty strings + for i, q := range step.Queries { + if q == "" { + return &ValidationError{ + Field: fmt.Sprintf("steps[%d].queries[%d]", index, i), + Message: "query in queries list must not be empty", + } + } + } + // Validate agent step has max_iterations + if step.MaxIterations <= 0 { + return &ValidationError{ + Field: fmt.Sprintf("steps[%d].max_iterations", index), + Message: "agent step must have max_iterations > 0", + } + } + // Validate agent step has agent_tools + if len(step.AgentTools) == 0 { + return &ValidationError{ + Field: fmt.Sprintf("steps[%d].agent_tools", index), + Message: "agent step must have at least one agent_tool", + } + } + // Validate individual agent tools + seen := make(map[string]bool) + for i, tool := range step.AgentTools { + if tool.IsPreset() { + // Validate preset tool exists + if _, ok := core.GetPresetTool(tool.Preset); !ok { + return &ValidationError{ + Field: fmt.Sprintf("steps[%d].agent_tools[%d].preset", index, i), + Message: fmt.Sprintf("unknown preset tool: %s", tool.Preset), + } + } + if seen[tool.Preset] { + return &ValidationError{ + Field: fmt.Sprintf("steps[%d].agent_tools[%d].preset", index, i), + Message: fmt.Sprintf("duplicate tool name: %s", tool.Preset), + } + } + seen[tool.Preset] = true + } else { + // Custom tool must have name and description + if tool.Name == "" { + return &ValidationError{ + Field: fmt.Sprintf("steps[%d].agent_tools[%d].name", index, i), + Message: "custom agent tool requires 'name' field", + } + } + if tool.Description == "" { + return &ValidationError{ + Field: fmt.Sprintf("steps[%d].agent_tools[%d].description", index, i), + Message: "custom agent tool requires 'description' field", + } + } + if seen[tool.Name] { + return &ValidationError{ + Field: fmt.Sprintf("steps[%d].agent_tools[%d].name", index, i), + Message: fmt.Sprintf("duplicate tool name: %s", tool.Name), + } + } + seen[tool.Name] = true + } + } default: return &ValidationError{ Field: fmt.Sprintf("steps[%d].type", index), diff --git a/internal/terminal/printer.go b/internal/terminal/printer.go index ed7750b..c0df382 100644 --- a/internal/terminal/printer.go +++ b/internal/terminal/printer.go @@ -150,7 +150,7 @@ func (p *Printer) StepSuccessWithCommand(stepName, typeSymbol, duration, command }) return } - _, _ = fmt.Fprintf(os.Stdout, "%s %s %s %s\n", StepSuccessSymbol(), typeSymbol, Gray("(finished in "+duration+")"), HiBlue(stepName)) + _, _ = fmt.Fprintf(os.Stdout, "%s %s %s %s\n", StepSuccessSymbol(), typeSymbol, Gray("(finished in ")+Magenta(duration)+Gray(")"), HiBlue(stepName)) // Command already shown when step started, no need to repeat } @@ -378,6 +378,8 @@ func TypeBadge(typ string) string { return Blue(typ) case "foreach": return Blue(typ) + case "agent": + return HiMagenta(typ) default: return typ } diff --git a/internal/terminal/symbols.go b/internal/terminal/symbols.go index a1165c6..d2eb730 100644 --- a/internal/terminal/symbols.go +++ b/internal/terminal/symbols.go @@ -19,6 +19,7 @@ const ( SymbolFunction = "ƒ" // Function step SymbolBash = "$" // Bash/command step SymbolForeach = "∀" // Foreach step (universal quantifier) + SymbolAgent = "⚙" // Agent step (agentic loop) SymbolDocker = "🐋" // Docker runner SymbolSSH = "❄" // SSH runner @@ -150,6 +151,8 @@ func StepTypeSymbol(stepType, runnerType string) string { switch stepType { case "llm": return Magenta(SymbolBowtie) + case "agent": + return HiMagenta(SymbolAgent) case "function": return Cyan(SymbolFunction) case "remote-bash": @@ -168,6 +171,8 @@ func StepCommandPrefix(stepType string) string { switch stepType { case "llm": return SymbolBowtie // "⋈" + case "agent": + return SymbolAgent // "⚙" case "function": return SymbolFunction // "ƒ" case "foreach": diff --git a/pkg/cli/root.go b/pkg/cli/root.go index eabcd9a..ceb7cd7 100644 --- a/pkg/cli/root.go +++ b/pkg/cli/root.go @@ -456,13 +456,14 @@ var versionCmd = &cobra.Command{ if globalJSON { // JSON output versionInfo := map[string]string{ - "name": core.BINARY, - "description": core.DESC, - "version": core.VERSION, - "build": buildTime, - "commit": commitHash, - "author": core.AUTHOR, - "docs": core.DOCS, + "name": core.BINARY, + "description": core.DESC, + "version": core.VERSION, + "build": buildTime, + "commit": commitHash, + "author": core.AUTHOR, + "docs": core.DOCS, + "generated_at": time.Now().Format(time.RFC3339), } jsonOut, _ := json.MarshalIndent(versionInfo, "", " ") fmt.Println(string(jsonOut)) diff --git a/test/e2e/agent_test.go b/test/e2e/agent_test.go new file mode 100644 index 0000000..4293e8d --- /dev/null +++ b/test/e2e/agent_test.go @@ -0,0 +1,709 @@ +package e2e + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRun_AgentModule_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent module in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + log.Info("Using workflow path: %s", workflowPath) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent", "-t", "test.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting stdout contains DRY-RUN indicator and agent steps") + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent") + + log.Success("agent module dry-run works correctly") +} + +func TestRun_AgentMinimalModule_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing minimal agent module in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + log.Info("Using workflow path: %s", workflowPath) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-minimal", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting stdout contains DRY-RUN indicator") + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-minimal") + + log.Success("minimal agent module dry-run works correctly") +} + +func TestRun_AgentModule_WorkflowValidate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent workflow validation") + + workflowPath := getAgentTestdataPath(t) + log.Info("Validating test-agent workflow") + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting validation passes") + assert.Contains(t, stdout, "test-agent") + + log.Success("agent workflow validates successfully") +} + +func TestRun_AgentMinimalModule_WorkflowValidate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing minimal agent workflow validation") + + workflowPath := getAgentTestdataPath(t) + log.Info("Validating test-agent-minimal workflow") + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-minimal", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting validation passes") + assert.Contains(t, stdout, "test-agent-minimal") + + log.Success("minimal agent workflow validates successfully") +} + +func TestRun_AgentModule_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent workflow show") + + workflowPath := getAgentTestdataPath(t) + log.Info("Showing test-agent workflow details") + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting workflow details contain agent step information") + assert.Contains(t, stdout, "test-agent") + assert.Contains(t, stdout, "agent") + + log.Success("agent workflow show works correctly") +} + +// --- Custom tools workflow tests --- + +func TestRun_AgentCustomTools_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent custom tools workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-custom-tools", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting dry-run mode and workflow name") + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-custom-tools") + + log.Success("agent custom tools dry-run works correctly") +} + +func TestRun_AgentCustomTools_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent custom tools workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-custom-tools", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting validation passes") + assert.Contains(t, stdout, "test-agent-custom-tools") + + log.Success("agent custom tools workflow validates successfully") +} + +func TestRun_AgentCustomTools_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent custom tools workflow show") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent-custom-tools", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting workflow show contains agent steps and bash step") + assert.Contains(t, stdout, "agent-custom-handler") + assert.Contains(t, stdout, "verify-custom-tool") + assert.Contains(t, stdout, "agent") + assert.Contains(t, stdout, "bash") + + log.Success("agent custom tools workflow show works correctly") +} + +// --- Exports workflow tests --- + +func TestRun_AgentExports_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent exports workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-exports", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting dry-run mode and workflow name") + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-exports") + + log.Success("agent exports dry-run works correctly") +} + +func TestRun_AgentExports_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent exports workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-exports", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting validation passes") + assert.Contains(t, stdout, "test-agent-exports") + + log.Success("agent exports workflow validates successfully") +} + +func TestRun_AgentExports_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent exports workflow show") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent-exports", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting workflow show contains agent and verify steps") + assert.Contains(t, stdout, "agent-with-exports") + assert.Contains(t, stdout, "verify-exports") + assert.Contains(t, stdout, "agent") + + log.Success("agent exports workflow show works correctly") +} + +// --- Validation failure tests --- + +func TestRun_AgentValidationFail_DuplicateTools(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent workflow validation fails for duplicate tool names") + + workflowPath := getAgentTestdataPath(t) + + _, stderr, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-validation-fail", "-F", workflowPath) + + log.Info("Asserting validation returns error") + assert.Error(t, err) + + log.Info("Asserting error mentions duplicate tool") + assert.Contains(t, stderr, "duplicate") + + log.Success("agent workflow correctly rejects duplicate tool names") +} + +func TestRun_AgentValidationFail_UnknownPreset(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent workflow validation fails for unknown preset tool") + + workflowPath := getAgentTestdataPath(t) + + _, stderr, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-unknown-preset", "-F", workflowPath) + + log.Info("Asserting validation returns error") + assert.Error(t, err) + + log.Info("Asserting error mentions unknown preset") + assert.Contains(t, stderr, "unknown preset tool") + + log.Success("agent workflow correctly rejects unknown preset tools") +} + +// --- Planning workflow tests --- + +func TestRun_AgentPlanning_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent planning workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-planning", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting dry-run mode and workflow name") + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-planning") + + log.Success("agent planning dry-run works correctly") +} + +func TestRun_AgentPlanning_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent planning workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-planning", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting validation passes") + assert.Contains(t, stdout, "test-agent-planning") + + log.Success("agent planning workflow validates successfully") +} + +func TestRun_AgentPlanning_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent planning workflow show") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent-planning", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting workflow show contains planning agent steps") + assert.Contains(t, stdout, "agent-with-plan") + assert.Contains(t, stdout, "verify-plan") + assert.Contains(t, stdout, "agent") + + log.Success("agent planning workflow show works correctly") +} + +// --- Multi-goal workflow tests --- + +func TestRun_AgentMultiGoal_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent multi-goal workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-multi-goal", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting dry-run mode and workflow name") + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-multi-goal") + + log.Success("agent multi-goal dry-run works correctly") +} + +func TestRun_AgentMultiGoal_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent multi-goal workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-multi-goal", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting validation passes") + assert.Contains(t, stdout, "test-agent-multi-goal") + + log.Success("agent multi-goal workflow validates successfully") +} + +func TestRun_AgentMultiGoal_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent multi-goal workflow show") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent-multi-goal", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting workflow show contains multi-goal steps") + assert.Contains(t, stdout, "multi-goal-agent") + assert.Contains(t, stdout, "verify-goals") + assert.Contains(t, stdout, "agent") + + log.Success("agent multi-goal workflow show works correctly") +} + +// --- Structured output workflow tests --- + +func TestRun_AgentStructured_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent structured output workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-structured", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting dry-run mode and workflow name") + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-structured") + + log.Success("agent structured output dry-run works correctly") +} + +func TestRun_AgentStructured_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent structured output workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-structured", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting validation passes") + assert.Contains(t, stdout, "test-agent-structured") + + log.Success("agent structured output workflow validates successfully") +} + +// --- Tracing hooks workflow tests --- + +func TestRun_AgentTracing_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent tracing hooks workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-tracing", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting dry-run mode and workflow name") + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-tracing") + + log.Success("agent tracing dry-run works correctly") +} + +func TestRun_AgentTracing_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent tracing hooks workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-tracing", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting validation passes") + assert.Contains(t, stdout, "test-agent-tracing") + + log.Success("agent tracing workflow validates successfully") +} + +func TestRun_AgentTracing_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent tracing hooks workflow show") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent-tracing", "-F", workflowPath) + require.NoError(t, err) + + log.Info("Asserting workflow show contains tracing agent steps") + assert.Contains(t, stdout, "traced-agent") + assert.Contains(t, stdout, "agent") + + log.Success("agent tracing workflow show works correctly") +} + +// --- Validation failure tests --- + +// --- File tools workflow tests --- + +func TestRun_AgentFileTools_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent file tools workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-file-tools", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-file-tools") + + log.Success("agent file tools dry-run works correctly") +} + +func TestRun_AgentFileTools_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent file tools workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-file-tools", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "test-agent-file-tools") + + log.Success("agent file tools workflow validates successfully") +} + +func TestRun_AgentFileTools_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent file tools workflow show") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent-file-tools", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "agent-file-tools") + assert.Contains(t, stdout, "verify-file-tools") + assert.Contains(t, stdout, "agent") + + log.Success("agent file tools workflow show works correctly") +} + +// --- Orchestration workflow tests --- + +func TestRun_AgentOrchestration_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent orchestration workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-orchestration", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-orchestration") + + log.Success("agent orchestration dry-run works correctly") +} + +func TestRun_AgentOrchestration_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent orchestration workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-orchestration", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "test-agent-orchestration") + + log.Success("agent orchestration workflow validates successfully") +} + +func TestRun_AgentOrchestration_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent orchestration workflow show") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent-orchestration", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "agent-orchestrator") + assert.Contains(t, stdout, "verify-orchestration") + assert.Contains(t, stdout, "agent") + + log.Success("agent orchestration workflow show works correctly") +} + +// --- Python tools workflow tests --- + +func TestRun_AgentPythonTools_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent Python tools workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-python-tools", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-python-tools") + + log.Success("agent Python tools dry-run works correctly") +} + +func TestRun_AgentPythonTools_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent Python tools workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-python-tools", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "test-agent-python-tools") + + log.Success("agent Python tools workflow validates successfully") +} + +func TestRun_AgentPythonTools_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent Python tools workflow show") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent-python-tools", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "agent-python") + assert.Contains(t, stdout, "verify-python-tools") + assert.Contains(t, stdout, "agent") + + log.Success("agent Python tools workflow show works correctly") +} + +// --- Sub-agents workflow tests --- + +func TestRun_AgentSubAgents_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent sub-agents workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-sub-agents", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-sub-agents") + + log.Success("agent sub-agents dry-run works correctly") +} + +func TestRun_AgentSubAgents_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent sub-agents workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-sub-agents", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "test-agent-sub-agents") + + log.Success("agent sub-agents workflow validates successfully") +} + +func TestRun_AgentSubAgents_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent sub-agents workflow show") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent-sub-agents", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "orchestrator") + assert.Contains(t, stdout, "agent") + + log.Success("agent sub-agents workflow show works correctly") +} + +// --- Nested sub-agents workflow tests --- + +func TestRun_AgentSubAgentsNested_DryRun(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing nested sub-agents workflow in dry-run mode") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "run", "-m", "test-agent-sub-agents-nested", "-t", "example.com", "--dry-run", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "DRY-RUN") + assert.Contains(t, stdout, "test-agent-sub-agents-nested") + + log.Success("nested sub-agents dry-run works correctly") +} + +func TestRun_AgentSubAgentsNested_Validate(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing nested sub-agents workflow validation") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-sub-agents-nested", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "test-agent-sub-agents-nested") + + log.Success("nested sub-agents workflow validates successfully") +} + +func TestRun_AgentSubAgentsNested_WorkflowShow(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing nested sub-agents workflow show") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "show", "test-agent-sub-agents-nested", "-F", workflowPath) + require.NoError(t, err) + + assert.Contains(t, stdout, "nested-orchestrator") + assert.Contains(t, stdout, "agent") + + log.Success("nested sub-agents workflow show works correctly") +} + +// --- Sub-agents validation failure test (duplicate names) --- + +func TestRun_AgentSubAgentsValidationFail_DuplicateNames(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing sub-agent validation fails for duplicate sub-agent names") + + workflowPath := getAgentTestdataPath(t) + + stdout, _, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-sub-agents-validation-fail", "-F", workflowPath, "--check") + + log.Info("Asserting validation returns error") + assert.Error(t, err) + + log.Info("Asserting output mentions duplicate sub-agent") + assert.Contains(t, stdout, "duplicate") + + log.Success("agent workflow correctly rejects duplicate sub-agent names") +} + +// --- Validation failure tests --- + +func TestRun_AgentValidationFail_MissingDescription(t *testing.T) { + log := NewTestLogger(t) + log.Step("Testing agent validation fails for custom tool missing description") + + tmpDir := t.TempDir() + + invalidYAML := `kind: module +name: test-agent-no-desc +description: Agent with custom tool missing description +params: + - name: target + required: true + default: example.com +steps: + - name: bad-agent + type: agent + query: "test query" + max_iterations: 3 + agent_tools: + - name: my_tool + parameters: + type: object + properties: + input: + type: string + handler: 'args.input' +` + err := os.WriteFile(filepath.Join(tmpDir, "test-agent-no-desc.yaml"), []byte(invalidYAML), 0644) + require.NoError(t, err) + + _, stderr, err := runCLIWithLog(t, log, "workflow", "validate", "test-agent-no-desc", "-F", tmpDir) + + log.Info("Asserting validation returns error") + assert.Error(t, err) + + log.Info("Asserting error mentions missing description") + assert.Contains(t, stderr, "description") + + log.Success("agent workflow correctly rejects custom tools without description") +} diff --git a/test/e2e/e2e_test.go b/test/e2e/e2e_test.go index 7090e7c..03a5550 100644 --- a/test/e2e/e2e_test.go +++ b/test/e2e/e2e_test.go @@ -71,6 +71,12 @@ func getTestdataPath(t *testing.T) string { return filepath.Join(filepath.Dir(filename), "..", "testdata", "workflows") } +// getAgentTestdataPath returns the path to agent/LLM test workflow fixtures +func getAgentTestdataPath(t *testing.T) string { + t.Helper() + return filepath.Join(getTestdataPath(t), "agent-and-llm") +} + // TestLogger provides verbose logging for E2E tests matching project style type TestLogger struct { t *testing.T diff --git a/test/integration/workflow_test.go b/test/integration/workflow_test.go index 14db029..4b91e17 100644 --- a/test/integration/workflow_test.go +++ b/test/integration/workflow_test.go @@ -47,15 +47,20 @@ func TestLoadAllWorkflows(t *testing.T) { // Skip files that use experimental features or have validation issues skipFiles := map[string]string{ - "test-remote-bash.yaml": "uses remote-bash step type (requires Docker)", - "test-remote-bash-ssh.yaml": "uses remote-bash step type (requires SSH)", - "test-remote-bash-docker.yaml": "uses remote-bash step type (requires Docker)", - "test-docker-file-outputs.yaml": "uses remote-bash step type (requires Docker)", + "test-remote-bash.yaml": "uses remote-bash step type (requires Docker)", + "test-remote-bash-ssh.yaml": "uses remote-bash step type (requires SSH)", + "test-remote-bash-docker.yaml": "uses remote-bash step type (requires Docker)", + "test-docker-file-outputs.yaml": "uses remote-bash step type (requires Docker)", + "test-agent-validation-fail.yaml": "intentionally invalid (duplicate agent tools)", + "test-agent-unknown-preset.yaml": "intentionally invalid (unknown preset tool)", } - // Get all workflow files + // Get all workflow files (top-level + agent-and-llm subdirectory) files, err := filepath.Glob(filepath.Join(workflowsPath, "*.yaml")) require.NoError(t, err) + subFiles, err := filepath.Glob(filepath.Join(workflowsPath, "agent-and-llm", "*.yaml")) + require.NoError(t, err) + files = append(files, subFiles...) require.Greater(t, len(files), 0, "No workflow files found") t.Logf("Found %d workflow files to load", len(files)) @@ -81,17 +86,22 @@ func TestLoadAllWorkflows(t *testing.T) { func TestValidateAllWorkflows(t *testing.T) { workflowsPath := getWorkflowsPath() - // Get all workflow files + // Get all workflow files (top-level + agent-and-llm subdirectory) files, err := filepath.Glob(filepath.Join(workflowsPath, "*.yaml")) require.NoError(t, err) + subFiles, err := filepath.Glob(filepath.Join(workflowsPath, "agent-and-llm", "*.yaml")) + require.NoError(t, err) + files = append(files, subFiles...) // Skip validation test files that are meant to fail or use experimental features skipFiles := map[string]string{ - "test-requirements-fail.yaml": "expected to fail validation", - "test-remote-bash.yaml": "uses remote-bash step type", - "test-remote-bash-ssh.yaml": "uses remote-bash step type", - "test-remote-bash-docker.yaml": "uses remote-bash step type", - "test-docker-file-outputs.yaml": "uses remote-bash step type (requires Docker)", + "test-requirements-fail.yaml": "expected to fail validation", + "test-remote-bash.yaml": "uses remote-bash step type", + "test-remote-bash-ssh.yaml": "uses remote-bash step type", + "test-remote-bash-docker.yaml": "uses remote-bash step type", + "test-docker-file-outputs.yaml": "uses remote-bash step type (requires Docker)", + "test-agent-validation-fail.yaml": "intentionally invalid (duplicate agent tools)", + "test-agent-unknown-preset.yaml": "intentionally invalid (unknown preset tool)", } for _, file := range files { diff --git a/test/testdata/sample-jsonl-output/semgrep-data.json b/test/testdata/sample-jsonl-output/semgrep-data.json new file mode 100644 index 0000000..b2085ef --- /dev/null +++ b/test/testdata/sample-jsonl-output/semgrep-data.json @@ -0,0 +1 @@ +{"version":"1.151.0","results":[{"check_id":"dockerfile.security.missing-user-entrypoint.missing-user-entrypoint","path":"Dockerfile","start":{"line":16,"col":1,"offset":429},"end":{"line":16,"col":22,"offset":450},"extra":{"metavars":{"$...VARS":{"start":{"line":16,"col":12,"offset":440},"end":{"line":16,"col":22,"offset":450},"abstract_content":"[\"python\"]"}},"message":"By not specifying a USER, a program in the container may run as 'root'. This is a security hazard. If an attacker can control a process running as root, they may have control over the container. Ensure that the last USER in a Dockerfile is a USER other than 'root'.","fix":"USER non-root\nENTRYPOINT [\"python\"]","metadata":{"cwe":["CWE-269: Improper Privilege Management"],"category":"security","technology":["dockerfile"],"confidence":"MEDIUM","owasp":["A04:2021 - Insecure Design","A06:2025 - Insecure Design"],"references":["https://owasp.org/Top10/A04_2021-Insecure_Design"],"subcategory":["audit"],"likelihood":"LOW","impact":"MEDIUM","license":"Semgrep Rules License v1.0. For more details, visit semgrep.dev/legal/rules-license","vulnerability_class":["Improper Authorization"],"source":"https://semgrep.dev/r/dockerfile.security.missing-user-entrypoint.missing-user-entrypoint","shortlink":"https://sg.run/k281","semgrep.dev":{"rule":{"origin":"community","r_id":47272,"rule_id":"ReUW9E","rv_id":1262659,"url":"https://semgrep.dev/playground/r/o5TbD21/dockerfile.security.missing-user-entrypoint.missing-user-entrypoint","version_id":"o5TbD21"}}},"severity":"ERROR","fingerprint":"9ebfc0d726acf772708e18784a26daaf7d2d3b542c4ef4c465778b02a5312329de5ef18d4f8974fe4819eb9db1645a0157a09fb0d5a3f370f6459b964f0b01a7_0","lines":"ENTRYPOINT [\"python\"]","is_ignored":false,"validation_state":"NO_VALIDATOR","engine_kind":"OSS"}},{"check_id":"dockerfile.security.missing-user.missing-user","path":"Dockerfile","start":{"line":17,"col":1,"offset":451},"end":{"line":17,"col":15,"offset":465},"extra":{"metavars":{"$...VARS":{"start":{"line":17,"col":5,"offset":455},"end":{"line":17,"col":15,"offset":465},"abstract_content":"[\"app.py\"]"}},"message":"By not specifying a USER, a program in the container may run as 'root'. This is a security hazard. If an attacker can control a process running as root, they may have control over the container. Ensure that the last USER in a Dockerfile is a USER other than 'root'.","fix":"USER non-root\nCMD [\"app.py\"]","metadata":{"cwe":["CWE-250: Execution with Unnecessary Privileges"],"category":"security","technology":["dockerfile"],"confidence":"MEDIUM","owasp":["A04:2021 - Insecure Design","A06:2025 - Insecure Design"],"references":["https://owasp.org/Top10/A04_2021-Insecure_Design"],"subcategory":["audit"],"likelihood":"LOW","impact":"MEDIUM","license":"Semgrep Rules License v1.0. For more details, visit semgrep.dev/legal/rules-license","vulnerability_class":["Improper Authorization"],"source":"https://semgrep.dev/r/dockerfile.security.missing-user.missing-user","shortlink":"https://sg.run/Gbvn","semgrep.dev":{"rule":{"origin":"community","r_id":20148,"rule_id":"AbUN06","rv_id":1262660,"url":"https://semgrep.dev/playground/r/zyTb2n2/dockerfile.security.missing-user.missing-user","version_id":"zyTb2n2"}}},"severity":"ERROR","fingerprint":"b4e6a9c3caa7f197bfaa44479e111925e766af68481dacdab05ecbf56974c149c1ab2681427306574ab31f35f34bb56157a364823b3df9582974ca121c473863_0","lines":"CMD [\"app.py\"]","is_ignored":false,"validation_state":"NO_VALIDATOR","engine_kind":"OSS"}},{"check_id":"python.flask.security.audit.hardcoded-config.avoid_hardcoded_config_SECRET_KEY","path":"config.py","start":{"line":13,"col":1,"offset":455},"end":{"line":13,"col":45,"offset":499},"extra":{"metavars":{"$M":{"start":{"line":13,"col":1,"offset":455},"end":{"line":13,"col":13,"offset":467},"abstract_content":"vuln_app.app"}},"message":"Hardcoded variable `SECRET_KEY` detected. Use environment variables or config files instead","metadata":{"likelihood":"LOW","impact":"LOW","confidence":"LOW","category":"security","cwe":["CWE-489: Active Debug Code"],"owasp":["A05:2021 - Security Misconfiguration","A02:2025 - Security Misconfiguration"],"references":["https://bento.dev/checks/flask/avoid-hardcoded-config/","https://flask.palletsprojects.com/en/1.1.x/config/?highlight=configuration#builtin-configuration-values","https://flask.palletsprojects.com/en/1.1.x/config/?highlight=configuration#environment-and-debug-features"],"subcategory":["audit"],"technology":["flask"],"license":"Semgrep Rules License v1.0. For more details, visit semgrep.dev/legal/rules-license","vulnerability_class":["Active Debug Code"],"source":"https://semgrep.dev/r/python.flask.security.audit.hardcoded-config.avoid_hardcoded_config_SECRET_KEY","shortlink":"https://sg.run/Ekde","semgrep.dev":{"rule":{"origin":"community","r_id":9537,"rule_id":"4bUkX0","rv_id":1263420,"url":"https://semgrep.dev/playground/r/l4TJRA9/python.flask.security.audit.hardcoded-config.avoid_hardcoded_config_SECRET_KEY","version_id":"l4TJRA9"}}},"severity":"ERROR","fingerprint":"a6697907ce3488675b4572624002d6856a5ba4e3908eabb52b17e83fb147f3437af6da4e8b52b2073cbe7ab2acac6e2b656aa2a363d2d8744d715e66a0c6b660_0","lines":"vuln_app.app.config['SECRET_KEY'] = 'random'","is_ignored":false,"validation_state":"NO_VALIDATOR","engine_kind":"OSS"}},{"check_id":"generic.secrets.security.detected-jwt-token.detected-jwt-token","path":"openapi_specs/openapi3.yml","start":{"line":193,"col":33,"offset":5902},"end":{"line":193,"col":141,"offset":6010},"extra":{"metavars":{},"message":"JWT token detected","metadata":{"source-rule-url":"https://github.com/Yelp/detect-secrets/blob/master/detect_secrets/plugins/jwt.py","category":"security","technology":["secrets","jwt"],"confidence":"LOW","references":["https://semgrep.dev/blog/2020/hardcoded-secrets-unverified-tokens-and-other-common-jwt-mistakes/"],"cwe":["CWE-321: Use of Hard-coded Cryptographic Key"],"owasp":["A02:2021 - Cryptographic Failures","A04:2025 - Cryptographic Failures"],"subcategory":["audit"],"likelihood":"LOW","impact":"MEDIUM","license":"Semgrep Rules License v1.0. For more details, visit semgrep.dev/legal/rules-license","vulnerability_class":["Cryptographic Issues"],"source":"https://semgrep.dev/r/generic.secrets.security.detected-jwt-token.detected-jwt-token","shortlink":"https://sg.run/05N5","semgrep.dev":{"rule":{"origin":"community","r_id":12854,"rule_id":"kxU8E8","rv_id":1262879,"url":"https://semgrep.dev/playground/r/d6Tyxvg/generic.secrets.security.detected-jwt-token.detected-jwt-token","version_id":"d6Tyxvg"}}},"severity":"ERROR","fingerprint":"851d262b7d2f01e95b7d06e6f5135a879b472779a00cc92abf565de9fe04bd3afbd5fa34d7580cdce761c37f35799eb374a3241489e74af86194771f76daee93_0","lines":" example: 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJleHAiOjE2NzAxNjA2MTcsImlhdCI6MTY3MDE2MDU1Nywic3ViIjoiSm9obi5Eb2UifQ.n17N4AxTbL4_z65-NR46meoytauPDjImUxrLiUMSTQw'","is_ignored":false,"validation_state":"NO_VALIDATOR","engine_kind":"OSS"}}],"errors":[],"paths":{"scanned":[".dockerignore",".github/FUNDING.yml",".github/workflows/docker-image.yml",".gitignore","Dockerfile","LICENSE","README.md","api_views/__init__.py","api_views/books.py","api_views/json_schemas.py","api_views/main.py","api_views/users.py","app.py","config.py","database/__init__.py","docker-compose.yaml","models/__init__.py","models/books_model.py","models/user_model.py","openapi_specs/VAmPI.postman_collection.json","openapi_specs/openapi3.yml","requirements.txt"]},"time":{"rules":[],"rules_parse_time":6.271836042404175,"profiling_times":{"config_time":7.641101121902466,"core_time":7.604748964309692,"ignores_time":0.001773834228515625,"total_time":15.253854990005493},"parsing_time":{"total_time":0.0,"per_file_time":{"mean":0.0,"std_dev":0.0},"very_slow_stats":{"time_ratio":0.0,"count_ratio":0.0},"very_slow_files":[]},"scanning_time":{"total_time":2.827584743499756,"per_file_time":{"mean":0.04635384825409437,"std_dev":0.017019160223435907},"very_slow_stats":{"time_ratio":0.0,"count_ratio":0.0},"very_slow_files":[]},"matching_time":{"total_time":0.0,"per_file_and_rule_time":{"mean":0.0,"std_dev":0.0},"very_slow_stats":{"time_ratio":0.0,"count_ratio":0.0},"very_slow_rules_on_files":[]},"tainting_time":{"total_time":0.0,"per_def_and_rule_time":{"mean":0.0,"std_dev":0.0},"very_slow_stats":{"time_ratio":0.0,"count_ratio":0.0},"very_slow_rules_on_defs":[]},"fixpoint_timeouts":[],"prefiltering":{"project_level_time":0.0,"file_level_time":0.0,"rules_with_project_prefilters_ratio":0.0,"rules_with_file_prefilters_ratio":0.9729977116704805,"rules_selected_ratio":0.06453089244851258,"rules_matched_ratio":0.06453089244851258},"targets":[],"total_bytes":0,"max_memory_bytes":9950453248},"engine_requested":"OSS","interfile_languages_used":[],"skipped_rules":[],"profiling_results":[]} \ No newline at end of file diff --git a/test/testdata/workflows/agent-and-llm/test-agent-custom-tools.yaml b/test/testdata/workflows/agent-and-llm/test-agent-custom-tools.yaml new file mode 100644 index 0000000..b5a6ff0 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-custom-tools.yaml @@ -0,0 +1,33 @@ +kind: module +name: test-agent-custom-tools +description: Test agent step with custom tool handlers and template rendering +tags: test,agent + +params: + - name: target + required: true + default: example.com + +steps: + - name: agent-custom-handler + type: agent + query: "Use the classify tool to classify the target '{{Target}}' and report the result." + max_iterations: 3 + agent_tools: + - name: classify + description: "Classify a string as domain, ip, or unknown" + parameters: + type: object + properties: + input: + type: string + description: "The string to classify" + required: + - input + handler: 'contains(args.input, ".") ? "domain" : "unknown"' + exports: + classification: "{{agent_content}}" + + - name: verify-custom-tool + type: bash + command: echo "Classification result for {{Target}} is - {{classification}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-exports.yaml b/test/testdata/workflows/agent-and-llm/test-agent-exports.yaml new file mode 100644 index 0000000..62153ee --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-exports.yaml @@ -0,0 +1,30 @@ +kind: module +name: test-agent-exports +description: Test agent step exports propagation and token tracking +tags: test,agent + +params: + - name: target + required: true + default: example.com + +steps: + - name: agent-with-exports + type: agent + system_prompt: "You are a helpful assistant. Always respond concisely." + query: "Check if the file /tmp/nonexistent-{{Target}} exists using file_exists tool, then report the result." + max_iterations: 3 + agent_tools: + - preset: file_exists + - preset: bash + exports: + agent_result: "{{agent_content}}" + iterations_used: "{{agent_iterations}}" + tokens_used: "{{agent_total_tokens}}" + + - name: verify-exports + type: bash + commands: + - echo "Agent result is - {{agent_result}}" + - echo "Iterations used - {{iterations_used}}" + - echo "Tokens used - {{tokens_used}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-file-tools.yaml b/test/testdata/workflows/agent-and-llm/test-agent-file-tools.yaml new file mode 100644 index 0000000..5881e09 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-file-tools.yaml @@ -0,0 +1,31 @@ +kind: module +name: test-agent-file-tools +description: Test agent with file-related preset tools (glob, read_file, file_exists, save_content, grep_string) +tags: test,agent,file-tools + +params: + - name: target + required: true + default: example.com + +steps: + - name: agent-file-tools + type: agent + log: "Running agent with file tool presets for {{Target}}" + query: "List files matching *.txt in the output directory, check if any exist, read their contents, and save a summary." + system_prompt: "You are a file analysis assistant. Use the provided tools to inspect files." + max_iterations: 5 + agent_tools: + - preset: glob + - preset: read_file + - preset: file_exists + - preset: save_content + - preset: grep_string + - preset: file_length + exports: + file_result: "{{agent_content}}" + + - name: verify-file-tools + type: bash + log: "Verifying file tools agent output" + command: echo "File tools agent result - {{file_result}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-minimal.yaml b/test/testdata/workflows/agent-and-llm/test-agent-minimal.yaml new file mode 100644 index 0000000..09f9296 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-minimal.yaml @@ -0,0 +1,19 @@ +kind: module +name: test-agent-minimal +description: Minimal agent step test +tags: test,agent,quick + +params: + - name: target + required: true + default: example.com + +steps: + - name: quick-agent + type: agent + query: "Run 'echo hello {{Target}}' and tell me the output." + max_iterations: 5 + agent_tools: + - preset: bash + exports: + result: "{{agent_content}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-multi-goal.yaml b/test/testdata/workflows/agent-and-llm/test-agent-multi-goal.yaml new file mode 100644 index 0000000..2d262b3 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-multi-goal.yaml @@ -0,0 +1,28 @@ +kind: module +name: test-agent-multi-goal +description: Test agent with multiple queries executed sequentially +tags: test,agent,multi-goal + +params: + - name: target + required: true + default: example.com + +steps: + - name: multi-goal-agent + type: agent + log: "Running multi-goal agent for {{Target}}" + queries: + - "List files in the current directory" + - "Summarize what you found" + max_iterations: 5 + agent_tools: + - preset: bash + exports: + result: "{{agent_content}}" + goals: "{{agent_goal_results}}" + + - name: verify-goals + type: bash + log: "Verifying multi-goal results" + command: echo "Results - {{result}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-orchestration.yaml b/test/testdata/workflows/agent-and-llm/test-agent-orchestration.yaml new file mode 100644 index 0000000..df2fb5b --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-orchestration.yaml @@ -0,0 +1,29 @@ +kind: module +name: test-agent-orchestration +description: Test agent with run_module and run_flow preset tools +tags: test,agent,orchestration + +params: + - name: target + required: true + default: example.com + +steps: + - name: agent-orchestrator + type: agent + log: "Running orchestration agent for {{Target}}" + query: "Analyze {{Target}} by running appropriate modules and flows." + system_prompt: "You are a workflow orchestrator. Use run_module and run_flow to delegate scanning tasks." + max_iterations: 5 + agent_tools: + - preset: bash + - preset: run_module + - preset: run_flow + - preset: read_file + exports: + orchestration_result: "{{agent_content}}" + + - name: verify-orchestration + type: bash + log: "Verifying orchestration agent output" + command: echo "Orchestration result - {{orchestration_result}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-planning.yaml b/test/testdata/workflows/agent-and-llm/test-agent-planning.yaml new file mode 100644 index 0000000..f168ff3 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-planning.yaml @@ -0,0 +1,28 @@ +kind: module +name: test-agent-planning +description: Test agent with planning stage before execution +tags: test,agent,planning + +params: + - name: target + required: true + default: example.com + +steps: + - name: agent-with-plan + type: agent + log: "Running agent with planning stage for {{Target}}" + plan_prompt: "Create a reconnaissance plan for {{Target}}. List 3 steps." + plan_max_tokens: 500 + query: "Execute the plan you created." + max_iterations: 5 + agent_tools: + - preset: bash + exports: + plan: "{{agent_plan}}" + result: "{{agent_content}}" + + - name: verify-plan + type: bash + log: "Verifying plan was created" + command: echo "Plan - {{plan}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-python-tools.yaml b/test/testdata/workflows/agent-and-llm/test-agent-python-tools.yaml new file mode 100644 index 0000000..827df79 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-python-tools.yaml @@ -0,0 +1,29 @@ +kind: module +name: test-agent-python-tools +description: Test agent with exec_python and exec_python_file preset tools +tags: test,agent,python + +params: + - name: target + required: true + default: example.com + +steps: + - name: agent-python + type: agent + log: "Running Python tools agent for {{Target}}" + query: "Use Python to parse and analyze data about {{Target}}. Write a small Python script to check if the target is a valid domain." + system_prompt: "You are a Python-powered analyst. Use exec_python for inline code and exec_python_file for scripts." + max_iterations: 5 + agent_tools: + - preset: exec_python + - preset: exec_python_file + - preset: save_content + - preset: bash + exports: + python_result: "{{agent_content}}" + + - name: verify-python-tools + type: bash + log: "Verifying Python tools agent output" + command: echo "Python tools result - {{python_result}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-structured.yaml b/test/testdata/workflows/agent-and-llm/test-agent-structured.yaml new file mode 100644 index 0000000..fb4b753 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-structured.yaml @@ -0,0 +1,21 @@ +kind: module +name: test-agent-structured +description: Test agent with structured output schema enforcement +tags: test,agent,structured + +params: + - name: target + required: true + default: example.com + +steps: + - name: structured-agent + type: agent + log: "Running structured output agent for {{Target}}" + query: "Find subdomains of {{Target}} and classify them" + output_schema: '{"type":"object","properties":{"subdomains":{"type":"array","items":{"type":"object","properties":{"name":{"type":"string"},"status":{"type":"string"}}}}}}' + max_iterations: 5 + agent_tools: + - preset: bash + exports: + results: "{{agent_content}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-sub-agents-nested.yaml b/test/testdata/workflows/agent-and-llm/test-agent-sub-agents-nested.yaml new file mode 100644 index 0000000..7367877 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-sub-agents-nested.yaml @@ -0,0 +1,57 @@ +kind: module +name: test-agent-sub-agents-nested +description: Test 3-level nested sub-agent orchestration (parent -> child -> grandchild) +tags: test,agent,sub-agents,nested + +params: + - name: target + required: true + default: example.com + +steps: + - name: nested-orchestrator + type: agent + query: "Coordinate a 3-level analysis of {{Target}} by delegating to specialist agents" + system_prompt: "You are a top-level orchestrator. Delegate tasks to sub-agents for deeper analysis." + max_iterations: 10 + max_agent_depth: 3 + agent_tools: + - preset: bash + sub_agents: + - name: recon_coordinator + description: "Coordinates reconnaissance by delegating to leaf agents" + system_prompt: "You coordinate recon tasks. Delegate to specialists." + max_iterations: 5 + agent_tools: + - preset: bash + - preset: http_get + sub_agents: + - name: port_scanner + description: "Scans ports on a target" + system_prompt: "You are a port scanning specialist" + max_iterations: 3 + agent_tools: + - preset: bash + - name: dns_resolver + description: "Resolves DNS records for a target" + system_prompt: "You are a DNS resolution specialist" + max_iterations: 3 + agent_tools: + - preset: bash + - name: vuln_coordinator + description: "Coordinates vulnerability analysis" + system_prompt: "You coordinate vulnerability scanning tasks." + max_iterations: 5 + agent_tools: + - preset: bash + - preset: read_file + sub_agents: + - name: web_scanner + description: "Scans web applications for vulnerabilities" + system_prompt: "You are a web vulnerability specialist" + max_iterations: 3 + agent_tools: + - preset: bash + - preset: http_get + exports: + nested_result: "{{agent_content}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-sub-agents-validation-fail.yaml b/test/testdata/workflows/agent-and-llm/test-agent-sub-agents-validation-fail.yaml new file mode 100644 index 0000000..923dde6 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-sub-agents-validation-fail.yaml @@ -0,0 +1,28 @@ +kind: module +name: test-agent-sub-agents-validation-fail +description: Agent with duplicate sub-agent names (should fail validation) +tags: test,agent,validation-fail + +params: + - name: target + required: true + default: example.com + +steps: + - name: bad-orchestrator + type: agent + query: "Analyze {{Target}}" + max_iterations: 5 + agent_tools: + - preset: bash + sub_agents: + - name: recon_agent + description: "First recon agent" + max_iterations: 3 + agent_tools: + - preset: bash + - name: recon_agent + description: "Duplicate name - should cause validation error" + max_iterations: 3 + agent_tools: + - preset: bash diff --git a/test/testdata/workflows/agent-and-llm/test-agent-sub-agents.yaml b/test/testdata/workflows/agent-and-llm/test-agent-sub-agents.yaml new file mode 100644 index 0000000..b6f15da --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-sub-agents.yaml @@ -0,0 +1,36 @@ +kind: module +name: test-agent-sub-agents +description: Test multi-agent orchestration with sub-agent spawning +tags: test,agent,sub-agents + +params: + - name: target + required: true + default: example.com + +steps: + - name: orchestrator + type: agent + query: "Analyze {{Target}} by coordinating specialists" + system_prompt: "You are an orchestrator. Delegate tasks to sub-agents." + max_iterations: 10 + max_agent_depth: 3 + agent_tools: + - preset: bash + sub_agents: + - name: recon_agent + description: "Specialized agent for reconnaissance" + system_prompt: "You are a recon specialist" + max_iterations: 5 + agent_tools: + - preset: bash + - preset: http_get + - name: vuln_scanner + description: "Specialized agent for vulnerability scanning" + system_prompt: "You are a vulnerability analyst" + max_iterations: 5 + agent_tools: + - preset: bash + - preset: read_file + exports: + result: "{{agent_content}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-tracing.yaml b/test/testdata/workflows/agent-and-llm/test-agent-tracing.yaml new file mode 100644 index 0000000..630723d --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-tracing.yaml @@ -0,0 +1,22 @@ +kind: module +name: test-agent-tracing +description: Test agent with tool tracing hooks +tags: test,agent,tracing + +params: + - name: target + required: true + default: example.com + +steps: + - name: traced-agent + type: agent + log: "Running traced agent for {{Target}}" + query: "Run 'echo hello' and report what you see" + on_tool_start: 'log_info("Starting tool: " + tool_name)' + on_tool_end: 'log_info("Tool " + tool_name + " completed in " + duration + "ms")' + max_iterations: 5 + agent_tools: + - preset: bash + exports: + result: "{{agent_content}}" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-unknown-preset.yaml b/test/testdata/workflows/agent-and-llm/test-agent-unknown-preset.yaml new file mode 100644 index 0000000..c048ec4 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-unknown-preset.yaml @@ -0,0 +1,21 @@ +kind: module +name: test-agent-unknown-preset +description: Agent workflow with unknown preset tool for validation testing +tags: test,agent,invalid + +params: + - name: target + required: true + default: example.com + +steps: + - name: agent-bad-preset + type: agent + query: "This should fail validation due to unknown preset" + max_iterations: 3 + agent_tools: + - preset: nonexistent_tool + + - name: unreachable + type: bash + command: echo "should not reach here" diff --git a/test/testdata/workflows/agent-and-llm/test-agent-validation-fail.yaml b/test/testdata/workflows/agent-and-llm/test-agent-validation-fail.yaml new file mode 100644 index 0000000..ab7860b --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent-validation-fail.yaml @@ -0,0 +1,22 @@ +kind: module +name: test-agent-validation-fail +description: Agent workflow with invalid configuration for validation testing +tags: test,agent,invalid + +params: + - name: target + required: true + default: example.com + +steps: + - name: agent-duplicate-tools + type: agent + query: "This should fail validation due to duplicate tool names" + max_iterations: 5 + agent_tools: + - preset: bash + - preset: bash + + - name: unreachable + type: bash + command: echo "should not reach here" diff --git a/test/testdata/workflows/agent-and-llm/test-agent.yaml b/test/testdata/workflows/agent-and-llm/test-agent.yaml new file mode 100644 index 0000000..21de1f7 --- /dev/null +++ b/test/testdata/workflows/agent-and-llm/test-agent.yaml @@ -0,0 +1,96 @@ +kind: module +name: test-agent +description: Test agent step execution with tool calling loop +tags: test,agent,llm + +params: + - name: target + required: true + default: example.com + +steps: + # Basic agent with preset tools + - name: basic-agent + type: agent + log: "Running basic agent for {{Target}}" + query: "List the files in the current directory using exec_cmd, then summarize what you found." + max_iterations: 5 + agent_tools: + - preset: bash + - preset: file_exists + exports: + agent_output: "{{agent_content}}" + + # Agent with system prompt and custom tool + - name: agent-with-system-prompt + type: agent + log: "Running agent with system prompt" + system_prompt: "You are a security analyst. Be concise and precise." + query: "Check if the target {{Target}} is a valid domain." + max_iterations: 3 + agent_tools: + - preset: bash + - name: check_domain + description: "Validate if a string is a valid domain name" + parameters: + type: object + properties: + domain: + type: string + description: "Domain to validate" + required: + - domain + handler: 'contains(args.domain, ".")' + exports: + analysis: "{{agent_content}}" + + # Agent with stop condition + - name: agent-with-stop-condition + type: agent + log: "Running agent with stop condition" + query: "Say DONE when you are finished." + max_iterations: 10 + stop_condition: 'contains(agent_content, "DONE")' + agent_tools: + - preset: bash + + # Agent with memory configuration + - name: agent-with-memory + type: agent + log: "Running agent with memory config" + query: "Execute 'echo hello' and report the result." + max_iterations: 5 + memory: + max_messages: 20 + persist_path: "{{Output}}/agent/conversation.json" + agent_tools: + - preset: bash + - preset: save_content + + # Agent with LLM config override + - name: agent-with-config + type: agent + log: "Running agent with LLM config override" + query: "What is 2+2?" + max_iterations: 3 + llm_config: + max_tokens: 100 + temperature: 0.1 + agent_tools: + - preset: bash + + # Agent with parallel tool calls disabled + - name: agent-sequential-tools + type: agent + log: "Running agent with sequential tool calls" + query: "Run 'echo hello' then 'echo world' sequentially." + max_iterations: 5 + parallel_tool_calls: false + agent_tools: + - preset: bash + + # Use previous agent export in next step + - name: use-agent-output + type: bash + log: "Using agent output from previous step" + command: echo "Agent said - {{agent_output}}" diff --git a/test/testdata/workflows/test-llm.yaml b/test/testdata/workflows/agent-and-llm/test-llm.yaml similarity index 100% rename from test/testdata/workflows/test-llm.yaml rename to test/testdata/workflows/agent-and-llm/test-llm.yaml