From 2e0abdd78a41ffd968b3c3e7e24afd789e30e3f1 Mon Sep 17 00:00:00 2001 From: Kayvan Sylvan Date: Wed, 25 Mar 2026 16:09:10 -0700 Subject: [PATCH] refactor: propagate `context.Context` through `Vendor` interface methods - Add `context.Context` parameter to `ListModels` interface method - Add `context.Context` parameter to `SendStream` interface method - Thread caller context through `Chatter.Send` instead of using `context.Background()` - Update all vendor implementations to accept context parameter - Introduce `publicError` type wrapping Codex provider errors - Normalize Codex error messages to lowercase for consistency - Add context-aware `sendStreamUpdate` helper in Codex client - Remove stale `context.Background()` calls from Anthropic and Azure AI Gateway - Update all vendor test mocks and stubs with new signatures - Pass HTTP request context from server handler into `chatter.Send` --- internal/cli/chat.go | 3 +- internal/core/chatter.go | 6 +-- internal/core/chatter_test.go | 14 +++---- internal/core/plugin_registry_test.go | 16 ++++---- internal/plugins/ai/anthropic/anthropic.go | 6 +-- .../plugins/ai/anthropic/anthropic_test.go | 5 ++- internal/plugins/ai/azure/azure.go | 3 +- internal/plugins/ai/azure/azure_test.go | 3 +- .../plugins/ai/azure_entra/azure_entra.go | 3 +- .../ai/azure_entra/azure_entra_test.go | 3 +- .../ai/azureaigateway/azureaigateway.go | 15 +++---- .../ai/azureaigateway/azureaigateway_test.go | 12 +++--- .../ai/azureaigateway/backend_azure_openai.go | 3 +- .../ai/azureaigateway/backend_bedrock.go | 3 +- .../ai/azureaigateway/backend_vertex_ai.go | 3 +- internal/plugins/ai/bedrock/bedrock.go | 4 +- internal/plugins/ai/bedrock/bedrock_test.go | 6 +-- internal/plugins/ai/codex/codex.go | 28 +++++++++---- internal/plugins/ai/codex/codex_test.go | 25 +++++++---- internal/plugins/ai/codex/errors.go | 41 ++++++++++++++----- internal/plugins/ai/codex/oauth.go | 10 ++--- internal/plugins/ai/codex/token.go | 4 +- internal/plugins/ai/copilot/copilot.go | 4 +- .../plugins/ai/digitalocean/digitalocean.go | 4 +- internal/plugins/ai/dryrun/dryrun.go | 4 +- internal/plugins/ai/dryrun/dryrun_test.go | 5 ++- internal/plugins/ai/exolab/exolab.go | 3 +- internal/plugins/ai/gemini/gemini.go | 4 +- internal/plugins/ai/lmstudio/lmstudio.go | 4 +- internal/plugins/ai/lmstudio/lmstudio_test.go | 4 +- internal/plugins/ai/ollama/ollama.go | 4 +- .../plugins/ai/openai/chat_completions.go | 4 +- internal/plugins/ai/openai/openai.go | 16 ++++---- .../ai/openai_compatible/providers_config.go | 6 +-- internal/plugins/ai/perplexity/perplexity.go | 4 +- internal/plugins/ai/vendor.go | 4 +- internal/plugins/ai/vendors.go | 2 +- internal/plugins/ai/vendors_test.go | 16 ++++---- internal/plugins/ai/vertexai/vertexai.go | 4 +- internal/server/chat.go | 2 +- 40 files changed, 178 insertions(+), 132 deletions(-) diff --git a/internal/cli/chat.go b/internal/cli/chat.go index 3489aa70..47f30b54 100644 --- a/internal/cli/chat.go +++ b/internal/cli/chat.go @@ -1,6 +1,7 @@ package cli import ( + "context" "errors" "fmt" "os" @@ -88,7 +89,7 @@ func handleChatProcessing(currentFlags *Flags, registry *core.PluginRegistry, me chatOptions.AudioFormat = "wav" // Default to WAV format } - if session, err = chatter.Send(chatReq, chatOptions); err != nil { + if session, err = chatter.Send(context.Background(), chatReq, chatOptions); err != nil { return } diff --git a/internal/core/chatter.go b/internal/core/chatter.go index fa3387dd..ff807de3 100644 --- a/internal/core/chatter.go +++ b/internal/core/chatter.go @@ -57,7 +57,7 @@ func joinPromptSections(parts ...string) string { } // Send processes a chat request and applies file changes for create_coding_feature pattern -func (o *Chatter) Send(request *domain.ChatRequest, opts *domain.ChatOptions) (session *fsdb.Session, err error) { +func (o *Chatter) Send(ctx context.Context, request *domain.ChatRequest, opts *domain.ChatOptions) (session *fsdb.Session, err error) { // Use o.model (normalized) for NeedsRawMode check instead of opts.Model // This ensures case-insensitive model names work correctly (e.g., "GPT-5" → "gpt-5") if o.vendor.NeedsRawMode(o.model) { @@ -107,7 +107,7 @@ func (o *Chatter) Send(request *domain.ChatRequest, opts *domain.ChatOptions) (s go func() { defer close(done) - if streamErr := o.vendor.SendStream(session.GetVendorMessages(), opts, responseChan); streamErr != nil { + if streamErr := o.vendor.SendStream(ctx, session.GetVendorMessages(), opts, responseChan); streamErr != nil { recordFirstStreamError(errChan, streamErr) } }() @@ -168,7 +168,7 @@ func (o *Chatter) Send(request *domain.ChatRequest, opts *domain.ChatOptions) (s // No errors, continue } } else { - if message, err = o.vendor.Send(context.Background(), session.GetVendorMessages(), opts); err != nil { + if message, err = o.vendor.Send(ctx, session.GetVendorMessages(), opts); err != nil { return } if debuglog.GetLevel() >= debuglog.Wire { diff --git a/internal/core/chatter_test.go b/internal/core/chatter_test.go index a02f855a..06f26be2 100644 --- a/internal/core/chatter_test.go +++ b/internal/core/chatter_test.go @@ -44,11 +44,11 @@ func (m *mockVendor) Setup() error { func (m *mockVendor) SetupFillEnvFileContent(*bytes.Buffer) { } -func (m *mockVendor) ListModels() ([]string, error) { +func (m *mockVendor) ListModels(context.Context) ([]string, error) { return []string{"test-model"}, nil } -func (m *mockVendor) SendStream(messages []*chat.ChatCompletionMessage, opts *domain.ChatOptions, responseChan chan domain.StreamUpdate) error { +func (m *mockVendor) SendStream(_ context.Context, messages []*chat.ChatCompletionMessage, opts *domain.ChatOptions, responseChan chan domain.StreamUpdate) error { // Send chunks if provided (for successful streaming test) if m.streamChunks != nil { for _, chunk := range m.streamChunks { @@ -180,7 +180,7 @@ func TestChatter_Send_SuppressThink(t *testing.T) { return "hidden visible", nil } - session, err := chatter.Send(request, opts) + session, err := chatter.Send(context.Background(), request, opts) if err != nil { t.Fatalf("Send returned error: %v", err) } @@ -296,7 +296,7 @@ func TestChatter_Send_StreamingErrorPropagation(t *testing.T) { } // Call Send and expect it to return the streaming error - session, err := chatter.Send(request, opts) + session, err := chatter.Send(context.Background(), request, opts) // Verify that the error from SendStream is propagated if err == nil { @@ -353,7 +353,7 @@ func TestChatter_Send_StreamingErrorUpdateAndReturnDoesNotDeadlock(t *testing.T) done := make(chan sendResult, 1) go func() { - session, err := chatter.Send(request, opts) + session, err := chatter.Send(context.Background(), request, opts) done <- sendResult{session: session, err: err} }() @@ -411,7 +411,7 @@ func TestChatter_Send_StreamingSuccessfulAggregation(t *testing.T) { } // Call Send and expect successful aggregation - session, err := chatter.Send(request, opts) + session, err := chatter.Send(context.Background(), request, opts) // Verify no error occurred if err != nil { @@ -494,7 +494,7 @@ func TestChatter_Send_StreamingMetadataPropagation(t *testing.T) { } // Call Send - _, err := chatter.Send(request, opts) + _, err := chatter.Send(context.Background(), request, opts) if err != nil { t.Fatalf("Expected no error, but got: %v", err) } diff --git a/internal/core/plugin_registry_test.go b/internal/core/plugin_registry_test.go index 87441899..654a1bb6 100644 --- a/internal/core/plugin_registry_test.go +++ b/internal/core/plugin_registry_test.go @@ -36,14 +36,14 @@ type testVendor struct { models []string } -func (m *testVendor) GetName() string { return m.name } -func (m *testVendor) GetSetupDescription() string { return m.name } -func (m *testVendor) IsConfigured() bool { return true } -func (m *testVendor) Configure() error { return nil } -func (m *testVendor) Setup() error { return nil } -func (m *testVendor) SetupFillEnvFileContent(*bytes.Buffer) {} -func (m *testVendor) ListModels() ([]string, error) { return m.models, nil } -func (m *testVendor) SendStream([]*chat.ChatCompletionMessage, *domain.ChatOptions, chan domain.StreamUpdate) error { +func (m *testVendor) GetName() string { return m.name } +func (m *testVendor) GetSetupDescription() string { return m.name } +func (m *testVendor) IsConfigured() bool { return true } +func (m *testVendor) Configure() error { return nil } +func (m *testVendor) Setup() error { return nil } +func (m *testVendor) SetupFillEnvFileContent(*bytes.Buffer) {} +func (m *testVendor) ListModels(context.Context) ([]string, error) { return m.models, nil } +func (m *testVendor) SendStream(context.Context, []*chat.ChatCompletionMessage, *domain.ChatOptions, chan domain.StreamUpdate) error { return nil } func (m *testVendor) Send(context.Context, []*chat.ChatCompletionMessage, *domain.ChatOptions) (string, error) { diff --git a/internal/plugins/ai/anthropic/anthropic.go b/internal/plugins/ai/anthropic/anthropic.go index cd9a60da..14e0825d 100644 --- a/internal/plugins/ai/anthropic/anthropic.go +++ b/internal/plugins/ai/anthropic/anthropic.go @@ -125,7 +125,7 @@ func (an *Client) configure() (err error) { return } -func (an *Client) ListModels() (ret []string, err error) { +func (an *Client) ListModels(context.Context) (ret []string, err error) { return an.models, nil } @@ -150,7 +150,7 @@ func parseThinking(level domain.ThinkingLevel) (anthropic.ThinkingConfigParamUni } func (an *Client) SendStream( - msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate, + ctx context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate, ) (err error) { messages := an.toMessages(msgs) if len(messages) == 0 { @@ -159,8 +159,6 @@ func (an *Client) SendStream( return } - ctx := context.Background() - params := an.buildMessageParams(messages, opts) betas := an.modelBetas[opts.Model] var reqOpts []option.RequestOption diff --git a/internal/plugins/ai/anthropic/anthropic_test.go b/internal/plugins/ai/anthropic/anthropic_test.go index ccbd9efd..3e012075 100644 --- a/internal/plugins/ai/anthropic/anthropic_test.go +++ b/internal/plugins/ai/anthropic/anthropic_test.go @@ -1,6 +1,7 @@ package anthropic import ( + "context" "strings" "testing" @@ -34,7 +35,7 @@ func TestNewClient_DefaultInitialization(t *testing.T) { func TestClientListModels(t *testing.T) { client := NewClient() - models, err := client.ListModels() + models, err := client.ListModels(context.Background()) if err != nil { t.Fatalf("Expected no error, got %v", err) } @@ -52,7 +53,7 @@ func TestClientListModels(t *testing.T) { func TestClient_ListModels_ReturnsCorrectModels(t *testing.T) { client := NewClient() - models, err := client.ListModels() + models, err := client.ListModels(context.Background()) if err != nil { t.Fatalf("Expected no error, got %v", err) diff --git a/internal/plugins/ai/azure/azure.go b/internal/plugins/ai/azure/azure.go index b2e78da5..6ec3b4c3 100644 --- a/internal/plugins/ai/azure/azure.go +++ b/internal/plugins/ai/azure/azure.go @@ -1,6 +1,7 @@ package azure import ( + "context" "errors" "strings" @@ -66,7 +67,7 @@ func (oi *Client) configure() error { return nil } -func (oi *Client) ListModels() (ret []string, err error) { +func (oi *Client) ListModels(context.Context) (ret []string, err error) { ret = oi.apiDeployments return } diff --git a/internal/plugins/ai/azure/azure_test.go b/internal/plugins/ai/azure/azure_test.go index cbaf4c0e..af548ab0 100644 --- a/internal/plugins/ai/azure/azure_test.go +++ b/internal/plugins/ai/azure/azure_test.go @@ -2,6 +2,7 @@ package azure import ( "bytes" + "context" "io" "net/http" "testing" @@ -78,7 +79,7 @@ func TestListModels(t *testing.T) { client := NewClient() client.apiDeployments = []string{"deployment1", "deployment2"} - models, err := client.ListModels() + models, err := client.ListModels(context.Background()) if err != nil { t.Fatalf("Expected no error, got %v", err) } diff --git a/internal/plugins/ai/azure_entra/azure_entra.go b/internal/plugins/ai/azure_entra/azure_entra.go index 2ee2d591..21c5f112 100644 --- a/internal/plugins/ai/azure_entra/azure_entra.go +++ b/internal/plugins/ai/azure_entra/azure_entra.go @@ -1,6 +1,7 @@ package azure_entra import ( + "context" "errors" "fmt" "strings" @@ -71,7 +72,7 @@ func (c *Client) configure() error { return nil } -func (c *Client) ListModels() (ret []string, err error) { +func (c *Client) ListModels(context.Context) (ret []string, err error) { ret = c.apiDeployments return } diff --git a/internal/plugins/ai/azure_entra/azure_entra_test.go b/internal/plugins/ai/azure_entra/azure_entra_test.go index d568d8da..6c6a2bf1 100644 --- a/internal/plugins/ai/azure_entra/azure_entra_test.go +++ b/internal/plugins/ai/azure_entra/azure_entra_test.go @@ -1,6 +1,7 @@ package azure_entra import ( + "context" "testing" ) @@ -47,7 +48,7 @@ func TestListModels(t *testing.T) { client := NewClient() client.apiDeployments = []string{"gpt-4o", "gpt-5"} - models, err := client.ListModels() + models, err := client.ListModels(context.Background()) if err != nil { t.Fatalf("Expected no error, got %v", err) } diff --git a/internal/plugins/ai/azureaigateway/azureaigateway.go b/internal/plugins/ai/azureaigateway/azureaigateway.go index 73fda949..a645f0b7 100644 --- a/internal/plugins/ai/azureaigateway/azureaigateway.go +++ b/internal/plugins/ai/azureaigateway/azureaigateway.go @@ -35,7 +35,7 @@ var _ ai.Vendor = (*Client)(nil) // are handled by the Client. type Backend interface { // ListModels returns the list of models available for this backend - ListModels() ([]string, error) + ListModels(context.Context) ([]string, error) // BuildEndpoint constructs the full API endpoint URL for the given model BuildEndpoint(baseURL, model string) string @@ -132,11 +132,11 @@ func (c *Client) IsConfigured() bool { } // ListModels delegates to the active backend -func (c *Client) ListModels() ([]string, error) { +func (c *Client) ListModels(ctx context.Context) ([]string, error) { if c.backend == nil { return nil, errors.New(i18n.T("azureaigateway_backend_not_initialized")) } - return c.backend.ListModels() + return c.backend.ListModels(ctx) } // Send sends a non-streaming request through the APIM gateway. @@ -199,18 +199,13 @@ func (c *Client) Send(ctx context.Context, msgs []*chat.ChatCompletionMessage, o } // SendStream falls back to non-streaming (APIM gateway doesn't support SSE pass-through). -// -// NOTE: This method uses context.Background() because the ai.Vendor interface does not -// accept a context parameter for SendStream. If the caller disconnects, this request will -// continue until the gateway timeout (300s). A future update to the ai.Vendor interface -// should add context propagation to SendStream. -func (c *Client) SendStream(msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) error { +func (c *Client) SendStream(ctx context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) error { defer close(channel) if c.backend == nil { return errors.New(i18n.T("azureaigateway_backend_not_initialized")) } - ctx, cancel := context.WithTimeout(context.Background(), gatewayTimeout) + ctx, cancel := context.WithTimeout(ctx, gatewayTimeout) defer cancel() result, err := c.Send(ctx, msgs, opts) diff --git a/internal/plugins/ai/azureaigateway/azureaigateway_test.go b/internal/plugins/ai/azureaigateway/azureaigateway_test.go index ea85271e..d6674ed5 100644 --- a/internal/plugins/ai/azureaigateway/azureaigateway_test.go +++ b/internal/plugins/ai/azureaigateway/azureaigateway_test.go @@ -60,7 +60,7 @@ func TestBedrockAuthHeader(t *testing.T) { func TestBedrockListModels(t *testing.T) { b := NewBedrockBackend("key") - models, err := b.ListModels() + models, err := b.ListModels(context.Background()) if err != nil { t.Fatalf("ListModels() error = %v", err) } @@ -253,7 +253,7 @@ func TestAzureOpenAIAuthHeader(t *testing.T) { func TestAzureOpenAIListModels(t *testing.T) { b := NewAzureOpenAIBackend("key", "") - models, err := b.ListModels() + models, err := b.ListModels(context.Background()) if err != nil { t.Fatalf("ListModels() error = %v", err) } @@ -336,7 +336,7 @@ func TestVertexAIAuthHeader(t *testing.T) { func TestVertexAIListModels(t *testing.T) { b := NewVertexAIBackend("key") - models, err := b.ListModels() + models, err := b.ListModels(context.Background()) if err != nil { t.Fatalf("ListModels() error = %v", err) } @@ -549,7 +549,7 @@ func TestConfigureInvalidBackend(t *testing.T) { func TestListModelsWithoutInit(t *testing.T) { c := NewClient() - _, err := c.ListModels() + _, err := c.ListModels(context.Background()) if err == nil { t.Error("ListModels() expected error when backend not initialized") } @@ -848,7 +848,7 @@ func TestSendStreamWithoutBackendInit(t *testing.T) { } channel := make(chan domain.StreamUpdate, 1) - err := c.SendStream(msgs, opts, channel) + err := c.SendStream(context.Background(), msgs, opts, channel) if err == nil { t.Fatal("SendStream() expected error when backend not initialized") } @@ -901,7 +901,7 @@ func TestSendStreamFallback(t *testing.T) { } channel := make(chan domain.StreamUpdate, 10) - err := c.SendStream(msgs, opts, channel) + err := c.SendStream(context.Background(), msgs, opts, channel) if err != nil { t.Fatalf("SendStream() error = %v", err) } diff --git a/internal/plugins/ai/azureaigateway/backend_azure_openai.go b/internal/plugins/ai/azureaigateway/backend_azure_openai.go index cb28614b..b6935645 100644 --- a/internal/plugins/ai/azureaigateway/backend_azure_openai.go +++ b/internal/plugins/ai/azureaigateway/backend_azure_openai.go @@ -2,6 +2,7 @@ package azureaigateway import ( + "context" "encoding/json" "errors" "fmt" @@ -34,7 +35,7 @@ func NewAzureOpenAIBackend(subscriptionKey, apiVersion string) *AzureOpenAIBacke // ListModels returns the list of models available through Azure OpenAI. // These are deployment names that must exist in your Azure OpenAI resource. -func (b *AzureOpenAIBackend) ListModels() ([]string, error) { +func (b *AzureOpenAIBackend) ListModels(_ context.Context) ([]string, error) { return []string{ "DeepSeek-R1", "gpt-4o", diff --git a/internal/plugins/ai/azureaigateway/backend_bedrock.go b/internal/plugins/ai/azureaigateway/backend_bedrock.go index f928e629..672e2029 100644 --- a/internal/plugins/ai/azureaigateway/backend_bedrock.go +++ b/internal/plugins/ai/azureaigateway/backend_bedrock.go @@ -2,6 +2,7 @@ package azureaigateway import ( + "context" "encoding/json" "errors" "fmt" @@ -27,7 +28,7 @@ func NewBedrockBackend(subscriptionKey string) *BedrockBackend { } // ListModels returns the list of available Bedrock inference profiles -func (b *BedrockBackend) ListModels() ([]string, error) { +func (b *BedrockBackend) ListModels(_ context.Context) ([]string, error) { return []string{ "us.anthropic.claude-3-haiku-20240307-v1:0", "us.anthropic.claude-3-opus-20240229-v1:0", diff --git a/internal/plugins/ai/azureaigateway/backend_vertex_ai.go b/internal/plugins/ai/azureaigateway/backend_vertex_ai.go index 59a50f5e..c8964f49 100644 --- a/internal/plugins/ai/azureaigateway/backend_vertex_ai.go +++ b/internal/plugins/ai/azureaigateway/backend_vertex_ai.go @@ -2,6 +2,7 @@ package azureaigateway import ( + "context" "encoding/json" "errors" "fmt" @@ -26,7 +27,7 @@ func NewVertexAIBackend(subscriptionKey string) *VertexAIBackend { } // ListModels returns the list of Gemini models available through Vertex AI -func (b *VertexAIBackend) ListModels() ([]string, error) { +func (b *VertexAIBackend) ListModels(_ context.Context) ([]string, error) { return []string{ "gemini-3-pro-preview", "gemini-2.5-pro", diff --git a/internal/plugins/ai/bedrock/bedrock.go b/internal/plugins/ai/bedrock/bedrock.go index 0a0b6c7b..27e1ea2a 100644 --- a/internal/plugins/ai/bedrock/bedrock.go +++ b/internal/plugins/ai/bedrock/bedrock.go @@ -443,7 +443,7 @@ func (c *BedrockClient) configure() error { // from AWS Bedrock that can be used with this plugin. // When using bearer token auth, the API may not be accessible, so a static // fallback list of common models is returned instead. -func (c *BedrockClient) ListModels() ([]string, error) { +func (c *BedrockClient) ListModels(_ context.Context) ([]string, error) { models, err := c.listModelsFromAPI() if err != nil && c.bedrockAPIKey.Value != "" { // Bearer token auth may lack ListFoundationModels permissions; @@ -488,7 +488,7 @@ func (c *BedrockClient) listModelsFromAPI() ([]string, error) { } // SendStream sends the messages to the Bedrock ConverseStream API -func (c *BedrockClient) SendStream(msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) (err error) { +func (c *BedrockClient) SendStream(_ context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) (err error) { // Ensure channel is closed on all exit paths to prevent goroutine leaks defer func() { if r := recover(); r != nil { diff --git a/internal/plugins/ai/bedrock/bedrock_test.go b/internal/plugins/ai/bedrock/bedrock_test.go index 0947bb7d..0adbf532 100644 --- a/internal/plugins/ai/bedrock/bedrock_test.go +++ b/internal/plugins/ai/bedrock/bedrock_test.go @@ -207,7 +207,7 @@ func TestListModels_NilClient_WithApiKey_ReturnsFallback(t *testing.T) { client.bedrockAPIKey.Value = "test-absk-token" // Don't call configure() — clients are nil - models, err := client.ListModels() + models, err := client.ListModels(context.Background()) assert.NoError(t, err, "ListModels should not error when falling back to static list") assert.Equal(t, defaultBedrockModels, models, "should return default models as fallback") } @@ -216,7 +216,7 @@ func TestListModels_NilClient_NoApiKey_ReturnsError(t *testing.T) { client := NewClient() // Don't call configure() and no API key — should propagate error - _, err := client.ListModels() + _, err := client.ListModels(context.Background()) assert.Error(t, err, "ListModels should error when client is nil and no API key for fallback") } @@ -227,7 +227,7 @@ func TestSendStream_NilClient_ReturnsError(t *testing.T) { ch := make(chan domain.StreamUpdate, 10) opts := &domain.ChatOptions{Model: "test-model", Temperature: 0.7, TopP: 0.9} - err := client.SendStream(nil, opts, ch) + err := client.SendStream(context.Background(), nil, opts, ch) assert.Error(t, err, "SendStream should return error when client is nil") assert.Contains(t, err.Error(), i18n.T("bedrock_client_not_initialized")) } diff --git a/internal/plugins/ai/codex/codex.go b/internal/plugins/ai/codex/codex.go index c8422abf..5fe6453f 100644 --- a/internal/plugins/ai/codex/codex.go +++ b/internal/plugins/ai/codex/codex.go @@ -45,6 +45,7 @@ const ( const oauthScope = "openid profile email offline_access api.connectors.read api.connectors.invoke" +// Client implements the Codex-backed AI vendor. type Client struct { *openaivendor.Client @@ -158,14 +159,14 @@ func (c *Client) configure() error { } // ListModels returns the Codex models available to the configured account. -func (c *Client) ListModels() ([]string, error) { +func (c *Client) ListModels(ctx context.Context) ([]string, error) { if c.apiHTTPClient == nil { if err := c.configure(); err != nil { return nil, err } } - ctx, cancel := context.WithTimeout(context.Background(), modelsRequestTimeout) + ctx, cancel := context.WithTimeout(ctx, modelsRequestTimeout) defer cancel() modelsURL := strings.TrimRight(c.ApiBaseURL.Value, "/") + "/models" @@ -194,7 +195,7 @@ func (c *Client) ListModels() ([]string, error) { var decoded modelsResponse if err := json.Unmarshal(body, &decoded); err != nil { - return nil, fmt.Errorf("failed to decode Codex models response: %w", err) + return nil, fmt.Errorf("failed to decode codex models response: %w", err) } models := make([]string, 0, len(decoded.Models)) @@ -250,7 +251,7 @@ func (c *Client) Send(ctx context.Context, msgs []*chat.ChatCompletionMessage, o // SendStream sends a request to Codex and streams the response text updates. func (c *Client) SendStream( - msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate, + ctx context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate, ) error { defer close(channel) @@ -264,15 +265,17 @@ func (c *Client) SendStream( } req := c.buildCodexResponseParams(msgs, opts) - stream := c.ApiClient.Responses.NewStreaming(context.Background(), req) + stream := c.ApiClient.Responses.NewStreaming(ctx, req) defer stream.Close() for stream.Next() { event := stream.Current() switch event.Type { case string(constant.ResponseOutputTextDelta("").Default()): - channel <- domain.StreamUpdate{ + if err := sendStreamUpdate(ctx, channel, domain.StreamUpdate{ Type: domain.StreamTypeContent, Content: event.AsResponseOutputTextDelta().Delta, + }); err != nil { + return err } case string(constant.ResponseOutputTextDone("").Default()): continue @@ -280,15 +283,26 @@ func (c *Client) SendStream( } if stream.Err() == nil { - channel <- domain.StreamUpdate{ + if err := sendStreamUpdate(ctx, channel, domain.StreamUpdate{ Type: domain.StreamTypeContent, Content: "\n", + }); err != nil { + return err } } return c.mapRequestError(stream.Err()) } +func sendStreamUpdate(ctx context.Context, channel chan domain.StreamUpdate, update domain.StreamUpdate) error { + select { + case <-ctx.Done(): + return ctx.Err() + case channel <- update: + return nil + } +} + func (c *Client) buildCodexResponseParams( msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, ) responses.ResponseNewParams { diff --git a/internal/plugins/ai/codex/codex_test.go b/internal/plugins/ai/codex/codex_test.go index 384b03f8..59899561 100644 --- a/internal/plugins/ai/codex/codex_test.go +++ b/internal/plugins/ai/codex/codex_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/base64" "encoding/json" + "errors" "fmt" "io" "net" @@ -176,7 +177,7 @@ func TestListModelsFiltersSupportedVisibleModels(t *testing.T) { client := newConfiguredTestClient(t, modelsServer.URL, "acct_models", testJWT("acct_models", time.Now().Add(time.Hour))) - models, err := client.ListModels() + models, err := client.ListModels(context.Background()) if err != nil { t.Fatalf("ListModels() error = %v", err) } @@ -220,8 +221,11 @@ func TestMapRequestErrorPreservesCodexAPIErrorMessage(t *testing.T) { if err == nil { t.Fatal("mapRequestError() returned nil") } - if got := err.Error(); got != "The requested model is not supported." { - t.Fatalf("mapRequestError() = %q, want %q", got, "The requested model is not supported.") + if got := err.Error(); got != "codex request failed with status 400" { + t.Fatalf("mapRequestError() = %q, want %q", got, "codex request failed with status 400") + } + if unwrapped := errors.Unwrap(err); unwrapped == nil || !strings.Contains(unwrapped.Error(), "The requested model is not supported.") { + t.Fatalf("wrapped error = %v, want provider detail", unwrapped) } } @@ -238,8 +242,11 @@ func TestMapRequestErrorReadsAPIErrorResponseBodyWhenRawJSONMissing(t *testing.T if err == nil { t.Fatal("mapRequestError() returned nil") } - if got := err.Error(); got != "The requested model is not supported for Codex." { - t.Fatalf("mapRequestError() = %q, want %q", got, "The requested model is not supported for Codex.") + if got := err.Error(); got != "codex request failed with status 400" { + t.Fatalf("mapRequestError() = %q, want %q", got, "codex request failed with status 400") + } + if unwrapped := errors.Unwrap(err); unwrapped == nil || !strings.Contains(unwrapped.Error(), "The requested model is not supported for Codex.") { + t.Fatalf("wrapped error = %v, want provider detail", unwrapped) } } @@ -469,7 +476,7 @@ func TestSendStreamReadsCodexSSE(t *testing.T) { client := newConfiguredTestClient(t, apiServer.URL, "acct_stream", testJWT("acct_stream", time.Now().Add(time.Hour))) updates := make(chan domain.StreamUpdate, 8) - err := client.SendStream([]*chat.ChatCompletionMessage{ + err := client.SendStream(context.Background(), []*chat.ChatCompletionMessage{ {Role: chat.ChatMessageRoleSystem, Content: "Follow the system prompt"}, {Role: "user", Content: "Hello"}, }, &domain.ChatOptions{ @@ -505,7 +512,7 @@ func TestSendStreamClosesChannelAndMapsHTTPError(t *testing.T) { client := newConfiguredTestClient(t, apiServer.URL, "acct_stream_error", testJWT("acct_stream_error", time.Now().Add(time.Hour))) updates := make(chan domain.StreamUpdate, 1) - err := client.SendStream([]*chat.ChatCompletionMessage{ + err := client.SendStream(context.Background(), []*chat.ChatCompletionMessage{ {Role: chat.ChatMessageRoleUser, Content: "Hello"}, }, &domain.ChatOptions{ Model: "gpt-5.4", @@ -513,8 +520,8 @@ func TestSendStreamClosesChannelAndMapsHTTPError(t *testing.T) { if err == nil { t.Fatal("SendStream() error = nil, want mapped HTTP error") } - if got := err.Error(); got != "usage limit reached" { - t.Fatalf("SendStream() error = %q, want %q", got, "usage limit reached") + if got := err.Error(); got != "codex usage limit reached" { + t.Fatalf("SendStream() error = %q, want %q", got, "codex usage limit reached") } update, ok := <-updates diff --git a/internal/plugins/ai/codex/errors.go b/internal/plugins/ai/codex/errors.go index fa825e74..dfd555f3 100644 --- a/internal/plugins/ai/codex/errors.go +++ b/internal/plugins/ai/codex/errors.go @@ -12,18 +12,28 @@ import ( openaiapi "github.com/openai/openai-go" ) +type publicError struct { + message string + cause error +} + +func (e *publicError) Error() string { + return e.message +} + +func (e *publicError) Unwrap() error { + return e.cause +} + func (c *Client) errorFromHTTPResponse(statusCode int, body []byte) error { message := extractErrorMessage(body) if statusCode == http.StatusUnauthorized { return errors.New(i18n.T("codex_login_invalid")) } if isUsageLimitMessage(message) { - return errors.New(message) + return wrapPublicError("codex usage limit reached", statusCode, message) } - if message == "" { - message = fmt.Sprintf("Codex request failed with status %d", statusCode) - } - return errors.New(message) + return wrapPublicError(fmt.Sprintf("codex request failed with status %d", statusCode), statusCode, message) } func (c *Client) refreshErrorFromResponse(statusCode int, body []byte) error { @@ -39,10 +49,7 @@ func (c *Client) refreshErrorFromResponse(statusCode int, body []byte) error { } } - if message == "" { - message = fmt.Sprintf("failed to refresh Codex login (status %d)", statusCode) - } - return errors.New(message) + return wrapPublicError(fmt.Sprintf("failed to refresh codex login (status %d)", statusCode), statusCode, message) } func (c *Client) mapRequestError(err error) error { @@ -69,12 +76,26 @@ func (c *Client) mapRequestError(err error) error { strings.Contains(lower, "chatgpt login"): return errors.New(i18n.T("codex_login_invalid")) case isUsageLimitMessage(message): - return errors.New(message) + return &publicError{ + message: "codex usage limit reached", + cause: fmt.Errorf("codex request failed: %w", err), + } default: return err } } +func wrapPublicError(message string, statusCode int, providerMessage string) error { + if providerMessage == "" { + return errors.New(message) + } + + return &publicError{ + message: message, + cause: fmt.Errorf("codex provider error (status %d): %s", statusCode, providerMessage), + } +} + func readAPIErrorBody(apiErr *openaiapi.Error) []byte { if apiErr == nil || apiErr.Response == nil || apiErr.Response.Body == nil { return nil diff --git a/internal/plugins/ai/codex/oauth.go b/internal/plugins/ai/codex/oauth.go index 1d4cf049..0c6673ce 100644 --- a/internal/plugins/ai/codex/oauth.go +++ b/internal/plugins/ai/codex/oauth.go @@ -56,7 +56,7 @@ func (c *Client) runOAuthFlow( ) (oauthTokens, error) { listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", defaultCallbackPort)) if err != nil { - return oauthTokens{}, fmt.Errorf("failed to start local OAuth callback server: %w", err) + return oauthTokens{}, fmt.Errorf("failed to start local oauth callback server: %w", err) } defer listener.Close() debuglog.Debug(debuglog.Detailed, "Codex OAuth callback listener started on 127.0.0.1:%d\n", defaultCallbackPort) @@ -216,7 +216,7 @@ func (c *Client) exchangeCodeForTokens( resp, err := c.authHTTPClient.Do(req) if err != nil { - return oauthTokens{}, fmt.Errorf("Codex token exchange failed: %w", err) + return oauthTokens{}, fmt.Errorf("codex token exchange failed: %w", err) } defer resp.Body.Close() @@ -230,7 +230,7 @@ func (c *Client) exchangeCodeForTokens( var tokens oauthTokens if err := json.Unmarshal(body, &tokens); err != nil { - return oauthTokens{}, fmt.Errorf("failed to decode Codex token exchange response: %w", err) + return oauthTokens{}, fmt.Errorf("failed to decode codex token exchange response: %w", err) } if strings.TrimSpace(tokens.AccessToken) == "" || strings.TrimSpace(tokens.RefreshToken) == "" { return oauthTokens{}, errors.New(i18n.T("codex_login_missing_tokens")) @@ -242,7 +242,7 @@ func (c *Client) exchangeCodeForTokens( func buildAuthorizeURL(authBaseURL string, callbackURL string, pkce pkceCodes, state string) (string, error) { issuer, err := url.Parse(strings.TrimRight(authBaseURL, "/")) if err != nil { - return "", fmt.Errorf("invalid Codex auth base URL: %w", err) + return "", fmt.Errorf("invalid codex auth base url: %w", err) } issuer.Path = strings.TrimRight(issuer.Path, "/") + "/oauth/authorize" @@ -278,7 +278,7 @@ func generatePKCECodes() (pkceCodes, error) { func randomBase64URL(size int) (string, error) { buf := make([]byte, size) if _, err := rand.Read(buf); err != nil { - return "", fmt.Errorf("failed to generate secure random OAuth state: %w", err) + return "", fmt.Errorf("failed to generate secure random oauth state: %w", err) } return base64.RawURLEncoding.EncodeToString(buf), nil } diff --git a/internal/plugins/ai/codex/token.go b/internal/plugins/ai/codex/token.go index 1ed55dc9..1068a2f7 100644 --- a/internal/plugins/ai/codex/token.go +++ b/internal/plugins/ai/codex/token.go @@ -42,7 +42,7 @@ func extractExpiryFromJWT(jwt string) (time.Time, error) { return time.Time{}, err } if claims.Exp == 0 { - return time.Time{}, errors.New("JWT did not include an exp claim") + return time.Time{}, errors.New("jwt did not include an exp claim") } return time.Unix(claims.Exp, 0), nil } @@ -58,7 +58,7 @@ func extractAccountIDFromJWT(jwt string) (string, error) { func parseTokenClaims(jwt string) (tokenClaims, error) { parts := strings.Split(jwt, ".") if len(parts) < 2 { - return tokenClaims{}, errors.New("invalid JWT format") + return tokenClaims{}, errors.New("invalid jwt format") } payload, err := base64.RawURLEncoding.DecodeString(parts[1]) diff --git a/internal/plugins/ai/copilot/copilot.go b/internal/plugins/ai/copilot/copilot.go index ab496027..5f3fb379 100644 --- a/internal/plugins/ai/copilot/copilot.go +++ b/internal/plugins/ai/copilot/copilot.go @@ -159,7 +159,7 @@ func (c *Client) IsConfigured() bool { // ListModels returns the available models. // Microsoft 365 Copilot exposes a single model - the Copilot service itself. -func (c *Client) ListModels() ([]string, error) { +func (c *Client) ListModels(_ context.Context) ([]string, error) { // Copilot doesn't expose multiple models - it's a unified service // We expose it as a single "model" for consistency with Fabric's architecture return []string{copilotModelName}, nil @@ -186,7 +186,7 @@ func (c *Client) Send(ctx context.Context, msgs []*chat.ChatCompletionMessage, o } // SendStream sends a message to Copilot and streams the response. -func (c *Client) SendStream(msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) error { +func (c *Client) SendStream(_ context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) error { defer close(channel) ctx := context.Background() diff --git a/internal/plugins/ai/digitalocean/digitalocean.go b/internal/plugins/ai/digitalocean/digitalocean.go index 45059093..803937d2 100644 --- a/internal/plugins/ai/digitalocean/digitalocean.go +++ b/internal/plugins/ai/digitalocean/digitalocean.go @@ -52,9 +52,9 @@ func NewClient() *Client { return client } -func (c *Client) ListModels() ([]string, error) { +func (c *Client) ListModels(ctx context.Context) ([]string, error) { if c.ControlPlaneToken.Value == "" { - models, err := c.Client.ListModels() + models, err := c.Client.ListModels(ctx) if err == nil && len(models) > 0 { return models, nil } diff --git a/internal/plugins/ai/dryrun/dryrun.go b/internal/plugins/ai/dryrun/dryrun.go index 0189117a..f8ff51db 100644 --- a/internal/plugins/ai/dryrun/dryrun.go +++ b/internal/plugins/ai/dryrun/dryrun.go @@ -22,7 +22,7 @@ func NewClient() *Client { return &Client{PluginBase: &plugins.PluginBase{Name: "DryRun"}} } -func (c *Client) ListModels() ([]string, error) { +func (c *Client) ListModels(_ context.Context) ([]string, error) { return []string{"dry-run-model"}, nil } @@ -108,7 +108,7 @@ func (c *Client) constructRequest(msgs []*chat.ChatCompletionMessage, opts *doma return builder.String() } -func (c *Client) SendStream(msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) error { +func (c *Client) SendStream(_ context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) error { defer close(channel) request := c.constructRequest(msgs, opts) channel <- domain.StreamUpdate{ diff --git a/internal/plugins/ai/dryrun/dryrun_test.go b/internal/plugins/ai/dryrun/dryrun_test.go index a5db58c3..720b02f5 100644 --- a/internal/plugins/ai/dryrun/dryrun_test.go +++ b/internal/plugins/ai/dryrun/dryrun_test.go @@ -1,6 +1,7 @@ package dryrun import ( + "context" "reflect" "testing" @@ -11,7 +12,7 @@ import ( // Test generated using Keploy func TestListModels_ReturnsExpectedModel(t *testing.T) { client := NewClient() - models, err := client.ListModels() + models, err := client.ListModels(context.Background()) if err != nil { t.Fatalf("Expected no error, got %v", err) } @@ -41,7 +42,7 @@ func TestSendStream_SendsMessages(t *testing.T) { } channel := make(chan domain.StreamUpdate) go func() { - err := client.SendStream(msgs, opts, channel) + err := client.SendStream(context.Background(), msgs, opts, channel) if err != nil { t.Errorf("Expected no error, got %v", err) } diff --git a/internal/plugins/ai/exolab/exolab.go b/internal/plugins/ai/exolab/exolab.go index b3f4fabc..9964ef28 100644 --- a/internal/plugins/ai/exolab/exolab.go +++ b/internal/plugins/ai/exolab/exolab.go @@ -1,6 +1,7 @@ package exolab import ( + "context" "strings" "github.com/danielmiessler/fabric/internal/plugins" @@ -42,7 +43,7 @@ func (oi *Client) configure() (err error) { return } -func (oi *Client) ListModels() (ret []string, err error) { +func (oi *Client) ListModels(context.Context) (ret []string, err error) { ret = oi.apiModels return } diff --git a/internal/plugins/ai/gemini/gemini.go b/internal/plugins/ai/gemini/gemini.go index 92ee8263..908047fb 100644 --- a/internal/plugins/ai/gemini/gemini.go +++ b/internal/plugins/ai/gemini/gemini.go @@ -60,7 +60,7 @@ type Client struct { ApiKey *plugins.SetupQuestion } -func (o *Client) ListModels() (ret []string, err error) { +func (o *Client) ListModels(_ context.Context) (ret []string, err error) { ctx := context.Background() var client *genai.Client if client, err = genai.NewClient(ctx, &genai.ClientConfig{ @@ -124,7 +124,7 @@ func (o *Client) Send(ctx context.Context, msgs []*chat.ChatCompletionMessage, o return } -func (o *Client) SendStream(msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) (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) diff --git a/internal/plugins/ai/lmstudio/lmstudio.go b/internal/plugins/ai/lmstudio/lmstudio.go index 3e874cfd..3241e76a 100644 --- a/internal/plugins/ai/lmstudio/lmstudio.go +++ b/internal/plugins/ai/lmstudio/lmstudio.go @@ -52,7 +52,7 @@ func (c *Client) configure() error { } // ListModels returns a list of available models. -func (c *Client) ListModels() ([]string, error) { +func (c *Client) ListModels(_ context.Context) ([]string, error) { url := fmt.Sprintf("%s/models", c.ApiUrl.Value) req, err := http.NewRequest("GET", url, nil) @@ -89,7 +89,7 @@ func (c *Client) ListModels() ([]string, error) { return models, nil } -func (c *Client) SendStream(msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) (err error) { +func (c *Client) SendStream(_ context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) (err error) { url := fmt.Sprintf("%s/chat/completions", c.ApiUrl.Value) payload := map[string]any{ diff --git a/internal/plugins/ai/lmstudio/lmstudio_test.go b/internal/plugins/ai/lmstudio/lmstudio_test.go index 2bb1d96b..e86288c3 100644 --- a/internal/plugins/ai/lmstudio/lmstudio_test.go +++ b/internal/plugins/ai/lmstudio/lmstudio_test.go @@ -27,7 +27,7 @@ func TestListModelsUsesBearerTokenWhenConfigured(t *testing.T) { client.ApiKey.Value = "secret" client.HttpClient = server.Client() - models, err := client.ListModels() + models, err := client.ListModels(context.Background()) require.NoError(t, err) require.Equal(t, []string{"model-1"}, models) } @@ -86,7 +86,7 @@ func TestListModelsDoesNotSendBearerForWhitespaceOnlyKey(t *testing.T) { client.ApiKey.Value = " " client.HttpClient = server.Client() - models, err := client.ListModels() + models, err := client.ListModels(context.Background()) require.NoError(t, err) require.Equal(t, []string{"model-1"}, models) } diff --git a/internal/plugins/ai/ollama/ollama.go b/internal/plugins/ai/ollama/ollama.go index ec3cd9f1..20ef92e7 100644 --- a/internal/plugins/ai/ollama/ollama.go +++ b/internal/plugins/ai/ollama/ollama.go @@ -90,7 +90,7 @@ func (o *Client) configure() (err error) { return } -func (o *Client) ListModels() (ret []string, err error) { +func (o *Client) ListModels(_ context.Context) (ret []string, err error) { ctx := context.Background() var listResp *ollamaapi.ListResponse @@ -104,7 +104,7 @@ func (o *Client) ListModels() (ret []string, err error) { return } -func (o *Client) SendStream(msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) (err error) { +func (o *Client) SendStream(_ context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) (err error) { ctx := context.Background() var req ollamaapi.ChatRequest diff --git a/internal/plugins/ai/openai/chat_completions.go b/internal/plugins/ai/openai/chat_completions.go index 3503a91b..1d32e9ed 100644 --- a/internal/plugins/ai/openai/chat_completions.go +++ b/internal/plugins/ai/openai/chat_completions.go @@ -30,7 +30,7 @@ func (o *Client) sendChatCompletions(ctx context.Context, msgs []*chat.ChatCompl // sendStreamChatCompletions sends a streaming request using the Chat Completions API func (o *Client) sendStreamChatCompletions( - msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate, + ctx context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate, ) (err error) { defer close(channel) @@ -39,7 +39,7 @@ func (o *Client) sendStreamChatCompletions( req.StreamOptions = openai.ChatCompletionStreamOptionsParam{ IncludeUsage: openai.Bool(true), } - stream := o.ApiClient.Chat.Completions.NewStreaming(context.Background(), req) + stream := o.ApiClient.Chat.Completions.NewStreaming(ctx, req) for stream.Next() { chunk := stream.Current() if len(chunk.Choices) > 0 && chunk.Choices[0].Delta.Content != "" { diff --git a/internal/plugins/ai/openai/openai.go b/internal/plugins/ai/openai/openai.go index fa703331..db147d91 100644 --- a/internal/plugins/ai/openai/openai.go +++ b/internal/plugins/ai/openai/openai.go @@ -96,9 +96,9 @@ func (o *Client) configure() (ret error) { return } -func (o *Client) ListModels() (ret []string, err error) { +func (o *Client) ListModels(ctx context.Context) (ret []string, err error) { var page *pagination.Page[openai.Model] - if page, err = o.ApiClient.Models.List(context.Background()); err == nil { + if page, err = o.ApiClient.Models.List(ctx); err == nil { for _, mod := range page.Data { ret = append(ret, mod.ID) } @@ -110,26 +110,26 @@ func (o *Client) ListModels() (ret []string, err error) { // Some providers (e.g., GitHub Models) return non-standard response formats // that the SDK fails to parse. debuglog.Debug(debuglog.Basic, "SDK Models.List failed for %s: %v, falling back to direct API fetch\n", o.GetName(), err) - return FetchModelsDirectly(context.Background(), o.ApiBaseURL.Value, o.ApiKey.Value, o.GetName(), o.httpClient) + return FetchModelsDirectly(ctx, o.ApiBaseURL.Value, o.ApiKey.Value, o.GetName(), o.httpClient) } func (o *Client) SendStream( - msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate, + ctx context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate, ) (err error) { // Use Responses API for OpenAI, Chat Completions API for other providers if o.supportsResponsesAPI() { - return o.sendStreamResponses(msgs, opts, channel) + return o.sendStreamResponses(ctx, msgs, opts, channel) } - return o.sendStreamChatCompletions(msgs, opts, channel) + return o.sendStreamChatCompletions(ctx, msgs, opts, channel) } func (o *Client) sendStreamResponses( - msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate, + ctx context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate, ) (err error) { defer close(channel) req := o.buildResponseParams(msgs, opts) - stream := o.ApiClient.Responses.NewStreaming(context.Background(), req) + stream := o.ApiClient.Responses.NewStreaming(ctx, req) for stream.Next() { event := stream.Current() switch event.Type { diff --git a/internal/plugins/ai/openai_compatible/providers_config.go b/internal/plugins/ai/openai_compatible/providers_config.go index 6dd14176..17601de8 100644 --- a/internal/plugins/ai/openai_compatible/providers_config.go +++ b/internal/plugins/ai/openai_compatible/providers_config.go @@ -44,7 +44,7 @@ func NewClient(providerConfig ProviderConfig) *Client { } // ListModels overrides the default ListModels to handle different response formats -func (c *Client) ListModels() ([]string, error) { +func (c *Client) ListModels(ctx context.Context) ([]string, error) { // If a custom models URL is provided, handle it if c.modelsURL != "" { if c.modelsURL == "static:abacus" { @@ -65,13 +65,13 @@ func (c *Client) ListModels() ([]string, error) { } // First try the standard OpenAI SDK approach - models, err := c.Client.ListModels() + models, err := c.Client.ListModels(ctx) if err == nil && len(models) > 0 { // only return if OpenAI SDK returns models return models, nil } // Fall back to direct API fetch - return c.DirectlyGetModels(context.Background()) + return c.DirectlyGetModels(ctx) } func (c *Client) fetchAbacusModels() ([]string, error) { diff --git a/internal/plugins/ai/perplexity/perplexity.go b/internal/plugins/ai/perplexity/perplexity.go index 58261a30..46e01d1c 100644 --- a/internal/plugins/ai/perplexity/perplexity.go +++ b/internal/plugins/ai/perplexity/perplexity.go @@ -53,7 +53,7 @@ func (c *Client) Configure() error { return nil } -func (c *Client) ListModels() ([]string, error) { +func (c *Client) ListModels(_ context.Context) ([]string, error) { // Perplexity API does not have a ListModels endpoint. // We return a predefined list. return models, nil @@ -119,7 +119,7 @@ func (c *Client) Send(ctx context.Context, msgs []*chat.ChatCompletionMessage, o return content.String(), nil } -func (c *Client) SendStream(msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) error { +func (c *Client) SendStream(_ context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) error { if c.client == nil { if err := c.Configure(); err != nil { close(channel) // Ensure channel is closed on error diff --git a/internal/plugins/ai/vendor.go b/internal/plugins/ai/vendor.go index c3134618..baa2ca92 100644 --- a/internal/plugins/ai/vendor.go +++ b/internal/plugins/ai/vendor.go @@ -11,8 +11,8 @@ import ( type Vendor interface { plugins.Plugin - ListModels() ([]string, error) - SendStream([]*chat.ChatCompletionMessage, *domain.ChatOptions, chan domain.StreamUpdate) error + ListModels(context.Context) ([]string, error) + SendStream(context.Context, []*chat.ChatCompletionMessage, *domain.ChatOptions, chan domain.StreamUpdate) error Send(context.Context, []*chat.ChatCompletionMessage, *domain.ChatOptions) (string, error) NeedsRawMode(modelName string) bool } diff --git a/internal/plugins/ai/vendors.go b/internal/plugins/ai/vendors.go index 796b7723..c5d5f159 100644 --- a/internal/plugins/ai/vendors.go +++ b/internal/plugins/ai/vendors.go @@ -117,7 +117,7 @@ func (o *VendorsManager) fetchVendorModels( defer wg.Done() - models, err := vendor.ListModels() + models, err := vendor.ListModels(ctx) select { case <-ctx.Done(): // Context canceled, don't send the result diff --git a/internal/plugins/ai/vendors_test.go b/internal/plugins/ai/vendors_test.go index 4534712d..eba9f91c 100644 --- a/internal/plugins/ai/vendors_test.go +++ b/internal/plugins/ai/vendors_test.go @@ -13,14 +13,14 @@ type stubVendor struct { name string } -func (v *stubVendor) GetName() string { return v.name } -func (v *stubVendor) GetSetupDescription() string { return "" } -func (v *stubVendor) IsConfigured() bool { return true } -func (v *stubVendor) Configure() error { return nil } -func (v *stubVendor) Setup() error { return nil } -func (v *stubVendor) SetupFillEnvFileContent(*bytes.Buffer) {} -func (v *stubVendor) ListModels() ([]string, error) { return nil, nil } -func (v *stubVendor) SendStream([]*chat.ChatCompletionMessage, *domain.ChatOptions, chan domain.StreamUpdate) error { +func (v *stubVendor) GetName() string { return v.name } +func (v *stubVendor) GetSetupDescription() string { return "" } +func (v *stubVendor) IsConfigured() bool { return true } +func (v *stubVendor) Configure() error { return nil } +func (v *stubVendor) Setup() error { return nil } +func (v *stubVendor) SetupFillEnvFileContent(*bytes.Buffer) {} +func (v *stubVendor) ListModels(context.Context) ([]string, error) { return nil, nil } +func (v *stubVendor) SendStream(context.Context, []*chat.ChatCompletionMessage, *domain.ChatOptions, chan domain.StreamUpdate) error { return nil } func (v *stubVendor) Send(context.Context, []*chat.ChatCompletionMessage, *domain.ChatOptions) (string, error) { diff --git a/internal/plugins/ai/vertexai/vertexai.go b/internal/plugins/ai/vertexai/vertexai.go index 73f177e7..c9a04766 100644 --- a/internal/plugins/ai/vertexai/vertexai.go +++ b/internal/plugins/ai/vertexai/vertexai.go @@ -61,7 +61,7 @@ func (c *Client) configure() error { return nil } -func (c *Client) ListModels() ([]string, error) { +func (c *Client) ListModels(_ context.Context) ([]string, error) { ctx := context.Background() // Get ADC credentials for API authentication @@ -179,7 +179,7 @@ func (c *Client) sendClaude(ctx context.Context, msgs []*chat.ChatCompletionMess return strings.Join(textParts, ""), nil } -func (c *Client) SendStream(msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) error { +func (c *Client) SendStream(_ context.Context, msgs []*chat.ChatCompletionMessage, opts *domain.ChatOptions, channel chan domain.StreamUpdate) error { if isGeminiModel(opts.Model) { return c.sendStreamGemini(msgs, opts, channel) } diff --git a/internal/server/chat.go b/internal/server/chat.go index b830e2f7..51b09fa1 100755 --- a/internal/server/chat.go +++ b/internal/server/chat.go @@ -132,7 +132,7 @@ func (h *ChatHandler) HandleChat(c *gin.Context) { Quiet: true, } - _, err = chatter.Send(chatReq, opts) + _, err = chatter.Send(c.Request.Context(), chatReq, opts) if err != nil { log.Printf("Error from chatter.Send: %v", err) // Error already sent to streamChan via domain.StreamTypeError if occurred in Send loop