diff --git a/internal/agent/tools/bash.go b/internal/agent/tools/bash.go index 855e09f6..5eba60fe 100644 --- a/internal/agent/tools/bash.go +++ b/internal/agent/tools/bash.go @@ -292,9 +292,7 @@ func (t *BashTool) executeBashWithStreaming(ctx context.Context, cmd *exec.Cmd, detached = true detachedMux.Unlock() - time.Sleep(20 * time.Millisecond) - - shellID, err := t.backgroundShellService.DetachToBackground(ctx, cmd, result.Command, outputBuffer) + shellID, err := t.backgroundShellService.DetachToBackground(ctx, cmd, result.Command, outputBuffer, done) if err != nil { logger.Debug("bash: DetachToBackground failed", "error", err) result.ExitCode = -1 diff --git a/internal/domain/shell.go b/internal/domain/shell.go index 1d827b89..653dc7c0 100644 --- a/internal/domain/shell.go +++ b/internal/domain/shell.go @@ -38,6 +38,8 @@ type BackgroundShell struct { OutputBuffer OutputRingBuffer CancelFunc context.CancelFunc ReadOffset int64 + + ReadersDone <-chan struct{} } // OutputRingBuffer defines the interface for the circular output buffer. @@ -110,8 +112,10 @@ func NewShellInfo(shell *BackgroundShell) *ShellInfo { // BackgroundShellService defines the interface for managing background shells type BackgroundShellService interface { - // DetachToBackground moves a running command to background - DetachToBackground(ctx context.Context, cmd *exec.Cmd, command string, outputBuffer OutputRingBuffer) (string, error) + // DetachToBackground moves a running command to background. readersDone, when + // non-nil, is closed once the caller's pipe readers reach EOF; the supervisor + // waits on it before reaping so trailing output is not truncated. + DetachToBackground(ctx context.Context, cmd *exec.Cmd, command string, outputBuffer OutputRingBuffer, readersDone <-chan struct{}) (string, error) // GetShellOutput retrieves output from a shell GetShellOutput(shellID string, fromOffset int64) (string, int64, ShellState, error) diff --git a/internal/services/background_shell_service.go b/internal/services/background_shell_service.go index 9f8b5158..1c50b047 100644 --- a/internal/services/background_shell_service.go +++ b/internal/services/background_shell_service.go @@ -51,6 +51,7 @@ func (s *BackgroundShellService) DetachToBackground( cmd *exec.Cmd, command string, outputBuffer domain.OutputRingBuffer, + readersDone <-chan struct{}, ) (string, error) { if !s.config.Tools.Bash.BackgroundShells.Enabled { return "", fmt.Errorf("background shells are disabled in configuration") @@ -66,6 +67,7 @@ func (s *BackgroundShellService) DetachToBackground( State: domain.ShellStateRunning, OutputBuffer: outputBuffer, ReadOffset: 0, + ReadersDone: readersDone, } if err := s.shellTracker.Add(shell); err != nil { diff --git a/internal/services/directexec/bash.go b/internal/services/directexec/bash.go index 9e377ac9..062408cc 100644 --- a/internal/services/directexec/bash.go +++ b/internal/services/directexec/bash.go @@ -367,18 +367,25 @@ func (s *Service) executeBashCommandInBackground(commandText, command string) te } }() + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + shellID, err := s.backgroundShellService.DetachToBackground( ctx, cmd, command, outputBuffer, + done, ) if err != nil { return } - wg.Wait() + <-done assistantEntry := domain.ConversationEntry{ Message: sdk.Message{ diff --git a/internal/services/jobs/shell_job.go b/internal/services/jobs/shell_job.go index 59333e56..464e3767 100644 --- a/internal/services/jobs/shell_job.go +++ b/internal/services/jobs/shell_job.go @@ -40,7 +40,17 @@ func (j *shellJob) Meta() domain.JobMeta { } // Run waits for the shell process to exit and records the outcome on the shell. +// It first waits for the pipe readers to drain (ReadersDone) so Cmd.Wait doesn't +// close the pipes mid-read and lose output; the wait honours ctx so a kill or +// shutdown can't wedge when a grandchild is holding the pipe open. func (j *shellJob) Run(ctx context.Context, _ func(domain.JobSignal)) domain.ToolExecutionResult { + if j.shell.ReadersDone != nil { + select { + case <-j.shell.ReadersDone: + case <-ctx.Done(): + } + } + waitErr := j.shell.Cmd.Wait() now := time.Now() @@ -49,7 +59,6 @@ func (j *shellJob) Run(ctx context.Context, _ func(domain.JobSignal)) domain.Too switch { case ctx.Err() != nil: - // Terminated by Wind(WindStop) / supervisor shutdown rather than finishing. code := -1 j.shell.ExitCode = &code j.shell.State = domain.ShellStateCancelled diff --git a/internal/services/jobs/shell_job_test.go b/internal/services/jobs/shell_job_test.go index 024359af..576f0aa8 100644 --- a/internal/services/jobs/shell_job_test.go +++ b/internal/services/jobs/shell_job_test.go @@ -45,6 +45,71 @@ func waitTerminal(t *testing.T, sup *Supervisor, id string) domain.JobStatus { return "" } +// isTerminal reports whether the job is currently in a terminal state without +// blocking, for asserting that a job has NOT yet finished. +func isTerminal(sup *Supervisor, id string) bool { + for _, j := range sup.Snapshot() { + if j.Meta.ID == id { + return j.Status.IsTerminal() + } + } + return false +} + +// TestShellJob_WaitsForReadersDoneBeforeReaping asserts Run does not call +// Cmd.Wait until the pipe readers have drained (ReadersDone closed), even after +// the process has already exited — otherwise Wait closes the pipes mid-drain and +// truncates trailing output. +func TestShellJob_WaitsForReadersDoneBeforeReaping(t *testing.T) { + tracker := utils.NewShellTracker(10) + sup := NewSupervisor(&domainmocks.FakeMessageQueue{}, &domainmocks.FakeConversationRepository{}, nil) + defer sup.Stop() + + // `true` exits immediately; the job must still park on ReadersDone. + shell := startShell(t, "s-drain", "true") + readersDone := make(chan struct{}) + shell.ReadersDone = readersDone + _ = tracker.Add(shell) + sup.Submit(NewShellJob(shell, tracker)) + + time.Sleep(50 * time.Millisecond) + if isTerminal(sup, "s-drain") { + t.Fatal("shell reaped before the pipe readers signalled done") + } + + close(readersDone) + if got := waitTerminal(t, sup, "s-drain"); got != domain.JobCompleted { + t.Fatalf("status = %s, want completed", got) + } + if shell.State != domain.ShellStateCompleted { + t.Fatalf("shell state = %s, want completed", shell.State) + } +} + +// TestShellJob_WindStopUnblocksReadersWait guards the cancellable-wait escape: +// when ReadersDone never closes (e.g. a grandchild holds the pipe open), +// Wind(WindStop) must still reap the job via ctx cancellation rather than +// wedging Run — which would also hang Supervisor.Stop(). +func TestShellJob_WindStopUnblocksReadersWait(t *testing.T) { + tracker := utils.NewShellTracker(10) + sup := NewSupervisor(&domainmocks.FakeMessageQueue{}, &domainmocks.FakeConversationRepository{}, nil) + defer sup.Stop() + + shell := startShell(t, "s-stuck", "sleep", "60") + shell.ReadersDone = make(chan struct{}) // never closed + _ = tracker.Add(shell) + sup.Submit(NewShellJob(shell, tracker)) + + time.Sleep(20 * time.Millisecond) + if err := sup.Wind("s-stuck", domain.WindStop); err != nil { + t.Fatalf("Wind: %v", err) + } + waitTerminal(t, sup, "s-stuck") + if shell.State != domain.ShellStateCancelled { + t.Fatalf("shell state = %s, want cancelled", shell.State) + } +} + func TestShellJob_CompletesAndNotifies(t *testing.T) { tracker := utils.NewShellTracker(10) queue := &domainmocks.FakeMessageQueue{} diff --git a/tests/mocks/domain/fake_background_shell_service.go b/tests/mocks/domain/fake_background_shell_service.go index 6a476569..6e4eaa41 100644 --- a/tests/mocks/domain/fake_background_shell_service.go +++ b/tests/mocks/domain/fake_background_shell_service.go @@ -21,13 +21,14 @@ type FakeBackgroundShellService struct { cancelShellReturnsOnCall map[int]struct { result1 error } - DetachToBackgroundStub func(context.Context, *exec.Cmd, string, domain.OutputRingBuffer) (string, error) + DetachToBackgroundStub func(context.Context, *exec.Cmd, string, domain.OutputRingBuffer, <-chan struct{}) (string, error) detachToBackgroundMutex sync.RWMutex detachToBackgroundArgsForCall []struct { arg1 context.Context arg2 *exec.Cmd arg3 string arg4 domain.OutputRingBuffer + arg5 <-chan struct{} } detachToBackgroundReturns struct { result1 string @@ -171,7 +172,7 @@ func (fake *FakeBackgroundShellService) CancelShellReturnsOnCall(i int, result1 }{result1} } -func (fake *FakeBackgroundShellService) DetachToBackground(arg1 context.Context, arg2 *exec.Cmd, arg3 string, arg4 domain.OutputRingBuffer) (string, error) { +func (fake *FakeBackgroundShellService) DetachToBackground(arg1 context.Context, arg2 *exec.Cmd, arg3 string, arg4 domain.OutputRingBuffer, arg5 <-chan struct{}) (string, error) { fake.detachToBackgroundMutex.Lock() ret, specificReturn := fake.detachToBackgroundReturnsOnCall[len(fake.detachToBackgroundArgsForCall)] fake.detachToBackgroundArgsForCall = append(fake.detachToBackgroundArgsForCall, struct { @@ -179,13 +180,14 @@ func (fake *FakeBackgroundShellService) DetachToBackground(arg1 context.Context, arg2 *exec.Cmd arg3 string arg4 domain.OutputRingBuffer - }{arg1, arg2, arg3, arg4}) + arg5 <-chan struct{} + }{arg1, arg2, arg3, arg4, arg5}) stub := fake.DetachToBackgroundStub fakeReturns := fake.detachToBackgroundReturns - fake.recordInvocation("DetachToBackground", []interface{}{arg1, arg2, arg3, arg4}) + fake.recordInvocation("DetachToBackground", []interface{}{arg1, arg2, arg3, arg4, arg5}) fake.detachToBackgroundMutex.Unlock() if stub != nil { - return stub(arg1, arg2, arg3, arg4) + return stub(arg1, arg2, arg3, arg4, arg5) } if specificReturn { return ret.result1, ret.result2 @@ -199,17 +201,17 @@ func (fake *FakeBackgroundShellService) DetachToBackgroundCallCount() int { return len(fake.detachToBackgroundArgsForCall) } -func (fake *FakeBackgroundShellService) DetachToBackgroundCalls(stub func(context.Context, *exec.Cmd, string, domain.OutputRingBuffer) (string, error)) { +func (fake *FakeBackgroundShellService) DetachToBackgroundCalls(stub func(context.Context, *exec.Cmd, string, domain.OutputRingBuffer, <-chan struct{}) (string, error)) { fake.detachToBackgroundMutex.Lock() defer fake.detachToBackgroundMutex.Unlock() fake.DetachToBackgroundStub = stub } -func (fake *FakeBackgroundShellService) DetachToBackgroundArgsForCall(i int) (context.Context, *exec.Cmd, string, domain.OutputRingBuffer) { +func (fake *FakeBackgroundShellService) DetachToBackgroundArgsForCall(i int) (context.Context, *exec.Cmd, string, domain.OutputRingBuffer, <-chan struct{}) { fake.detachToBackgroundMutex.RLock() defer fake.detachToBackgroundMutex.RUnlock() argsForCall := fake.detachToBackgroundArgsForCall[i] - return argsForCall.arg1, argsForCall.arg2, argsForCall.arg3, argsForCall.arg4 + return argsForCall.arg1, argsForCall.arg2, argsForCall.arg3, argsForCall.arg4, argsForCall.arg5 } func (fake *FakeBackgroundShellService) DetachToBackgroundReturns(result1 string, result2 error) {