Skip to content

feat: make MaxRetries configurable in OpenAI provider #271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mtojek merged 3 commits into from
Apr 17, 2026
Merged

feat: make MaxRetries configurable in OpenAI provider #271

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8ee1e8b
feat: make MaxRetries configurable in OpenAI provider (#115)
mtojek Apr 17, 2026
49bdf32
feat: read MaxRetries from OPENAI_MAX_RETRIES env var
mtojek Apr 17, 2026
0313016
feat: extend MaxRetries support to Anthropic and Copilot providers
mtojek Apr 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions config/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,10 @@ type Anthropic struct {
// with a access token. When set, the access token is used for upstream
// LLM requests instead of the API key.
BYOKBearerToken string
// MaxRetries controls the number of automatic retries the SDK will perform
// on transient errors. If nil, the SDK default (2) is used.
// Set to 0 to disable retries entirely.
MaxRetries *int
}

type AWSBedrock struct {
Expand All @@ -43,6 +47,10 @@ type OpenAI struct {
CircuitBreaker *CircuitBreaker
SendActorHeaders bool
ExtraHeaders map[string]string
// MaxRetries controls the number of automatic retries the SDK will perform
// on transient errors. If nil, the SDK default (2) is used.
// Set to 0 to disable retries entirely.
MaxRetries *int
}

type Copilot struct {
Expand All @@ -51,6 +59,10 @@ type Copilot struct {
BaseURL string
APIDumpDir string
CircuitBreaker *CircuitBreaker
// MaxRetries controls the number of automatic retries the SDK will perform
// on transient errors. If nil, the SDK default (2) is used.
// Set to 0 to disable retries entirely.
MaxRetries *int
}

// CircuitBreaker holds configuration for circuit breakers.
Expand Down
3 changes: 3 additions & 0 deletions intercept/chatcompletions/base.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,9 @@ type interceptionBase struct {

func (i *interceptionBase) newCompletionsService() openai.ChatCompletionService {
opts := []option.RequestOption{option.WithAPIKey(i.cfg.Key), option.WithBaseURL(i.cfg.BaseURL)}
if i.cfg.MaxRetries != nil {
opts = append(opts, option.WithMaxRetries(*i.cfg.MaxRetries))
}

// Add extra headers if configured.
// Some providers require additional headers that are not added by the SDK.
Expand Down
3 changes: 3 additions & 0 deletions intercept/messages/base.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,9 @@ func (i *interceptionBase) newMessagesService(ctx context.Context, opts ...optio
opts = append(opts, option.WithAPIKey(i.cfg.Key))
}
opts = append(opts, option.WithBaseURL(i.cfg.BaseURL))
if i.cfg.MaxRetries != nil {
opts = append(opts, option.WithMaxRetries(*i.cfg.MaxRetries))
}

// Add extra headers if configured.
// Some providers require additional headers that are not added by the SDK.
Expand Down
3 changes: 3 additions & 0 deletions intercept/responses/base.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ type responsesInterceptionBase struct {

func (i *responsesInterceptionBase) newResponsesService() responses.ResponseService {
opts := []option.RequestOption{option.WithBaseURL(i.cfg.BaseURL), option.WithAPIKey(i.cfg.Key)}
if i.cfg.MaxRetries != nil {
opts = append(opts, option.WithMaxRetries(*i.cfg.MaxRetries))
}

// Add extra headers if configured.
// Some providers require additional headers that are not added by the SDK.
Expand Down
18 changes: 10 additions & 8 deletions internal/integrationtest/responses_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -596,9 +596,6 @@ func TestResponsesParallelToolsOverwritten(t *testing.T) {
}
}

// TODO set MaxRetries to speed up this test
// option.WithMaxRetries(0), in base responses interceptor
// https://github.com/coder/aibridge/issues/115
func TestClientAndConnectionError(t *testing.T) {
t.Parallel()

Expand Down Expand Up @@ -642,7 +639,11 @@ func TestClientAndConnectionError(t *testing.T) {
t.Cleanup(cancel)

// tc.addr may be an intentionally invalid URL; use withCustomProvider.
bridgeServer := newBridgeTestServer(ctx, t, tc.addr, withCustomProvider(provider.NewOpenAI(openAICfg(tc.addr, apiKey))))
// MaxRetries is set to 0 to disable SDK retries and speed up the test.
cfg := openAICfg(tc.addr, apiKey)
maxRetries := 0
cfg.MaxRetries = &maxRetries
bridgeServer := newBridgeTestServer(ctx, t, tc.addr, withCustomProvider(provider.NewOpenAI(cfg)))

reqBytes := responsesRequestBytes(t, tc.streaming)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIResponses, reqBytes)
Expand All @@ -660,9 +661,6 @@ func TestClientAndConnectionError(t *testing.T) {
}
}

// TODO set MaxRetries to speed up this test
// option.WithMaxRetries(0), in base responses interceptor
// https://github.com/coder/aibridge/issues/115
func TestUpstreamError(t *testing.T) {
t.Parallel()

Expand Down Expand Up @@ -721,7 +719,11 @@ func TestUpstreamError(t *testing.T) {
}))
t.Cleanup(upstream.Close)

bridgeServer := newBridgeTestServer(ctx, t, upstream.URL)
// MaxRetries is set to 0 to disable SDK retries and speed up the test.
cfg := openAICfg(upstream.URL, apiKey)
maxRetries := 0
cfg.MaxRetries = &maxRetries
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, withCustomProvider(provider.NewOpenAI(cfg)))

reqBytes := responsesRequestBytes(t, tc.streaming)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIResponses, reqBytes)
Expand Down
8 changes: 8 additions & 0 deletions provider/anthropic.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"io"
"net/http"
"os"
"strconv"
"strings"

"github.com/google/uuid"
Expand Down Expand Up @@ -62,6 +63,13 @@ func NewAnthropic(cfg config.Anthropic, bedrockCfg *config.AWSBedrock) *Anthropi
if cfg.APIDumpDir == "" {
cfg.APIDumpDir = os.Getenv("BRIDGE_DUMP_DIR")
}
if cfg.MaxRetries == nil {
if v := os.Getenv("ANTHROPIC_MAX_RETRIES"); v != "" {
if n, err := strconv.Atoi(v); err == nil {
cfg.MaxRetries = &n
}
}
}
if cfg.CircuitBreaker != nil {
cfg.CircuitBreaker.IsFailure = anthropicIsFailure
cfg.CircuitBreaker.OpenErrorResponse = anthropicOpenErrorResponse
Expand Down
9 changes: 9 additions & 0 deletions provider/copilot.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"io"
"net/http"
"os"
"strconv"
"strings"

"github.com/google/uuid"
Expand Down Expand Up @@ -63,6 +64,13 @@ func NewCopilot(cfg config.Copilot) *Copilot {
if cfg.APIDumpDir == "" {
cfg.APIDumpDir = os.Getenv("BRIDGE_DUMP_DIR")
}
if cfg.MaxRetries == nil {
if v := os.Getenv("COPILOT_MAX_RETRIES"); v != "" {
Copy link
Copy Markdown
Contributor

@pawbana pawbana Apr 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a bit weird in Copilot case as we've recently added support for multiple providers and multiple Copilot providers (normal / business / enterprise) in coder/coder but looking at what configuration options are possible in aibridge I think it is ok.

if n, err := strconv.Atoi(v); err == nil {
cfg.MaxRetries = &n
}
}
}
if cfg.CircuitBreaker != nil {
cfg.CircuitBreaker.OpenErrorResponse = copilotOpenErrorResponse
}
Expand Down Expand Up @@ -145,6 +153,7 @@ func (p *Copilot) CreateInterceptor(_ http.ResponseWriter, r *http.Request, trac
APIDumpDir: p.cfg.APIDumpDir,
CircuitBreaker: p.cfg.CircuitBreaker,
ExtraHeaders: extractCopilotHeaders(r),
MaxRetries: p.cfg.MaxRetries,
}

cred := intercept.NewCredentialInfo(intercept.CredentialKindBYOK, key)
Expand Down
8 changes: 8 additions & 0 deletions provider/openai.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"io"
"net/http"
"os"
"strconv"
"strings"

"github.com/google/uuid"
Expand Down Expand Up @@ -51,6 +52,13 @@ func NewOpenAI(cfg config.OpenAI) *OpenAI {
if cfg.APIDumpDir == "" {
cfg.APIDumpDir = os.Getenv("BRIDGE_DUMP_DIR")
}
if cfg.MaxRetries == nil {
if v := os.Getenv("OPENAI_MAX_RETRIES"); v != "" {
if n, err := strconv.Atoi(v); err == nil {
cfg.MaxRetries = &n
}
}
}
if cfg.CircuitBreaker != nil {
cfg.CircuitBreaker.OpenErrorResponse = openAIOpenErrorResponse
}
Expand Down
Loading