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:
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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"`
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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():
|
||||||
|
|||||||
@@ -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():
|
||||||
|
|||||||
@@ -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():
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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():
|
||||||
|
|||||||
@@ -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}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user