scriptorium/internal/report/report_test.go

package report

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

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

func TestCollectorProjectsOnlyTerminalAggregates(t *testing.T) {
	collector := NewCollector()
	if collector.Snapshot() != nil {
		t.Fatal("empty collector produced a report")
	}
	collector.Observe(promptkit.ExecutionEvent{Kind: promptkit.ExecutionAttempt, Attempts: 99})
	if collector.Snapshot() != nil {
		t.Fatal("nonterminal event produced a report")
	}
	collector.Observe(promptkit.ExecutionEvent{
		Kind: promptkit.ExecutionFinished, RunID: "run-1", PromptID: "prompt-1",
		PromptVersion: "1", ProfileID: "profile-1", BackendID: "backend-1", Model: "model-1",
		Outcome: promptkit.ExecutionOutcomeError, Failure: promptkit.ExecutionFailureProviderStatus,
		StatusCode: 503, RetriesExhausted: true, Attempts: 3, TransientRetries: 2, RepairCalls: 1,
		Duration: 1500 * time.Millisecond, CapacityWait: 20 * time.Millisecond,
		GenerationDuration: 300 * time.Millisecond, RetryDelay: 700 * time.Millisecond,
		Usage: promptkit.ObservedTokenUsage{
			PromptTokens:     promptkit.ObservedTokenCount{Value: 0, Known: true, Complete: true},
			CompletionTokens: promptkit.ObservedTokenCount{Value: 0, Known: false, Complete: false},
			TotalTokens:      promptkit.ObservedTokenCount{Value: 11, Known: true, Complete: false},
		},
	})
	collector.Observe(promptkit.ExecutionEvent{Kind: promptkit.ExecutionRetryDelay, Attempts: 50})
	report := collector.Snapshot()
	if report == nil || report.Attempts != 3 || report.TransientRetries != 2 || report.RepairCalls != 1 || report.StatusCode != 503 || !report.RetriesExhausted {
		t.Fatalf("terminal aggregates changed: %+v", report)
	}
	if report.DurationMS != 1500 || report.CapacityWaitMS != 20 || report.GenerationDurationMS != 300 || report.RetryDelayMS != 700 {
		t.Fatalf("terminal timings changed: %+v", report)
	}
	if report.Usage.PromptTokens != (TokenCount{Known: true, Complete: true}) || report.Usage.CompletionTokens != (TokenCount{}) || report.Usage.TotalTokens != (TokenCount{Value: 11, Known: true}) {
		t.Fatalf("presence-aware usage changed: %+v", report.Usage)
	}
	encoded, err := json.Marshal(report)
	if err != nil {
		t.Fatal(err)
	}
	var fields map[string]any
	if err := json.Unmarshal(encoded, &fields); err != nil {
		t.Fatal(err)
	}
	for _, key := range []string{"run_id", "prompt_id", "prompt_version", "profile_id", "backend_id", "model", "outcome", "failure", "status_code", "retries_exhausted", "attempts", "transient_retries", "repair_calls", "duration_ms", "capacity_wait_ms", "generation_duration_ms", "retry_delay_ms", "usage"} {
		if _, ok := fields[key]; !ok {
			t.Fatalf("report missing %q: %s", key, encoded)
		}
	}
	for _, key := range []string{"prompt_tokens", "completion_tokens", "total_tokens", "cached_tokens", "cache_write_tokens"} {
		item := fields["usage"].(map[string]any)[key].(map[string]any)
		if len(item) != 3 {
			t.Fatalf("incomplete token count %q: %#v", key, item)
		}
	}
	for _, forbidden := range []string{"session_id", "headers", "endpoint", "api_key", "error", "output", "input"} {
		if _, present := fields[forbidden]; present {
			t.Fatalf("report exposed %q: %s", forbidden, encoded)
		}
	}
}

func TestCollectorHandlesUnknownValuesAndKeepsInvocationsSeparate(t *testing.T) {
	first := NewCollector()
	second := NewCollector()
	first.Observe(promptkit.ExecutionEvent{Kind: promptkit.ExecutionFinished, RunID: "first", Outcome: "future", Failure: "future"})
	second.Observe(promptkit.ExecutionEvent{Kind: promptkit.ExecutionFinished, RunID: "second", Outcome: promptkit.ExecutionOutcomeSuccess})
	one := first.Snapshot()
	two := second.Snapshot()
	if one.RunID != "first" || one.Outcome != "other" || one.Failure != "other" || two.RunID != "second" || two.Failure != "none" {
		t.Fatalf("collector isolation or classification failed: first=%+v second=%+v", one, two)
	}
	one.RunID = "changed"
	if first.Snapshot().RunID != "first" {
		t.Fatal("snapshot mutation changed stored report")
	}
	encoded, err := json.Marshal(two)
	if err != nil {
		t.Fatal(err)
	}
	for _, absent := range []string{"prompt_id", "prompt_version", "profile_id", "backend_id", "model", "status_code"} {
		if strings.Contains(string(encoded), `"`+absent+`"`) {
			t.Fatalf("unavailable optional field %q was included: %s", absent, encoded)
		}
	}
}

type localClient struct{}

func (localClient) Generate(ctx context.Context, req promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
	return &promptkit.GenerateResponse{Content: "ok"}, nil
}

func TestCollectorReceivesRealEngineTerminalEvent(t *testing.T) {
	promptDir := t.TempDir()
	profileDir := t.TempDir()
	if err := os.WriteFile(filepath.Join(promptDir, "prompt.yaml"), []byte("id: p\nversion: \"1\"\ndefault_profile: local\nmessages:\n  - role: user\n    content: hi\noutput:\n  format: text\n  validation_mode: none\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(profileDir, "profile.yaml"), []byte("id: local\nendpoint: http://127.0.0.1:1/v1\nmodel: local-model\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	engine, err := promptkit.NewEngine(promptkit.Config{PromptDir: promptDir, ProfileDir: profileDir}, promptkit.WithLLMClient(localClient{}))
	if err != nil {
		t.Fatal(err)
	}
	collector := NewCollector()
	if _, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "p", Observer: collector.Observe}); err != nil {
		t.Fatal(err)
	}
	value := collector.Snapshot()
	if value == nil || value.PromptID != "p" || value.ProfileID != "local" || value.Outcome != "success" || value.Attempts != 1 {
		t.Fatalf("unexpected terminal event projection: %+v", value)
	}
}