xinmi/internal/models/requestbudget/graph.go

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
}