189 lines
7.7 KiB
Go
189 lines
7.7 KiB
Go
package generation
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net/http"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
|
)
|
|
|
|
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)
|
|
}
|
|
}
|
|
|
|
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()
|
|
}
|