443 lines
19 KiB
Go
443 lines
19 KiB
Go
package openai
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"go.uber.org/zap"
|
|
"go.uber.org/zap/zaptest/observer"
|
|
|
|
edgeservice "iop/apps/edge/internal/service"
|
|
"iop/packages/go/config"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
func TestGeminiIngressAuthenticatesAndStreamsThroughChatRoute(t *testing.T) {
|
|
fake := &fakeRunService{events: bufferedRunEvents(
|
|
&iop.RunEvent{Type: "delta", Delta: "hello"},
|
|
&iop.RunEvent{Type: "complete", Usage: &iop.Usage{InputTokens: 3, OutputTokens: 2}, Metadata: map[string]string{"finish_reason": "length"}},
|
|
)}
|
|
srv := NewServer(config.EdgeOpenAIConf{BearerToken: "iop-principal", Adapter: "ollama", Target: "provider-model"}, fake, nil)
|
|
body := `{
|
|
"systemInstruction":{"role":"user","parts":[{"text":"be concise"}]},
|
|
"contents":[{"role":"user","parts":[{"text":"say hello"}]}],
|
|
"generationConfig":{"candidateCount":1,"maxOutputTokens":16,"temperature":0.2,"topK":8,"topP":0.9,"thinkingConfig":{"includeThoughts":true,"thinkingBudget":-1}}
|
|
}`
|
|
req := httptest.NewRequest(http.MethodPost, "/gemini/gemini-direct/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", strings.NewReader(body))
|
|
req.Header.Set("X-Goog-Api-Key", "iop-principal")
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK || w.Header().Get("Content-Type") != "text/event-stream" {
|
|
t.Fatalf("status/header: %d %q body=%s", w.Code, w.Header().Get("Content-Type"), w.Body.String())
|
|
}
|
|
if fake.req.ModelGroupKey != "gemini-direct" || fake.req.Target != "provider-model" {
|
|
t.Fatalf("route bypassed Chat admission: %+v", fake.req)
|
|
}
|
|
if !strings.Contains(fake.req.Prompt, "system: be concise") || !strings.Contains(fake.req.Prompt, "user: say hello") {
|
|
t.Fatalf("prompt conversion mismatch: %q", fake.req.Prompt)
|
|
}
|
|
options := fake.req.Input["options"].(map[string]any)
|
|
if options["max_tokens"] != 16 {
|
|
t.Fatalf("generation config mismatch: input=%+v", fake.req.Input)
|
|
}
|
|
if _, exists := options["top_k"]; exists {
|
|
t.Fatalf("deprecated Gemini sampling option reached Chat upstream: input=%+v", fake.req.Input)
|
|
}
|
|
response := w.Body.String()
|
|
for _, want := range []string{`"text":"hello"`, `"finishReason":"MAX_TOKENS"`} {
|
|
if !strings.Contains(response, want) {
|
|
t.Fatalf("missing %s in %s", want, response)
|
|
}
|
|
}
|
|
if strings.Contains(response, "iop-principal") {
|
|
t.Fatal("principal token leaked to response")
|
|
}
|
|
}
|
|
|
|
func TestGeminiStreamPathAcceptsOfficialCallerModelLabel(t *testing.T) {
|
|
req := httptest.NewRequest(
|
|
http.MethodPost,
|
|
"/gemini/gemini-direct/v1beta/models/Gemini%203.6%20Flash:streamGenerateContent?alt=sse",
|
|
nil,
|
|
)
|
|
routeID, callerModel, err := parseGeminiStreamPath(req)
|
|
if err != nil {
|
|
t.Fatalf("parseGeminiStreamPath: %v", err)
|
|
}
|
|
if routeID != "gemini-direct" || callerModel != "Gemini 3.6 Flash" {
|
|
t.Fatalf("path = %q/%q", routeID, callerModel)
|
|
}
|
|
|
|
for _, path := range []string{
|
|
"/gemini/gemini-direct/v1beta/models/%20Gemini:streamGenerateContent?alt=sse",
|
|
"/gemini/gemini-direct/v1beta/models/Gemini%20:streamGenerateContent?alt=sse",
|
|
} {
|
|
bad := httptest.NewRequest(http.MethodPost, path, nil)
|
|
if _, _, err := parseGeminiStreamPath(bad); err == nil {
|
|
t.Fatalf("expected invalid caller model path: %s", path)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestGeminiIngressRejectsAuthenticationAndShapeBeforeDispatch(t *testing.T) {
|
|
base := `{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`
|
|
for _, tc := range []struct {
|
|
name, path, body, bearer, key string
|
|
status int
|
|
}{
|
|
{"missing key", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", base, "", "", 401},
|
|
{"conflicting auth", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", base, "Bearer other", "iop-principal", 401},
|
|
{"wrong query", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent", base, "", "iop-principal", 400},
|
|
{"duplicate member", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", `{"contents":[],"contents":[]}`, "", "iop-principal", 400},
|
|
{"two candidates", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", `{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"candidateCount":2}}`, "", "iop-principal", 400},
|
|
{"invalid thinking budget", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", `{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":-2}}}`, "", "iop-principal", 400},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
fake := &fakeRunService{}
|
|
srv := NewServer(config.EdgeOpenAIConf{BearerToken: "iop-principal", Adapter: "ollama"}, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, tc.path, strings.NewReader(tc.body))
|
|
if tc.bearer != "" {
|
|
req.Header.Set("Authorization", tc.bearer)
|
|
}
|
|
if tc.key != "" {
|
|
req.Header.Set("X-Goog-Api-Key", tc.key)
|
|
}
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
if w.Code != tc.status {
|
|
t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
|
|
}
|
|
var failure geminiErrorResponse
|
|
if json.Unmarshal(w.Body.Bytes(), &failure) != nil || failure.Error.Code != tc.status {
|
|
t.Fatalf("not a Gemini error: %s", w.Body.String())
|
|
}
|
|
if fake.req.ModelGroupKey != "" {
|
|
t.Fatalf("unexpected dispatch: %+v", fake.req)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgePreservesToolsAndThoughtSignature(t *testing.T) {
|
|
body := []byte(`{
|
|
"contents":[
|
|
{"role":"model","parts":[{"text":"thinking","thought":true},{"functionCall":{"name":"lookup","args":{"q":"x"}},"thoughtSignature":"opaque"}]},
|
|
{"role":"user","parts":[{"functionResponse":{"name":"lookup","response":{"value":1}}}]},
|
|
{"role":"user","parts":[{"text":"continue"}]}
|
|
],
|
|
"tools":[{"functionDeclarations":[{"name":"lookup","description":"find","parametersJsonSchema":{"type":"object"}}]}],
|
|
"toolConfig":{"functionCallingConfig":{"mode":"ANY"}}
|
|
}`)
|
|
converted, err := prepareGeminiChatBridge(body, "preset-gemini-hybrid")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
text := string(converted)
|
|
for _, want := range []string{`"model":"preset-gemini-hybrid"`, `"reasoning_content":"thinking"`, `"thought_signature":"opaque"`, `"tool_call_id":"gemini_call_0_1"`, `"tool_choice":"required"`} {
|
|
if !strings.Contains(text, want) {
|
|
t.Fatalf("missing %s in %s", want, text)
|
|
}
|
|
}
|
|
if strings.Contains(text, `"tool_name"`) {
|
|
t.Fatalf("nonstandard tool result field reached Chat bridge: %s", text)
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeMatchesExplicitToolCallIDsOutOfOrder(t *testing.T) {
|
|
converted, err := prepareGeminiChatBridge([]byte(`{
|
|
"contents":[
|
|
{"role":"model","parts":[
|
|
{"functionCall":{"id":"call_first","name":"lookup","args":{"q":"first"}}},
|
|
{"functionCall":{"id":"call_second","name":"lookup","args":{"q":"second"}}}
|
|
]},
|
|
{"role":"user","parts":[
|
|
{"functionResponse":{"id":"call_second","name":"lookup","response":{"value":2}}},
|
|
{"functionResponse":{"id":"call_first","name":"lookup","response":{"value":1}}}
|
|
]}
|
|
]
|
|
}`), "gemini-route")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var body struct {
|
|
Messages []map[string]any `json:"messages"`
|
|
}
|
|
if err := json.Unmarshal(converted, &body); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(body.Messages) != 3 {
|
|
t.Fatalf("messages=%+v", body.Messages)
|
|
}
|
|
if body.Messages[1]["tool_call_id"] != "call_second" || body.Messages[2]["tool_call_id"] != "call_first" {
|
|
t.Fatalf("explicit response order was not preserved: %+v", body.Messages)
|
|
}
|
|
for _, message := range body.Messages[1:] {
|
|
if _, exists := message["tool_name"]; exists {
|
|
t.Fatalf("nonstandard tool_name field survived: %+v", message)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeKeepsIDLessFIFOCompatibility(t *testing.T) {
|
|
converted, err := prepareGeminiChatBridge([]byte(`{
|
|
"contents":[
|
|
{"role":"model","parts":[
|
|
{"functionCall":{"name":"lookup","args":{"q":"first"}}},
|
|
{"functionCall":{"name":"lookup","args":{"q":"second"}}}
|
|
]},
|
|
{"role":"user","parts":[
|
|
{"functionResponse":{"name":"lookup","response":{"value":1}}},
|
|
{"functionResponse":{"name":"lookup","response":{"value":2}}}
|
|
]}
|
|
]
|
|
}`), "gemini-route")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var body struct {
|
|
Messages []map[string]any `json:"messages"`
|
|
}
|
|
if err := json.Unmarshal(converted, &body); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(body.Messages) != 3 || body.Messages[1]["tool_call_id"] != "gemini_call_0_0" || body.Messages[2]["tool_call_id"] != "gemini_call_0_1" {
|
|
t.Fatalf("ID-less FIFO changed: %+v", body.Messages)
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeAcceptsOfficialAgyModelRoleFunctionResponse(t *testing.T) {
|
|
converted, err := prepareGeminiChatBridge([]byte(`{
|
|
"contents":[
|
|
{"role":"user","parts":[{"text":"create index.html"}]},
|
|
{"role":"model","parts":[
|
|
{"text":"I will inspect the workspace."},
|
|
{"functionCall":{"name":"list_dir","args":{"path":"."}},"thoughtSignature":"opaque"}
|
|
]},
|
|
{"role":"model","parts":[
|
|
{"functionResponse":{"name":"list_dir","response":{"entries":[]}}}
|
|
]}
|
|
]
|
|
}`), "gemini-3.6-flash")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var body struct {
|
|
Messages []map[string]any `json:"messages"`
|
|
}
|
|
if err := json.Unmarshal(converted, &body); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(body.Messages) != 3 {
|
|
t.Fatalf("messages=%+v", body.Messages)
|
|
}
|
|
if body.Messages[1]["role"] != "assistant" || body.Messages[2]["role"] != "tool" {
|
|
t.Fatalf("agy tool continuation roles changed: %+v", body.Messages)
|
|
}
|
|
if body.Messages[2]["tool_call_id"] != "gemini_call_1_1" {
|
|
t.Fatalf("agy tool response did not match its call: %+v", body.Messages[2])
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeRejectsModelRoleFunctionResponseMixedWithAssistantContent(t *testing.T) {
|
|
_, err := prepareGeminiChatBridge([]byte(`{
|
|
"contents":[
|
|
{"role":"model","parts":[{"functionCall":{"name":"lookup","args":{}}}]},
|
|
{"role":"model","parts":[
|
|
{"text":"mixed"},
|
|
{"functionResponse":{"name":"lookup","response":{}}}
|
|
]}
|
|
]
|
|
}`), "gemini-3.6-flash")
|
|
if err == nil {
|
|
t.Fatal("model functionResponse mixed with assistant content must fail")
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeRejectsInvalidToolCallIdentity(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
body string
|
|
}{
|
|
{
|
|
name: "duplicate id",
|
|
body: `{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call_1","name":"a","args":{}}},{"functionCall":{"id":"call_1","name":"b","args":{}}}]}]}`,
|
|
},
|
|
{
|
|
name: "mismatched name",
|
|
body: `{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call_1","name":"a","args":{}}}]},{"role":"user","parts":[{"functionResponse":{"id":"call_1","name":"b","response":{}}}]}]}`,
|
|
},
|
|
{
|
|
name: "missing response id for explicit call",
|
|
body: `{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call_1","name":"a","args":{}}}]},{"role":"user","parts":[{"functionResponse":{"name":"a","response":{}}}]}]}`,
|
|
},
|
|
{
|
|
name: "invalid id",
|
|
body: `{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call id","name":"a","args":{}}}]}]}`,
|
|
},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
if _, err := prepareGeminiChatBridge([]byte(tc.body), "gemini-route"); err == nil {
|
|
t.Fatal("invalid identity must fail")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeUsesProviderNativeThinkingEnvelope(t *testing.T) {
|
|
converted, err := prepareGeminiChatBridge([]byte(`{
|
|
"contents":[{"role":"user","parts":[{"text":"hello"}]}],
|
|
"generationConfig":{"temperature":1,"topK":50,"topP":1,"thinkingConfig":{"includeThoughts":true,"thinkingBudget":-1}}
|
|
}`), "gemini-3.6-flash")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var body map[string]any
|
|
if json.Unmarshal(converted, &body) != nil {
|
|
t.Fatal("converted body is invalid")
|
|
}
|
|
for _, forbidden := range []string{"temperature", "top_k", "top_p", "think", "include_reasoning", "thinking_token_budget"} {
|
|
if _, exists := body[forbidden]; exists {
|
|
t.Fatalf("unsupported field %q in converted body: %s", forbidden, converted)
|
|
}
|
|
}
|
|
extra := body["extra_body"].(map[string]any)
|
|
google := extra["google"].(map[string]any)
|
|
thinking := google["thinking_config"].(map[string]any)
|
|
if thinking["include_thoughts"] != true || thinking["thinking_budget"] != float64(-1) {
|
|
t.Fatalf("thinking envelope mismatch: %+v", thinking)
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeAcceptsOfficialAgyPlannerStructuredOutput(t *testing.T) {
|
|
converted, err := prepareGeminiChatBridge([]byte(`{
|
|
"contents":[{"role":"user","parts":[{"text":"plan"}]}],
|
|
"generationConfig":{
|
|
"responseMimeType":"application/json",
|
|
"responseJsonSchema":{"type":"OBJECT","properties":{"steps":{"type":"ARRAY"}},"required":["steps"]}
|
|
},
|
|
"tools":[{"functionDeclarations":[{
|
|
"name":"read_file","parameters":{"type":"OBJECT","properties":{"path":{"type":"STRING"}}},
|
|
"responseJsonSchema":{"type":"object"}
|
|
}]}]
|
|
}`), "gemini-3.6-flash")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var body map[string]any
|
|
if json.Unmarshal(converted, &body) != nil {
|
|
t.Fatal("converted body is invalid")
|
|
}
|
|
format := body["response_format"].(map[string]any)
|
|
jsonSchema := format["json_schema"].(map[string]any)
|
|
schema := jsonSchema["schema"].(map[string]any)
|
|
if format["type"] != "json_schema" || jsonSchema["name"] != "agy_response" || jsonSchema["strict"] != true || schema["type"] != "object" {
|
|
t.Fatalf("structured output mismatch: %+v", format)
|
|
}
|
|
tools := body["tools"].([]any)
|
|
function := tools[0].(map[string]any)["function"].(map[string]any)
|
|
if function["name"] != "read_file" {
|
|
t.Fatalf("function conversion mismatch: %+v", function)
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeRejectsConflictingStructuredSchemas(t *testing.T) {
|
|
for _, body := range []string{
|
|
`{"contents":[{"role":"user","parts":[{"text":"x"}]}],"generationConfig":{"responseMimeType":"application/json","responseSchema":{"type":"object"},"responseJsonSchema":{"type":"object"}}}`,
|
|
`{"contents":[{"role":"user","parts":[{"text":"x"}]}],"tools":[{"functionDeclarations":[{"name":"f","parameters":{"type":"object"},"parametersJsonSchema":{"type":"object"}}]}]}`,
|
|
} {
|
|
if _, err := prepareGeminiChatBridge([]byte(body), "gemini-3.6-flash"); err == nil {
|
|
t.Fatal("conflicting schema alternatives must fail")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestGeminiStreamBridgeEmitsBoundedToolOnce(t *testing.T) {
|
|
w := httptest.NewRecorder()
|
|
stream := newGeminiBridgeStream(w, "gemini-3.6-flash")
|
|
input := "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_provider_1\",\"function\":{\"name\":\"lookup\",\"arguments\":\"{\\\"q\\\":\"}}]}}]}\n\n" +
|
|
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"x\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n" +
|
|
"data: [DONE]\n\n"
|
|
if err := stream.Feed([]byte(input)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if strings.Count(w.Body.String(), `"functionCall"`) != 1 || !strings.Contains(w.Body.String(), `"finishReason":"STOP"`) {
|
|
t.Fatalf("tool terminal mismatch: %s", w.Body.String())
|
|
}
|
|
if !strings.Contains(w.Body.String(), `"id":"call_provider_1"`) {
|
|
t.Fatalf("provider tool call id was not preserved: %s", w.Body.String())
|
|
}
|
|
usageWriter := httptest.NewRecorder()
|
|
usageStream := newGeminiBridgeStream(usageWriter, "m")
|
|
if err := usageStream.Feed([]byte("data: {\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":2,\"total_tokens\":5,\"prompt_tokens_details\":{\"cached_tokens\":1},\"completion_tokens_details\":{\"reasoning_tokens\":1}}}\n\ndata: [DONE]\n\n")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for _, want := range []string{`"promptTokenCount":3`, `"candidatesTokenCount":2`, `"cachedContentTokenCount":1`, `"thoughtsTokenCount":1`, `"totalTokenCount":5`} {
|
|
if !strings.Contains(usageWriter.Body.String(), want) {
|
|
t.Fatalf("missing usage %s in %s", want, usageWriter.Body.String())
|
|
}
|
|
}
|
|
oversize := newGeminiBridgeStream(httptest.NewRecorder(), "m")
|
|
state := &geminiBridgeToolState{name: "lookup"}
|
|
state.arguments.WriteString(strings.Repeat("x", geminiToolArgumentLimit))
|
|
oversize.tools[0] = state
|
|
chunk := `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"x"}}]}}]}` + "\n\n"
|
|
if err := oversize.Feed([]byte(chunk)); err == nil {
|
|
t.Fatal("oversize tool arguments must fail")
|
|
}
|
|
}
|
|
|
|
func TestGeminiRejectionObservationIsClassificationOnly(t *testing.T) {
|
|
core, observed := observer.New(zap.InfoLevel)
|
|
secret := "secret-fixture-must-not-be-logged"
|
|
|
|
authServer := NewServer(config.EdgeOpenAIConf{BearerToken: "principal"}, &fakeRunService{}, zap.New(core))
|
|
authReq := httptest.NewRequest(http.MethodPost, "/gemini/r/v1beta/models/m:streamGenerateContent?alt=sse", strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"`+secret+`"}]}]}`))
|
|
authServer.routes().ServeHTTP(httptest.NewRecorder(), authReq)
|
|
|
|
managedServer := managedCredentialMigrationServer(t, &providerFakeRunService{})
|
|
managedServer.logger = zap.New(core)
|
|
managedReq := httptest.NewRequest(http.MethodPost, "/gemini/public-route/v1beta/models/m:streamGenerateContent?alt=sse", strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`))
|
|
managedReq.Header.Set("X-Goog-Api-Key", "managed-iop-token")
|
|
managedReq.Header.Set(legacyProviderCredentialHeader, secret)
|
|
managedServer.routes().ServeHTTP(httptest.NewRecorder(), managedReq)
|
|
|
|
candidate := anthropicTestCandidate(t, "openai")
|
|
candidate.ActualModel = "served-chat"
|
|
providerBody := []byte(`{"error":{"message":"` + secret + `"}}`)
|
|
fake := &providerFakeRunService{
|
|
poolDispatchPath: string(edgeservice.ProviderPoolPathTunnel),
|
|
poolSelectedCandidate: candidate,
|
|
tunnelFrames: anthropicTunnelFrames(http.StatusBadRequest, "application/json", providerBody),
|
|
}
|
|
providerServer := NewServer(config.EdgeOpenAIConf{BearerToken: "principal"}, fake, zap.New(core))
|
|
providerServer.SetModelCatalog([]config.ModelCatalogEntry{{ID: "r", Providers: map[string]string{"chat": "served-chat"}}})
|
|
providerReq := httptest.NewRequest(http.MethodPost, "/gemini/r/v1beta/models/m:streamGenerateContent?alt=sse", strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`))
|
|
providerReq.Header.Set("X-Goog-Api-Key", "principal")
|
|
providerServer.routes().ServeHTTP(httptest.NewRecorder(), providerReq)
|
|
|
|
localFailureServer := NewServer(config.EdgeOpenAIConf{BearerToken: "principal", Adapter: "ollama", Target: "provider-model"}, &fakeRunService{submitErrAfter: 1}, zap.New(core))
|
|
localFailureReq := httptest.NewRequest(http.MethodPost, "/gemini/r/v1beta/models/m:streamGenerateContent?alt=sse", strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"`+secret+`"}]}]}`))
|
|
localFailureReq.Header.Set("X-Goog-Api-Key", "principal")
|
|
localFailureServer.routes().ServeHTTP(httptest.NewRecorder(), localFailureReq)
|
|
|
|
entries := observed.FilterMessage(geminiRejectionLogMessage).All()
|
|
if len(entries) != 3 {
|
|
t.Fatalf("rejection observations=%d logs=%+v", len(entries), observed.All())
|
|
}
|
|
wantClasses := []string{string(geminiRejectionPreIngress), string(geminiRejectionPreIngress), string(geminiRejectionProviderHTTP)}
|
|
wantStatuses := []int64{http.StatusUnauthorized, http.StatusBadRequest, http.StatusBadRequest}
|
|
for index, entry := range entries {
|
|
context := entry.ContextMap()
|
|
if len(context) != 4 || context["surface"] != "gemini" || context["bridge"] != "chat" || context["rejection_class"] != wantClasses[index] || context["http_status"] != wantStatuses[index] {
|
|
t.Fatalf("unsafe or incomplete observation: %+v", context)
|
|
}
|
|
if strings.Contains(fmt.Sprint(context), secret) {
|
|
t.Fatalf("secret leaked to observation: %+v", context)
|
|
}
|
|
}
|
|
}
|