Separate frozen HTTP preparation from attempts

2026-09-23 16:43:51 +00:00
parent d31a563ca1
commit 65d5894648
4 changed files with 380 additions and 51 deletions
--- a/internal/llm/openai_compatible_client.go
+++ b/internal/llm/openai_compatible_client.go
@@ -7,9 +7,11 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"net/http"
 	"net/url"
 	"os"
+	"strconv"
 	"strings"
 	"time"

@@ -87,6 +89,40 @@ func NewOpenAICompatibleClient(cfg OpenAICompatibleConfig) (*OpenAICompatibleCli
 }

 func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
+	prepared, err := c.PrepareGeneration(req)
+	if err != nil {
+		return nil, err
+	}
+	outcome := prepared.Execute(ctx)
+	if outcome.Err != nil {
+		return nil, outcome.Err
+	}
+	return outcome.Response, nil
+}
+
+// PreparedGeneration holds one encoded request and resolved authentication for
+// repeated attempts. Each Execute call creates a fresh HTTP request and context.
+type PreparedGeneration struct {
+	endpoint       string
+	payload        []byte
+	apiKey         string
+	timeoutSeconds int
+	httpClient     *http.Client
+}
+
+// AttemptOutcome describes one HTTP execution without exposing response bodies
+// or headers to the orchestration layer.
+type AttemptOutcome struct {
+	Response       *domain.GenerateResponse
+	Err            error
+	StatusCode     int
+	Usage          domain.TokenUsage
+	UsagePresence  *domain.TokenUsagePresence
+	RetryNotBefore time.Time
+}
+
+// PrepareGeneration validates and freezes one built-in request before dispatch.
+func (c *OpenAICompatibleClient) PrepareGeneration(req domain.GenerateRequest) (*PreparedGeneration, error) {
 	if err := domain.ValidateExecutionTargetSettings(req.Target); err != nil {
 		return nil, fmt.Errorf("%w: %v", ErrInvalidRequest, err)
 	}
@@ -118,22 +154,6 @@ func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.Genera
 	if err != nil {
 		return nil, fmt.Errorf("%w: failed to encode request: %v", ErrRequestFailed, err)
 	}
-
-	requestContext := ctx
-	if req.Target.TimeoutSeconds > 0 {
-		var cancel context.CancelFunc
-		requestContext, cancel = context.WithTimeout(
-			ctx,
-			time.Duration(req.Target.TimeoutSeconds)*time.Second,
-		)
-		defer cancel()
-	}
-
-	httpReq, err := http.NewRequestWithContext(requestContext, http.MethodPost, endpoint, bytes.NewReader(payload))
-	if err != nil {
-		return nil, fmt.Errorf("%w: failed to create request: %v", ErrRequestFailed, err)
-	}
-	httpReq.Header.Set("Content-Type", "application/json")
 	apiKey := strings.TrimSpace(req.Target.APIKey)
 	envName := strings.TrimSpace(req.Target.APIKeyEnv)
 	if apiKey == "" && envName != "" {
@@ -145,62 +165,133 @@ func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.Genera
 		}
 		return nil, fmt.Errorf("%w: api key is required", ErrInvalidRequest)
 	}
-	if apiKey != "" {
-		httpReq.Header.Set("Authorization", "Bearer "+apiKey)
-	}
-
 	httpClient := c.httpClient
 	if httpClient == nil {
 		httpClient = &http.Client{Timeout: defaults.LLMRequestTimeoutDefault}
 	}
+	return &PreparedGeneration{
+		endpoint: endpoint, payload: payload, apiKey: apiKey,
+		timeoutSeconds: req.Target.TimeoutSeconds, httpClient: httpClient,
+	}, nil
+}

-	httpResp, err := httpClient.Do(httpReq)
-	if err != nil {
-		return nil, &requestFailedError{cause: err}
+// Execute dispatches one attempt. The returned response is non-nil only when
+// the successful envelope contains a usable first choice.
+func (p *PreparedGeneration) Execute(ctx context.Context) AttemptOutcome {
+	requestContext := ctx
+	if p.timeoutSeconds > 0 {
+		var cancel context.CancelFunc
+		requestContext, cancel = context.WithTimeout(ctx, time.Duration(p.timeoutSeconds)*time.Second)
+		defer cancel()
 	}
+	httpReq, err := http.NewRequestWithContext(requestContext, http.MethodPost, p.endpoint, bytes.NewReader(p.payload))
+	if err != nil {
+		return AttemptOutcome{Err: fmt.Errorf("%w: failed to create request: %v", ErrRequestFailed, err)}
+	}
+	httpReq.Header.Set("Content-Type", "application/json")
+	if p.apiKey != "" {
+		httpReq.Header.Set("Authorization", "Bearer "+p.apiKey)
+	}
+
+	httpResp, err := p.httpClient.Do(httpReq)
+	if err != nil {
+		return AttemptOutcome{Err: &requestFailedError{cause: err}}
+	}
+	receivedAt := time.Now()
 	defer httpResp.Body.Close()
+	outcome := AttemptOutcome{StatusCode: httpResp.StatusCode}

 	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
-		return nil, providerHTTPErrorFromBody(
+		outcome.RetryNotBefore = retryNotBefore(httpResp.Header.Values("Retry-After"), receivedAt)
+		outcome.Err = providerHTTPErrorFromBody(
 			httpResp.StatusCode,
 			httpResp.ContentLength,
 			httpResp.Body,
 		)
+		if err := observedContextError(requestContext, ctx, outcome.Err); err != nil {
+			outcome.Err = err
+		}
+		return outcome
 	}
 	if httpResp.ContentLength > maxOpenAIChatResponseBytes {
-		return nil, openAIChatResponseTooLargeError()
+		outcome.Err = openAIChatResponseTooLargeError()
+		return outcome
 	}

 	wireResp, err := decodeOpenAIChatResponse(httpResp.Body)
 	if err != nil {
-		return nil, err
+		outcome.Err = err
+		if contextErr := observedContextError(requestContext, ctx, err); contextErr != nil {
+			outcome.Err = contextErr
+		}
+		return outcome
+	}
+	outcome.Usage = domain.TokenUsage{
+		PromptTokens: tokenCount(wireResp.Usage.PromptTokens), CompletionTokens: tokenCount(wireResp.Usage.CompletionTokens),
+		TotalTokens: tokenCount(wireResp.Usage.TotalTokens), CachedTokens: tokenCount(wireResp.Usage.PromptTokensDetails.CachedTokens),
+		CacheWriteTokens: tokenCount(wireResp.Usage.CacheWriteTokens),
+	}
+	outcome.UsagePresence = &domain.TokenUsagePresence{
+		PromptTokens: wireResp.Usage.PromptTokens != nil, CompletionTokens: wireResp.Usage.CompletionTokens != nil,
+		TotalTokens: wireResp.Usage.TotalTokens != nil, CachedTokens: wireResp.Usage.PromptTokensDetails.CachedTokens != nil,
+		CacheWriteTokens: wireResp.Usage.CacheWriteTokens != nil,
 	}

 	if len(wireResp.Choices) == 0 {
-		return nil, fmt.Errorf("%w: no choices returned", ErrMalformedResponse)
+		outcome.Err = fmt.Errorf("%w: no choices returned", ErrMalformedResponse)
+		return outcome
 	}
 	content := wireResp.Choices[0].Message.Content
 	if content == nil {
-		return nil, fmt.Errorf("%w: first choice has missing message content", ErrMalformedResponse)
+		outcome.Err = fmt.Errorf("%w: first choice has missing message content", ErrMalformedResponse)
+		return outcome
 	}
+	outcome.Response = &domain.GenerateResponse{Content: *content, Usage: outcome.Usage, UsagePresence: outcome.UsagePresence}
+	return outcome
+}

-	return &domain.GenerateResponse{
-		Content: *content,
-		Usage: domain.TokenUsage{
-			PromptTokens:     tokenCount(wireResp.Usage.PromptTokens),
-			CompletionTokens: tokenCount(wireResp.Usage.CompletionTokens),
-			TotalTokens:      tokenCount(wireResp.Usage.TotalTokens),
-			CachedTokens:     tokenCount(wireResp.Usage.PromptTokensDetails.CachedTokens),
-			CacheWriteTokens: tokenCount(wireResp.Usage.CacheWriteTokens),
-		},
-		UsagePresence: &domain.TokenUsagePresence{
-			PromptTokens:     wireResp.Usage.PromptTokens != nil,
-			CompletionTokens: wireResp.Usage.CompletionTokens != nil,
-			TotalTokens:      wireResp.Usage.TotalTokens != nil,
-			CachedTokens:     wireResp.Usage.PromptTokensDetails.CachedTokens != nil,
-			CacheWriteTokens: wireResp.Usage.CacheWriteTokens != nil,
-		},
-	}, nil
+func observedContextError(attempt, caller context.Context, cause error) error {
+	if err := caller.Err(); err != nil {
+		return &requestFailedError{cause: err}
+	}
+	if err := attempt.Err(); err != nil {
+		return &requestFailedError{cause: err}
+	}
+	if errors.Is(cause, context.Canceled) || errors.Is(cause, context.DeadlineExceeded) {
+		return &requestFailedError{cause: cause}
+	}
+	return nil
+}
+
+func retryNotBefore(values []string, receivedAt time.Time) time.Time {
+	if len(values) == 0 {
+		return time.Time{}
+	}
+	value := strings.TrimSpace(values[0])
+	if value == "" {
+		return time.Time{}
+	}
+	digits := true
+	for _, r := range value {
+		if r < '0' || r > '9' {
+			digits = false
+			break
+		}
+	}
+	if digits {
+		seconds, err := strconv.ParseUint(value, 10, 64)
+		if err != nil || seconds > uint64(math.MaxInt64/int64(time.Second)) {
+			return receivedAt.Add(time.Duration(math.MaxInt64))
+		}
+		return receivedAt.Add(time.Duration(seconds) * time.Second)
+	}
+	if when, err := http.ParseTime(value); err == nil {
+		if when.Before(receivedAt) {
+			return receivedAt
+		}
+		return when
+	}
+	return time.Time{}
 }

 func tokenCount(value *int) int {
@@ -219,6 +310,9 @@ func decodeOpenAIChatResponse(body io.Reader) (openAIChatResponse, error) {

 	var response openAIChatResponse
 	if err := decoder.Decode(&response); err != nil {
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			return openAIChatResponse{}, &requestFailedError{cause: err}
+		}
 		if limited.N == 0 {
 			return openAIChatResponse{}, openAIChatResponseTooLargeError()
 		}
@@ -230,6 +324,9 @@ func decodeOpenAIChatResponse(body io.Reader) (openAIChatResponse, error) {

 	var trailing any
 	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			return openAIChatResponse{}, &requestFailedError{cause: err}
+		}
 		if limited.N == 0 {
 			return openAIChatResponse{}, openAIChatResponseTooLargeError()
 		}
--- a/internal/llm/prepared_generation_test.go
+++ b/internal/llm/prepared_generation_test.go
@@ -0,0 +1,225 @@
+package llm
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net/http"
+	"reflect"
+	"strings"
+	"testing"
+	"time"
+
+	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
+)
+
+func TestPreparedGenerationFreezesPayloadAndAuthentication(t *testing.T) {
+	for _, tc := range []struct {
+		name       string
+		initialKey string
+		wantAuth   string
+	}{
+		{name: "absent optional key"},
+		{name: "present key", initialKey: "original-key", wantAuth: "Bearer original-key"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			const envName = "PROMPTKIT_PREPARED_GENERATION_TEST_KEY"
+			t.Setenv(envName, tc.initialKey)
+			var bodies, auth, urls []string
+			client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
+				BaseURL: "https://example.invalid/v1", Model: "default-model",
+				HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+					body, err := io.ReadAll(req.Body)
+					if err != nil {
+						t.Fatal(err)
+					}
+					bodies = append(bodies, string(body))
+					auth = append(auth, req.Header.Get("Authorization"))
+					urls = append(urls, req.URL.String())
+					return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`))}, nil
+				})},
+			})
+			if err != nil {
+				t.Fatal(err)
+			}
+			req := domain.GenerateRequest{
+				Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "original"}}},
+				Target: domain.ExecutionTarget{APIKeyEnv: envName},
+			}
+			prepared, err := client.PrepareGeneration(req)
+			if err != nil {
+				t.Fatal(err)
+			}
+			req.Prompt.Messages[0].Content = "changed"
+			t.Setenv(envName, "later-key")
+			for range 2 {
+				outcome := prepared.Execute(context.Background())
+				if outcome.Err != nil || outcome.Response == nil || outcome.Response.Content != "ok" {
+					t.Fatalf("Execute = %+v", outcome)
+				}
+			}
+			if len(bodies) != 2 || bodies[0] != bodies[1] || !strings.Contains(bodies[0], "original") || strings.Contains(bodies[0], "changed") || !reflect.DeepEqual(auth, []string{tc.wantAuth, tc.wantAuth}) || !reflect.DeepEqual(urls, []string{"https://example.invalid/v1/chat/completions", "https://example.invalid/v1/chat/completions"}) {
+				t.Fatalf("bodies=%v auth=%v urls=%v", bodies, auth, urls)
+			}
+		})
+	}
+}
+
+func TestPreparedGenerationAttemptsHaveIndependentContextsAndCloseBodies(t *testing.T) {
+	var contexts []context.Context
+	var closed int
+	client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
+		BaseURL: "https://example.invalid/v1", Model: "model",
+		HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+			if err := req.Context().Err(); err != nil {
+				return nil, err
+			}
+			contexts = append(contexts, req.Context())
+			return &http.Response{StatusCode: http.StatusOK, Body: &countingBody{
+				Reader: strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`),
+				closed: &closed,
+			}}, nil
+		})},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	prepared, err := client.PrepareGeneration(domain.GenerateRequest{
+		Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
+		Target: domain.ExecutionTarget{TimeoutSeconds: 1},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	for range 2 {
+		if outcome := prepared.Execute(context.Background()); outcome.Err != nil {
+			t.Fatal(outcome.Err)
+		}
+	}
+	if len(contexts) != 2 || contexts[0] == contexts[1] || closed != 2 {
+		t.Fatalf("contexts=%d, closed=%d", len(contexts), closed)
+	}
+	for _, ctx := range contexts {
+		if _, ok := ctx.Deadline(); !ok || !errors.Is(ctx.Err(), context.Canceled) {
+			t.Fatalf("attempt context deadline/error = %v/%v", ok, ctx.Err())
+		}
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if outcome := prepared.Execute(ctx); !errors.Is(outcome.Err, context.Canceled) {
+		t.Fatalf("canceled attempt = %+v", outcome)
+	}
+	if outcome := prepared.Execute(context.Background()); outcome.Err != nil || outcome.Response == nil {
+		t.Fatalf("later attempt = %+v", outcome)
+	}
+}
+
+func TestPreparedGenerationRetainsUsageFromMalformedChoice(t *testing.T) {
+	client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
+		BaseURL: "https://example.invalid/v1", Model: "model",
+		HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
+			return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"usage":{"prompt_tokens":0},"choices":[]}`))}, nil
+		})},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	prepared, err := client.PrepareGeneration(domain.GenerateRequest{
+		Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	outcome := prepared.Execute(context.Background())
+	if !errors.Is(outcome.Err, ErrMalformedResponse) || outcome.Response != nil || outcome.UsagePresence == nil || !outcome.UsagePresence.PromptTokens || outcome.Usage.PromptTokens != 0 {
+		t.Fatalf("malformed choice = %+v", outcome)
+	}
+}
+
+func TestPreparedGenerationDoesNotUseInvalidlyFramedUsage(t *testing.T) {
+	client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
+		BaseURL: "https://example.invalid/v1", Model: "model",
+		HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
+			return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"usage":{"prompt_tokens":9},"choices":[]}{}`))}, nil
+		})},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	prepared, err := client.PrepareGeneration(domain.GenerateRequest{
+		Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	outcome := prepared.Execute(context.Background())
+	if !errors.Is(outcome.Err, ErrMalformedResponse) || outcome.UsagePresence != nil || outcome.Usage != (domain.TokenUsage{}) {
+		t.Fatalf("invalid framing = %+v", outcome)
+	}
+}
+
+func TestPreparedGenerationRetainsReadCancellation(t *testing.T) {
+	closed := false
+	client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
+		BaseURL: "https://example.invalid/v1", Model: "model",
+		HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
+			return &http.Response{StatusCode: http.StatusOK, Body: &readErrorBody{err: context.Canceled, closed: &closed}}, nil
+		})},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	prepared, err := client.PrepareGeneration(domain.GenerateRequest{
+		Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	outcome := prepared.Execute(context.Background())
+	if !errors.Is(outcome.Err, context.Canceled) || !errors.Is(outcome.Err, ErrTransport) || errors.Is(outcome.Err, ErrMalformedResponse) || !closed {
+		t.Fatalf("read cancellation = %+v, closed=%t", outcome, closed)
+	}
+}
+
+func TestRetryNotBefore(t *testing.T) {
+	received := time.Date(2026, 9, 23, 12, 0, 0, 0, time.UTC)
+	for _, tc := range []struct {
+		name  string
+		value string
+		want  time.Time
+	}{
+		{name: "seconds", value: " 12 ", want: received.Add(12 * time.Second)},
+		{name: "date", value: "Wed, 23 Sep 2026 12:00:15 GMT", want: received.Add(15 * time.Second)},
+		{name: "past date", value: "Wed, 23 Sep 2026 11:59:59 GMT", want: received},
+		{name: "negative", value: "-2"},
+		{name: "fractional", value: "1.5"},
+		{name: "malformed", value: "tomorrow"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := retryNotBefore([]string{tc.value, "999"}, received); !got.Equal(tc.want) {
+				t.Fatalf("not before = %v, want %v", got, tc.want)
+			}
+		})
+	}
+	if got := retryNotBefore([]string{"999999999999999999999999999999999999"}, received); !got.After(received) {
+		t.Fatalf("large valid delay wrapped: %v", got)
+	}
+}
+
+type countingBody struct {
+	io.Reader
+	closed *int
+}
+
+func (b *countingBody) Close() error {
+	*b.closed++
+	return nil
+}
+
+type readErrorBody struct {
+	err    error
+	closed *bool
+}
+
+func (b *readErrorBody) Read([]byte) (int, error) { return 0, b.err }
+func (b *readErrorBody) Close() error             { *b.closed = true; return nil }