WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/models/requestbudget/context.go

@@ -0,0 +1,88 @@
+// Package requestbudget carries a durable request reservation to provider transports.
+package requestbudget
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+)
+
+var ErrStopped = errors.New("wiki recovery budget stopped")
+
+type Guard interface {
+	Key() string
+	Check(context.Context) error
+	Reserve(context.Context, string, string) error
+}
+
+type guardKey struct{}
+type phaseKey struct{}
+type modelKey struct{}
+
+func WithGuard(ctx context.Context, guard Guard) context.Context {
+	return context.WithValue(ctx, guardKey{}, guard)
+}
+
+func FromContext(ctx context.Context) Guard {
+	guard, _ := ctx.Value(guardKey{}).(Guard)
+	return guard
+}
+
+func WithPhase(ctx context.Context, phase string) context.Context {
+	return context.WithValue(ctx, phaseKey{}, phase)
+}
+
+func Phase(ctx context.Context) string {
+	phase, _ := ctx.Value(phaseKey{}).(string)
+	return phase
+}
+
+func WithModel(ctx context.Context, modelID string) context.Context {
+	return context.WithValue(ctx, modelKey{}, modelID)
+}
+
+func Check(ctx context.Context) error {
+	if guard := FromContext(ctx); guard != nil {
+		return guard.Check(ctx)
+	}
+	return ctx.Err()
+}
+
+// RequireSupported prevents a controlled request from reaching an adapter
+// whose actual sends are not instrumented. Ordinary requests are unaffected.
+func RequireSupported(ctx context.Context, supported bool) error {
+	if FromContext(ctx) != nil && !supported {
+		return fmt.Errorf("%w: unsupported provider transport", ErrStopped)
+	}
+	return nil
+}
+
+// Transport reserves each RoundTrip, including redirects and SDK fallbacks.
+// Failed or cancelled sends intentionally consume their reservation.
+type Transport struct{ Base http.RoundTripper }
+
+func (t Transport) RoundTrip(req *http.Request) (*http.Response, error) {
+	if err := req.Context().Err(); err != nil {
+		return nil, err
+	}
+	if guard := FromContext(req.Context()); guard != nil {
+		modelID, _ := req.Context().Value(modelKey{}).(string)
+		if modelID == "" {
+			return nil, fmt.Errorf("%w: missing provider model identity", ErrStopped)
+		}
+		if err := guard.Reserve(req.Context(), modelID, Phase(req.Context())); err != nil {
+			return nil, fmt.Errorf("%w: %w", ErrStopped, err)
+		}
+		// A replayable POST with an idempotency header can otherwise be retried
+		// inside http.Transport without reserving again. Clone so Client.Do
+		// keeps the original GetBody for separately accounted 307/308 redirects.
+		req = req.Clone(req.Context())
+		req.GetBody = nil
+	}
+	base := t.Base
+	if base == nil {
+		base = http.DefaultTransport
+	}
+	return base.RoundTrip(req)
+}

internal/models/requestbudget/context_test.go

@@ -0,0 +1,205 @@
+package requestbudget
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+)
+
+type testGuard struct {
+	mu              sync.Mutex
+	limit, reserved int
+	err             error
+	model, phase    string
+}
+
+func (g *testGuard) Key() string                     { return "local-test" }
+func (g *testGuard) Check(ctx context.Context) error { return ctx.Err() }
+func (g *testGuard) Reserve(ctx context.Context, model, phase string) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	if g.err != nil {
+		return g.err
+	}
+	if g.reserved >= g.limit {
+		return ErrStopped
+	}
+	g.reserved++
+	g.model, g.phase = model, phase
+	return nil
+}
+
+type roundTripFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
+func budgetTestContext(g *testGuard) context.Context {
+	return WithModel(WithPhase(WithGuard(context.Background(), g), "phase"), "adapter-model")
+}
+
+func TestTransportBlocksBeforeBaseOnGuardFailure(t *testing.T) {
+	for _, reason := range []string{"budget", "storage error", "missing model", "cancelled"} {
+		t.Run(reason, func(t *testing.T) {
+			g := &testGuard{limit: 0}
+			ctx := budgetTestContext(g)
+			if reason == "storage error" {
+				g.err = errors.New("durable ledger unavailable")
+			}
+			if reason == "missing model" {
+				ctx = WithGuard(context.Background(), g)
+			}
+			if reason == "cancelled" {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithCancel(ctx)
+				cancel()
+			}
+			sends := 0
+			transport := Transport{Base: roundTripFunc(func(*http.Request) (*http.Response, error) { sends++; return nil, errors.New("should not send") })}
+			req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://local.test", nil)
+			_, err := transport.RoundTrip(req)
+			if err == nil || sends != 0 || g.reserved != 0 {
+				t.Fatalf("failed closed: sends=%d reserved=%d err=%v", sends, g.reserved, err)
+			}
+		})
+	}
+}
+
+func TestTransportNeverRefundsFailedSend(t *testing.T) {
+	g := &testGuard{limit: 1}
+	sends := 0
+	transport := Transport{Base: roundTripFunc(func(*http.Request) (*http.Response, error) { sends++; return nil, io.ErrUnexpectedEOF })}
+	req, _ := http.NewRequestWithContext(budgetTestContext(g), http.MethodPost, "http://local.test", nil)
+	if _, err := transport.RoundTrip(req); !errors.Is(err, io.ErrUnexpectedEOF) {
+		t.Fatal(err)
+	}
+	if _, err := transport.RoundTrip(req); !errors.Is(err, ErrStopped) {
+		t.Fatal(err)
+	}
+	if sends != 1 || g.reserved != 1 || g.model != "adapter-model" || g.phase != "phase" {
+		t.Fatalf("reservation lost: %+v sends=%d", g, sends)
+	}
+}
+
+func TestTransportCountsRedirectsAndPreservesContext(t *testing.T) {
+	for _, limit := range []int{1, 2} {
+		t.Run(fmt.Sprint(limit), func(t *testing.T) {
+			var sends atomic.Int32
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				sends.Add(1)
+				if r.URL.Path == "/first" {
+					http.Redirect(w, r, "/second", http.StatusTemporaryRedirect)
+					return
+				}
+				fmt.Fprint(w, "ok")
+			}))
+			defer server.Close()
+			g := &testGuard{limit: limit}
+			client := &http.Client{Transport: Transport{Base: http.DefaultTransport}}
+			req, _ := http.NewRequestWithContext(budgetTestContext(g), http.MethodPost, server.URL+"/first", strings.NewReader("body"))
+			resp, err := client.Do(req)
+			if resp != nil {
+				resp.Body.Close()
+			}
+			if limit == 1 && !errors.Is(err, ErrStopped) {
+				t.Fatalf("redirect escaped budget: %v", err)
+			}
+			if limit == 2 && err != nil {
+				t.Fatal(err)
+			}
+			if int(sends.Load()) != limit || g.reserved != limit || g.model != "adapter-model" || g.phase != "phase" {
+				t.Fatalf("redirect context lost: sends=%d guard=%+v", sends.Load(), g)
+			}
+		})
+	}
+}
+
+func TestTransportConcurrentSendsCannotExceedGuard(t *testing.T) {
+	g := &testGuard{limit: 3}
+	var sends atomic.Int32
+	transport := Transport{Base: roundTripFunc(func(*http.Request) (*http.Response, error) { sends.Add(1); return nil, io.ErrUnexpectedEOF })}
+	var wg sync.WaitGroup
+	for i := 0; i < 20; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			req, _ := http.NewRequestWithContext(budgetTestContext(g), http.MethodPost, "http://local.test", nil)
+			transport.RoundTrip(req)
+		}()
+	}
+	wg.Wait()
+	if sends.Load() != 3 || g.reserved != 3 {
+		t.Fatalf("concurrent sends=%d reservations=%d", sends.Load(), g.reserved)
+	}
+}
+
+func TestTransportPreventsUncountedPostReplay(t *testing.T) {
+	for _, controlled := range []bool{false, true} {
+		t.Run(fmt.Sprintf("controlled=%t", controlled), func(t *testing.T) {
+			var sends atomic.Int32
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				io.Copy(io.Discard, r.Body)
+				if r.URL.Path == "/bill" && sends.Add(1) == 1 {
+					connection, _, err := w.(http.Hijacker).Hijack()
+					if err == nil {
+						connection.Close()
+					}
+					return
+				}
+				fmt.Fprint(w, "ok")
+			}))
+			defer server.Close()
+			base := &http.Transport{}
+			defer base.CloseIdleConnections()
+			client := &http.Client{Transport: Transport{Base: base}}
+			// A reused connection is necessary to exercise Go's transparent retry.
+			warm, err := client.Get(server.URL + "/warm")
+			if err != nil {
+				t.Fatal(err)
+			}
+			io.Copy(io.Discard, warm.Body)
+			warm.Body.Close()
+			ctx := context.Background()
+			g := &testGuard{limit: 1}
+			if controlled {
+				ctx = budgetTestContext(g)
+			}
+			req, _ := http.NewRequestWithContext(ctx, http.MethodPost, server.URL+"/bill", strings.NewReader("body"))
+			req.Header.Set("Idempotency-Key", "local-test")
+			resp, err := client.Do(req)
+			if resp != nil {
+				resp.Body.Close()
+			}
+			if req.GetBody == nil {
+				t.Fatal("caller request was mutated")
+			}
+			if controlled {
+				if err == nil || sends.Load() != 1 || g.reserved != 1 {
+					t.Fatalf("hidden retry: sends=%d reserved=%d err=%v", sends.Load(), g.reserved, err)
+				}
+			} else if err != nil || sends.Load() != 2 {
+				t.Fatalf("test did not exercise underlying retry: sends=%d err=%v", sends.Load(), err)
+			}
+		})
+	}
+}
+
+func TestRequireSupportedOnlyRestrictsControlledContexts(t *testing.T) {
+	if err := RequireSupported(context.Background(), false); err != nil {
+		t.Fatal(err)
+	}
+	if err := RequireSupported(budgetTestContext(&testGuard{}), true); err != nil {
+		t.Fatal(err)
+	}
+	if err := RequireSupported(budgetTestContext(&testGuard{}), false); !errors.Is(err, ErrStopped) {
+		t.Fatal(err)
+	}
+}

internal/models/requestbudget/graph.go

@@ -0,0 +1,73 @@
+package requestbudget
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+)
+
+const GraphMaxSends = 3
+const GraphTimeout = 180 * time.Second
+const GraphPhase = "graph_extraction"
+
+// GraphExecution is shared by extraction and its one correction. Callers must
+// consume a durable single-use receipt before binding this in-memory allowance.
+type GraphExecution struct {
+	mu       sync.Mutex
+	key      string
+	model    string
+	reserved int
+	stopped  error
+	cancel   context.CancelCauseFunc
+}
+
+func BindGraphExecution(ctx context.Context, key, model string) (context.Context, *GraphExecution, context.CancelFunc, error) {
+	if key == "" || model == "" || FromContext(ctx) != nil {
+		return ctx, nil, func() {}, fmt.Errorf("%w: invalid graph execution scope", ErrStopped)
+	}
+	if err := ctx.Err(); err != nil {
+		return ctx, nil, func() {}, err
+	}
+	ctx, deadlineCancel := context.WithTimeout(ctx, GraphTimeout)
+	ctx, cancel := context.WithCancelCause(ctx)
+	g := &GraphExecution{key: "graph:" + key, model: model, cancel: cancel}
+	ctx = WithPhase(WithGuard(ctx, g), GraphPhase)
+	return ctx, g, func() { deadlineCancel(); cancel(nil) }, nil
+}
+
+func (g *GraphExecution) Key() string { return g.key }
+
+func (g *GraphExecution) checkLocked(ctx context.Context) error {
+	if g.stopped != nil {
+		return g.stopped
+	}
+	return ctx.Err()
+}
+
+func (g *GraphExecution) Check(ctx context.Context) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	return g.checkLocked(ctx)
+}
+
+func (g *GraphExecution) Reserve(ctx context.Context, model, phase string) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if err := g.checkLocked(ctx); err != nil {
+		return err
+	}
+	if model != g.model || phase != GraphPhase || g.reserved >= GraphMaxSends {
+		g.stopped = fmt.Errorf("%w: graph send limit or scope mismatch", ErrStopped)
+		g.cancel(g.stopped)
+		return g.stopped
+	}
+	g.reserved++
+	return nil
+}
+
+func (g *GraphExecution) Reservations() int {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	return g.reserved
+}

internal/models/requestbudget/graph_test.go

@@ -0,0 +1,98 @@
+package requestbudget
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestGraphExecutionRefusesFourthSendAndScopeChanges(t *testing.T) {
+	for _, mismatch := range []string{"limit", "model", "phase"} {
+		t.Run(mismatch, func(t *testing.T) {
+			ctx, guard, cancel, err := BindGraphExecution(context.Background(), "spent-receipt", "model")
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer cancel()
+			ctx = WithModel(ctx, "model")
+			if mismatch == "model" {
+				ctx = WithModel(ctx, "changed")
+			}
+			if mismatch == "phase" {
+				ctx = WithPhase(ctx, "changed")
+			}
+			sends := 0
+			transport := Transport{Base: roundTripFunc(func(*http.Request) (*http.Response, error) {
+				sends++
+				return nil, errors.New("provider failure")
+			})}
+			for i := 0; i < 4; i++ {
+				req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://local.test", nil)
+				_, _ = transport.RoundTrip(req)
+			}
+			want := 0
+			if mismatch == "limit" {
+				want = 3
+			}
+			if sends != want || guard.Reservations() != want || !errors.Is(guard.Check(ctx), ErrStopped) {
+				t.Fatalf("sends=%d reserved=%d stop=%v", sends, guard.Reservations(), guard.Check(ctx))
+			}
+		})
+	}
+}
+
+func TestGraphExecutionCancelsBlockedHTTPAtInheritedDeadline(t *testing.T) {
+	var sends atomic.Int32
+	arrived := make(chan struct{})
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		sends.Add(1)
+		close(arrived)
+		<-r.Context().Done()
+	}))
+	defer server.Close()
+	parent, parentCancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
+	defer parentCancel()
+	ctx, guard, cancel, err := BindGraphExecution(parent, "spent-receipt", "model")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cancel()
+	wantDeadline, _ := parent.Deadline()
+	if deadline, _ := ctx.Deadline(); !deadline.Equal(wantDeadline) {
+		t.Fatal("deadline was replenished")
+	}
+	req, _ := http.NewRequestWithContext(WithModel(ctx, "model"), http.MethodGet, server.URL, nil)
+	started := time.Now()
+	_, err = (&http.Client{Transport: Transport{}}).Do(req)
+	if !errors.Is(err, context.DeadlineExceeded) || time.Since(started) > 2*time.Second {
+		t.Fatalf("blocked send did not cancel at deadline: %v", err)
+	}
+	select {
+	case <-arrived:
+	default:
+		t.Fatal("request never reached local server")
+	}
+	if sends.Load() != 1 || guard.Reservations() != 1 {
+		t.Fatal("blocked send was not counted")
+	}
+}
+
+func TestGraphExecutionLastAllowedResponseRemainsUsable(t *testing.T) {
+	ctx, guard, cancel, err := BindGraphExecution(context.Background(), "spent-receipt", "model")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cancel()
+	for i := 0; i < GraphMaxSends; i++ {
+		if err := guard.Reserve(ctx, "model", GraphPhase); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := guard.Check(ctx); err != nil {
+		t.Fatalf("third valid response rejected: %v", err)
+	}
+}