Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 101 additions & 25 deletions internal/runtime/executor/aistudio_executor.go
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,11 @@ func (e *AIStudioExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth,
if helps.ApplyPatchTranslationError(param) != nil || len(out) == 0 {
return cliproxyexecutor.Response{}, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
}
if errPayload, okErr := helps.TranslatedClaudeErrorPayload(responseFormat, out); okErr {
err = statusErr{code: http.StatusBadGateway, msg: errPayload}
reporter.PublishFailureWithDetail(ctx, helps.ParseGeminiUsage(wsResp.Body), err)
return cliproxyexecutor.Response{}, err
}
reporter.Publish(ctx, helps.ParseGeminiUsage(wsResp.Body))
if responseFormat == sdktranslator.FormatOpenAIResponse {
out = helps.EnsureResponsesUsageDetails(out)
Expand Down Expand Up @@ -310,7 +315,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
}
return nil, statusErr{code: firstEvent.Status, msg: body.String()}
}
out := make(chan cliproxyexecutor.StreamChunk)
out := make(chan cliproxyexecutor.StreamChunk, 1)
go func(first wsrelay.StreamEvent) {
defer close(out)
defer cancelStream()
Expand All @@ -323,15 +328,28 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
claudeInputTokens := helps.NewClaudeInputTokenState(opts.SourceFormat, body.toFormat, responseFormat, originalRequest)
var streamUsage helps.StreamUsageBuffer
defer streamUsage.Publish(ctx, reporter)
var streamErr error
defer func() {
_ = helps.StopTranslatedClaudeStream(ctx, &streamErr, reporter, &streamUsage, out)
}()
var param any
helps.InitializeApplyPatchStream(ctx, body.toFormat, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), translatedReq, &param)
metadataLogged := false
finishStream := func() {
if helps.EndApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) || ctx.Err() != nil {
if helps.EndApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) {
return
}
if helps.StopTranslatedClaudeStream(ctx, &streamErr, reporter, &streamUsage, out) {
return
}
if ctx.Err() != nil {
return
}
lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), translatedReq, []byte("[DONE]"), &param, claudeInputTokens)
helps.RecordApplyPatchStreamFailure(ctx, param, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error {
return statusErr{code: http.StatusBadGateway, msg: msg}
})
for _, line := range lines {
select {
case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(line)}:
Expand All @@ -340,9 +358,13 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
}
}
_ = helps.StopApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
_ = helps.StopTranslatedClaudeStream(ctx, &streamErr, reporter, &streamUsage, out)
}
processEvent := func(event wsrelay.StreamEvent) bool {
if event.Err != nil {
if helps.StopTranslatedClaudeStream(ctx, &streamErr, reporter, &streamUsage, out) {
return false
}
helps.RecordAPIResponseError(ctx, e.cfg, event.Err)
reporter.PublishFailure(ctx, event.Err)
select {
Expand All @@ -362,23 +384,30 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
if len(event.Payload) > 0 {
reporter.MarkFirstResponseByte()
helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload)
reporter.ObserveResponseModel(event.Payload)
filtered := helps.FilterSSEUsageMetadata(event.Payload)
if detail, ok := helps.ParseGeminiStreamUsage(filtered); ok {
streamUsage.Observe(detail, true)
}
lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), translatedReq, filtered, &param, claudeInputTokens)
helps.RecordApplyPatchStreamFailure(ctx, param, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
for i := range lines {
select {
case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}:
case <-ctx.Done():
// A relay chunk may hold several SSE frames; the translators read one frame per call (gjson parses the first
// JSON object only), so each frame goes through the whole per-frame path: model, usage, translation, errors.
for _, frame := range splitRelayFrames(event.Payload) {
reporter.ObserveResponseModel(frame)
if detail, ok := helps.ParseGeminiStreamUsage(frame); ok {
streamUsage.Observe(detail, true)
}
filtered := helps.FilterSSEUsageMetadata(frame)
lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), translatedReq, filtered, &param, claudeInputTokens)
helps.RecordApplyPatchStreamFailure(ctx, param, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error {
return statusErr{code: http.StatusBadGateway, msg: msg}
})
for i := range lines {
select {
case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}:
case <-ctx.Done():
return false
}
}
if helps.StopApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) {
return false
}
}
if helps.StopApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) {
return false
}
break
}
case wsrelay.MessageTypeStreamEnd:
Expand All @@ -404,8 +433,13 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
}
return false
}
reporter.ObserveResponseModel(event.Payload)
streamUsage.Observe(helps.ParseGeminiUsage(event.Payload), true)
lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), translatedReq, event.Payload, &param, claudeInputTokens)
helps.RecordApplyPatchStreamFailure(ctx, param, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error {
return statusErr{code: http.StatusBadGateway, msg: msg}
})
for i := range lines {
select {
case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}:
Expand All @@ -417,10 +451,11 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
return false
}
finishStream()
reporter.ObserveResponseModel(event.Payload)
streamUsage.Observe(helps.ParseGeminiUsage(event.Payload), true)
return false
case wsrelay.MessageTypeError:
if helps.StopTranslatedClaudeStream(ctx, &streamErr, reporter, &streamUsage, out) {
return false
}
helps.RecordAPIResponseError(ctx, e.cfg, event.Err)
reporter.PublishFailure(ctx, event.Err)
select {
Expand All @@ -440,6 +475,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
}
}
_ = helps.EndApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
_ = helps.StopTranslatedClaudeStream(ctx, &streamErr, reporter, &streamUsage, out)
}(firstEvent)
return &cliproxyexecutor.StreamResult{Headers: firstEvent.Headers.Clone(), Chunks: out}, nil
}
Expand All @@ -460,10 +496,6 @@ func (e *AIStudioExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.A
return cliproxyexecutor.Response{}, err
}

body.payload, _ = sjson.DeleteBytes(body.payload, "generationConfig")
body.payload, _ = sjson.DeleteBytes(body.payload, "tools")
body.payload, _ = sjson.DeleteBytes(body.payload, "safetySettings")

endpoint := e.buildEndpoint(baseModel, "countTokens", "")
wsReq := &wsrelay.HTTPRequest{
Method: http.MethodPost,
Expand Down Expand Up @@ -534,15 +566,17 @@ func (e *AIStudioExecutor) translateRequest(ctx context.Context, req cliproxyexe
}
originalPayload := originalPayloadSource
originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, stream)
payload := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, stream)
payload, err := helps.ApplyThinkingWithSourcePayload(payload, req.Payload, originalPayloadSource, req.Model, from.String(), to.String(), e.Identifier())
payload, err := helps.TranslateRequestReturningError(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, stream, false)
if err != nil {
return nil, translatedPayload{}, err
}
payload, err = helps.ApplyThinkingWithSourcePayload(payload, req.Payload, originalPayloadSource, req.Model, from.String(), to.String(), e.Identifier())
if err != nil {
return nil, translatedPayload{}, err
}
payload = fixGeminiImageAspectRatio(baseModel, payload)
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
requestPath := helps.PayloadRequestPath(opts)
payload = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", payload, originalTranslated, requestedModel, requestPath, opts.Headers)
payload, _ = sjson.DeleteBytes(payload, "generationConfig.maxOutputTokens")
payload, _ = sjson.DeleteBytes(payload, "generationConfig.responseMimeType")
payload, _ = sjson.DeleteBytes(payload, "generationConfig.responseJsonSchema")
Expand All @@ -562,6 +596,12 @@ func (e *AIStudioExecutor) translateRequest(ctx context.Context, req cliproxyexe
payload = helps.EnsureGeminiTrailingUserContent(payload, "contents")
}
payload = normalizeAIStudioThinkingLevel(payload)
if action == "countTokens" {
payload, _ = sjson.DeleteBytes(payload, "generationConfig")
payload, _ = sjson.DeleteBytes(payload, "tools")
payload, _ = sjson.DeleteBytes(payload, "safetySettings")
}
payload = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", payload, originalTranslated, requestedModel, requestPath, opts.Headers)
return payload, translatedPayload{payload: payload, action: action, toFormat: to}, nil
}

Expand Down Expand Up @@ -606,6 +646,42 @@ func (e *AIStudioExecutor) buildEndpoint(model, action, alt string) string {
return base
}

// splitRelayFrames splits one relay stream chunk into the SSE frames it holds. A chunk that starts with an SSE field (data:, event:,
// id:, retry: or a comment) is cut at its blank lines and each frame is reduced to its data: lines (the translators read data
// lines; event:, id:, retry: and comment lines carry nothing they use, and a frame without data, a keep-alive comment, is dropped),
// each frame ending with the blank line. A chunk that starts anything else (a bare or pretty-printed JSON value, "[DONE]") is one
// frame as it is, never cut.
func splitRelayFrames(payload []byte) [][]byte {
trimmed := bytes.TrimLeft(payload, " \t\r\n")
if !startsWithSSEField(trimmed) {
return [][]byte{payload}
}
normalized := bytes.ReplaceAll(payload, []byte("\r\n"), []byte("\n"))
var frames [][]byte
for _, block := range bytes.Split(normalized, []byte("\n\n")) {
var frame []byte
for _, line := range bytes.Split(block, []byte("\n")) {
if bytes.HasPrefix(bytes.TrimLeft(line, " \t"), []byte("data:")) {
frame = append(frame, bytes.TrimLeft(line, " \t")...)
frame = append(frame, '\n')
}
}
if len(frame) > 0 {
frames = append(frames, append(frame, '\n'))
}
}
return frames
}

func startsWithSSEField(b []byte) bool {
for _, prefix := range []string{"data:", "event:", "id:", "retry:", ":"} {
if bytes.HasPrefix(b, []byte(prefix)) {
return true
}
}
return false
}

// ensureColonSpacedJSON normalizes JSON objects so that colons are followed by a single space while
// keeping the payload otherwise compact. Non-JSON inputs are returned unchanged.
func ensureColonSpacedJSON(payload []byte) []byte {
Expand Down
13 changes: 9 additions & 4 deletions internal/runtime/executor/antigravity_executor.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import (
"github.com/router-for-me/CLIProxyAPI/v8/internal/util"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor"
"github.com/router-for-me/CLIProxyAPI/v8/sdk/httptransport"
"github.com/router-for-me/CLIProxyAPI/v8/sdk/proxyutil"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator"
log "github.com/sirupsen/logrus"
Expand All @@ -41,7 +42,7 @@ const (
antigravityClientSecret = "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"
antigravityAuthType = "antigravity"
antigravityRequestTokenSafetyWindow = 5 * time.Minute
antigravityCreditsHintRefreshInterval = 10 * time.Minute
antigravityCreditsHintRefreshInterval = cliproxyauth.AntigravityCreditsHintRefreshInterval
antigravityCreditsHintRefreshTimeout = 5 * time.Second
antigravityShortQuotaCooldownThreshold = 5 * time.Minute
antigravityInstantRetryThreshold = 3 * time.Second
Expand Down Expand Up @@ -436,7 +437,8 @@ func antigravityCredentialScope(prefix, secret string) string {
// negotiates TLS 1.3 without advertising an ALPN protocol and therefore never uses h2.
// The underlying Transport is always shared so keep-alive connections survive across
// requests instead of forcing a fresh TCP + TLS handshake every time.
func newAntigravityHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {
func newAntigravityHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) (clientOut *http.Client) {
defer func() { clientOut = httptransport.Decorate(ctx, clientOut) }()
// Native Antigravity reuses one transport across requests. Opt into a
// credential-scoped proxy transport only here so other providers keep their
// existing lifecycle and different OAuth identities remain isolated.
Expand All @@ -448,7 +450,7 @@ func newAntigravityHTTPClient(ctx context.Context, cfg *config.Config, auth *cli
// context transport fallback, preserving the previous behavior.
}

client := helps.NewProxyAwareHTTPClient(ctx, cfg, auth, timeout)
client := helps.NewProxyAwareHTTPClient(httptransport.WithoutDecorator(ctx), cfg, auth, timeout)
// Direct requests share an HTTP/1.1 pool only within the selected credential.
if client.Transport == nil {
client.Transport = antigravityHTTP11Transport(auth, antigravityBaseTransport, cfg)
Expand Down Expand Up @@ -517,11 +519,14 @@ func ensureAntigravityGeminiTrailingUserContent(modelName string, payload []byte
}

// ensureAntigravityGeminiBoundaryUserContent normalizes both leading and trailing
// turns for Gemini targets. Claude targets are left unchanged.
// turns for Gemini targets. A trailing reasoning-only model turn is dropped first
// so the request does not rely only on a synthetic empty user turn (#6530).
// Claude targets are left unchanged.
func ensureAntigravityGeminiBoundaryUserContent(modelName string, payload []byte) []byte {
if strings.Contains(strings.ToLower(modelName), "claude") {
return payload
}
payload = helps.DropGeminiTrailingThoughtOnlyContent(payload, "request.contents")
return helps.EnsureGeminiBoundaryUserContent(payload, "request.contents")
}

Expand Down
15 changes: 10 additions & 5 deletions internal/runtime/executor/antigravity_executor_credits.go
Original file line number Diff line number Diff line change
Expand Up @@ -247,7 +247,7 @@ func decideAntigravity429(body []byte) antigravity429Decision {
reason := strings.TrimSpace(detail.Get("reason").String())
decision.reason = reason
switch {
case strings.EqualFold(reason, "QUOTA_EXHAUSTED"):
case strings.EqualFold(reason, "QUOTA_EXHAUSTED"), strings.EqualFold(reason, "INSUFFICIENT_G1_CREDITS_BALANCE"):
decision.kind = antigravity429DecisionFullQuotaExhausted
return decision
case strings.EqualFold(reason, "RATE_LIMIT_EXCEEDED"):
Expand Down Expand Up @@ -289,7 +289,13 @@ func clearAntigravityCreditsFailureState(auth *cliproxyauth.Auth) {
if auth == nil || strings.TrimSpace(auth.ID) == "" {
return
}
antigravityCreditsFailureByAuth.Delete(strings.TrimSpace(auth.ID))
authID := strings.TrimSpace(auth.ID)
_, hadFailure := antigravityCreditsFailureByAuth.LoadAndDelete(authID)
if hint, ok := cliproxyauth.GetAntigravityCreditsHint(authID); ok && hint.Known && !hint.Available {
if hadFailure || !cliproxyauth.IsAntigravityCreditsHintFresh(hint, time.Now()) {
cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{})
}
}
}
func markAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) {
if auth == nil || strings.TrimSpace(auth.ID) == "" || antigravityCoolingDisabled(auth, nil) {
Expand Down Expand Up @@ -363,7 +369,7 @@ func (e *AntigravityExecutor) maybeRefreshAntigravityCreditsHint(ctx context.Con
if authID == "" {
return
}
if hint, ok := cliproxyauth.GetAntigravityCreditsHint(authID); ok && hint.Known {
if hint, ok := cliproxyauth.GetAntigravityCreditsHint(authID); ok && cliproxyauth.IsAntigravityCreditsHintFresh(hint, time.Now()) {
return
}
if strings.TrimSpace(accessToken) == "" {
Expand Down Expand Up @@ -516,7 +522,7 @@ func (e *AntigravityExecutor) updateAntigravityCreditsBalanceForTask(ctx context
return
}

userAgent := resolveUserAgent(auth)
userAgent := resolveLoadCodeAssistUserAgent(auth)
loadReqBody, errMarshal := json.Marshal(map[string]any{
"metadata": map[string]string{
"ideType": "ANTIGRAVITY",
Expand All @@ -534,7 +540,6 @@ func (e *AntigravityExecutor) updateAntigravityCreditsBalanceForTask(ctx context
return
}
httpReq.Header.Set("Authorization", "Bearer "+token)
httpReq.Header.Set("Accept", "*/*")
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("User-Agent", userAgent)

Expand Down
Loading
Loading