diff --git a/internal/connector/driver/codex/codex.go b/internal/connector/driver/codex/codex.go index 0b6901504..eb1fad190 100644 --- a/internal/connector/driver/codex/codex.go +++ b/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 diff --git a/internal/connector/driver/codex/codex_test.go b/internal/connector/driver/codex/codex_test.go index 7a3b2c8d0..2747982dd 100644 --- a/internal/connector/driver/codex/codex_test.go +++ b/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) { diff --git a/internal/connector/driver/codex/fake_test.go b/internal/connector/driver/codex/fake_test.go index be04c2ca8..04f526ab3 100644 --- a/internal/connector/driver/codex/fake_test.go +++ b/internal/connector/driver/codex/fake_test.go @@ -45,6 +45,12 @@ type scenario struct { TurnContextAfterEvents bool `json:"turn_context_after_events"` // Events are written to stdout after thread.started. Events []string `json:"events"` + // LateEvents are written once LateAfter exists, and before any + // turn_context held back to the end. A test that needs the reader, or + // the turn, to be somewhere in particular when a line arrives puts it + // here and makes the file when it is there, rather than hoping. + LateEvents []string `json:"late_events"` + LateAfter string `json:"late_after"` // NoThread skips thread.started. NoThread bool `json:"no_thread"` // RunMCP starts each MCP server as Codex would and waits for it. @@ -185,6 +191,27 @@ func fakeCodex() int { for _, e := range sc.Events { fmt.Println(e) } + if sc.LateAfter != "" { + // The barrier is the point of these events: said without it, they + // prove nothing about the ordering the test claims, and a test that + // passes on a timer is the very thing this harness is for catching. + // A trigger that never comes is a broken test, and the fake says so + // rather than going ahead. + deadline := time.Now().Add(30 * time.Second) + for { + if _, err := os.Stat(sc.LateAfter); err == nil { + break + } + if time.Now().After(deadline) { + fmt.Fprintln(os.Stderr, "fake codex: the late events were never triggered:", sc.LateAfter) + return 2 + } + time.Sleep(5 * time.Millisecond) + } + for _, e := range sc.LateEvents { + fmt.Println(e) + } + } if sc.TurnContextAfterEvents { // The policy the driver judges is readable only once everything this // turn had to say is on the stream: the verdict, and the kill it