From 16d057741c478b5592774f6455860c6f76266a79 Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Fri, 17 Jul 2026 16:43:07 +0200 Subject: [PATCH] execute: add process detach mechanism Allows the user to detach a running process from the agent mid-execution. The process continues running; output flows to a 64KB ring buffer. Detached processes can be listed, signaled (TERM/KILL), and dismissed. - DetachedProcess struct with ring buffer, signal, output methods - Server.Detach() signals the currently executing process to detach - Execute() rewritten: cmd.Start() + select on wait/ctx/detach - cleanupDetached() on Close() sends SIGTERM to all - OnDetach/OnExit hooks for integration with D-Bus layer --- pkg/tools/execute/code.go | 4 + pkg/tools/execute/detach.go | 96 +++++++++++++++ pkg/tools/execute/execute_test.go | 59 +++++++++ pkg/tools/execute/server.go | 196 ++++++++++++++++++++++++++++-- 4 files changed, 344 insertions(+), 11 deletions(-) create mode 100644 pkg/tools/execute/detach.go diff --git a/pkg/tools/execute/code.go b/pkg/tools/execute/code.go index 340e157..b3a9ea5 100644 --- a/pkg/tools/execute/code.go +++ b/pkg/tools/execute/code.go @@ -7,6 +7,7 @@ import ( "io" "os" "path/filepath" + "sync" "ollie/internal/sandbox" "ollie/pkg/paths" @@ -158,6 +159,7 @@ func loadSandboxConfig(name string) (*sandbox.Config, error) { } type limitedWriter struct { + mu sync.Mutex w io.Writer written int limit int @@ -165,6 +167,8 @@ type limitedWriter struct { } func (lw *limitedWriter) Write(p []byte) (n int, err error) { + lw.mu.Lock() + defer lw.mu.Unlock() if lw.written >= lw.limit { lw.truncated = true return len(p), nil diff --git a/pkg/tools/execute/detach.go b/pkg/tools/execute/detach.go new file mode 100644 index 0000000..077b4a6 --- /dev/null +++ b/pkg/tools/execute/detach.go @@ -0,0 +1,96 @@ +package execute + +import ( + "fmt" + "os" + "sync" + "syscall" + "time" +) + +const ringBufSize = 64 * 1024 // 64KB ring buffer per detached process + +// DetachedProcess represents a process that the agent has detached from +// but which continues running. The user can view its output and signal it. +type DetachedProcess struct { + PID int + Command string + Started time.Time + Exited bool + ExitCode int + + ring *ringBuffer + cmd *os.Process + done chan struct{} + mu sync.Mutex +} + +// Output returns the current contents of the ring buffer. +func (p *DetachedProcess) Output() string { + p.mu.Lock() + defer p.mu.Unlock() + return p.ring.String() +} + +// Signal sends a signal to the detached process. +func (p *DetachedProcess) Signal(sig syscall.Signal) error { + p.mu.Lock() + defer p.mu.Unlock() + if p.Exited { + return fmt.Errorf("process %d already exited", p.PID) + } + if p.cmd == nil { + return fmt.Errorf("process %d has no handle", p.PID) + } + // Send to process group + return syscall.Kill(-p.PID, sig) +} + +// ringBuffer is a fixed-size circular byte buffer. +type ringBuffer struct { + buf []byte + size int + pos int + full bool +} + +func newRingBuffer(size int) *ringBuffer { + return &ringBuffer{buf: make([]byte, size), size: size} +} + +// Write implements io.Writer. +func (r *ringBuffer) Write(p []byte) (int, error) { + n := len(p) + if n >= r.size { + // Data larger than buffer: just keep the tail + copy(r.buf, p[n-r.size:]) + r.pos = 0 + r.full = true + return n, nil + } + if r.pos+n <= r.size { + copy(r.buf[r.pos:], p) + } else { + first := r.size - r.pos + copy(r.buf[r.pos:], p[:first]) + copy(r.buf, p[first:]) + r.full = true + } + r.pos = (r.pos + n) % r.size + if r.pos == 0 && n > 0 { + r.full = true + } + return n, nil +} + +// String returns the buffer contents in order. +func (r *ringBuffer) String() string { + if !r.full { + return string(r.buf[:r.pos]) + } + // Buffer has wrapped: data from pos..end + 0..pos + out := make([]byte, r.size) + copy(out, r.buf[r.pos:]) + copy(out[r.size-r.pos:], r.buf[:r.pos]) + return string(out) +} \ No newline at end of file diff --git a/pkg/tools/execute/execute_test.go b/pkg/tools/execute/execute_test.go index cb246bb..f8d3d22 100644 --- a/pkg/tools/execute/execute_test.go +++ b/pkg/tools/execute/execute_test.go @@ -7,6 +7,7 @@ import ( "path/filepath" "strings" "testing" + "time" ) // TestMain points OLLIE_CFG_PATH at testdata so loadSandboxConfig uses the @@ -515,3 +516,61 @@ func TestRunReadBatch_WithLockDir(t *testing.T) { t.Errorf("output = %q; want 'locked'", out) } } + +// ---- detach ---- + +func TestDetach(t *testing.T) { + s := newServer(t) + ctx := context.Background() + + // Start a long-running process, then detach it + var detachedPID int + s.OnDetach = func(pid int, cmd string) { + detachedPID = pid + } + + done := make(chan string, 1) + go func() { + out, _ := s.Execute(ctx, "echo started; sleep 30", "bash", 60, "default", true) + done <- out + }() + + // Give the process time to start and reach the select + for i := 0; i < 50; i++ { + if s.Detach() { + break + } + time.Sleep(10 * time.Millisecond) + } + + out := <-done + if !strings.Contains(out, "[detached: pid") { + t.Fatalf("expected detach marker in output, got: %q", out) + } + if detachedPID == 0 { + t.Fatal("OnDetach not called") + } + + // Verify process is in detached list + procs := s.ListDetached() + if len(procs) != 1 { + t.Fatalf("expected 1 detached process, got %d", len(procs)) + } + if procs[0].PID != detachedPID { + t.Errorf("PID mismatch: got %d, want %d", procs[0].PID, detachedPID) + } + if procs[0].Exited { + t.Error("process should still be running") + } + + // Kill it + if err := s.SignalDetached(detachedPID, 15); err != nil { // SIGTERM + t.Fatalf("SignalDetached: %v", err) + } + + // Wait for exit + <-procs[0].done + if !procs[0].Exited { + t.Error("process should have exited after SIGTERM") + } +} diff --git a/pkg/tools/execute/server.go b/pkg/tools/execute/server.go index ebffc24..8bc50bf 100644 --- a/pkg/tools/execute/server.go +++ b/pkg/tools/execute/server.go @@ -65,6 +65,13 @@ type Server struct { validationFailures int lastFailure time.Time blockedUntil time.Time + + // Detached process management + detachMu sync.Mutex + detachCh chan struct{} // signal to detach the currently running process + detached []*DetachedProcess + OnDetach func(pid int, cmd string) // hook: called when a process is detached + OnExit func(pid int, exitCode int) // hook: called when a detached process exits } // Option configures a Server. @@ -386,6 +393,7 @@ func (e *Server) SetEnv(key, value string) { } // Close is called when the session ends. Calls OnClose hook if registered. func (e *Server) Close() { + e.cleanupDetached() if e.OnClose != nil { e.OnClose() } @@ -505,7 +513,9 @@ func (e *Server) executeWithStdin(ctx context.Context, code, language string, ti } ctx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second) - defer cancel() + // NOTE: do NOT defer cancel() here — the detach path must avoid cancelling + // the context (which would kill the detached process via cmd.Cancel). + // Each exit path calls cancel() explicitly except detach. var cmd *exec.Cmd var interpreter []string @@ -606,19 +616,106 @@ func (e *Server) executeWithStdin(ctx context.Context, code, language string, ti cmd.Stdout = lw cmd.Stderr = lw - err = cmd.Run() - output := outputBuf.Bytes() - if lw.truncated { - output = append(output, []byte("\n[output truncated at 10MB]")...) + // Set up detach channel for this execution + detachCh := make(chan struct{}, 1) + e.detachMu.Lock() + e.detachCh = detachCh + e.detachMu.Unlock() + defer func() { + e.detachMu.Lock() + e.detachCh = nil + e.detachMu.Unlock() + }() + + if err = cmd.Start(); err != nil { + return "", fmt.Errorf("execution failed: %v", err) } - if ctx.Err() == context.DeadlineExceeded { - return "", fmt.Errorf("execution timeout after %d seconds", timeout) + // Wait for completion, context cancellation, or detach signal + waitCh := make(chan error, 1) + go func() { waitCh <- cmd.Wait() }() + + select { + case err = <-waitCh: + // Normal completion + cancel() + output := outputBuf.Bytes() + if lw.truncated { + output = append(output, []byte("\n[output truncated at 10MB]")...) + } + if ctx.Err() == context.DeadlineExceeded { + return "", fmt.Errorf("execution timeout after %d seconds", timeout) + } + if err != nil { + return string(output), fmt.Errorf("execution failed: %v\nOutput: %s", err, string(output)) + } + return string(output), nil + + case <-ctx.Done(): + // Context cancelled (interrupt or timeout) — cancel() already fired via ctx + syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) + <-waitCh // reap + output := outputBuf.Bytes() + if ctx.Err() == context.DeadlineExceeded { + return string(output), fmt.Errorf("execution timeout after %d seconds", timeout) + } + return string(output), fmt.Errorf("execution interrupted") + + case <-detachCh: + // Detach: move process to registry, continue running + // Neutralize cmd.Cancel so context cancellation won't SIGKILL the process + cmd.Cancel = nil + // Cancel the timeout context (harmless now that cmd.Cancel is nil) + cancel() + + cmdStr := strings.Join(interpreter, " ") + if len(cmdStr) > 80 { + cmdStr = cmdStr[:77] + "..." + } + ring := newRingBuffer(ringBufSize) + // Splice: future output goes to ring buffer instead of outputBuf. + lw.mu.Lock() + lw.w = ring + lw.mu.Unlock() + + proc := &DetachedProcess{ + PID: cmd.Process.Pid, + Command: cmdStr, + Started: time.Now(), + ring: ring, + cmd: cmd.Process, + done: make(chan struct{}), + } + e.detachMu.Lock() + e.detached = append(e.detached, proc) + e.detachMu.Unlock() + + // Monitor for exit in background + go func() { + waitErr := <-waitCh + proc.mu.Lock() + proc.Exited = true + if waitErr != nil { + if exitErr, ok := waitErr.(*exec.ExitError); ok { + proc.ExitCode = exitErr.ExitCode() + } else { + proc.ExitCode = -1 + } + } + proc.mu.Unlock() + close(proc.done) + if e.OnExit != nil { + e.OnExit(proc.PID, proc.ExitCode) + } + }() + + if e.OnDetach != nil { + e.OnDetach(proc.PID, cmdStr) + } + + partial := outputBuf.String() + return partial + fmt.Sprintf("\n[detached: pid %d]", proc.PID), nil } - if err != nil { - return string(output), fmt.Errorf("execution failed: %v\nOutput: %s", err, string(output)) - } - return string(output), nil } // prependOlliePath returns env with $OLLIE_CFG_PATH/scripts/x/ prepended to PATH so wrapper @@ -662,4 +759,81 @@ func plan9Namespace(env map[string]string) string { return "/tmp/ns." + user + "." + disp } +// Detach signals the currently running process to be detached from the agent. +// The process continues running; its output is captured in a ring buffer. +// Returns false if no process is currently running. +func (e *Server) Detach() bool { + e.detachMu.Lock() + ch := e.detachCh + e.detachMu.Unlock() + if ch == nil { + return false + } + select { + case ch <- struct{}{}: + return true + default: + return false + } +} + +// ListDetached returns all detached processes (running and exited). +func (e *Server) ListDetached() []*DetachedProcess { + e.detachMu.Lock() + defer e.detachMu.Unlock() + out := make([]*DetachedProcess, len(e.detached)) + copy(out, e.detached) + return out +} + +// SignalDetached sends a signal to a detached process by PID. +func (e *Server) SignalDetached(pid int, sig syscall.Signal) error { + e.detachMu.Lock() + defer e.detachMu.Unlock() + for _, p := range e.detached { + if p.PID == pid { + return p.Signal(sig) + } + } + return fmt.Errorf("no detached process with pid %d", pid) +} + +// GetDetachedOutput returns the ring buffer contents for a detached process. +func (e *Server) GetDetachedOutput(pid int) (string, error) { + e.detachMu.Lock() + defer e.detachMu.Unlock() + for _, p := range e.detached { + if p.PID == pid { + return p.Output(), nil + } + } + return "", fmt.Errorf("no detached process with pid %d", pid) +} + +// DismissDetached removes an exited process from the list. +func (e *Server) DismissDetached(pid int) bool { + e.detachMu.Lock() + defer e.detachMu.Unlock() + for i, p := range e.detached { + if p.PID == pid && p.Exited { + e.detached = append(e.detached[:i], e.detached[i+1:]...) + return true + } + } + return false +} + +// cleanupDetached sends SIGTERM to all running detached processes. +func (e *Server) cleanupDetached() { + e.detachMu.Lock() + defer e.detachMu.Unlock() + for _, p := range e.detached { + p.mu.Lock() + if !p.Exited && p.cmd != nil { + syscall.Kill(-p.PID, syscall.SIGTERM) + } + p.mu.Unlock() + } +} + var _ tools.Server = (*Server)(nil) // compile-time interface check