diff --git a/tool/shelltool/localshell_internal_test.go b/tool/shelltool/localshell_internal_test.go index 31c88378..1be2b1b4 100644 --- a/tool/shelltool/localshell_internal_test.go +++ b/tool/shelltool/localshell_internal_test.go @@ -3,9 +3,12 @@ package shelltool import ( + "context" + "io" "slices" "strings" "testing" + "time" ) func TestPreserveEnvironmentValuesKeepsOnlyAllowlist(t *testing.T) { @@ -78,6 +81,68 @@ func TestHeadTailBuffer_MultiByteUTF8OverflowPreservesHeadTailOrder(t *testing.T } } +// scriptedReader delivers a fixed sequence of chunks, sleeping delay before +// each so successive stdout chunks land on distinct readLoop signals. +type scriptedReader struct { + chunks [][]byte + delay time.Duration + i int +} + +func (r *scriptedReader) Read(p []byte) (int, error) { + if r.i >= len(r.chunks) { + return 0, io.EOF + } + if r.delay > 0 { + time.Sleep(r.delay) + } + n := copy(p, r.chunks[r.i]) + r.i++ + return n, nil +} + +// TestReadExitCodeWakesOnLiveStdoutSignal exercises the exit-code read path +// when the sentinel and the exit-code line arrive in two separate stdout +// reads. readExitCode must wake on the live stdout signal for the second read +// rather than waiting out the 100ms fallback poll. +func TestReadExitCodeWakesOnLiveStdoutSignal(t *testing.T) { + s := &persistentSession{stdoutSignal: newSignal()} + s.readerWG.Add(1) + + reader := &scriptedReader{ + chunks: [][]byte{ + []byte("hello\nSENTINEL"), + []byte("_0\n"), + }, + delay: 5 * time.Millisecond, + } + go s.readLoop(reader, &s.stdoutBuf, true) + + start := time.Now() + idx, rc, timedOut, overflow, err := s.waitForSentinel(context.Background(), []byte("SENTINEL"), 0, 1<<20) + elapsed := time.Since(start) + s.readerWG.Wait() + + if err != nil { + t.Fatalf("waitForSentinel: %v", err) + } + if timedOut || overflow { + t.Fatalf("timedOut=%v overflow=%v, want both false", timedOut, overflow) + } + if idx < 0 { + t.Fatalf("sentinel not found: idx=%d", idx) + } + if rc != 0 { + t.Fatalf("rc = %d, want 0", rc) + } + // The exit-code line arrives ~5ms after the sentinel on a live signal. + // Before the fix readExitCode discarded that signal and only noticed the + // line on the 100ms fallback poll, so elapsed would be ~100ms. + if elapsed >= 80*time.Millisecond { + t.Fatalf("readExitCode gated on 100ms poll: elapsed %v", elapsed) + } +} + func TestResolvedShellArgvIncludesExtraArgv(t *testing.T) { shell := resolvedShell{ binary: "/custom/bash", diff --git a/tool/shelltool/session.go b/tool/shelltool/session.go index 770cdcf1..6553e22e 100644 --- a/tool/shelltool/session.go +++ b/tool/shelltool/session.go @@ -313,7 +313,6 @@ func (s *persistentSession) readExitCode(ctx context.Context, afterIdx int) (int tail = snapshotRange(s.stdoutBuf, afterIdx, lenTail) } signal := s.stdoutSignal - s.stdoutSignal = newSignal() s.bufferGate.Unlock() if nl := bytes.IndexByte(tail, '\n'); nl >= 0 {