airkjw/CLIProxyAPI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files
internal/runtime/executor/aistudio_executor.go
+2 -2
View File
@@ -284,7 +284,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
processEvent := func(event wsrelay.StreamEvent) bool {
if event.Err != nil {
helps.RecordAPIResponseError(ctx, e.cfg, event.Err)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, event.Err)
select {
case out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("wsrelay: %v", event.Err)}:
case <-ctx.Done():
@@ -336,7 +336,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
return false
case wsrelay.MessageTypeError:
helps.RecordAPIResponseError(ctx, e.cfg, event.Err)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, event.Err)
select {
case out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("wsrelay: %v", event.Err)}:
case <-ctx.Done():
internal/runtime/executor/antigravity_executor.go
+2 -2
View File
@@ -898,7 +898,7 @@ attemptLoop:
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
out <- cliproxyexecutor.StreamChunk{Err: errScan}
} else {
reporter.EnsurePublished(ctx)
@@ -1374,7 +1374,7 @@ attemptLoop:
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
internal/runtime/executor/claude_executor.go
+2 -2
View File
@@ -472,7 +472,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
@@ -512,7 +512,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
internal/runtime/executor/codex_executor.go
+1 -1
View File
@@ -524,7 +524,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
internal/runtime/executor/codex_websockets_executor.go
+3 -3
View File
@@ -580,7 +580,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
terminateReason = "read_error"
terminateErr = errRead
helps.RecordAPIWebsocketError(ctx, e.cfg, "read", errRead)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errRead)
_ = send(cliproxyexecutor.StreamChunk{Err: errRead})
return
}
@@ -590,7 +590,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
terminateReason = "unexpected_binary"
terminateErr = err
helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", err)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, err)
if sess != nil {
e.invalidateUpstreamConn(sess, conn, "unexpected_binary", err)
}
@@ -610,7 +610,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
terminateReason = "upstream_error"
terminateErr = wsErr
helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, wsErr)
if sess != nil {
e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr)
}
internal/runtime/executor/gemini_cli_executor.go
+2 -2
View File
@@ -430,7 +430,7 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
@@ -444,7 +444,7 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
data, errRead := io.ReadAll(resp.Body)
if errRead != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errRead)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errRead)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errRead}:
case <-ctx.Done():
internal/runtime/executor/gemini_executor.go
+1 -1
View File
@@ -341,7 +341,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
internal/runtime/executor/gemini_vertex_executor.go
+2 -2
View File
@@ -679,7 +679,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
@@ -821,7 +821,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
internal/runtime/executor/helps/usage_helpers.go
+37 -14
View File
@@ -3,6 +3,7 @@ package helps
import (
"bytes"
"context"
"errors"
"fmt"
"strings"
"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) {
r.publishWithOutcome(ctx, detail, false)
r.publishWithOutcome(ctx, detail, false, usage.Failure{})
}
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) {
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) {
r.publishWithOutcome(ctx, usage.Detail{}, true)
func (r *UsageReporter) PublishFailure(ctx context.Context, errs ...error) {
r.publishWithOutcome(ctx, usage.Detail{}, true, failFromErrors(errs...))
}
func (r *UsageReporter) TrackFailure(ctx context.Context, errPtr *error) {
@@ -86,17 +87,17 @@ func (r *UsageReporter) TrackFailure(ctx context.Context, errPtr *error) {
return
}
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 {
return
}
detail = normalizeUsageDetailTotal(detail)
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
}
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 {
if r == nil {
return usage.Record{Detail: detail, Failed: failed}
func (r *UsageReporter) buildRecord(detail usage.Detail, failed bool, failures ...usage.Failure) usage.Record {
var fail usage.Failure
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 {
return usage.Record{Model: model, Detail: detail, Failed: failed}
return usage.Record{Model: model, Detail: detail, Failed: failed, Fail: fail}
}
return usage.Record{
Provider: r.provider,
@@ -154,10 +159,28 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f
RequestedAt: r.requestedAt,
Latency: r.latency(),
Failed: failed,
Fail: fail,
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 {
if r == nil || r.requestedAt.IsZero() {
return 0
internal/runtime/executor/kimi_executor.go
+1 -1
View File
@@ -307,7 +307,7 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
internal/runtime/executor/openai_compat_executor.go
+2 -2
View File
@@ -296,7 +296,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
if bytes.HasPrefix(trimmedLine, []byte("{")) || bytes.HasPrefix(trimmedLine, []byte("[")) {
streamErr := statusErr{code: http.StatusBadGateway, msg: string(trimmedLine)}
helps.RecordAPIResponseError(ctx, e.cfg, streamErr)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, streamErr)
select {
case out <- cliproxyexecutor.StreamChunk{Err: streamErr}:
case <-ctx.Done():
@@ -318,7 +318,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
reporter.PublishFailure(ctx, errScan)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():