888 lines · 30.6 KB · Go
RawDownload
package pluginhost import (	"context"	"errors"	"net/http"	"testing"	"time" 	"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"	"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{

Lines 1–500 of 888.