feat(aiplugin): add TransitionBudget for runaway prevention

Directly addresses the qmuntal/stateless #77 pattern: workflow
execution that needs to halt after N transitions to prevent runaway
loops. The author there had to fight the library — statekit users
can drop in TransitionBudget instead.

Behavior:
- AfterTransition counts only state-changing transitions (XState
  semantics — self-transitions don't tick the counter).
- OnEvent rewrites the event type to a configurable halt-event
  once the budget is exhausted. Payload is preserved so the halt
  state can read why it fired.
- atomic counter — safe under concurrent OnEvent + AfterTransition.

Tests cover budget tracking, self-transition skipping, event rewriting
under exhaustion, Reset, concurrent safety, plugin-interface
satisfaction, and an end-to-end runaway-machine scenario that
confirms the budget halts a real loop within the configured budget.

91.8% coverage.

Source: ICP signal sweep — verbatim quote from qmuntal/stateless #77.

Author: felixgeelhaar

aiplugin/budget.go

@@ -0,0 +1,98 @@
+package aiplugin
+
+import (
+	"sync/atomic"
+
+	"github.com/felixgeelhaar/statekit/plugin"
+)
+
+// TransitionBudget caps the number of transitions an interpreter is
+// allowed to perform. When the budget is exhausted, subsequent events
+// are rewritten to a configurable halt-event type, signalling
+// runaway prevention without forcing the machine into a bespoke
+// error state.
+//
+// Why this exists: a recurring pain in incumbent Go FSM libraries
+// (e.g. https://github.com/qmuntal/stateless/issues/77) is the
+// inability to halt execution after N transitions. Wire a
+// TransitionBudget into the plugin slot and your machine gets a
+// bounded-blast-radius guarantee out of the box.
+//
+// Usage:
+//
+//	budget := aiplugin.NewTransitionBudget[Ctx](100, "BUDGET_EXHAUSTED")
+//	interp.Use(budget)
+//	// Add a transition for "BUDGET_EXHAUSTED" that routes to a
+//	// halted/failed state from anywhere it can fire.
+//
+// TransitionBudget tracks transitions via the AfterTransition hook,
+// not via OnEvent — this counts only events that actually changed
+// state, matching the qmuntal-issue's "transition count" semantics.
+// OnEvent rewrites the event type once the budget is exhausted so
+// the machine can route to its halt state on the same Send call.
+type TransitionBudget[C any] struct {
+	limit       int64
+	haltEvent   plugin.EventType
+	transitions atomic.Int64
+}
+
+// NewTransitionBudget constructs a TransitionBudget[C] capped at
+// `limit` transitions. Once the budget is reached, OnEvent rewrites
+// every subsequent event to `haltEvent`.
+func NewTransitionBudget[C any](limit int64, haltEvent plugin.EventType) *TransitionBudget[C] {
+	return &TransitionBudget[C]{
+		limit:     limit,
+		haltEvent: haltEvent,
+	}
+}
+
+// Name implements plugin.Plugin.
+func (*TransitionBudget[C]) Name() string { return "ai-transition-budget" }
+
+// AfterTransition implements plugin.OnTransitionHook. Every observed
+// transition increments the counter.
+func (b *TransitionBudget[C]) AfterTransition(_ plugin.Context[C], from, to plugin.StateID, _ plugin.Event) {
+	if from == to {
+		// Internal/no-op transitions don't count against the budget;
+		// XState semantics: only state-changing transitions tick.
+		return
+	}
+	b.transitions.Add(1)
+}
+
+// BeforeTransition is required to satisfy plugin.OnTransitionHook
+// (paired with AfterTransition). No-op.
+func (*TransitionBudget[C]) BeforeTransition(_ plugin.Context[C], _, _ plugin.StateID, _ plugin.Event) {
+}
+
+// OnEvent implements plugin.OnEventHook. Once the budget is
+// exhausted, the event type is rewritten to the halt event.
+func (b *TransitionBudget[C]) OnEvent(_ plugin.Context[C], event plugin.Event) plugin.Event {
+	if b.Exhausted() {
+		// Preserve payload — caller can read why the halt fired.
+		return plugin.Event{Type: b.haltEvent, Payload: event.Payload}
+	}
+	return event
+}
+
+// Used returns the running transition count.
+func (b *TransitionBudget[C]) Used() int64 { return b.transitions.Load() }
+
+// Remaining returns how many transitions are still allowed before
+// the halt event is forced. A negative value means the budget has
+// already been overshot (race conditions on counter increments are
+// possible but bounded by the number of in-flight events).
+func (b *TransitionBudget[C]) Remaining() int64 {
+	return b.limit - b.transitions.Load()
+}
+
+// Exhausted reports whether the budget has been reached.
+func (b *TransitionBudget[C]) Exhausted() bool {
+	return b.transitions.Load() >= b.limit
+}
+
+// Reset zeroes the transition counter. Useful between independent
+// runs of the same interpreter, or after a manual recovery flow.
+func (b *TransitionBudget[C]) Reset() {
+	b.transitions.Store(0)
+}

aiplugin/budget_test.go

@@ -0,0 +1,152 @@
+package aiplugin_test
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/felixgeelhaar/statekit"
+	"github.com/felixgeelhaar/statekit/aiplugin"
+	"github.com/felixgeelhaar/statekit/plugin"
+)
+
+type budgetCtx struct{}
+
+func TestTransitionBudget_CountsAndExhausts(t *testing.T) {
+	t.Parallel()
+	b := aiplugin.NewTransitionBudget[budgetCtx](2, "HALT")
+
+	// AfterTransition counts state-changing transitions.
+	b.AfterTransition(plugin.Context[budgetCtx]{}, "a", "b", plugin.Event{Type: "X"})
+	if b.Used() != 1 {
+		t.Errorf("Used = %d, want 1", b.Used())
+	}
+	b.AfterTransition(plugin.Context[budgetCtx]{}, "b", "c", plugin.Event{Type: "Y"})
+	if b.Used() != 2 {
+		t.Errorf("Used = %d, want 2", b.Used())
+	}
+	if !b.Exhausted() {
+		t.Error("expected Exhausted")
+	}
+	if b.Remaining() != 0 {
+		t.Errorf("Remaining = %d, want 0", b.Remaining())
+	}
+}
+
+func TestTransitionBudget_SkipsSelfTransition(t *testing.T) {
+	t.Parallel()
+	b := aiplugin.NewTransitionBudget[budgetCtx](100, "HALT")
+
+	// Self-transitions don't count.
+	b.AfterTransition(plugin.Context[budgetCtx]{}, "x", "x", plugin.Event{Type: "X"})
+	if b.Used() != 0 {
+		t.Errorf("self-transition counted: Used = %d", b.Used())
+	}
+}
+
+func TestTransitionBudget_RewritesEventWhenExhausted(t *testing.T) {
+	t.Parallel()
+	b := aiplugin.NewTransitionBudget[budgetCtx](1, "HALT")
+
+	// Below budget — pass-through.
+	out := b.OnEvent(plugin.Context[budgetCtx]{}, plugin.Event{Type: "GO", Payload: 42})
+	if out.Type != "GO" {
+		t.Errorf("below budget, event rewritten to %q", out.Type)
+	}
+
+	// Tick the budget.
+	b.AfterTransition(plugin.Context[budgetCtx]{}, "a", "b", plugin.Event{Type: "GO"})
+
+	// Above budget — rewritten.
+	out = b.OnEvent(plugin.Context[budgetCtx]{}, plugin.Event{Type: "STILL_GO", Payload: 99})
+	if out.Type != "HALT" {
+		t.Errorf("above budget, event = %q, want HALT", out.Type)
+	}
+	if out.Payload != 99 {
+		t.Errorf("payload not preserved: %v", out.Payload)
+	}
+}
+
+func TestTransitionBudget_Reset(t *testing.T) {
+	t.Parallel()
+	b := aiplugin.NewTransitionBudget[budgetCtx](2, "HALT")
+	b.AfterTransition(plugin.Context[budgetCtx]{}, "a", "b", plugin.Event{})
+	b.AfterTransition(plugin.Context[budgetCtx]{}, "b", "c", plugin.Event{})
+	if !b.Exhausted() {
+		t.Fatal("expected exhausted before Reset")
+	}
+	b.Reset()
+	if b.Exhausted() {
+		t.Error("expected not exhausted after Reset")
+	}
+	if b.Used() != 0 {
+		t.Errorf("Used after Reset = %d, want 0", b.Used())
+	}
+}
+
+func TestTransitionBudget_Concurrent(t *testing.T) {
+	t.Parallel()
+	b := aiplugin.NewTransitionBudget[budgetCtx](100_000, "HALT")
+	var wg sync.WaitGroup
+	const goroutines = 16
+	const perG = 1000
+	wg.Add(goroutines)
+	for i := 0; i < goroutines; i++ {
+		go func() {
+			defer wg.Done()
+			for j := 0; j < perG; j++ {
+				b.AfterTransition(plugin.Context[budgetCtx]{}, "a", "b", plugin.Event{})
+			}
+		}()
+	}
+	wg.Wait()
+	if got := b.Used(); got != goroutines*perG {
+		t.Errorf("Used = %d, want %d", got, goroutines*perG)
+	}
+}
+
+// TestTransitionBudget_PluginInterface verifies the type satisfies
+// the relevant plugin hook interfaces at compile time.
+func TestTransitionBudget_PluginInterface(t *testing.T) {
+	t.Parallel()
+	var _ plugin.OnTransitionHook[budgetCtx] = aiplugin.NewTransitionBudget[budgetCtx](1, "X")
+	var _ plugin.OnEventHook[budgetCtx] = aiplugin.NewTransitionBudget[budgetCtx](1, "X")
+	if aiplugin.NewTransitionBudget[budgetCtx](1, "X").Name() != "ai-transition-budget" {
+		t.Error("name mismatch")
+	}
+}
+
+// TestTransitionBudget_EndToEnd_HaltsRunawayMachine drops a budget
+// into a real interpreter and confirms it halts a runaway loop —
+// the pattern qmuntal/stateless #77 author wanted.
+func TestTransitionBudget_EndToEnd_HaltsRunawayMachine(t *testing.T) {
+	t.Parallel()
+
+	machine, err := statekit.NewMachine[budgetCtx]("loop").
+		WithInitial("a").
+		State("a").On("TICK").Target("b").On("HALT").Target("done").Done().
+		State("b").On("TICK").Target("a").On("HALT").Target("done").Done().
+		State("done").Final().Done().
+		Build()
+	if err != nil {
+		t.Fatalf("build: %v", err)
+	}
+
+	interp := statekit.NewInterpreter(machine)
+	defer func() { _ = interp.Close() }()
+
+	budget := aiplugin.NewTransitionBudget[budgetCtx](3, "HALT")
+	interp.Use(budget)
+	interp.Start()
+
+	// Send 10 TICKs — should halt by the 4th transition.
+	for i := 0; i < 10; i++ {
+		interp.Send(statekit.Event{Type: "TICK"})
+	}
+
+	if !interp.Done() {
+		t.Errorf("expected interpreter Done (in halt final state); state = %q", interp.State().Value)
+	}
+	if got := budget.Used(); got > 4 {
+		t.Errorf("budget overshot: Used = %d, expected <= 4", got)
+	}
+}