diff --git a/internal/plugins/ai/ollama/ollama.go b/internal/plugins/ai/ollama/ollama.go index 20ef92e7..b967e735 100644 --- a/internal/plugins/ai/ollama/ollama.go +++ b/internal/plugins/ai/ollama/ollama.go @@ -106,6 +106,7 @@ func (o *Client) ListModels(_ context.Context) (ret []string, err error) { func (o *Client) SendStream(_ context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) (err error) { ctx := context.Background() + defer close(channel) var req ollamaapi.ChatRequest if req, err = o.createChatRequest(ctx, msgs, opts); err != nil { @@ -135,7 +136,6 @@ func (o *Client) SendStream(_ context.Context, msgs []*chat.ChatCompletionMessag return } - close(channel) return } diff --git a/internal/plugins/ai/ollama/ollama_test.go b/internal/plugins/ai/ollama/ollama_test.go index 4d814db6..25dbee53 100644 --- a/internal/plugins/ai/ollama/ollama_test.go +++ b/internal/plugins/ai/ollama/ollama_test.go @@ -6,10 +6,14 @@ import ( "fmt" "net/http" "net/http/httptest" + "net/url" "strings" "testing" + "github.com/danielmiessler/fabric/internal/chat" + "github.com/danielmiessler/fabric/internal/domain" "github.com/danielmiessler/fabric/internal/i18n" + ollamaapi "github.com/ollama/ollama/api" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -60,3 +64,28 @@ func TestLoadImageBytes_DataURLSuccess(t *testing.T) { require.NoError(t, err) assert.Equal(t, expected, got) } + +func TestSendStreamClosesChannelOnChatError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error":"ollama failed"}` + "\n")) + })) + t.Cleanup(server.Close) + + baseURL, err := url.Parse(server.URL) + require.NoError(t, err) + + client := &Client{client: ollamaapi.NewClient(baseURL, server.Client())} + channel := make(chan domain.StreamUpdate) + + err = client.SendStream( + context.Background(), + []*chat.ChatCompletionMessage{{Role: chat.ChatMessageRoleUser, Content: "hello"}}, + &domain.ChatOptions{Model: "missing-model"}, + channel, + ) + + require.Error(t, err) + _, ok := <-channel + assert.False(t, ok, "stream channel should be closed when Ollama chat returns an error") +}