805 lines · 25.6 KB · Go
RawDownload
package executor import (	"bufio"	"bytes"	"context"	"fmt"	"io"	"net/http"	"os"	"path/filepath"	"runtime"	"strings"	"time" 	kimiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/kimi"	"github.com/router-for-me/CLIProxyAPI/v7/internal/buildinfo"	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"	"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"	"github.com/router-for-me/CLIProxyAPI/v7/internal/util"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"	log "github.com/sirupsen/logrus"	"github.com/tidwall/gjson"	"github.com/tidwall/sjson") // KimiExecutor is a stateless executor for Kimi API using OpenAI-compatible chat completions.type KimiExecutor struct {	ClaudeExecutor	cfg *config.Config} // NewKimiExecutor creates a new Kimi executor.func NewKimiExecutor(cfg *config.Config) *KimiExecutor {	return &KimiExecutor{		ClaudeExecutor: ClaudeExecutor{			cfg:                     cfg,			requestLogProvider:      "kimi",			upstreamModelNormalizer: normalizeKimiUpstreamModel,		},		cfg: cfg,	}} // Identifier returns the executor identifier.func (e *KimiExecutor) Identifier() string { return "kimi" } // PrepareRequest injects Kimi credentials into the outgoing HTTP request.func (e *KimiExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error {	if req == nil {		return nil	}	token := kimiCreds(auth)	if strings.TrimSpace(token) != "" {		req.Header.Set("Authorization", "Bearer "+token)	}	var attrs map[string]string	if auth != nil {		attrs = auth.Attributes	}	util.ApplyCustomHeadersFromAttrs(req, attrs)	return nil} // HttpRequest injects Kimi credentials into the request and executes it.func (e *KimiExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) {	if req == nil {		return nil, fmt.Errorf("kimi executor: request is nil")	}	if ctx == nil {		ctx = req.Context()	}	httpReq := req.WithContext(ctx)	if err := e.PrepareRequest(httpReq, auth); err != nil {		return nil, err	}	httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0)	return httpClient.Do(httpReq)} // Execute performs a non-streaming chat completion request to Kimi.func (e *KimiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {	from := opts.SourceFormat	if from.String() == "claude" {		auth.Attributes["base_url"] = kimiauth.KimiAPIBaseURL		return e.ClaudeExecutor.Execute(ctx, auth, req, opts)	}	responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) 	baseModel := thinking.ParseSuffix(req.Model).ModelName 	token := kimiCreds(auth) 	reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)	defer reporter.TrackFailure(ctx, &err) 	to := sdktranslator.FromString("openai")	originalPayloadSource := req.Payload	if len(opts.OriginalRequest) > 0 {		originalPayloadSource = opts.OriginalRequest	}	originalPayload := bytes.Clone(originalPayloadSource)	originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, false)	body := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, bytes.Clone(req.Payload), false) 	// 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)	} 	body, err = helps.ApplyThinkingWithSourcePayload(body, req.Payload, originalPayloadSource, req.Model, from.String(), "kimi", e.Identifier())	if err != nil {		return resp, err	} 	requestedModel := helps.PayloadRequestedModel(opts, req.Model)	requestPath := helps.PayloadRequestPath(opts)	body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)	body, err = normalizeKimiToolMessageLinks(body)	if err != nil {		return resp, err	}	reporter.SetTranslatedReasoningEffort(body, e.Identifier()) 	url := kimiauth.KimiAPIBaseURL + "/v1/chat/completions"	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))	if err != nil {		return resp, err	}	applyKimiHeadersWithAuth(httpReq, token, false, auth)	var attrs map[string]string	if auth != nil {		attrs = auth.Attributes	}	util.ApplyCustomHeadersFromAttrs(httpReq, attrs)	var authID, authLabel, authType, authValue string	if auth != nil {		authID = auth.ID		authLabel = auth.Label		authType, authValue = auth.AccountInfo()	}	helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{		URL:       url,		Method:    http.MethodPost,		Headers:   httpReq.Header.Clone(),		Body:      body,		Provider:  e.Identifier(),		AuthID:    authID,		AuthLabel: authLabel,		AuthType:  authType,		AuthValue: authValue,	}) 	httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0)	httpClient = reporter.TrackHTTPClient(httpClient)	httpResp, err := httpClient.Do(httpReq)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return resp, err	}	defer func() {		if errClose := httpResp.Body.Close(); errClose != nil {			log.Errorf("kimi executor: close response body error: %v", errClose)		}	}()	helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {		b, _ := io.ReadAll(httpResp.Body)		helps.AppendAPIResponseChunk(ctx, e.cfg, b)		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b))		err = statusErr{code: httpResp.StatusCode, msg: string(b)}		return resp, err	}	data, err := io.ReadAll(httpResp.Body)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return resp, err	}	helps.AppendAPIResponseChunk(ctx, e.cfg, data)	reporter.Publish(ctx, helps.ParseOpenAIUsage(data))	var param any	// Note: TranslateNonStream uses req.Model (original with suffix) to preserve	// the original model name in the response for client compatibility.	out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, data, &param)	resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()}	return resp, nil} // ExecuteStream performs a streaming chat completion request to Kimi.func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) {	from := opts.SourceFormat	if from.String() == "claude" {		auth.Attributes["base_url"] = kimiauth.KimiAPIBaseURL		return e.ClaudeExecutor.ExecuteStream(ctx, auth, req, opts)	}	responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) 	baseModel := thinking.ParseSuffix(req.Model).ModelName	token := kimiCreds(auth) 	reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)	defer reporter.TrackFailure(ctx, &err) 	to := sdktranslator.FromString("openai")	originalPayloadSource := req.Payload	if len(opts.OriginalRequest) > 0 {		originalPayloadSource = opts.OriginalRequest	}	originalPayload := bytes.Clone(originalPayloadSource)	originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, true)	body := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, bytes.Clone(req.Payload), true) 	// 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)	} 	body, err = helps.ApplyThinkingWithSourcePayload(body, req.Payload, originalPayloadSource, req.Model, from.String(), "kimi", e.Identifier())	if err != nil {		return nil, err	} 	body, err = sjson.SetBytes(body, "stream_options.include_usage", true)	if err != nil {		return nil, fmt.Errorf("kimi executor: failed to set stream_options in payload: %w", err)	}	requestedModel := helps.PayloadRequestedModel(opts, req.Model)	requestPath := helps.PayloadRequestPath(opts)	body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)	body, err = normalizeKimiToolMessageLinks(body)	if err != nil {		return nil, err	}	reporter.SetTranslatedReasoningEffort(body, e.Identifier()) 	url := kimiauth.KimiAPIBaseURL + "/v1/chat/completions"	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))	if err != nil {		return nil, err	}	applyKimiHeadersWithAuth(httpReq, token, true, auth)	var attrs map[string]string	if auth != nil {		attrs = auth.Attributes	}	util.ApplyCustomHeadersFromAttrs(httpReq, attrs)	var authID, authLabel, authType, authValue string	if auth != nil {		authID = auth.ID		authLabel = auth.Label		authType, authValue = auth.AccountInfo()	}	helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{		URL:       url,		Method:    http.MethodPost,		Headers:   httpReq.Header.Clone(),		Body:      body,		Provider:  e.Identifier(),		AuthID:    authID,		AuthLabel: authLabel,		AuthType:  authType,		AuthValue: authValue,	}) 	httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0)	httpClient = reporter.TrackHTTPClient(httpClient)	httpResp, err := httpClient.Do(httpReq)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return nil, err	}	helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {		b, _ := io.ReadAll(httpResp.Body)		helps.AppendAPIResponseChunk(ctx, e.cfg, b)		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b))		if errClose := httpResp.Body.Close(); errClose != nil {			log.Errorf("kimi executor: close response body error: %v", errClose)		}		err = statusErr{code: httpResp.StatusCode, msg: string(b)}		return nil, err	}	out := make(chan cliproxyexecutor.StreamChunk)	go func() {		defer close(out)		defer func() {			if errClose := httpResp.Body.Close(); errClose != nil {				log.Errorf("kimi executor: close response body error: %v", errClose)			}		}()		scanner := bufio.NewScanner(httpResp.Body)		scanner.Buffer(nil, 1_048_576) // 1MB		claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload)		var param any		var streamUsage helps.StreamUsageBuffer		defer streamUsage.Publish(ctx, reporter)		for scanner.Scan() {			line := scanner.Bytes()			helps.AppendAPIResponseChunk(ctx, e.cfg, line)			streamUsage.ObserveOpenAIStream(line)			chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, bytes.Clone(line), &param, claudeInputTokens)			for i := range chunks {				select {				case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}:				case <-ctx.Done():					return				}			}		}		doneChunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, []byte("[DONE]"), &param, claudeInputTokens)		for i := range doneChunks {			select {			case out <- cliproxyexecutor.StreamChunk{Payload: doneChunks[i]}:			case <-ctx.Done():				return			}		}		if errScan := scanner.Err(); errScan != nil {			helps.RecordAPIResponseError(ctx, e.cfg, errScan)			reporter.PublishFailure(ctx, errScan)			select {			case out <- cliproxyexecutor.StreamChunk{Err: errScan}:			case <-ctx.Done():			}		}	}()	return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil} // CountTokens estimates token count for Kimi requests.func (e *KimiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {	auth.Attributes["base_url"] = kimiauth.KimiAPIBaseURL	return e.ClaudeExecutor.countTokensUpstream(ctx, auth, req, opts)} func normalizeKimiToolMessageLinks(body []byte) ([]byte, error) {	if len(body) == 0 || !gjson.ValidBytes(body) {		return body, nil	} 	messages := util.GetGJSONBytesNoCopy(body, "messages")	if !messages.Exists() || !messages.IsArray() {		return body, nil	} 	type messagePatch struct {		index        int		path         string		value        string		errorContext string	} 	msgs := messages.Array()	droppedMessages := make([]bool, len(msgs))	patches := make([]messagePatch, 0)	pending := make([]string, 0)	dropped := 0	patched := 0	patchedReasoning := 0	ambiguous := 0	latestReasoning := ""	hasLatestReasoning := false 	removePending := func(id string) {		for idx := range pending {			if pending[idx] != id {				continue			}			pending = append(pending[:idx], pending[idx+1:]...)			return		}	} 	for msgIndex, msg := range msgs {		if shouldDropKimiAssistantMessage(msg) {			droppedMessages[msgIndex] = true			dropped++			continue		} 		role := strings.TrimSpace(msg.Get("role").String())		switch role {		case "assistant":			reasoning := msg.Get("reasoning_content")			if reasoning.Exists() {				reasoningText := reasoning.String()				if strings.TrimSpace(reasoningText) != "" {					latestReasoning = reasoningText					hasLatestReasoning = true				}			} 			toolCalls := msg.Get("tool_calls")			if toolCalls.Exists() && toolCalls.IsArray() {				toolCallItems := toolCalls.Array()				if len(toolCallItems) > 0 {					if !reasoning.Exists() || strings.TrimSpace(reasoning.String()) == "" {						patches = append(patches, messagePatch{							index:        msgIndex,							path:         "reasoning_content",							value:        fallbackAssistantReasoning(msg, hasLatestReasoning, latestReasoning),							errorContext: "failed to set assistant reasoning_content",						})						patchedReasoning++					}					for _, toolCall := range toolCallItems {						id := strings.TrimSpace(toolCall.Get("id").String())						if id != "" {							pending = append(pending, id)						}					}				}			}		case "tool":			toolCallID := strings.TrimSpace(msg.Get("tool_call_id").String())			if toolCallID == "" {				toolCallID = strings.TrimSpace(msg.Get("call_id").String())				if toolCallID != "" {					patches = append(patches, messagePatch{index: msgIndex, path: "tool_call_id", value: toolCallID, errorContext: "failed to set tool_call_id from call_id"})					patched++				}			}			if toolCallID == "" {				if len(pending) == 1 {					toolCallID = pending[0]					patches = append(patches, messagePatch{index: msgIndex, path: "tool_call_id", value: toolCallID, errorContext: "failed to infer tool_call_id"})					patched++				} else if len(pending) > 1 {					ambiguous++				}			}			if toolCallID != "" {				removePending(toolCallID)			}		}	} 	if dropped > 0 {		log.WithField("dropped_assistant_messages", dropped).Debug("kimi executor: dropped empty assistant messages")	}	if dropped == 0 && len(patches) == 0 {		if ambiguous > 0 {			log.WithFields(log.Fields{				"ambiguous_tool_messages": ambiguous,				"pending_tool_calls":      len(pending),			}).Warn("kimi executor: tool messages missing tool_call_id with ambiguous candidates")		}		return body, nil	} 	var out []byte	if dropped == 0 && len(patches) == 1 {		patch := patches[0]		path := fmt.Sprintf("messages.%d.%s", patch.index, patch.path)		updated, errSet := sjson.SetBytes(body, path, patch.value)		if errSet != nil {			return body, fmt.Errorf("kimi executor: %s: %w", patch.errorContext, errSet)		}		out = updated	} else {		messageItems := make([]string, 0, len(msgs)-dropped)		patchIndex := 0		for msgIndex, msg := range msgs {			if droppedMessages[msgIndex] {				continue			}			messageJSON := msg.Raw			for patchIndex < len(patches) && patches[patchIndex].index == msgIndex {				patch := patches[patchIndex]				next, errSet := sjson.SetBytes([]byte(messageJSON), patch.path, patch.value)				if errSet != nil {					return body, fmt.Errorf("kimi executor: %s: %w", patch.errorContext, errSet)				}				messageJSON = string(next)				patchIndex++			}			messageItems = append(messageItems, messageJSON)		}		updated, errSet := sjson.SetRawBytes(body, "messages", helps.JoinRawJSONStrings(messageItems))		if errSet != nil {			if dropped > 0 {				return body, fmt.Errorf("kimi executor: failed to drop empty assistant messages: %w", errSet)			}			return body, fmt.Errorf("kimi executor: %s: %w", patches[0].errorContext, errSet)		}		out = updated	} 	if patched > 0 || patchedReasoning > 0 {		log.WithFields(log.Fields{			"patched_tool_messages":      patched,			"patched_reasoning_messages": patchedReasoning,		}).Debug("kimi executor: normalized tool message fields")

Lines 1–500 of 805.