Concurrency guards: single-flight, semaphore and timeouts (#61)
Co-authored-by: dev_agent <dev_agent@cubecraftlabs.com>
This commit was merged in pull request #61.
This commit is contained in:
@@ -162,6 +162,16 @@ automatic** (D-021).
|
|||||||
with key-like query parameters are removed, inside JSON strings too.
|
with key-like query parameters are removed, inside JSON strings too.
|
||||||
`harness.Purge` deletes rows older than `harness.RunRetention` (30 days),
|
`harness.Purge` deletes rows older than `harness.RunRetention` (30 days),
|
||||||
clearing the `suggestion_batches` and `week_plans` references first.
|
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
|
- **Prompt layering** (D-007): each skill version's prompt is a
|
||||||
`text/template` file embedded from `ai/skills/prompts/`
|
`text/template` file embedded from `ai/skills/prompts/`
|
||||||
(`<skill>.v<version>.tmpl`). The system prompt is always, in order:
|
(`<skill>.v<version>.tmpl`). The system prompt is always, in order:
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"strings"
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
"code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/ai/llm"
|
"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)
|
out.Echo = strings.TrimSpace(out.Echo)
|
||||||
return out, nil
|
return out, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (echoSkill) Timeout() time.Duration { return DefaultTimeout }
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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))
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -5,6 +5,8 @@
|
|||||||
// and on failure sends exactly one repair turn quoting the problems. If the
|
// 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
|
// 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).
|
// 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
|
package harness
|
||||||
|
|
||||||
import (
|
import (
|
||||||
@@ -15,6 +17,7 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"strings"
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
"code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/ai/llm"
|
"code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/ai/llm"
|
||||||
"code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/store"
|
"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.
|
// is validated again; problems found then are repaired like any other.
|
||||||
// Its error fails the run without a repair turn.
|
// Its error fails the run without a repair turn.
|
||||||
PostProcess(out Out) (Out, error)
|
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
|
// NewProvider builds the provider a run sends its requests to from the
|
||||||
@@ -63,6 +70,9 @@ type Config struct {
|
|||||||
Persona string
|
Persona string
|
||||||
// Store receives the run's ai_runs row. It must not be nil.
|
// Store receives the run's ai_runs row. It must not be nil.
|
||||||
Store *store.Store
|
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.
|
// 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")
|
return errors.New("config has no model")
|
||||||
case c.Store == nil:
|
case c.Store == nil:
|
||||||
return errors.New("config has no store")
|
return errors.New("config has no store")
|
||||||
|
case c.Guard == nil:
|
||||||
|
return errors.New("config has no guard")
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -103,6 +115,19 @@ func (e *BadOutputError) Unwrap() error { return llm.ErrBadOutput }
|
|||||||
// output. A provider error other than llm.ErrBadOutput is returned
|
// output. A provider error other than llm.ErrBadOutput is returned
|
||||||
// unchanged and never repaired.
|
// 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,
|
// 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
|
// covering both turns, with its RunStatus and the request and response
|
||||||
// bodies redacted (see redact). A render failure is recorded as
|
// 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
|
// 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
|
// written without one. When the row cannot be written the run fails with
|
||||||
// that error too, even if its output was usable.
|
// 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
|
var zero Out
|
||||||
if err := cfg.validate(); err != nil {
|
if err := cfg.validate(); err != nil {
|
||||||
return zero, fmt.Errorf("harness: skill %s v%d: %w", s.Name(), s.Version(), err)
|
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()}
|
log := &runLog{start: store.Now()}
|
||||||
out, status, err := execute(ctx, cfg, s, in, log)
|
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 {
|
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))
|
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 {
|
if err != nil {
|
||||||
return zero, err
|
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
|
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 {
|
if err != nil {
|
||||||
return zero, StatusProviderError, fmt.Errorf("harness: skill %s v%d: provider: %w", s.Name(), s.Version(), err)
|
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()
|
schema := s.Schema()
|
||||||
temperature := cfg.Settings.Temperature
|
temperature := cfg.Settings.Temperature
|
||||||
req := llm.Request{
|
req := llm.Request{
|
||||||
|
|||||||
@@ -23,12 +23,22 @@ func testConfig(t *testing.T, f *llm.Fake) Config {
|
|||||||
Settings: llm.Settings{Provider: "anthropic", Model: "claude-test", Temperature: 0.7},
|
Settings: llm.Settings{Provider: "anthropic", Model: "claude-test", Temperature: 0.7},
|
||||||
NewProvider: func(llm.Settings, string) (llm.Provider, error) { return f, nil },
|
NewProvider: func(llm.Settings, string) (llm.Provider, error) { return f, nil },
|
||||||
Store: emptyHousehold(t),
|
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) {
|
func run(t *testing.T, f *llm.Fake) (echoOutput, error) {
|
||||||
t.Helper()
|
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 {
|
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) {
|
func TestRunRenderError(t *testing.T) {
|
||||||
var f llm.Fake
|
var f llm.Fake
|
||||||
cfg := testConfig(t, &f)
|
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") {
|
if err == nil || !strings.Contains(err.Error(), "render") {
|
||||||
t.Fatalf("err = %v, want a render error", err)
|
t.Fatalf("err = %v, want a render error", err)
|
||||||
}
|
}
|
||||||
@@ -235,7 +245,7 @@ func TestRunProviderBuildError(t *testing.T) {
|
|||||||
var f llm.Fake
|
var f llm.Fake
|
||||||
cfg := testConfig(t, &f)
|
cfg := testConfig(t, &f)
|
||||||
cfg.NewProvider = func(llm.Settings, string) (llm.Provider, error) { return nil, llm.ErrMissingKey }
|
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) {
|
if !errors.Is(err, llm.ErrMissingKey) {
|
||||||
t.Fatalf("err = %v, want ErrMissingKey", err)
|
t.Fatalf("err = %v, want ErrMissingKey", err)
|
||||||
}
|
}
|
||||||
@@ -254,11 +264,12 @@ func TestRunConfigError(t *testing.T) {
|
|||||||
"provider name": func(c *Config) { c.Settings.Provider = "" },
|
"provider name": func(c *Config) { c.Settings.Provider = "" },
|
||||||
"model": func(c *Config) { c.Settings.Model = "" },
|
"model": func(c *Config) { c.Settings.Model = "" },
|
||||||
"store": func(c *Config) { c.Store = nil },
|
"store": func(c *Config) { c.Store = nil },
|
||||||
|
"guard": func(c *Config) { c.Guard = nil },
|
||||||
} {
|
} {
|
||||||
t.Run(name, func(t *testing.T) {
|
t.Run(name, func(t *testing.T) {
|
||||||
cfg := full
|
cfg := full
|
||||||
edit(&cfg)
|
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) {
|
if err == nil || !strings.Contains(err.Error(), "config has no "+name) {
|
||||||
t.Fatalf("err = %v, want missing %s", err, name)
|
t.Fatalf("err = %v, want missing %s", err, name)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -119,7 +119,7 @@ func TestRunSendsLayeredPrompt(t *testing.T) {
|
|||||||
in := echoInput{Word: "Taco", Context: tacoContext}
|
in := echoInput{Word: "Taco", Context: tacoContext}
|
||||||
cfg := testConfig(t, &f)
|
cfg := testConfig(t, &f)
|
||||||
cfg.Persona = persona
|
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)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
req := assertCalls(t, &f, 1)[0]
|
req := assertCalls(t, &f, 1)[0]
|
||||||
|
|||||||
@@ -128,7 +128,7 @@ func TestRunWritesOneRowWithStatus(t *testing.T) {
|
|||||||
f.QueueGenerate(tc.script...)
|
f.QueueGenerate(tc.script...)
|
||||||
cfg := testConfig(t, &f)
|
cfg := testConfig(t, &f)
|
||||||
before := store.Now()
|
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) {
|
if (runErr == nil) != (tc.status == StatusOK || tc.status == StatusRepaired) {
|
||||||
t.Fatalf("Run err = %v for status %s", runErr, tc.status)
|
t.Fatalf("Run err = %v for status %s", runErr, tc.status)
|
||||||
}
|
}
|
||||||
@@ -186,7 +186,7 @@ func TestRunRowPerRun(t *testing.T) {
|
|||||||
cfg := testConfig(t, &f)
|
cfg := testConfig(t, &f)
|
||||||
f.QueueGenerate(reply(`{"echo":"taco"}`), reply(`{"echo":"TACO"}`), reply(`{"echo":"taco"}`), llm.GenerateResult{Err: llm.ErrAuth})
|
f.QueueGenerate(reply(`{"echo":"taco"}`), reply(`{"echo":"TACO"}`), reply(`{"echo":"taco"}`), llm.GenerateResult{Err: llm.ErrAuth})
|
||||||
for range 3 {
|
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)
|
assertCalls(t, &f, 4)
|
||||||
var got []string
|
var got []string
|
||||||
@@ -204,7 +204,7 @@ func TestRunRecordsCancelledRun(t *testing.T) {
|
|||||||
ctx, cancel := context.WithCancel(context.Background())
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
cancel()
|
cancel()
|
||||||
f.QueueGenerate(llm.GenerateResult{Err: context.Canceled})
|
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)
|
t.Fatalf("err = %v, want context.Canceled", err)
|
||||||
}
|
}
|
||||||
if rows := runRows(t, cfg.Store); len(rows) != 1 || rows[0].Status != string(StatusProviderError) {
|
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 := testConfig(t, &f)
|
||||||
cfg.Settings.Provider = "gemini" // refused by the provider CHECK
|
cfg.Settings.Provider = "gemini" // refused by the provider CHECK
|
||||||
f.QueueGenerate(reply(`{"echo":"taco"}`))
|
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") {
|
if err == nil || !strings.Contains(err.Error(), "record ai_runs row") {
|
||||||
t.Fatalf("err = %v, want the row write failure", err)
|
t.Fatalf("err = %v, want the row write failure", err)
|
||||||
}
|
}
|
||||||
@@ -276,7 +276,7 @@ func TestRunRowsHoldNoSecrets(t *testing.T) {
|
|||||||
return &f, nil
|
return &f, nil
|
||||||
}
|
}
|
||||||
cfg.Persona = "Our kitchen. " + leaky
|
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))
|
assertCalls(t, &f, len(script))
|
||||||
|
|
||||||
rows := runRows(t, cfg.Store)
|
rows := runRows(t, cfg.Store)
|
||||||
@@ -334,8 +334,9 @@ func TestRunRowsHoldNoProviderKey(t *testing.T) {
|
|||||||
APIKey: seededKey,
|
APIKey: seededKey,
|
||||||
NewProvider: presets.FromSettings,
|
NewProvider: presets.FromSettings,
|
||||||
Store: emptyHousehold(t),
|
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 {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
|||||||
+27
-5
@@ -11,7 +11,9 @@ import (
|
|||||||
// Fake is a scripted Provider for tests. Each method returns its queued
|
// Fake is a scripted Provider for tests. Each method returns its queued
|
||||||
// results in order, one per call, and records every Request Generate
|
// results in order, one per call, and records every Request Generate
|
||||||
// receives. A call with nothing left in its queue fails, so an unexpected
|
// 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 {
|
type Fake struct {
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
generate []GenerateResult
|
generate []GenerateResult
|
||||||
@@ -26,6 +28,13 @@ var _ Provider = (*Fake)(nil)
|
|||||||
type GenerateResult struct {
|
type GenerateResult struct {
|
||||||
Response Response
|
Response Response
|
||||||
Err error
|
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.
|
// ListModelsResult is one scripted ListModels outcome.
|
||||||
@@ -55,16 +64,29 @@ func (f *Fake) QueueVerify(results ...error) {
|
|||||||
f.verify = append(f.verify, results...)
|
f.verify = append(f.verify, results...)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Generate records req and returns the next scripted result.
|
// Generate records req and returns the next scripted result, blocking
|
||||||
func (f *Fake) Generate(_ context.Context, req Request) (Response, error) {
|
// first when the result has a Release channel.
|
||||||
|
func (f *Fake) Generate(ctx context.Context, req Request) (Response, error) {
|
||||||
f.mu.Lock()
|
f.mu.Lock()
|
||||||
defer f.mu.Unlock()
|
|
||||||
f.requests = append(f.requests, cloneRequest(req))
|
f.requests = append(f.requests, cloneRequest(req))
|
||||||
if len(f.generate) == 0 {
|
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]
|
r := f.generate[0]
|
||||||
f.generate = f.generate[1:]
|
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
|
return r.Response, r.Err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -136,3 +136,37 @@ func TestFakeConcurrentUse(t *testing.T) {
|
|||||||
t.Fatalf("%d Generate results unused", g)
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user