Feature 09: AI harness and skills framework (#62)

Co-authored-by: dev_agent <dev_agent@cubecraftlabs.com>

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)
+}