345 lines
11 KiB
Go
345 lines
11 KiB
Go
// Package observation owns invocation-local execution identities, metadata
|
|
// accounting, and synchronous event delivery.
|
|
package observation
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"fmt"
|
|
"math"
|
|
"time"
|
|
|
|
"go.maximumdirect.net/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
|
|
parentRunID 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 UsageAccumulator
|
|
}
|
|
|
|
// Selection contains resolved identities only. Unresolved fields stay empty.
|
|
type Selection struct {
|
|
PromptID 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}
|
|
}
|
|
|
|
// 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 }
|
|
|
|
// SetParentRunID records the immediate source identity before start delivery.
|
|
func (r *Recorder) SetParentRunID(id string) { r.parentRunID = id }
|
|
|
|
// 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.Observed()
|
|
r.deliver(event)
|
|
}
|
|
|
|
func (r *Recorder) event(kind domain.ExecutionEventKind) domain.ExecutionEvent {
|
|
return domain.ExecutionEvent{
|
|
Kind: kind, RunID: r.runID, ParentRunID: r.parentRunID, Time: r.now(), PromptID: r.selection.PromptID,
|
|
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{}
|
|
|
|
// WithoutExecutionState starts a public operation without inheriting another
|
|
// invocation's recorder or prepared-claim sink. Other context values and all
|
|
// cancellation/deadline behavior remain with the caller's context.
|
|
func WithoutExecutionState(ctx context.Context) context.Context {
|
|
return executionBoundary{Context: ctx}
|
|
}
|
|
|
|
type executionBoundary struct{ context.Context }
|
|
|
|
func (c executionBoundary) Value(key any) any {
|
|
switch key.(type) {
|
|
case contextKey, claimSinkKey:
|
|
return nil
|
|
default:
|
|
return c.Context.Value(key)
|
|
}
|
|
}
|
|
|
|
// 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 claimSinkKey struct{}
|
|
|
|
// WithClaimSink lets the facade learn which invocation actually claimed a
|
|
// prepared handle, so it can emit the terminal event after public mapping.
|
|
func WithClaimSink(ctx context.Context, sink **Recorder) context.Context {
|
|
return context.WithValue(ctx, claimSinkKey{}, sink)
|
|
}
|
|
|
|
// RegisterClaim records an invocation after its prepared handle is claimed.
|
|
func RegisterClaim(ctx context.Context, recorder *Recorder) {
|
|
if sink, ok := ctx.Value(claimSinkKey{}).(**Recorder); ok && sink != nil {
|
|
*sink = recorder
|
|
}
|
|
}
|
|
|
|
type usageField struct {
|
|
value int
|
|
known, complete bool
|
|
}
|
|
type usageTotals struct {
|
|
prompt, completion, total, cached, cacheWrite usageField
|
|
}
|
|
|
|
// UsageAccumulator combines presence-aware usage across dispatched attempts.
|
|
// Its zero value represents no dispatched attempts.
|
|
type UsageAccumulator struct {
|
|
totals usageTotals
|
|
attempts uint64
|
|
}
|
|
|
|
// Add includes the usage reported by one dispatched attempt.
|
|
func (u *UsageAccumulator) Add(usage domain.TokenUsage, presence *domain.TokenUsagePresence) {
|
|
if u.attempts == 0 {
|
|
u.totals = completeUsageTotals()
|
|
}
|
|
u.attempts++
|
|
u.totals.add(usage, presence)
|
|
}
|
|
|
|
// Observed returns aggregate counts, marking every field known and complete
|
|
// when no attempt has been added.
|
|
func (u UsageAccumulator) Observed() domain.ObservedTokenUsage {
|
|
return u.totals.snapshot(u.attempts)
|
|
}
|
|
|
|
func completeUsageTotals() 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
|
|
}
|