diff --git a/pkg/agent/core.go b/pkg/agent/core.go index 8637bfa..0baaea6 100644 --- a/pkg/agent/core.go +++ b/pkg/agent/core.go @@ -246,6 +246,18 @@ func (s *agent) pushSessionEnv() { } } +// pushLockDir sets the flock directory on the execute server to the session tmpdir. +func (s *agent) pushLockDir() { + if s.dispatcher == nil || s.sessionID == "" { + return + } + if srv, ok := s.dispatcher.GetServer("execute"); ok { + if ls, ok := srv.(tools.LockDirSetter); ok { + ls.SetLockDir(filepath.Join(ollieTmpDir(), s.sessionID)) + } + } +} + var _ Core = (*agent)(nil) // compile-time interface check var sweepTmpOnce sync.Once @@ -313,6 +325,7 @@ func NewAgentCore(cfg AgentCoreConfig) Core { } a.changeCond = sync.NewCond(&a.changeMu) a.pushSessionEnv() + a.pushLockDir() return a } diff --git a/pkg/tools/execute/code.go b/pkg/tools/execute/code.go index 0101d09..84b1e2d 100644 --- a/pkg/tools/execute/code.go +++ b/pkg/tools/execute/code.go @@ -207,14 +207,38 @@ func dispatchExecuteCode(ctx context.Context, e *Server, args json.RawMessage) ( return e.executeWithStdin(ctx, code, lang, timeout, sandboxName, trusted, "") } - // Pipeline: run stages sequentially, chaining stdout → stdin. + // Pipeline: run stages sequentially, auto-batching consecutive read-safe + // tool steps into parallel groups. Write/global steps serialize. var input string - for i, stage := range stages { - out, err := e.runStage(ctx, i, stage, timeout, sandboxName, input) - if err != nil { - return out, fmt.Errorf("step %d: %w", i, err) + for i := 0; i < len(stages); { + if classifyStep(stages[i]) == lockClassRead { + // Collect run of consecutive read-safe stages. + j := i + 1 + for j < len(stages) && classifyStep(stages[j]) == lockClassRead { + j++ + } + out, err := e.runReadBatch(ctx, i, stages[i:j], timeout, sandboxName, input) + if err != nil { + return out, err + } + input = out + i = j + } else { + // Write or global: acquire exclusive lock, run single stage. + lf, err := acquireFlock(e.lockDir, "rw", true) + if err != nil { + return "", fmt.Errorf("step %d: lock: %w", i, err) + } + out, runErr := e.runStage(ctx, i, stages[i], timeout, sandboxName, input) + if lf != nil { + lf.Close() + } + if runErr != nil { + return out, fmt.Errorf("step %d: %w", i, runErr) + } + input = out + i++ } - input = out } return input, nil } diff --git a/pkg/tools/execute/flock.go b/pkg/tools/execute/flock.go new file mode 100644 index 0000000..e45d9b5 --- /dev/null +++ b/pkg/tools/execute/flock.go @@ -0,0 +1,108 @@ +package execute + +import ( + "os" + "path/filepath" + "strings" + "syscall" +) + +// lockClass classifies a CodeStep by its concurrency profile. +type lockClass int + +const ( + lockClassRead lockClass = iota // file_read, file_glob, file_grep — safe to parallelize + lockClassWrite // file_write, file_edit — exclusive per path + lockClassGlobal // inline code, bash, unknown tools — global serialize +) + +// classifyStep returns the lock class for a single CodeStep. +// Inline code, parallel groups, and elevated steps are always lockClassGlobal. +// Tool steps are classified by reading the script and checking for an +// ollie:parallel annotation in the first 10 lines (see detectParallelClass). +// Unknown or unreadable tools default to lockClassGlobal. +func classifyStep(step CodeStep) lockClass { + if step.Code != "" || len(step.Parallel) > 0 || step.Elevated || step.Tool == "" { + return lockClassGlobal + } + code, err := ReadTool(step.Tool) + if err != nil { + return lockClassGlobal + } + return detectParallelClass(code) +} + +// detectParallelClass scans the first 10 lines of a tool script for an +// ollie:parallel annotation: +// +// # ollie:parallel read → lockClassRead (safe to run concurrently with other reads) +// # ollie:parallel write → lockClassWrite (needs exclusive access to the file it writes) +// +// Absence of the annotation → lockClassGlobal (serialize). +func detectParallelClass(code string) lockClass { + lines := strings.SplitN(code, "\n", 11) + if len(lines) > 10 { + lines = lines[:10] + } + for _, line := range lines { + line = strings.TrimSpace(line) + if !strings.HasPrefix(line, "# ollie:parallel") { + continue + } + rest := strings.TrimSpace(strings.TrimPrefix(line, "# ollie:parallel")) + switch { + case rest == "read" || strings.HasPrefix(rest, "read "): + return lockClassRead + case rest == "write" || strings.HasPrefix(rest, "write "): + return lockClassWrite + } + } + return lockClassGlobal +} + +// acquireFlock opens (or creates) a lock file in dir named name and acquires +// LOCK_SH (exclusive=false) or LOCK_EX (exclusive=true). +// Returns nil, nil when dir is empty (locking disabled). +// Caller must Close the returned file to release the lock. +func acquireFlock(dir, name string, exclusive bool) (*os.File, error) { + if dir == "" { + return nil, nil + } + if err := os.MkdirAll(dir, 0700); err != nil { + return nil, err + } + how := syscall.LOCK_SH + if exclusive { + how = syscall.LOCK_EX + } + path := filepath.Join(dir, sanitizeLockName(name)+".lock") + f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, 0600) + if err != nil { + return nil, err + } + if err := syscall.Flock(int(f.Fd()), how); err != nil { + f.Close() + return nil, err + } + return f, nil +} + +func sanitizeLockName(s string) string { + var b strings.Builder + for _, r := range s { + switch r { + case '/', '\\', ':', '*', '?', '<', '>', '|', '"', ' ': + b.WriteRune('_') + default: + b.WriteRune(r) + } + } + name := b.String() + if len(name) > 64 { + name = name[:64] + } + if name == "" { + return "unnamed" + } + return name +} diff --git a/pkg/tools/execute/pipe.go b/pkg/tools/execute/pipe.go index c97560d..cf1bf7a 100644 --- a/pkg/tools/execute/pipe.go +++ b/pkg/tools/execute/pipe.go @@ -60,3 +60,53 @@ func (e *Server) runStage(ctx context.Context, idx int, stage CodeStep, timeout return "", fmt.Errorf("stage %d: requires code, tool, or parallel", idx) } + +// runReadBatch runs a contiguous slice of read-safe stages concurrently, +// concatenating their outputs in submission order. All stages receive the +// same stdinData (read-safe tools don't depend on stdin chaining). +// If e.lockDir is set, each goroutine acquires LOCK_SH before running. +func (e *Server) runReadBatch(ctx context.Context, startIdx int, stages []CodeStep, timeout int, sandboxName, stdinData string) (string, error) { + if len(stages) == 1 { + lf, err := acquireFlock(e.lockDir, "rw", false) + if err != nil { + return "", fmt.Errorf("step %d: lock: %w", startIdx, err) + } + if lf != nil { + defer lf.Close() + } + out, err := e.runStage(ctx, startIdx, stages[0], timeout, sandboxName, stdinData) + if err != nil { + return out, fmt.Errorf("step %d: %w", startIdx, err) + } + return out, nil + } + + results := make([]string, len(stages)) + errs := make([]error, len(stages)) + var wg sync.WaitGroup + for i, stage := range stages { + wg.Add(1) + go func(j int, s CodeStep) { + defer wg.Done() + lf, lerr := acquireFlock(e.lockDir, "rw", false) + if lerr != nil { + errs[j] = fmt.Errorf("step %d: lock: %w", startIdx+j, lerr) + return + } + if lf != nil { + defer lf.Close() + } + results[j], errs[j] = e.runStage(ctx, startIdx+j, s, timeout, sandboxName, stdinData) + }(i, stage) + } + wg.Wait() + + var out string + for i, r := range results { + out += r + if errs[i] != nil { + return out, errs[i] + } + } + return out, nil +} diff --git a/pkg/tools/execute/server.go b/pkg/tools/execute/server.go index 2d144aa..44d0ae5 100644 --- a/pkg/tools/execute/server.go +++ b/pkg/tools/execute/server.go @@ -39,6 +39,10 @@ type Server struct { envMu sync.RWMutex envExtra map[string]string + // lockDir is the directory for advisory flock files used during parallel + // step dispatch. Set once at session init via SetLockDir; empty disables locking. + lockDir string + // rate limiting state (per-Server) rateLimitMu sync.Mutex validationFailures int @@ -169,6 +173,10 @@ func (e *Server) allowed(tool, detail string) bool { return e.Confirm(detail) } +// SetLockDir sets the directory used for advisory flock files during parallel +// step dispatch. Must be called before the Server handles concurrent requests. +func (e *Server) SetLockDir(dir string) { e.lockDir = dir } + // SetEnv adds a session-scoped environment variable injected into all // subsequent subprocess invocations for this session. func (e *Server) SetEnv(key, value string) { diff --git a/pkg/tools/tools.go b/pkg/tools/tools.go index d1c0485..e71d38a 100644 --- a/pkg/tools/tools.go +++ b/pkg/tools/tools.go @@ -106,3 +106,9 @@ type EnvSetter interface { type TrustedToolsSetter interface { SetTrustedTools(tools []string) } + +// LockDirSetter is implemented by tool servers that use advisory flock files +// for parallel-step coordination. SetLockDir must be called once at session init. +type LockDirSetter interface { + SetLockDir(dir string) +}