Add invocation-local execution accounting
This commit is contained in:
285
internal/observation/recorder.go
Normal file
285
internal/observation/recorder.go
Normal file
@@ -0,0 +1,285 @@
|
||||
// Package observation owns invocation-local execution identities, metadata
|
||||
// accounting, and synchronous event delivery.
|
||||
package observation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"fmt"
|
||||
"math"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
||||
)
|
||||
|
||||
// Recorder accumulates one execution's measured attempts and usage. Its owner
|
||||
// calls methods in execution order; it has no shared history or background work.
|
||||
type Recorder struct {
|
||||
runID string
|
||||
startedAt time.Time
|
||||
now func() time.Time
|
||||
observer domain.ExecutionObserver
|
||||
selection Selection
|
||||
attempts uint64
|
||||
retries uint64
|
||||
repairCalls uint64
|
||||
lastRepairGeneration uint64
|
||||
finished bool
|
||||
capacityWait time.Duration
|
||||
generationDuration time.Duration
|
||||
retryDelay time.Duration
|
||||
usage usageTotals
|
||||
}
|
||||
|
||||
// Selection contains resolved identities only. Unresolved fields stay empty.
|
||||
type Selection struct {
|
||||
PromptID string
|
||||
PromptVersion string
|
||||
ProfileID string
|
||||
BackendID string
|
||||
Model string
|
||||
}
|
||||
|
||||
// AttemptData describes one completed active-permit acquisition and optional
|
||||
// generation dispatch. Non-dispatched attempts contribute no token usage.
|
||||
type AttemptData struct {
|
||||
Generation uint64
|
||||
Attempt uint64
|
||||
Dispatched bool
|
||||
CapacityWait time.Duration
|
||||
GenerationDuration time.Duration
|
||||
Outcome domain.ExecutionOutcome
|
||||
Failure domain.ExecutionFailureKind
|
||||
StatusCode int
|
||||
Usage domain.TokenUsage
|
||||
UsagePresence *domain.TokenUsagePresence
|
||||
}
|
||||
|
||||
// DelayData describes one completed or interrupted retry wait.
|
||||
type DelayData struct {
|
||||
Generation uint64
|
||||
Attempt uint64
|
||||
Elapsed time.Duration
|
||||
Failure domain.ExecutionFailureKind
|
||||
}
|
||||
|
||||
// FinishData describes the terminal outcome after resource release.
|
||||
type FinishData struct {
|
||||
Outcome domain.ExecutionOutcome
|
||||
Failure domain.ExecutionFailureKind
|
||||
StatusCode int
|
||||
RetriesExhausted bool
|
||||
}
|
||||
|
||||
// New creates an invocation-local recorder and run ID. It returns the recorder
|
||||
// even if ID creation fails, allowing an error terminal event with an empty ID.
|
||||
func New(observer domain.ExecutionObserver) (*Recorder, error) {
|
||||
runID, err := NewRunID()
|
||||
return newRecorder(runID, time.Now, observer), err
|
||||
}
|
||||
|
||||
func newRecorder(runID string, now func() time.Time, observer domain.ExecutionObserver) *Recorder {
|
||||
return &Recorder{runID: runID, startedAt: now(), now: now, observer: observer,
|
||||
usage: newUsageTotals()}
|
||||
}
|
||||
|
||||
// NewRunID creates the identity shared by observations and successful results.
|
||||
func NewRunID() (string, error) {
|
||||
var b [16]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
return "", err
|
||||
}
|
||||
b[6] = (b[6] & 0x0f) | 0x40
|
||||
b[8] = (b[8] & 0x3f) | 0x80
|
||||
return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16]), nil
|
||||
}
|
||||
|
||||
// RunID returns this invocation's identity.
|
||||
func (r *Recorder) RunID() string { return r.runID }
|
||||
|
||||
// Start delivers the opening event with only identities already resolved.
|
||||
func (r *Recorder) Start() {
|
||||
event := r.event(domain.ExecutionStarted)
|
||||
r.deliver(event)
|
||||
}
|
||||
|
||||
// UpdateSelection replaces the known resolved identities for later events.
|
||||
func (r *Recorder) UpdateSelection(selection Selection) { r.selection = selection }
|
||||
|
||||
// Attempt records one completed acquisition and optional dispatch.
|
||||
func (r *Recorder) Attempt(data AttemptData) {
|
||||
r.capacityWait = addDuration(r.capacityWait, data.CapacityWait)
|
||||
if data.Dispatched {
|
||||
r.attempts++
|
||||
if data.Attempt > 1 {
|
||||
r.retries++
|
||||
}
|
||||
if data.Generation > 1 && data.Generation != r.lastRepairGeneration {
|
||||
r.repairCalls++
|
||||
r.lastRepairGeneration = data.Generation
|
||||
}
|
||||
r.generationDuration = addDuration(r.generationDuration, data.GenerationDuration)
|
||||
r.usage.add(data.Usage, data.UsagePresence)
|
||||
}
|
||||
event := r.event(domain.ExecutionAttempt)
|
||||
event.Generation = data.Generation
|
||||
event.Attempt = data.Attempt
|
||||
event.Dispatched = data.Dispatched
|
||||
event.CapacityWait = data.CapacityWait
|
||||
if data.Dispatched {
|
||||
event.GenerationDuration = data.GenerationDuration
|
||||
}
|
||||
event.Outcome = data.Outcome
|
||||
event.Failure = data.Failure
|
||||
event.StatusCode = data.StatusCode
|
||||
if data.Dispatched {
|
||||
event.Usage = observedUsage(data.Usage, data.UsagePresence)
|
||||
}
|
||||
r.deliver(event)
|
||||
}
|
||||
|
||||
// Delay records the measured time spent waiting after one failed attempt.
|
||||
func (r *Recorder) Delay(data DelayData) {
|
||||
r.retryDelay = addDuration(r.retryDelay, data.Elapsed)
|
||||
event := r.event(domain.ExecutionRetryDelay)
|
||||
event.Generation = data.Generation
|
||||
event.Attempt = data.Attempt
|
||||
event.RetryDelay = data.Elapsed
|
||||
event.Failure = data.Failure
|
||||
r.deliver(event)
|
||||
}
|
||||
|
||||
// Finish delivers terminal totals once on a normal return. Its caller owns the
|
||||
// operation's resource cleanup and failure classification.
|
||||
func (r *Recorder) Finish(data FinishData) {
|
||||
if r.finished {
|
||||
return
|
||||
}
|
||||
r.finished = true
|
||||
event := r.event(domain.ExecutionFinished)
|
||||
event.Outcome = data.Outcome
|
||||
event.Failure = data.Failure
|
||||
event.StatusCode = data.StatusCode
|
||||
event.RetriesExhausted = data.RetriesExhausted
|
||||
event.Attempts = r.attempts
|
||||
event.TransientRetries = r.retries
|
||||
event.RepairCalls = r.repairCalls
|
||||
event.CapacityWait = r.capacityWait
|
||||
event.GenerationDuration = r.generationDuration
|
||||
event.RetryDelay = r.retryDelay
|
||||
event.Duration = event.Time.Sub(r.startedAt)
|
||||
if event.Duration < 0 {
|
||||
event.Duration = 0
|
||||
}
|
||||
event.Usage = r.usage.snapshot(r.attempts)
|
||||
r.deliver(event)
|
||||
}
|
||||
|
||||
func (r *Recorder) event(kind domain.ExecutionEventKind) domain.ExecutionEvent {
|
||||
return domain.ExecutionEvent{
|
||||
Kind: kind, RunID: r.runID, Time: r.now(), PromptID: r.selection.PromptID,
|
||||
PromptVersion: r.selection.PromptVersion, ProfileID: r.selection.ProfileID,
|
||||
BackendID: r.selection.BackendID, Model: r.selection.Model,
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Recorder) deliver(event domain.ExecutionEvent) {
|
||||
if r.observer != nil {
|
||||
r.observer(event)
|
||||
}
|
||||
}
|
||||
|
||||
type contextKey struct{}
|
||||
|
||||
// WithRecorder carries one invocation's recorder through internal client calls.
|
||||
func WithRecorder(ctx context.Context, recorder *Recorder) context.Context {
|
||||
return context.WithValue(ctx, contextKey{}, recorder)
|
||||
}
|
||||
|
||||
// FromContext returns the invocation-local recorder, if one was installed.
|
||||
func FromContext(ctx context.Context) *Recorder {
|
||||
recorder, _ := ctx.Value(contextKey{}).(*Recorder)
|
||||
return recorder
|
||||
}
|
||||
|
||||
type usageField struct {
|
||||
value int
|
||||
known, complete bool
|
||||
}
|
||||
type usageTotals struct {
|
||||
prompt, completion, total, cached, cacheWrite usageField
|
||||
}
|
||||
|
||||
func newUsageTotals() usageTotals {
|
||||
return usageTotals{prompt: usageField{complete: true}, completion: usageField{complete: true},
|
||||
total: usageField{complete: true}, cached: usageField{complete: true}, cacheWrite: usageField{complete: true}}
|
||||
}
|
||||
|
||||
func (u *usageTotals) add(usage domain.TokenUsage, presence *domain.TokenUsagePresence) {
|
||||
known := usagePresence(usage, presence)
|
||||
u.prompt.add(usage.PromptTokens, known.PromptTokens)
|
||||
u.completion.add(usage.CompletionTokens, known.CompletionTokens)
|
||||
u.total.add(usage.TotalTokens, known.TotalTokens)
|
||||
u.cached.add(usage.CachedTokens, known.CachedTokens)
|
||||
u.cacheWrite.add(usage.CacheWriteTokens, known.CacheWriteTokens)
|
||||
}
|
||||
|
||||
func (u usageTotals) snapshot(attempts uint64) domain.ObservedTokenUsage {
|
||||
return domain.ObservedTokenUsage{
|
||||
PromptTokens: u.prompt.snapshot(attempts), CompletionTokens: u.completion.snapshot(attempts),
|
||||
TotalTokens: u.total.snapshot(attempts), CachedTokens: u.cached.snapshot(attempts),
|
||||
CacheWriteTokens: u.cacheWrite.snapshot(attempts),
|
||||
}
|
||||
}
|
||||
|
||||
func (f *usageField) add(value int, known bool) {
|
||||
if known {
|
||||
f.value += value
|
||||
f.known = true
|
||||
} else {
|
||||
f.complete = false
|
||||
}
|
||||
}
|
||||
|
||||
func (f usageField) snapshot(attempts uint64) domain.ObservedTokenCount {
|
||||
if attempts == 0 {
|
||||
return domain.ObservedTokenCount{Known: true, Complete: true}
|
||||
}
|
||||
return domain.ObservedTokenCount{Value: f.value, Known: f.known, Complete: f.complete}
|
||||
}
|
||||
|
||||
func observedUsage(usage domain.TokenUsage, presence *domain.TokenUsagePresence) domain.ObservedTokenUsage {
|
||||
known := usagePresence(usage, presence)
|
||||
return domain.ObservedTokenUsage{
|
||||
PromptTokens: observedCount(usage.PromptTokens, known.PromptTokens),
|
||||
CompletionTokens: observedCount(usage.CompletionTokens, known.CompletionTokens),
|
||||
TotalTokens: observedCount(usage.TotalTokens, known.TotalTokens),
|
||||
CachedTokens: observedCount(usage.CachedTokens, known.CachedTokens),
|
||||
CacheWriteTokens: observedCount(usage.CacheWriteTokens, known.CacheWriteTokens),
|
||||
}
|
||||
}
|
||||
|
||||
func usagePresence(usage domain.TokenUsage, presence *domain.TokenUsagePresence) domain.TokenUsagePresence {
|
||||
if presence != nil {
|
||||
return *presence
|
||||
}
|
||||
return domain.TokenUsagePresence{
|
||||
PromptTokens: usage.PromptTokens != 0, CompletionTokens: usage.CompletionTokens != 0,
|
||||
TotalTokens: usage.TotalTokens != 0, CachedTokens: usage.CachedTokens != 0,
|
||||
CacheWriteTokens: usage.CacheWriteTokens != 0,
|
||||
}
|
||||
}
|
||||
|
||||
func observedCount(value int, known bool) domain.ObservedTokenCount {
|
||||
if !known {
|
||||
return domain.ObservedTokenCount{}
|
||||
}
|
||||
return domain.ObservedTokenCount{Value: value, Known: true, Complete: true}
|
||||
}
|
||||
|
||||
func addDuration(total, next time.Duration) time.Duration {
|
||||
if next > 0 && total > time.Duration(math.MaxInt64)-next {
|
||||
return time.Duration(math.MaxInt64)
|
||||
}
|
||||
return total + next
|
||||
}
|
||||
129
internal/observation/recorder_test.go
Normal file
129
internal/observation/recorder_test.go
Normal file
@@ -0,0 +1,129 @@
|
||||
package observation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
||||
)
|
||||
|
||||
func TestRecorderAccountsDispatchedAttemptsAndReportedUsage(t *testing.T) {
|
||||
start := time.Date(2026, 9, 23, 12, 0, 0, 0, time.UTC)
|
||||
tick := 0
|
||||
now := func() time.Time {
|
||||
value := start.Add(time.Duration(tick) * time.Second)
|
||||
tick++
|
||||
return value
|
||||
}
|
||||
var events []domain.ExecutionEvent
|
||||
recorder := newRecorder("run-1", now, func(event domain.ExecutionEvent) { events = append(events, event) })
|
||||
recorder.Start()
|
||||
recorder.UpdateSelection(Selection{PromptID: "prompt", PromptVersion: "v1", ProfileID: "profile", BackendID: "backend", Model: "model"})
|
||||
recorder.Attempt(AttemptData{
|
||||
Generation: 1, Attempt: 1, Dispatched: true,
|
||||
CapacityWait: 2 * time.Millisecond, GenerationDuration: 10 * time.Millisecond,
|
||||
Outcome: domain.ExecutionOutcomeSuccess,
|
||||
UsagePresence: &domain.TokenUsagePresence{PromptTokens: true},
|
||||
})
|
||||
recorder.Delay(DelayData{Generation: 1, Attempt: 1, Elapsed: 7 * time.Millisecond})
|
||||
recorder.Attempt(AttemptData{
|
||||
Generation: 1, Attempt: 2, Dispatched: true,
|
||||
CapacityWait: 3 * time.Millisecond, GenerationDuration: 20 * time.Millisecond,
|
||||
Outcome: domain.ExecutionOutcomeSuccess,
|
||||
Usage: domain.TokenUsage{PromptTokens: 2, TotalTokens: 5},
|
||||
UsagePresence: &domain.TokenUsagePresence{PromptTokens: true, TotalTokens: true},
|
||||
})
|
||||
recorder.Attempt(AttemptData{
|
||||
Generation: 2, Attempt: 1, Dispatched: false,
|
||||
CapacityWait: 4 * time.Millisecond, GenerationDuration: 99 * time.Millisecond,
|
||||
Outcome: domain.ExecutionOutcomeError, Failure: domain.ExecutionFailureCanceled,
|
||||
Usage: domain.TokenUsage{PromptTokens: 99},
|
||||
})
|
||||
recorder.Attempt(AttemptData{
|
||||
Generation: 2, Attempt: 1, Dispatched: true,
|
||||
CapacityWait: 5 * time.Millisecond, GenerationDuration: 30 * time.Millisecond,
|
||||
Outcome: domain.ExecutionOutcomeError, Failure: domain.ExecutionFailureProviderStatus,
|
||||
StatusCode: 503, Usage: domain.TokenUsage{PromptTokens: 3},
|
||||
})
|
||||
recorder.Finish(FinishData{Outcome: domain.ExecutionOutcomeError, Failure: domain.ExecutionFailureProviderStatus, StatusCode: 503, RetriesExhausted: true})
|
||||
recorder.Finish(FinishData{})
|
||||
|
||||
if len(events) != 7 {
|
||||
t.Fatalf("event count=%d", len(events))
|
||||
}
|
||||
if events[0].PromptID != "" || events[1].PromptID != "prompt" || events[0].RunID != "run-1" || events[1].RunID != "run-1" {
|
||||
t.Fatalf("identity snapshots=%+v, %+v", events[0], events[1])
|
||||
}
|
||||
if events[1].Usage.PromptTokens != (domain.ObservedTokenCount{Known: true, Complete: true}) ||
|
||||
events[4].Dispatched || events[4].GenerationDuration != 0 || events[4].Usage.PromptTokens.Known {
|
||||
t.Fatalf("attempt usage or undispatched event=%+v, %+v", events[1], events[4])
|
||||
}
|
||||
last := events[len(events)-1]
|
||||
if last.Attempts != 3 || last.TransientRetries != 1 || last.RepairCalls != 1 ||
|
||||
last.CapacityWait != 14*time.Millisecond || last.GenerationDuration != 60*time.Millisecond ||
|
||||
last.RetryDelay != 7*time.Millisecond || last.Duration != 7*time.Second {
|
||||
t.Fatalf("terminal counters=%+v", last)
|
||||
}
|
||||
if last.Usage.PromptTokens != (domain.ObservedTokenCount{Value: 5, Known: true, Complete: true}) ||
|
||||
last.Usage.TotalTokens != (domain.ObservedTokenCount{Value: 5, Known: true, Complete: false}) ||
|
||||
last.Usage.CachedTokens != (domain.ObservedTokenCount{}) {
|
||||
t.Fatalf("aggregate usage=%+v", last.Usage)
|
||||
}
|
||||
if last.Failure != domain.ExecutionFailureProviderStatus || !last.RetriesExhausted || last.StatusCode != 503 {
|
||||
t.Fatalf("terminal failure=%+v", last)
|
||||
}
|
||||
if events[1].Attempts != 0 || events[1].TransientRetries != 0 || events[1].RepairCalls != 0 {
|
||||
t.Fatalf("attempt event carried terminal totals=%+v", events[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecorderZeroAttemptsAndExplicitUnknownUsage(t *testing.T) {
|
||||
var events []domain.ExecutionEvent
|
||||
recorder := newRecorder("", time.Now, func(event domain.ExecutionEvent) { events = append(events, event) })
|
||||
recorder.Finish(FinishData{Outcome: domain.ExecutionOutcomeError, Failure: domain.ExecutionFailureOther})
|
||||
if len(events) != 1 || events[0].RunID != "" || events[0].Attempts != 0 ||
|
||||
events[0].Usage.PromptTokens != (domain.ObservedTokenCount{Known: true, Complete: true}) ||
|
||||
events[0].Usage.CacheWriteTokens != (domain.ObservedTokenCount{Known: true, Complete: true}) {
|
||||
t.Fatalf("zero-attempt terminal=%+v", events)
|
||||
}
|
||||
|
||||
recorder = newRecorder("run-2", time.Now, func(event domain.ExecutionEvent) { events = append(events, event) })
|
||||
recorder.Attempt(AttemptData{Dispatched: true, Generation: 1, Attempt: 1,
|
||||
Usage: domain.TokenUsage{PromptTokens: 7}, UsagePresence: &domain.TokenUsagePresence{}})
|
||||
recorder.Finish(FinishData{Outcome: domain.ExecutionOutcomeSuccess})
|
||||
last := events[len(events)-1]
|
||||
if last.Usage.PromptTokens != (domain.ObservedTokenCount{}) {
|
||||
t.Fatalf("explicitly unknown nonzero field=%+v", last.Usage.PromptTokens)
|
||||
}
|
||||
_ = newRecorder("run-3", time.Now, nil) // A nil callback still permits recording.
|
||||
}
|
||||
|
||||
func TestRecorderOrderContextAndCallbackPanic(t *testing.T) {
|
||||
var kinds []domain.ExecutionEventKind
|
||||
recorder := newRecorder("run", time.Now, func(event domain.ExecutionEvent) { kinds = append(kinds, event.Kind) })
|
||||
ctx := WithRecorder(context.Background(), recorder)
|
||||
if FromContext(ctx) != recorder || FromContext(context.Background()) != nil {
|
||||
t.Fatal("recorder context propagation failed")
|
||||
}
|
||||
recorder.Start()
|
||||
recorder.Attempt(AttemptData{Generation: 1, Attempt: 1})
|
||||
recorder.Delay(DelayData{Generation: 1, Attempt: 1})
|
||||
recorder.Finish(FinishData{Outcome: domain.ExecutionOutcomeSuccess})
|
||||
if !reflect.DeepEqual(kinds, []domain.ExecutionEventKind{
|
||||
domain.ExecutionStarted, domain.ExecutionAttempt, domain.ExecutionRetryDelay, domain.ExecutionFinished,
|
||||
}) {
|
||||
t.Fatalf("callback order=%v", kinds)
|
||||
}
|
||||
|
||||
recorder = newRecorder("panic", time.Now, func(domain.ExecutionEvent) { panic("observer panic") })
|
||||
func() {
|
||||
defer func() {
|
||||
if recover() == nil {
|
||||
t.Fatal("observer panic was swallowed")
|
||||
}
|
||||
}()
|
||||
recorder.Start()
|
||||
}()
|
||||
}
|
||||
Reference in New Issue
Block a user