Merge pull request #2182 from drawliin/fix/ollama-stream-channel-close

cmd/generate_changelog/incoming/2182.txt

@@ -0,0 +1,3 @@
+### PR [#2182](https://github.com/danielmiessler/Fabric/pull/2182) by [drawliin](https://github.com/drawliin): fix(ollama): close stream channel on errors
+
+- Fix(ollama): close stream channel on errors

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
 }
 

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")
+}