ai-agent/internal/ai/openai_client_test.go

package ai

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

	"code.tczkiot.com/wlw/ai-agent/internal/models"

	openai "github.com/openai/openai-go/v3"
	"github.com/openai/openai-go/v3/shared"
)

type requestIDRecordingTransport struct {
	mu         sync.Mutex
	attempts   int
	requestIDs []string
}

func (t *requestIDRecordingTransport) RoundTrip(request *http.Request) (*http.Response, error) {
	t.mu.Lock()
	t.attempts++
	attempt := t.attempts
	t.requestIDs = append(t.requestIDs, request.Header.Get("X-AI-Request-ID"))
	t.mu.Unlock()

	status := http.StatusInternalServerError
	body := `{"error":{"message":"retry","type":"server_error"}}`
	if attempt > 1 {
		status = http.StatusOK
		body = `{"id":"chatcmpl-test","object":"chat.completion","created":1,"model":"platform-default","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`
	}
	return &http.Response{
		StatusCode: status,
		Status:     http.StatusText(status),
		Header:     http.Header{"Content-Type": []string{"application/json"}},
		Body:       io.NopCloser(strings.NewReader(body)),
		Request:    request,
	}, nil
}

func TestPlatformOpenAIClientKeepsRequestIDAcrossRetries(t *testing.T) {
	transport := &requestIDRecordingTransport{}
	config := models.AIConfig{
		APIKey:        "platform-license",
		BaseURL:       "https://platform.example/v1",
		ModelName:     "platform-default",
		MaxRetryCount: 1,
		Platform:      true,
		HTTPClient:    &http.Client{Transport: transport},
	}
	client := newOpenAIClient(config)
	requestContext := WithPlatformAIRequestScope(context.Background(), "conversation:10:message:20:revision:30")
	params := openai.ChatCompletionNewParams{
		Model: shared.ChatModel("platform-default"),
		Messages: []openai.ChatCompletionMessageParamUnion{{
			OfUser: &openai.ChatCompletionUserMessageParam{
				Content: openai.ChatCompletionUserMessageParamContentUnion{OfString: openai.String("hello")},
			},
		}},
	}

	_, err := client.Chat.Completions.New(
		requestContext,
		params,
		platformRequestOptions(requestContext, config, "chat.completion")...,
	)
	if err != nil {
		t.Fatalf("chat completion after retry: %v", err)
	}
	_, err = client.Chat.Completions.New(
		requestContext,
		params,
		platformRequestOptions(requestContext, config, "chat.completion")...,
	)
	if err != nil {
		t.Fatalf("second logical chat completion: %v", err)
	}

	transport.mu.Lock()
	defer transport.mu.Unlock()
	if transport.attempts != 3 {
		t.Fatalf("attempts = %d, want 3", transport.attempts)
	}
	if transport.requestIDs[0] == "" || transport.requestIDs[0] != transport.requestIDs[1] {
		t.Fatalf("request IDs = %q, want one stable non-empty ID", transport.requestIDs)
	}
	if transport.requestIDs[2] == "" || transport.requestIDs[2] == transport.requestIDs[0] {
		t.Fatalf("request IDs = %q, want a fresh ID for the next logical call", transport.requestIDs)
	}
}

func TestPlatformRequestIDsAreStableAcrossRecoveryAndSeparatePurposeAndOrdinal(t *testing.T) {
	firstRun := WithPlatformAIRequestScope(context.Background(), "conversation:10:message:20:revision:30")
	firstQuery := nextPlatformAIRequestID(WithPlatformAIRequestPurpose(firstRun, "embedding.knowledge-query"), "embedding")
	secondQuery := nextPlatformAIRequestID(WithPlatformAIRequestPurpose(firstRun, "embedding.knowledge-query"), "embedding")
	chat := nextPlatformAIRequestID(firstRun, "chat.completion")

	recovered := WithPlatformAIRequestScope(context.Background(), "conversation:10:message:20:revision:30")
	recoveredFirstQuery := nextPlatformAIRequestID(WithPlatformAIRequestPurpose(recovered, "embedding.knowledge-query"), "embedding")
	recoveredSecondQuery := nextPlatformAIRequestID(WithPlatformAIRequestPurpose(recovered, "embedding.knowledge-query"), "embedding")
	recoveredChat := nextPlatformAIRequestID(recovered, "chat.completion")

	if firstQuery == secondQuery {
		t.Fatalf("embedding ordinals collided: %q", firstQuery)
	}
	if firstQuery == chat {
		t.Fatalf("embedding and chat purposes collided: %q", firstQuery)
	}
	if firstQuery != recoveredFirstQuery || secondQuery != recoveredSecondQuery || chat != recoveredChat {
		t.Fatalf("recovery IDs changed: first=(%q,%q,%q) recovered=(%q,%q,%q)", firstQuery, secondQuery, chat, recoveredFirstQuery, recoveredSecondQuery, recoveredChat)
	}
}

func TestCustomModelDoesNotReceivePlatformRequestOptions(t *testing.T) {
	config := models.AIConfig{Platform: false}
	ctx := WithPlatformAIRequestScope(context.Background(), "conversation:10:message:20:revision:30")
	if options := platformRequestOptions(ctx, config, "embedding"); len(options) != 0 {
		t.Fatalf("custom model options = %d, want 0", len(options))
	}
}