From d03068ace5a8222ca909fe90d2b1134288aa4bbf Mon Sep 17 00:00:00 2001 From: dev_agent Date: Mon, 5 Oct 2026 10:00:20 -0400 Subject: [PATCH] Add concurrency guards to harness.Run: single-flight, semaphore and timeouts harness.Run now takes a scope key. Concurrent runs of one skill with one scope share a single in-flight run and its result, errors included, so only one Generate call is made; different scopes run independently. A process-wide Guard (NewGuard, DefaultMaxCalls = 2) bounds concurrent LLM calls with a context-aware semaphore. Each skill declares Timeout() (DefaultTimeout = 60 s), enforced via context over the whole run; a run past it fails with an error matching llm.ErrTimeout. A caller whose context ends gets the context error; the shared run is cancelled, and its slot freed, once no caller waits for it. llm.Fake results can now block (Release) and signal entry (Started), so tests hold calls in flight. Maestro task: 8f533721-32c9-4f73-af9f-9a9f0e94abf8 Co-Authored-By: Claude Opus 5.5 --- docs/ARCHITECTURE.md | 10 + internal/ai/harness/echo_test.go | 3 + internal/ai/harness/guard.go | 133 ++++++++++++ internal/ai/harness/guard_test.go | 305 ++++++++++++++++++++++++++++ internal/ai/harness/harness.go | 57 +++++- internal/ai/harness/harness_test.go | 19 +- internal/ai/harness/prompt_test.go | 2 +- internal/ai/harness/runs_test.go | 13 +- internal/ai/llm/fake.go | 32 ++- internal/ai/llm/fake_test.go | 34 ++++ 10 files changed, 591 insertions(+), 17 deletions(-) create mode 100644 internal/ai/harness/guard.go create mode 100644 internal/ai/harness/guard_test.go diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 8d3cc87..750478f 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -162,6 +162,16 @@ automatic** (D-021). with key-like query parameters are removed, inside JSON strings too. `harness.Purge` deletes rows older than `harness.RunRetention` (30 days), clearing the `suggestion_batches` and `week_plans` references first. +- **Concurrency guards** (`harness.Guard`, one per process): `harness.Run` + takes a scope key, and concurrent runs of one skill with one scope share a + single in-flight run and its result, error included (single-flight); runs + with different scopes are independent. A semaphore lets at most + `harness.DefaultMaxCalls` (2, set by `NewGuard`) LLM calls run at once; + waiting for a slot honours the context. Each skill declares a timeout + (`harness.DefaultTimeout`, 60 s, unless it needs another) bounding the + whole run, slot waits and repair turn included; a run past it fails with + `llm.ErrTimeout`. A caller whose context ends gets the context error; the + shared run is cancelled, freeing its slot, once no caller waits for it. - **Prompt layering** (D-007): each skill version's prompt is a `text/template` file embedded from `ai/skills/prompts/` (`.v.tmpl`). The system prompt is always, in order: diff --git a/internal/ai/harness/echo_test.go b/internal/ai/harness/echo_test.go index 2a02ede..8ac6e74 100644 --- a/internal/ai/harness/echo_test.go +++ b/internal/ai/harness/echo_test.go @@ -4,6 +4,7 @@ import ( "encoding/json" "errors" "strings" + "time" "code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/ai/llm" ) @@ -53,3 +54,5 @@ func (echoSkill) PostProcess(out echoOutput) (echoOutput, error) { out.Echo = strings.TrimSpace(out.Echo) return out, nil } + +func (echoSkill) Timeout() time.Duration { return DefaultTimeout } diff --git a/internal/ai/harness/guard.go b/internal/ai/harness/guard.go new file mode 100644 index 0000000..0993c97 --- /dev/null +++ b/internal/ai/harness/guard.go @@ -0,0 +1,133 @@ +package harness + +import ( + "context" + "fmt" + "sync" + "time" + + "code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/ai/llm" +) + +// DefaultMaxCalls is how many LLM calls a Guard lets run at once unless it +// is built with another limit. +const DefaultMaxCalls = 2 + +// DefaultTimeout is the run timeout a skill declares unless it needs +// another. +const DefaultTimeout = 60 * time.Second + +// Guard bounds the process's runs. Runs of one skill with one scope share a +// single in-flight execution and its result (single-flight), and at most +// its limit of LLM calls run at once across every run (the semaphore). +// Build one with NewGuard and share it between every run of the process. +type Guard struct { + slots chan struct{} + + mu sync.Mutex + flights map[flightKey]*flight +} + +// flightKey identifies the runs that share one execution. +type flightKey struct { + skill string + scope string +} + +// flight is one in-flight execution and the callers waiting for it. +type flight struct { + done chan struct{} + val any + err error + // waiters counts the callers still waiting; when the last one gives + // up, cancel stops the execution. + waiters int + cancel context.CancelFunc +} + +// NewGuard returns a Guard that lets maxCalls LLM calls run at once. +func NewGuard(maxCalls int) (*Guard, error) { + if maxCalls < 1 { + return nil, fmt.Errorf("harness: guard: max concurrent LLM calls %d is not positive", maxCalls) + } + return &Guard{slots: make(chan struct{}, maxCalls), flights: map[flightKey]*flight{}}, nil +} + +// do returns the result of fn for key, starting fn unless a call for key is +// already in flight, in which case it waits for that call's result +// instead. fn runs on a context that keeps the first caller's values but +// not its cancellation, so one caller giving up does not fail the others. +// When ctx is done first, do returns its error; when no caller is left +// waiting, fn's context is cancelled and do returns only once fn has +// ended, so the abandoned execution has released its slot by then. +func (g *Guard) do(ctx context.Context, key flightKey, fn func(context.Context) (any, error)) (any, error) { + g.mu.Lock() + f, ok := g.flights[key] + if !ok { + fctx, cancel := context.WithCancel(context.WithoutCancel(ctx)) + f = &flight{done: make(chan struct{}), cancel: cancel} + g.flights[key] = f + go func() { + defer cancel() + f.val, f.err = fn(fctx) + g.forget(key, f) + close(f.done) + }() + } + f.waiters++ + g.mu.Unlock() + + select { + case <-f.done: + return f.val, f.err + case <-ctx.Done(): + } + g.mu.Lock() + f.waiters-- + last := f.waiters == 0 + if last { + // No caller may join an execution that is being abandoned. + g.forgetLocked(key, f) + f.cancel() + } + g.mu.Unlock() + if last { + <-f.done + } + return nil, ctx.Err() +} + +// forget removes f from the in-flight executions if it is still there. +func (g *Guard) forget(key flightKey, f *flight) { + g.mu.Lock() + defer g.mu.Unlock() + g.forgetLocked(key, f) +} + +func (g *Guard) forgetLocked(key flightKey, f *flight) { + if g.flights[key] == f { + delete(g.flights, key) + } +} + +// limited is a Provider whose Generate calls hold one of its Guard's slots. +type limited struct { + llm.Provider + slots chan struct{} +} + +// Generate waits for a free slot, or for ctx to be done, then sends req +// and frees the slot when the call returns. +func (p limited) Generate(ctx context.Context, req llm.Request) (llm.Response, error) { + // A done context never takes a slot, even when one is free. + if err := ctx.Err(); err != nil { + return llm.Response{}, err + } + select { + case p.slots <- struct{}{}: + case <-ctx.Done(): + return llm.Response{}, ctx.Err() + } + defer func() { <-p.slots }() + return p.Provider.Generate(ctx, req) +} diff --git a/internal/ai/harness/guard_test.go b/internal/ai/harness/guard_test.go new file mode 100644 index 0000000..e9b96bc --- /dev/null +++ b/internal/ai/harness/guard_test.go @@ -0,0 +1,305 @@ +package harness + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/ai/llm" +) + +// timeoutSkill is echoSkill with its own timeout. +type timeoutSkill struct { + echoSkill + timeout time.Duration +} + +func (s timeoutSkill) Timeout() time.Duration { return s.timeout } + +// blocked is one scripted Generate call that blocks until released. +type blocked struct { + started chan struct{} + release chan struct{} +} + +// queueBlocked queues r on f as a call that blocks until released. +func queueBlocked(f *llm.Fake, r llm.GenerateResult) blocked { + b := blocked{started: make(chan struct{}), release: make(chan struct{})} + r.Started, r.Release = b.started, b.release + f.QueueGenerate(r) + return b +} + +// result is what one Run returned. +type result struct { + out echoOutput + err error +} + +// goRun starts Run of s with scope on cfg and returns where its result +// arrives. +func goRun(ctx context.Context, cfg Config, s Skill[echoInput, echoOutput], scope string) <-chan result { + ch := make(chan result, 1) + go func() { + out, err := Run(ctx, cfg, s, scope, echoInput{Word: "taco"}) + ch <- result{out, err} + }() + return ch +} + +// await returns what arrives on ch, failing the test after a second. +func await[T any](t *testing.T, ch <-chan T, what string) T { + t.Helper() + select { + case v := <-ch: + return v + case <-time.After(time.Second): + t.Fatalf("timed out waiting for %s", what) + panic("unreachable") + } +} + +// waitStarted waits for b's Generate call to start. +func waitStarted(t *testing.T, b blocked, what string) { + t.Helper() + select { + case <-b.started: + case <-time.After(time.Second): + t.Fatalf("timed out waiting for %s to call Generate", what) + } +} + +// assertNotStarted checks that b's Generate call has not started after a +// grace period. +func assertNotStarted(t *testing.T, b blocked, what string) { + t.Helper() + select { + case <-b.started: + t.Fatalf("%s called Generate while no slot was free", what) + case <-time.After(50 * time.Millisecond): + } +} + +// waitWaiters waits until n callers wait for the run of skill with scope. +func waitWaiters(t *testing.T, g *Guard, skill, scope string, n int) { + t.Helper() + deadline := time.Now().Add(time.Second) + for { + g.mu.Lock() + got := 0 + if f := g.flights[flightKey{skill: skill, scope: scope}]; f != nil { + got = f.waiters + } + g.mu.Unlock() + if got == n { + return + } + if time.Now().After(deadline) { + t.Fatalf("%d callers wait for %s/%s, want %d", got, skill, scope, n) + } + time.Sleep(time.Millisecond) + } +} + +func TestRunSameScopeSharesOneExecution(t *testing.T) { + for name, script := range map[string]llm.GenerateResult{ + "output": reply(`{"echo":"taco"}`), + "error": {Err: llm.ErrAuth}, + } { + t.Run(name, func(t *testing.T) { + var f llm.Fake + cfg := testConfig(t, &f) + b := queueBlocked(&f, script) + + first := goRun(context.Background(), cfg, echoSkill{}, "week-1") + waitStarted(t, b, "the first run") + second := goRun(context.Background(), cfg, echoSkill{}, "week-1") + waitWaiters(t, cfg.Guard, "echo", "week-1", 2) + close(b.release) + + r1, r2 := await(t, first, "the first run"), await(t, second, "the second run") + if r1 != r2 { + t.Fatalf("results differ: %+v and %+v", r1, r2) + } + if name == "output" && (r1.err != nil || r1.out.Echo != "taco") { + t.Fatalf("result = %+v, want taco", r1) + } + if name == "error" && !errors.Is(r1.err, llm.ErrAuth) { + t.Fatalf("err = %v, want ErrAuth", r1.err) + } + assertCalls(t, &f, 1) + if rows := runRows(t, cfg.Store); len(rows) != 1 { + t.Fatalf("%d ai_runs rows, want 1 for the shared run", len(rows)) + } + }) + } +} + +func TestRunDifferentScopesRunIndependently(t *testing.T) { + var f llm.Fake + cfg := testConfig(t, &f) + b1 := queueBlocked(&f, reply(`{"echo":"taco"}`)) + b2 := queueBlocked(&f, reply(`{"echo":"taco"}`)) + + week1 := goRun(context.Background(), cfg, echoSkill{}, "week-1") + week2 := goRun(context.Background(), cfg, echoSkill{}, "week-2") + // Both calls are in Generate at once: neither waits for the other. + waitStarted(t, b1, "a run") + waitStarted(t, b2, "the other run") + close(b1.release) + close(b2.release) + + for _, ch := range []<-chan result{week1, week2} { + if r := await(t, ch, "a run"); r.err != nil || r.out.Echo != "taco" { + t.Fatalf("result = %+v", r) + } + } + assertCalls(t, &f, 2) +} + +func TestRunThirdCallWaitsForSlot(t *testing.T) { + var f llm.Fake + cfg := testConfig(t, &f) + b1 := queueBlocked(&f, reply(`{"echo":"taco"}`)) + b2 := queueBlocked(&f, reply(`{"echo":"taco"}`)) + b3 := queueBlocked(&f, reply(`{"echo":"taco"}`)) + + first := goRun(context.Background(), cfg, echoSkill{}, "week-1") + waitStarted(t, b1, "the first run") + second := goRun(context.Background(), cfg, echoSkill{}, "week-2") + waitStarted(t, b2, "the second run") + third := goRun(context.Background(), cfg, echoSkill{}, "week-3") + assertNotStarted(t, b3, "the third run") + if n := len(f.Requests()); n != 2 { + t.Fatalf("Generate called %d times while both slots were taken, want 2", n) + } + + close(b1.release) + waitStarted(t, b3, "the third run") + close(b2.release) + close(b3.release) + for _, ch := range []<-chan result{first, second, third} { + if r := await(t, ch, "a run"); r.err != nil || r.out.Echo != "taco" { + t.Fatalf("result = %+v", r) + } + } + assertCalls(t, &f, 3) +} + +func TestRunCancelledContextReleasesSlot(t *testing.T) { + var f llm.Fake + cfg := testConfig(t, &f) + b1 := queueBlocked(&f, reply(`{"echo":"taco"}`)) + b2 := queueBlocked(&f, reply(`{"echo":"taco"}`)) + b3 := queueBlocked(&f, reply(`{"echo":"taco"}`)) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + first := goRun(ctx, cfg, echoSkill{}, "week-1") + waitStarted(t, b1, "the first run") + second := goRun(context.Background(), cfg, echoSkill{}, "week-2") + waitStarted(t, b2, "the second run") + third := goRun(context.Background(), cfg, echoSkill{}, "week-3") + assertNotStarted(t, b3, "the third run") + + cancel() + if r := await(t, first, "the cancelled run"); !errors.Is(r.err, context.Canceled) { + t.Fatalf("cancelled run err = %v, want context.Canceled", r.err) + } + // The cancelled run's slot goes to the third run; the first call's + // release is never closed. + waitStarted(t, b3, "the third run") + close(b2.release) + close(b3.release) + for _, ch := range []<-chan result{second, third} { + if r := await(t, ch, "a run"); r.err != nil || r.out.Echo != "taco" { + t.Fatalf("result = %+v", r) + } + } + assertCalls(t, &f, 3) + var statuses []string + for _, r := range runRows(t, cfg.Store) { + statuses = append(statuses, r.Status) + } + if strings.Join(statuses, ",") != "provider_error,ok,ok" { + t.Fatalf("statuses = %v, want the cancelled run recorded as provider_error", statuses) + } +} + +func TestRunCallerGivingUpLeavesSharedRun(t *testing.T) { + var f llm.Fake + cfg := testConfig(t, &f) + b := queueBlocked(&f, reply(`{"echo":"taco"}`)) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + leader := goRun(ctx, cfg, echoSkill{}, "week-1") + waitStarted(t, b, "the first run") + follower := goRun(context.Background(), cfg, echoSkill{}, "week-1") + waitWaiters(t, cfg.Guard, "echo", "week-1", 2) + + cancel() + if r := await(t, leader, "the cancelled caller"); !errors.Is(r.err, context.Canceled) { + t.Fatalf("cancelled caller err = %v, want context.Canceled", r.err) + } + close(b.release) + if r := await(t, follower, "the remaining caller"); r.err != nil || r.out.Echo != "taco" { + t.Fatalf("remaining caller = %+v, want the shared run's output", r) + } + assertCalls(t, &f, 1) +} + +func TestRunTimeout(t *testing.T) { + t.Run("in Generate", func(t *testing.T) { + var f llm.Fake + cfg := testConfig(t, &f) + queueBlocked(&f, reply(`{"echo":"taco"}`)) + _, err := Run(context.Background(), cfg, timeoutSkill{timeout: 20 * time.Millisecond}, "week-1", echoInput{Word: "taco"}) + if !errors.Is(err, llm.ErrTimeout) { + t.Fatalf("err = %v, want ErrTimeout", err) + } + assertCalls(t, &f, 1) + rows := runRows(t, cfg.Store) + if len(rows) != 1 || rows[0].Status != string(StatusProviderError) || !strings.Contains(rows[0].Error.String, llm.ErrTimeout.Error()) { + t.Fatalf("rows = %+v, want one provider_error row naming the timeout", rows) + } + }) + t.Run("waiting for a slot", func(t *testing.T) { + var f llm.Fake + cfg := testConfig(t, &f) + cfg.Guard = testGuard(t, 1) + b := queueBlocked(&f, reply(`{"echo":"taco"}`)) + holder := goRun(context.Background(), cfg, echoSkill{}, "week-1") + waitStarted(t, b, "the slot holder") + + _, err := Run(context.Background(), cfg, timeoutSkill{timeout: 20 * time.Millisecond}, "week-2", echoInput{Word: "taco"}) + if !errors.Is(err, llm.ErrTimeout) { + t.Fatalf("err = %v, want ErrTimeout", err) + } + close(b.release) + if r := await(t, holder, "the slot holder"); r.err != nil { + t.Fatal(r.err) + } + assertCalls(t, &f, 1) + }) +} + +func TestRunRefusesBadGuardInput(t *testing.T) { + if _, err := NewGuard(0); err == nil { + t.Fatal("NewGuard(0) succeeded") + } + var f llm.Fake + cfg := testConfig(t, &f) + if _, err := Run(context.Background(), cfg, echoSkill{}, "", echoInput{Word: "taco"}); err == nil || !strings.Contains(err.Error(), "no scope") { + t.Fatalf("empty scope err = %v", err) + } + if _, err := Run(context.Background(), cfg, timeoutSkill{}, "week-1", echoInput{Word: "taco"}); err == nil || !strings.Contains(err.Error(), "not positive") { + t.Fatalf("zero timeout err = %v", err) + } + assertCalls(t, &f, 0) + if rows := runRows(t, cfg.Store); len(rows) != 0 { + t.Fatalf("refused runs wrote %d ai_runs rows, want none", len(rows)) + } +} diff --git a/internal/ai/harness/harness.go b/internal/ai/harness/harness.go index 5ad82d1..5c3bbe6 100644 --- a/internal/ai/harness/harness.go +++ b/internal/ai/harness/harness.go @@ -5,6 +5,8 @@ // and on failure sends exactly one repair turn quoting the problems. If the // repair fails too the run fails with llm.ErrBadOutput; there is no silent // fallback (D-007). Every run is recorded in ai_runs (see Run and Purge). +// A Guard shares in-flight runs of one skill and scope, bounds concurrent +// LLM calls, and each run ends at its skill's timeout. package harness import ( @@ -15,6 +17,7 @@ import ( "fmt" "io" "strings" + "time" "code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/ai/llm" "code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/store" @@ -39,6 +42,10 @@ type Skill[In, Out any] interface { // is validated again; problems found then are repaired like any other. // Its error fails the run without a repair turn. PostProcess(out Out) (Out, error) + // Timeout bounds a whole run, waiting for LLM call slots and the + // repair turn included. It is DefaultTimeout unless the skill needs + // another, and must be positive. + Timeout() time.Duration } // NewProvider builds the provider a run sends its requests to from the @@ -63,6 +70,9 @@ type Config struct { Persona string // Store receives the run's ai_runs row. It must not be nil. Store *store.Store + // Guard is the process's Guard, shared by every run. It must not be + // nil. + Guard *Guard } // validate reports a Config missing what every run needs. @@ -76,6 +86,8 @@ func (c Config) validate() error { return errors.New("config has no model") case c.Store == nil: return errors.New("config has no store") + case c.Guard == nil: + return errors.New("config has no guard") } return nil } @@ -103,6 +115,19 @@ func (e *BadOutputError) Unwrap() error { return llm.ErrBadOutput } // output. A provider error other than llm.ErrBadOutput is returned // unchanged and never repaired. // +// scope names what the run is for, e.g. a week, and must not be empty. +// While a run of s with scope is in flight on cfg.Guard, another Run of s +// with the same scope does not start a run of its own: it waits for that +// run and returns its output or error, whatever its own in and cfg. So a +// scope must identify the input. Runs with different scopes, or of +// different skills, run independently. Every Generate call holds one of +// cfg.Guard's slots while it runs. +// +// A run that has not ended within s.Timeout() fails with an error matching +// llm.ErrTimeout. When ctx is done first, Run returns ctx.Err(); the run +// goes on for the callers still waiting for it, and is cancelled when none +// is left. +// // Every run with a valid Config writes one ai_runs row when it ends, // covering both turns, with its RunStatus and the request and response // bodies redacted (see redact). A render failure is recorded as @@ -111,13 +136,42 @@ func (e *BadOutputError) Unwrap() error { return llm.ErrBadOutput } // before anything else, and writes no row since the row cannot be // written without one. When the row cannot be written the run fails with // that error too, even if its output was usable. -func Run[In, Out any](ctx context.Context, cfg Config, s Skill[In, Out], in In) (Out, error) { +func Run[In, Out any](ctx context.Context, cfg Config, s Skill[In, Out], scope string, in In) (Out, error) { var zero Out if err := cfg.validate(); err != nil { return zero, fmt.Errorf("harness: skill %s v%d: %w", s.Name(), s.Version(), err) } + if scope == "" { + return zero, fmt.Errorf("harness: skill %s v%d: no scope", s.Name(), s.Version()) + } + timeout := s.Timeout() + if timeout <= 0 { + return zero, fmt.Errorf("harness: skill %s v%d: timeout %v is not positive", s.Name(), s.Version(), timeout) + } + val, err := cfg.Guard.do(ctx, flightKey{skill: s.Name(), scope: scope}, func(ctx context.Context) (any, error) { + return runOnce(ctx, cfg, s, timeout, in) + }) + if err != nil { + return zero, err + } + out, ok := val.(Out) + if !ok { + return zero, fmt.Errorf("harness: skill %s v%d: scope %q is shared with a run whose output is %T", s.Name(), s.Version(), scope, val) + } + return out, nil +} + +// runOnce is one execution of s on in, ended at timeout, with its ai_runs row. +func runOnce[In, Out any](ctx context.Context, cfg Config, s Skill[In, Out], timeout time.Duration, in In) (Out, error) { + var zero Out + expired := fmt.Errorf("harness: skill %s v%d: %w: no result within %v", s.Name(), s.Version(), llm.ErrTimeout, timeout) + ctx, cancel := context.WithTimeoutCause(ctx, timeout, expired) + defer cancel() log := &runLog{start: store.Now()} out, status, err := execute(ctx, cfg, s, in, log) + if err != nil && context.Cause(ctx) == expired { + err = fmt.Errorf("%w: %w", expired, err) + } if saveErr := log.save(ctx, cfg, s.Name(), s.Version(), status, err); saveErr != nil { return zero, errors.Join(err, fmt.Errorf("harness: skill %s v%d: record ai_runs row: %w", s.Name(), s.Version(), saveErr)) } @@ -139,6 +193,7 @@ func execute[In, Out any](ctx context.Context, cfg Config, s Skill[In, Out], in if err != nil { return zero, StatusProviderError, fmt.Errorf("harness: skill %s v%d: provider: %w", s.Name(), s.Version(), err) } + provider = limited{Provider: provider, slots: cfg.Guard.slots} schema := s.Schema() temperature := cfg.Settings.Temperature req := llm.Request{ diff --git a/internal/ai/harness/harness_test.go b/internal/ai/harness/harness_test.go index 9e0470b..5ed8347 100644 --- a/internal/ai/harness/harness_test.go +++ b/internal/ai/harness/harness_test.go @@ -23,12 +23,22 @@ func testConfig(t *testing.T, f *llm.Fake) Config { Settings: llm.Settings{Provider: "anthropic", Model: "claude-test", Temperature: 0.7}, NewProvider: func(llm.Settings, string) (llm.Provider, error) { return f, nil }, Store: emptyHousehold(t), + Guard: testGuard(t, DefaultMaxCalls), } } +func testGuard(t *testing.T, maxCalls int) *Guard { + t.Helper() + g, err := NewGuard(maxCalls) + if err != nil { + t.Fatal(err) + } + return g +} + func run(t *testing.T, f *llm.Fake) (echoOutput, error) { t.Helper() - return Run(context.Background(), testConfig(t, f), echoSkill{}, echoInput{Word: "taco"}) + return Run(context.Background(), testConfig(t, f), echoSkill{}, "week-1", echoInput{Word: "taco"}) } func assertCalls(t *testing.T, f *llm.Fake, want int) []llm.Request { @@ -215,7 +225,7 @@ func TestRunProviderErrorOnRepairPassesThrough(t *testing.T) { func TestRunRenderError(t *testing.T) { var f llm.Fake cfg := testConfig(t, &f) - _, err := Run(context.Background(), cfg, echoSkill{}, echoInput{}) + _, err := Run(context.Background(), cfg, echoSkill{}, "week-1", echoInput{}) if err == nil || !strings.Contains(err.Error(), "render") { t.Fatalf("err = %v, want a render error", err) } @@ -235,7 +245,7 @@ func TestRunProviderBuildError(t *testing.T) { var f llm.Fake cfg := testConfig(t, &f) cfg.NewProvider = func(llm.Settings, string) (llm.Provider, error) { return nil, llm.ErrMissingKey } - _, err := Run(context.Background(), cfg, echoSkill{}, echoInput{Word: "taco"}) + _, err := Run(context.Background(), cfg, echoSkill{}, "week-1", echoInput{Word: "taco"}) if !errors.Is(err, llm.ErrMissingKey) { t.Fatalf("err = %v, want ErrMissingKey", err) } @@ -254,11 +264,12 @@ func TestRunConfigError(t *testing.T) { "provider name": func(c *Config) { c.Settings.Provider = "" }, "model": func(c *Config) { c.Settings.Model = "" }, "store": func(c *Config) { c.Store = nil }, + "guard": func(c *Config) { c.Guard = nil }, } { t.Run(name, func(t *testing.T) { cfg := full edit(&cfg) - _, err := Run(context.Background(), cfg, echoSkill{}, echoInput{Word: "taco"}) + _, err := Run(context.Background(), cfg, echoSkill{}, "week-1", echoInput{Word: "taco"}) if err == nil || !strings.Contains(err.Error(), "config has no "+name) { t.Fatalf("err = %v, want missing %s", err, name) } diff --git a/internal/ai/harness/prompt_test.go b/internal/ai/harness/prompt_test.go index deda41e..7a8345b 100644 --- a/internal/ai/harness/prompt_test.go +++ b/internal/ai/harness/prompt_test.go @@ -119,7 +119,7 @@ func TestRunSendsLayeredPrompt(t *testing.T) { in := echoInput{Word: "Taco", Context: tacoContext} cfg := testConfig(t, &f) cfg.Persona = persona - if _, err := Run(context.Background(), cfg, echoSkill{}, in); err != nil { + if _, err := Run(context.Background(), cfg, echoSkill{}, "week-1", in); err != nil { t.Fatal(err) } req := assertCalls(t, &f, 1)[0] diff --git a/internal/ai/harness/runs_test.go b/internal/ai/harness/runs_test.go index 339f40d..963207a 100644 --- a/internal/ai/harness/runs_test.go +++ b/internal/ai/harness/runs_test.go @@ -128,7 +128,7 @@ func TestRunWritesOneRowWithStatus(t *testing.T) { f.QueueGenerate(tc.script...) cfg := testConfig(t, &f) before := store.Now() - _, runErr := Run(context.Background(), cfg, tc.skill, echoInput{Word: "taco"}) + _, runErr := Run(context.Background(), cfg, tc.skill, "week-1", echoInput{Word: "taco"}) if (runErr == nil) != (tc.status == StatusOK || tc.status == StatusRepaired) { t.Fatalf("Run err = %v for status %s", runErr, tc.status) } @@ -186,7 +186,7 @@ func TestRunRowPerRun(t *testing.T) { cfg := testConfig(t, &f) f.QueueGenerate(reply(`{"echo":"taco"}`), reply(`{"echo":"TACO"}`), reply(`{"echo":"taco"}`), llm.GenerateResult{Err: llm.ErrAuth}) for range 3 { - _, _ = Run(context.Background(), cfg, echoSkill{}, echoInput{Word: "taco"}) + _, _ = Run(context.Background(), cfg, echoSkill{}, "week-1", echoInput{Word: "taco"}) } assertCalls(t, &f, 4) var got []string @@ -204,7 +204,7 @@ func TestRunRecordsCancelledRun(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() f.QueueGenerate(llm.GenerateResult{Err: context.Canceled}) - if _, err := Run(ctx, cfg, echoSkill{}, echoInput{Word: "taco"}); !errors.Is(err, context.Canceled) { + if _, err := Run(ctx, cfg, echoSkill{}, "week-1", echoInput{Word: "taco"}); !errors.Is(err, context.Canceled) { t.Fatalf("err = %v, want context.Canceled", err) } if rows := runRows(t, cfg.Store); len(rows) != 1 || rows[0].Status != string(StatusProviderError) { @@ -217,7 +217,7 @@ func TestRunFailsWhenRowCannotBeWritten(t *testing.T) { cfg := testConfig(t, &f) cfg.Settings.Provider = "gemini" // refused by the provider CHECK f.QueueGenerate(reply(`{"echo":"taco"}`)) - _, err := Run(context.Background(), cfg, echoSkill{}, echoInput{Word: "taco"}) + _, err := Run(context.Background(), cfg, echoSkill{}, "week-1", echoInput{Word: "taco"}) if err == nil || !strings.Contains(err.Error(), "record ai_runs row") { t.Fatalf("err = %v, want the row write failure", err) } @@ -276,7 +276,7 @@ func TestRunRowsHoldNoSecrets(t *testing.T) { return &f, nil } cfg.Persona = "Our kitchen. " + leaky - _, _ = Run(context.Background(), cfg, echoSkill{}, echoInput{Word: "tacos " + seededWebhook}) + _, _ = Run(context.Background(), cfg, echoSkill{}, "week-1", echoInput{Word: "tacos " + seededWebhook}) assertCalls(t, &f, len(script)) rows := runRows(t, cfg.Store) @@ -334,8 +334,9 @@ func TestRunRowsHoldNoProviderKey(t *testing.T) { APIKey: seededKey, NewProvider: presets.FromSettings, Store: emptyHousehold(t), + Guard: testGuard(t, DefaultMaxCalls), } - out, err := Run(context.Background(), cfg, echoSkill{}, echoInput{Word: "taco"}) + out, err := Run(context.Background(), cfg, echoSkill{}, "week-1", echoInput{Word: "taco"}) if err != nil { t.Fatal(err) } diff --git a/internal/ai/llm/fake.go b/internal/ai/llm/fake.go index b418d82..af50192 100644 --- a/internal/ai/llm/fake.go +++ b/internal/ai/llm/fake.go @@ -11,7 +11,9 @@ import ( // Fake is a scripted Provider for tests. Each method returns its queued // results in order, one per call, and records every Request Generate // receives. A call with nothing left in its queue fails, so an unexpected -// call is never silently answered. It is safe for concurrent use. +// call is never silently answered. A Generate result can block until the +// test releases it, so tests can hold calls in flight. It is safe for +// concurrent use. type Fake struct { mu sync.Mutex generate []GenerateResult @@ -26,6 +28,13 @@ var _ Provider = (*Fake)(nil) type GenerateResult struct { Response Response Err error + // Started, when set, is closed once the call has taken this result and + // recorded its Request, before it blocks on Release. + Started chan struct{} + // Release, when set, makes the call block until it is closed. When + // the call's context is done first, the call returns ctx.Err() instead + // of Response and Err, as a real provider does. + Release <-chan struct{} } // ListModelsResult is one scripted ListModels outcome. @@ -55,16 +64,29 @@ func (f *Fake) QueueVerify(results ...error) { f.verify = append(f.verify, results...) } -// Generate records req and returns the next scripted result. -func (f *Fake) Generate(_ context.Context, req Request) (Response, error) { +// Generate records req and returns the next scripted result, blocking +// first when the result has a Release channel. +func (f *Fake) Generate(ctx context.Context, req Request) (Response, error) { f.mu.Lock() - defer f.mu.Unlock() f.requests = append(f.requests, cloneRequest(req)) if len(f.generate) == 0 { - return Response{}, fmt.Errorf("llm: fake: unscripted Generate call %d", len(f.requests)) + n := len(f.requests) + f.mu.Unlock() + return Response{}, fmt.Errorf("llm: fake: unscripted Generate call %d", n) } r := f.generate[0] f.generate = f.generate[1:] + f.mu.Unlock() + if r.Started != nil { + close(r.Started) + } + if r.Release != nil { + select { + case <-r.Release: + case <-ctx.Done(): + return Response{}, ctx.Err() + } + } return r.Response, r.Err } diff --git a/internal/ai/llm/fake_test.go b/internal/ai/llm/fake_test.go index bba1b3c..2e0880b 100644 --- a/internal/ai/llm/fake_test.go +++ b/internal/ai/llm/fake_test.go @@ -136,3 +136,37 @@ func TestFakeConcurrentUse(t *testing.T) { t.Fatalf("%d Generate results unused", g) } } + +func TestFakeGenerateBlocksUntilReleased(t *testing.T) { + var f Fake + started, release := make(chan struct{}), make(chan struct{}) + f.QueueGenerate(GenerateResult{Response: Response{Text: "done"}, Started: started, Release: release}) + got := make(chan string) + go func() { + resp, err := f.Generate(context.Background(), Request{}) + if err != nil { + t.Error(err) + } + got <- resp.Text + }() + <-started + select { + case text := <-got: + t.Fatalf("Generate returned %q before release", text) + case <-time.After(20 * time.Millisecond): + } + close(release) + if text := <-got; text != "done" { + t.Fatalf("Generate = %q, want the scripted reply", text) + } +} + +func TestFakeBlockedGenerateHonoursContext(t *testing.T) { + var f Fake + f.QueueGenerate(GenerateResult{Response: Response{Text: "never"}, Release: make(chan struct{})}) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + if _, err := f.Generate(ctx, Request{}); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("err = %v, want context.DeadlineExceeded", err) + } +}