promptkit/internal/generation/observation_test.go

package generation

import (
	"context"
	"errors"
	"net/http"
	"testing"
	"time"

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

func recordEvents(t *testing.T, ctx context.Context, events *[]domain.ExecutionEvent) context.Context {
	t.Helper()
	recorder, err := observation.New(func(event domain.ExecutionEvent) { *events = append(*events, event) })
	if err != nil {
		t.Fatal(err)
	}
	return observation.WithRecorder(ctx, recorder)
}

func TestObservedRetryAndExhaustion(t *testing.T) {
	for _, status := range []int{200, 503} {
		calls := 0
		coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
			calls++
			if calls == 1 || status == 503 {
				return response(503, `{"error":{}}`), nil
			}
			return response(200, `{"choices":[{"message":{"content":"ok"}}],"usage":{"prompt_tokens":0,"completion_tokens":2}}`), nil
		})
		var events []domain.ExecutionEvent
		ctx := recordEvents(t, context.Background(), &events)
		_, err := coordinator.Generate(ctx, testRequest(1))
		if status == 200 && err != nil {
			t.Fatal(err)
		}
		if status == 503 {
			var exhausted *ExhaustedError
			if !errors.As(err, &exhausted) {
				t.Fatalf("error=%v", err)
			}
		}
		if len(events) != 3 || events[0].Kind != domain.ExecutionAttempt || events[1].Kind != domain.ExecutionRetryDelay || events[2].Kind != domain.ExecutionAttempt {
			t.Fatalf("events=%+v", events)
		}
		if events[0].Attempt != 1 || events[1].Attempt != 1 || events[2].Attempt != 2 || !events[0].Dispatched || !events[2].Dispatched {
			t.Fatalf("indexes=%+v", events)
		}
		if events[0].Failure != domain.ExecutionFailureProviderStatus || events[1].RetryDelay < 0 {
			t.Fatalf("first=%+v delay=%+v", events[0], events[1])
		}
		if status == 200 && (!events[2].Usage.PromptTokens.Known || events[2].Usage.PromptTokens.Value != 0 || events[2].Usage.TotalTokens.Known) {
			t.Fatalf("usage=%+v", events[2].Usage)
		}
	}
}

func TestObservedInterruptedWaitAndCanceledCapacity(t *testing.T) {
	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) { return response(503, `{"error":{}}`), nil })
	ctx, cancel := context.WithCancel(context.Background())
	var events []domain.ExecutionEvent
	ctx = recordEvents(t, ctx, &events)
	coordinator.wait = func(ctx context.Context, _ time.Duration) error { cancel(); return ctx.Err() }
	_, err := coordinator.Generate(ctx, testRequest(1))
	if !errors.Is(err, context.Canceled) || len(events) != 2 || events[1].Kind != domain.ExecutionRetryDelay || events[1].Failure != domain.ExecutionFailureCanceled {
		t.Fatalf("error=%v events=%+v", err, events)
	}

	coordinator = testCoordinator(t, func(*http.Request) (*http.Response, error) { t.Fatal("unexpected dispatch"); return nil, nil })
	release, err := coordinator.manager.Acquire(context.Background(), "limited")
	if err != nil {
		t.Fatal(err)
	}
	defer release()
	ctx, cancel = context.WithCancel(context.Background())
	events = nil
	ctx = recordEvents(t, ctx, &events)
	entered := make(chan struct{})
	ctx = &notifyingContext{Context: ctx, entered: entered}
	done := make(chan error, 1)
	go func() { _, err := coordinator.Generate(ctx, testRequest(0)); done <- err }()
	<-entered
	cancel()
	if err := <-done; !errors.Is(err, context.Canceled) {
		t.Fatalf("error=%v", err)
	}
	if len(events) != 1 || events[0].Dispatched || events[0].Failure != domain.ExecutionFailureCanceled {
		t.Fatalf("events=%+v", events)
	}
}

type notifyingContext struct {
	context.Context
	entered chan struct{}
}

func (c *notifyingContext) Done() <-chan struct{} {
	select {
	case <-c.entered:
	default:
		close(c.entered)
	}
	return c.Context.Done()
}

func TestObservedCallbackReleasesPermitAndCanCancelRetry(t *testing.T) {
	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) { return response(503, `{"error":{}}`), nil })
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var events []domain.ExecutionEvent
	recorder, err := observation.New(func(event domain.ExecutionEvent) {
		events = append(events, event)
		if event.Kind != domain.ExecutionAttempt {
			return
		}
		permitCtx, stop := context.WithTimeout(context.Background(), time.Second)
		defer stop()
		release, err := coordinator.manager.Acquire(permitCtx, "limited")
		if err != nil {
			t.Errorf("callback could not acquire permit: %v", err)
		} else {
			release()
		}
		cancel()
	})
	if err != nil {
		t.Fatal(err)
	}
	_, err = coordinator.Generate(observation.WithRecorder(ctx, recorder), testRequest(2))
	if !errors.Is(err, context.Canceled) || len(events) != 1 {
		t.Fatalf("error=%v events=%+v", err, events)
	}
}

type fixedClient struct {
	response *domain.GenerateResponse
	err      error
}

func (c fixedClient) Generate(context.Context, domain.GenerateRequest) (*domain.GenerateResponse, error) {
	return c.response, c.err
}

func TestObservedInjectedUsageAndRepairGeneration(t *testing.T) {
	manager, err := capacity.NewManager(nil)
	if err != nil {
		t.Fatal(err)
	}
	var events []domain.ExecutionEvent
	ctx := recordEvents(t, context.Background(), &events)
	legacy := NewInjected(manager, fixedClient{response: &domain.GenerateResponse{Usage: domain.TokenUsage{PromptTokens: 3}}})
	if _, err := legacy.Generate(ctx, testRequest(4)); err != nil {
		t.Fatal(err)
	}
	presence := &domain.TokenUsagePresence{CompletionTokens: true}
	repair := NewInjected(manager, fixedClient{response: &domain.GenerateResponse{UsagePresence: presence}})
	req := testRequest(4)
	req.Generation = 2
	if _, err := repair.Generate(ctx, req); err != nil {
		t.Fatal(err)
	}
	if len(events) != 2 || !events[0].Usage.PromptTokens.Known || events[0].Usage.CompletionTokens.Known || !events[1].Usage.CompletionTokens.Known || events[1].Usage.CompletionTokens.Value != 0 || events[1].Usage.PromptTokens.Known {
		t.Fatalf("events=%+v", events)
	}
	if events[1].Generation != 2 || events[1].Attempt != 1 {
		t.Fatalf("repair indexes=%+v", events[1])
	}
}

func TestObservedSuccessfulGenerationBeforeRepairFailure(t *testing.T) {
	manager, err := capacity.NewManager(nil)
	if err != nil {
		t.Fatal(err)
	}
	var events []domain.ExecutionEvent
	recorder, err := observation.New(func(event domain.ExecutionEvent) { events = append(events, event) })
	if err != nil {
		t.Fatal(err)
	}
	ctx := observation.WithRecorder(context.Background(), recorder)
	initial := NewInjected(manager, fixedClient{response: &domain.GenerateResponse{Usage: domain.TokenUsage{PromptTokens: 5}}})
	if _, err := initial.Generate(ctx, testRequest(0)); err != nil {
		t.Fatal(err)
	}
	repair := NewInjected(manager, fixedClient{err: errors.New("repair unavailable")})
	req := testRequest(0)
	req.Generation = 2
	if _, err := repair.Generate(ctx, req); err == nil {
		t.Fatal("expected repair failure")
	}
	recorder.Finish(observation.FinishData{Outcome: domain.ExecutionOutcomeError, Failure: domain.ExecutionFailureGeneration})
	if len(events) != 3 || events[0].Outcome != domain.ExecutionOutcomeSuccess || events[1].Outcome != domain.ExecutionOutcomeError || events[2].Attempts != 2 || events[2].RepairCalls != 1 {
		t.Fatalf("events=%+v", events)
	}
	if got := events[2].Usage.PromptTokens; got.Value != 5 || !got.Known || got.Complete {
		t.Fatalf("prompt usage=%+v", got)
	}
}