From b935fa4007c2ef84126c6a440ca9791d148703fa Mon Sep 17 00:00:00 2001 From: phantom5099 <1011668688@qq.com> Date: Fri, 3 Apr 2026 21:39:02 +0800 Subject: [PATCH 1/2] =?UTF-8?q?pref:=E4=BA=8B=E4=BB=B6=E6=B5=81=E4=BF=A1?= =?UTF-8?q?=E6=81=AF=E7=BC=BA=E5=8F=A3=E4=B8=8E=E5=BA=9F=E5=BC=83=E5=87=BD?= =?UTF-8?q?=E6=95=B0=E6=B8=85=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/provider/openai/openai.go | 105 +++-- internal/provider/openai/openai_test.go | 522 ++++++++++++++---------- internal/provider/provider.go | 17 - internal/provider/types.go | 31 ++ 4 files changed, 412 insertions(+), 263 deletions(-) diff --git a/internal/provider/openai/openai.go b/internal/provider/openai/openai.go index c0f891c3..e19489b0 100644 --- a/internal/provider/openai/openai.go +++ b/internal/provider/openai/openai.go @@ -240,9 +240,37 @@ func (p *Provider) consumeStream(ctx context.Context, body io.Reader, events cha } } - for _, discovered := range mergeToolCallDeltas(toolCalls, choice.Delta.ToolCalls) { - if err := emitToolCallStart(ctx, events, discovered.ID, discovered.Name); err != nil { - return err + // 处理 tool call 增量 + for _, delta := range choice.Delta.ToolCalls { + call, exists := toolCalls[delta.Index] + if !exists { + call = &domain.ToolCall{} + toolCalls[delta.Index] = call + } + + // 首次发现且有名称,发送 tool_call_start + if !exists && strings.TrimSpace(delta.Function.Name) != "" { + call.ID = delta.ID + call.Name = delta.Function.Name + if err := emitToolCallStart(ctx, events, delta.Index, delta.ID, delta.Function.Name); err != nil { + return err + } + } else { + // 更新现有 tool call + if strings.TrimSpace(delta.ID) != "" { + call.ID = delta.ID + } + if strings.TrimSpace(delta.Function.Name) != "" { + call.Name = delta.Function.Name + } + } + + // 发送 tool_call_delta + if delta.Function.Arguments != "" { + call.Arguments += delta.Function.Arguments + if err := emitToolCallDelta(ctx, events, delta.Index, delta.Function.Arguments); err != nil { + return err + } } } } @@ -267,6 +295,10 @@ func (p *Provider) consumeStream(ctx context.Context, body io.Reader, events cha return domain.ChatResponse{}, flushErr } if done { + // 发送 message_done 事件 + if err := emitMessageDone(ctx, events, finishReason, &usage); err != nil { + return domain.ChatResponse{}, err + } return finalizeResponse(contentBuilder.String(), toolCalls, finishReason, usage), nil } case strings.HasPrefix(trimmed, ":"): @@ -277,6 +309,10 @@ func (p *Provider) consumeStream(ctx context.Context, body io.Reader, events cha if flushErr := flushEvent(); flushErr != nil { return domain.ChatResponse{}, flushErr } + // 发送 message_done 事件 + if err := emitMessageDone(ctx, events, finishReason, &usage); err != nil { + return domain.ChatResponse{}, err + } return finalizeResponse(contentBuilder.String(), toolCalls, finishReason, usage), nil } } @@ -342,16 +378,17 @@ func emitTextDelta(ctx context.Context, events chan<- domain.StreamEvent, text s } } -func emitToolCallStart(ctx context.Context, events chan<- domain.StreamEvent, id, name string) error { +func emitToolCallStart(ctx context.Context, events chan<- domain.StreamEvent, index int, id, name string) error { if events == nil || name == "" { return nil } select { case events <- domain.StreamEvent{ - Type: domain.StreamEventToolCallStart, - ToolCallID: id, - ToolName: name, + Type: domain.StreamEventToolCallStart, + ToolCallID: id, + ToolName: name, + ToolCallIndex: index, }: return nil case <-ctx.Done(): @@ -359,32 +396,40 @@ func emitToolCallStart(ctx context.Context, events chan<- domain.StreamEvent, id } } -// mergeToolCallDeltas 将流式增量合并到 target map 中,并返回本次新发现的 tool call 列表。 -// 当一个新的 tool call 索引首次出现且拥有非空工具名称时,视为新发现。 -func mergeToolCallDeltas(target map[int]*domain.ToolCall, deltas []toolCallDelta) []domain.ToolCall { - var discovered []domain.ToolCall - for _, delta := range deltas { - call, exists := target[delta.Index] - if !exists { - call = &domain.ToolCall{} - target[delta.Index] = call - } +// emitToolCallDelta 发送工具调用参数增量事件。 +func emitToolCallDelta(ctx context.Context, events chan<- domain.StreamEvent, index int, argumentsDelta string) error { + if events == nil || argumentsDelta == "" { + return nil + } - if strings.TrimSpace(delta.ID) != "" { - call.ID = delta.ID - } - if strings.TrimSpace(delta.Function.Name) != "" { - call.Name = delta.Function.Name - } - if delta.Function.Arguments != "" { - call.Arguments += delta.Function.Arguments - } + select { + case events <- domain.StreamEvent{ + Type: domain.StreamEventToolCallDelta, + ToolCallIndex: index, + ToolArgumentsDelta: argumentsDelta, + }: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} - if !exists && strings.TrimSpace(call.Name) != "" { - discovered = append(discovered, *call) - } +// emitMessageDone 发送消息完成事件。 +func emitMessageDone(ctx context.Context, events chan<- domain.StreamEvent, finishReason string, usage *domain.Usage) error { + if events == nil { + return nil + } + + select { + case events <- domain.StreamEvent{ + Type: domain.StreamEventMessageDone, + FinishReason: finishReason, + Usage: usage, + }: + return nil + case <-ctx.Done(): + return ctx.Err() } - return discovered } func finalizeResponse(content string, toolCalls map[int]*domain.ToolCall, finishReason string, usage domain.Usage) domain.ChatResponse { diff --git a/internal/provider/openai/openai_test.go b/internal/provider/openai/openai_test.go index 89c16c0c..5b8cc716 100644 --- a/internal/provider/openai/openai_test.go +++ b/internal/provider/openai/openai_test.go @@ -32,84 +32,6 @@ func resolvedConfig(baseURL string, model string) config.ResolvedProviderConfig } } -func TestMergeToolCallDeltas(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - deltas []toolCallDelta - assert func(t *testing.T, calls map[int]*domain.ToolCall) - assertDisc func(t *testing.T, disc []domain.ToolCall) - }{ - { - name: "single tool call fragments are merged by index", - deltas: []toolCallDelta{ - {Index: 0, ID: "call_1", Function: openAIFunctionCall{Name: "filesystem_edit", Arguments: `{"path":"a.go",`}}, - {Index: 0, Function: openAIFunctionCall{Arguments: `"search_string":"old",`}}, - {Index: 0, Function: openAIFunctionCall{Arguments: `"replace_string":"new"}`}}, - }, - assert: func(t *testing.T, calls map[int]*domain.ToolCall) { - t.Helper() - call := calls[0] - if call == nil { - t.Fatalf("expected call at index 0") - } - if call.Name != "filesystem_edit" { - t.Fatalf("expected name filesystem_edit, got %q", call.Name) - } - if call.Arguments != `{"path":"a.go","search_string":"old","replace_string":"new"}` { - t.Fatalf("unexpected arguments: %q", call.Arguments) - } - }, - assertDisc: func(t *testing.T, disc []domain.ToolCall) { - t.Helper() - if len(disc) != 1 { - t.Fatalf("expected 1 discovered tool call, got %d", len(disc)) - } - if disc[0].ID != "call_1" || disc[0].Name != "filesystem_edit" { - t.Fatalf("unexpected discovered call: %+v", disc[0]) - } - }, - }, - { - name: "multiple indices stay isolated", - deltas: []toolCallDelta{ - {Index: 0, ID: "call_1", Function: openAIFunctionCall{Name: "filesystem_read_file", Arguments: `{"path":"a.go"}`}}, - {Index: 1, ID: "call_2", Function: openAIFunctionCall{Name: "filesystem_write_file", Arguments: `{"path":"b.go"`}}, - {Index: 1, Function: openAIFunctionCall{Arguments: `,"content":"ok"}`}}, - }, - assert: func(t *testing.T, calls map[int]*domain.ToolCall) { - t.Helper() - if len(calls) != 2 { - t.Fatalf("expected 2 calls, got %d", len(calls)) - } - if calls[1].Arguments != `{"path":"b.go","content":"ok"}` { - t.Fatalf("unexpected second arguments: %q", calls[1].Arguments) - } - }, - assertDisc: func(t *testing.T, disc []domain.ToolCall) { - t.Helper() - if len(disc) != 2 { - t.Fatalf("expected 2 discovered tool calls, got %d", len(disc)) - } - }, - }, - } - - for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - calls := map[int]*domain.ToolCall{} - discovered := mergeToolCallDeltas(calls, tt.deltas) - tt.assert(t, calls) - if tt.assertDisc != nil { - tt.assertDisc(t, discovered) - } - }) - } -} - func TestProviderChatConsumesSSEAndMergesToolCalls(t *testing.T) { t.Setenv(config.OpenAIDefaultAPIKeyEnv, "test-key") @@ -498,16 +420,7 @@ func (r *readCloser) Close() error { return nil } -func containsString(items []string, target string) bool { - for _, item := range items { - if item == target { - return true - } - } - return false -} - -// --- emitToolCallStart / mergeToolCallDeltas 边界补充测试 --- +// --- emitToolCallStart 边界测试 --- func TestEmitToolCallStartGuards(t *testing.T) { t.Parallel() @@ -515,13 +428,13 @@ func TestEmitToolCallStartGuards(t *testing.T) { ctx := context.Background() // nil events 守卫 - if err := emitToolCallStart(ctx, nil, "call-1", "filesystem_edit"); err != nil { + if err := emitToolCallStart(ctx, nil, 0, "call-1", "filesystem_edit"); err != nil { t.Fatalf("expected nil events guard to return nil, got %v", err) } // 空 name 守卫 events := make(chan domain.StreamEvent, 1) - if err := emitToolCallStart(ctx, events, "call-1", ""); err != nil { + if err := emitToolCallStart(ctx, events, 0, "call-1", ""); err != nil { t.Fatalf("expected empty name guard to return nil, got %v", err) } select { @@ -531,145 +444,22 @@ func TestEmitToolCallStartGuards(t *testing.T) { } // 正常发送 - if err := emitToolCallStart(ctx, events, "call-1", "filesystem_edit"); err != nil { + if err := emitToolCallStart(ctx, events, 2, "call-1", "filesystem_edit"); err != nil { t.Fatalf("emitToolCallStart() error = %v", err) } got := <-events - if got.Type != domain.StreamEventToolCallStart || got.ToolName != "filesystem_edit" || got.ToolCallID != "call-1" { + if got.Type != domain.StreamEventToolCallStart || got.ToolName != "filesystem_edit" || got.ToolCallID != "call-1" || got.ToolCallIndex != 2 { t.Fatalf("unexpected event: %+v", got) } // context 取消 cancelledCtx, cancel := context.WithCancel(context.Background()) cancel() - if err := emitToolCallStart(cancelledCtx, make(chan domain.StreamEvent), "call-1", "filesystem_edit"); err == nil { + if err := emitToolCallStart(cancelledCtx, make(chan domain.StreamEvent), 0, "call-1", "filesystem_edit"); err == nil { t.Fatalf("expected cancellation error") } } -func TestMergeToolCallDeltasEdgeCases(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - deltas []toolCallDelta - assert func(t *testing.T, calls map[int]*domain.ToolCall) - assertDisc func(t *testing.T, disc []domain.ToolCall) - }{ - { - name: "empty deltas returns empty discovered", - deltas: []toolCallDelta{}, - assert: func(t *testing.T, calls map[int]*domain.ToolCall) { - t.Helper() - if len(calls) != 0 { - t.Fatalf("expected empty calls, got %d", len(calls)) - } - }, - assertDisc: func(t *testing.T, disc []domain.ToolCall) { - t.Helper() - if len(disc) != 0 { - t.Fatalf("expected empty discovered, got %d", len(disc)) - } - }, - }, - { - name: "only ID without Name is not discovered", - deltas: []toolCallDelta{ - {Index: 0, ID: "call-1"}, - }, - assert: func(t *testing.T, calls map[int]*domain.ToolCall) { - t.Helper() - if calls[0].ID != "call-1" { - t.Fatalf("expected ID to be set") - } - }, - assertDisc: func(t *testing.T, disc []domain.ToolCall) { - t.Helper() - if len(disc) != 0 { - t.Fatalf("expected 0 discovered when only ID, got %d", len(disc)) - } - }, - }, - { - name: "only Name without ID is discovered", - deltas: []toolCallDelta{ - {Index: 0, Function: openAIFunctionCall{Name: "filesystem_edit"}}, - }, - assert: func(t *testing.T, calls map[int]*domain.ToolCall) { - t.Helper() - if calls[0].Name != "filesystem_edit" { - t.Fatalf("expected Name to be set") - } - }, - assertDisc: func(t *testing.T, disc []domain.ToolCall) { - t.Helper() - if len(disc) != 1 { - t.Fatalf("expected 1 discovered, got %d", len(disc)) - } - if disc[0].Name != "filesystem_edit" { - t.Fatalf("expected discovered Name %q, got %q", "filesystem_edit", disc[0].Name) - } - }, - }, - { - name: "only Arguments without Name is not discovered", - deltas: []toolCallDelta{ - {Index: 0, Function: openAIFunctionCall{Arguments: `{"path":"x"}`}}, - }, - assert: func(t *testing.T, calls map[int]*domain.ToolCall) { - t.Helper() - if calls[0].Arguments != `{"path":"x"}` { - t.Fatalf("expected Arguments to be set") - } - }, - assertDisc: func(t *testing.T, disc []domain.ToolCall) { - t.Helper() - if len(disc) != 0 { - t.Fatalf("expected 0 discovered when only Arguments, got %d", len(disc)) - } - }, - }, - { - name: "cumulative calls across multiple rounds", - deltas: []toolCallDelta{ - {Index: 0, ID: "call-1", Function: openAIFunctionCall{Name: "tool_a"}}, - }, - assert: func(t *testing.T, calls map[int]*domain.ToolCall) { - t.Helper() - // 第二轮:添加新工具 - discovered := mergeToolCallDeltas(calls, []toolCallDelta{ - {Index: 1, ID: "call-2", Function: openAIFunctionCall{Name: "tool_b"}}, - }) - if len(discovered) != 1 || discovered[0].Name != "tool_b" { - t.Fatalf("expected 1 new discovered in second round, got %+v", discovered) - } - if len(calls) != 2 { - t.Fatalf("expected 2 calls after cumulative merge, got %d", len(calls)) - } - }, - assertDisc: func(t *testing.T, disc []domain.ToolCall) { - t.Helper() - if len(disc) != 1 || disc[0].ID != "call-1" { - t.Fatalf("expected first round discovered ID %q, got %+v", "call-1", disc) - } - }, - }, - } - - for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - calls := map[int]*domain.ToolCall{} - discovered := mergeToolCallDeltas(calls, tt.deltas) - tt.assert(t, calls) - if tt.assertDisc != nil { - tt.assertDisc(t, discovered) - } - }) - } -} - func TestProviderChatEmitsToolCallStartEvent(t *testing.T) { t.Setenv(config.OpenAIDefaultAPIKeyEnv, "test-key") @@ -735,3 +525,303 @@ func TestProviderChatEmitsToolCallStartEvent(t *testing.T) { t.Fatalf("expected StreamEventToolCallStart event in stream") } } + +// TestEmitToolCallDeltaGuards 测试 emitToolCallDelta 的边界保护。 +func TestEmitToolCallDeltaGuards(t *testing.T) { + t.Parallel() + + t.Run("nil events channel returns nil", func(t *testing.T) { + t.Parallel() + err := emitToolCallDelta(context.Background(), nil, 0, `{"path":"a.go"}`) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + }) + + t.Run("empty arguments returns nil", func(t *testing.T) { + t.Parallel() + events := make(chan domain.StreamEvent, 1) + err := emitToolCallDelta(context.Background(), events, 0, "") + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + if len(events) != 0 { + t.Fatalf("expected no events, got %d", len(events)) + } + }) + + t.Run("valid arguments sends event with correct fields", func(t *testing.T) { + t.Parallel() + events := make(chan domain.StreamEvent, 1) + err := emitToolCallDelta(context.Background(), events, 2, `{"path":"a.go"}`) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + + select { + case evt := <-events: + if evt.Type != domain.StreamEventToolCallDelta { + t.Fatalf("expected type %s, got %s", domain.StreamEventToolCallDelta, evt.Type) + } + if evt.ToolCallIndex != 2 { + t.Fatalf("expected ToolCallIndex %d, got %d", 2, evt.ToolCallIndex) + } + if evt.ToolArgumentsDelta != `{"path":"a.go"}` { + t.Fatalf("expected ToolArgumentsDelta %q, got %q", `{"path":"a.go"}`, evt.ToolArgumentsDelta) + } + default: + t.Fatal("expected event in channel") + } + }) + + t.Run("context cancellation returns error", func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + events := make(chan domain.StreamEvent, 1) + err := emitToolCallDelta(ctx, events, 0, `{"path":"a.go"}`) + if err == nil { + t.Fatal("expected error from canceled context") + } + if err != context.Canceled { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) +} + +// TestEmitMessageDoneGuards 测试 emitMessageDone 的边界保护。 +func TestEmitMessageDoneGuards(t *testing.T) { + t.Parallel() + + t.Run("nil events channel returns nil", func(t *testing.T) { + t.Parallel() + usage := &domain.Usage{InputTokens: 100, OutputTokens: 50, TotalTokens: 150} + err := emitMessageDone(context.Background(), nil, "stop", usage) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + }) + + t.Run("valid message_done sends event with correct fields", func(t *testing.T) { + t.Parallel() + events := make(chan domain.StreamEvent, 1) + usage := &domain.Usage{InputTokens: 100, OutputTokens: 50, TotalTokens: 150} + err := emitMessageDone(context.Background(), events, "tool_calls", usage) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + + select { + case evt := <-events: + if evt.Type != domain.StreamEventMessageDone { + t.Fatalf("expected type %s, got %s", domain.StreamEventMessageDone, evt.Type) + } + if evt.FinishReason != "tool_calls" { + t.Fatalf("expected FinishReason %q, got %q", "tool_calls", evt.FinishReason) + } + if evt.Usage == nil { + t.Fatal("expected Usage to be non-nil") + } + if evt.Usage.InputTokens != 100 || evt.Usage.OutputTokens != 50 || evt.Usage.TotalTokens != 150 { + t.Fatalf("expected Usage %+v, got %+v", usage, evt.Usage) + } + default: + t.Fatal("expected event in channel") + } + }) + + t.Run("nil usage is allowed", func(t *testing.T) { + t.Parallel() + events := make(chan domain.StreamEvent, 1) + err := emitMessageDone(context.Background(), events, "stop", nil) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + + select { + case evt := <-events: + if evt.Usage != nil { + t.Fatalf("expected Usage to be nil, got %+v", evt.Usage) + } + default: + t.Fatal("expected event in channel") + } + }) + + t.Run("context cancellation returns error", func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + events := make(chan domain.StreamEvent, 1) + usage := &domain.Usage{InputTokens: 100, OutputTokens: 50, TotalTokens: 150} + err := emitMessageDone(ctx, events, "stop", usage) + if err == nil { + t.Fatal("expected error from canceled context") + } + if err != context.Canceled { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) +} + +// TestProviderChatEmitsFullEventStream 测试完整的事件流(包括 tool_call_delta 和 message_done)。 +func TestProviderChatEmitsFullEventStream(t *testing.T) { + t.Setenv(config.OpenAIDefaultAPIKeyEnv, "test-key") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + // 发送文本 delta + writeSSEChunk(t, w, map[string]any{ + "choices": []map[string]any{ + { + "index": 0, + "delta": map[string]any{ + "content": "Hello", + }, + }, + }, + }) + + // 发送 tool call start 和 delta + writeSSEChunk(t, w, map[string]any{ + "choices": []map[string]any{ + { + "index": 0, + "delta": map[string]any{ + "tool_calls": []map[string]any{ + { + "index": 0, + "id": "call_tool_1", + "type": "function", + "function": map[string]any{ + "name": "filesystem_edit", + "arguments": `{"path":"a.`, + }, + }, + }, + }, + }, + }, + }) + + // 发送 tool call delta(参数增量) + writeSSEChunk(t, w, map[string]any{ + "choices": []map[string]any{ + { + "index": 0, + "delta": map[string]any{ + "tool_calls": []map[string]any{ + { + "index": 0, + "function": map[string]any{ + "arguments": `go"}`, + }, + }, + }, + }, + }, + }, + }) + + // 发送 usage 和 finish_reason + writeSSEChunk(t, w, map[string]any{ + "choices": []map[string]any{ + { + "index": 0, + "finish_reason": "tool_calls", + }, + }, + "usage": map[string]any{ + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": 150, + }, + }) + + _, _ = w.Write([]byte("data: [DONE]\n\n")) + })) + defer server.Close() + + provider, err := New(resolvedConfig(server.URL, config.OpenAIDefaultModel)) + if err != nil { + t.Fatalf("New() error = %v", err) + } + provider.client = server.Client() + + events := make(chan domain.StreamEvent, 16) + _, err = provider.Chat(context.Background(), domain.ChatRequest{ + Model: config.OpenAIDefaultModel, + Messages: []domain.Message{{Role: "user", Content: "test"}}, + }, events) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + close(events) + + var ( + foundTextDelta bool + foundToolCallStart bool + foundToolCallDelta bool + foundMessageDone bool + toolCallDeltaContent string + messageDoneEvt *domain.StreamEvent + ) + + for evt := range events { + switch evt.Type { + case domain.StreamEventTextDelta: + foundTextDelta = true + case domain.StreamEventToolCallStart: + foundToolCallStart = true + if evt.ToolName != "filesystem_edit" { + t.Fatalf("expected ToolName %q, got %q", "filesystem_edit", evt.ToolName) + } + if evt.ToolCallIndex != 0 { + t.Fatalf("expected ToolCallIndex %d for tool_call_start, got %d", 0, evt.ToolCallIndex) + } + case domain.StreamEventToolCallDelta: + foundToolCallDelta = true + toolCallDeltaContent += evt.ToolArgumentsDelta + case domain.StreamEventMessageDone: + foundMessageDone = true + messageDoneEvt = &evt + } + } + + if !foundTextDelta { + t.Fatal("expected StreamEventTextDelta event") + } + if !foundToolCallStart { + t.Fatal("expected StreamEventToolCallStart event") + } + if !foundToolCallDelta { + t.Fatal("expected StreamEventToolCallDelta event") + } + if !foundMessageDone { + t.Fatal("expected StreamEventMessageDone event") + } + + // 验证 tool_call_delta 内容被正确累加 + expectedDelta := `{"path":"a.go"}` + if toolCallDeltaContent != expectedDelta { + t.Fatalf("expected tool call delta content %q, got %q", expectedDelta, toolCallDeltaContent) + } + + // 验证 message_done 事件包含正确的字段 + if messageDoneEvt == nil { + t.Fatal("message_done event is nil") + } + if messageDoneEvt.FinishReason != "tool_calls" { + t.Fatalf("expected FinishReason %q, got %q", "tool_calls", messageDoneEvt.FinishReason) + } + if messageDoneEvt.Usage == nil { + t.Fatal("expected Usage in message_done event") + } + if messageDoneEvt.Usage.TotalTokens != 150 { + t.Fatalf("expected TotalTokens %d, got %d", 150, messageDoneEvt.Usage.TotalTokens) + } +} diff --git a/internal/provider/provider.go b/internal/provider/provider.go index 531dba4f..40a838b9 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -5,20 +5,3 @@ import "context" type Provider interface { Chat(ctx context.Context, req ChatRequest, events chan<- StreamEvent) (ChatResponse, error) } - -type StreamEventType string - -const ( - // StreamEventTextDelta 表示模型输出的文本片段。 - StreamEventTextDelta StreamEventType = "text_delta" - // StreamEventToolCallStart 表示模型开始请求工具调用,TUI 可据此展示过渡提示。 - StreamEventToolCallStart StreamEventType = "tool_call_start" -) - -// StreamEvent 表示 provider 驱动层向 runtime 推送的流式事件。 -type StreamEvent struct { - Type StreamEventType - Text string // 文本片段(text_delta 时使用) - ToolName string // 工具名称(tool_call_start 时使用) - ToolCallID string // 工具调用 ID(tool_call_start 时使用) -} diff --git a/internal/provider/types.go b/internal/provider/types.go index ee5b8443..87b34207 100644 --- a/internal/provider/types.go +++ b/internal/provider/types.go @@ -66,3 +66,34 @@ type ProviderSelection struct { ProviderID string `json:"provider_id"` ModelID string `json:"model_id"` } + +type StreamEventType string + +const ( + // StreamEventTextDelta 表示模型输出的文本片段。 + StreamEventTextDelta StreamEventType = "text_delta" + // StreamEventToolCallStart 表示模型开始请求工具调用,TUI 可据此展示过渡提示。 + StreamEventToolCallStart StreamEventType = "tool_call_start" + // StreamEventToolCallDelta 表示工具调用参数的增量片段。 + StreamEventToolCallDelta StreamEventType = "tool_call_delta" + // StreamEventMessageDone 表示本轮消息完成,包含最终统计信息。 + StreamEventMessageDone StreamEventType = "message_done" +) + +// StreamEvent 表示 provider 驱动层向 runtime 推送的流式事件。 +type StreamEvent struct { + Type StreamEventType + + // text_delta + Text string `json:"text,omitempty"` // 文本片段 + + // tool_call_start / tool_call_delta + ToolCallIndex int `json:"tool_call_index,omitempty"` // 工具调用索引 + ToolCallID string `json:"tool_call_id,omitempty"` // 工具调用 ID(tool_call_start 时使用) + ToolName string `json:"tool_name,omitempty"` // 工具名称(tool_call_start 时使用) + ToolArgumentsDelta string `json:"tool_arguments_delta,omitempty"` // 参数增量片段(tool_call_delta 时使用) + + // message_done + FinishReason string `json:"finish_reason,omitempty"` // 结束原因(仅 message_done 时有效) + Usage *Usage `json:"usage,omitempty"` // 使用统计(仅 message_done 时有效) +} From 9c336e3742cc9ae37a50325260b92ef8d5b39f47 Mon Sep 17 00:00:00 2001 From: phantom5099 <1011668688@qq.com> Date: Fri, 3 Apr 2026 22:02:05 +0800 Subject: [PATCH 2/2] =?UTF-8?q?test(provider):=E8=A1=A5=E5=85=A8=E6=B5=8B?= =?UTF-8?q?=E8=AF=95=E5=88=86=E6=94=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/provider/builtin/builtin_test.go | 28 ++ internal/provider/metadata_test.go | 171 ++++++++++++ internal/provider/openai/openai_test.go | 293 ++++++++++---------- internal/provider/registry_test.go | 170 ++++++++++++ internal/provider/selection/service_test.go | 147 ++++++++++ 5 files changed, 669 insertions(+), 140 deletions(-) diff --git a/internal/provider/builtin/builtin_test.go b/internal/provider/builtin/builtin_test.go index 8d39d296..fb230a0c 100644 --- a/internal/provider/builtin/builtin_test.go +++ b/internal/provider/builtin/builtin_test.go @@ -3,6 +3,7 @@ package builtin import ( "testing" + "neo-code/internal/provider" "neo-code/internal/provider/openai" ) @@ -20,3 +21,30 @@ func TestNewRegistry(t *testing.T) { t.Fatalf("expected registry to include %q driver", openai.DriverName) } } + +func TestRegister(t *testing.T) { + t.Parallel() + + t.Run("nil registry", func(t *testing.T) { + t.Parallel() + err := Register(nil) + if err == nil { + t.Fatal("expected error for nil registry") + } + if err.Error() != "builtin provider registry is nil" { + t.Fatalf("unexpected error: %v", err) + } + }) + + t.Run("valid registry", func(t *testing.T) { + t.Parallel() + registry := provider.NewRegistry() + err := Register(registry) + if err != nil { + t.Fatalf("Register() error = %v", err) + } + if !registry.Supports(openai.DriverName) { + t.Fatalf("expected registry to include %q driver", openai.DriverName) + } + }) +} diff --git a/internal/provider/metadata_test.go b/internal/provider/metadata_test.go index d205e196..0223d3b6 100644 --- a/internal/provider/metadata_test.go +++ b/internal/provider/metadata_test.go @@ -2,6 +2,177 @@ package provider import "testing" +func TestModelDescriptorsFromIDs(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + modelIDs []string + want []ModelDescriptor + }{ + { + name: "empty slice", + modelIDs: []string{}, + want: nil, + }, + { + name: "nil slice", + modelIDs: nil, + want: nil, + }, + { + name: "single model ID", + modelIDs: []string{"gpt-4"}, + want: []ModelDescriptor{ + {ID: "gpt-4", Name: "gpt-4"}, + }, + }, + { + name: "multiple model IDs", + modelIDs: []string{"gpt-4", "gpt-3.5-turbo"}, + want: []ModelDescriptor{ + {ID: "gpt-4", Name: "gpt-4"}, + {ID: "gpt-3.5-turbo", Name: "gpt-3.5-turbo"}, + }, + }, + { + name: "IDs with whitespace", + modelIDs: []string{" gpt-4 ", "\tgpt-3.5-turbo\t"}, + want: []ModelDescriptor{ + {ID: "gpt-4", Name: "gpt-4"}, + {ID: "gpt-3.5-turbo", Name: "gpt-3.5-turbo"}, + }, + }, + { + name: "empty strings are skipped", + modelIDs: []string{"gpt-4", "", " ", "gpt-3.5-turbo"}, + want: []ModelDescriptor{ + {ID: "gpt-4", Name: "gpt-4"}, + {ID: "gpt-3.5-turbo", Name: "gpt-3.5-turbo"}, + }, + }, + { + name: "all empty strings", + modelIDs: []string{"", " ", "\t"}, + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := modelDescriptorsFromIDs(tt.modelIDs) + if len(got) != len(tt.want) { + t.Fatalf("expected %d descriptors, got %d", len(tt.want), len(got)) + } + for i := range got { + if got[i].ID != tt.want[i].ID || got[i].Name != tt.want[i].Name { + t.Fatalf("descriptor %d: expected %+v, got %+v", i, tt.want[i], got[i]) + } + } + }) + } +} + +func TestFirstPositiveInt(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + values []any + want int + }{ + { + name: "empty values", + values: []any{}, + want: 0, + }, + { + name: "int positive", + values: []any{int(100)}, + want: 100, + }, + { + name: "int negative returns zero", + values: []any{int(-100)}, + want: 0, + }, + { + name: "int zero returns zero", + values: []any{int(0)}, + want: 0, + }, + { + name: "int32 positive", + values: []any{int32(200)}, + want: 200, + }, + { + name: "int32 negative returns zero", + values: []any{int32(-200)}, + want: 0, + }, + { + name: "int64 positive", + values: []any{int64(300)}, + want: 300, + }, + { + name: "int64 negative returns zero", + values: []any{int64(-300)}, + want: 0, + }, + { + name: "float32 positive", + values: []any{float32(400.5)}, + want: 400, + }, + { + name: "float32 negative returns zero", + values: []any{float32(-400.5)}, + want: 0, + }, + { + name: "float64 positive", + values: []any{float64(500.7)}, + want: 500, + }, + { + name: "float64 negative returns zero", + values: []any{float64(-500.7)}, + want: 0, + }, + { + name: "first positive wins", + values: []any{0, -10, 100, 200}, + want: 100, + }, + { + name: "all negative or zero", + values: []any{-1, -2, 0, -3}, + want: 0, + }, + { + name: "mixed types", + values: []any{int(0), int32(-10), int64(1000), float64(2000)}, + want: 1000, + }, + { + name: "non-numeric type ignored", + values: []any{"100", 200}, + want: 200, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := firstPositiveInt(tt.values...) + if got != tt.want { + t.Fatalf("expected %d, got %d", tt.want, got) + } + }) + } +} + func TestDescriptorFromRawModelNormalizesUsefulFields(t *testing.T) { t.Parallel() diff --git a/internal/provider/openai/openai_test.go b/internal/provider/openai/openai_test.go index 5b8cc716..4d9d33de 100644 --- a/internal/provider/openai/openai_test.go +++ b/internal/provider/openai/openai_test.go @@ -12,6 +12,159 @@ import ( domain "neo-code/internal/provider" ) +func TestDriver(t *testing.T) { + t.Parallel() + + driver := Driver() + if driver.Name != DriverName { + t.Fatalf("expected driver name %q, got %q", DriverName, driver.Name) + } + if driver.Build == nil { + t.Fatal("expected Build function to be non-nil") + } + if driver.Discover == nil { + t.Fatal("expected Discover function to be non-nil") + } +} + +func TestWithTransport(t *testing.T) { + t.Parallel() + + customTransport := &http.Transport{} + cfg := resolvedConfig("", "") + + provider, err := New(cfg, WithTransport(customTransport)) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + if provider.client.Transport != customTransport { + t.Fatal("expected custom transport to be set") + } +} + +func TestDefaultRetryTransport(t *testing.T) { + t.Parallel() + + transport := defaultRetryTransport() + if transport == nil { + t.Fatal("expected non-nil transport") + } +} + +func TestDiscoverModels(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/models" { + t.Fatalf("unexpected path: %s", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "gpt-4", "name": "GPT-4"}, + {"id": "gpt-3.5-turbo"}, + }, + }) + })) + defer server.Close() + + provider, err := New(resolvedConfig(server.URL, "")) + if err != nil { + t.Fatalf("New() error = %v", err) + } + provider.client = server.Client() + + models, err := provider.DiscoverModels(context.Background()) + if err != nil { + t.Fatalf("DiscoverModels() error = %v", err) + } + if len(models) != 2 { + t.Fatalf("expected 2 models, got %d", len(models)) + } + if models[0].ID != "gpt-4" || models[0].Name != "GPT-4" { + t.Fatalf("unexpected first model: %+v", models[0]) + } +} + +func TestEmitToolCallDelta(t *testing.T) { + t.Parallel() + + t.Run("nil events guard", func(t *testing.T) { + t.Parallel() + if err := emitToolCallDelta(context.Background(), nil, 0, "args"); err != nil { + t.Fatalf("expected nil events guard to return nil, got %v", err) + } + }) + + t.Run("empty arguments guard", func(t *testing.T) { + t.Parallel() + events := make(chan domain.StreamEvent, 1) + if err := emitToolCallDelta(context.Background(), events, 0, ""); err != nil { + t.Fatalf("expected empty arguments guard to return nil, got %v", err) + } + select { + case <-events: + t.Fatal("expected no event for empty arguments") + default: + } + }) + + t.Run("normal send", func(t *testing.T) { + t.Parallel() + events := make(chan domain.StreamEvent, 1) + if err := emitToolCallDelta(context.Background(), events, 3, `{"path":"main.go"}`); err != nil { + t.Fatalf("emitToolCallDelta() error = %v", err) + } + got := <-events + if got.Type != domain.StreamEventToolCallDelta || got.ToolCallIndex != 3 || got.ToolArgumentsDelta != `{"path":"main.go"}` { + t.Fatalf("unexpected event: %+v", got) + } + }) + + t.Run("context cancellation", func(t *testing.T) { + t.Parallel() + cancelledCtx, cancel := context.WithCancel(context.Background()) + cancel() + if err := emitToolCallDelta(cancelledCtx, make(chan domain.StreamEvent), 0, "args"); err == nil { + t.Fatal("expected cancellation error") + } + }) +} + +func TestEmitMessageDone(t *testing.T) { + t.Parallel() + + t.Run("nil events guard", func(t *testing.T) { + t.Parallel() + if err := emitMessageDone(context.Background(), nil, "stop", nil); err != nil { + t.Fatalf("expected nil events guard to return nil, got %v", err) + } + }) + + t.Run("normal send", func(t *testing.T) { + t.Parallel() + events := make(chan domain.StreamEvent, 1) + usage := &domain.Usage{TotalTokens: 100} + if err := emitMessageDone(context.Background(), events, "stop", usage); err != nil { + t.Fatalf("emitMessageDone() error = %v", err) + } + got := <-events + if got.Type != domain.StreamEventMessageDone || got.FinishReason != "stop" || got.Usage.TotalTokens != 100 { + t.Fatalf("unexpected event: %+v", got) + } + }) + + t.Run("context cancellation", func(t *testing.T) { + t.Parallel() + cancelledCtx, cancel := context.WithCancel(context.Background()) + cancel() + if err := emitMessageDone(cancelledCtx, make(chan domain.StreamEvent), "stop", nil); err == nil { + t.Fatal("expected cancellation error") + } + }) +} + func resolvedConfig(baseURL string, model string) config.ResolvedProviderConfig { if strings.TrimSpace(baseURL) == "" { baseURL = config.OpenAIDefaultBaseURL @@ -526,146 +679,6 @@ func TestProviderChatEmitsToolCallStartEvent(t *testing.T) { } } -// TestEmitToolCallDeltaGuards 测试 emitToolCallDelta 的边界保护。 -func TestEmitToolCallDeltaGuards(t *testing.T) { - t.Parallel() - - t.Run("nil events channel returns nil", func(t *testing.T) { - t.Parallel() - err := emitToolCallDelta(context.Background(), nil, 0, `{"path":"a.go"}`) - if err != nil { - t.Fatalf("expected nil error, got %v", err) - } - }) - - t.Run("empty arguments returns nil", func(t *testing.T) { - t.Parallel() - events := make(chan domain.StreamEvent, 1) - err := emitToolCallDelta(context.Background(), events, 0, "") - if err != nil { - t.Fatalf("expected nil error, got %v", err) - } - if len(events) != 0 { - t.Fatalf("expected no events, got %d", len(events)) - } - }) - - t.Run("valid arguments sends event with correct fields", func(t *testing.T) { - t.Parallel() - events := make(chan domain.StreamEvent, 1) - err := emitToolCallDelta(context.Background(), events, 2, `{"path":"a.go"}`) - if err != nil { - t.Fatalf("expected nil error, got %v", err) - } - - select { - case evt := <-events: - if evt.Type != domain.StreamEventToolCallDelta { - t.Fatalf("expected type %s, got %s", domain.StreamEventToolCallDelta, evt.Type) - } - if evt.ToolCallIndex != 2 { - t.Fatalf("expected ToolCallIndex %d, got %d", 2, evt.ToolCallIndex) - } - if evt.ToolArgumentsDelta != `{"path":"a.go"}` { - t.Fatalf("expected ToolArgumentsDelta %q, got %q", `{"path":"a.go"}`, evt.ToolArgumentsDelta) - } - default: - t.Fatal("expected event in channel") - } - }) - - t.Run("context cancellation returns error", func(t *testing.T) { - t.Parallel() - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - events := make(chan domain.StreamEvent, 1) - err := emitToolCallDelta(ctx, events, 0, `{"path":"a.go"}`) - if err == nil { - t.Fatal("expected error from canceled context") - } - if err != context.Canceled { - t.Fatalf("expected context.Canceled, got %v", err) - } - }) -} - -// TestEmitMessageDoneGuards 测试 emitMessageDone 的边界保护。 -func TestEmitMessageDoneGuards(t *testing.T) { - t.Parallel() - - t.Run("nil events channel returns nil", func(t *testing.T) { - t.Parallel() - usage := &domain.Usage{InputTokens: 100, OutputTokens: 50, TotalTokens: 150} - err := emitMessageDone(context.Background(), nil, "stop", usage) - if err != nil { - t.Fatalf("expected nil error, got %v", err) - } - }) - - t.Run("valid message_done sends event with correct fields", func(t *testing.T) { - t.Parallel() - events := make(chan domain.StreamEvent, 1) - usage := &domain.Usage{InputTokens: 100, OutputTokens: 50, TotalTokens: 150} - err := emitMessageDone(context.Background(), events, "tool_calls", usage) - if err != nil { - t.Fatalf("expected nil error, got %v", err) - } - - select { - case evt := <-events: - if evt.Type != domain.StreamEventMessageDone { - t.Fatalf("expected type %s, got %s", domain.StreamEventMessageDone, evt.Type) - } - if evt.FinishReason != "tool_calls" { - t.Fatalf("expected FinishReason %q, got %q", "tool_calls", evt.FinishReason) - } - if evt.Usage == nil { - t.Fatal("expected Usage to be non-nil") - } - if evt.Usage.InputTokens != 100 || evt.Usage.OutputTokens != 50 || evt.Usage.TotalTokens != 150 { - t.Fatalf("expected Usage %+v, got %+v", usage, evt.Usage) - } - default: - t.Fatal("expected event in channel") - } - }) - - t.Run("nil usage is allowed", func(t *testing.T) { - t.Parallel() - events := make(chan domain.StreamEvent, 1) - err := emitMessageDone(context.Background(), events, "stop", nil) - if err != nil { - t.Fatalf("expected nil error, got %v", err) - } - - select { - case evt := <-events: - if evt.Usage != nil { - t.Fatalf("expected Usage to be nil, got %+v", evt.Usage) - } - default: - t.Fatal("expected event in channel") - } - }) - - t.Run("context cancellation returns error", func(t *testing.T) { - t.Parallel() - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - events := make(chan domain.StreamEvent, 1) - usage := &domain.Usage{InputTokens: 100, OutputTokens: 50, TotalTokens: 150} - err := emitMessageDone(ctx, events, "stop", usage) - if err == nil { - t.Fatal("expected error from canceled context") - } - if err != context.Canceled { - t.Fatalf("expected context.Canceled, got %v", err) - } - }) -} - // TestProviderChatEmitsFullEventStream 测试完整的事件流(包括 tool_call_delta 和 message_done)。 func TestProviderChatEmitsFullEventStream(t *testing.T) { t.Setenv(config.OpenAIDefaultAPIKeyEnv, "test-key") diff --git a/internal/provider/registry_test.go b/internal/provider/registry_test.go index 031c2a99..3aefddf7 100644 --- a/internal/provider/registry_test.go +++ b/internal/provider/registry_test.go @@ -80,3 +80,173 @@ func TestRegistryRejectsDuplicateDriverRegistration(t *testing.T) { t.Fatalf("expected duplicate driver registration to fail") } } + +func TestRegistryDiscoverModels(t *testing.T) { + t.Parallel() + + t.Run("driver with discovery function", func(t *testing.T) { + t.Parallel() + + expectedModels := []provider.ModelDescriptor{ + {ID: "model-1", Name: "Model 1"}, + {ID: "model-2", Name: "Model 2"}, + } + + registry := provider.NewRegistry() + driver := provider.DriverDefinition{ + Name: "test-driver", + Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) { + return stubProvider{}, nil + }, + Discover: func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]provider.ModelDescriptor, error) { + return expectedModels, nil + }, + } + if err := registry.Register(driver); err != nil { + t.Fatalf("Register() error = %v", err) + } + + got, err := registry.DiscoverModels(context.Background(), config.ResolvedProviderConfig{ + ProviderConfig: config.ProviderConfig{Driver: "test-driver"}, + }) + if err != nil { + t.Fatalf("DiscoverModels() error = %v", err) + } + if len(got) != len(expectedModels) { + t.Fatalf("expected %d models, got %d", len(expectedModels), len(got)) + } + }) + + t.Run("driver without discovery function", func(t *testing.T) { + t.Parallel() + + registry := provider.NewRegistry() + driver := provider.DriverDefinition{ + Name: "test-driver", + Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) { + return stubProvider{}, nil + }, + Discover: nil, + } + if err := registry.Register(driver); err != nil { + t.Fatalf("Register() error = %v", err) + } + + got, err := registry.DiscoverModels(context.Background(), config.ResolvedProviderConfig{ + ProviderConfig: config.ProviderConfig{Driver: "test-driver"}, + }) + if err != nil { + t.Fatalf("DiscoverModels() error = %v", err) + } + if got != nil { + t.Fatalf("expected nil models, got %v", got) + } + }) + + t.Run("unknown driver", func(t *testing.T) { + t.Parallel() + + registry := provider.NewRegistry() + _, err := registry.DiscoverModels(context.Background(), config.ResolvedProviderConfig{ + ProviderConfig: config.ProviderConfig{Driver: "missing"}, + }) + if !errors.Is(err, provider.ErrDriverNotFound) { + t.Fatalf("expected ErrDriverNotFound, got %v", err) + } + }) +} + +func TestRegistrySupports(t *testing.T) { + t.Parallel() + + registry := provider.NewRegistry() + if err := registry.Register(stubDriver("openai")); err != nil { + t.Fatalf("Register() error = %v", err) + } + + tests := []struct { + driverType string + want bool + }{ + {"openai", true}, + {"OPENAI", true}, + {"OpenAI", true}, + {"missing", false}, + {"MISSING", false}, + {"", false}, + } + + for _, tt := range tests { + t.Run(tt.driverType, func(t *testing.T) { + got := registry.Supports(tt.driverType) + if got != tt.want { + t.Fatalf("Supports(%q) = %v, want %v", tt.driverType, got, tt.want) + } + }) + } +} + +func TestRegistryRegisterErrors(t *testing.T) { + t.Parallel() + + t.Run("nil registry", func(t *testing.T) { + t.Parallel() + + var registry *provider.Registry + err := registry.Register(stubDriver("test")) + if err == nil { + t.Fatal("expected error for nil registry") + } + if err.Error() != "provider: registry is nil" { + t.Fatalf("unexpected error: %v", err) + } + }) + + t.Run("empty driver name", func(t *testing.T) { + t.Parallel() + + registry := provider.NewRegistry() + err := registry.Register(provider.DriverDefinition{ + Name: " ", + Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) { + return nil, nil + }, + }) + if err == nil { + t.Fatal("expected error for empty driver name") + } + if err.Error() != "provider: driver name is empty" { + t.Fatalf("unexpected error: %v", err) + } + }) + + t.Run("nil build function", func(t *testing.T) { + t.Parallel() + + registry := provider.NewRegistry() + err := registry.Register(provider.DriverDefinition{ + Name: "test", + Build: nil, + }) + if err == nil { + t.Fatal("expected error for nil build function") + } + if err.Error() != `provider: driver "test" build func is nil` { + t.Fatalf("unexpected error: %v", err) + } + }) +} + +func TestRegistryDriverWithNilMap(t *testing.T) { + t.Parallel() + + registry := &provider.Registry{} + err := registry.Register(stubDriver("test")) + if err != nil { + t.Fatalf("Register() error = %v", err) + } + + if !registry.Supports("test") { + t.Fatal("expected registry to support test driver") + } +} diff --git a/internal/provider/selection/service_test.go b/internal/provider/selection/service_test.go index 36072b9a..33c68d8a 100644 --- a/internal/provider/selection/service_test.go +++ b/internal/provider/selection/service_test.go @@ -193,6 +193,153 @@ func TestServiceSelectProviderRejectsUnsupportedDriver(t *testing.T) { } } +func TestServiceListModelsSnapshot(t *testing.T) { + t.Parallel() + + manager := newTestManager(t, config.DefaultConfig()) + registry := newRegistry(t, config.OpenAIName, nil) + service := NewService(manager, registry, providercatalog.NewService("", registry, newCatalogStoreStub())) + + models, err := service.ListModelsSnapshot(context.Background()) + if err != nil { + t.Fatalf("ListModelsSnapshot() error = %v", err) + } + if len(models) != 1 || models[0].ID != config.OpenAIDefaultModel { + t.Fatalf("expected default model fallback, got %+v", models) + } +} + +func TestServiceValidateErrors(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + service *Service + errMsg string + }{ + { + name: "nil service", + service: nil, + errMsg: "provider selection: service is nil", + }, + { + name: "nil manager", + service: &Service{ + manager: nil, + registry: provider.NewRegistry(), + catalogs: providercatalog.NewService("", provider.NewRegistry(), newCatalogStoreStub()), + }, + errMsg: "provider selection: config manager is nil", + }, + { + name: "nil registry", + service: &Service{ + manager: config.NewManager(config.NewLoader(t.TempDir(), config.DefaultConfig())), + registry: nil, + catalogs: providercatalog.NewService("", provider.NewRegistry(), newCatalogStoreStub()), + }, + errMsg: "provider selection: registry is nil", + }, + { + name: "nil catalogs", + service: &Service{ + manager: config.NewManager(config.NewLoader(t.TempDir(), config.DefaultConfig())), + registry: provider.NewRegistry(), + catalogs: nil, + }, + errMsg: "provider selection: catalog service is nil", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + _, err := tt.service.ListProviders(context.Background()) + if err == nil || err.Error() != tt.errMsg { + t.Fatalf("expected error %q, got %v", tt.errMsg, err) + } + }) + } +} + +func TestServiceOperationsWithCanceledContext(t *testing.T) { + t.Parallel() + + manager := newTestManager(t, config.DefaultConfig()) + registry := newRegistry(t, config.OpenAIName, nil) + service := NewService(manager, registry, providercatalog.NewService("", registry, newCatalogStoreStub())) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + tests := []struct { + name string + fn func(context.Context) error + }{ + { + name: "ListProviders", + fn: func(ctx context.Context) error { + _, err := service.ListProviders(ctx) + return err + }, + }, + { + name: "ListModels", + fn: func(ctx context.Context) error { + _, err := service.ListModels(ctx) + return err + }, + }, + { + name: "ListModelsSnapshot", + fn: func(ctx context.Context) error { + _, err := service.ListModelsSnapshot(ctx) + return err + }, + }, + { + name: "EnsureSelection", + fn: func(ctx context.Context) error { + _, err := service.EnsureSelection(ctx) + return err + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := tt.fn(ctx) + if err == nil { + t.Fatalf("expected error for canceled context") + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) + } +} + +func TestServiceSetCurrentModelEmptyModelID(t *testing.T) { + t.Parallel() + + manager := newTestManager(t, config.DefaultConfig()) + registry := newRegistry(t, config.OpenAIName, nil) + service := NewService(manager, registry, providercatalog.NewService("", registry, newCatalogStoreStub())) + + _, err := service.SetCurrentModel(context.Background(), "") + if !errors.Is(err, provider.ErrModelNotFound) { + t.Fatalf("expected ErrModelNotFound for empty model ID, got %v", err) + } + + _, err = service.SetCurrentModel(context.Background(), " ") + if !errors.Is(err, provider.ErrModelNotFound) { + t.Fatalf("expected ErrModelNotFound for whitespace model ID, got %v", err) + } +} + func TestResolveCurrentModelHelper(t *testing.T) { t.Parallel()