feat(home): implement count for home auth dispatch requests and enable usage statistics

- Added `count` attribute to `homeAuthCount` requests to improve home message batching.
- Enabled usage statistics for home mode by default and added config-level enforcement.
- Adjusted failure logging to include detailed metadata in `UsageReporter`.
- Updated multiple executors to pass error details to `PublishFailure` for better debugging.
- Enhanced unit tests to validate `count` behavior and usage statistics enforcement across components.
This commit is contained in:
Luis Pater
2026-05-10 01:30:43 +08:00
parent 1abf8625d8
commit 66c3dae06b
21 changed files with 281 additions and 52 deletions
+1
View File
@@ -290,6 +290,7 @@ func main() {
} }
parsed.Home = homeCfg parsed.Home = homeCfg
parsed.Port = 8317 // Default to 8317 for home mode, can be overridden by home config parsed.Port = 8317 // Default to 8317 for home mode, can be overridden by home config
parsed.UsageStatisticsEnabled = true
cfg = parsed cfg = parsed
// Keep a non-empty config path for downstream components (log paths, management assets, etc), // Keep a non-empty config path for downstream components (log paths, management assets, etc),
+15 -7
View File
@@ -190,7 +190,20 @@ func headersToLowerMap(headers http.Header) map[string]string {
return out return out
} }
func (c *Client) RPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header) ([]byte, error) { func newAuthDispatchRequest(requestedModel string, sessionID string, headers http.Header, count int) authDispatchRequest {
if count <= 0 {
count = 1
}
return authDispatchRequest{
Type: "auth",
Model: requestedModel,
Count: count,
SessionID: strings.TrimSpace(sessionID),
Headers: headersToLowerMap(headers),
}
}
func (c *Client) RPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int) ([]byte, error) {
if err := c.ensureClients(); err != nil { if err := c.ensureClients(); err != nil {
return nil, err return nil, err
} }
@@ -198,12 +211,7 @@ func (c *Client) RPopAuth(ctx context.Context, requestedModel string, sessionID
if requestedModel == "" { if requestedModel == "" {
return nil, fmt.Errorf("home: requested model is empty") return nil, fmt.Errorf("home: requested model is empty")
} }
req := authDispatchRequest{ req := newAuthDispatchRequest(requestedModel, sessionID, headers, count)
Type: "auth",
Model: requestedModel,
SessionID: strings.TrimSpace(sessionID),
Headers: headersToLowerMap(headers),
}
keyBytes, err := json.Marshal(&req) keyBytes, err := json.Marshal(&req)
if err != nil { if err != nil {
return nil, err return nil, err
+32
View File
@@ -0,0 +1,32 @@
package home
import (
"encoding/json"
"net/http"
"testing"
)
func TestAuthDispatchRequestIncludesCount(t *testing.T) {
req := newAuthDispatchRequest("gpt-5.4", "session-1", http.Header{"Authorization": {"Bearer test"}}, 2)
raw, err := json.Marshal(&req)
if err != nil {
t.Fatalf("marshal auth dispatch request: %v", err)
}
var payload map[string]any
if err := json.Unmarshal(raw, &payload); err != nil {
t.Fatalf("unmarshal auth dispatch request: %v", err)
}
if got := int(payload["count"].(float64)); got != 2 {
t.Fatalf("count = %d, want 2", got)
}
}
func TestAuthDispatchRequestDefaultsCountToOne(t *testing.T) {
req := newAuthDispatchRequest("gpt-5.4", "", nil, 0)
if req.Count != 1 {
t.Fatalf("count = %d, want 1", req.Count)
}
}
+1
View File
@@ -3,6 +3,7 @@ package home
type authDispatchRequest struct { type authDispatchRequest struct {
Type string `json:"type"` Type string `json:"type"`
Model string `json:"model"` Model string `json:"model"`
Count int `json:"count"`
SessionID string `json:"session_id,omitempty"` SessionID string `json:"session_id,omitempty"`
Headers map[string]string `json:"headers,omitempty"` Headers map[string]string `json:"headers,omitempty"`
} }
+25
View File
@@ -66,6 +66,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
if !failed { if !failed {
failed = !resolveSuccess(ctx) failed = !resolveSuccess(ctx)
} }
fail := resolveFail(ctx, record, failed)
detail := requestDetail{ detail := requestDetail{
Timestamp: timestamp, Timestamp: timestamp,
@@ -74,6 +75,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
AuthIndex: record.AuthIndex, AuthIndex: record.AuthIndex,
Tokens: tokens, Tokens: tokens,
Failed: failed, Failed: failed,
Fail: fail,
} }
payload, err := json.Marshal(queuedUsageDetail{ payload, err := json.Marshal(queuedUsageDetail{
@@ -110,6 +112,7 @@ type requestDetail struct {
AuthIndex string `json:"auth_index"` AuthIndex string `json:"auth_index"`
Tokens tokenStats `json:"tokens"` Tokens tokenStats `json:"tokens"`
Failed bool `json:"failed"` Failed bool `json:"failed"`
Fail failDetail `json:"fail"`
} }
type tokenStats struct { type tokenStats struct {
@@ -120,6 +123,28 @@ type tokenStats struct {
TotalTokens int64 `json:"total_tokens"` TotalTokens int64 `json:"total_tokens"`
} }
type failDetail struct {
StatusCode int `json:"status_code"`
Body string `json:"body"`
}
func resolveFail(ctx context.Context, record coreusage.Record, failed bool) failDetail {
fail := failDetail{
StatusCode: record.Fail.StatusCode,
Body: strings.TrimSpace(record.Fail.Body),
}
if !failed {
return failDetail{StatusCode: 200}
}
if fail.StatusCode <= 0 {
fail.StatusCode = internallogging.GetResponseStatus(ctx)
}
if fail.StatusCode <= 0 {
fail.StatusCode = 500
}
return fail
}
func resolveSuccess(ctx context.Context) bool { func resolveSuccess(ctx context.Context) bool {
status := internallogging.GetResponseStatus(ctx) status := internallogging.GetResponseStatus(ctx)
if status == 0 { if status == 0 {
+41 -3
View File
@@ -44,9 +44,10 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
requireStringField(t, payload, "alias", "client-gpt") requireStringField(t, payload, "alias", "client-gpt")
requireStringField(t, payload, "endpoint", "POST /v1/chat/completions") requireStringField(t, payload, "endpoint", "POST /v1/chat/completions")
requireStringField(t, payload, "auth_type", "apikey") requireStringField(t, payload, "auth_type", "apikey")
requireStringField(t, payload, "user_api_key", "test-key") requireMissingField(t, payload, "user_api_key")
requireStringField(t, payload, "request_id", "ctx-request-id") requireStringField(t, payload, "request_id", "ctx-request-id")
requireBoolField(t, payload, "failed", false) requireBoolField(t, payload, "failed", false)
requireFailField(t, payload, http.StatusOK, "")
}) })
} }
@@ -68,6 +69,10 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndFailureAndGinRequestID(t
Source: "user@example.com", Source: "user@example.com",
RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC),
Latency: 2500 * time.Millisecond, Latency: 2500 * time.Millisecond,
Fail: coreusage.Failure{
StatusCode: http.StatusInternalServerError,
Body: "upstream failed",
},
Detail: coreusage.Detail{ Detail: coreusage.Detail{
InputTokens: 10, InputTokens: 10,
OutputTokens: 20, OutputTokens: 20,
@@ -81,9 +86,10 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndFailureAndGinRequestID(t
requireStringField(t, payload, "alias", "client-mini") requireStringField(t, payload, "alias", "client-mini")
requireStringField(t, payload, "endpoint", "GET /v1/responses") requireStringField(t, payload, "endpoint", "GET /v1/responses")
requireStringField(t, payload, "auth_type", "apikey") requireStringField(t, payload, "auth_type", "apikey")
requireStringField(t, payload, "user_api_key", "test-key") requireMissingField(t, payload, "user_api_key")
requireStringField(t, payload, "request_id", "gin-request-id") requireStringField(t, payload, "request_id", "gin-request-id")
requireBoolField(t, payload, "failed", true) requireBoolField(t, payload, "failed", true)
requireFailField(t, payload, http.StatusInternalServerError, "upstream failed")
}) })
} }
@@ -115,6 +121,10 @@ func TestUsageQueuePluginAsyncIgnoresRecycledGinContext(t *testing.T) {
Source: "user@example.com", Source: "user@example.com",
RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC),
Latency: 1500 * time.Millisecond, Latency: 1500 * time.Millisecond,
Fail: coreusage.Failure{
StatusCode: http.StatusBadGateway,
Body: "bad gateway",
},
Detail: coreusage.Detail{ Detail: coreusage.Detail{
InputTokens: 10, InputTokens: 10,
OutputTokens: 20, OutputTokens: 20,
@@ -125,9 +135,10 @@ func TestUsageQueuePluginAsyncIgnoresRecycledGinContext(t *testing.T) {
payload := waitForSinglePayload(t, 2*time.Second) payload := waitForSinglePayload(t, 2*time.Second)
requireStringField(t, payload, "endpoint", "POST /v1/chat/completions") requireStringField(t, payload, "endpoint", "POST /v1/chat/completions")
requireStringField(t, payload, "alias", "client-gpt") requireStringField(t, payload, "alias", "client-gpt")
requireStringField(t, payload, "user_api_key", "test-key") requireMissingField(t, payload, "user_api_key")
requireStringField(t, payload, "request_id", "ctx-request-id") requireStringField(t, payload, "request_id", "ctx-request-id")
requireBoolField(t, payload, "failed", true) requireBoolField(t, payload, "failed", true)
requireFailField(t, payload, http.StatusBadGateway, "bad gateway")
}) })
} }
@@ -217,6 +228,14 @@ func requireStringField(t *testing.T, payload map[string]json.RawMessage, key, w
} }
} }
func requireMissingField(t *testing.T, payload map[string]json.RawMessage, key string) {
t.Helper()
if _, ok := payload[key]; ok {
t.Fatalf("payload unexpectedly contains %q", key)
}
}
type pluginFunc func(context.Context, coreusage.Record) type pluginFunc func(context.Context, coreusage.Record)
func (fn pluginFunc) HandleUsage(ctx context.Context, record coreusage.Record) { func (fn pluginFunc) HandleUsage(ctx context.Context, record coreusage.Record) {
@@ -238,3 +257,22 @@ func requireBoolField(t *testing.T, payload map[string]json.RawMessage, key stri
t.Fatalf("%s = %t, want %t", key, got, want) t.Fatalf("%s = %t, want %t", key, got, want)
} }
} }
func requireFailField(t *testing.T, payload map[string]json.RawMessage, wantStatus int, wantBody string) {
t.Helper()
raw, ok := payload["fail"]
if !ok {
t.Fatalf("payload missing %q", "fail")
}
var got struct {
StatusCode int `json:"status_code"`
Body string `json:"body"`
}
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("unmarshal fail: %v", err)
}
if got.StatusCode != wantStatus || got.Body != wantBody {
t.Fatalf("fail = {status_code:%d body:%q}, want {status_code:%d body:%q}", got.StatusCode, got.Body, wantStatus, wantBody)
}
}
@@ -284,7 +284,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
processEvent := func(event wsrelay.StreamEvent) bool { processEvent := func(event wsrelay.StreamEvent) bool {
if event.Err != nil { if event.Err != nil {
helps.RecordAPIResponseError(ctx, e.cfg, event.Err) helps.RecordAPIResponseError(ctx, e.cfg, event.Err)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, event.Err)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("wsrelay: %v", event.Err)}: case out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("wsrelay: %v", event.Err)}:
case <-ctx.Done(): case <-ctx.Done():
@@ -336,7 +336,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
return false return false
case wsrelay.MessageTypeError: case wsrelay.MessageTypeError:
helps.RecordAPIResponseError(ctx, e.cfg, event.Err) helps.RecordAPIResponseError(ctx, e.cfg, event.Err)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, event.Err)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("wsrelay: %v", event.Err)}: case out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("wsrelay: %v", event.Err)}:
case <-ctx.Done(): case <-ctx.Done():
@@ -898,7 +898,7 @@ attemptLoop:
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
out <- cliproxyexecutor.StreamChunk{Err: errScan} out <- cliproxyexecutor.StreamChunk{Err: errScan}
} else { } else {
reporter.EnsurePublished(ctx) reporter.EnsurePublished(ctx)
@@ -1374,7 +1374,7 @@ attemptLoop:
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done(): case <-ctx.Done():
+2 -2
View File
@@ -472,7 +472,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done(): case <-ctx.Done():
@@ -512,7 +512,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done(): case <-ctx.Done():
+1 -1
View File
@@ -524,7 +524,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done(): case <-ctx.Done():
@@ -580,7 +580,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
terminateReason = "read_error" terminateReason = "read_error"
terminateErr = errRead terminateErr = errRead
helps.RecordAPIWebsocketError(ctx, e.cfg, "read", errRead) helps.RecordAPIWebsocketError(ctx, e.cfg, "read", errRead)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errRead)
_ = send(cliproxyexecutor.StreamChunk{Err: errRead}) _ = send(cliproxyexecutor.StreamChunk{Err: errRead})
return return
} }
@@ -590,7 +590,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
terminateReason = "unexpected_binary" terminateReason = "unexpected_binary"
terminateErr = err terminateErr = err
helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", err) helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", err)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, err)
if sess != nil { if sess != nil {
e.invalidateUpstreamConn(sess, conn, "unexpected_binary", err) e.invalidateUpstreamConn(sess, conn, "unexpected_binary", err)
} }
@@ -610,7 +610,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
terminateReason = "upstream_error" terminateReason = "upstream_error"
terminateErr = wsErr terminateErr = wsErr
helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, wsErr)
if sess != nil { if sess != nil {
e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr)
} }
@@ -430,7 +430,7 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done(): case <-ctx.Done():
@@ -444,7 +444,7 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
data, errRead := io.ReadAll(resp.Body) data, errRead := io.ReadAll(resp.Body)
if errRead != nil { if errRead != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errRead) helps.RecordAPIResponseError(ctx, e.cfg, errRead)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errRead)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errRead}: case out <- cliproxyexecutor.StreamChunk{Err: errRead}:
case <-ctx.Done(): case <-ctx.Done():
+1 -1
View File
@@ -341,7 +341,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done(): case <-ctx.Done():
@@ -679,7 +679,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done(): case <-ctx.Done():
@@ -821,7 +821,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done(): case <-ctx.Done():
@@ -3,6 +3,7 @@ package helps
import ( import (
"bytes" "bytes"
"context" "context"
"errors"
"fmt" "fmt"
"strings" "strings"
"sync" "sync"
@@ -51,7 +52,7 @@ func NewUsageReporter(ctx context.Context, provider, model string, auth *cliprox
} }
func (r *UsageReporter) Publish(ctx context.Context, detail usage.Detail) { func (r *UsageReporter) Publish(ctx context.Context, detail usage.Detail) {
r.publishWithOutcome(ctx, detail, false) r.publishWithOutcome(ctx, detail, false, usage.Failure{})
} }
func (r *UsageReporter) PublishAdditionalModel(ctx context.Context, model string, detail usage.Detail) { func (r *UsageReporter) PublishAdditionalModel(ctx context.Context, model string, detail usage.Detail) {
@@ -74,11 +75,11 @@ func (r *UsageReporter) buildAdditionalModelRecord(model string, detail usage.De
if !hasNonZeroTokenUsage(detail) { if !hasNonZeroTokenUsage(detail) {
return usage.Record{}, false return usage.Record{}, false
} }
return r.buildRecordForModel(model, detail, false), true return r.buildRecordForModel(model, detail, false, usage.Failure{}), true
} }
func (r *UsageReporter) PublishFailure(ctx context.Context) { func (r *UsageReporter) PublishFailure(ctx context.Context, errs ...error) {
r.publishWithOutcome(ctx, usage.Detail{}, true) r.publishWithOutcome(ctx, usage.Detail{}, true, failFromErrors(errs...))
} }
func (r *UsageReporter) TrackFailure(ctx context.Context, errPtr *error) { func (r *UsageReporter) TrackFailure(ctx context.Context, errPtr *error) {
@@ -86,17 +87,17 @@ func (r *UsageReporter) TrackFailure(ctx context.Context, errPtr *error) {
return return
} }
if *errPtr != nil { if *errPtr != nil {
r.PublishFailure(ctx) r.PublishFailure(ctx, *errPtr)
} }
} }
func (r *UsageReporter) publishWithOutcome(ctx context.Context, detail usage.Detail, failed bool) { func (r *UsageReporter) publishWithOutcome(ctx context.Context, detail usage.Detail, failed bool, fail usage.Failure) {
if r == nil { if r == nil {
return return
} }
detail = normalizeUsageDetailTotal(detail) detail = normalizeUsageDetailTotal(detail)
r.once.Do(func() { r.once.Do(func() {
usage.PublishRecord(ctx, r.buildRecord(detail, failed)) usage.PublishRecord(ctx, r.buildRecord(detail, failed, fail))
}) })
} }
@@ -127,20 +128,24 @@ func (r *UsageReporter) EnsurePublished(ctx context.Context) {
return return
} }
r.once.Do(func() { r.once.Do(func() {
usage.PublishRecord(ctx, r.buildRecord(usage.Detail{}, false)) usage.PublishRecord(ctx, r.buildRecord(usage.Detail{}, false, usage.Failure{}))
}) })
} }
func (r *UsageReporter) buildRecord(detail usage.Detail, failed bool) usage.Record { func (r *UsageReporter) buildRecord(detail usage.Detail, failed bool, failures ...usage.Failure) usage.Record {
if r == nil { var fail usage.Failure
return usage.Record{Detail: detail, Failed: failed} if len(failures) > 0 {
fail = failures[0]
} }
return r.buildRecordForModel(r.model, detail, failed) if r == nil {
return usage.Record{Detail: detail, Failed: failed, Fail: fail}
}
return r.buildRecordForModel(r.model, detail, failed, fail)
} }
func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, failed bool) usage.Record { func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, failed bool, fail usage.Failure) usage.Record {
if r == nil { if r == nil {
return usage.Record{Model: model, Detail: detail, Failed: failed} return usage.Record{Model: model, Detail: detail, Failed: failed, Fail: fail}
} }
return usage.Record{ return usage.Record{
Provider: r.provider, Provider: r.provider,
@@ -154,10 +159,28 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f
RequestedAt: r.requestedAt, RequestedAt: r.requestedAt,
Latency: r.latency(), Latency: r.latency(),
Failed: failed, Failed: failed,
Fail: fail,
Detail: detail, Detail: detail,
} }
} }
func failFromErrors(errs ...error) usage.Failure {
for _, err := range errs {
if err == nil {
continue
}
fail := usage.Failure{
Body: strings.TrimSpace(err.Error()),
}
var se interface{ StatusCode() int }
if errors.As(err, &se) && se != nil {
fail.StatusCode = se.StatusCode()
}
return fail
}
return usage.Failure{}
}
func (r *UsageReporter) latency() time.Duration { func (r *UsageReporter) latency() time.Duration {
if r == nil || r.requestedAt.IsZero() { if r == nil || r.requestedAt.IsZero() {
return 0 return 0
+1 -1
View File
@@ -307,7 +307,7 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done(): case <-ctx.Done():
@@ -296,7 +296,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
if bytes.HasPrefix(trimmedLine, []byte("{")) || bytes.HasPrefix(trimmedLine, []byte("[")) { if bytes.HasPrefix(trimmedLine, []byte("{")) || bytes.HasPrefix(trimmedLine, []byte("[")) {
streamErr := statusErr{code: http.StatusBadGateway, msg: string(trimmedLine)} streamErr := statusErr{code: http.StatusBadGateway, msg: string(trimmedLine)}
helps.RecordAPIResponseError(ctx, e.cfg, streamErr) helps.RecordAPIResponseError(ctx, e.cfg, streamErr)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, streamErr)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: streamErr}: case out <- cliproxyexecutor.StreamChunk{Err: streamErr}:
case <-ctx.Done(): case <-ctx.Done():
@@ -318,7 +318,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
} }
if errScan := scanner.Err(); errScan != nil { if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan) helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx) reporter.PublishFailure(ctx, errScan)
select { select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done(): case <-ctx.Done():
+73 -10
View File
@@ -51,6 +51,7 @@ type ExecutionSessionCloser interface {
} }
const ( const (
homeAuthCountMetadataKey = "__cliproxy_home_auth_count"
// CloseAllExecutionSessionsID asks an executor to release all active execution sessions. // CloseAllExecutionSessionsID asks an executor to release all active execution sessions.
// Executors that do not support this marker may ignore it. // Executors that do not support this marker may ignore it.
CloseAllExecutionSessionsID = "__all_execution_sessions__" CloseAllExecutionSessionsID = "__all_execution_sessions__"
@@ -1316,19 +1317,25 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req
} }
routeModel := req.Model routeModel := req.Model
opts = ensureRequestedModelMetadata(opts, routeModel) opts = ensureRequestedModelMetadata(opts, routeModel)
homeMode := m.HomeEnabled()
homeAuthCount := 1
tried := make(map[string]struct{}) tried := make(map[string]struct{})
attempted := make(map[string]struct{}) attempted := make(map[string]struct{})
var lastErr error var lastErr error
for { for {
if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
if lastErr != nil { if lastErr != nil {
return cliproxyexecutor.Response{}, lastErr return cliproxyexecutor.Response{}, lastErr
} }
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
} }
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried) pickOpts := opts
if homeMode {
pickOpts = withHomeAuthCount(opts, homeAuthCount)
}
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried)
if errPick != nil { if errPick != nil {
if lastErr != nil { if lastErr != nil && !homeMode {
return cliproxyexecutor.Response{}, lastErr return cliproxyexecutor.Response{}, lastErr
} }
return cliproxyexecutor.Response{}, errPick return cliproxyexecutor.Response{}, errPick
@@ -1384,6 +1391,9 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req
return cliproxyexecutor.Response{}, authErr return cliproxyexecutor.Response{}, authErr
} }
lastErr = authErr lastErr = authErr
if homeMode {
homeAuthCount++
}
continue continue
} }
} }
@@ -1395,19 +1405,25 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string,
} }
routeModel := req.Model routeModel := req.Model
opts = ensureRequestedModelMetadata(opts, routeModel) opts = ensureRequestedModelMetadata(opts, routeModel)
homeMode := m.HomeEnabled()
homeAuthCount := 1
tried := make(map[string]struct{}) tried := make(map[string]struct{})
attempted := make(map[string]struct{}) attempted := make(map[string]struct{})
var lastErr error var lastErr error
for { for {
if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
if lastErr != nil { if lastErr != nil {
return cliproxyexecutor.Response{}, lastErr return cliproxyexecutor.Response{}, lastErr
} }
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
} }
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried) pickOpts := opts
if homeMode {
pickOpts = withHomeAuthCount(opts, homeAuthCount)
}
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried)
if errPick != nil { if errPick != nil {
if lastErr != nil { if lastErr != nil && !homeMode {
return cliproxyexecutor.Response{}, lastErr return cliproxyexecutor.Response{}, lastErr
} }
return cliproxyexecutor.Response{}, errPick return cliproxyexecutor.Response{}, errPick
@@ -1463,6 +1479,9 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string,
return cliproxyexecutor.Response{}, authErr return cliproxyexecutor.Response{}, authErr
} }
lastErr = authErr lastErr = authErr
if homeMode {
homeAuthCount++
}
continue continue
} }
} }
@@ -1474,19 +1493,25 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string
} }
routeModel := req.Model routeModel := req.Model
opts = ensureRequestedModelMetadata(opts, routeModel) opts = ensureRequestedModelMetadata(opts, routeModel)
homeMode := m.HomeEnabled()
homeAuthCount := 1
tried := make(map[string]struct{}) tried := make(map[string]struct{})
attempted := make(map[string]struct{}) attempted := make(map[string]struct{})
var lastErr error var lastErr error
for { for {
if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
if lastErr != nil { if lastErr != nil {
return nil, lastErr return nil, lastErr
} }
return nil, &Error{Code: "auth_not_found", Message: "no auth available"} return nil, &Error{Code: "auth_not_found", Message: "no auth available"}
} }
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried) pickOpts := opts
if homeMode {
pickOpts = withHomeAuthCount(opts, homeAuthCount)
}
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried)
if errPick != nil { if errPick != nil {
if lastErr != nil { if lastErr != nil && !homeMode {
return nil, lastErr return nil, lastErr
} }
return nil, errPick return nil, errPick
@@ -1516,6 +1541,9 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string
return nil, errStream return nil, errStream
} }
lastErr = errStream lastErr = errStream
if homeMode {
homeAuthCount++
}
continue continue
} }
return streamResult, nil return streamResult, nil
@@ -1543,6 +1571,40 @@ func ensureRequestedModelMetadata(opts cliproxyexecutor.Options, requestedModel
return opts return opts
} }
func withHomeAuthCount(opts cliproxyexecutor.Options, count int) cliproxyexecutor.Options {
if count <= 0 {
count = 1
}
meta := make(map[string]any, len(opts.Metadata)+1)
for k, v := range opts.Metadata {
meta[k] = v
}
meta[homeAuthCountMetadataKey] = count
opts.Metadata = meta
return opts
}
func homeAuthCountFromMetadata(meta map[string]any) int {
if len(meta) == 0 {
return 1
}
switch value := meta[homeAuthCountMetadataKey].(type) {
case int:
if value > 0 {
return value
}
case int64:
if value > 0 {
return int(value)
}
case float64:
if value > 0 {
return int(value)
}
}
return 1
}
func hasRequestedModelMetadata(meta map[string]any) bool { func hasRequestedModelMetadata(meta map[string]any) bool {
if len(meta) == 0 { if len(meta) == 0 {
return false return false
@@ -3099,8 +3161,9 @@ func (m *Manager) pickNextViaHome(ctx context.Context, model string, opts clipro
requestedModel := requestedModelFromMetadata(opts.Metadata, model) requestedModel := requestedModelFromMetadata(opts.Metadata, model)
sessionID := ExtractSessionID(opts.Headers, opts.OriginalRequest, opts.Metadata) sessionID := ExtractSessionID(opts.Headers, opts.OriginalRequest, opts.Metadata)
count := homeAuthCountFromMetadata(opts.Metadata)
raw, err := client.RPopAuth(ctx, requestedModel, sessionID, opts.Headers) raw, err := client.RPopAuth(ctx, requestedModel, sessionID, opts.Headers, count)
if err != nil { if err != nil {
return nil, nil, "", &Error{Code: "auth_not_found", Message: err.Error(), HTTPStatus: http.StatusServiceUnavailable} return nil, nil, "", &Error{Code: "auth_not_found", Message: err.Error(), HTTPStatus: http.StatusServiceUnavailable}
} }
+2
View File
@@ -561,6 +561,7 @@ func forceHomeRuntimeConfig(cfg *config.Config) {
return return
} }
cfg.APIKeys = nil cfg.APIKeys = nil
cfg.UsageStatisticsEnabled = true
cfg.DisableCooling = true cfg.DisableCooling = true
cfg.WebsocketAuth = false cfg.WebsocketAuth = false
cfg.EnableGeminiCLIEndpoint = false cfg.EnableGeminiCLIEndpoint = false
@@ -732,6 +733,7 @@ func (s *Service) Run(ctx context.Context) error {
homeEnabled := s.cfg != nil && s.cfg.Home.Enabled homeEnabled := s.cfg != nil && s.cfg.Home.Enabled
if homeEnabled { if homeEnabled {
forceHomeRuntimeConfig(s.cfg) forceHomeRuntimeConfig(s.cfg)
redisqueue.SetUsageStatisticsEnabled(true)
} }
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
+29
View File
@@ -99,3 +99,32 @@ func TestServiceApplyCoreAuthAddOrUpdate_DeleteReAddDoesNotInheritStaleRuntimeSt
t.Fatalf("expected re-added auth to re-register models in global registry") t.Fatalf("expected re-added auth to re-register models in global registry")
} }
} }
func TestForceHomeRuntimeConfigEnablesUsageStatistics(t *testing.T) {
cfg := &config.Config{
UsageStatisticsEnabled: false,
}
forceHomeRuntimeConfig(cfg)
if !cfg.UsageStatisticsEnabled {
t.Fatal("expected home runtime config to force usage statistics enabled")
}
}
func TestApplyHomeOverlayForcesUsageStatisticsEnabled(t *testing.T) {
baseCfg := &config.Config{}
baseCfg.Home.Enabled = true
service := &Service{cfg: baseCfg}
service.applyHomeOverlay(&config.Config{
UsageStatisticsEnabled: false,
})
if service.cfg == nil || !service.cfg.UsageStatisticsEnabled {
t.Fatal("expected home overlay to force usage statistics enabled")
}
if !service.cfg.Home.Enabled {
t.Fatal("expected home overlay to preserve local home settings")
}
}
+7
View File
@@ -22,9 +22,16 @@ type Record struct {
RequestedAt time.Time RequestedAt time.Time
Latency time.Duration Latency time.Duration
Failed bool Failed bool
Fail Failure
Detail Detail Detail Detail
} }
// Failure holds HTTP failure metadata for an upstream request attempt.
type Failure struct {
StatusCode int
Body string
}
// Detail holds the token usage breakdown. // Detail holds the token usage breakdown.
type Detail struct { type Detail struct {
InputTokens int64 InputTokens int64