promptkit/internal/llm/prepared_generation_evidence_test.go

package llm

import (
	"context"
	"errors"
	"io"
	"net/http"
	"strings"
	"testing"

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

func TestPreparedGenerationEvidenceSeparatesResponseAndGenerationFacts(t *testing.T) {
	tests := []struct {
		name               string
		status             int
		body               string
		structured         bool
		wantDecoded        bool
		wantGeneration     bool
		wantResponse       bool
		wantContent        string
		wantFailure        error
		wantUsagePresence  bool
		wantStructuredFact bool
	}{
		{
			name:               "valid content",
			status:             http.StatusOK,
			body:               `{"choices":[{"message":{"content":"ok"}}],"usage":{"prompt_tokens":0}}`,
			structured:         true,
			wantDecoded:        true,
			wantGeneration:     true,
			wantResponse:       true,
			wantContent:        "ok",
			wantUsagePresence:  true,
			wantStructuredFact: true,
		},
		{
			name:               "empty content",
			status:             http.StatusOK,
			body:               `{"choices":[{"message":{"content":""}}]}`,
			wantDecoded:        true,
			wantGeneration:     true,
			wantResponse:       true,
			wantStructuredFact: false,
		},
		{
			name:               "decoded envelope without generated content",
			status:             http.StatusOK,
			body:               `{"choices":[],"usage":{"prompt_tokens":0}}`,
			wantDecoded:        true,
			wantUsagePresence:  true,
			wantFailure:        ErrMalformedResponse,
			wantStructuredFact: false,
		},
		{
			name:               "malformed envelope",
			status:             http.StatusOK,
			body:               `not json`,
			wantFailure:        ErrMalformedResponse,
			wantStructuredFact: false,
		},
		{
			name:               "provider status",
			status:             http.StatusTooManyRequests,
			body:               `{"error":{"message":"try again"}}`,
			wantFailure:        ErrUnexpectedStatus,
			wantStructuredFact: false,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
				BaseURL: "https://provider.example/v1",
				Model:   "model",
				HTTPClient: &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
					response := responseFor(request, test.status, "")
					response.Body = io.NopCloser(strings.NewReader(test.body))
					response.ContentLength = int64(len(test.body))
					return response, nil
				})},
			})
			if err != nil {
				t.Fatal(err)
			}
			request := ordinaryGenerateRequest()
			if test.structured {
				request.StructuredOutput = &domain.StructuredOutputSpec{
					Type: domain.StructuredOutputJSONSchema,
					JSONSchema: &domain.StructuredOutputJSONSpec{
						Name:   "smoke",
						Schema: map[string]any{"type": "object"},
					},
				}
			}
			prepared, err := client.PrepareGeneration(request)
			if err != nil {
				t.Fatal(err)
			}

			outcome := prepared.ExecuteWithEvidence(context.Background())
			if (outcome.Err != nil) != (test.wantFailure != nil) || (test.wantFailure != nil && !errors.Is(outcome.Err, test.wantFailure)) {
				t.Fatalf("error = %v, want failure %v", outcome.Err, test.wantFailure)
			}
			if !outcome.Evidence.DispatchInvoked || !outcome.Evidence.RequestObserved || !outcome.Evidence.ResponseReceived {
				t.Fatalf("response evidence = %+v", outcome.Evidence)
			}
			if outcome.Evidence.HTTPStatus == nil || *outcome.Evidence.HTTPStatus != test.status {
				t.Fatalf("HTTP evidence status = %v, want %d", outcome.Evidence.HTTPStatus, test.status)
			}
			if outcome.GenerationCompleted != test.wantGeneration {
				t.Fatalf("generation completed = %t, want %t", outcome.GenerationCompleted, test.wantGeneration)
			}
			if outcome.ResponseDecoded != test.wantDecoded {
				t.Fatalf("response decoded = %t, want %t", outcome.ResponseDecoded, test.wantDecoded)
			}
			if (outcome.Response != nil) != test.wantResponse {
				t.Fatalf("response presence = %t, want %t", outcome.Response != nil, test.wantResponse)
			}
			if test.wantResponse && outcome.Response.Content != test.wantContent {
				t.Fatalf("generated content = %q, want %q", outcome.Response.Content, test.wantContent)
			}
			if (outcome.UsagePresence != nil && outcome.UsagePresence.PromptTokens) != test.wantUsagePresence {
				t.Fatalf("usage presence = %+v, want prompt token presence %t", outcome.UsagePresence, test.wantUsagePresence)
			}
			if outcome.StructuredOutputRequested == nil || *outcome.StructuredOutputRequested != test.wantStructuredFact {
				t.Fatalf("structured-output evidence = %v, want %t", outcome.StructuredOutputRequested, test.wantStructuredFact)
			}
		})
	}
}

func TestPreparedGenerationEvidenceRetainsDecodedResponseOnCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	body := `{"choices":[{"message":{"content":"ok"}}],"usage":{"prompt_tokens":0}}`
	client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
		BaseURL: "https://provider.example/v1",
		Model:   "model",
		HTTPClient: &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
			return &http.Response{
				StatusCode:    http.StatusOK,
				Header:        make(http.Header),
				Body:          &cancelAfterReadBody{Reader: strings.NewReader(body), cancel: cancel},
				ContentLength: int64(len(body)),
				Request:       request,
			}, nil
		})},
	})
	if err != nil {
		t.Fatal(err)
	}
	prepared, err := client.PrepareGeneration(ordinaryGenerateRequest())
	if err != nil {
		t.Fatal(err)
	}

	outcome := prepared.ExecuteWithEvidence(ctx)
	if !errors.Is(outcome.Err, context.Canceled) || !errors.Is(outcome.Err, ErrTransport) {
		t.Fatalf("canceled outcome = %+v", outcome)
	}
	if !outcome.Evidence.ResponseReceived || outcome.Evidence.HTTPStatus == nil || *outcome.Evidence.HTTPStatus != http.StatusOK {
		t.Fatalf("response evidence was lost after cancellation: %+v", outcome.Evidence)
	}
	if !outcome.ResponseDecoded || !outcome.GenerationCompleted || outcome.Response == nil || outcome.Response.Content != "ok" {
		t.Fatalf("decoded generation was lost after cancellation: %+v", outcome)
	}
	if outcome.UsagePresence == nil || !outcome.UsagePresence.PromptTokens || outcome.Usage.PromptTokens != 0 {
		t.Fatalf("decoded usage was lost after cancellation: %+v", outcome)
	}
}

func TestPreparedGenerationEvidenceRecordsStructuredRequestAtDispatch(t *testing.T) {
	transportErr := errors.New("provider unavailable")
	client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
		BaseURL: "https://provider.example/v1",
		Model:   "model",
		HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
			return nil, transportErr
		})},
	})
	if err != nil {
		t.Fatal(err)
	}
	request := ordinaryGenerateRequest()
	request.StructuredOutput = &domain.StructuredOutputSpec{
		Type: domain.StructuredOutputJSONSchema,
		JSONSchema: &domain.StructuredOutputJSONSpec{
			Name:   "smoke",
			Schema: map[string]any{"type": "object"},
		},
	}
	prepared, err := client.PrepareGeneration(request)
	if err != nil {
		t.Fatal(err)
	}

	outcome := prepared.ExecuteWithEvidence(context.Background())
	if !errors.Is(outcome.Err, transportErr) || !outcome.Evidence.DispatchInvoked || !outcome.Evidence.RequestObserved {
		t.Fatalf("dispatch outcome = %+v", outcome)
	}
	if outcome.Evidence.ResponseReceived || outcome.StructuredOutputRequested == nil || !*outcome.StructuredOutputRequested {
		t.Fatalf("structured request evidence = %+v, request=%v", outcome.Evidence, outcome.StructuredOutputRequested)
	}
}

type cancelAfterReadBody struct {
	io.Reader
	cancel context.CancelFunc
}

func (b *cancelAfterReadBody) Read(buffer []byte) (int, error) {
	n, err := b.Reader.Read(buffer)
	if n > 0 {
		b.cancel()
	}
	return n, err
}

func (b *cancelAfterReadBody) Close() error { return nil }