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..e292526 --- /dev/null +++ b/internal/ai/harness/guard_test.go @@ -0,0 +1,362 @@ +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 } + +// slowSkill is timeoutSkill whose PostProcess succeeds only after its +// timeout has passed. +type slowSkill struct { + timeoutSkill +} + +func (s slowSkill) PostProcess(out echoOutput) (echoOutput, error) { + time.Sleep(3 * s.timeout) + return s.timeoutSkill.PostProcess(out) +} + +// 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("work that succeeds late", func(t *testing.T) { + var f llm.Fake + cfg := testConfig(t, &f) + f.QueueGenerate(reply(`{"echo":"taco"}`)) + s := slowSkill{timeoutSkill{timeout: 20 * time.Millisecond}} + out, err := Run(context.Background(), cfg, s, "week-1", echoInput{Word: "taco"}) + if !errors.Is(err, llm.ErrTimeout) { + t.Fatalf("err = %v, want ErrTimeout", err) + } + if out != (echoOutput{}) { + t.Fatalf("out = %+v, want no output", out) + } + 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("row saved late", func(t *testing.T) { + var f llm.Fake + cfg := testConfig(t, &f) + f.QueueGenerate(reply(`{"echo":"taco"}`)) + timeout := 100 * time.Millisecond + // Holding the store's only writer connection makes the save wait + // past the deadline, after the work itself has ended in time. + conn, err := cfg.Store.Writer.Conn(context.Background()) + if err != nil { + t.Fatal(err) + } + released := time.AfterFunc(3*timeout, func() { conn.Close() }) + defer released.Stop() + out, err := Run(context.Background(), cfg, timeoutSkill{timeout: timeout}, "week-1", echoInput{Word: "taco"}) + if !errors.Is(err, llm.ErrTimeout) { + t.Fatalf("err = %v, want ErrTimeout", err) + } + if out != (echoOutput{}) { + t.Fatalf("out = %+v, want no output", out) + } + assertCalls(t, &f, 1) + // An ok row with no error proves the work ended in time and only + // the save ran past the deadline. + rows := runRows(t, cfg.Store) + if len(rows) != 1 || rows[0].Status != string(StatusOK) || rows[0].Error.Valid { + t.Fatalf("rows = %+v, want one ok row with no error", 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..0ec9ee3 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,19 +136,62 @@ 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) + deadline := time.Now().Add(timeout) + ctx, cancel := context.WithDeadlineCause(ctx, deadline, expired) + defer cancel() log := &runLog{start: store.Now()} out, status, err := execute(ctx, cfg, s, in, log) + // Work that only ends after the deadline, such as a slow PostProcess, + // may return no error; it is still a timeout, not a success. The clock + // is checked as well as ctx because ctx's timer may not have fired yet. + if context.Cause(ctx) == expired || !time.Now().Before(deadline) { + if err == nil { + err, status = expired, StatusProviderError + } else { + 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)) } if err != nil { return zero, err } + // The save may wait past the deadline for the store's single writer. + // The run has then not ended in time, so it fails, though its row + // records the work's outcome as it stood before the save. + if !time.Now().Before(deadline) { + return zero, expired + } return out, nil } @@ -139,6 +207,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) + } +}