diff --git a/docs/changes/unreleased/1268-a-dollar-limit-holds-while-a-worker-works.md b/docs/changes/unreleased/1268-a-dollar-limit-holds-while-a-worker-works.md new file mode 100644 index 0000000000..3278d7d298 --- /dev/null +++ b/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. diff --git a/internal/manual/chat/starting-codeaf.md b/internal/manual/chat/starting-codeaf.md index 9c11c7b533..ff460212f3 100644 --- a/internal/manual/chat/starting-codeaf.md +++ b/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 diff --git a/internal/run/bashworker.go b/internal/run/bashworker.go index e59e2fe525..25049318f1 100644 --- a/internal/run/bashworker.go +++ b/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 diff --git a/internal/run/bashworker_test.go b/internal/run/bashworker_test.go index 251899548d..0a29b5937e 100644 --- a/internal/run/bashworker_test.go +++ b/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) + } +} diff --git a/internal/run/cost_limit_test.go b/internal/run/cost_limit_test.go new file mode 100644 index 0000000000..0f59ebf0ed --- /dev/null +++ b/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) + } + }) + } +} diff --git a/internal/run/run.go b/internal/run/run.go index 84c233be44..fc58afe73c 100644 --- a/internal/run/run.go +++ b/internal/run/run.go @@ -100,9 +100,28 @@ type Supervisor struct { // and drain is its only reader: every road out of Run waits on it before // answering, so the store, the working copy and the process are the // caller's alone the moment the run is over. - workers sync.WaitGroup - inFlight int - spent float64 + workers sync.WaitGroup + inFlight int + spent float64 + // THE DOLLAR LIMIT IS READ WHILE THE WORK IS GOING, not only when a worker + // comes home. live is what each worker in flight says it has banked so far, + // a cumulative figure its own goroutine writes under liveMu; counted is how + // much of that figure the loop has already added to spent, and it belongs to + // the loop alone. Both are keyed by task, both are cumulative, and a + // worker's return reconciles them against its report, so a dollar is + // counted once whether the loop saw it on the way or only at the end. + // + // A WORKER NEVER WAITS ON THE RUN TO SAY WHAT IT SPENT. It writes its figure + // and drops a token into liveMoved if there is room; a token that does not + // fit is lost and nothing else is, because the figure is cumulative and the + // next read takes all of it. That matters most on the way out: drain waits + // for workers and reads nothing, so a worker that had to be heard before it + // could go on would hold the run open for good. + liveMu sync.Mutex + live map[string]float64 + counted map[string]float64 + liveMoved chan struct{} + nodes int steps int rootResult string @@ -148,6 +167,7 @@ func NewSupervisor(store *plandb.Store, workspace string, slots int, limits Limi factory: factory, after: time.After, finished: make(chan workerReturn, slots+1), + liveMoved: make(chan struct{}, 1), cancels: make(map[string]context.CancelFunc), checkOf: make(map[string]string), } @@ -183,6 +203,10 @@ func (s *Supervisor) Run(ctx context.Context) Outcome { return OutcomeCannotRun } s.spent = 0 + s.liveMu.Lock() + s.live = make(map[string]float64) + s.liveMu.Unlock() + s.counted = make(map[string]float64) s.nodes = 0 s.steps = 0 s.rootResult = "" @@ -219,6 +243,8 @@ func (s *Supervisor) Run(ctx context.Context) Outcome { case ret := <-s.finished: s.inFlight-- s.absorb(ret) + case <-s.liveMoved: + s.countLiveSpend() case <-timer.C: timer.Reset(passInterval) case <-elapsed: @@ -226,9 +252,9 @@ func (s *Supervisor) Run(ctx context.Context) Outcome { // spend counter marks, so the next pass launches nothing and answers // the existing OutcomeLimit. // - // THE DIFFERENCE IS THE WORK IN FLIGHT. A cost limit lets it finish, - // because what it will spend is already committed; a time limit - // cannot, because the time is gone. So every worker is ended here, and + // TIME ENDS THE WORK IN FLIGHT, AS DOLLARS DO ([countLiveSpend]): the + // time is gone whatever a worker was in the middle of. So every worker + // is ended here, and // ITS ENDING IS ABSORBED, the way the caller's wall below reads its // endings: a return that was dropped would leave its task claimed and // reading as running on a run that is over, and what the worker spent @@ -349,6 +375,17 @@ func (s *Supervisor) launch(ctx context.Context, task plandb.Task, wake string) taskCtx = WithWakeClause(taskCtx, wake) } s.cancels[task.ID] = cancel + s.counted[task.ID] = 0 + if s.limits.CostUSD > 0 { + // ONLY A RUN WITH A DOLLAR LIMIT LISTENS. A worker launched again for a + // wake starts its own count from nothing, so the figure a former run of + // the same task left behind is cleared before this one can write. + id := task.ID + s.liveMu.Lock() + s.live[id] = 0 + s.liveMu.Unlock() + taskCtx = WithSpendBank(taskCtx, func(usd float64) { s.bankLive(id, usd) }) + } s.inFlight++ s.nodes++ // THE GOROUTINE IS COUNTED BEFORE IT STARTS, so a drain that begins the @@ -380,15 +417,105 @@ func (s *Supervisor) launch(ctx context.Context, task plandb.Task, wake string) // a worker whose context a pass already ended — a cancelled task, the caller's // wall — is in the map no more and has only to be waited for. The endings are // dropped rather than absorbed: the run is over, the store already carries the -// ending that stands, and the return of a worker the run outlived has no +// ending that stands, and the words of a worker the run outlived have no // reader — the same reading absorb makes of a worker whose task the store // cancelled under it. +// +// THE DOLLARS ARE NOT DROPPED, for absorb's own reason: what a run counts as +// spent includes every paid call whatever way its task ended, and a worker the +// run outlived was paid for like any other. Every worker has returned by the +// time the wait is over and each left exactly one return in the channel, which +// is deep enough to hold them all, so they are read here without waiting and +// only their spend is settled. It also leaves the channel empty, so a return +// from this run can never be read as one of the next. func (s *Supervisor) drain() { for id, cancel := range s.cancels { cancel() delete(s.cancels, id) } s.workers.Wait() + for { + select { + case ret := <-s.finished: + s.inFlight-- + s.settleSpend(ret) + default: + return + } + } +} + +// bankLive is the one thing a worker's goroutine does to the run's account: it +// writes what the worker has banked so far and lets the loop know there is +// something to read. IT NEVER BLOCKS, for the reason the fields state. The +// figure only ever rises, so a late or repeated call cannot take a dollar back. +func (s *Supervisor) bankLive(id string, usd float64) { + s.liveMu.Lock() + if usd > s.live[id] { + s.live[id] = usd + } + s.liveMu.Unlock() + select { + case s.liveMoved <- struct{}{}: + default: + } +} + +// countLiveSpend adds to the run's spend whatever the workers in flight have +// banked since the loop last looked, and ends the work in flight the moment the +// sum reaches the person's dollar limit. +// +// A DOLLAR LIMIT MEANS THE RUN'S DOLLARS, NOT THE DOLLARS OF WORKERS THAT HAVE +// COME HOME. Read only at a return, one long worker could pass the limit many +// times over before anyone counted, and the limit would be a word about the +// run after it rather than a hold on this one. The ending is the time limit's: +// the limit state is marked, every worker's context is ended, the loop goes on +// reading their returns and absorbing each one, and the pass that finds nothing +// in flight answers the limit. Nothing is dropped, because drain is not the +// road: a dropped return would leave its task claimed on a run that is over. +func (s *Supervisor) countLiveSpend() { + if s.limits.CostUSD <= 0 { + return + } + s.liveMu.Lock() + for id, counted := range s.counted { + if now := s.live[id]; now > counted { + s.spent += now - counted + s.counted[id] = now + } + } + s.liveMu.Unlock() + if s.spent >= s.limits.CostUSD && !s.limitHit { + s.limitHit = true + for _, cancel := range s.cancels { + cancel() + } + } +} + +// settleSpend reconciles a returned worker's report with what the loop counted +// while it worked, so its dollars are in the run's spend exactly once. The +// report is the receipt; a report smaller than what was already counted (a +// worker that lost its figure on the way out) takes nothing back, because the +// money was spent. A return can itself be what reaches the limit, as it always +// could. +func (s *Supervisor) settleSpend(ret workerReturn) { + if counted := s.counted[ret.task.ID]; ret.report.USD > counted { + s.spent += ret.report.USD - counted + } + s.forgetLive(ret.task.ID) + if s.limits.CostUSD > 0 && s.spent >= s.limits.CostUSD { + s.limitHit = true + } +} + +// forgetLive drops a returned worker's live figure and the loop's count of it, +// so the set the loop reads is the set of workers in flight. +func (s *Supervisor) forgetLive(id string) { + delete(s.counted, id) + s.liveMu.Lock() + delete(s.live, id) + s.liveMu.Unlock() } // absorb writes one worker's ending into the store and keeps the run's @@ -402,9 +529,12 @@ func (s *Supervisor) absorb(ret workerReturn) { // ending the store carries — a cancellation that landed while the worker // ran, which the same write cleared the claim and cascaded down — is the // one that stands, and no Done or Fail of this run may speak over it. The - // report is dropped whole, spend included: a worker stopped mid-flight - // hands back no account this run counts. The completion check below still - // runs, because the cancellation may be the write that finished the tree. + // report's words and its steps are dropped: a worker stopped mid-flight + // hands back no work this run counts. ITS DOLLARS ARE NOT DROPPED. What a + // run counts as spent never goes down and includes every paid call whatever + // way its task ended, because a person's limit is about dollars and not + // about how the work came out. The completion check below still runs, + // because the cancellation may be the write that finished the tree. delete(s.cancels, ret.task.ID) if ret.report.Waiting { // A PARKED RETURN IS NOT AN ENDING. The worker called `plandb wait`; the @@ -412,15 +542,13 @@ func (s *Supervisor) absorb(ret workerReturn) { // task stays open until the run wakes it ([launchWaits]). Nothing is // written over the store's own park, and the spend the worker made is // still the run's. - s.spent += ret.report.USD + s.settleSpend(ret) s.steps += ret.report.Steps - if s.limits.CostUSD > 0 && s.spent >= s.limits.CostUSD { - s.limitHit = true - } s.completeTree() return } if task := s.store.Task(ret.task.ID); task == nil || task.Status == plandb.StatusCancelled { + s.settleSpend(ret) s.completeTree() return } @@ -429,11 +557,8 @@ func (s *Supervisor) absorb(ret workerReturn) { // counts what the worker spent, and the root's own report is still the // run's result. endedByStore := ret.err == nil && s.store.Task(ret.task.ID).Status == plandb.StatusDone - s.spent += ret.report.USD + s.settleSpend(ret) s.steps += ret.report.Steps - if s.limits.CostUSD > 0 && s.spent >= s.limits.CostUSD { - s.limitHit = true - } if ret.task.ID == s.store.RootID() { // The harness owns the root, and the store refuses worker writes to it. // Its worker's report is kept for the completion; its error is what diff --git a/internal/run/supervisor_test.go b/internal/run/supervisor_test.go index 43d3a0ccd0..895c8b9ec8 100644 --- a/internal/run/supervisor_test.go +++ b/internal/run/supervisor_test.go @@ -609,7 +609,7 @@ func TestStartRefusesADoorBuiltWithoutItsStoreOrItsFactory(t *testing.T) { } } -func TestStartCountsNothingFromAWorkerWhoseTaskWasCancelled(t *testing.T) { +func TestStartCountsOnlyTheDollarsOfAWorkerWhoseTaskWasCancelled(t *testing.T) { store := startOpenStore(t, "the run's own title") ctx := runContext(t) seat := newFakeSeat() @@ -617,8 +617,9 @@ func TestStartCountsNothingFromAWorkerWhoseTaskWasCancelled(t *testing.T) { // The leaf is cancelled mid-flight and comes home well afterwards, so its // report is a late one carrying a completion, steps and spend. The store // would refuse the write on its own law; what only the drop in absorb does - // is keep the run's counters out of it, and the counters are what Start - // answers with. + // is keep the work's counters out of it, and the counters are what Start + // answers with. THE DOLLARS ARE THE EXCEPTION: they were paid whatever way + // the task ended, so they are in the run's account. seat.actions["l1"] = func(ctx context.Context, task plandb.Task) (run.Report, error) { <-ctx.Done() return run.Report{Result: "late " + task.ID, Steps: 4, USD: 0.30}, nil @@ -637,14 +638,14 @@ func TestStartCountsNothingFromAWorkerWhoseTaskWasCancelled(t *testing.T) { if outcome != run.OutcomeIncomplete { t.Fatalf("outcome = %q, want %q", outcome, run.OutcomeIncomplete) } - // The root worker's own figures and the wake's are the whole account: three - // workers went out (the root, the cancelled leaf, and the root's wake), and - // the cancelled one's 0.30 and four steps counted for nothing. + // Three workers went out (the root, the cancelled leaf, and the root's + // wake). The cancelled one's four steps count for nothing, and its 0.30 is + // counted beside the root's 0.1 and the wake's 0.05. if summary.Nodes != 3 { t.Fatalf("summary nodes = %d, want the root, the cancelled leaf, and the root's wake", summary.Nodes) } - if !usdClose(summary.USD, 0.15) { - t.Fatalf("summary usd = %v, want the root's 0.1 and the wake's 0.05 alone", summary.USD) + if !usdClose(summary.USD, 0.45) { + t.Fatalf("summary usd = %v, want the root's 0.1, the wake's 0.05 and the cancelled leaf's 0.30", summary.USD) } if summary.Steps != 2 { t.Fatalf("summary steps = %d, want the root's 1 and the wake's 1 alone", summary.Steps) diff --git a/internal/run/worker.go b/internal/run/worker.go index f96bfde4be..039a5a0abb 100644 --- a/internal/run/worker.go +++ b/internal/run/worker.go @@ -49,9 +49,10 @@ type WorkerFactory func(task plandb.Task) Worker // Elapsed of zero (or less) sets no corresponding run limit, a StepsPerTask // of zero hands the worker no cap, and ReviewRound's false is the run every caller had before it. type Limits struct { - // CostUSD is what the whole run may spend, as the sum of every Report's - // USD. When the counter has reached it no new worker starts and the run - // ends on the limit word of the outcome ladder. + // CostUSD is what the whole run may spend. Banked calls count while a worker + // is still working and reconcile with its final Report. When the counter has + // reached it no new worker starts, work in flight ends, and the run ends on + // the limit word of the outcome ladder. CostUSD float64 // Elapsed is how long the whole run may remain active. When it passes, // workers already in flight are ended and drained, no new worker starts, @@ -84,6 +85,24 @@ type Limits struct { // collide with. type stepsPerTaskKey struct{} +// spendBankKey carries the run-owned observer of a worker banking cumulative +// spend. The cumulative figure is reconciled with the worker return, so the +// live limit and the final receipt share one account. +type spendBankKey struct{} + +// WithSpendBank returns a context that reports cumulative spend as a worker +// banks calls. Workers without this property remain valid and report at return. +func WithSpendBank(ctx context.Context, bank func(float64)) context.Context { + return context.WithValue(ctx, spendBankKey{}, bank) +} + +// bankSpend publishes the cumulative spend banked by this worker. +func bankSpend(ctx context.Context, usd float64) { + if bank, _ := ctx.Value(spendBankKey{}).(func(float64)); bank != nil { + bank(usd) + } +} + // wakeClauseKey is the type behind the context value that carries a woken // parent's resume clause, for the same reason as the step cap beside it: a // typed lookup no other package can collide with.