From 89348b55123fddb1a137f699e9a8d844f9e03a93 Mon Sep 17 00:00:00 2001 From: BaaRai <130491490+BaaRai@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:00:16 +0200 Subject: [PATCH] feat(engine): per-intent LLM chain-head rotation (opt-in llm_rotate) Implements the per-intent rotation shape from #85: a task created with llm_rotate=true distributes worker intents round-robin over its LLM chain heads, so lane A can hunt on one model while lane B hunts on another (different lenses per intent, no extra tokens). Default stays the classic ordered failover from the active head; planner/goals/mainagent keep the active head always. Mechanics: the engine pins the next chain head into the run context when a worker claims an intent; the pin travels through the provider resolver (hooks.current now takes ctx). A pin is only a dispatch preference - it must belong to the task chain and never bypasses failover state (an exhausted chain stays exhausted). Schema: tasks.llm_rotate boolean, default false, idempotent ALTER; exposed on the create payload and the task DTO. Closes #85 (the ensemble/P4 shape stays parked as before) --- db/schema.sql | 2 ++ db/tasks.go | 22 +++++++++++++---- server/dto.go | 2 ++ server/engine.go | 8 +++++++ server/manager.go | 22 ++++++++++++++++- server/server.go | 4 +++- server/task_llm.go | 41 ++++++++++++++++++++++++++++---- server/task_llm_rotation_test.go | 32 +++++++++++++++++++++++++ server/task_llm_test.go | 20 ++++++++-------- 9 files changed, 132 insertions(+), 21 deletions(-) create mode 100644 server/task_llm_rotation_test.go diff --git a/db/schema.sql b/db/schema.sql index 8f0c634..8efe57c 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -522,6 +522,7 @@ CREATE TABLE IF NOT EXISTS tasks ( llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL, active_llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL, llm_chain_revision BIGINT NOT NULL DEFAULT 0, + llm_rotate BOOLEAN NOT NULL DEFAULT false, company_id BIGINT REFERENCES companies(id) ON DELETE SET NULL, parent_ref TEXT, timeout_seconds INTEGER NOT NULL DEFAULT 0, @@ -545,6 +546,7 @@ CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks -- planner heartbeat trigger interval (seconds); backfills old databases. Default 300s (5min). See docs/planner-trigger-impl-plan.md ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300; ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_quiet_mode BOOLEAN NOT NULL DEFAULT true; +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS llm_rotate BOOLEAN NOT NULL DEFAULT false; -- Concurrency-limit hold state; backfills old databases. true=queued due to the concurrency limit, auto-starts when a slot frees up. ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued BOOLEAN NOT NULL DEFAULT false; -- Asset coverage feature switch; backfills old databases. true (default)=compute/show test coverage, auto-accumulate test scope, diff --git a/db/tasks.go b/db/tasks.go index afc52c2..413b7ef 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -33,6 +33,9 @@ type Task struct { LLMChainRevision int64 `json:"-"` LLMFailoverState string `json:"llm_failover_state,omitempty"` LLMFailoverReason string `json:"llm_failover_reason,omitempty"` + // LLMRotate: per-intent chain-head rotation switch (#85); false = classic + // ordered failover from the active head. + LLMRotate bool `json:"llm_rotate"` SourceTaskIDs []int64 `json:"source_task_ids,omitempty"` CompanyIDs []int64 `json:"company_ids,omitempty"` ParentRef string `json:"parent_ref,omitempty"` // parent task id (recorded by orchestration spawn; empty=top level) @@ -165,6 +168,11 @@ type TaskCreateOptions struct { CoverageEnabled *bool // InterceptRules are task-level asset interception rules, written into task_intercept_rules in the same transaction as the task at creation. InterceptRules []TaskInterceptRuleInput + // LLMRotate enables per-intent round-robin rotation of the task's LLM chain + // heads (#85): each claimed worker intent runs on the next profile in the + // chain instead of the single active head. nil = off (default), preserving + // the ordered-failover semantics. + LLMRotate *bool } // CreateTaskWithOptions creates an exploration, task, direct source relations, @@ -239,10 +247,14 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); PlanQuietMode: quietMode, CoverageEnabled: coverageEnabled, } + rotate := false + if opts.LLMRotate != nil { + rotate = *opts.LLMRotate + } if err := tx.QueryRow(` -INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, plan_quiet_mode, coverage_enabled) -VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9,$10) -RETURNING id, status, paused, created_at`, opts.Name, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, quietMode, coverageEnabled).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { +INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, plan_quiet_mode, coverage_enabled, llm_rotate) +VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9,$10,$11) +RETURNING id, status, paused, created_at`, opts.Name, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, quietMode, coverageEnabled, rotate).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { return nil, err } if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil { @@ -358,11 +370,11 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error { const taskCols = `id, COALESCE(name,''), category_id, COALESCE((SELECT category.name FROM task_categories category WHERE category.id=tasks.category_id),''), -description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), pinned_at, created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(plan_quiet_mode,true), COALESCE(coverage_enabled,true), first_run_at, deadline_at` +description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), pinned_at, created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(plan_quiet_mode,true), COALESCE(coverage_enabled,true), COALESCE(llm_rotate,false), first_run_at, deadline_at` func scanTask(sc interface{ Scan(...any) error }) (*Task, error) { var t Task - if err := sc.Scan(&t.ID, &t.Name, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.PinnedAt, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.PlanQuietMode, &t.CoverageEnabled, &t.FirstRunAt, &t.DeadlineAt); err != nil { + if err := sc.Scan(&t.ID, &t.Name, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.PinnedAt, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.PlanQuietMode, &t.CoverageEnabled, &t.LLMRotate, &t.FirstRunAt, &t.DeadlineAt); err != nil { return nil, err } t.Pinned = t.PinnedAt != nil diff --git a/server/dto.go b/server/dto.go index 1fc9e66..f9511e5 100644 --- a/server/dto.go +++ b/server/dto.go @@ -55,6 +55,7 @@ type TaskDTO struct { LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // LLM profile used for this task; nil = default LLMProfileIDs []int64 `json:"llm_profile_ids"` ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"` + LLMRotate bool `json:"llm_rotate"` LLMFailoverState string `json:"llm_failover_state"` LLMFailoverReason string `json:"llm_failover_reason,omitempty"` SourceTaskIDs []string `json:"source_task_ids"` @@ -130,6 +131,7 @@ func taskDTO(t *Task, status string) TaskDTO { LLMProfileID: llmState.ProfileID, LLMProfileIDs: profileIDs, ActiveLLMProfileID: llmState.ActiveID, + LLMRotate: t.LLMRotate, LLMFailoverState: llmState.FailoverState, LLMFailoverReason: llmState.FailoverReason, SourceTaskIDs: sourceIDs, diff --git a/server/engine.go b/server/engine.go index 6e7dc29..8ccff08 100644 --- a/server/engine.go +++ b/server/engine.go @@ -1017,6 +1017,14 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker // must already hold one task-operation admission for the whole sequence so a delete // cannot observe quiescence between the LLM return and the final DB writes. func (e *Engine) runIntent(ctx context.Context, t *Task, name string, worker *agent.Worker, intent *db.Node, requestID, message string) bool { + // Per-intent LLM rotation (#85): each claimed worker intent runs on the + // next chain head. The pin travels through the run context into the + // provider resolver; non-worker roles and failover keep the active head. + if t.LLMRotate { + if head, ok := t.nextRotatedProfile(); ok { + ctx = withLLMProfilePin(ctx, head) + } + } hasChatMessage := message != "" e.stampFirstRun(t) // first real execution → stamp first_run_at + compute the deadline (only for tasks with a timeout) e.touch(t.ID) diff --git a/server/manager.go b/server/manager.go index bb73ef7..888f238 100644 --- a/server/manager.go +++ b/server/manager.go @@ -51,6 +51,10 @@ type Task struct { ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"` LLMChainRevision int64 `json:"-"` LLMFailoverState string `json:"llm_failover_state,omitempty"` + // LLMRotate (#85): per-intent round-robin over the task's LLM chain heads. + LLMRotate bool `json:"llm_rotate,omitempty"` + rotateMu sync.Mutex + rotateNext int LLMFailoverReason string `json:"llm_failover_reason,omitempty"` SourceTaskIDs []int64 `json:"source_task_ids,omitempty"` CompanyIDs []int64 `json:"company_ids,omitempty"` @@ -178,6 +182,22 @@ func (t *Task) llmStateSnapshot() taskLLMState { } } +// nextRotatedProfile returns the chain head for the NEXT claimed worker +// intent under per-intent rotation (#85): strict round-robin over the task's +// chain. The active head stays authoritative for every non-worker role and for +// failover ordering; rotation only distributes worker lanes across heads. +func (t *Task) nextRotatedProfile() (int64, bool) { + t.rotateMu.Lock() + defer t.rotateMu.Unlock() + ids := t.LLMProfileIDs + if len(ids) == 0 { + return 0, false + } + pick := ids[t.rotateNext%len(ids)] + t.rotateNext++ + return pick, true +} + func (t *Task) setLLMState(profileID, activeID *int64, profileIDs []int64, revision int64, state, reason string) bool { t.llmMu.Lock() defer t.llmMu.Unlock() @@ -840,7 +860,7 @@ func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Inter QueuedAt: unixNanoOrZero(pt.QueuedAt), QueueMode: pt.QueueMode, CompletedAt: unixOrZero(pt.CompletedAt), Status: pt.Status, ParentRef: pt.ParentRef, LLMProfileID: pt.LLMProfileID, - LLMProfileIDs: append([]int64(nil), pt.LLMProfileIDs...), ActiveLLMProfileID: pt.ActiveLLMProfileID, + LLMProfileIDs: append([]int64(nil), pt.LLMProfileIDs...), ActiveLLMProfileID: pt.ActiveLLMProfileID, LLMRotate: pt.LLMRotate, LLMChainRevision: pt.LLMChainRevision, LLMFailoverState: pt.LLMFailoverState, LLMFailoverReason: pt.LLMFailoverReason, SourceTaskIDs: append([]int64(nil), pt.SourceTaskIDs...), diff --git a/server/server.go b/server/server.go index ef79c7b..0130b75 100644 --- a/server/server.go +++ b/server/server.go @@ -1506,7 +1506,8 @@ type createTaskReq struct { LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` // ordered task-level profile chain; the first entry takes effect initially SourceTaskIDs []string `json:"source_task_ids,omitempty"` // direct, read-only inherited source tasks CompanyIDs []int64 `json:"company_ids,omitempty"` // links the company scope and snapshots its current assets; does not copy assets or force intent generation - TimeoutSeconds int `json:"timeout_seconds"` // task-level timeout (seconds); 0/omitted = no limit + TimeoutSeconds int `json:"timeout_seconds"` + LLMRotate *bool `json:"llm_rotate"` // task-level timeout (seconds); 0/omitted = no limit PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner heartbeat interval (seconds); 0/omitted = default 600 (10min); floor = default = 600, lower values are raised to 600 PlanQuietMode *bool `json:"plan_quiet_mode,omitempty"` // planner quiet-mode switch (#86); omitted = default on; false = edge rounds always plan SeedFirstIntent *bool `json:"seed_first_intent,omitempty"` // issue one seed intent at creation (content = description + goal) so a worker can start right away without waiting for the first planner round; omitted/null = default off, the standard plan-then-execute flow. Pass true explicitly to enable (for CTF tasks often solved in one work, this skips the pre-start planner round). @@ -1580,6 +1581,7 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) { TimeoutSeconds: req.TimeoutSeconds, PlanHeartbeatSeconds: req.PlanHeartbeatSeconds, PlanQuietMode: req.PlanQuietMode, CoverageEnabled: req.CoverageEnabled, InterceptRules: interceptRules, + LLMRotate: req.LLMRotate, }) if err != nil { if errors.Is(err, db.ErrTaskCategoryInvalid) || errors.Is(err, db.ErrTaskCategoryNotFound) { diff --git a/server/task_llm.go b/server/task_llm.go index df37e5d..5062fe0 100644 --- a/server/task_llm.go +++ b/server/task_llm.go @@ -95,8 +95,25 @@ type taskLLMSelection struct { retry agent.RetryConfig } +// llmProfilePin carries a per-run chain-head override through the call stack +// (#85 per-intent rotation). The engine sets it when launching a worker run; +// resolution stays within the task's chain and never bypasses failover state. +type llmProfilePinKey struct{} + +func withLLMProfilePin(ctx context.Context, profileID int64) context.Context { + return context.WithValue(ctx, llmProfilePinKey{}, profileID) +} + +func llmProfilePinFrom(ctx context.Context) (int64, bool) { + if ctx == nil { + return 0, false + } + id, ok := ctx.Value(llmProfilePinKey{}).(int64) + return id, ok +} + type taskLLMStreamHooks struct { - current func() (taskLLMSelection, error) + current func(context.Context) (taskLLMSelection, error) exhaust func(taskLLMSelection, error) (db.TaskLLMTransition, error) transition func(taskLLMSelection, db.TaskLLMTransition, error) } @@ -108,6 +125,14 @@ type taskLLMStreamHooks struct { // to the task chain, and to the global configuration when the chain is empty. // The returned profile id is non-zero only on the task-chain path — streamTaskLLM uses it to decide whether a quota error should advance the task's failover state (binding/global paths leave the chain state untouched, preserving existing semantics). func (r *taskLLMRuntime) current() (taskLLMSelection, error) { + return r.currentFor(context.Background()) +} + +// currentFor resolves like current, honoring a per-run profile pin set by the +// engine for per-intent rotation (#85). A pin is only a dispatch preference: +// the chain must still be in a ready state (an exhausted chain with no active +// head stays exhausted) and the pinned id must belong to the task's chain. +func (r *taskLLMRuntime) currentFor(ctx context.Context) (taskLLMSelection, error) { taskNum, err := parseTaskID(r.taskID) if err != nil { return taskLLMSelection{}, err @@ -131,6 +156,14 @@ func (r *taskLLMRuntime) current() (taskLLMSelection, error) { return sel, &taskLLMError{taskID: r.taskID, chainExhausted: true, cause: errors.New("all selected profiles are quota exhausted")} } sel.profileID = *pt.ActiveLLMProfileID + if pin, ok := llmProfilePinFrom(ctx); ok && pin != sel.profileID { + for _, id := range pt.LLMProfileIDs { + if id == pin { + sel.profileID = pin + break + } + } + } prov, cfg, ok := r.s.providerForProfile(sel.profileID) if !ok { return sel, fmt.Errorf("LLM profile #%d is missing or invalid", sel.profileID) @@ -201,7 +234,7 @@ func parseTaskID(id string) (int64, error) { // exhausted, and how to emit a failover transition. func (r *taskLLMRuntime) streamHooks() taskLLMStreamHooks { return taskLLMStreamHooks{ - current: r.current, + current: func(ctx context.Context) (taskLLMSelection, error) { return r.currentFor(ctx) }, exhaust: func(selection taskLLMSelection, cause error) (db.TaskLLMTransition, error) { taskNum, _ := parseTaskID(r.taskID) transition, err := r.s.m.pg.MarkTaskLLMProfileQuotaExhaustedAtRevision(taskNum, selection.profileID, selection.revision, cause.Error()) @@ -236,7 +269,7 @@ func (r *taskLLMRuntime) Complete(ctx context.Context, req llm.CompletionRequest func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest, hooks taskLLMStreamHooks) (llm.Message, string, llm.Usage, error) { for { - selection, err := hooks.current() + selection, err := hooks.current(ctx) if err != nil { return llm.Message{}, "", llm.Usage{}, err } @@ -287,7 +320,7 @@ func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionReque func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest, hooks taskLLMStreamHooks) iter.Seq2[llm.StreamEvent, error] { return func(yield func(llm.StreamEvent, error) bool) { for { - selection, err := hooks.current() + selection, err := hooks.current(ctx) if err != nil { yield(llm.StreamEvent{}, err) return diff --git a/server/task_llm_rotation_test.go b/server/task_llm_rotation_test.go new file mode 100644 index 0000000..09f77bc --- /dev/null +++ b/server/task_llm_rotation_test.go @@ -0,0 +1,32 @@ +package server + +import ( + "context" + "testing" +) + +func TestLLMProfilePinRoundTrip(t *testing.T) { + ctx := context.Background() + if _, ok := llmProfilePinFrom(ctx); ok { + t.Fatal("empty context must carry no pin") + } + pinned := withLLMProfilePin(ctx, 42) + got, ok := llmProfilePinFrom(pinned) + if !ok || got != 42 { + t.Fatalf("pin roundtrip got=%d ok=%v", got, ok) + } +} + +func TestNextRotatedProfileRoundRobin(t *testing.T) { + tk := &Task{LLMProfileIDs: []int64{3, 1, 2}} + want := []int64{3, 1, 2, 3, 1} + for i, w := range want { + got, ok := tk.nextRotatedProfile() + if !ok || got != w { + t.Fatalf("rotation %d: got=%d ok=%v want=%d", i, got, ok, w) + } + } + if _, ok := (&Task{}).nextRotatedProfile(); ok { + t.Fatal("empty chain must not rotate") + } +} diff --git a/server/task_llm_test.go b/server/task_llm_test.go index be6fdd0..9f97bff 100644 --- a/server/task_llm_test.go +++ b/server/task_llm_test.go @@ -246,7 +246,7 @@ func TestTaskLLMStreamRetriesPreStreamQuotaOnNextProfile(t *testing.T) { active := int64(11) transitions := 0 hooks := taskLLMStreamHooks{ - current: func() (taskLLMSelection, error) { + current: func(ctx context.Context) (taskLLMSelection, error) { if active == 11 { return taskLLMSelection{profileID: 11, provider: first}, nil } @@ -288,7 +288,7 @@ func TestTaskLLMStreamRetriesAfterNonContentStartEvent(t *testing.T) { }} active := int64(11) hooks := taskLLMStreamHooks{ - current: func() (taskLLMSelection, error) { + current: func(ctx context.Context) (taskLLMSelection, error) { if active == 11 { return taskLLMSelection{profileID: 11, provider: first}, nil } @@ -319,7 +319,7 @@ func TestTaskLLMStreamRetriesStalePreStreamFailureAgainstReplacementChain(t *tes second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "replacement"}}} current := first hooks := taskLLMStreamHooks{ - current: func() (taskLLMSelection, error) { + current: func(ctx context.Context) (taskLLMSelection, error) { profileID := int64(11) if current == second { profileID = 22 @@ -353,7 +353,7 @@ func TestTaskLLMStreamDoesNotSwitchForOrdinaryErrors(t *testing.T) { provider := &scriptedLLMProvider{err: errors.New(message)} exhausted := 0 hooks := taskLLMStreamHooks{ - current: func() (taskLLMSelection, error) { + current: func(ctx context.Context) (taskLLMSelection, error) { return taskLLMSelection{profileID: 11, provider: provider}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { @@ -382,7 +382,7 @@ func TestTaskLLMStreamAdvancesAfterMidStreamQuotaWithoutReplay(t *testing.T) { second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "next-call"}}} active := int64(11) hooks := taskLLMStreamHooks{ - current: func() (taskLLMSelection, error) { + current: func(ctx context.Context) (taskLLMSelection, error) { if active == 11 { return taskLLMSelection{profileID: 11, provider: first}, nil } @@ -418,7 +418,7 @@ func TestTaskLLMStreamMarkFailureAfterToolStartSkipsWorkerReplay(t *testing.T) { err: quota, } hooks := taskLLMStreamHooks{ - current: func() (taskLLMSelection, error) { + current: func(ctx context.Context) (taskLLMSelection, error) { return taskLLMSelection{profileID: 11, provider: provider}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { @@ -467,7 +467,7 @@ func TestTaskLLMStreamRetriesTransientPreCommitOnSameProvider(t *testing.T) { } exhausted := 0 hooks := taskLLMStreamHooks{ - current: func() (taskLLMSelection, error) { + current: func(ctx context.Context) (taskLLMSelection, error) { return taskLLMSelection{profileID: 11, provider: provider}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { @@ -499,7 +499,7 @@ func TestTaskLLMStreamSurfacesTransientAfterRetriesExhausted(t *testing.T) { failErr: errors.New("dial tcp: connection reset by peer"), } hooks := taskLLMStreamHooks{ - current: func() (taskLLMSelection, error) { + current: func(ctx context.Context) (taskLLMSelection, error) { return taskLLMSelection{profileID: 11, provider: provider}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { @@ -524,7 +524,7 @@ func TestTaskLLMStreamDoesNotRetryDeterministicRejection(t *testing.T) { failErr: errors.New("anthropic: status 400: messages.1: invalid request"), } hooks := taskLLMStreamHooks{ - current: func() (taskLLMSelection, error) { + current: func(ctx context.Context) (taskLLMSelection, error) { return taskLLMSelection{profileID: 11, provider: provider}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { @@ -620,7 +620,7 @@ func TestTaskLLMStreamStopsWhenChainExhausted(t *testing.T) { provider := &scriptedLLMProvider{err: errors.New("billing_not_active")} transitions := 0 hooks := taskLLMStreamHooks{ - current: func() (taskLLMSelection, error) { + current: func(ctx context.Context) (taskLLMSelection, error) { return taskLLMSelection{profileID: 11, provider: provider}, nil }, exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {