An unsafe session could report fewer refusals than it made

internal/connector/driver/codex/codex.go

@@ -474,6 +474,12 @@ type turn struct {
 	result   driver.PromptResult
 	err      error
 	canceled bool
+	// settling is a turn an ending has claimed: it is the one that finishes
+	// it, and every other ending leaves it alone. The turn stays the
+	// session's while it settles — a refusal read in the meantime still
+	// lands on it, which is the whole reason the claim is a flag and not a
+	// removal.
+	settling bool
 	refusals []driver.Refusal
 }
 
@@ -549,7 +555,12 @@ func (s *session) Prompt(ctx context.Context, prompt string) (driver.PromptResul
 		if canceled {
 			s.finishCanceled(t)
 		} else {
-			s.finish(t, driver.PromptResult{}, fmt.Errorf("%w: %w", driver.ErrSessionEnded, err))
+			// The refusals the turn has made are on its result, as they are
+			// on every other ending's: a session that ended under the turn
+			// still refused what it refused. They are what the reader has
+			// recorded by now, which on this ending — the reader is still
+			// reading — may be fewer than it goes on to record.
+			s.finish(t, driver.PromptResult{Refusals: s.refusalsOf(t)}, fmt.Errorf("%w: %w", driver.ErrSessionEnded, err))
 		}
 	}()
 	select {
@@ -601,25 +612,61 @@ func (s *session) Close() error {
 	case <-time.After(s.grace):
 	}
 	s.worker.Terminate(s.grace)
+	s.readerDone()
+	return nil
+}
+
+// readerDone waits for the reader to be through the worker's output. It is
+// called only from goroutines that are not the reader — Close, and the policy
+// check ending an unsafe session — because the reader cannot wait for itself.
+//
+// Past the grace, a descendant outside the worker's group is holding the
+// output open: the reading is ended rather than waited on, rather than hold
+// the attempt, its working directory and the connector's shutdown open
+// forever.
+func (s *session) readerDone() {
 	select {
 	case <-s.readerEnd:
 	case <-time.After(s.grace):
-		// The worker is gone and a descendant outside its group still holds
-		// the output: stop reading it, rather than hold the attempt, its
-		// working directory and the connector's shutdown open forever.
 		s.worker.CloseStdout()
 		<-s.readerEnd
 	}
-	return nil
 }
 
+// claim takes the turn in flight for the ending that is about to finish it,
+// so that no other ending can. It is what an ending does when it must be the
+// one reported and cannot be the first to get there — the unsafe verdict,
+// whose own act of ending the worker wakes the prompt's writer.
+//
+// It returns nil when there is no turn, or when another ending has it.
+func (s *session) claim() *turn {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.turn == nil || s.turn.settling {
+		return nil
+	}
+	s.turn.settling = true
+	return s.turn
+}
+
+// finish ends t with its result, unless an ending has claimed it.
 func (s *session) finish(t *turn, result driver.PromptResult, err error) {
+	s.end(t, result, err, false)
+}
+
+// settle ends a turn this ending claimed, giving up the claim in the same
+// step so that nothing can come between the two.
+func (s *session) settle(t *turn, result driver.PromptResult, err error) {
+	s.end(t, result, err, true)
+}
+
+func (s *session) end(t *turn, result driver.PromptResult, err error, claimed bool) {
 	s.mu.Lock()
-	if s.turn != t {
+	if s.turn != t || (t.settling && !claimed) {
 		s.mu.Unlock()
 		return
 	}
-	s.turn = nil
+	s.turn, t.settling = nil, false
 	s.mu.Unlock()
 	t.result, t.err = result, err
 	close(t.done)
@@ -785,21 +832,41 @@ func (s *session) threadStarted(id string) {
 }
 
 // unsafe ends the turn in flight with err and the process group.
+//
+// It runs on the policy check's own goroutine, which is the one ending that
+// is not the reader's, and the one ending that must be reported: an unsafe
+// session is reported as unsafe (invariant 3). So it claims the turn before
+// it touches the worker. Ending the worker closes the worker's stdin, and a
+// prompt still blocked writing that stdin — a worker that stopped reading
+// its input — wakes on the close and would otherwise finish this turn itself,
+// with its own ErrSessionEnded, over a session already known to be unsafe.
+//
+// Then the reader is given the worker's output. The refusals Codex put on
+// its stream are the reader's to record, and a worker that is gone is not a
+// stream that has been read: asking the turn what it refused at the instant
+// the worker died answers short of what the ledger goes on to hold, which
+// for this ending is the one thing it is for.
 func (s *session) unsafe(err error) {
-	s.mu.Lock()
-	t := s.turn
-	s.mu.Unlock()
+	t := s.claim()
+	s.worker.Terminate(0)
 	if t == nil {
-		s.worker.Terminate(0)
 		return
 	}
-	s.finishUnsafe(t, err)
+	s.readerDone()
+	s.lastWord()
+	s.settle(t, driver.PromptResult{Refusals: s.refusalsOf(t)}, err)
 }
 
 // finishUnsafe ends a turn whose session did not run under the policy it was
 // asked to: the worker goes first, then its last word is read, so the result
 // carries the refusals it made and logged before it was stopped, as every
 // other ending does.
+//
+// This is the reader's way of ending such a turn, and it says what the
+// turn's refusals are at the moment it runs — which on the reader is every
+// refusal of the turn, because the reader has parsed the stream up to here.
+// The policy check's own goroutine has unsafe, which waits for the reader
+// first, for that same reason.
 func (s *session) finishUnsafe(t *turn, err error) {
 	s.worker.Terminate(0)
 	s.lastWord()
@@ -824,6 +891,12 @@ func (s *session) failedVerification() error {
 // refusals it only logged. Whatever ends a turn ends it after this, so a
 // refusal Codex wrote on its way out is in the turn's result and not only in
 // the ledger.
+//
+// That is the stderr, which this reads itself. It says nothing about the
+// refusals Codex put on its JSON stream: those are the reader's to record,
+// and a worker that is gone is not a stream that has been read. An ending
+// that is not the reader's own waits for the reader too — see unsafe — before
+// it asks a turn what its refusals are.
 func (s *session) lastWord() {
 	if s.worker != nil {
 		select {
@@ -840,8 +913,12 @@ func (s *session) lastWord() {
 // being ended by the cancel anyway.
 func (s *session) finishCanceled(t *turn) {
 	s.lastWord()
-	// The turn's refusals are read after the worker's last word, so the
-	// result carries what the ledger carries.
+	// The turn's refusals are read after the worker's last word: the ones it
+	// only logged are read by lastWord just above, and the ones it put on its
+	// stream are already the turn's when the reader is the one ending it,
+	// which it is for every cancel the worker answered. A prompt whose write
+	// failed ends its own turn from a goroutine of its own, and that one can
+	// be short of a refusal the reader has not reached.
 	refusals := s.refusalsOf(t)
 	s.mu.Lock()
 	done := s.verifyDone

internal/connector/driver/codex/codex_test.go

@@ -11,6 +11,7 @@ import (
 	"slices"
 	"strconv"
 	"strings"
+	"sync"
 	"syscall"
 	"testing"
 	"time"
@@ -1032,6 +1033,190 @@ func TestAnUnsafeSessionStillReportsItsRefusals(t *testing.T) {
 	assert.Len(t, result.Refusals, 1, "the result carries what the ledger carries")
 }
 
+// heldLedger is a ledger that takes its time over the first refusal it is
+// given, as a ledger that writes to a database can. It holds the reader on
+// the line it is reading, with the rest of the stream unread behind it, until
+// the test lets it go.
+type heldLedger struct {
+	drivertest.Refusals
+	writing chan struct{} // one send per write begun
+	release chan struct{}
+	once    sync.Once
+}
+
+func (l *heldLedger) RecordRefusal(ctx context.Context, refusal driver.Refusal) error {
+	l.once.Do(func() {
+		l.writing <- struct{}{}
+		<-l.release
+	})
+	return l.Refusals.RecordRefusal(ctx, refusal)
+}
+
+// An unsafe verdict does not end the turn before the reader has read what the
+// worker put on its stream. The verdict is reached on the policy check's own
+// goroutine, which is the one ending that is not the reader's, and the
+// refusals on the stream are the reader's to record: a turn ended the instant
+// the worker died carries a result short of what the ledger goes on to hold,
+// and for a refusal that is the whole point of the ending.
+//
+// The ordering is made, not waited for, at both ends. The refusals are held
+// behind a barrier the test releases once it has seen the turn in flight, so
+// the reader cannot reach them before there is a turn for them to be on; and
+// the ledger is held over its first write, so the reader is provably stopped
+// on the first of the two with the second still unread when the verdict
+// lands.
+func TestAnUnsafeVerdictWaitsForWhatTheWorkerPutOnItsStream(t *testing.T) {
+	denial := func(id string) string {
+		return `{"type":"item.completed","item":{"id":"` + id + `","type":"mcp_tool_call","server":"other","tool":"write",` +
+			`"error":{"message":"MCP tool call requires approval, but approval policy is never"},"status":"failed"}}`
+	}
+	unsafe := safeTurnContext()
+	unsafe["approval_policy"] = "on-request"
+	for name, deaf := range map[string]bool{
+		"a worker that took its prompt": false,
+		// Ending the worker closes its stdin, which wakes a prompt still
+		// blocked writing it. That writer is an ending of its own, and it
+		// would report the session as ended rather than as unsafe.
+		"a worker that stopped reading its input": true,
+	} {
+		t.Run(name, func(t *testing.T) {
+			ledger := &heldLedger{writing: make(chan struct{}, 1), release: make(chan struct{})}
+			refusals := filepath.Join(t.TempDir(), "say-the-refusals")
+			h := newHarness(t, scenario{
+				TurnContext:            unsafe,
+				TurnContextAfterEvents: true,
+				Events:                 []string{`{"type":"turn.started"}`},
+				// Held until the turn is in flight: a refusal read before
+				// there is a turn is on no turn, and the test would be
+				// asking about something that never happened.
+				LateEvents: []string{denial("item_1"), denial("item_2")},
+				LateAfter:  refusals,
+				Deaf:       deaf,
+				// No turn.completed: the worker is still working when the
+				// check ends it, which is what this ending is for.
+				Hang: true,
+			})
+			cfg := h.config()
+			cfg.Refusals = ledger
+			s, err := h.drv.NewSession(context.Background(), cfg)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = s.Close() })
+
+			type answer struct {
+				result driver.PromptResult
+				err    error
+			}
+			// A prompt long enough to fill the pipe, so a worker that does
+			// not read it leaves the write blocked.
+			prompt := "Task 1. Event 2."
+			if deaf {
+				prompt = strings.Repeat("Event 1. ", 200_000)
+			}
+			answers := make(chan answer, 1)
+			go func() {
+				result, err := s.Prompt(context.Background(), prompt)
+				answers <- answer{result, err}
+			}()
+
+			// The turn is in flight, and only then are the refusals said.
+			session := s.(*session)
+			var inFlight chan struct{}
+			require.Eventually(t, func() bool {
+				session.mu.Lock()
+				defer session.mu.Unlock()
+				if session.turn == nil {
+					return false
+				}
+				inFlight = session.turn.done
+				return true
+			}, 20*time.Second, 5*time.Millisecond, "the turn never started")
+			require.NoError(t, os.WriteFile(refusals, nil, 0o600))
+			<-ledger.writing
+
+			// The turn must not finish while the reader is held — and if it
+			// does, that is the failure, so the ledger is let go the moment
+			// it does. The half second is only what lets the run finish when
+			// it does not.
+			go func() {
+				select {
+				case <-inFlight:
+				case <-time.After(500 * time.Millisecond):
+				}
+				close(ledger.release)
+			}()
+
+			got := <-answers
+			require.ErrorIs(t, got.err, driver.ErrUnsafeMode, "an unsafe session is reported as unsafe, whoever gets to the turn first")
+			require.ErrorContains(t, got.err, `Codex applied "on-request"`)
+			assert.Len(t, ledger.Recorded(), 2, "both refusals reach the ledger")
+			assert.Len(t, got.result.Refusals, 2,
+				"the result carries what the ledger carries, and carries %d of the %d recorded",
+				len(got.result.Refusals), len(ledger.Recorded()))
+		})
+	}
+}
+
+// A session closed under a turn still reports the refusals that turn made.
+// A worker that has stopped reading its input leaves the prompt's write
+// blocked, and closing the session fails that write: the turn is then
+// finished by the prompt's own writer rather than by the reader, and that
+// ending used to hand back an empty result over a turn that had refusals on
+// it.
+//
+// The reader is held in the ledger so that it cannot be the one to finish the
+// turn: the ending under test is the writer's, and nothing else.
+func TestASessionClosedUnderATurnStillReportsItsRefusals(t *testing.T) {
+	ledger := &heldLedger{writing: make(chan struct{}, 1), release: make(chan struct{})}
+	denial := `{"type":"item.completed","item":{"id":"item_1","type":"mcp_tool_call","server":"other","tool":"write",` +
+		`"error":{"message":"MCP tool call requires approval, but approval policy is never"},"status":"failed"}}`
+	refused := filepath.Join(t.TempDir(), "say-the-refusal")
+	h := newHarness(t, scenario{
+		TurnContext: safeTurnContext(),
+		Deaf:        true,
+		Hang:        true,
+		Events:      []string{`{"type":"turn.started"}`},
+		// Held until the turn is in flight: a refusal read before there is
+		// a turn is on no turn, and this test is about what a turn carries.
+		LateEvents: []string{denial},
+		LateAfter:  refused,
+	})
+	cfg := h.config()
+	cfg.Refusals = ledger
+	s, err := h.drv.NewSession(context.Background(), cfg)
+	require.NoError(t, err)
+
+	type answer struct {
+		result driver.PromptResult
+		err    error
+	}
+	answers := make(chan answer, 1)
+	go func() {
+		result, err := s.Prompt(context.Background(), strings.Repeat("Event 1. ", 200_000))
+		answers <- answer{result, err}
+	}()
+	waitDeaf(t, h)
+	session := s.(*session)
+	require.Eventually(t, func() bool {
+		session.mu.Lock()
+		defer session.mu.Unlock()
+		return session.turn != nil
+	}, 20*time.Second, 5*time.Millisecond, "the turn never started")
+	require.NoError(t, os.WriteFile(refused, nil, 0o600))
+	// The refusal is the turn's: refused puts it there before it writes it,
+	// and the write is where the reader is now held.
+	<-ledger.writing
+
+	closed := make(chan struct{})
+	go func() { _ = s.Close(); close(closed) }()
+	got := <-answers
+	close(ledger.release)
+	<-closed
+
+	require.ErrorIs(t, got.err, driver.ErrSessionEnded)
+	assert.Len(t, got.result.Refusals, 1,
+		"the turn's refusals are on its result, and it carries %d", len(got.result.Refusals))
+}
+
 // Codex logs its sandbox refusals and keeps writing: each one is recorded,
 // not only whatever it said last.
 func TestEveryRefusalCodexOnlyLogsIsRecorded(t *testing.T) {