promptkit/internal/observation/recorder_test.go

package observation

import (
	"context"
	"errors"
	"reflect"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
)

func TestExecutionBoundaryPreservesCallerContext(t *testing.T) {
	type callerKey struct{}
	parent, stop := context.WithTimeout(context.WithValue(context.Background(), callerKey{}, "caller"), time.Minute)
	defer stop()
	parent, cancel := context.WithCancelCause(parent)
	defer cancel(nil)
	recorder := newRecorder("outer", time.Now, nil)
	sink := recorder
	parent = WithClaimSink(WithRecorder(parent, recorder), &sink)
	child := WithoutExecutionState(parent)
	if FromContext(child) != nil || FromContext(parent) != recorder {
		t.Fatal("execution boundary did not isolate recorder state")
	}
	RegisterClaim(child, nil)
	if sink != recorder {
		t.Fatal("execution boundary retained the parent claim sink")
	}
	wantDeadline, _ := parent.Deadline()
	gotDeadline, ok := child.Deadline()
	if !ok || gotDeadline != wantDeadline || child.Value(callerKey{}) != "caller" || child.Done() != parent.Done() {
		t.Fatal("execution boundary changed caller context")
	}
	cause := errors.New("caller canceled")
	cancel(cause)
	if !errors.Is(child.Err(), context.Canceled) || context.Cause(child) != cause {
		t.Fatal("execution boundary lost cancellation or its cause")
	}
}

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()
	version := "v1"
	recorder.UpdateSelection(Selection{PromptID: "prompt", PromptVersion: &version, 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()
	}()
}