A run's dollar limit holds while a worker is still working

docs/changes/unreleased/1268-a-dollar-limit-holds-while-a-worker-works.md

@@ -0,0 +1,15 @@
+---
+kind: fixed
+title: a run's dollar limit now holds while a worker is still working
+pr: 1268
+surface: [chat, engine]
+invalidates:
+  - "A run's dollar limit was read only when a worker came home. A worker's spend reached the run's count in its report, so one long worker could pass the limit many times over before anything counted it, and the limit was a word about the run after it instead of a hold on this one."
+  - "A worker now says what it has spent as each model call is paid for, and the run adds it while the worker is still working. When the sum reaches the limit the run launches nothing new, ends the work in flight, absorbs every ending, and answers `a limit you set stopped it`. That is the time limit's road, so the sentence in the entry for the time limit that said a cost limit lets work in flight finish is no longer true."
+  - "A worker never waits on the run to report its spend. It writes a running figure and leaves a token the run may miss; the figure is the whole of what the worker has spent, so a missed token loses nothing, and a run that is draining is never held open by a worker that has something to say."
+  - "What a run counts as spent never goes down. A worker whose task the store cancelled under it handed back no spend the run counted, so those dollars were missing from the run's figure and from the limit. Its words and steps are still dropped; its dollars are now counted, whatever way the task ended. The same holds for a worker the run outlived, ended when the tree finished under it: its return was dropped whole, and its dollars are now settled."
+  - "The call that reaches the limit is already paid for, so a run can end a little over its figure. A run with no dollar limit counts as it always did, at each return."
+---
+
+`--max-cost` still does not reach a run: a run's dollar limit is the conversation's
+spend ceiling.

internal/manual/chat/starting-codeaf.md

@@ -141,9 +141,12 @@ comes home the ordinary way. The run's row ends `incomplete` and says
 `a limit you set stopped it`, the same sentence a run says when it reaches the
 conversation's spend ceiling. A task that was cut part way through reads `incomplete` too.
 
-The dollar limit on a run is the conversation's spend ceiling, not `--max-cost`. A run
-that reaches it lets the work already going finish, because that money is already
-committed; a time limit cannot wait, so it ends that work.
+The dollar limit on a run is the conversation's spend ceiling, not `--max-cost`. A run's
+spend is counted as each model call is paid for, while its tasks are still working, so
+one long task cannot carry the run far past the figure. A run that reaches it ends the
+work still going and says `a limit you set stopped it`, just as when its elapsed-time
+limit is up. The call that reached the limit is already paid for, so the run can end a
+little over it.
 
 ## Leaving it running on its own · leaving a headless run going with a budget · --once yolo · no screen · unattended · overnight · nobody watching
 

internal/run/bashworker.go

@@ -127,6 +127,13 @@ func (w *BashWorker) Run(ctx context.Context, task plandb.Task) (Report, error)
 	// did call a tool resets the run to one — its own trailing words are the
 	// first of the new run — and the fourth in a row fails the task.
 	noAction := 0
+	// banked is the last spend figure this worker told its run. THE RUN'S DOLLAR
+	// LIMIT IS READ WHILE THE WORKER WORKS, so the figure has to move when a
+	// call is paid for and not when a turn ends: a turn is the worker's whole
+	// working life, and its usage arrives with its ending. The agent's own
+	// running total is the one account that moves per paid call, so it is read
+	// as each event comes by and told to the run whenever it has risen.
+	banked := 0.0
 	for {
 		events, err := agent.Submit(runCtx, brief)
 		if err != nil {
@@ -145,6 +152,10 @@ func (w *BashWorker) Run(ctx context.Context, task plandb.Task) (Report, error)
 			ending     storeEnding
 		)
 		for event := range events {
+			if spent := agent.Usage().CostUSD; spent > banked {
+				banked = spent
+				bankSpend(runCtx, spent)
+			}
 			switch event.Kind {
 			case session.EventToolBegin:
 				// A LIVE STEP IS TRUE ONLY WHILE ITS COMMAND RUNS. The begin event is

internal/run/bashworker_test.go

@@ -22,6 +22,7 @@ import (
 	"path/filepath"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -509,3 +510,65 @@ func TestBashWorkerEndsItsLoopAtTheStepCap(t *testing.T) {
 		t.Fatalf("the ending reads %d steps with %q, want the cap's count and reason", end.Steps, end.Reason)
 	}
 }
+
+// THE WORKER SAYS WHAT IT HAS SPENT WHILE IT IS STILL WORKING, not only when it
+// comes home. A run's dollar limit is read from these figures, so the first
+// paid call must be told to the run before the worker's turn is over, every
+// figure must be the whole of what the worker has spent so far, and the last
+// one must be the figure the report carries: two accounts of one worker's money
+// would drift.
+//
+// THE SEAT HOLDS THE SECOND CALL UNTIL THE FIRST IS REPORTED, which is what makes
+// "while it is still working" a fact and not a race: the figures are readings of
+// a running total and a fast worker may fold several calls into one reading, so
+// their count is not the property. A worker that reports only at its turn's
+// end never frees the seat, and the first figure it then gives is not the first
+// call's.
+func TestBashWorkerReportsItsSpendWhileItIsStillWorking(t *testing.T) {
+	t.Setenv("CODEAF_TASK_BELT", "bash")
+	t.Setenv("CODEAF_PLANDB_BIN", stubCLI(t))
+	store := runOpenStore(t)
+	firstReported := make(chan struct{})
+	var calls atomic.Int32
+	seat := &seat{ever: func(ctx context.Context, _ []ai.Message) (*ai.Response, error) {
+		if calls.Add(1) > 1 {
+			select {
+			case <-firstReported:
+			case <-ctx.Done():
+			case <-time.After(10 * time.Second):
+			}
+		}
+		reply := toolReply(`{"command":"true"}`)
+		cost := 0.25
+		reply.Usage.Cost = &cost
+		return reply, nil
+	}}
+	worker := run.NewBashWorker(store, t.TempDir(), "test/model", seat)
+	var mu sync.Mutex
+	var figures []float64
+	ctx := run.WithSpendBank(run.WithStepsPerTask(runContext(t), 3), func(usd float64) {
+		mu.Lock()
+		figures = append(figures, usd)
+		first := len(figures) == 1
+		mu.Unlock()
+		if first {
+			close(firstReported)
+		}
+	})
+
+	report, _ := worker.Run(ctx, *store.Task(store.RootID()))
+
+	mu.Lock()
+	defer mu.Unlock()
+	if len(figures) == 0 || figures[0] != 0.25 {
+		t.Fatalf("spend figures = %v, want the first paid call's 0.25 told before the second call", figures)
+	}
+	for i := 1; i < len(figures); i++ {
+		if figures[i] <= figures[i-1] {
+			t.Fatalf("spend figures are not a rising whole: %v", figures)
+		}
+	}
+	if last := figures[len(figures)-1]; last != report.USD {
+		t.Fatalf("the last figure reported = %v and the report carries %v, want one account", last, report.USD)
+	}
+}

internal/run/cost_limit_test.go

@@ -0,0 +1,267 @@
+package run
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/Agent-Field/codeaf/internal/plandb"
+)
+
+func TestCostLimitEndsAndAbsorbsWorkingWorker(t *testing.T) {
+	store, err := plandb.Open(t.TempDir()+"/plan.json", "cost-limit", "root", "root", "run until the limit")
+	if err != nil {
+		t.Fatalf("open store: %v", err)
+	}
+	defer store.Close()
+
+	ended := make(chan struct{})
+	var launches atomic.Int32
+	factory := func(task plandb.Task) Worker {
+		launches.Add(1)
+		return workerFunc(func(ctx context.Context, task plandb.Task) (Report, error) {
+			bankSpend(ctx, 1.25)
+			<-ctx.Done()
+			if _, err := store.AddMany([]plandb.TaskSpec{{ID: "after-limit", Title: "must not launch", ParentID: task.ID}}); err != nil {
+				t.Errorf("add post-limit task: %v", err)
+			}
+			close(ended)
+			return Report{USD: 1.25}, ctx.Err()
+		})
+	}
+	supervisor := NewSupervisor(store, t.TempDir(), 2, Limits{CostUSD: 1}, factory)
+	if outcome := supervisor.Run(context.Background()); outcome != OutcomeLimit {
+		t.Fatalf("outcome = %q, want %q", outcome, OutcomeLimit)
+	}
+	select {
+	case <-ended:
+	default:
+		t.Fatal("Run returned before its in-flight worker ended and was absorbed")
+	}
+	if supervisor.spent != 1.25 {
+		t.Fatalf("run spend = %v, want 1.25 exactly once", supervisor.spent)
+	}
+	if supervisor.inFlight != 0 || len(supervisor.cancels) != 0 {
+		t.Fatalf("run retained workers: inFlight=%d cancels=%d", supervisor.inFlight, len(supervisor.cancels))
+	}
+	if got := launches.Load(); got != 1 {
+		t.Fatalf("worker launches = %d, want only the worker cut at the limit", got)
+	}
+	if task := store.Task("after-limit"); task == nil || task.Status != plandb.StatusReady {
+		t.Fatalf("post-limit task = %#v, want ready and unlaunched", task)
+	}
+}
+
+func TestCostLimitSumsWorkingWorkers(t *testing.T) {
+	store, err := plandb.Open(t.TempDir()+"/plan.json", "cost-limit-sum", "root", "root", "run until the limit")
+	if err != nil {
+		t.Fatalf("open store: %v", err)
+	}
+	defer store.Close()
+	if _, err := store.AddMany([]plandb.TaskSpec{{ID: "peer", Title: "peer", ParentID: store.RootID()}}); err != nil {
+		t.Fatalf("add peer: %v", err)
+	}
+
+	bothStarted := make(chan struct{})
+	var started atomic.Int32
+	var ended atomic.Int32
+	factory := func(task plandb.Task) Worker {
+		return workerFunc(func(ctx context.Context, task plandb.Task) (Report, error) {
+			if started.Add(1) == 2 {
+				close(bothStarted)
+			}
+			<-bothStarted
+			bankSpend(ctx, 0.6)
+			<-ctx.Done()
+			ended.Add(1)
+			return Report{USD: 0.6}, ctx.Err()
+		})
+	}
+	supervisor := NewSupervisor(store, t.TempDir(), 2, Limits{CostUSD: 1}, factory)
+	if outcome := supervisor.Run(context.Background()); outcome != OutcomeLimit {
+		t.Fatalf("outcome = %q, want %q", outcome, OutcomeLimit)
+	}
+	if got := ended.Load(); got != 2 {
+		t.Fatalf("absorbed worker returns = %d, want 2", got)
+	}
+	if supervisor.spent != 1.2 {
+		t.Fatalf("run spend = %v, want 1.2 exactly once", supervisor.spent)
+	}
+	if supervisor.inFlight != 0 || len(supervisor.cancels) != 0 {
+		t.Fatalf("run retained workers: inFlight=%d cancels=%d", supervisor.inFlight, len(supervisor.cancels))
+	}
+}
+
+func TestSpendBankWithoutLimitLetsWorkerReturn(t *testing.T) {
+	store, err := plandb.Open(t.TempDir()+"/plan.json", "unlimited", "root", "root", "finish")
+	if err != nil {
+		t.Fatalf("open store: %v", err)
+	}
+	defer store.Close()
+	factory := func(task plandb.Task) Worker {
+		return workerFunc(func(ctx context.Context, task plandb.Task) (Report, error) {
+			bankSpend(ctx, 3)
+			return Report{Result: "done", USD: 3}, nil
+		})
+	}
+	supervisor := NewSupervisor(store, t.TempDir(), 1, Limits{}, factory)
+	if outcome := supervisor.Run(context.Background()); outcome != OutcomeDone {
+		t.Fatalf("outcome = %q, want %q", outcome, OutcomeDone)
+	}
+	if supervisor.spent != 3 {
+		t.Fatalf("run spend = %v, want returned 3 exactly once", supervisor.spent)
+	}
+}
+
+func TestCostLimitCountsReturnedWorkerBelowCeilingOnce(t *testing.T) {
+	store, err := plandb.Open(t.TempDir()+"/plan.json", "cost-limit-under", "root", "root", "finish below the limit")
+	if err != nil {
+		t.Fatalf("open store: %v", err)
+	}
+	defer store.Close()
+
+	factory := func(task plandb.Task) Worker {
+		return workerFunc(func(ctx context.Context, task plandb.Task) (Report, error) {
+			bankSpend(ctx, 0.4)
+			bankSpend(ctx, 0.9)
+			return Report{Result: "done", USD: 0.9}, nil
+		})
+	}
+	supervisor := NewSupervisor(store, t.TempDir(), 1, Limits{CostUSD: 1}, factory)
+	if outcome := supervisor.Run(context.Background()); outcome != OutcomeDone {
+		t.Fatalf("outcome = %q, want %q", outcome, OutcomeDone)
+	}
+	if supervisor.spent != 0.9 {
+		t.Fatalf("run spend = %v, want cumulative 0.9 exactly once", supervisor.spent)
+	}
+}
+
+// A WORKER NEVER WAITS ON THE RUN TO SAY WHAT IT SPENT. The road this pins is
+// the one on which nobody is listening: the tree finished while one worker was
+// still out, so the run drains, and drain waits for that worker and reads
+// nothing. A worker that had to be heard before it could go on would hold the
+// run open for good, which is an engine that never answers. The worker here
+// reports its spend far more often than any buffer is deep, after its context
+// has ended, and the run must still come home.
+func TestAWorkerReportingSpendWhileTheRunDrainsDoesNotHoldTheRunOpen(t *testing.T) {
+	store, err := plandb.Open(t.TempDir()+"/plan.json", "cost-limit-drain", "root", "root", "finish while a worker is still out")
+	if err != nil {
+		t.Fatalf("open store: %v", err)
+	}
+	defer store.Close()
+	if _, err := store.AddMany([]plandb.TaskSpec{{ID: "peer", Title: "peer", ParentID: store.RootID()}}); err != nil {
+		t.Fatalf("add peer: %v", err)
+	}
+
+	peerEnded := make(chan struct{})
+	var reports atomic.Int32
+	factory := func(task plandb.Task) Worker {
+		if task.ID == "peer" {
+			return workerFunc(func(ctx context.Context, task plandb.Task) (Report, error) {
+				if _, err := store.Done(task.ID, task.ID, "the peer wrote its own ending", nil, nil); err != nil {
+					t.Errorf("peer done: %v", err)
+				}
+				close(peerEnded)
+				<-ctx.Done()
+				for i := 1; i <= 64; i++ {
+					bankSpend(ctx, float64(i)/1000)
+					reports.Add(1)
+				}
+				return Report{USD: 0.064}, ctx.Err()
+			})
+		}
+		return workerFunc(func(ctx context.Context, task plandb.Task) (Report, error) {
+			select {
+			case <-peerEnded:
+			case <-ctx.Done():
+				return Report{}, ctx.Err()
+			}
+			return Report{Result: "the tree is finished"}, nil
+		})
+	}
+	supervisor := NewSupervisor(store, t.TempDir(), 2, Limits{CostUSD: 100}, factory)
+	answered := make(chan Outcome, 1)
+	go func() { answered <- supervisor.Run(context.Background()) }()
+	select {
+	case outcome := <-answered:
+		if outcome != OutcomeDone {
+			t.Fatalf("outcome = %q, want %q", outcome, OutcomeDone)
+		}
+	case <-time.After(20 * time.Second):
+		t.Fatalf("the run never answered: a worker reporting its spend is holding the drain open (%d of 64 reports made)", reports.Load())
+	}
+	if got := reports.Load(); got != 64 {
+		t.Fatalf("spend reports made = %d, want all 64", got)
+	}
+	// AND THE WORKER THE RUN OUTLIVED WAS STILL PAID FOR. Its words are dropped
+	// with its return; its dollars are in the run's count, once.
+	if supervisor.spent != 0.064 {
+		t.Fatalf("run spend = %v, want the outlived worker's 0.064 counted exactly once", supervisor.spent)
+	}
+	if supervisor.inFlight != 0 || len(supervisor.finished) != 0 {
+		t.Fatalf("the run left work behind it: inFlight=%d, returns unread=%d", supervisor.inFlight, len(supervisor.finished))
+	}
+}
+
+// WHAT A RUN COUNTS AS SPENT NEVER GOES DOWN, and it includes every paid call
+// whatever way its task ended. A person's limit is about dollars, not about how
+// the work came out: a task the store cancelled under its worker hands back no
+// work the run counts, and its dollars were still paid. With a limit the run
+// has already counted them on the way and must not take them back at the
+// return; without one the return is where they are counted.
+func TestSpendOfAWorkerWhoseTaskTheStoreCancelledStaysInTheRunsCount(t *testing.T) {
+	for _, test := range []struct {
+		name  string
+		limit float64
+	}{
+		{name: "with a dollar limit", limit: 100},
+		{name: "with no dollar limit", limit: 0},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			store, err := plandb.Open(t.TempDir()+"/plan.json", "cancelled-spend", "root", "root", "cancel one task part way")
+			if err != nil {
+				t.Fatalf("open store: %v", err)
+			}
+			defer store.Close()
+			if _, err := store.AddMany([]plandb.TaskSpec{{ID: "peer", Title: "peer", ParentID: store.RootID()}}); err != nil {
+				t.Fatalf("add peer: %v", err)
+			}
+			peerPaid := make(chan struct{})
+			var paidOnce atomic.Bool
+			factory := func(task plandb.Task) Worker {
+				if task.ID == "peer" {
+					return workerFunc(func(ctx context.Context, task plandb.Task) (Report, error) {
+						bankSpend(ctx, 0.30)
+						if paidOnce.CompareAndSwap(false, true) {
+							close(peerPaid)
+						}
+						<-ctx.Done()
+						return Report{Result: "late", Steps: 4, USD: 0.30}, ctx.Err()
+					})
+				}
+				return workerFunc(func(ctx context.Context, task plandb.Task) (Report, error) {
+					select {
+					case <-peerPaid:
+					case <-ctx.Done():
+						return Report{}, ctx.Err()
+					}
+					// Already cancelled on a wake of the root: the first cancellation stands.
+					_, _ = store.Cancel("peer", "stopped by hand")
+					return Report{Result: "the root's word"}, nil
+				})
+			}
+			supervisor := NewSupervisor(store, t.TempDir(), 2, Limits{CostUSD: test.limit}, factory)
+			supervisor.Run(context.Background())
+			if peer := store.Task("peer"); peer == nil || peer.Status != plandb.StatusCancelled {
+				t.Fatalf("peer = %#v, want cancelled by the store", peer)
+			}
+			if supervisor.spent != 0.30 {
+				t.Fatalf("run spend = %v, want the cancelled task's 0.30 counted exactly once", supervisor.spent)
+			}
+			if supervisor.steps != 0 {
+				t.Fatalf("run steps = %d, want none of the cancelled task's steps", supervisor.steps)
+			}
+		})
+	}
+}