promptkit/internal/observation/recorder_test.go

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()
	}()
}