iop/apps/edge/internal/openai/provider_model_rewrite_test.go

172 lines
8.8 KiB
Go

package openai
import (
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
edgeservice "iop/apps/edge/internal/service"
"iop/packages/go/config"
)
func TestGeminiChatProviderThoughtSignatureRoundTrip(t *testing.T) {
response := []byte(`{"id":"chat-1","model":"served","choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"call-1","type":"function","function":{"name":"glob","arguments":"{\"pattern\":\"*\"}"},"extra_content":{"google":{"thought_signature":"opaque-signature"}}}]},"finish_reason":"tool_calls"}]}`)
rewritten := rewriteProviderJSONResponse(response, "public-gemini", config.ProtocolToolCallWireGeminiChat)
var decoded map[string]any
if err := json.Unmarshal(rewritten, &decoded); err != nil {
t.Fatal(err)
}
choice := anySlice(decoded["choices"])[0].(map[string]any)
message := choice["message"].(map[string]any)
call := anySlice(message["tool_calls"])[0].(map[string]any)
encodedID := call["id"].(string)
if !strings.HasPrefix(encodedID, geminiThoughtSignatureToolIDPrefix) {
t.Fatalf("thought signature was not encoded in tool id: %+v", call)
}
if _, exists := call["extra_content"]; exists {
t.Fatalf("Gemini extension leaked after normalization: %+v", call)
}
if decoded["model"] != "public-gemini" {
t.Fatalf("model echo was not preserved through normalization: %+v", decoded)
}
request, err := normalizeGeminiChatProviderRequest([]byte(`{"model":"served","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"` + encodedID + `","type":"function","function":{"name":"glob","arguments":"{\"pattern\":\"*\"}"}}]},{"role":"tool","tool_call_id":"` + encodedID + `","content":"[]"}]}`))
if err != nil {
t.Fatal(err)
}
var replay map[string]any
if err := json.Unmarshal(request, &replay); err != nil {
t.Fatal(err)
}
messages := anySlice(replay["messages"])
replayedCall := anySlice(messages[0].(map[string]any)["tool_calls"])[0].(map[string]any)
if replayedCall["id"] != "call-1" || messages[1].(map[string]any)["tool_call_id"] != "call-1" {
t.Fatalf("original tool ids were not restored: %+v", messages)
}
extra := replayedCall["extra_content"].(map[string]any)
google := extra["google"].(map[string]any)
if google["thought_signature"] != "opaque-signature" {
t.Fatalf("thought signature was not restored: %+v", replayedCall)
}
}
func TestProviderThoughtSignatureNormalizationIsGeminiProfileOnly(t *testing.T) {
response := []byte(`{"model":"served","choices":[{"message":{"tool_calls":[{"id":"call-1","extra_content":{"google":{"thought_signature":"opaque"}}}]}}]}`)
if got := rewriteProviderJSONResponse(response, "public", ""); !strings.Contains(string(got), `"id":"call-1"`) || !strings.Contains(string(got), `"thought_signature":"opaque"`) {
t.Fatalf("non-Gemini response was normalized: %s", got)
}
profile, err := config.ResolveProtocolProfile("openai", "", config.BuiltInProtocolProfileCatalog())
if err != nil {
t.Fatal(err)
}
baseBody := []byte(`{"model":"served","messages":[{"role":"user","content":"hi"}]}`)
tunnel := edgeservice.SubmitProviderTunnelRequest{BuildBody: func(string) ([]byte, error) { return baseBody, nil }}
prepared, err := prepareProviderChatToolCallNormalization(tunnel, edgeservice.ProviderPoolCandidate{ProtocolProfile: &profile})
if err != nil {
t.Fatal(err)
}
got, err := prepared.BuildBody("served")
if err != nil || string(got) != string(baseBody) {
t.Fatalf("OpenAI profile body changed: %s err=%v", got, err)
}
}
func TestGeminiChatProviderStreamingThoughtSignatureNormalization(t *testing.T) {
line := []byte("data: {\"model\":\"served\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"call-1\",\"type\":\"function\",\"function\":{\"name\":\"glob\",\"arguments\":\"{}\"},\"extra_content\":{\"google\":{\"thought_signature\":\"opaque\"}}}]}}]}\n\n")
rewriter := newProviderModelRewriterWithToolCallWire(true, "public", config.ProtocolToolCallWireGeminiChat)
got := append(rewriter.AppendStream(line), rewriter.FlushStream()...)
if strings.Contains(string(got), "thought_signature") || !strings.Contains(string(got), geminiThoughtSignatureToolIDPrefix) {
t.Fatalf("streaming signature was not normalized: %s", got)
}
}
func TestGeminiChatProviderRejectsMalformedOpaqueToolID(t *testing.T) {
_, err := normalizeGeminiChatProviderRequest([]byte(`{"messages":[{"role":"assistant","tool_calls":[{"id":"iop_gts_not-base64"}]}]}`))
if err == nil {
t.Fatal("malformed IOP Gemini tool id must fail closed")
}
}
func TestGeminiChatProviderHTTPToolContinuationRoundTrip(t *testing.T) {
var providerRequests []map[string]any
provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, _ := io.ReadAll(r.Body)
var request map[string]any
if err := json.Unmarshal(body, &request); err != nil {
t.Errorf("decode provider request: %v", err)
w.WriteHeader(http.StatusBadRequest)
return
}
providerRequests = append(providerRequests, request)
w.Header().Set("Content-Type", "application/json")
if len(providerRequests) == 1 {
_, _ = w.Write([]byte(`{"id":"chat-1","model":"served-gemini","choices":[{"index":0,"message":{"role":"assistant","content":null,"tool_calls":[{"id":"call-1","type":"function","function":{"name":"glob","arguments":"{\"pattern\":\"*\"}"},"extra_content":{"google":{"thought_signature":"opaque-signature"}}}]},"finish_reason":"tool_calls"}]}`))
return
}
_, _ = w.Write([]byte(`{"id":"chat-2","model":"served-gemini","choices":[{"index":0,"message":{"role":"assistant","content":"done"},"finish_reason":"stop"}]}`))
}))
defer provider.Close()
profile, err := config.ResolveProtocolProfile("gemini", "", config.BuiltInProtocolProfileCatalog())
if err != nil {
t.Fatal(err)
}
fake := &providerFakeRunService{
poolDispatchPath: string(edgeservice.ProviderPoolPathTunnel),
tunnelProviderURL: provider.URL,
tunnelServedTarget: "served-gemini",
poolSelectedCandidate: edgeservice.ProviderPoolCandidate{
ActualModel: "served-gemini", ProviderID: "gemini-provider",
ExecutionPath: string(edgeservice.ProviderPoolPathTunnel), ProfileID: profile.ID,
ProfileDriver: string(profile.Driver), ProfileCapabilities: append([]string(nil), profile.Capabilities...),
ProtocolProfile: &profile,
},
}
srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: "gemini-route", Providers: map[string]string{"gemini-provider": "served-gemini"}}})
first := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gemini-route","messages":[{"role":"user","content":"list files"}],"tools":[{"type":"function","function":{"name":"glob","parameters":{"type":"object"}}}]}`))
firstResponse := httptest.NewRecorder()
srv.handleChatCompletions(firstResponse, first)
if firstResponse.Code != http.StatusOK {
t.Fatalf("first status=%d body=%s", firstResponse.Code, firstResponse.Body.String())
}
if got := fake.lastTunnelHandle.Dispatch().ProfileToolCallWire; got != config.ProtocolToolCallWireGeminiChat {
t.Fatalf("selected dispatch tool-call wire=%q", got)
}
var firstBody map[string]any
if err := json.Unmarshal(firstResponse.Body.Bytes(), &firstBody); err != nil {
t.Fatal(err)
}
firstChoice := anySlice(firstBody["choices"])[0].(map[string]any)
firstMessage := firstChoice["message"].(map[string]any)
encodedID := anySlice(firstMessage["tool_calls"])[0].(map[string]any)["id"].(string)
if !strings.HasPrefix(encodedID, geminiThoughtSignatureToolIDPrefix) {
t.Fatalf("caller did not receive opaque normalized id: %s", firstResponse.Body.String())
}
secondBody := `{"model":"gemini-route","messages":[{"role":"user","content":"list files"},{"role":"assistant","content":null,"tool_calls":[{"id":"` + encodedID + `","type":"function","function":{"name":"glob","arguments":"{\"pattern\":\"*\"}"}}]},{"role":"tool","tool_call_id":"` + encodedID + `","content":"[]"}]}`
second := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(secondBody))
secondResponse := httptest.NewRecorder()
srv.handleChatCompletions(secondResponse, second)
if secondResponse.Code != http.StatusOK || !strings.Contains(secondResponse.Body.String(), `"content":"done"`) {
t.Fatalf("second status=%d body=%s", secondResponse.Code, secondResponse.Body.String())
}
if len(providerRequests) != 2 {
t.Fatalf("provider requests=%d", len(providerRequests))
}
messages := anySlice(providerRequests[1]["messages"])
replayedCall := anySlice(messages[1].(map[string]any)["tool_calls"])[0].(map[string]any)
if replayedCall["id"] != "call-1" || messages[2].(map[string]any)["tool_call_id"] != "call-1" {
t.Fatalf("provider did not receive original tool ids: %+v", messages)
}
google := replayedCall["extra_content"].(map[string]any)["google"].(map[string]any)
if google["thought_signature"] != "opaque-signature" {
t.Fatalf("provider did not receive restored signature: %+v", replayedCall)
}
}