Wrap up passes that run out of time instead of discarding them

README.md

@@ -20,7 +20,7 @@ An LLM code reviewer for git changes, built on [Goose](https://github.com/aaif-g
 - **The judge** re-examines only the dropped candidates and restores the ones that are real defects. Use a different model than the coordinator, so the two don't share blind spots.
 - **Delivery** merges the coordinator's and judge's findings and applies a gate: by default P0–P2 with confidence of at least 0.8. Calibrate the gate per model; models differ in how they rate confidence.
 
-If a panelist fails, the coordinator covers its lens and the review is marked degraded. If the judge fails, the coordinator's findings still stand.
+A pass that is still working two minutes before its timeout is stopped and asked, in the same session, for the findings it has so far; a pass that ends without a JSON answer gets one more turn to give it. If a panelist still fails, the coordinator covers its lens and the review is marked degraded. If the judge fails, the coordinator's findings still stand.
 
 ## Usage
 

cmd/review/main.go

@@ -102,7 +102,8 @@ func firstLine(s string) string {
 // reviewBenchCmd implements https://github.com/review-bench/ReviewBench/blob/main/AGENT_CONTRACT.md.
 // The container is the sandbox, and the whole review must fit the 15-minute
 // limit per pull request: 7 minutes for the panel, then the coordinator and
-// judge share the rest of a 14-minute budget.
+// judge share the rest of a 14-minute budget. Each pass holds back its last
+// two minutes to report what it has if it runs long.
 func reviewBenchCmd(ctx context.Context) error {
 	env := func(k string) string { return strings.TrimSpace(os.Getenv(k)) }
 	for _, k := range []string{"RB_NWO", "RB_PR_NUMBER", "RB_BASE", "RB_HEAD", "RB_OUT"} {
@@ -140,7 +141,7 @@ func reviewBenchCmd(ctx context.Context) error {
 		CoordinatorTimeout: 6 * time.Minute,
 		JudgeTimeout:       4 * time.Minute,
 		Budget:             14 * time.Minute,
-		JudgeReserve:       2*time.Minute + 30*time.Second,
+		JudgeReserve:       90 * time.Second,
 	}
 	for name, d := range map[string]*time.Duration{
 		"REVIEW_PANELIST_TIMEOUT": &cfg.PanelistTimeout, "REVIEW_COORDINATOR_TIMEOUT": &cfg.CoordinatorTimeout,

review/goose.go

@@ -35,9 +35,14 @@ type GooseConfig struct {
 	// under StateDir/<role>. Otherwise they go to a temporary directory.
 	StateDir string
 	// SharedHome runs every pass in the caller's Goose home instead of a fresh
-	// one, and records sessions there, for hosts that stage their own Goose
-	// config and read usage from Goose's session store. StateDir is ignored.
+	// one, for hosts that stage their own Goose config and read usage from
+	// Goose's session store. StateDir is ignored.
 	SharedHome bool
+	// WrapUp is how much of a pass's timeout is held back for one more turn
+	// when the pass runs out of time or ends without a JSON answer. Zero uses
+	// two minutes, about what a long JSON answer takes to write. Passes with
+	// a timeout under twice WrapUp get no wrap-up.
+	WrapUp time.Duration
 	// Env is the process environment. Nil inherits the caller's environment.
 	Env []string
 }
@@ -49,7 +54,7 @@ type GooseRunner struct {
 	Config GooseConfig
 }
 
-func (r GooseRunner) args(model string) []string {
+func (r GooseRunner) args(model, session string, resume bool) []string {
 	c := r.Config
 	if model != "" {
 		c.Model = model
@@ -58,9 +63,11 @@ func (r GooseRunner) args(model string) []string {
 	if maxTurns == 0 {
 		maxTurns = 300
 	}
-	args := []string{"run"}
-	if !c.SharedHome {
-		args = append(args, "--no-session")
+	// Every pass records a named session so it can be resumed for a wrap-up
+	// turn; without SharedHome the session lives in the pass's private home.
+	args := []string{"run", "--name", session}
+	if resume {
+		args = append(args, "--resume")
 	}
 	args = append(args, "--no-profile", "--with-builtin", "developer",
 		"--output-format", "json", "--max-turns", fmt.Sprint(maxTurns), "--instructions", "-")
@@ -108,6 +115,18 @@ func (r GooseRunner) env(root, effort string) []string {
 	return out
 }
 
+const (
+	wrapUpTimeUp  = "Time is up. Do not call any more tools. Reply now with only the final JSON object your instructions describe, covering what you have verified so far."
+	wrapUpNoJSON  = "Your last reply did not contain the required JSON object. Do not call any more tools. Reply now with only that JSON object."
+	defaultWrapUp = 2 * time.Minute
+)
+
+var errPassDeadline = errors.New("pass deadline")
+
+// Run runs the pass until its timeout less WrapUp. If the pass is still
+// working then, or finishes without a JSON object, Run resumes its session
+// for one more turn that asks for the answer, so a slow pass reports what it
+// has instead of nothing.
 func (r GooseRunner) Run(ctx context.Context, pass Pass) (string, error) {
 	root := ""
 	if r.Config.SharedHome {
@@ -130,14 +149,58 @@ func (r GooseRunner) Run(ctx context.Context, pass Pass) (string, error) {
 		ctx, cancel = context.WithTimeout(ctx, pass.Timeout)
 		defer cancel()
 	}
+	wrapUp := r.Config.WrapUp
+	if wrapUp == 0 {
+		wrapUp = defaultWrapUp
+	}
+	work := ctx
+	if pass.Timeout > 2*wrapUp {
+		var cancel context.CancelFunc
+		work, cancel = context.WithTimeout(ctx, pass.Timeout-wrapUp)
+		defer cancel()
+	}
+	session := fmt.Sprintf("review-%s-%d", pass.Role, time.Now().UnixNano())
+
+	out, err := r.exec(work, root, session, false, pass.Prompt, pass)
+	var nudge string
+	switch {
+	case errors.Is(err, errPassDeadline) && work != ctx && ctx.Err() == nil:
+		nudge = wrapUpTimeUp
+	case errors.Is(err, errPassDeadline):
+		return "", fmt.Errorf("%s pass exceeded %s", pass.Role, pass.Timeout)
+	case err != nil:
+		return "", err
+	default:
+		if _, jsonErr := extractJSONObject(out); jsonErr == nil {
+			return out, nil
+		}
+		nudge = wrapUpNoJSON
+	}
+
+	wrapCtx, cancel := context.WithTimeout(ctx, wrapUp)
+	defer cancel()
+	wrapped, wrapErr := r.exec(wrapCtx, root, session, true, nudge, pass)
+	if nudge == wrapUpNoJSON {
+		if wrapErr != nil {
+			return out, nil
+		}
+		return wrapped, nil
+	}
+	if wrapErr != nil {
+		return "", fmt.Errorf("%s pass exceeded %s and its wrap-up failed: %w", pass.Role, pass.Timeout-wrapUp, wrapErr)
+	}
+	return wrapped, nil
+}
+
+func (r GooseRunner) exec(ctx context.Context, root, session string, resume bool, input string, pass Pass) (string, error) {
 	bin := r.Config.Bin
 	if bin == "" {
 		bin = "goose"
 	}
-	cmd := exec.CommandContext(ctx, bin, r.args(pass.Model)...)
+	cmd := exec.CommandContext(ctx, bin, r.args(pass.Model, session, resume)...)
 	cmd.Dir = pass.RepoDir
 	cmd.Env = r.env(root, pass.Effort)
-	cmd.Stdin = strings.NewReader(pass.Prompt)
+	cmd.Stdin = strings.NewReader(input)
 	var stdout, stderr bytes.Buffer
 	cmd.Stdout = &stdout
 	cmd.Stderr = &stderr
@@ -147,7 +210,7 @@ func (r GooseRunner) Run(ctx context.Context, pass Pass) (string, error) {
 
 	runErr := cmd.Run()
 	if errors.Is(ctx.Err(), context.DeadlineExceeded) {
-		return "", fmt.Errorf("%s pass exceeded %s", pass.Role, pass.Timeout)
+		return "", errPassDeadline
 	}
 	if runErr != nil {
 		return "", fmt.Errorf("%s pass: goose: %w: %s", pass.Role, runErr, tail(stderr.String()+stdout.String(), 4000))

review/review_test.go

@@ -258,8 +258,8 @@ func TestGooseFinalMessage(t *testing.T) {
 
 func TestGooseRunIsIsolated(t *testing.T) {
 	r := GooseRunner{Config: GooseConfig{Provider: "openai", Model: "m", Effort: "high", Env: []string{"GOOSE_MODE=approve", "CONTEXT_FILE_NAMES=[\"AGENTS.md\"]", "KEEP=1"}}}
-	args := strings.Join(r.args("override"), " ")
-	for _, want := range []string{"--no-session --no-profile --with-builtin developer", "--output-format json", "--instructions -", "--provider openai --model override"} {
+	args := strings.Join(r.args("override", "s1", false), " ")
+	for _, want := range []string{"run --name s1 --no-profile --with-builtin developer", "--output-format json", "--instructions -", "--provider openai --model override"} {
 		if !strings.Contains(args, want) {
 			t.Errorf("args missing %q: %s", want, args)
 		}
@@ -277,15 +277,53 @@ func TestGooseRunIsIsolated(t *testing.T) {
 
 func TestGooseSharedHomeKeepsCallerHomeAndSessions(t *testing.T) {
 	r := GooseRunner{Config: GooseConfig{SharedHome: true, Env: []string{"GOOSE_PATH_ROOT=/caller", "XDG_DATA_HOME=/data"}}}
-	if args := strings.Join(r.args(""), " "); strings.Contains(args, "--no-session") {
-		t.Errorf("shared home must record sessions: %s", args)
-	}
 	env := strings.Join(r.env("", ""), "\n")
 	if !strings.Contains(env, "GOOSE_PATH_ROOT=/caller") || !strings.Contains(env, "XDG_DATA_HOME=/data") {
 		t.Errorf("shared home must keep the caller's Goose paths:\n%s", env)
 	}
-	if args := strings.Join(GooseRunner{}.args(""), " "); !strings.Contains(args, "--no-session") {
-		t.Errorf("default runner must not record sessions: %s", args)
+}
+
+// fakeGoose answers like `goose run --output-format json`. A first run in
+// mode "slow" outlives the soft deadline and one in mode "prose" ends without
+// JSON; a resumed run answers with JSON. It logs each call's args and input.
+const fakeGoose = `#!/bin/sh
+log="$FAKE_DIR/calls"
+input=$(cat)
+printf '%s\t%s\n' "$*" "$input" >> "$log"
+answer() { printf '  banner\n{\n"messages":[{"role":"assistant","content":[{"type":"text","text":"%s"}]}],"metadata":{"status":"completed"}}\n' "$1"; }
+case " $* " in
+*" --resume "*) answer '{\"findings\":[\"wrapped\"]}' ;;
+*) if [ "$FAKE_MODE" = slow ]; then sleep 30; fi; answer 'I reviewed it and it looks fine.' ;;
+esac
+`
+
+func TestGooseWrapsUpSlowAndJSONlessPasses(t *testing.T) {
+	bin := filepath.Join(t.TempDir(), "goose")
+	if err := os.WriteFile(bin, []byte(fakeGoose), 0o755); err != nil {
+		t.Fatal(err)
+	}
+	for _, tc := range []struct{ mode, nudge string }{{"slow", wrapUpTimeUp}, {"prose", wrapUpNoJSON}} {
+		t.Run(tc.mode, func(t *testing.T) {
+			dir := t.TempDir()
+			r := GooseRunner{Config: GooseConfig{Bin: bin, WrapUp: time.Second, Env: []string{"PATH=" + os.Getenv("PATH"), "FAKE_DIR=" + dir, "FAKE_MODE=" + tc.mode}}}
+			started := time.Now()
+			out, err := r.Run(context.Background(), Pass{Role: "p", Prompt: "review", RepoDir: dir, Timeout: 4 * time.Second})
+			if err != nil || out != `{"findings":["wrapped"]}` {
+				t.Fatalf("got %q, %v", out, err)
+			}
+			if elapsed := time.Since(started); elapsed > 4*time.Second {
+				t.Errorf("took %s, past the 4s timeout", elapsed)
+			}
+			raw, _ := os.ReadFile(filepath.Join(dir, "calls"))
+			calls := strings.Split(strings.TrimSpace(string(raw)), "\n")
+			if len(calls) != 2 {
+				t.Fatalf("want a run and one wrap-up, got %q", calls)
+			}
+			first, second := strings.Fields(calls[0]), strings.Fields(calls[1])
+			if first[2] != second[2] || !strings.Contains(calls[1], "--resume") || !strings.HasSuffix(calls[1], "\t"+tc.nudge) {
+				t.Errorf("wrap-up must resume the same session with %q:\n%s\n%s", tc.nudge, calls[0], calls[1])
+			}
+		})
 	}
 }