diff --git a/pkg/tools/execute/code.go b/pkg/tools/execute/code.go index cd11e17..f032c70 100644 --- a/pkg/tools/execute/code.go +++ b/pkg/tools/execute/code.go @@ -8,8 +8,45 @@ import ( "ollie/internal/sandbox" "regexp" "strings" + "sync" ) +// CodeStep is one step in a parallel execute_code call or a parallel pipe stage. +// Exactly one of Code or Tool must be set. +type CodeStep struct { + Code string `json:"code"` + Language string `json:"language"` + Tool string `json:"tool"` + Args []string `json:"args"` +} + +// resolveCodeStep loads a CodeStep into executable (code, language, trusted). +// Tool steps are read from the tools directory and treated as trusted. +// Inline code steps default to bash and are validated by the caller. +func resolveCodeStep(s CodeStep) (code, language string, trusted bool, err error) { + if s.Tool != "" { + toolCode, terr := ReadTool(s.Tool) + if terr != nil { + return "", "", false, terr + } + language = detectLanguage(toolCode) + code = toolCode + if len(s.Args) > 0 { + code = injectArgs(language, s.Tool, s.Args, toolCode) + switch language { + case "awk", "sed", "jq", "ed", "expect", "bc": + language = "bash" + } + } + return code, language, true, nil + } + language = s.Language + if language == "" { + language = "bash" + } + return s.Code, language, false, nil +} + // execute_code is implemented directly by Executor.Execute with trusted=false. // See executor.go for the implementation. @@ -130,41 +167,99 @@ func (lw *limitedWriter) Write(p []byte) (n int, err error) { } func dispatchExecuteCode(ctx context.Context, e *Server, args json.RawMessage) (string, error) { - code, language, sandbox, timeout, err := execArgs(args) + steps, timeout, sandboxName, err := execCodeArgs(args) if err != nil { return "", fmt.Errorf("execute_code: bad args: %w", err) } - if code == "" { - return "", fmt.Errorf("execute_code: 'code' is required") + if len(steps) == 0 { + return "", fmt.Errorf("execute_code: at least one step is required") } - if !e.allowed("execute_code", fmt.Sprintf("execute_code: %s", code)) { + + // Build confirmation label and check upfront (before any goroutines). + stepLabel := func(s CodeStep) string { + if s.Tool != "" { + return fmt.Sprintf("tool:%s %s", s.Tool, strings.Join(s.Args, " ")) + } + return s.Code + } + var label string + if len(steps) == 1 { + label = fmt.Sprintf("execute_code: %s", stepLabel(steps[0])) + } else { + var sb strings.Builder + fmt.Fprintf(&sb, "execute_code: %d parallel steps", len(steps)) + for i, s := range steps { + fmt.Fprintf(&sb, "\n [%d]: %s", i, stepLabel(s)) + } + label = sb.String() + } + if !e.allowed("execute_code", label) { return "", fmt.Errorf("execute_code: denied by user") } - return e.Execute(ctx, code, language, timeout, sandbox, false) + + if len(steps) == 1 { + code, lang, trusted, err := resolveCodeStep(steps[0]) + if err != nil { + return "", err + } + return e.executeWithStdin(ctx, code, lang, timeout, sandboxName, trusted, "") + } + + // Parallel fan-out: all steps run concurrently; results collected in submission order. + type result struct { + output string + err error + } + results := make([]result, len(steps)) + var wg sync.WaitGroup + for i, step := range steps { + wg.Add(1) + go func(idx int, s CodeStep) { + defer wg.Done() + code, lang, trusted, err := resolveCodeStep(s) + if err != nil { + results[idx] = result{err: err} + return + } + out, err := e.executeWithStdin(ctx, code, lang, timeout, sandboxName, trusted, "") + results[idx] = result{output: out, err: err} + }(i, step) + } + wg.Wait() + + var sb strings.Builder + var firstErr error + for i, r := range results { + fmt.Fprintf(&sb, "=== step %d ===\n", i) + sb.WriteString(r.output) + if r.err != nil { + fmt.Fprintf(&sb, "\nerror: %v\n", r.err) + if firstErr == nil { + firstErr = r.err + } + } + sb.WriteByte('\n') + } + return sb.String(), firstErr } -func execArgs(args json.RawMessage) (code, language, sandbox string, timeout int, err error) { +func execCodeArgs(args json.RawMessage) (steps []CodeStep, timeout int, sandboxName string, err error) { var a struct { - Code string `json:"code"` - Language string `json:"language"` - Timeout int `json:"timeout"` - Sandbox string `json:"sandbox"` + Steps []CodeStep `json:"steps"` + Timeout int `json:"timeout"` + Sandbox string `json:"sandbox"` } if err = json.Unmarshal(args, &a); err != nil { return } - code = a.Code - language = a.Language - if language == "" { - language = "bash" - } + steps = a.Steps timeout = a.Timeout if timeout <= 0 { timeout = 30 } - sandbox = a.Sandbox - if sandbox == "" { - sandbox = "default" + sandboxName = a.Sandbox + if sandboxName == "" { + sandboxName = "default" } return } diff --git a/pkg/tools/execute/pipe.go b/pkg/tools/execute/pipe.go index f51aa8f..25ad772 100644 --- a/pkg/tools/execute/pipe.go +++ b/pkg/tools/execute/pipe.go @@ -8,14 +8,16 @@ import ( "encoding/json" "fmt" "strings" + "sync" ) // PipeStep is one stage in a tool pipeline. -// Exactly one of Tool or Code must be set. +// Exactly one of Tool, Code, or Parallel must be set. type PipeStep struct { - Tool string // named tool read from 9P (trusted) - Code string // inline bash code (untrusted, validated) - Args []string + Tool string // named tool read from 9P (trusted) + Code string // inline bash code (untrusted, validated) + Args []string + Parallel []CodeStep // parallel fan-out; outputs concatenated in submission order } // BuildPipeline constructs a single bash pipeline string from the given steps. @@ -75,9 +77,14 @@ func BuildPipeline(steps []PipeStep) (string, bool, error) { func dispatchExecutePipe(ctx context.Context, e *Server, args json.RawMessage) (string, error) { var a struct { - Pipe []PipeStep `json:"pipe"` - Timeout int `json:"timeout"` - Sandbox string `json:"sandbox"` + Pipe []struct { + Tool string `json:"tool"` + Args []string `json:"args"` + Code string `json:"code"` + Parallel []CodeStep `json:"parallel"` + } `json:"pipe"` + Timeout int `json:"timeout"` + Sandbox string `json:"sandbox"` } if err := json.Unmarshal(args, &a); err != nil { return "", fmt.Errorf("execute_pipe: bad args: %w", err) @@ -85,20 +92,94 @@ func dispatchExecutePipe(ctx context.Context, e *Server, args json.RawMessage) ( if len(a.Pipe) == 0 { return "", fmt.Errorf("execute_pipe: 'pipe' is required") } - code, _, err := BuildPipeline(a.Pipe) - if err != nil { - return "", err + + stages := make([]PipeStep, len(a.Pipe)) + for i, s := range a.Pipe { + stages[i] = PipeStep{Tool: s.Tool, Code: s.Code, Args: s.Args, Parallel: s.Parallel} } - if !e.allowed("execute_pipe", fmt.Sprintf("execute_pipe: %s", code)) { + + if !e.allowed("execute_pipe", fmt.Sprintf("execute_pipe: %d stage(s)", len(stages))) { return "", fmt.Errorf("execute_pipe: denied by user") } + timeout := a.Timeout if timeout <= 0 { timeout = 30 } - sandbox := a.Sandbox - if sandbox == "" { - sandbox = "default" + sandboxName := a.Sandbox + if sandboxName == "" { + sandboxName = "default" } - return e.Execute(ctx, code, "bash", timeout, sandbox, true) + + // Execute stages sequentially, feeding each stage's stdout as the next stage's stdin. + var input string + for i, stage := range stages { + out, err := e.runPipeStage(ctx, i, stage, timeout, sandboxName, input) + if err != nil { + return out, fmt.Errorf("pipe stage %d: %w", i, err) + } + input = out + } + return input, nil +} + +// runPipeStage executes one stage of an execute_pipe pipeline, feeding stdinData +// as stdin. Parallel stages fan out concurrently and concatenate in submission order. +func (e *Server) runPipeStage(ctx context.Context, idx int, stage PipeStep, timeout int, sandboxName, stdinData string) (string, error) { + // Parallel stage: fan out, collect in submission order, concatenate. + if len(stage.Parallel) > 0 { + results := make([]string, len(stage.Parallel)) + errs := make([]error, len(stage.Parallel)) + var wg sync.WaitGroup + for i, step := range stage.Parallel { + wg.Add(1) + go func(j int, s CodeStep) { + defer wg.Done() + code, lang, trusted, err := resolveCodeStep(s) + if err != nil { + errs[j] = err + return + } + results[j], errs[j] = e.executeWithStdin(ctx, code, lang, timeout, sandboxName, trusted, stdinData) + }(i, step) + } + wg.Wait() + var sb strings.Builder + for i, out := range results { + sb.WriteString(out) + if errs[i] != nil { + return sb.String(), errs[i] + } + } + return sb.String(), nil + } + + // Tool stage. + if stage.Tool != "" { + toolCode, err := ReadTool(stage.Tool) + if err != nil { + return "", err + } + language := detectLanguage(toolCode) + code := toolCode + if len(stage.Args) > 0 { + code = injectArgs(language, stage.Tool, stage.Args, toolCode) + switch language { + case "awk", "sed", "jq", "ed", "expect", "bc": + language = "bash" + } + } + return e.executeWithStdin(ctx, code, language, timeout, sandboxName, true, stdinData) + } + + // Inline code stage (bash). + if stage.Code != "" { + validator := &Server{} + if err := validator.ValidateCode(stage.Code, "bash"); err != nil { + return "", err + } + return e.executeWithStdin(ctx, stage.Code, "bash", timeout, sandboxName, false, stdinData) + } + + return "", fmt.Errorf("stage %d: requires code, tool, or parallel", idx) } diff --git a/pkg/tools/execute/server.go b/pkg/tools/execute/server.go index eb8d337..b83b971 100644 --- a/pkg/tools/execute/server.go +++ b/pkg/tools/execute/server.go @@ -53,30 +53,43 @@ func (e *Server) ListTools() ([]tools.ToolInfo, error) { return []tools.ToolInfo{ { Name: "execute_code", - Description: `Run inline shell code in a sandboxed environment. + Description: `Run one or more code snippets in a sandboxed environment. + +Steps run in parallel when more than one is provided; results are returned in +submission order under === step N === headers. Single-step calls return raw output. Usage: -- Run shell commands: grep, cat, sed, find, etc. -- Default timeout: 30 seconds -- Default language: bash -- Output includes stdout, stderr, and exit code +- Default language: bash. Supported: bash, python3, perl, lua, awk, sed, jq, ed, expect, bc. +- Default timeout: 30 seconds (applied per step). +- Output includes stdout, stderr, and exit code. Security: - Dangerous commands blocked: rm -rf, fork bombs, sudo, etc. -- File system access limited to sandbox -- Network access restricted +- File system access limited to sandbox. +- Network access restricted. Examples: -- List files: code='ls -la' -- Search content: code='grep -r "TODO" ./src/' -- Count lines: code='wc -l file.txt'`, +- Single step: steps=[{code: "ls -la"}] +- Parallel: steps=[{code: "wc -l a.txt"}, {code: "wc -l b.txt"}] +- Mixed langs: steps=[{code: "...", language: "python3"}, {code: "...", language: "bash"}]`, InputSchema: json.RawMessage(`{ "type": "object", - "required": ["code"], + "required": ["steps"], "properties": { - "code": {"type": "string", "description": "Code to execute."}, - "language": {"type": "string", "description": "Language interpreter (default: bash)."}, - "timeout": {"type": "integer", "description": "Timeout in seconds (default: 30)."}, + "steps": { + "type": "array", + "description": "One or more steps. Multiple steps run in parallel. Each step is either inline code or a named tool.", + "items": { + "type": "object", + "properties": { + "code": {"type": "string", "description": "Inline code to execute."}, + "language": {"type": "string", "description": "Language interpreter (default: bash). Ignored when tool is set."}, + "tool": {"type": "string", "description": "Named tool script from the tools directory. Takes precedence over code."}, + "args": {"type": "array", "items": {"type": "string"}, "description": "Arguments for the tool script."} + } + } + }, + "timeout": {"type": "integer", "description": "Timeout in seconds per step (default: 30)."}, "sandbox": {"type": "string", "description": "Sandbox name (default: default)."} } }`), @@ -111,23 +124,24 @@ Examples: }, { Name: "execute_pipe", - Description: `Run a pipeline of commands, chaining stdout to stdin. + Description: `Run a sequential pipeline, chaining each stage's stdout to the next stage's stdin. -Usage: -- Pipe output between multiple commands -- Each step: {code: "cmd"} or {tool: "name", args: [...]} -- Default timeout: 30 seconds per pipeline -- Steps execute sequentially +Each stage is one of: +- {code: "..."} — inline bash +- {tool: "name", args: [...]} — named script from the tools directory +- {parallel: [{code,language}, ...]} — fan-out: N steps run concurrently, outputs + concatenated in submission order, result fed to next stage -Structure: -- pipe: array of step objects -- Step with code: shell command string -- Step with tool: named script from tools directory -- Use code for shell commands, tool for scripts +Use parallel stages when N independent operations produce the same output schema and +their combined output should flow into the next stage as a single stream. For disparate +schemas, normalize with an inner pipe stage before merging. + +Timeout applies per stage. Stages execute sequentially; a failed stage aborts the pipeline. Examples: -- Filter and count: pipe=[{code: "grep error log.txt"}, {code: "wc -l"}] -- Process with script: pipe=[{tool: "parse.py", args: ["data.json"]}, {code: "jq .result"}]`, +- Filter and count: pipe=[{code: "grep error log.txt"}, {code: "wc -l"}] +- Parallel then sort: pipe=[{parallel: [{code:"cat a.txt"},{code:"cat b.txt"}]}, {code:"sort"}] +- Tool transform: pipe=[{tool: "parse.py", args: ["data.json"]}, {code: "jq .result"}]`, InputSchema: json.RawMessage(`{ "type": "object", "required": ["pipe"], @@ -137,13 +151,26 @@ Examples: "items": { "type": "object", "properties": { - "tool": {"type": "string"}, - "args": {"type": "array", "items": {"type": "string"}}, - "code": {"type": "string"} + "tool": {"type": "string"}, + "args": {"type": "array", "items": {"type": "string"}}, + "code": {"type": "string"}, + "parallel": { + "type": "array", + "description": "Steps to run in parallel; outputs concatenated in submission order. Each step is inline code or a named tool.", + "items": { + "type": "object", + "properties": { + "code": {"type": "string", "description": "Inline code to execute."}, + "language": {"type": "string", "description": "Interpreter for inline code (default: bash). Ignored when tool is set."}, + "tool": {"type": "string", "description": "Named tool script; language detected from shebang."}, + "args": {"type": "array", "items": {"type": "string"}, "description": "Arguments for the tool script."} + } + } + } } } }, - "timeout": {"type": "integer", "description": "Timeout in seconds (default: 30)."}, + "timeout": {"type": "integer", "description": "Timeout in seconds per stage (default: 30)."}, "sandbox": {"type": "string", "description": "Sandbox name (default: default)."} } }`), @@ -254,6 +281,12 @@ func (e *Server) recordValidationFailure() { // Execute runs code in a sandbox and returns combined stdout+stderr. func (e *Server) Execute(ctx context.Context, code, language string, timeout int, sandboxName string, trusted bool) (string, error) { + return e.executeWithStdin(ctx, code, language, timeout, sandboxName, trusted, "") +} + +// executeWithStdin is like Execute but feeds stdinData to the command's stdin. +// For languages where code is itself passed via stdin (ed, expect, bc), stdinData is ignored. +func (e *Server) executeWithStdin(ctx context.Context, code, language string, timeout int, sandboxName string, trusted bool, stdinData string) (string, error) { if timeout <= 0 { timeout = 30 } @@ -287,7 +320,9 @@ func (e *Server) Execute(ctx context.Context, code, language string, timeout int var cmd *exec.Cmd var interpreter []string - var stdin string // non-empty: pipe this to cmd.Stdin instead of embedding in args + // codeStdin: non-empty means the code itself is fed via stdin (ed, expect, bc). + // In these cases stdinData cannot be used simultaneously. + var codeStdin string switch language { case "bash", "": interpreter = []string{"bash", "-c", code} @@ -305,21 +340,24 @@ func (e *Server) Execute(ctx context.Context, code, language string, timeout int interpreter = []string{"jq", code} case "ed": interpreter = []string{"ed", "-s"} - stdin = code + codeStdin = code case "expect": interpreter = []string{"expect", "-"} - stdin = code + codeStdin = code case "bc": interpreter = []string{"bc", "-ql"} - stdin = code + codeStdin = code default: return "", fmt.Errorf("unsupported language: %s (supported: bash, python3, perl, awk, sed, ed, jq, expect, bc, lua)", language) } wrapped := sandbox.WrapCommand(cfg, interpreter, workDir) cmd = exec.CommandContext(ctx, wrapped[0], wrapped[1:]...) cmd.Dir = workDir - if stdin != "" { - cmd.Stdin = strings.NewReader(stdin) + switch { + case codeStdin != "": + cmd.Stdin = strings.NewReader(codeStdin) + case stdinData != "": + cmd.Stdin = strings.NewReader(stdinData) } // Inject per-session env vars so they're visible inside the sandbox.