diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index fec8a3b4..89efeed9 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -243,6 +243,9 @@ func (e *XAIExecutor) executeCompactRequest(ctx context.Context, auth *cliproxya } prepared.body, _ = sjson.DeleteBytes(prepared.body, "stream") prepared.body, _ = sjson.DeleteBytes(prepared.body, "tools") + for _, field := range []string{"max_output_tokens", "temperature", "top_p", "top_k", "stop"} { + prepared.body, _ = sjson.DeleteBytes(prepared.body, field) + } prepared.body = xaiRemoveInputItemsByType(prepared.body, "compaction_trigger") reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth) @@ -1011,7 +1014,9 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox } originalPayload := bytes.Clone(originalPayloadSource) originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, stream) + originalTranslated = preserveXAIResponsesOutputControls(originalTranslated, originalPayload, from) body := sdktranslator.TranslateRequest(from, to, baseModel, bytes.Clone(req.Payload), stream) + body = preserveXAIResponsesOutputControls(body, req.Payload, from) var err error body, err = thinking.ApplyThinking(body, req.Model, from.String(), e.Identifier(), e.Identifier()) @@ -1437,7 +1442,35 @@ func xaiMetadataString(meta map[string]any, key string) string { } } +func preserveXAIResponsesOutputControls(body, source []byte, from sdktranslator.Format) []byte { + var maxOutputTokens gjson.Result + switch from { + case sdktranslator.FormatOpenAI: + maxOutputTokens = gjson.GetBytes(source, "max_completion_tokens") + if !maxOutputTokens.Exists() || maxOutputTokens.Type == gjson.Null { + maxOutputTokens = gjson.GetBytes(source, "max_tokens") + } + case sdktranslator.FormatOpenAIResponse: + maxOutputTokens = gjson.GetBytes(source, "max_output_tokens") + default: + return body + } + + if maxOutputTokens.Exists() && maxOutputTokens.Type != gjson.Null { + body, _ = sjson.SetRawBytes(body, "max_output_tokens", []byte(maxOutputTokens.Raw)) + } + for _, field := range []string{"temperature", "top_p", "top_k"} { + value := gjson.GetBytes(source, field) + if value.Exists() && value.Type != gjson.Null { + body, _ = sjson.SetRawBytes(body, field, []byte(value.Raw)) + } + } + return body +} + func sanitizeXAIResponsesBody(body []byte, model string) []byte { + // stop is supported by Chat Completions but not by xAI's Responses API. + body, _ = sjson.DeleteBytes(body, "stop") if !xaiSupportsReasoningEffort(model) { if gjson.GetBytes(body, "reasoning.effort").Exists() { log.Debugf("xai: stripping reasoning.effort for model %s (no thinking levels in model registry)", model) diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 1da22114..09ca38d2 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -814,6 +814,137 @@ func TestXAIExecutorPrepareDropsOrphanedToolChoiceBeforeXSearchInject(t *testing } } +func TestXAIExecutorPrepareResponsesRequestPreservesSupportedOutputControls(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + sourceFormat sdktranslator.Format + payload []byte + want map[string]string + absent []string + }{ + { + name: "Chat Completions prefers max_completion_tokens", + sourceFormat: sdktranslator.FormatOpenAI, + payload: []byte(`{ + "model":"grok-4.5", + "messages":[{"role":"user","content":"hello"}], + "max_completion_tokens":64, + "max_tokens":128, + "temperature":0, + "top_p":0.25, + "top_k":7, + "stop":["END"] + }`), + want: map[string]string{ + "max_output_tokens": "64", + "temperature": "0", + "top_p": "0.25", + "top_k": "7", + }, + absent: []string{"max_completion_tokens", "max_tokens", "stop"}, + }, + { + name: "Chat Completions falls back to max_tokens", + sourceFormat: sdktranslator.FormatOpenAI, + payload: []byte(`{ + "model":"grok-4.5", + "messages":[{"role":"user","content":"hello"}], + "max_completion_tokens":null, + "max_tokens":128 + }`), + want: map[string]string{ + "max_output_tokens": "128", + }, + absent: []string{"max_completion_tokens", "max_tokens", "temperature", "top_p", "top_k"}, + }, + { + name: "Responses preserves native controls", + sourceFormat: sdktranslator.FormatOpenAIResponse, + payload: []byte(`{ + "model":"grok-4.5", + "input":"hello", + "max_output_tokens":256, + "temperature":0.4, + "top_p":0.8, + "top_k":20, + "stop":["END"] + }`), + want: map[string]string{ + "max_output_tokens": "256", + "temperature": "0.4", + "top_p": "0.8", + "top_k": "20", + }, + absent: []string{"stop"}, + }, + { + name: "No controls remain absent", + sourceFormat: sdktranslator.FormatOpenAI, + payload: []byte(`{"model":"grok-4.5","messages":[{"role":"user","content":"hello"}]}`), + absent: []string{"max_output_tokens", "temperature", "top_p", "top_k", "stop"}, + }, + } + + exec := NewXAIExecutor(&config.Config{}) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + prepared, errPrepare := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{ + Model: "grok-4.5", + Payload: tt.payload, + }, cliproxyexecutor.Options{ + SourceFormat: tt.sourceFormat, + Stream: true, + }, true) + if errPrepare != nil { + t.Fatalf("prepareResponsesRequest() error = %v", errPrepare) + } + + for path, want := range tt.want { + if got := gjson.GetBytes(prepared.body, path).Raw; got != want { + t.Fatalf("%s = %s, want %s; body=%s", path, got, want, prepared.body) + } + } + for _, path := range tt.absent { + if gjson.GetBytes(prepared.body, path).Exists() { + t.Fatalf("%s should be absent; body=%s", path, prepared.body) + } + } + }) + } +} + +func TestXAIExecutorPrepareResponsesRequestDropsPayloadStopOverride(t *testing.T) { + t.Parallel() + + exec := NewXAIExecutor(&config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "grok-4.5"}}, + Params: map[string]any{"stop": []string{"END"}}, + }, + }, + }, + }) + prepared, errPrepare := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{ + Model: "grok-4.5", + Payload: []byte(`{"model":"grok-4.5","input":"hello"}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + }, true) + if errPrepare != nil { + t.Fatalf("prepareResponsesRequest() error = %v", errPrepare) + } + if gjson.GetBytes(prepared.body, "stop").Exists() { + t.Fatalf("stop should be removed after payload config; body=%s", prepared.body) + } +} + func TestXAIExecutorPrepareResponsesRequestAddsObjectTypeToRootUnionBranches(t *testing.T) { t.Parallel() @@ -1648,7 +1779,16 @@ func TestXAIExecutorCompactUsesCompactEndpoint(t *testing.T) { })) defer server.Close() - exec := NewXAIExecutor(&config.Config{}) + exec := NewXAIExecutor(&config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "grok-4.3"}}, + Params: map[string]any{"top_k": 10}, + }, + }, + }, + }) auth := &cliproxyauth.Auth{ Provider: "xai", Attributes: map[string]string{ @@ -1657,7 +1797,7 @@ func TestXAIExecutorCompactUsesCompactEndpoint(t *testing.T) { }, } - payload := []byte(`{"model":"grok-4.3","stream":true,"input":[{"type":"compaction","encrypted_content":""},{"role":"user","content":"hello"}]}`) + payload := []byte(`{"model":"grok-4.3","stream":true,"max_output_tokens":64,"temperature":0.3,"top_p":0.8,"stop":["END"],"input":[{"type":"compaction","encrypted_content":""},{"role":"user","content":"hello"}]}`) payload, _ = sjson.SetBytes(payload, "input.0.encrypted_content", validEncryptedContent) resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "grok-4.3", @@ -1679,8 +1819,10 @@ func TestXAIExecutorCompactUsesCompactEndpoint(t *testing.T) { if gotAccept != "application/json" { t.Fatalf("Accept = %q, want application/json", gotAccept) } - if gjson.GetBytes(gotBody, "stream").Exists() { - t.Fatalf("stream exists in compact body: %s", string(gotBody)) + for _, field := range []string{"stream", "max_output_tokens", "temperature", "top_p", "top_k", "stop"} { + if gjson.GetBytes(gotBody, field).Exists() { + t.Fatalf("%s exists in compact body: %s", field, string(gotBody)) + } } if got := gjson.GetBytes(gotBody, "input.0.encrypted_content").String(); got != validEncryptedContent { t.Fatalf("input.0.encrypted_content = %q, want valid sample; body=%s", got, string(gotBody))