diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index a5984385..8c7bb207 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -38,8 +38,9 @@ import ( // ClaudeExecutor is a stateless executor for Anthropic Claude over the messages API. // If api_key is unavailable on auth, it falls back to legacy via ClientAdapter. type ClaudeExecutor struct { - cfg *config.Config - requestLogProvider string + cfg *config.Config + requestLogProvider string + upstreamModelNormalizer func(string) string } // claudeToolPrefix is empty to match real Claude Code behavior (no tool name prefix). @@ -161,6 +162,65 @@ func (e *ClaudeExecutor) upstreamRequestLogProvider() string { return e.Identifier() } +func (e *ClaudeExecutor) upstreamModel(baseModel string) string { + if e.upstreamModelNormalizer != nil { + return e.upstreamModelNormalizer(baseModel) + } + return baseModel +} + +func (e *ClaudeExecutor) restoreResponseModel(payload []byte, model string) []byte { + if e.upstreamModelNormalizer == nil || strings.TrimSpace(model) == "" { + return payload + } + return restoreClaudeResponseModel(payload, model) +} + +func restoreClaudeResponseModel(payload []byte, model string) []byte { + if updated, changed := setClaudeResponseModel(payload, model); changed { + return updated + } + + trimmed := bytes.TrimSpace(payload) + if !bytes.HasPrefix(trimmed, []byte("data:")) { + return payload + } + dataIndex := bytes.Index(payload, []byte("data:")) + if dataIndex < 0 { + return payload + } + rawJSON := bytes.TrimSpace(payload[dataIndex+len("data:"):]) + updated, changed := setClaudeResponseModel(rawJSON, model) + if !changed { + return payload + } + rebuilt := make([]byte, 0, dataIndex+len("data: ")+len(updated)) + rebuilt = append(rebuilt, payload[:dataIndex]...) + rebuilt = append(rebuilt, []byte("data: ")...) + rebuilt = append(rebuilt, updated...) + return rebuilt +} + +func setClaudeResponseModel(payload []byte, model string) ([]byte, bool) { + if !gjson.ValidBytes(payload) { + return payload, false + } + updated := payload + changed := false + for _, path := range []string{"model", "message.model"} { + if !gjson.GetBytes(updated, path).Exists() { + continue + } + next, errSet := sjson.SetBytes(updated, path, model) + if errSet != nil { + continue + } + updated = next + changed = true + } + return updated, changed +} + // PrepareRequest injects Claude credentials into the outgoing HTTP request. func (e *ClaudeExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error { if req == nil { @@ -208,6 +268,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} } baseModel := thinking.ParseSuffix(req.Model).ModelName + upstreamModel := e.upstreamModel(baseModel) apiKey, baseURL := claudeCreds(auth) if baseURL == "" { @@ -228,7 +289,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r originalPayload := originalPayloadSource originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, stream) body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, stream) - body, _ = sjson.SetBytes(body, "model", baseModel) + body, _ = sjson.SetBytes(body, "model", upstreamModel) body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) if err != nil { @@ -383,6 +444,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r reporter.Publish(ctx, helps.ParseClaudeUsage(data)) } data = restoreClaudeOAuthToolNamesFromResponse(data, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap) + data = e.restoreResponseModel(data, req.Model) var param any out := sdktranslator.TranslateNonStream( ctx, @@ -403,6 +465,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} } baseModel := thinking.ParseSuffix(req.Model).ModelName + upstreamModel := e.upstreamModel(baseModel) apiKey, baseURL := claudeCreds(auth) if baseURL == "" { @@ -421,7 +484,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A originalPayload := originalPayloadSource originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true) body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true) - body, _ = sjson.SetBytes(body, "model", baseModel) + body, _ = sjson.SetBytes(body, "model", upstreamModel) body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) if err != nil { @@ -580,6 +643,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A reporter.Publish(ctx, detail) } line = restoreClaudeOAuthToolNamesFromStreamLine(line, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap) + line = e.restoreResponseModel(line, req.Model) event.Write(line) event.WriteByte('\n') if len(bytes.TrimSpace(line)) == 0 && !flushEvent() { @@ -611,6 +675,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A reporter.Publish(ctx, detail) } line = restoreClaudeOAuthToolNamesFromStreamLine(line, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap) + line = e.restoreResponseModel(line, req.Model) chunks := sdktranslator.TranslateStream( ctx, to, @@ -701,6 +766,7 @@ func validateClaudeStreamingResponse(data []byte) error { func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { baseModel := thinking.ParseSuffix(req.Model).ModelName + upstreamModel := e.upstreamModel(baseModel) apiKey, baseURL := claudeCreds(auth) if baseURL == "" { @@ -713,7 +779,7 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut // Use streaming translation to preserve function calling, except for claude. stream := from != to body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, stream) - body, _ = sjson.SetBytes(body, "model", baseModel) + body, _ = sjson.SetBytes(body, "model", upstreamModel) if rebuildMidSystemMessageEnabled(e.cfg, auth) { body = rebuildMidSystemMessagesToTopLevel(body) } diff --git a/internal/runtime/executor/kimi_executor.go b/internal/runtime/executor/kimi_executor.go index 13966b7e..53d19f81 100644 --- a/internal/runtime/executor/kimi_executor.go +++ b/internal/runtime/executor/kimi_executor.go @@ -36,8 +36,12 @@ type KimiExecutor struct { // NewKimiExecutor creates a new Kimi executor. func NewKimiExecutor(cfg *config.Config) *KimiExecutor { return &KimiExecutor{ - ClaudeExecutor: ClaudeExecutor{cfg: cfg, requestLogProvider: "kimi"}, - cfg: cfg, + ClaudeExecutor: ClaudeExecutor{ + cfg: cfg, + requestLogProvider: "kimi", + upstreamModelNormalizer: normalizeKimiUpstreamModel, + }, + cfg: cfg, } } @@ -102,8 +106,8 @@ func (e *KimiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, false) body := sdktranslator.TranslateRequest(from, to, baseModel, bytes.Clone(req.Payload), false) - // Strip kimi- prefix for upstream API - upstreamModel := stripKimiPrefix(baseModel) + // Strip kimi- prefix and any [1m] suffix for upstream API + upstreamModel := normalizeKimiUpstreamModel(baseModel) body, err = sjson.SetBytes(body, "model", upstreamModel) if err != nil { return resp, fmt.Errorf("kimi executor: failed to set model in payload: %w", err) @@ -211,8 +215,8 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true) body := sdktranslator.TranslateRequest(from, to, baseModel, bytes.Clone(req.Payload), true) - // Strip kimi- prefix for upstream API - upstreamModel := stripKimiPrefix(baseModel) + // Strip kimi- prefix and any [1m] suffix for upstream API + upstreamModel := normalizeKimiUpstreamModel(baseModel) body, err = sjson.SetBytes(body, "model", upstreamModel) if err != nil { return nil, fmt.Errorf("kimi executor: failed to set model in payload: %w", err) @@ -759,3 +763,21 @@ func stripKimiPrefix(model string) string { } return model } + +// normalizeKimiUpstreamModel returns the canonical upstream model ID for Kimi. +// It strips the CLIProxyAPI "kimi-" prefix and any Claude Code "[1m]" context +// suffix while preserving a trailing thinking suffix (e.g. "(1024)"), so that +// the upstream API receives IDs such as "k3(1024)" instead of "kimi-k3[1m](1024)". +func normalizeKimiUpstreamModel(model string) string { + model = strings.TrimSpace(model) + parsed := thinking.ParseSuffix(model) + base := parsed.ModelName + if strings.HasSuffix(strings.ToLower(base), "[1m]") { + base = base[:len(base)-len("[1m]")] + } + normalized := strings.ToLower(stripKimiPrefix(strings.TrimSpace(base))) + if parsed.HasSuffix { + return normalized + "(" + parsed.RawSuffix + ")" + } + return normalized +} diff --git a/internal/runtime/executor/kimi_executor_test.go b/internal/runtime/executor/kimi_executor_test.go index 824fb192..3caa862b 100644 --- a/internal/runtime/executor/kimi_executor_test.go +++ b/internal/runtime/executor/kimi_executor_test.go @@ -28,6 +28,92 @@ func TestNewKimiExecutorInitializesDelegatedClaudeConfig(t *testing.T) { } } +func TestKimiExecutorClaudeRequestPreservesInternalModelSemantics(t *testing.T) { + var upstreamBody []byte + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", kimiRoundTripperFunc(func(req *http.Request) (*http.Response, error) { + var errRead error + upstreamBody, errRead = io.ReadAll(req.Body) + if errRead != nil { + return nil, errRead + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader( + `{"id":"msg_test","type":"message","role":"assistant","model":"k2.5","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}`, + )), + }, nil + })) + + executor := NewKimiExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Attributes: map[string]string{}, + Metadata: map[string]any{"access_token": "test-token"}, + } + const model = "kimi-k2.5(max)" + payload := []byte(`{"model":"kimi-k2.5(max)","max_tokens":32,"messages":[{"role":"user","content":"hello"}]}`) + response, err := executor.Execute(ctx, auth, cliproxyexecutor.Request{ + Model: model, + Payload: payload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + OriginalRequest: payload, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + if got := gjson.GetBytes(upstreamBody, "model").String(); got != "k2.5" { + t.Fatalf("upstream model = %q, want k2.5", got) + } + if got := gjson.GetBytes(upstreamBody, "output_config.effort").String(); got != "high" { + t.Fatalf("upstream output_config.effort = %q, want high", got) + } + if got := gjson.GetBytes(response.Payload, "model").String(); got != model { + t.Fatalf("response model = %q, want %q", got, model) + } +} + +func TestKimiExecutorCountTokensUsesCanonicalUpstreamModel(t *testing.T) { + var upstreamRequest *http.Request + var upstreamBody []byte + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", kimiRoundTripperFunc(func(req *http.Request) (*http.Response, error) { + upstreamRequest = req.Clone(req.Context()) + var errRead error + upstreamBody, errRead = io.ReadAll(req.Body) + if errRead != nil { + return nil, errRead + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"input_tokens":42}`)), + }, nil + })) + + executor := NewKimiExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Attributes: map[string]string{}, + Metadata: map[string]any{"access_token": "test-token"}, + } + payload := []byte(`{"model":"kimi-k3[1m](high)","messages":[{"role":"user","content":"hello"}]}`) + _, err := executor.CountTokens(ctx, auth, cliproxyexecutor.Request{ + Model: "kimi-k3[1m](high)", + Payload: payload, + }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) + if err != nil { + t.Fatalf("CountTokens() error = %v", err) + } + if upstreamRequest == nil { + t.Fatal("upstream request was not captured") + } + if got := upstreamRequest.URL.String(); got != "https://api.kimi.com/coding/v1/messages/count_tokens?beta=true" { + t.Fatalf("upstream URL = %q, want Kimi count tokens endpoint", got) + } + if got := gjson.GetBytes(upstreamBody, "model").String(); got != "k3" { + t.Fatalf("upstream model = %q, want k3", got) + } +} + func TestKimiExecutorClaudeStreamForwardsAnthropicBetaAndLogsUpstream(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() @@ -44,7 +130,7 @@ func TestKimiExecutorClaudeStreamForwardsAnthropicBetaAndLogsUpstream(t *testing Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader( "event: message_start\n" + - `data: {"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"kimi-k3","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":0}}}` + "\n\n" + + `data: {"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"k3","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":0}}}` + "\n\n" + "event: message_stop\n" + `data: {"type":"message_stop"}` + "\n\n", )), @@ -73,10 +159,15 @@ func TestKimiExecutorClaudeStreamForwardsAnthropicBetaAndLogsUpstream(t *testing if err != nil { t.Fatalf("ExecuteStream() error = %v", err) } + var output strings.Builder for chunk := range result.Chunks { if chunk.Err != nil { t.Fatalf("stream chunk error = %v", chunk.Err) } + output.Write(chunk.Payload) + } + if !strings.Contains(output.String(), `"model":"kimi-k3"`) { + t.Fatalf("stream output = %q, want requested model kimi-k3", output.String()) } if upstreamRequest == nil { t.Fatal("upstream request was not captured") @@ -102,7 +193,7 @@ func TestKimiExecutorClaudeStreamForwardsAnthropicBetaAndLogsUpstream(t *testing "Upstream URL: https://api.kimi.com/coding/v1/messages?beta=true", "Auth: provider=kimi", "Anthropic-Beta: " + upstreamBetas, - `"model":"kimi-k3"`, + `"model":"k3"`, } { if !strings.Contains(apiRequestText, want) { t.Fatalf("API_REQUEST = %q, want %q", apiRequestText, want) @@ -395,3 +486,29 @@ func TestNormalizeKimiToolMessageLinks_PreservesAssistantWithToolLinkOrReasoning t.Fatalf("messages.3.content.0.text = %q, want %q", got, " visible ") } } + +func TestNormalizeKimiUpstreamModel(t *testing.T) { + cases := []struct { + in string + want string + }{ + {"kimi-k3[1m]", "k3"}, + {"kimi-k3", "k3"}, + {"Kimi-K3[1M]", "k3"}, + {"k3[1m]", "k3"}, + {"k3", "k3"}, + {"kimi-k2.6", "k2.6"}, + {"kimi-k2.6[1m]", "k2.6"}, + {"kimi-k3(1024)", "k3(1024)"}, + {"kimi-k3[1m](1024)", "k3(1024)"}, + {"kimi-k2.6(high)", "k2.6(high)"}, + {"kimi-k2.6[1m](high)", "k2.6(high)"}, + } + + for _, c := range cases { + got := normalizeKimiUpstreamModel(c.in) + if got != c.want { + t.Errorf("normalizeKimiUpstreamModel(%q) = %q, want %q", c.in, got, c.want) + } + } +}