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

scripts/graph_extraction_probe/budget.go

@@ -0,0 +1,73 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"sync"
+
+	"github.com/Tencent/WeKnora/internal/models/requestbudget"
+)
+
+func probeSendLimit(requested *int, environment string, envSet bool) (int, error) {
+	limit := requestbudget.GraphMaxSends
+	if requested != nil {
+		limit = *requested
+	}
+	if limit < 1 || limit > requestbudget.GraphMaxSends {
+		return 0, fmt.Errorf("invalid probe HTTP send cap")
+	}
+	if envSet && environment != strconv.Itoa(limit) {
+		return 0, fmt.Errorf("probe HTTP send cap differs between request and environment")
+	}
+	return limit, nil
+}
+
+// This tighter probe allowance delegates model, phase and deadline enforcement
+// to the existing graph guard; worker construction and defaults are untouched.
+type probeExecutionBudget struct {
+	*requestbudget.GraphExecution
+	mu      sync.Mutex
+	limit   int
+	stopped error
+	cancel  context.CancelCauseFunc
+}
+
+func bindProbeExecution(ctx context.Context, key, model string, limit int) (context.Context, *probeExecutionBudget, context.CancelFunc, error) {
+	if _, err := probeSendLimit(&limit, "", false); err != nil {
+		return ctx, nil, func() {}, err
+	}
+	ctx, base, baseCancel, err := requestbudget.BindGraphExecution(ctx, key, model)
+	if err != nil {
+		return ctx, nil, baseCancel, err
+	}
+	ctx, cancel := context.WithCancelCause(ctx)
+	guard := &probeExecutionBudget{GraphExecution: base, limit: limit, cancel: cancel}
+	return requestbudget.WithGuard(ctx, guard), guard, func() { cancel(nil); baseCancel() }, nil
+}
+
+func (g *probeExecutionBudget) Check(ctx context.Context) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if g.stopped != nil {
+		return g.stopped
+	}
+	return g.GraphExecution.Check(ctx)
+}
+
+func (g *probeExecutionBudget) Reserve(ctx context.Context, model, phase string) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if g.stopped != nil {
+		return g.stopped
+	}
+	if err := g.GraphExecution.Check(ctx); err != nil {
+		return err
+	}
+	if g.Reservations() >= g.limit {
+		g.stopped = fmt.Errorf("%w: probe HTTP send cap reached", requestbudget.ErrStopped)
+		g.cancel(g.stopped)
+		return g.stopped
+	}
+	return g.GraphExecution.Reserve(ctx, model, phase)
+}