From d31b15916d15b550bbf388fd6da4a47d4d864109 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sun, 30 Aug 2026 14:51:08 +0800 Subject: [PATCH] feat(executor): support token usage parsing for plugin executors - Add `ParsePluginExecutorResponseUsage` to extract token usage from non-streaming plugin responses across Claude, Gemini, Interactions, Antigravity, and OpenAI/Codex protocols. - Add `ObservePluginExecutorStreamUsage` to observe and aggregate token usage across streaming chunks. Closes: #5340 --- internal/pluginhost/adapters_executors.go | 172 +++- .../adapters_executors_usage_test.go | 836 ++++++++++++++++++ .../executor/helps/gemini_ttft_helpers.go | 8 +- .../executor/helps/plugin_executor_usage.go | 220 +++++ .../handlers_plugin_executor_usage.go | 187 +--- 5 files changed, 1235 insertions(+), 188 deletions(-) create mode 100644 internal/pluginhost/adapters_executors_usage_test.go create mode 100644 internal/runtime/executor/helps/plugin_executor_usage.go diff --git a/internal/pluginhost/adapters_executors.go b/internal/pluginhost/adapters_executors.go index a80f7f359..376171982 100644 --- a/internal/pluginhost/adapters_executors.go +++ b/internal/pluginhost/adapters_executors.go @@ -9,9 +9,11 @@ import ( "net/http" "sort" "strings" + "sync" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" @@ -657,11 +659,28 @@ func (a *executorAdapter) Execute(ctx context.Context, auth *coreauth.Auth, req if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { return coreexecutor.Response{}, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) } + + var reporter *helps.UsageReporter + if auth != nil { + modelName := strings.TrimSpace(thinking.ParseSuffix(req.Model).ModelName) + if modelName == "" { + modelName = req.Model + } + reporter = helps.NewExecutorUsageReporter(ctx, a, modelName, auth) + } + defer func() { if recovered := recover(); recovered != nil { a.host.fusePlugin(a.pluginID, "Executor.Execute", recovered) resp = coreexecutor.Response{} err = fmt.Errorf("plugin executor %s panic: %v", a.Identifier(), recovered) + if reporter != nil { + reporter.PublishFailure(ctx, err) + } + return + } + if err != nil && reporter != nil { + reporter.PublishFailure(ctx, err) } }() @@ -669,10 +688,24 @@ func (a *executorAdapter) Execute(ctx context.Context, auth *coreauth.Auth, req if errPrepare != nil { return coreexecutor.Response{}, errPrepare } + + if reporter != nil { + reporter.SetTranslatedReasoningEffort(prepared.req.Payload, prepared.inputFormat.String()) + reporter.StartResponseTTFT() + } + pluginResp, errExecute := a.executor.Execute(ctx, buildExecutorRequest(a.host, a.provider, auth, prepared.req, prepared.opts)) if errExecute != nil { return coreexecutor.Response{}, errExecute } + + if reporter != nil { + reporter.RecordFirstPacket() + detail := helps.ParsePluginExecutorResponseUsage(prepared.outputFormat.String(), pluginResp.Payload) + reporter.Publish(ctx, detail) + reporter.EnsurePublished(ctx) + } + return coreexecutor.Response{ Payload: a.translateExecutorResponse(ctx, prepared, pluginResp.Payload, false, nil), Metadata: cloneAnyMap(pluginResp.Metadata), @@ -684,11 +717,28 @@ func (a *executorAdapter) ExecuteStream(ctx context.Context, auth *coreauth.Auth if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { return nil, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) } + + var reporter *helps.UsageReporter + if auth != nil { + modelName := strings.TrimSpace(thinking.ParseSuffix(req.Model).ModelName) + if modelName == "" { + modelName = req.Model + } + reporter = helps.NewExecutorUsageReporter(ctx, a, modelName, auth) + } + defer func() { if recovered := recover(); recovered != nil { a.host.fusePlugin(a.pluginID, "Executor.ExecuteStream", recovered) result = nil err = fmt.Errorf("plugin executor %s stream panic: %v", a.Identifier(), recovered) + if reporter != nil { + reporter.PublishFailure(ctx, err) + } + return + } + if err != nil && reporter != nil { + reporter.PublishFailure(ctx, err) } }() @@ -696,16 +746,136 @@ func (a *executorAdapter) ExecuteStream(ctx context.Context, auth *coreauth.Auth if errPrepare != nil { return nil, errPrepare } + + if reporter != nil { + reporter.SetTranslatedReasoningEffort(prepared.req.Payload, prepared.inputFormat.String()) + reporter.StartResponseTTFT() + } + pluginResp, errExecuteStream := a.executor.ExecuteStream(ctx, buildExecutorRequest(a.host, a.provider, auth, prepared.req, prepared.opts)) if errExecuteStream != nil { return nil, errExecuteStream } + + chunks := a.observeAndTranslateExecutorStream(ctx, prepared, pluginResp.Chunks, reporter) return &coreexecutor.StreamResult{ Headers: cloneHeader(pluginResp.Headers), - Chunks: mapExecutorStreamChunks(ctx, a.translateExecutorStreamChunks(ctx, prepared, pluginResp.Chunks)), + Chunks: chunks, }, nil } +func (a *executorAdapter) observeAndTranslateExecutorStream(ctx context.Context, prepared preparedExecutorCall, in <-chan pluginapi.ExecutorStreamChunk, reporter *helps.UsageReporter) <-chan coreexecutor.StreamChunk { + if ctx == nil { + ctx = context.Background() + } + if in == nil { + out := make(chan coreexecutor.StreamChunk) + close(out) + if reporter != nil { + reporter.EnsurePublished(ctx) + } + return out + } + if reporter == nil { + return mapExecutorStreamChunks(ctx, a.translateExecutorStreamChunks(ctx, prepared, in)) + } + + observedIn := make(chan pluginapi.ExecutorStreamChunk) + var streamUsage helps.StreamUsageBuffer + var streamErr error + var publishOnce sync.Once + var lineBuffer []byte + const maxLineBufferSize = 64 * 1024 + + publishResult := func() { + publishOnce.Do(func() { + if len(lineBuffer) > 0 { + helps.ObservePluginExecutorStreamUsage(prepared.outputFormat.String(), lineBuffer, &streamUsage) + lineBuffer = nil + } + if streamErr != nil { + if !streamUsage.PublishFailure(ctx, reporter, streamErr) { + reporter.PublishFailure(ctx, streamErr) + } + reporter.EnsurePublished(ctx) + return + } + if ctx.Err() != nil { + if !streamUsage.PublishFailure(ctx, reporter, ctx.Err()) { + reporter.PublishFailure(ctx, ctx.Err()) + } + reporter.EnsurePublished(ctx) + return + } + if !streamUsage.Publish(ctx, reporter) { + reporter.EnsurePublished(ctx) + } + }) + } + + go func() { + defer close(observedIn) + defer publishResult() + + for { + select { + case <-ctx.Done(): + return + case chunk, ok := <-in: + if !ok { + return + } + if chunk.Err != nil { + if streamErr == nil { + streamErr = chunk.Err + } + publishResult() + select { + case observedIn <- chunk: + case <-ctx.Done(): + return + } + continue + } + + if len(chunk.Payload) > 0 { + helps.ObservePluginExecutorStreamTTFT(prepared.outputFormat.String(), reporter, chunk.Payload) + + lineBuffer = append(lineBuffer, chunk.Payload...) + for { + idx := bytes.IndexByte(lineBuffer, '\n') + if idx < 0 { + break + } + line := lineBuffer[:idx+1] + lineBuffer = lineBuffer[idx+1:] + helps.ObservePluginExecutorStreamUsage(prepared.outputFormat.String(), line, &streamUsage) + } + if len(lineBuffer) > 0 { + if jsonBytes := helps.ExtractStreamJSONPayload(lineBuffer); len(jsonBytes) > 0 && json.Valid(jsonBytes) { + helps.ObservePluginExecutorStreamUsage(prepared.outputFormat.String(), lineBuffer, &streamUsage) + lineBuffer = nil + } + } + if len(lineBuffer) > maxLineBufferSize { + helps.ObservePluginExecutorStreamUsage(prepared.outputFormat.String(), lineBuffer, &streamUsage) + lineBuffer = nil + } + } + + select { + case observedIn <- chunk: + case <-ctx.Done(): + return + } + } + } + }() + + translatedOut := a.translateExecutorStreamChunks(ctx, prepared, observedIn) + return mapExecutorStreamChunks(ctx, translatedOut) +} + func (a *executorAdapter) Refresh(ctx context.Context, auth *coreauth.Auth) (refreshed *coreauth.Auth, err error) { if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { return nil, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) diff --git a/internal/pluginhost/adapters_executors_usage_test.go b/internal/pluginhost/adapters_executors_usage_test.go new file mode 100644 index 000000000..c6dd11829 --- /dev/null +++ b/internal/pluginhost/adapters_executors_usage_test.go @@ -0,0 +1,836 @@ +package pluginhost + +import ( + "context" + "errors" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" +) + +type testUsageCapturePlugin struct { + targetProvider string + records chan coreusage.Record +} + +func newTestUsageCapturePlugin(targetProvider string) *testUsageCapturePlugin { + return &testUsageCapturePlugin{ + targetProvider: targetProvider, + records: make(chan coreusage.Record, 50), + } +} + +func (p *testUsageCapturePlugin) HandleUsage(_ context.Context, record coreusage.Record) { + if p.targetProvider != "" && record.Provider != p.targetProvider { + return + } + select { + case p.records <- record: + default: + } +} + +func (p *testUsageCapturePlugin) waitRecord(t *testing.T, timeout time.Duration) coreusage.Record { + t.Helper() + select { + case rec := <-p.records: + return rec + case <-time.After(timeout): + t.Fatal("timed out waiting for usage record") + return coreusage.Record{} + } +} + +func (p *testUsageCapturePlugin) assertNoRecord(t *testing.T, wait time.Duration) { + t.Helper() + select { + case rec := <-p.records: + t.Fatalf("expected no usage record for %q, got %+v", p.targetProvider, rec) + case <-time.After(wait): + } +} + +func registerTestUsagePlugin(t *testing.T, name string, plugin coreusage.Plugin) { + t.Helper() + coreusage.RegisterNamedPlugin(name, plugin) + t.Cleanup(func() { + coreusage.RegisterNamedPlugin(name, noopUsagePlugin{}) + }) +} + +type noopUsagePlugin struct{} + +func (noopUsagePlugin) HandleUsage(context.Context, coreusage.Record) {} + +func TestExecutorAdapterExecutePublishesUsage(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider") + registerTestUsagePlugin(t, "test-executor-adapter-execute-usage", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin"}) + host := newHostWithRecords(executorRecord) + + exec := &fakeExecutor{ + identifier: "plugin-provider", + execute: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { + return pluginapi.ExecutorResponse{ + Payload: []byte(`{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":10,"completion_tokens":20,"total_tokens":30}}`), + Headers: http.Header{"Content-Type": []string{"application/json"}}, + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + + auth := &coreauth.Auth{ + ID: "auth-1", + Provider: "plugin-provider", + FileName: "auth-1.json", + Attributes: map[string]string{"type": "oauth"}, + } + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model","messages":[{"role":"user","content":"hi"}]}`), + } + opts := coreexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + } + + resp, err := adapter.Execute(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("adapter.Execute returned unexpected error: %v", err) + } + if len(resp.Payload) == 0 { + t.Fatal("adapter.Execute returned empty payload") + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if rec.Provider != "plugin-provider" { + t.Errorf("got provider %q, want %q", rec.Provider, "plugin-provider") + } + if rec.Detail.InputTokens != 10 || rec.Detail.OutputTokens != 20 || rec.Detail.TotalTokens != 30 { + t.Errorf("got usage %+v, want 10 input, 20 output, 30 total", rec.Detail) + } +} + +func TestExecutorAdapterExecuteThroughAuthManagerPublishesUsage(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-mgr") + registerTestUsagePlugin(t, "test-executor-adapter-auth-manager-usage", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-mgr"}) + host := newHostWithRecords(executorRecord) + + exec := &fakeExecutor{ + identifier: "plugin-provider-mgr", + execute: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { + return pluginapi.ExecutorResponse{ + Payload: []byte(`{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":10,"completion_tokens":20,"total_tokens":30}}`), + Headers: http.Header{"Content-Type": []string{"application/json"}}, + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-mgr" + + authMgr := coreauth.NewManager(nil, nil, nil) + authMgr.RegisterExecutor(adapter) + + model := "test-model-mgr" + auth := &coreauth.Auth{ + ID: "auth-mgr-1", + Provider: "plugin-provider-mgr", + Status: coreauth.StatusActive, + FileName: "auth-mgr-1.json", + Attributes: map[string]string{"type": "oauth"}, + } + if _, err := authMgr.Register(context.Background(), auth); err != nil { + t.Fatalf("Register auth: %v", err) + } + + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(auth.ID) + }) + + req := coreexecutor.Request{ + Model: model, + Payload: []byte(`{"model":"test-model-mgr","messages":[{"role":"user","content":"hi"}]}`), + } + opts := coreexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + } + + resp, err := authMgr.Execute(context.Background(), []string{"plugin-provider-mgr"}, req, opts) + if err != nil { + t.Fatalf("authMgr.Execute returned unexpected error: %v", err) + } + if len(resp.Payload) == 0 { + t.Fatal("authMgr.Execute returned empty payload") + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if rec.Provider != "plugin-provider-mgr" { + t.Errorf("got provider %q, want %q", rec.Provider, "plugin-provider-mgr") + } + if rec.Detail.InputTokens != 10 || rec.Detail.OutputTokens != 20 || rec.Detail.TotalTokens != 30 { + t.Errorf("got usage %+v, want 10 input, 20 output, 30 total", rec.Detail) + } +} + +func TestExecutorAdapterExecuteNilAuthSkipsUsage(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-nil-auth") + registerTestUsagePlugin(t, "test-executor-adapter-execute-nil-auth", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-nil-auth"}) + host := newHostWithRecords(executorRecord) + + exec := &fakeExecutor{ + identifier: "plugin-provider-nil-auth", + execute: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { + return pluginapi.ExecutorResponse{ + Payload: []byte(`{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":10,"completion_tokens":20,"total_tokens":30}}`), + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-nil-auth" + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model","messages":[{"role":"user","content":"hi"}]}`), + } + opts := coreexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + } + + resp, err := adapter.Execute(context.Background(), nil, req, opts) + if err != nil { + t.Fatalf("adapter.Execute returned unexpected error: %v", err) + } + if len(resp.Payload) == 0 { + t.Fatal("adapter.Execute returned empty payload") + } + + plugin.assertNoRecord(t, 50*time.Millisecond) +} + +func TestExecutorAdapterExecuteErrorPublishesFailure(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-error") + registerTestUsagePlugin(t, "test-executor-adapter-execute-error", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-error"}) + host := newHostWithRecords(executorRecord) + + exec := &fakeExecutor{ + identifier: "plugin-provider-error", + execute: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { + return pluginapi.ExecutorResponse{}, errors.New("upstream failed") + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-error" + + auth := &coreauth.Auth{ + ID: "auth-1", + Provider: "plugin-provider-error", + } + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model"}`), + } + + _, err := adapter.Execute(context.Background(), auth, req, coreexecutor.Options{}) + if err == nil { + t.Fatal("expected error from adapter.Execute") + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if !rec.Failed { + t.Errorf("rec.Failed = false, want true") + } +} + +func TestExecutorAdapterExecutePanicPublishesFailure(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-panic") + registerTestUsagePlugin(t, "test-executor-adapter-execute-panic", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-panic"}) + host := newHostWithRecords(executorRecord) + + exec := &fakeExecutor{ + identifier: "plugin-provider-panic", + execute: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { + panic("execute panic boom") + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-panic" + + auth := &coreauth.Auth{ + ID: "auth-panic-1", + Provider: "plugin-provider-panic", + } + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model"}`), + } + + _, err := adapter.Execute(context.Background(), auth, req, coreexecutor.Options{}) + if err == nil { + t.Fatal("expected error from adapter.Execute on panic") + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if !rec.Failed { + t.Errorf("rec.Failed = false, want true") + } +} + +func TestExecutorAdapterExecuteStreamPublishesUsage(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-stream") + registerTestUsagePlugin(t, "test-executor-adapter-stream-usage", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-stream"}) + host := newHostWithRecords(executorRecord) + + streamChunks := make(chan pluginapi.ExecutorStreamChunk, 4) + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":15,\"completion_tokens\":25,\"total_tokens\":40}}\n\n")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: [DONE]\n\n")} + close(streamChunks) + + exec := &fakeExecutor{ + identifier: "plugin-provider-stream", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + return pluginapi.ExecutorStreamResponse{ + Chunks: streamChunks, + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-stream" + + auth := &coreauth.Auth{ + ID: "auth-stream-1", + Provider: "plugin-provider-stream", + } + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model","stream":true}`), + } + opts := coreexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + Stream: true, + } + + streamRes, err := adapter.ExecuteStream(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("ExecuteStream returned unexpected error: %v", err) + } + + var receivedChunks [][]byte + for chunk := range streamRes.Chunks { + if chunk.Err != nil { + t.Fatalf("unexpected chunk error: %v", chunk.Err) + } + receivedChunks = append(receivedChunks, chunk.Payload) + } + + if len(receivedChunks) == 0 { + t.Fatal("expected non-empty received chunks") + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if rec.Provider != "plugin-provider-stream" { + t.Errorf("got provider %q, want %q", rec.Provider, "plugin-provider-stream") + } + if rec.Detail.InputTokens != 15 || rec.Detail.OutputTokens != 25 || rec.Detail.TotalTokens != 40 { + t.Errorf("got usage %+v, want prompt=15 completion=25 total=40", rec.Detail) + } +} + +func TestExecutorAdapterExecuteStreamThroughAuthManagerPublishesUsage(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-stream-mgr") + registerTestUsagePlugin(t, "test-executor-adapter-auth-manager-stream-usage", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-stream-mgr"}) + host := newHostWithRecords(executorRecord) + + streamChunks := make(chan pluginapi.ExecutorStreamChunk, 4) + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":15,\"completion_tokens\":25,\"total_tokens\":40}}\n\n")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: [DONE]\n\n")} + close(streamChunks) + + exec := &fakeExecutor{ + identifier: "plugin-provider-stream-mgr", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + return pluginapi.ExecutorStreamResponse{ + Chunks: streamChunks, + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-stream-mgr" + + authMgr := coreauth.NewManager(nil, nil, nil) + authMgr.RegisterExecutor(adapter) + + model := "test-model-stream-mgr" + auth := &coreauth.Auth{ + ID: "auth-stream-mgr-1", + Provider: "plugin-provider-stream-mgr", + Status: coreauth.StatusActive, + FileName: "auth-stream-mgr-1.json", + Attributes: map[string]string{"type": "oauth"}, + } + if _, err := authMgr.Register(context.Background(), auth); err != nil { + t.Fatalf("Register auth: %v", err) + } + + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(auth.ID) + }) + + req := coreexecutor.Request{ + Model: model, + Payload: []byte(`{"model":"test-model-stream-mgr","stream":true}`), + } + opts := coreexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + Stream: true, + } + + streamRes, err := authMgr.ExecuteStream(context.Background(), []string{"plugin-provider-stream-mgr"}, req, opts) + if err != nil { + t.Fatalf("authMgr.ExecuteStream returned unexpected error: %v", err) + } + + for range streamRes.Chunks { + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if rec.Provider != "plugin-provider-stream-mgr" { + t.Errorf("got provider %q, want %q", rec.Provider, "plugin-provider-stream-mgr") + } + if rec.Detail.InputTokens != 15 || rec.Detail.OutputTokens != 25 || rec.Detail.TotalTokens != 40 { + t.Errorf("got usage %+v, want prompt=15 completion=25 total=40", rec.Detail) + } +} + +func TestExecutorAdapterExecuteStreamTwoCompleteChunksWithoutNewline(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-no-newline") + registerTestUsagePlugin(t, "test-executor-adapter-stream-no-newline", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-no-newline"}) + host := newHostWithRecords(executorRecord) + + streamChunks := make(chan pluginapi.ExecutorStreamChunk, 3) + // Two complete chunks without trailing \n + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":22,\"total_tokens\":33}}")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: [DONE]")} + close(streamChunks) + + exec := &fakeExecutor{ + identifier: "plugin-provider-no-newline", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + return pluginapi.ExecutorStreamResponse{ + Chunks: streamChunks, + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-no-newline" + + auth := &coreauth.Auth{ + ID: "auth-no-newline-1", + Provider: "plugin-provider-no-newline", + } + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model","stream":true}`), + } + + streamRes, err := adapter.ExecuteStream(context.Background(), auth, req, coreexecutor.Options{}) + if err != nil { + t.Fatalf("ExecuteStream returned unexpected error: %v", err) + } + + for range streamRes.Chunks { + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if rec.Detail.InputTokens != 11 || rec.Detail.OutputTokens != 22 || rec.Detail.TotalTokens != 33 { + t.Errorf("got usage %+v, want prompt=11 completion=22 total=33", rec.Detail) + } +} + +func TestExecutorAdapterExecuteStreamSplitChunks(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-split") + registerTestUsagePlugin(t, "test-executor-adapter-stream-split", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-split"}) + host := newHostWithRecords(executorRecord) + + streamChunks := make(chan pluginapi.ExecutorStreamChunk, 5) + // Split SSE frame across two chunks + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":{\"prompt_tok")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("ens\":18,\"completion_tokens\":22,\"total_tokens\":40}}\n\n")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: [DONE]\n\n")} + close(streamChunks) + + exec := &fakeExecutor{ + identifier: "plugin-provider-split", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + return pluginapi.ExecutorStreamResponse{ + Chunks: streamChunks, + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-split" + + auth := &coreauth.Auth{ + ID: "auth-split-1", + Provider: "plugin-provider-split", + } + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model","stream":true}`), + } + + streamRes, err := adapter.ExecuteStream(context.Background(), auth, req, coreexecutor.Options{}) + if err != nil { + t.Fatalf("ExecuteStream returned unexpected error: %v", err) + } + + for range streamRes.Chunks { + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if rec.Detail.InputTokens != 18 || rec.Detail.OutputTokens != 22 || rec.Detail.TotalTokens != 40 { + t.Errorf("got usage %+v, want prompt=18 completion=22 total=40", rec.Detail) + } +} + +func TestExecutorAdapterExecuteStreamErrorChunk(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-stream-err") + registerTestUsagePlugin(t, "test-executor-adapter-stream-err", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-stream-err"}) + host := newHostWithRecords(executorRecord) + + expectedErr := errors.New("mid-stream chunk failure") + streamChunks := make(chan pluginapi.ExecutorStreamChunk, 2) + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"start\"}}]}\n\n")} + streamChunks <- pluginapi.ExecutorStreamChunk{Err: expectedErr} + close(streamChunks) + + exec := &fakeExecutor{ + identifier: "plugin-provider-stream-err", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + return pluginapi.ExecutorStreamResponse{ + Chunks: streamChunks, + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-stream-err" + + auth := &coreauth.Auth{ + ID: "auth-stream-err-1", + Provider: "plugin-provider-stream-err", + } + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model","stream":true}`), + } + + streamRes, err := adapter.ExecuteStream(context.Background(), auth, req, coreexecutor.Options{}) + if err != nil { + t.Fatalf("ExecuteStream returned unexpected error: %v", err) + } + + var gotError error + for chunk := range streamRes.Chunks { + if chunk.Err != nil { + gotError = chunk.Err + } + } + + if gotError == nil || !errors.Is(gotError, expectedErr) { + t.Fatalf("expected error %v, got %v", expectedErr, gotError) + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if !rec.Failed { + t.Errorf("rec.Failed = false, want true") + } +} + +func TestExecutorAdapterExecuteStreamPanicPublishesFailure(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-stream-panic") + registerTestUsagePlugin(t, "test-executor-adapter-stream-panic", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-stream-panic"}) + host := newHostWithRecords(executorRecord) + + exec := &fakeExecutor{ + identifier: "plugin-provider-stream-panic", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + panic("execute stream panic boom") + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-stream-panic" + + auth := &coreauth.Auth{ + ID: "auth-stream-panic-1", + Provider: "plugin-provider-stream-panic", + } + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model","stream":true}`), + } + + _, err := adapter.ExecuteStream(context.Background(), auth, req, coreexecutor.Options{}) + if err == nil { + t.Fatal("expected error from adapter.ExecuteStream on panic") + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if !rec.Failed { + t.Errorf("rec.Failed = false, want true") + } +} + +func TestExecutorAdapterExecuteStreamNilAuthSkipsUsage(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-stream-nil") + registerTestUsagePlugin(t, "test-executor-adapter-stream-nil", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-stream-nil"}) + host := newHostWithRecords(executorRecord) + + streamChunks := make(chan pluginapi.ExecutorStreamChunk, 2) + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":20,\"total_tokens\":30}}\n\n")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: [DONE]\n\n")} + close(streamChunks) + + exec := &fakeExecutor{ + identifier: "plugin-provider-stream-nil", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + return pluginapi.ExecutorStreamResponse{ + Chunks: streamChunks, + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-stream-nil" + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model","stream":true}`), + } + + streamRes, err := adapter.ExecuteStream(context.Background(), nil, req, coreexecutor.Options{}) + if err != nil { + t.Fatalf("ExecuteStream returned unexpected error: %v", err) + } + + for range streamRes.Chunks { + } + + plugin.assertNoRecord(t, 50*time.Millisecond) +} + +func TestExecutorAdapterExecuteStreamSplitAcrossBraceBoundary(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-split-brace") + registerTestUsagePlugin(t, "test-executor-adapter-stream-split-brace", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-split-brace"}) + host := newHostWithRecords(executorRecord) + + streamChunks := make(chan pluginapi.ExecutorStreamChunk, 3) + // Split exactly before opening brace of usage object + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("{\"prompt_tokens\":18,\"completion_tokens\":22,\"total_tokens\":40}}\n\n")} + streamChunks <- pluginapi.ExecutorStreamChunk{Payload: []byte("data: [DONE]\n\n")} + close(streamChunks) + + exec := &fakeExecutor{ + identifier: "plugin-provider-split-brace", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + return pluginapi.ExecutorStreamResponse{ + Chunks: streamChunks, + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-split-brace" + + auth := &coreauth.Auth{ + ID: "auth-split-brace-1", + Provider: "plugin-provider-split-brace", + } + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model","stream":true}`), + } + + streamRes, err := adapter.ExecuteStream(context.Background(), auth, req, coreexecutor.Options{}) + if err != nil { + t.Fatalf("ExecuteStream returned unexpected error: %v", err) + } + + for range streamRes.Chunks { + } + + rec := plugin.waitRecord(t, 200*time.Millisecond) + if rec.Detail.InputTokens != 18 || rec.Detail.OutputTokens != 22 || rec.Detail.TotalTokens != 40 { + t.Errorf("got usage %+v, want prompt=18 completion=22 total=40", rec.Detail) + } +} + +func TestExecutorAdapterExecuteStreamErrorChunkImmediatePublishWithoutClose(t *testing.T) { + plugin := newTestUsageCapturePlugin("plugin-provider-stream-err-unclosed") + registerTestUsagePlugin(t, "test-executor-adapter-stream-err-unclosed", plugin) + + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-plugin-stream-err-unclosed"}) + host := newHostWithRecords(executorRecord) + + expectedErr := errors.New("mid-stream chunk failure without close") + streamChunks := make(chan pluginapi.ExecutorStreamChunk) + + exec := &fakeExecutor{ + identifier: "plugin-provider-stream-err-unclosed", + executeStream: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + return pluginapi.ExecutorStreamResponse{ + Chunks: streamChunks, + }, nil + }, + } + + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = "plugin-provider-stream-err-unclosed" + + auth := &coreauth.Auth{ + ID: "auth-stream-err-unclosed-1", + Provider: "plugin-provider-stream-err-unclosed", + } + + req := coreexecutor.Request{ + Model: "test-model", + Payload: []byte(`{"model":"test-model","stream":true}`), + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + streamRes, err := adapter.ExecuteStream(ctx, auth, req, coreexecutor.Options{}) + if err != nil { + t.Fatalf("ExecuteStream returned unexpected error: %v", err) + } + + // Send error chunk but do NOT close the channel + go func() { + streamChunks <- pluginapi.ExecutorStreamChunk{Err: expectedErr} + }() + + chunk := <-streamRes.Chunks + if chunk.Err == nil || !errors.Is(chunk.Err, expectedErr) { + t.Fatalf("expected chunk error %v, got %v", expectedErr, chunk.Err) + } + + // Verify that usage failure record was published immediately without channel close + rec := plugin.waitRecord(t, 200*time.Millisecond) + if !rec.Failed { + t.Errorf("rec.Failed = false, want true") + } +} + +func TestObservePluginExecutorStreamTTFT_Antigravity(t *testing.T) { + reporter := helps.NewUsageReporter(context.Background(), "antigravity", "test-model", nil) + reporter.StartResponseTTFT() + antigravityChunk := []byte(`data: {"response":{"candidates":[{"content":{"parts":[{"text":"hello"}]}}]}}`) + helps.ObservePluginExecutorStreamTTFT("antigravity", reporter, antigravityChunk) + if !reporter.IsTTFTSet() { + t.Errorf("ObservePluginExecutorStreamTTFT for antigravity should set TTFT on token event") + } +} diff --git a/internal/runtime/executor/helps/gemini_ttft_helpers.go b/internal/runtime/executor/helps/gemini_ttft_helpers.go index 005f82bfc..c4d531799 100644 --- a/internal/runtime/executor/helps/gemini_ttft_helpers.go +++ b/internal/runtime/executor/helps/gemini_ttft_helpers.go @@ -28,11 +28,15 @@ func IsGeminiTokenEvent(payload []byte) bool { } // Terminal error fallback - if gjson.GetBytes(payload, "error.message").Exists() || gjson.GetBytes(payload, "error").Exists() { + if gjson.GetBytes(payload, "error.message").Exists() || gjson.GetBytes(payload, "error").Exists() || gjson.GetBytes(payload, "response.error").Exists() { return true } - candidates := gjson.GetBytes(payload, "candidates").Array() + candidatesNode := gjson.GetBytes(payload, "candidates") + if !candidatesNode.Exists() { + candidatesNode = gjson.GetBytes(payload, "response.candidates") + } + candidates := candidatesNode.Array() if len(candidates) == 0 { return false } diff --git a/internal/runtime/executor/helps/plugin_executor_usage.go b/internal/runtime/executor/helps/plugin_executor_usage.go new file mode 100644 index 000000000..a0564ef03 --- /dev/null +++ b/internal/runtime/executor/helps/plugin_executor_usage.go @@ -0,0 +1,220 @@ +package helps + +import ( + "bytes" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" + "github.com/tidwall/gjson" +) + +// ParsePluginExecutorResponseUsage extracts token usage from a non-streaming plugin executor response. +func ParsePluginExecutorResponseUsage(protocol string, payload []byte) usage.Detail { + if len(payload) == 0 { + return usage.Detail{} + } + switch strings.ToLower(strings.TrimSpace(protocol)) { + case "claude": + return parseClaudePayloadUsage(payload) + case "gemini": + return ParseGeminiUsage(payload) + case "interactions", "interactions-response": + return ParseInteractionsUsage(payload) + case "antigravity": + return ParseAntigravityUsage(payload) + case "codex", "openai-response": + if detail, ok := ParseCodexUsage(payload); ok { + return detail + } + return ParseOpenAIUsage(payload) + default: + return ParseOpenAIUsage(payload) + } +} + +// ObservePluginExecutorStreamUsage parses streaming chunks and updates a stream usage buffer. +func ObservePluginExecutorStreamUsage(protocol string, payload []byte, buffer *StreamUsageBuffer) { + if buffer == nil || len(payload) == 0 { + return + } + switch strings.ToLower(strings.TrimSpace(protocol)) { + case "claude": + IterateStreamLines(payload, func(line []byte) { + if detail, ok := parseClaudeStreamLine(line); ok { + ObserveMergedStreamUsage(buffer, detail) + } + }) + case "gemini": + IterateStreamLines(payload, func(line []byte) { + if detail, ok := ParseGeminiStreamUsage(line); ok { + buffer.Observe(detail, ok) + } + }) + case "interactions", "interactions-response": + IterateStreamLines(payload, func(line []byte) { + if detail, ok := ParseInteractionsStreamUsage(line); ok { + ObserveMergedStreamUsage(buffer, detail) + } + }) + case "antigravity": + IterateStreamLines(payload, func(line []byte) { + if detail, ok := ParseAntigravityStreamUsage(line); ok { + buffer.Observe(detail, ok) + } + }) + case "codex", "openai-response": + IterateStreamLines(payload, func(line []byte) { + if jsonBytes := ExtractStreamJSONPayload(line); len(jsonBytes) > 0 { + if detail, ok := ParseCodexUsage(jsonBytes); ok { + buffer.Observe(detail, ok) + return + } + } + buffer.ObserveOpenAIStream(line) + }) + default: + IterateStreamLines(payload, func(line []byte) { + buffer.ObserveOpenAIStream(line) + }) + } +} + +// ObservePluginExecutorStreamTTFT inspects a streaming payload and records TTFT on the reporter. +func ObservePluginExecutorStreamTTFT(protocol string, reporter *UsageReporter, payload []byte) { + if reporter == nil || len(payload) == 0 { + return + } + reporter.RecordFirstPacket() + switch strings.ToLower(strings.TrimSpace(protocol)) { + case "claude": + ObserveClaudeTokenEvent(reporter, payload) + case "gemini", "antigravity", "interactions", "interactions-response": + ObserveGeminiTokenEvent(reporter, payload) + case "codex", "openai-response": + ObserveResponsesTokenEvent(reporter, payload) + default: + ObserveChatTokenEvent(reporter, payload) + } +} + +func parseClaudePayloadUsage(payload []byte) usage.Detail { + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return usage.Detail{} + } + usageNode := gjson.GetBytes(payload, "usage") + if !usageNode.Exists() { + usageNode = gjson.GetBytes(payload, "message.usage") + } + if !usageNode.Exists() { + return usage.Detail{} + } + return ParseClaudeUsage([]byte(`{"usage":` + usageNode.Raw + `}`)) +} + +func parseClaudeStreamLine(line []byte) (usage.Detail, bool) { + payload := ExtractStreamJSONPayload(line) + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return usage.Detail{}, false + } + usageNode := gjson.GetBytes(payload, "usage") + if !usageNode.Exists() { + usageNode = gjson.GetBytes(payload, "message.usage") + } + if !usageNode.Exists() { + return usage.Detail{}, false + } + detail := ParseClaudeUsage([]byte(`{"usage":` + usageNode.Raw + `}`)) + return detail, true +} + +// ObserveMergedStreamUsage updates buffer with merged usage details. +func ObserveMergedStreamUsage(buffer *StreamUsageBuffer, update usage.Detail) { + if buffer == nil { + return + } + if existing, ok := buffer.Detail(); ok { + merged := MergeStreamUsageDetail(existing, update) + buffer.Observe(merged, true) + return + } + buffer.Observe(update, true) +} + +// MergeStreamUsageDetail merges existing stream usage with a newer update. +func MergeStreamUsageDetail(existing, update usage.Detail) usage.Detail { + merged := update + if merged.InputTokens == 0 && existing.InputTokens > 0 { + merged.InputTokens = existing.InputTokens + } + if merged.CachedTokens == 0 && existing.CachedTokens > 0 { + merged.CachedTokens = existing.CachedTokens + } + if merged.CacheReadTokens == 0 && existing.CacheReadTokens > 0 { + merged.CacheReadTokens = existing.CacheReadTokens + } + if merged.CacheCreationTokens == 0 && existing.CacheCreationTokens > 0 { + merged.CacheCreationTokens = existing.CacheCreationTokens + } + if merged.OutputTokens == 0 && existing.OutputTokens > 0 { + merged.OutputTokens = existing.OutputTokens + } + if merged.ReasoningTokens == 0 && existing.ReasoningTokens > 0 { + merged.ReasoningTokens = existing.ReasoningTokens + } + if merged.ResponseServiceTier == "" { + merged.ResponseServiceTier = existing.ResponseServiceTier + } + cached := merged.CacheReadTokens + merged.CacheCreationTokens + if cached == 0 { + cached = merged.CachedTokens + } + calculatedTotal := merged.InputTokens + merged.OutputTokens + cached + if merged.TotalTokens == 0 || merged.TotalTokens < calculatedTotal { + merged.TotalTokens = calculatedTotal + } + nonReasoningOutput := merged.OutputTokens - merged.ReasoningTokens + if nonReasoningOutput < 0 { + nonReasoningOutput = 0 + } + merged.TokenBreakdown = usage.NewIndependentTokenBreakdown( + merged.InputTokens, + merged.CacheReadTokens, + merged.CacheCreationTokens, + nonReasoningOutput, + merged.ReasoningTokens, + merged.TotalTokens, + ) + return merged +} + +// IterateStreamLines splits payload by newline and invokes fn for non-empty lines. +func IterateStreamLines(payload []byte, fn func(line []byte)) { + for _, line := range bytes.Split(payload, []byte("\n")) { + trimmed := bytes.TrimSpace(line) + if len(trimmed) == 0 { + continue + } + fn(trimmed) + } +} + +// ExtractStreamJSONPayload extracts SSE data/json payload from a raw line. +func ExtractStreamJSONPayload(line []byte) []byte { + trimmed := bytes.TrimSpace(line) + if len(trimmed) == 0 { + return nil + } + if bytes.Equal(trimmed, []byte("[DONE]")) { + return nil + } + if bytes.HasPrefix(trimmed, []byte("event:")) { + return nil + } + if bytes.HasPrefix(trimmed, []byte("data:")) { + trimmed = bytes.TrimSpace(bytes.TrimPrefix(trimmed, []byte("data:"))) + } + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("[DONE]")) { + return nil + } + return trimmed +} diff --git a/sdk/api/handlers/handlers_plugin_executor_usage.go b/sdk/api/handlers/handlers_plugin_executor_usage.go index 0d6c5d1eb..4ab07bcae 100644 --- a/sdk/api/handlers/handlers_plugin_executor_usage.go +++ b/sdk/api/handlers/handlers_plugin_executor_usage.go @@ -1,197 +1,14 @@ package handlers import ( - "bytes" - "strings" - "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" - "github.com/tidwall/gjson" ) func parsePluginExecutorResponseUsage(protocol string, payload []byte) usage.Detail { - if len(payload) == 0 { - return usage.Detail{} - } - switch strings.ToLower(strings.TrimSpace(protocol)) { - case "claude": - return parseClaudePayloadUsage(payload) - case "gemini": - return helps.ParseGeminiUsage(payload) - case "interactions", "interactions-response": - return helps.ParseInteractionsUsage(payload) - case "antigravity": - return helps.ParseAntigravityUsage(payload) - case "codex", "openai-response": - if detail, ok := helps.ParseCodexUsage(payload); ok { - return detail - } - return helps.ParseOpenAIUsage(payload) - default: - return helps.ParseOpenAIUsage(payload) - } + return helps.ParsePluginExecutorResponseUsage(protocol, payload) } func observePluginExecutorStreamUsage(protocol string, payload []byte, buffer *helps.StreamUsageBuffer) { - if buffer == nil || len(payload) == 0 { - return - } - switch strings.ToLower(strings.TrimSpace(protocol)) { - case "claude": - iterateStreamLines(payload, func(line []byte) { - if detail, ok := parseClaudeStreamLine(line); ok { - observeMergedStreamUsage(buffer, detail) - } - }) - case "gemini": - iterateStreamLines(payload, func(line []byte) { - if detail, ok := helps.ParseGeminiStreamUsage(line); ok { - buffer.Observe(detail, ok) - } - }) - case "interactions", "interactions-response": - iterateStreamLines(payload, func(line []byte) { - if detail, ok := helps.ParseInteractionsStreamUsage(line); ok { - observeMergedStreamUsage(buffer, detail) - } - }) - case "antigravity": - iterateStreamLines(payload, func(line []byte) { - if detail, ok := helps.ParseAntigravityStreamUsage(line); ok { - buffer.Observe(detail, ok) - } - }) - case "codex", "openai-response": - iterateStreamLines(payload, func(line []byte) { - if jsonBytes := extractStreamJSONPayload(line); len(jsonBytes) > 0 { - if detail, ok := helps.ParseCodexUsage(jsonBytes); ok { - buffer.Observe(detail, ok) - return - } - } - buffer.ObserveOpenAIStream(line) - }) - default: - iterateStreamLines(payload, func(line []byte) { - buffer.ObserveOpenAIStream(line) - }) - } -} - -func parseClaudePayloadUsage(payload []byte) usage.Detail { - if len(payload) == 0 || !gjson.ValidBytes(payload) { - return usage.Detail{} - } - usageNode := gjson.GetBytes(payload, "usage") - if !usageNode.Exists() { - usageNode = gjson.GetBytes(payload, "message.usage") - } - if !usageNode.Exists() { - return usage.Detail{} - } - return helps.ParseClaudeUsage([]byte(`{"usage":` + usageNode.Raw + `}`)) -} - -func parseClaudeStreamLine(line []byte) (usage.Detail, bool) { - payload := extractStreamJSONPayload(line) - if len(payload) == 0 || !gjson.ValidBytes(payload) { - return usage.Detail{}, false - } - usageNode := gjson.GetBytes(payload, "usage") - if !usageNode.Exists() { - usageNode = gjson.GetBytes(payload, "message.usage") - } - if !usageNode.Exists() { - return usage.Detail{}, false - } - detail := helps.ParseClaudeUsage([]byte(`{"usage":` + usageNode.Raw + `}`)) - return detail, true -} - -func observeMergedStreamUsage(buffer *helps.StreamUsageBuffer, update usage.Detail) { - if buffer == nil { - return - } - if existing, ok := buffer.Detail(); ok { - merged := mergeStreamUsageDetail(existing, update) - buffer.Observe(merged, true) - return - } - buffer.Observe(update, true) -} - -func mergeStreamUsageDetail(existing, update usage.Detail) usage.Detail { - merged := update - if merged.InputTokens == 0 && existing.InputTokens > 0 { - merged.InputTokens = existing.InputTokens - } - if merged.CachedTokens == 0 && existing.CachedTokens > 0 { - merged.CachedTokens = existing.CachedTokens - } - if merged.CacheReadTokens == 0 && existing.CacheReadTokens > 0 { - merged.CacheReadTokens = existing.CacheReadTokens - } - if merged.CacheCreationTokens == 0 && existing.CacheCreationTokens > 0 { - merged.CacheCreationTokens = existing.CacheCreationTokens - } - if merged.OutputTokens == 0 && existing.OutputTokens > 0 { - merged.OutputTokens = existing.OutputTokens - } - if merged.ReasoningTokens == 0 && existing.ReasoningTokens > 0 { - merged.ReasoningTokens = existing.ReasoningTokens - } - if merged.ResponseServiceTier == "" { - merged.ResponseServiceTier = existing.ResponseServiceTier - } - cached := merged.CacheReadTokens + merged.CacheCreationTokens - if cached == 0 { - cached = merged.CachedTokens - } - calculatedTotal := merged.InputTokens + merged.OutputTokens + cached - if merged.TotalTokens == 0 || merged.TotalTokens < calculatedTotal { - merged.TotalTokens = calculatedTotal - } - nonReasoningOutput := merged.OutputTokens - merged.ReasoningTokens - if nonReasoningOutput < 0 { - nonReasoningOutput = 0 - } - merged.TokenBreakdown = usage.NewIndependentTokenBreakdown( - merged.InputTokens, - merged.CacheReadTokens, - merged.CacheCreationTokens, - nonReasoningOutput, - merged.ReasoningTokens, - merged.TotalTokens, - ) - return merged -} - -func iterateStreamLines(payload []byte, fn func(line []byte)) { - for _, line := range bytes.Split(payload, []byte("\n")) { - trimmed := bytes.TrimSpace(line) - if len(trimmed) == 0 { - continue - } - fn(trimmed) - } -} - -func extractStreamJSONPayload(line []byte) []byte { - trimmed := bytes.TrimSpace(line) - if len(trimmed) == 0 { - return nil - } - if bytes.Equal(trimmed, []byte("[DONE]")) { - return nil - } - if bytes.HasPrefix(trimmed, []byte("event:")) { - return nil - } - if bytes.HasPrefix(trimmed, []byte("data:")) { - trimmed = bytes.TrimSpace(bytes.TrimPrefix(trimmed, []byte("data:"))) - } - if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("[DONE]")) { - return nil - } - return trimmed + helps.ObservePluginExecutorStreamUsage(protocol, payload, buffer) }