lkn
/
ollie-core
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

execute: refactor execute_code to parallel steps; add parallel pipe stages 

- execute_code now takes steps:[{code,language}|{tool,args}] instead of
  a single code/language pair. Single-step calls return raw output;
  multi-step calls run in parallel and return results under === step N ===
  headers in submission order.
- execute_pipe stages may now include {parallel:[...]} to fan out
  concurrently within a pipeline. Parallel outputs are concatenated in
  submission order and fed as stdin to the next stage.
- Stage execution is now Go-level (stdin chaining) rather than a single
  bash pipeline string, enabling parallel stages and tool/code mixing.
- CodeStep and resolveCodeStep are the shared primitives for both
  execute_code and execute_pipe parallel stages.
- executeWithStdin added to Server for stdin chaining between stages.
- Tool steps use shebang detection; language field applies only to
  inline code steps (default: bash).
- ListTools schemas and descriptions updated throughout.
This commit is contained in:
Levi Neely 2026-04-16 20:06:05 +02:00
parent 54f20bce02
commit cf25268d80
3 changed files with 284 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

131
pkg/tools/execute/code.go
View File

@ -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
}

111
pkg/tools/execute/pipe.go
View File

@ -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)
}

112
pkg/tools/execute/server.go
View File

@ -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.