Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
30 kB · 724 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725package handlers
import ( "context" "errors" "fmt" "testing" "time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry" 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/cliproxy/usage" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi")
type noopUsagePlugin struct{}
func (noopUsagePlugin) HandleUsage(context.Context, usage.Record) {}
type capturePluginExecutorUsagePlugin struct { targetProvider string records chan usage.Record}
func newCapturePluginExecutorUsagePlugin(targetProvider string) *capturePluginExecutorUsagePlugin { return &capturePluginExecutorUsagePlugin{ targetProvider: targetProvider, records: make(chan usage.Record, 50), }}
func (p *capturePluginExecutorUsagePlugin) HandleUsage(_ context.Context, record usage.Record) { if p.targetProvider != "" && record.Provider != p.targetProvider { return } select { case p.records <- record: default: }}
func (p *capturePluginExecutorUsagePlugin) waitRecord(t *testing.T) usage.Record { t.Helper() select { case rec := <-p.records: return rec case <-time.After(2 * time.Second): t.Fatal("timed out waiting for usage record") return usage.Record{} }}
func (p *capturePluginExecutorUsagePlugin) assertNoRecord(t *testing.T) { t.Helper() select { case rec := <-p.records: t.Fatalf("expected no usage record for %q, got %+v", p.targetProvider, rec) case <-time.After(50 * time.Millisecond): }}
func registerUsagePluginForTest(t *testing.T, name string, plugin usage.Plugin) { t.Helper() usage.RegisterNamedPlugin(name, plugin) t.Cleanup(func() { usage.RegisterNamedPlugin(name, noopUsagePlugin{}) })}
func TestHandlerPluginExecutorPublishesUsageNonStreamOpenAI(t *testing.T) { targetPluginID := "custom-openai-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-nonstream-openai", plugin)
originalModel := "gpt-4o"
openAIResponseBody := []byte(`{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":12,"completion_tokens":34,"total_tokens":46}}`)
mockHost := &mockPluginUsageHost{ execResp: coreexecutor.Response{ Payload: openAIResponseBody, }, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "") if errMsg != nil { t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg) } if len(body) == 0 { t.Fatal("empty response body") }
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if record.Stream { t.Errorf("record.Stream = true, want false") } if record.Detail.InputTokens != 12 || record.Detail.OutputTokens != 34 || record.Detail.TotalTokens != 46 { t.Errorf("record.Detail = %+v, want prompt=12 completion=34 total=46", record.Detail) }}
func TestHandlerPluginExecutorPublishesUsageStreamOpenAI(t *testing.T) { targetPluginID := "custom-stream-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-openai", plugin)
originalModel := "gpt-4o"
chunks := make(chan coreexecutor.StreamChunk, 3) chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n")} chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":15,\"completion_tokens\":25,\"total_tokens\":40}}\n\n")} chunks <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} close(chunks)
mockHost := &mockPluginUsageHost{ streamResult: &coreexecutor.StreamResult{Chunks: chunks}, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "") for range dataChan { } for err := range errChan { if err != nil { t.Fatalf("stream error = %+v", err) } }
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if !record.Stream { t.Errorf("record.Stream = false, want true") } if record.Detail.InputTokens != 15 || record.Detail.OutputTokens != 25 || record.Detail.TotalTokens != 40 { t.Errorf("record.Detail = %+v, want prompt=15 completion=25 total=40", record.Detail) }}
func TestHandlerPluginExecutorPublishesUsageStreamCodex(t *testing.T) { targetPluginID := "custom-codex-stream-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-codex", plugin)
originalModel := "codex-5.2"
chunks := make(chan coreexecutor.StreamChunk, 2) chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":18,\"output_tokens\":22,\"total_tokens\":40}}}\n\n")} chunks <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} close(chunks)
mockHost := &mockPluginUsageHost{ streamResult: &coreexecutor.StreamResult{Chunks: chunks}, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai-response", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "") for range dataChan { } for err := range errChan { if err != nil { t.Fatalf("stream error = %+v", err) } }
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if record.Detail.InputTokens != 18 || record.Detail.OutputTokens != 22 || record.Detail.TotalTokens != 40 { t.Errorf("record.Detail = %+v, want input=18 output=22 total=40", record.Detail) }}
func TestHandlerPluginExecutorPublishesUsageNonStreamClaude(t *testing.T) { targetPluginID := "custom-claude-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-nonstream-claude", plugin)
originalModel := "claude-3-5-sonnet"
claudeResponseBody := []byte(`{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hello"}],"usage":{"input_tokens":50,"output_tokens":30,"output_tokens_details":{"thinking_tokens":10}}}`)
mockHost := &mockPluginUsageHost{ execResp: coreexecutor.Response{ Payload: claudeResponseBody, }, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "claude", originalModel, []byte(fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hi"}]}`, originalModel)), "") if errMsg != nil { t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg) } if len(body) == 0 { t.Fatal("empty response body") }
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if record.Detail.InputTokens != 50 || record.Detail.OutputTokens != 30 || record.Detail.ReasoningTokens != 10 || record.Detail.TotalTokens != 80 { t.Errorf("record.Detail = %+v, want input=50 output=30 reasoning=10 total=80", record.Detail) }}
func TestHandlerPluginExecutorPublishesUsageStreamClaude(t *testing.T) { targetPluginID := "custom-claude-stream-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-claude", plugin)
originalModel := "claude-3-5-sonnet"
// Claude streams split usage between message_start (input, cache) and message_delta (output, thinking) chunks := make(chan coreexecutor.StreamChunk, 4) chunks <- coreexecutor.StreamChunk{Payload: []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"usage\":{\"input_tokens\":100,\"cache_read_input_tokens\":50,\"cache_creation_input_tokens\":20,\"output_tokens\":1}}}\n\n")} chunks <- coreexecutor.StreamChunk{Payload: []byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n")} chunks <- coreexecutor.StreamChunk{Payload: []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":25,\"output_tokens_details\":{\"thinking_tokens\":5}}}\n\n")} chunks <- coreexecutor.StreamChunk{Payload: []byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")} close(chunks)
mockHost := &mockPluginUsageHost{ streamResult: &coreexecutor.StreamResult{Chunks: chunks}, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "claude", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "") for range dataChan { } for err := range errChan { if err != nil { t.Fatalf("stream error = %+v", err) } }
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if record.Detail.InputTokens != 100 || record.Detail.CacheReadTokens != 50 || record.Detail.CacheCreationTokens != 20 || record.Detail.OutputTokens != 25 || record.Detail.ReasoningTokens != 5 || record.Detail.TotalTokens != 195 { t.Errorf("record.Detail = %+v, want input=100 cache_read=50 cache_creation=20 output=25 reasoning=5 total=195", record.Detail) } tb := record.Detail.TokenBreakdown if !tb.Valid() || tb.TotalTokens != 195 || tb.Input.TotalTokens != 170 || tb.Input.UncachedTokens != 100 || tb.Input.CacheReadTokens != 50 || tb.Input.CacheWriteTokens != 20 || tb.Output.TotalTokens != 25 || tb.Output.NonReasoningTokens != 20 || tb.Output.ReasoningTokens != 5 { t.Errorf("record.Detail.TokenBreakdown = %+v, want valid independent breakdown with total=195 input=170 output=25", tb) }}
func TestHandlerPluginExecutorPublishesUsageGemini(t *testing.T) { targetPluginID := "custom-gemini-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-gemini", plugin)
originalModel := "gemini-2.5-flash"
geminiResponseBody := []byte(`{"candidates":[{"content":{"parts":[{"text":"hello"}]}}],"usageMetadata":{"promptTokenCount":40,"candidatesTokenCount":60,"totalTokenCount":100}}`)
mockHost := &mockPluginUsageHost{ execResp: coreexecutor.Response{ Payload: geminiResponseBody, }, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "gemini", originalModel, []byte(fmt.Sprintf(`{"contents":[{"parts":[{"text":"hi"}]}]}`)), "") if errMsg != nil { t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg) } if len(body) == 0 { t.Fatal("empty response body") }
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if record.Detail.InputTokens != 40 || record.Detail.OutputTokens != 60 || record.Detail.TotalTokens != 100 { t.Errorf("record.Detail = %+v, want prompt=40 candidates=60 total=100", record.Detail) }}
func TestHandlerPluginExecutorPublishesUsageStreamInteractions(t *testing.T) { targetPluginID := "custom-interactions-stream-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-interactions", plugin)
originalModel := "gemini-2.5-flash"
chunks := make(chan coreexecutor.StreamChunk, 2) chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"event_type\":\"finish\",\"metadata\":{\"total_usage\":{\"total_input_tokens\":30,\"total_output_tokens\":70,\"total_tokens\":100}}}\n\n")} chunks <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} close(chunks)
mockHost := &mockPluginUsageHost{ streamResult: &coreexecutor.StreamResult{Chunks: chunks}, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "interactions", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "") for range dataChan { } for err := range errChan { if err != nil { t.Fatalf("stream error = %+v", err) } }
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if record.Detail.InputTokens != 30 || record.Detail.OutputTokens != 70 || record.Detail.TotalTokens != 100 { t.Errorf("record.Detail = %+v, want input=30 output=70 total=100", record.Detail) }}
func TestHandlerPluginExecutorPublishesUsageStreamAntigravity(t *testing.T) { targetPluginID := "custom-antigravity-stream-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-antigravity", plugin)
originalModel := "claude-3-5-sonnet"
chunks := make(chan coreexecutor.StreamChunk, 2) chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"response\":{\"usageMetadata\":{\"promptTokenCount\":33,\"candidatesTokenCount\":67,\"totalTokenCount\":100}}}\n\n")} chunks <- coreexecutor.StreamChunk{Payload: []byte("data: [DONE]\n\n")} close(chunks)
mockHost := &mockPluginUsageHost{ streamResult: &coreexecutor.StreamResult{Chunks: chunks}, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "antigravity", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "") for range dataChan { } for err := range errChan { if err != nil { t.Fatalf("stream error = %+v", err) } }
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if record.Detail.InputTokens != 33 || record.Detail.OutputTokens != 67 || record.Detail.TotalTokens != 100 { t.Errorf("record.Detail = %+v, want prompt=33 candidates=67 total=100", record.Detail) }}
func TestHandlerPluginExecutorPublishesFailure(t *testing.T) { targetPluginID := "failing-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-failure", plugin)
originalModel := "gpt-4o"
mockHost := &mockPluginUsageHost{ execErr: errors.New("upstream plugin failure"), } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
_, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), "") if errMsg == nil { t.Fatal("expected ExecuteWithAuthManager() to fail") }
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if !record.Failed { t.Error("record.Failed = false, want true") } if record.Fail.Body != "upstream plugin failure" { t.Errorf("record.Fail.Body = %q, want %q", record.Fail.Body, "upstream plugin failure") }}
func TestHandlerPluginExecutorPublishesStreamFailure(t *testing.T) { targetPluginID := "failing-stream-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-failure", plugin)
originalModel := "gpt-4o"
chunks := make(chan coreexecutor.StreamChunk, 2) chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n")} chunks <- coreexecutor.StreamChunk{Err: errors.New("upstream stream broke")} close(chunks)
mockHost := &mockPluginUsageHost{ streamResult: &coreexecutor.StreamResult{Chunks: chunks}, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "") for range dataChan { } seenErr := false for err := range errChan { if err != nil { seenErr = true } } if !seenErr { t.Fatal("expected stream error") }
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if !record.Failed { t.Error("record.Failed = false, want true") }}
func TestHandlerPluginExecutorPublishesStreamCancellation(t *testing.T) { targetPluginID := "canceling-stream-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-stream-cancel", plugin)
originalModel := "gpt-4o"
ctx, cancel := context.WithCancel(context.Background()) chunks := make(chan coreexecutor.StreamChunk, 5) chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n")}
mockHost := &mockPluginUsageHost{ streamResult: &coreexecutor.StreamResult{Chunks: chunks}, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
dataChan, _, _ := handler.ExecuteStreamWithAuthManager(ctx, "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "") // Receive first chunk then cancel context <-dataChan cancel() close(chunks)
record := plugin.waitRecord(t) if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } if !record.Failed { t.Error("record.Failed = false, want true for canceled stream") }}
func TestHandlerPluginExecutorSkipsUsageForNestedExecution(t *testing.T) { targetPluginID := "nested-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-usage-nested", plugin)
originalModel := "gpt-4o"
openAIResponseBody := []byte(`{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":12,"completion_tokens":34,"total_tokens":46}}`)
mockHost := &mockPluginUsageHost{ execResp: coreexecutor.Response{ Payload: openAIResponseBody, }, } mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true }
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) handler.SetModelRouterHost(mockHost)
// ExecuteModel triggers execution with InternalSource = true (host.model.execute callback) resp, errMsg := handler.ExecuteModel(context.Background(), ModelExecutionRequest{ EntryProtocol: "openai", ExitProtocol: "openai", Model: originalModel, Body: []byte(fmt.Sprintf(`{"model":%q}`, originalModel)), }) if errMsg != nil { t.Fatalf("ExecuteModel() error = %+v", errMsg) } if len(resp.Body) == 0 { t.Fatal("empty response body") }
plugin.assertNoRecord(t)}
func TestHandlerPluginExecutorSkipsOuterUsageWhenPluginCallsHostModelExecute(t *testing.T) { targetPluginID := "agent-wrapper-plugin" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-nested-callback", plugin)
outerModel := "agent-wrapper-model" innerModel := "inner-model"
manager := coreauth.NewManager(nil, nil, nil) innerExecutor := &modelExecutionCaptureExecutor{ provider: "openai", execute: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{ Payload: []byte(`{"id":"chatcmpl-inner","choices":[{"message":{"role":"assistant","content":"inner"}}],"usage":{"prompt_tokens":5,"completion_tokens":5,"total_tokens":10}}`), }, nil }, } manager.RegisterExecutor(innerExecutor) auth := &coreauth.Auth{ ID: "auth-" + innerModel, Provider: innerExecutor.Identifier(), Status: coreauth.StatusActive, } if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { t.Fatalf("manager.Register(): %v", errRegister) } registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: innerModel}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) })
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
mockHost := &mockPluginUsageHost{} mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { if req.RequestedModel == outerModel { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true } return pluginapi.ModelRouteResponse{}, false } // When the plugin executor executes, it simulates calling back into the host via ExecuteModel mockHost.execFunc = func(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) { // Plugin calls back into host.model.execute using the provided ctx innerResp, errInner := handler.ExecuteModel(ctx, ModelExecutionRequest{ EntryProtocol: "openai", ExitProtocol: "openai", Model: innerModel, Body: []byte(fmt.Sprintf(`{"model":%q}`, innerModel)), }) if errInner != nil { return coreexecutor.Response{}, errInner.Error } // Return response payload back to outer caller return coreexecutor.Response{Payload: innerResp.Body}, nil }
handler.SetModelRouterHost(mockHost)
body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", outerModel, []byte(fmt.Sprintf(`{"model":%q}`, outerModel)), "") if errMsg != nil { t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg) } if len(body) == 0 { t.Fatal("empty response body") }
// Since inner ExecuteModel was executed, the outer plugin executor must NOT publish a duplicate record for targetPluginID plugin.assertNoRecord(t)}
func TestHandlerPluginExecutorSkipsOuterFailureWhenPluginCallsHostModelExecute(t *testing.T) { targetPluginID := "agent-wrapper-plugin-failure" plugin := newCapturePluginExecutorUsagePlugin(targetPluginID) registerUsagePluginForTest(t, "test-plugin-executor-nested-callback-failure", plugin)
outerModel := "agent-wrapper-model-fail" innerModel := "inner-model-fail"
manager := coreauth.NewManager(nil, nil, nil) innerExecutor := &modelExecutionCaptureExecutor{ provider: "openai", execute: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, errors.New("inner failure") }, } manager.RegisterExecutor(innerExecutor) auth := &coreauth.Auth{ ID: "auth-" + innerModel, Provider: innerExecutor.Identifier(), Status: coreauth.StatusActive, } if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { t.Fatalf("manager.Register(): %v", errRegister) } registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: innerModel}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) })
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
mockHost := &mockPluginUsageHost{} mockHost.hasRouters = true mockHost.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { if req.RequestedModel == outerModel { return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true } return pluginapi.ModelRouteResponse{}, false } mockHost.execFunc = func(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) { _, errInner := handler.ExecuteModel(ctx, ModelExecutionRequest{ EntryProtocol: "openai", ExitProtocol: "openai", Model: innerModel, Body: []byte(fmt.Sprintf(`{"model":%q}`, innerModel)), }) if errInner != nil { return coreexecutor.Response{}, errInner.Error } return coreexecutor.Response{Payload: []byte("ok")}, nil }
handler.SetModelRouterHost(mockHost)
_, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", outerModel, []byte(fmt.Sprintf(`{"model":%q}`, outerModel)), "") if errMsg == nil { t.Fatal("expected failure") }
// Since inner ExecuteModel was executed, the outer plugin executor must NOT publish a duplicate failure record plugin.assertNoRecord(t)}
type mockPluginUsageHost struct { handlerDirectExecutorRouteHost execResp coreexecutor.Response execErr error execFunc func(context.Context, string, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) streamResult *coreexecutor.StreamResult}
func (h *mockPluginUsageHost) ExecutePluginExecutor(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) { h.lastPluginID = pluginID h.lastRequest = req h.lastOptions = opts if h.execFunc != nil { return h.execFunc(ctx, pluginID, req, opts) } if h.execErr != nil { return coreexecutor.Response{}, h.execErr } return h.execResp, nil}
func (h *mockPluginUsageHost) ExecutePluginExecutorStream(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { h.lastPluginID = pluginID h.lastRequest = req h.lastOptions = opts return h.streamResult, nil}