Merge pull request #4419 from mikewong23571/fix/kimi-upstream-model-normalization

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)
 	}

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
+}

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)
+		}
+	}
+}