322 lines
13 KiB
Go
322 lines
13 KiB
Go
package generation
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net/http"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/llm"
|
|
)
|
|
|
|
func TestGenerateWithEvidenceRetainsAttemptsAndUsage(t *testing.T) {
|
|
var now = time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC)
|
|
calls := 0
|
|
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
if calls == 1 {
|
|
now = now.Add(4 * time.Millisecond)
|
|
return response(http.StatusServiceUnavailable, `{"error":{}}`), nil
|
|
}
|
|
now = now.Add(7 * time.Millisecond)
|
|
return response(http.StatusOK, `{"choices":[{"message":{"content":"ok"}}],"usage":{"prompt_tokens":3}}`), nil
|
|
})
|
|
coordinator.now = func() time.Time { return now }
|
|
coordinator.wait = func(context.Context, time.Duration) error {
|
|
now = now.Add(19 * time.Millisecond)
|
|
return nil
|
|
}
|
|
|
|
result, err, evidence := coordinator.GenerateWithEvidence(context.Background(), testRequest(1))
|
|
if err != nil || result == nil || result.Content != "ok" || calls != 2 {
|
|
t.Fatalf("result=%+v error=%v calls=%d", result, err, calls)
|
|
}
|
|
if len(evidence.Attempts) != 2 || len(evidence.RetryDelays) != 1 || evidence.RetryDelays[0] != 19*time.Millisecond {
|
|
t.Fatalf("attempts/delays = %#v / %#v", evidence.Attempts, evidence.RetryDelays)
|
|
}
|
|
first, second := evidence.Attempts[0], evidence.Attempts[1]
|
|
if !first.Dispatched || !first.HTTP.RequestObserved || !first.HTTP.ResponseReceived || first.HTTP.HTTPStatus == nil || *first.HTTP.HTTPStatus != http.StatusServiceUnavailable {
|
|
t.Fatalf("first attempt evidence = %+v", first)
|
|
}
|
|
if first.Failure != domain.ExecutionFailureProviderStatus || first.RequestDuration != 4*time.Millisecond {
|
|
t.Fatalf("first attempt failure/timing = %v / %v", first.Failure, first.RequestDuration)
|
|
}
|
|
if !second.Dispatched || !second.GenerationCompleted || !second.ResponseDecoded || second.HTTP.HTTPStatus == nil || *second.HTTP.HTTPStatus != http.StatusOK {
|
|
t.Fatalf("second attempt evidence = %+v", second)
|
|
}
|
|
if second.Failure != domain.ExecutionFailureNone || second.RequestDuration != 7*time.Millisecond {
|
|
t.Fatalf("second attempt failure/timing = %v / %v", second.Failure, second.RequestDuration)
|
|
}
|
|
if !second.Usage.PromptTokens.Known || !second.Usage.PromptTokens.Complete || second.Usage.PromptTokens.Value != 3 {
|
|
t.Fatalf("per-attempt usage = %+v", second.Usage.PromptTokens)
|
|
}
|
|
if !evidence.Usage.PromptTokens.Known || evidence.Usage.PromptTokens.Complete || evidence.Usage.PromptTokens.Value != 3 {
|
|
t.Fatalf("aggregate usage = %+v", evidence.Usage.PromptTokens)
|
|
}
|
|
}
|
|
|
|
func TestGenerateWithEvidenceRetainsExhaustedAttempts(t *testing.T) {
|
|
calls := 0
|
|
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
return response(http.StatusServiceUnavailable, `{"error":{}}`), nil
|
|
})
|
|
coordinator.wait = func(context.Context, time.Duration) error { return nil }
|
|
|
|
result, err, evidence := coordinator.GenerateWithEvidence(context.Background(), testRequest(1))
|
|
var exhausted *ExhaustedError
|
|
if result != nil || !errors.As(err, &exhausted) || calls != 2 || exhausted.Attempts() != 2 {
|
|
t.Fatalf("result=%+v error=%v calls=%d", result, err, calls)
|
|
}
|
|
if len(evidence.Attempts) != 2 || len(evidence.RetryDelays) != 1 {
|
|
t.Fatalf("attempt/delay evidence = %#v / %#v", evidence.Attempts, evidence.RetryDelays)
|
|
}
|
|
for _, attempt := range evidence.Attempts {
|
|
if !attempt.Dispatched || attempt.Failure != domain.ExecutionFailureProviderStatus || attempt.HTTP.HTTPStatus == nil || *attempt.HTTP.HTTPStatus != http.StatusServiceUnavailable {
|
|
t.Fatalf("exhausted attempt evidence = %+v", attempt)
|
|
}
|
|
}
|
|
if evidence.Usage.PromptTokens.Known || evidence.Usage.PromptTokens.Complete {
|
|
t.Fatalf("missing provider usage was claimed: %+v", evidence.Usage.PromptTokens)
|
|
}
|
|
}
|
|
|
|
func TestGenerateWithEvidenceRetainsInterruptedBackoffWithoutAttempt(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
var now = time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC)
|
|
calls := 0
|
|
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
now = now.Add(3 * time.Millisecond)
|
|
return response(http.StatusServiceUnavailable, `{"error":{}}`), nil
|
|
})
|
|
coordinator.now = func() time.Time { return now }
|
|
coordinator.wait = func(ctx context.Context, _ time.Duration) error {
|
|
now = now.Add(11 * time.Millisecond)
|
|
cancel()
|
|
return ctx.Err()
|
|
}
|
|
|
|
result, err, evidence := coordinator.GenerateWithEvidence(ctx, testRequest(2))
|
|
if result != nil || !errors.Is(err, context.Canceled) || calls != 1 {
|
|
t.Fatalf("result=%+v error=%v calls=%d", result, err, calls)
|
|
}
|
|
if len(evidence.Attempts) != 1 || !evidence.Attempts[0].Dispatched || len(evidence.RetryDelays) != 1 || evidence.RetryDelays[0] != 11*time.Millisecond {
|
|
t.Fatalf("interrupted evidence = %#v", evidence)
|
|
}
|
|
if evidence.Attempts[0].HTTP.HTTPStatus == nil || *evidence.Attempts[0].HTTP.HTTPStatus != http.StatusServiceUnavailable {
|
|
t.Fatalf("last dispatched response was lost: %+v", evidence.Attempts[0])
|
|
}
|
|
}
|
|
|
|
func TestGenerateWithEvidenceRecordsCanceledReacquisition(t *testing.T) {
|
|
var now = time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC)
|
|
calls := 0
|
|
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
now = now.Add(2 * time.Millisecond)
|
|
return response(http.StatusServiceUnavailable, `{"error":{}}`), nil
|
|
})
|
|
coordinator.now = func() time.Time { return now }
|
|
baseContext, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
contextWithCancelOnWait := &cancelOnDoneContext{Context: baseContext, cancel: cancel}
|
|
var releaseHeld func()
|
|
coordinator.wait = func(context.Context, time.Duration) error {
|
|
release, err := coordinator.manager.Acquire(context.Background(), "limited")
|
|
if err != nil {
|
|
t.Fatalf("hold permit: %v", err)
|
|
}
|
|
releaseHeld = release
|
|
contextWithCancelOnWait.armed.Store(true)
|
|
return nil
|
|
}
|
|
|
|
// Signal cancellation only when the second acquisition has queued behind
|
|
// the held permit, avoiding scheduler- or sleep-based synchronization.
|
|
result, err, evidence := coordinator.GenerateWithEvidence(contextWithCancelOnWait, testRequest(2))
|
|
if releaseHeld != nil {
|
|
releaseHeld()
|
|
}
|
|
if result != nil || !errors.Is(err, context.Canceled) || calls != 1 {
|
|
t.Fatalf("result=%+v error=%v calls=%d", result, err, calls)
|
|
}
|
|
if len(evidence.Attempts) != 2 || !evidence.Attempts[0].Dispatched || evidence.Attempts[1].Dispatched {
|
|
t.Fatalf("acquisition evidence = %#v", evidence.Attempts)
|
|
}
|
|
if evidence.Attempts[1].Failure != domain.ExecutionFailureCanceled {
|
|
t.Fatalf("canceled acquisition failure = %v", evidence.Attempts[1].Failure)
|
|
}
|
|
if evidence.Attempts[0].HTTP.HTTPStatus == nil || *evidence.Attempts[0].HTTP.HTTPStatus != http.StatusServiceUnavailable {
|
|
t.Fatalf("last dispatched evidence was lost: %+v", evidence.Attempts)
|
|
}
|
|
}
|
|
|
|
func TestGenerateWithEvidenceNoDispatchHasCompleteZeroUsage(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
|
|
t.Fatal("canceled generation dispatched")
|
|
return nil, nil
|
|
})
|
|
|
|
result, err, evidence := coordinator.GenerateWithEvidence(ctx, testRequest(0))
|
|
if result != nil || !errors.Is(err, context.Canceled) || len(evidence.Attempts) != 0 {
|
|
t.Fatalf("result=%+v error=%v evidence=%+v", result, err, evidence)
|
|
}
|
|
if !evidence.Usage.PromptTokens.Known || !evidence.Usage.PromptTokens.Complete || evidence.Usage.PromptTokens.Value != 0 {
|
|
t.Fatalf("no-dispatch usage = %+v", evidence.Usage.PromptTokens)
|
|
}
|
|
}
|
|
|
|
func TestGenerateWithEvidenceForInjectedClient(t *testing.T) {
|
|
opaqueFailure := errors.New("injected client failed")
|
|
presence := &domain.TokenUsagePresence{PromptTokens: true}
|
|
for _, tc := range []struct {
|
|
name string
|
|
response *domain.GenerateResponse
|
|
err error
|
|
wantResponse bool
|
|
wantGeneration bool
|
|
wantFailure domain.ExecutionFailureKind
|
|
wantUsage bool
|
|
wantStructured bool
|
|
}{
|
|
{
|
|
name: "success",
|
|
response: &domain.GenerateResponse{Content: "private output", Usage: domain.TokenUsage{PromptTokens: 4}, UsagePresence: presence},
|
|
wantResponse: true,
|
|
wantGeneration: true,
|
|
wantUsage: true,
|
|
wantStructured: true,
|
|
},
|
|
{
|
|
name: "nil response",
|
|
wantFailure: domain.ExecutionFailureMalformedResponse,
|
|
},
|
|
{
|
|
name: "opaque failure",
|
|
err: opaqueFailure,
|
|
wantFailure: domain.ExecutionFailureGeneration,
|
|
},
|
|
{
|
|
name: "response with error",
|
|
response: &domain.GenerateResponse{Content: "must not be evaluated", Usage: domain.TokenUsage{PromptTokens: 9}, UsagePresence: presence},
|
|
err: opaqueFailure,
|
|
wantResponse: true,
|
|
wantFailure: domain.ExecutionFailureGeneration,
|
|
wantStructured: true,
|
|
},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
client := &evidenceInjectedClient{response: tc.response, err: tc.err}
|
|
coordinator := newEvidenceInjectedCoordinator(t, client)
|
|
req := testRequest(7)
|
|
if tc.wantStructured {
|
|
req.StructuredOutput = &domain.StructuredOutputSpec{}
|
|
}
|
|
|
|
result, err, evidence := coordinator.GenerateWithEvidence(context.Background(), req)
|
|
if client.calls != 1 || len(evidence.Attempts) != 1 || len(evidence.RetryDelays) != 0 {
|
|
t.Fatalf("calls=%d attempts=%+v retry delays=%v", client.calls, evidence.Attempts, evidence.RetryDelays)
|
|
}
|
|
attempt := evidence.Attempts[0]
|
|
if !attempt.Dispatched || attempt.ResponseReceived != tc.wantResponse || attempt.GenerationCompleted != tc.wantGeneration || attempt.Failure != tc.wantFailure {
|
|
t.Fatalf("attempt evidence = %+v", attempt)
|
|
}
|
|
if attempt.HTTP != (llm.HTTPExecutionEvidence{}) {
|
|
t.Fatalf("injected call claimed HTTP evidence: %+v", attempt.HTTP)
|
|
}
|
|
if attempt.StructuredOutputRequested == nil || *attempt.StructuredOutputRequested != tc.wantStructured {
|
|
t.Fatalf("structured-output request fact = %v", attempt.StructuredOutputRequested)
|
|
}
|
|
if tc.wantFailure == domain.ExecutionFailureNone {
|
|
if err != nil || result == nil || result.Content != "private output" {
|
|
t.Fatalf("result=%+v error=%v", result, err)
|
|
}
|
|
} else if result != nil || err == nil {
|
|
t.Fatalf("result=%+v error=%v; expected failed call", result, err)
|
|
}
|
|
if attempt.Usage.PromptTokens.Known != tc.wantUsage || evidence.Usage.PromptTokens.Known != tc.wantUsage {
|
|
t.Fatalf("attempt/aggregate usage = %+v / %+v", attempt.Usage, evidence.Usage)
|
|
}
|
|
if tc.wantUsage && (attempt.Usage.PromptTokens.Value != 4 || evidence.Usage.PromptTokens.Value != 4) {
|
|
t.Fatalf("reported usage was not retained: %+v / %+v", attempt.Usage, evidence.Usage)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGenerateWithEvidenceRetainsInjectedSuccessAfterCancellation(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
client := &evidenceInjectedClient{
|
|
response: &domain.GenerateResponse{Content: "private output"},
|
|
afterGenerate: func() { cancel() },
|
|
}
|
|
coordinator := newEvidenceInjectedCoordinator(t, client)
|
|
req := testRequest(5)
|
|
req.StructuredOutput = &domain.StructuredOutputSpec{}
|
|
|
|
result, err, evidence := coordinator.GenerateWithEvidence(ctx, req)
|
|
if !errors.Is(err, context.Canceled) || result == nil || result.Content != "private output" {
|
|
t.Fatalf("result=%+v error=%v", result, err)
|
|
}
|
|
if client.calls != 1 || len(evidence.Attempts) != 1 {
|
|
t.Fatalf("calls=%d attempts=%+v", client.calls, evidence.Attempts)
|
|
}
|
|
attempt := evidence.Attempts[0]
|
|
if !attempt.ResponseReceived || !attempt.GenerationCompleted || attempt.Failure != domain.ExecutionFailureNone {
|
|
t.Fatalf("successful generation evidence was lost: %+v", attempt)
|
|
}
|
|
if attempt.HTTP != (llm.HTTPExecutionEvidence{}) {
|
|
t.Fatalf("injected call claimed HTTP evidence: %+v", attempt.HTTP)
|
|
}
|
|
}
|
|
|
|
type evidenceInjectedClient struct {
|
|
calls int
|
|
response *domain.GenerateResponse
|
|
err error
|
|
afterGenerate func()
|
|
}
|
|
|
|
func (c *evidenceInjectedClient) Generate(context.Context, domain.GenerateRequest) (*domain.GenerateResponse, error) {
|
|
c.calls++
|
|
if c.afterGenerate != nil {
|
|
c.afterGenerate()
|
|
}
|
|
return c.response, c.err
|
|
}
|
|
|
|
func newEvidenceInjectedCoordinator(t *testing.T, client llm.Client) *Coordinator {
|
|
t.Helper()
|
|
manager, err := capacity.NewManager(map[string]domain.BackendCapacityPolicy{
|
|
"limited": {ConcurrencyLimit: 1},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return NewInjected(manager, client)
|
|
}
|
|
|
|
type cancelOnDoneContext struct {
|
|
context.Context
|
|
cancel context.CancelFunc
|
|
once sync.Once
|
|
armed atomic.Bool
|
|
}
|
|
|
|
func (c *cancelOnDoneContext) Done() <-chan struct{} {
|
|
if c.armed.Load() {
|
|
c.once.Do(c.cancel)
|
|
}
|
|
return c.Context.Done()
|
|
}
|