ai/danielmiessler.fabric
1
0
Fork 0
mirror of https://github.com/danielmiessler/fabric.git synced 2026-09-10 07:36:44 -04:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor: propagate context.Context through Vendor interface methods

Browse source 
- 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`
This commit is contained in:
Kayvan Sylvan 2026-03-25 16:09:10 -07:00
parent 070c626b7b
commit 2e0abdd78a
40 changed files with 178 additions and 132 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
internal/cli/chat.go
View file

@ -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
}

6
internal/core/chatter.go
View file

@ -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 {

14
internal/core/chatter_test.go
View file

@ -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 "<think>hidden</think> 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)
}

16
internal/core/plugin_registry_test.go
View file

@ -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) {

6
internal/plugins/ai/anthropic/anthropic.go
View file

@ -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

5
internal/plugins/ai/anthropic/anthropic_test.go
View file

@ -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)

3
internal/plugins/ai/azure/azure.go
View file

@ -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
}

3
internal/plugins/ai/azure/azure_test.go
View file

@ -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)
}

3
internal/plugins/ai/azure_entra/azure_entra.go
View file

@ -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
}

3
internal/plugins/ai/azure_entra/azure_entra_test.go
View file

@ -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)
}

15
internal/plugins/ai/azureaigateway/azureaigateway.go
View file

@ -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)

12
internal/plugins/ai/azureaigateway/azureaigateway_test.go
View file

@ -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)
}

3
internal/plugins/ai/azureaigateway/backend_azure_openai.go
View file

@ -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",

3
internal/plugins/ai/azureaigateway/backend_bedrock.go
View file

@ -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",

3
internal/plugins/ai/azureaigateway/backend_vertex_ai.go
View file

@ -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",

4
internal/plugins/ai/bedrock/bedrock.go
View file

@ -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 {

6
internal/plugins/ai/bedrock/bedrock_test.go
View file

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

28
internal/plugins/ai/codex/codex.go
View file

@ -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 {

25
internal/plugins/ai/codex/codex_test.go
View file

@ -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

41
internal/plugins/ai/codex/errors.go
View file

@ -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

10
internal/plugins/ai/codex/oauth.go
View file

@ -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
}

4
internal/plugins/ai/codex/token.go
View file

@ -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])

4
internal/plugins/ai/copilot/copilot.go
View file

@ -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()

4
internal/plugins/ai/digitalocean/digitalocean.go
View file

@ -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
}

4
internal/plugins/ai/dryrun/dryrun.go
View file

@ -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{

5
internal/plugins/ai/dryrun/dryrun_test.go
View file

@ -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)
}

3
internal/plugins/ai/exolab/exolab.go
View file

@ -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
}

4
internal/plugins/ai/gemini/gemini.go
View file

@ -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)

4
internal/plugins/ai/lmstudio/lmstudio.go
View file

@ -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{

4
internal/plugins/ai/lmstudio/lmstudio_test.go
View file

@ -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)
}

4
internal/plugins/ai/ollama/ollama.go
View file

@ -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

4
internal/plugins/ai/openai/chat_completions.go
View file

@ -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 != "" {

16
internal/plugins/ai/openai/openai.go
View file

@ -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 {

6
internal/plugins/ai/openai_compatible/providers_config.go
View file

@ -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) {

4
internal/plugins/ai/perplexity/perplexity.go
View file

@ -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

4
internal/plugins/ai/vendor.go
View file

@ -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
}

2
internal/plugins/ai/vendors.go
View file

@ -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

16
internal/plugins/ai/vendors_test.go
View file

@ -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) {

4
internal/plugins/ai/vertexai/vertexai.go
View file

@ -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)
}

2
internal/server/chat.go
View file

@ -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