Skip to content

Commit 959067e

Browse files
committed
feat(usage): introduce executor type tracking in usage reporting
- Replaced `NewUsageReporter` with `NewExecutorUsageReporter` to include executor type in usage records. - Updated all executors to use the new reporter implementation. - Extended `UsageReporter` to track and publish executor type. - Added tests to validate proper executor type recording and handling. - Enhanced RedisQueue plugin and payload schema with executor type support.
1 parent 05b9724 commit 959067e

17 files changed

Lines changed: 110 additions & 54 deletions

internal/redisqueue/plugin.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,10 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
4242
if provider == "" {
4343
provider = "unknown"
4444
}
45+
executorType := strings.TrimSpace(record.ExecutorType)
46+
if executorType == "" {
47+
executorType = "unknown"
48+
}
4549
authType := strings.TrimSpace(record.AuthType)
4650
if authType == "" {
4751
authType = "unknown"
@@ -94,6 +98,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
9498
payload, err := json.Marshal(queuedUsageDetail{
9599
requestDetail: detail,
96100
Provider: provider,
101+
ExecutorType: executorType,
97102
Model: modelName,
98103
Alias: aliasName,
99104
Endpoint: resolveEndpoint(ctx),
@@ -112,6 +117,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
112117
type queuedUsageDetail struct {
113118
requestDetail
114119
Provider string `json:"provider"`
120+
ExecutorType string `json:"executor_type"`
115121
Model string `json:"model"`
116122
Alias string `json:"alias"`
117123
Endpoint string `json:"endpoint"`

internal/redisqueue/plugin_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
2626
plugin := &usageQueuePlugin{}
2727
plugin.HandleUsage(ctx, coreusage.Record{
2828
Provider: "openai",
29+
ExecutorType: "KimiExecutor",
2930
Model: "gpt-5.4",
3031
Alias: "client-gpt",
3132
APIKey: "test-key",
@@ -47,6 +48,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
4748

4849
payload := popSinglePayload(t)
4950
requireStringField(t, payload, "provider", "openai")
51+
requireStringField(t, payload, "executor_type", "KimiExecutor")
5052
requireStringField(t, payload, "model", "gpt-5.4")
5153
requireStringField(t, payload, "alias", "client-gpt")
5254
requireStringField(t, payload, "endpoint", "POST /v1/chat/completions")

internal/runtime/executor/aistudio_executor.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ func (e *AIStudioExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth,
128128
return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"}
129129
}
130130
baseModel := thinking.ParseSuffix(req.Model).ModelName
131-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
131+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
132132
defer reporter.TrackFailure(ctx, &err)
133133

134134
translatedReq, body, err := e.translateRequest(req, opts, false)
@@ -196,7 +196,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
196196
return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"}
197197
}
198198
baseModel := thinking.ParseSuffix(req.Model).ModelName
199-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
199+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
200200
defer reporter.TrackFailure(ctx, &err)
201201

202202
translatedReq, body, err := e.translateRequest(req, opts, true)

internal/runtime/executor/antigravity_executor.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -530,7 +530,7 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
530530
return e.executeClaudeNonStream(ctx, auth, req, opts)
531531
}
532532

533-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
533+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
534534
defer reporter.TrackFailure(ctx, &err)
535535

536536
from := opts.SourceFormat
@@ -730,7 +730,7 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
730730
return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
731731
}
732732

733-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
733+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
734734
defer reporter.TrackFailure(ctx, &err)
735735

736736
from := opts.SourceFormat
@@ -1192,7 +1192,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
11921192
return nil, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
11931193
}
11941194

1195-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
1195+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
11961196
defer reporter.TrackFailure(ctx, &err)
11971197

11981198
from := opts.SourceFormat

internal/runtime/executor/claude_executor.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
171171
baseURL = "https://api.anthropic.com"
172172
}
173173

174-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
174+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
175175
defer reporter.TrackFailure(ctx, &err)
176176
from := opts.SourceFormat
177177
to := sdktranslator.FromString("claude")
@@ -354,7 +354,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
354354
baseURL = "https://api.anthropic.com"
355355
}
356356

357-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
357+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
358358
defer reporter.TrackFailure(ctx, &err)
359359
from := opts.SourceFormat
360360
to := sdktranslator.FromString("claude")

internal/runtime/executor/codex_executor.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -264,7 +264,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re
264264
baseURL = "https://chatgpt.com/backend-api/codex"
265265
}
266266

267-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
267+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
268268
defer reporter.TrackFailure(ctx, &err)
269269

270270
from := opts.SourceFormat
@@ -431,7 +431,7 @@ func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.A
431431
baseURL = "https://chatgpt.com/backend-api/codex"
432432
}
433433

434-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
434+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
435435
defer reporter.TrackFailure(ctx, &err)
436436

437437
from := opts.SourceFormat
@@ -536,7 +536,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
536536
baseURL = "https://chatgpt.com/backend-api/codex"
537537
}
538538

539-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
539+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
540540
defer reporter.TrackFailure(ctx, &err)
541541

542542
from := opts.SourceFormat

internal/runtime/executor/codex_openai_images.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ func (e *CodexExecutor) executeOpenAIImage(ctx context.Context, auth *cliproxyau
8989
}
9090

9191
mainModel := e.resolveGPTImage2BaseModel()
92-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), mainModel, auth)
92+
reporter := helps.NewExecutorUsageReporter(ctx, e, mainModel, auth)
9393
defer reporter.TrackFailure(ctx, &err)
9494

9595
body, errBuild := e.prepareCodexOpenAIImageBody(prepared.Body, req, opts, mainModel)
@@ -182,7 +182,7 @@ func (e *CodexExecutor) executeOpenAIImageStream(ctx context.Context, auth *clip
182182
}
183183

184184
mainModel := e.resolveGPTImage2BaseModel()
185-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), mainModel, auth)
185+
reporter := helps.NewExecutorUsageReporter(ctx, e, mainModel, auth)
186186
defer reporter.TrackFailure(ctx, &err)
187187

188188
body, errBuild := e.prepareCodexOpenAIImageBody(prepared.Body, req, opts, mainModel)

internal/runtime/executor/codex_websockets_executor.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut
184184
baseURL = "https://chatgpt.com/backend-api/codex"
185185
}
186186

187-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
187+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
188188
defer reporter.TrackFailure(ctx, &err)
189189

190190
from := opts.SourceFormat
@@ -404,7 +404,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
404404
baseURL = "https://chatgpt.com/backend-api/codex"
405405
}
406406

407-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
407+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
408408
defer reporter.TrackFailure(ctx, &err)
409409

410410
from := opts.SourceFormat

internal/runtime/executor/gemini_cli_executor.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ func (e *GeminiCLIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth
118118
return resp, err
119119
}
120120

121-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
121+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
122122
defer reporter.TrackFailure(ctx, &err)
123123

124124
from := opts.SourceFormat
@@ -277,7 +277,7 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
277277
return nil, err
278278
}
279279

280-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
280+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
281281
defer reporter.TrackFailure(ctx, &err)
282282

283283
from := opts.SourceFormat

internal/runtime/executor/gemini_executor.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
112112

113113
apiKey, bearer := geminiCreds(auth)
114114

115-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
115+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
116116
defer reporter.TrackFailure(ctx, &err)
117117

118118
// Official Gemini API via API key or OAuth bearer
@@ -224,7 +224,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
224224

225225
apiKey, bearer := geminiCreds(auth)
226226

227-
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
227+
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
228228
defer reporter.TrackFailure(ctx, &err)
229229

230230
from := opts.SourceFormat

0 commit comments

Comments
 (0)