Merge branch 'router-for-me:main' into main

This commit is contained in:
Luis Pater
2026-01-13 03:15:46 +08:00
committed by GitHub
15 changed files with 209 additions and 59 deletions
+13 -6
View File
@@ -146,10 +146,12 @@ func (h *ClaudeCodeAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSO
c.Header("Content-Type", "application/json")
alt := h.GetAlt(c)
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
modelName := gjson.GetBytes(rawJSON, "model").String()
resp, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
stopKeepAlive()
if errMsg != nil {
h.WriteErrorResponse(c, errMsg)
cliCancel(errMsg.Error)
@@ -159,13 +161,18 @@ func (h *ClaudeCodeAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSO
// Decompress gzipped responses - Claude API sometimes returns gzip without Content-Encoding header
// This fixes title generation and other non-streaming responses that arrive compressed
if len(resp) >= 2 && resp[0] == 0x1f && resp[1] == 0x8b {
gzReader, err := gzip.NewReader(bytes.NewReader(resp))
if err != nil {
log.Warnf("failed to decompress gzipped Claude response: %v", err)
gzReader, errGzip := gzip.NewReader(bytes.NewReader(resp))
if errGzip != nil {
log.Warnf("failed to decompress gzipped Claude response: %v", errGzip)
} else {
defer gzReader.Close()
if decompressed, err := io.ReadAll(gzReader); err != nil {
log.Warnf("failed to read decompressed Claude response: %v", err)
defer func() {
if errClose := gzReader.Close(); errClose != nil {
log.Warnf("failed to close Claude gzip reader: %v", errClose)
}
}()
decompressed, errRead := io.ReadAll(gzReader)
if errRead != nil {
log.Warnf("failed to read decompressed Claude response: %v", errRead)
} else {
resp = decompressed
}
@@ -336,7 +336,9 @@ func (h *GeminiAPIHandler) handleGenerateContent(c *gin.Context, modelName strin
c.Header("Content-Type", "application/json")
alt := h.GetAlt(c)
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
resp, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
stopKeepAlive()
if errMsg != nil {
h.WriteErrorResponse(c, errMsg)
cliCancel(errMsg.Error)
+61
View File
@@ -9,6 +9,7 @@ import (
"fmt"
"net/http"
"strings"
"sync"
"time"
"github.com/gin-gonic/gin"
@@ -113,6 +114,19 @@ func StreamingKeepAliveInterval(cfg *config.SDKConfig) time.Duration {
return time.Duration(seconds) * time.Second
}
// NonStreamingKeepAliveInterval returns the keep-alive interval for non-streaming responses.
// Returning 0 disables keep-alives (default when unset).
func NonStreamingKeepAliveInterval(cfg *config.SDKConfig) time.Duration {
seconds := 0
if cfg != nil {
seconds = cfg.NonStreamKeepAliveInterval
}
if seconds <= 0 {
return 0
}
return time.Duration(seconds) * time.Second
}
// StreamingBootstrapRetries returns how many times a streaming request may be retried before any bytes are sent.
func StreamingBootstrapRetries(cfg *config.SDKConfig) int {
retries := defaultStreamingBootstrapRetries
@@ -294,6 +308,53 @@ func (h *BaseAPIHandler) GetContextWithCancel(handler interfaces.APIHandler, c *
}
}
// StartNonStreamingKeepAlive emits blank lines every 5 seconds while waiting for a non-streaming response.
// It returns a stop function that must be called before writing the final response.
func (h *BaseAPIHandler) StartNonStreamingKeepAlive(c *gin.Context, ctx context.Context) func() {
if h == nil || c == nil {
return func() {}
}
interval := NonStreamingKeepAliveInterval(h.Cfg)
if interval <= 0 {
return func() {}
}
flusher, ok := c.Writer.(http.Flusher)
if !ok {
return func() {}
}
if ctx == nil {
ctx = context.Background()
}
stopChan := make(chan struct{})
var stopOnce sync.Once
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-stopChan:
return
case <-ctx.Done():
return
case <-ticker.C:
_, _ = c.Writer.Write([]byte("\n"))
flusher.Flush()
}
}
}()
return func() {
stopOnce.Do(func() {
close(stopChan)
})
wg.Wait()
}
}
// appendAPIResponse preserves any previously captured API response and appends new data.
func appendAPIResponse(c *gin.Context, data []byte) {
if c == nil || len(data) == 0 {
@@ -56,6 +56,14 @@ func (e *failOnceStreamExecutor) CountTokens(context.Context, *coreauth.Auth, co
return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"}
}
func (e *failOnceStreamExecutor) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (*http.Response, error) {
return nil, &coreauth.Error{
Code: "not_implemented",
Message: "HttpRequest not implemented",
HTTPStatus: http.StatusNotImplemented,
}
}
func (e *failOnceStreamExecutor) Calls() int {
e.mu.Lock()
defer e.mu.Unlock()
@@ -524,7 +524,9 @@ func (h *OpenAIAPIHandler) handleCompletionsNonStreamingResponse(c *gin.Context,
modelName := gjson.GetBytes(chatCompletionsJSON, "model").String()
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
resp, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, chatCompletionsJSON, "")
stopKeepAlive()
if errMsg != nil {
h.WriteErrorResponse(c, errMsg)
cliCancel(errMsg.Error)
@@ -103,20 +103,17 @@ func (h *OpenAIResponsesAPIHandler) handleNonStreamingResponse(c *gin.Context, r
modelName := gjson.GetBytes(rawJSON, "model").String()
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
defer func() {
cliCancel()
}()
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
resp, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
stopKeepAlive()
if errMsg != nil {
h.WriteErrorResponse(c, errMsg)
cliCancel(errMsg.Error)
return
}
_, _ = c.Writer.Write(resp)
return
// no legacy fallback
cliCancel()
}
// handleStreamingResponse handles streaming responses for Gemini models.
+30 -31
View File
@@ -1247,7 +1247,7 @@ func applyOAuthModelMappings(cfg *config.Config, provider, authKind string, mode
fork bool
}
forward := make(map[string]mappingEntry, len(mappings))
forward := make(map[string][]mappingEntry, len(mappings))
for i := range mappings {
name := strings.TrimSpace(mappings[i].Name)
alias := strings.TrimSpace(mappings[i].Alias)
@@ -1258,14 +1258,12 @@ func applyOAuthModelMappings(cfg *config.Config, provider, authKind string, mode
continue
}
key := strings.ToLower(name)
if _, exists := forward[key]; exists {
continue
}
forward[key] = mappingEntry{alias: alias, fork: mappings[i].Fork}
forward[key] = append(forward[key], mappingEntry{alias: alias, fork: mappings[i].Fork})
}
if len(forward) == 0 {
return models
}
out := make([]*ModelInfo, 0, len(models))
seen := make(map[string]struct{}, len(models))
for _, model := range models {
@@ -1277,17 +1275,8 @@ func applyOAuthModelMappings(cfg *config.Config, provider, authKind string, mode
continue
}
key := strings.ToLower(id)
entry, ok := forward[key]
if !ok {
if _, exists := seen[key]; exists {
continue
}
seen[key] = struct{}{}
out = append(out, model)
continue
}
mappedID := strings.TrimSpace(entry.alias)
if mappedID == "" {
entries := forward[key]
if len(entries) == 0 {
if _, exists := seen[key]; exists {
continue
}
@@ -1296,11 +1285,29 @@ func applyOAuthModelMappings(cfg *config.Config, provider, authKind string, mode
continue
}
if entry.fork {
keepOriginal := false
for _, entry := range entries {
if entry.fork {
keepOriginal = true
break
}
}
if keepOriginal {
if _, exists := seen[key]; !exists {
seen[key] = struct{}{}
out = append(out, model)
}
}
addedAlias := false
for _, entry := range entries {
mappedID := strings.TrimSpace(entry.alias)
if mappedID == "" {
continue
}
if strings.EqualFold(mappedID, id) {
continue
}
aliasKey := strings.ToLower(mappedID)
if _, exists := seen[aliasKey]; exists {
continue
@@ -1312,24 +1319,16 @@ func applyOAuthModelMappings(cfg *config.Config, provider, authKind string, mode
clone.Name = rewriteModelInfoName(clone.Name, id, mappedID)
}
out = append(out, &clone)
continue
addedAlias = true
}
uniqueKey := strings.ToLower(mappedID)
if _, exists := seen[uniqueKey]; exists {
continue
}
seen[uniqueKey] = struct{}{}
if mappedID == id {
if !keepOriginal && !addedAlias {
if _, exists := seen[key]; exists {
continue
}
seen[key] = struct{}{}
out = append(out, model)
continue
}
clone := *model
clone.ID = mappedID
if clone.Name != "" {
clone.Name = rewriteModelInfoName(clone.Name, id, mappedID)
}
out = append(out, &clone)
}
return out
}
@@ -56,3 +56,37 @@ func TestApplyOAuthModelMappings_ForkAddsAlias(t *testing.T) {
t.Fatalf("expected forked model name %q, got %q", "models/g5", out[1].Name)
}
}
func TestApplyOAuthModelMappings_ForkAddsMultipleAliases(t *testing.T) {
cfg := &config.Config{
OAuthModelMappings: map[string][]config.ModelNameMapping{
"codex": {
{Name: "gpt-5", Alias: "g5", Fork: true},
{Name: "gpt-5", Alias: "g5-2", Fork: true},
},
},
}
models := []*ModelInfo{
{ID: "gpt-5", Name: "models/gpt-5"},
}
out := applyOAuthModelMappings(cfg, "codex", "oauth", models)
if len(out) != 3 {
t.Fatalf("expected 3 models, got %d", len(out))
}
if out[0].ID != "gpt-5" {
t.Fatalf("expected first model id %q, got %q", "gpt-5", out[0].ID)
}
if out[1].ID != "g5" {
t.Fatalf("expected second model id %q, got %q", "g5", out[1].ID)
}
if out[1].Name != "models/g5" {
t.Fatalf("expected forked model name %q, got %q", "models/g5", out[1].Name)
}
if out[2].ID != "g5-2" {
t.Fatalf("expected third model id %q, got %q", "g5-2", out[2].ID)
}
if out[2].Name != "models/g5-2" {
t.Fatalf("expected forked model name %q, got %q", "models/g5-2", out[2].Name)
}
}