Skip to content

Commit 65aa351

Browse files
conversunconversun
committed
Merge remote-tracking branch 'upstream/main' into feat/copilot
2 parents e8b1c58 + 6bfcb0c commit 65aa351

10 files changed

Lines changed: 300 additions & 79 deletions

File tree

‎internal/api/protocol_multiplexer.go‎

Lines changed: 67 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"net"
88
"net/http"
99
"strings"
10+
"time"
1011

1112
log "github.com/sirupsen/logrus"
1213
)
@@ -48,68 +49,89 @@ func (s *Server) acceptMuxConnections(listener net.Listener, httpListener *muxLi
4849
continue
4950
}
5051

51-
tlsConn, ok := conn.(*tls.Conn)
52-
if ok {
53-
if errHandshake := tlsConn.Handshake(); errHandshake != nil {
54-
if errClose := conn.Close(); errClose != nil {
55-
log.Errorf("failed to close connection after TLS handshake error: %v", errClose)
56-
}
57-
continue
58-
}
59-
proto := strings.TrimSpace(tlsConn.ConnectionState().NegotiatedProtocol)
60-
if proto == "h2" || proto == "http/1.1" {
61-
if httpListener == nil {
62-
if errClose := conn.Close(); errClose != nil {
63-
log.Errorf("failed to close connection: %v", errClose)
64-
}
65-
continue
66-
}
67-
if errPut := httpListener.Put(tlsConn); errPut != nil {
68-
if errClose := conn.Close(); errClose != nil {
69-
log.Errorf("failed to close connection after HTTP routing failure: %v", errClose)
70-
}
71-
}
72-
continue
73-
}
74-
}
52+
// Dispatch each connection to a goroutine so that slow/idle clients
53+
// cannot block the accept loop. Previously, TLS handshake and
54+
// reader.Peek(1) were performed inline; an idle TCP connection that
55+
// never sent bytes would block Peek indefinitely, preventing all
56+
// subsequent connections from being accepted (issue #3267).
57+
go s.routeMuxConnection(conn, httpListener)
58+
}
59+
}
60+
61+
// routeMuxConnection performs per-connection protocol detection and routing.
62+
func (s *Server) routeMuxConnection(conn net.Conn, httpListener *muxListener) {
63+
// Set a read deadline so that idle connections that never send bytes do not
64+
// leak goroutines and file descriptors. The deadline is cleared once the
65+
// connection is successfully routed to its handler.
66+
const muxSniffDeadline = 10 * time.Second
67+
_ = conn.SetReadDeadline(time.Now().Add(muxSniffDeadline))
7568

76-
reader := bufio.NewReader(conn)
77-
prefix, errPeek := reader.Peek(1)
78-
if errPeek != nil {
69+
tlsConn, ok := conn.(*tls.Conn)
70+
if ok {
71+
if errHandshake := tlsConn.Handshake(); errHandshake != nil {
7972
if errClose := conn.Close(); errClose != nil {
80-
log.Errorf("failed to close connection after protocol peek failure: %v", errClose)
73+
log.Errorf("failed to close connection after TLS handshake error: %v", errClose)
8174
}
82-
continue
75+
return
8376
}
84-
85-
if isRedisRESPPrefix(prefix[0]) {
86-
if s.cfg != nil && s.cfg.Home.Enabled {
77+
proto := strings.TrimSpace(tlsConn.ConnectionState().NegotiatedProtocol)
78+
if proto == "h2" || proto == "http/1.1" {
79+
if httpListener == nil {
8780
if errClose := conn.Close(); errClose != nil {
88-
log.Errorf("failed to close redis connection while home mode is enabled: %v", errClose)
81+
log.Errorf("failed to close connection: %v", errClose)
8982
}
90-
continue
83+
return
9184
}
92-
if !s.managementRoutesEnabled.Load() {
85+
if errPut := httpListener.Put(tlsConn); errPut != nil {
9386
if errClose := conn.Close(); errClose != nil {
94-
log.Errorf("failed to close redis connection while management is disabled: %v", errClose)
87+
log.Errorf("failed to close connection after HTTP routing failure: %v", errClose)
9588
}
96-
continue
89+
} else {
90+
_ = conn.SetReadDeadline(time.Time{})
9791
}
98-
go s.handleRedisConnection(conn, reader)
99-
continue
92+
return
10093
}
94+
}
95+
96+
reader := bufio.NewReader(conn)
97+
prefix, errPeek := reader.Peek(1)
98+
if errPeek != nil {
99+
if errClose := conn.Close(); errClose != nil {
100+
log.Errorf("failed to close connection after protocol peek failure: %v", errClose)
101+
}
102+
return
103+
}
101104

102-
if httpListener == nil {
105+
if isRedisRESPPrefix(prefix[0]) {
106+
if s.cfg != nil && s.cfg.Home.Enabled {
103107
if errClose := conn.Close(); errClose != nil {
104-
log.Errorf("failed to close connection without HTTP listener: %v", errClose)
108+
log.Errorf("failed to close redis connection while home mode is enabled: %v", errClose)
105109
}
106-
continue
110+
return
107111
}
108-
109-
if errPut := httpListener.Put(&bufferedConn{Conn: conn, reader: reader}); errPut != nil {
112+
if !s.managementRoutesEnabled.Load() {
110113
if errClose := conn.Close(); errClose != nil {
111-
log.Errorf("failed to close connection after HTTP routing failure: %v", errClose)
114+
log.Errorf("failed to close redis connection while management is disabled: %v", errClose)
112115
}
116+
return
117+
}
118+
_ = conn.SetReadDeadline(time.Time{})
119+
s.handleRedisConnection(conn, reader)
120+
return
121+
}
122+
123+
if httpListener == nil {
124+
if errClose := conn.Close(); errClose != nil {
125+
log.Errorf("failed to close connection without HTTP listener: %v", errClose)
126+
}
127+
return
128+
}
129+
130+
if errPut := httpListener.Put(&bufferedConn{Conn: conn, reader: reader}); errPut != nil {
131+
if errClose := conn.Close(); errClose != nil {
132+
log.Errorf("failed to close connection after HTTP routing failure: %v", errClose)
113133
}
134+
} else {
135+
_ = conn.SetReadDeadline(time.Time{})
114136
}
115137
}

‎internal/api/protocol_multiplexer_test.go‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
package api
2+
3+
import (
4+
"net"
5+
"net/http"
6+
"net/http/httptest"
7+
"sync/atomic"
8+
"testing"
9+
"time"
10+
)
11+
12+
func TestAcceptMuxNotBlockedByIdleConnection(t *testing.T) {
13+
listener, err := net.Listen("tcp", "127.0.0.1:0")
14+
if err != nil {
15+
t.Fatalf("failed to listen: %v", err)
16+
}
17+
defer listener.Close()
18+
19+
var routed atomic.Int32
20+
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
21+
routed.Add(1)
22+
w.WriteHeader(http.StatusOK)
23+
})
24+
srv := httptest.NewUnstartedServer(handler)
25+
defer srv.Close()
26+
27+
muxLn := newMuxListener(listener.Addr(), 1024)
28+
server := &Server{managementRoutesEnabled: atomic.Bool{}}
29+
server.managementRoutesEnabled.Store(false)
30+
31+
errCh := make(chan error, 1)
32+
go func() {
33+
errCh <- server.acceptMuxConnections(listener, muxLn)
34+
}()
35+
36+
srv.Listener = muxLn
37+
srv.Start()
38+
39+
// Open an idle TCP connection that never sends any bytes.
40+
idleConn, err := net.DialTimeout("tcp", listener.Addr().String(), 2*time.Second)
41+
if err != nil {
42+
t.Fatalf("failed to dial idle connection: %v", err)
43+
}
44+
defer idleConn.Close()
45+
46+
// Give the accept loop time to pick up the idle connection.
47+
time.Sleep(50 * time.Millisecond)
48+
49+
// Send a real HTTP request. Before the fix, the accept loop would be
50+
// blocked on Peek(1) for the idle connection, causing this request to
51+
// time out.
52+
client := &http.Client{Timeout: 3 * time.Second}
53+
resp, err := client.Get("http://" + listener.Addr().String() + "/")
54+
if err != nil {
55+
listener.Close()
56+
t.Fatalf("HTTP request failed (accept loop may be blocked by idle connection): %v", err)
57+
}
58+
resp.Body.Close()
59+
60+
listener.Close()
61+
62+
if routed.Load() == 0 {
63+
t.Error("expected at least one request to be routed")
64+
}
65+
}

‎internal/redisqueue/plugin.go‎

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,11 +49,13 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
4949
requestID := strings.TrimSpace(internallogging.GetRequestID(ctx))
5050

5151
tokens := tokenStats{
52-
InputTokens: record.Detail.InputTokens,
53-
OutputTokens: record.Detail.OutputTokens,
54-
ReasoningTokens: record.Detail.ReasoningTokens,
55-
CachedTokens: record.Detail.CachedTokens,
56-
TotalTokens: record.Detail.TotalTokens,
52+
InputTokens: record.Detail.InputTokens,
53+
OutputTokens: record.Detail.OutputTokens,
54+
ReasoningTokens: record.Detail.ReasoningTokens,
55+
CachedTokens: record.Detail.CachedTokens,
56+
CacheReadTokens: record.Detail.CacheReadTokens,
57+
CacheCreationTokens: record.Detail.CacheCreationTokens,
58+
TotalTokens: record.Detail.TotalTokens,
5759
}
5860
if tokens.TotalTokens == 0 {
5961
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens
@@ -116,11 +118,13 @@ type requestDetail struct {
116118
}
117119

118120
type tokenStats struct {
119-
InputTokens int64 `json:"input_tokens"`
120-
OutputTokens int64 `json:"output_tokens"`
121-
ReasoningTokens int64 `json:"reasoning_tokens"`
122-
CachedTokens int64 `json:"cached_tokens"`
123-
TotalTokens int64 `json:"total_tokens"`
121+
InputTokens int64 `json:"input_tokens"`
122+
OutputTokens int64 `json:"output_tokens"`
123+
ReasoningTokens int64 `json:"reasoning_tokens"`
124+
CachedTokens int64 `json:"cached_tokens"`
125+
CacheReadTokens int64 `json:"cache_read_tokens"`
126+
CacheCreationTokens int64 `json:"cache_creation_tokens"`
127+
TotalTokens int64 `json:"total_tokens"`
124128
}
125129

126130
type failDetail struct {

‎internal/runtime/executor/helps/home_refresh.go‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ func RefreshAuthViaHome(ctx context.Context, cfg *config.Config, auth *cliproxya
7878
if msg == "" {
7979
msg = "home returned error"
8080
}
81-
return nil, true, homeStatusErr{code: http.StatusBadGateway, msg: msg}
81+
return nil, true, homeStatusErr{code: statusFromHomeErrorCode(code), msg: msg}
8282
}
8383

8484
var updated cliproxyauth.Auth
@@ -89,3 +89,14 @@ func RefreshAuthViaHome(ctx context.Context, cfg *config.Config, auth *cliproxya
8989
updated.EnsureIndex()
9090
return &updated, true, nil
9191
}
92+
93+
func statusFromHomeErrorCode(code string) int {
94+
switch strings.ToLower(strings.TrimSpace(code)) {
95+
case "authentication_error", "unauthorized":
96+
return http.StatusUnauthorized
97+
case "model_not_found":
98+
return http.StatusNotFound
99+
default:
100+
return http.StatusBadGateway
101+
}
102+
}

‎internal/runtime/executor/helps/home_refresh_test.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package helps
2+
3+
import (
4+
"net/http"
5+
"testing"
6+
)
7+
8+
func TestStatusFromHomeErrorCodeMapsAuthenticationErrorToUnauthorized(t *testing.T) {
9+
if got := statusFromHomeErrorCode("authentication_error"); got != http.StatusUnauthorized {
10+
t.Fatalf("statusFromHomeErrorCode(authentication_error) = %d, want %d", got, http.StatusUnauthorized)
11+
}
12+
if got := statusFromHomeErrorCode("unauthorized"); got != http.StatusUnauthorized {
13+
t.Fatalf("statusFromHomeErrorCode(unauthorized) = %d, want %d", got, http.StatusUnauthorized)
14+
}
15+
}

‎internal/runtime/executor/helps/usage_helpers.go‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,8 @@ func hasNonZeroTokenUsage(detail usage.Detail) bool {
116116
detail.OutputTokens != 0 ||
117117
detail.ReasoningTokens != 0 ||
118118
detail.CachedTokens != 0 ||
119+
detail.CacheReadTokens != 0 ||
120+
detail.CacheCreationTokens != 0 ||
119121
detail.TotalTokens != 0
120122
}
121123

@@ -356,17 +358,7 @@ func ParseClaudeUsage(data []byte) usage.Detail {
356358
if !usageNode.Exists() {
357359
return usage.Detail{}
358360
}
359-
detail := usage.Detail{
360-
InputTokens: usageNode.Get("input_tokens").Int(),
361-
OutputTokens: usageNode.Get("output_tokens").Int(),
362-
CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),
363-
}
364-
if detail.CachedTokens == 0 {
365-
// fall back to creation tokens when read tokens are absent
366-
detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()
367-
}
368-
detail.TotalTokens = detail.InputTokens + detail.OutputTokens
369-
return detail
361+
return parseClaudeUsageNode(usageNode)
370362
}
371363

372364
func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
@@ -378,16 +370,24 @@ func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
378370
if !usageNode.Exists() {
379371
return usage.Detail{}, false
380372
}
373+
return parseClaudeUsageNode(usageNode), true
374+
}
375+
376+
func parseClaudeUsageNode(usageNode gjson.Result) usage.Detail {
377+
cacheReadTokens := usageNode.Get("cache_read_input_tokens").Int()
378+
cacheCreationTokens := usageNode.Get("cache_creation_input_tokens").Int()
381379
detail := usage.Detail{
382-
InputTokens: usageNode.Get("input_tokens").Int(),
383-
OutputTokens: usageNode.Get("output_tokens").Int(),
384-
CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),
380+
InputTokens: usageNode.Get("input_tokens").Int(),
381+
OutputTokens: usageNode.Get("output_tokens").Int(),
382+
CachedTokens: cacheReadTokens,
383+
CacheReadTokens: cacheReadTokens,
384+
CacheCreationTokens: cacheCreationTokens,
385385
}
386386
if detail.CachedTokens == 0 {
387-
detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()
387+
detail.CachedTokens = detail.CacheCreationTokens
388388
}
389389
detail.TotalTokens = detail.InputTokens + detail.OutputTokens
390-
return detail, true
390+
return detail
391391
}
392392

393393
func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail {

‎sdk/cliproxy/auth/auto_refresh_loop.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,9 @@ func nextRefreshCheckAt(now time.Time, auth *Auth, interval time.Duration) (time
339339
if auth == nil {
340340
return time.Time{}, false
341341
}
342+
if hasUnauthorizedAuthFailure(auth) {
343+
return time.Time{}, false
344+
}
342345

343346
accountType, _ := auth.AccountInfo()
344347
if accountType == "api_key" {

0 commit comments

Comments
 (0)