From 708dc964df2ed06d48b656754522691e8a37a8d7 Mon Sep 17 00:00:00 2001 From: rcourtman Date: Wed, 24 Dec 2025 10:52:11 +0000 Subject: [PATCH] feat: Add configurable AI request timeout for slow hardware. Related to #880 Adds RequestTimeoutSeconds to AI config (default 300s / 5 min). Users with low-power hardware running Ollama can increase this value in Settings to prevent timeouts on slower inference. --- internal/ai/providers/anthropic.go | 14 +++++--- internal/ai/providers/anthropic_oauth.go | 13 ++++--- internal/ai/providers/anthropic_oauth_test.go | 4 ++- internal/ai/providers/anthropic_test.go | 14 ++++---- internal/ai/providers/factory.go | 28 +++++++++------ internal/ai/providers/gemini.go | 9 +++-- internal/ai/providers/gemini_test.go | 34 +++++++++--------- internal/ai/providers/integration_test.go | 36 +++++++++---------- internal/ai/providers/ollama.go | 8 +++-- internal/ai/providers/ollama_test.go | 28 +++++++-------- internal/ai/providers/openai.go | 9 +++-- internal/ai/providers/openai_test.go | 22 ++++++------ internal/config/ai.go | 13 +++++++ 13 files changed, 137 insertions(+), 95 deletions(-) diff --git a/internal/ai/providers/anthropic.go b/internal/ai/providers/anthropic.go index 2ae85f06d..ef2e84352 100644 --- a/internal/ai/providers/anthropic.go +++ b/internal/ai/providers/anthropic.go @@ -29,23 +29,27 @@ type AnthropicClient struct { } // NewAnthropicClient creates a new Anthropic API client -func NewAnthropicClient(apiKey, model string) *AnthropicClient { - return NewAnthropicClientWithBaseURL(apiKey, model, anthropicAPIURL) +// timeout is optional - pass 0 to use the default 5 minute timeout +func NewAnthropicClient(apiKey, model string, timeout time.Duration) *AnthropicClient { + return NewAnthropicClientWithBaseURL(apiKey, model, anthropicAPIURL, timeout) } // NewAnthropicClientWithBaseURL creates a new Anthropic client using a custom messages endpoint. // This is useful for testing and for deployments that route requests through a proxy. -func NewAnthropicClientWithBaseURL(apiKey, model, baseURL string) *AnthropicClient { +// timeout is optional - pass 0 to use the default 5 minute timeout +func NewAnthropicClientWithBaseURL(apiKey, model, baseURL string, timeout time.Duration) *AnthropicClient { if baseURL == "" { baseURL = anthropicAPIURL } + if timeout <= 0 { + timeout = 300 * time.Second // Default 5 minutes + } return &AnthropicClient{ apiKey: apiKey, model: model, baseURL: baseURL, client: &http.Client{ - // 5 minutes - Opus and other large models can take a very long time - Timeout: 300 * time.Second, + Timeout: timeout, }, } } diff --git a/internal/ai/providers/anthropic_oauth.go b/internal/ai/providers/anthropic_oauth.go index e23b48223..4f17ae4a9 100644 --- a/internal/ai/providers/anthropic_oauth.go +++ b/internal/ai/providers/anthropic_oauth.go @@ -320,16 +320,21 @@ type AnthropicOAuthClient struct { } // NewAnthropicOAuthClient creates a new Anthropic client using OAuth tokens -func NewAnthropicOAuthClient(accessToken, refreshToken string, expiresAt time.Time, model string) *AnthropicOAuthClient { - return NewAnthropicOAuthClientWithBaseURL(accessToken, refreshToken, expiresAt, model, "https://api.anthropic.com/v1/messages?beta=true") +// timeout is optional - pass 0 to use the default 5 minute timeout +func NewAnthropicOAuthClient(accessToken, refreshToken string, expiresAt time.Time, model string, timeout time.Duration) *AnthropicOAuthClient { + return NewAnthropicOAuthClientWithBaseURL(accessToken, refreshToken, expiresAt, model, "https://api.anthropic.com/v1/messages?beta=true", timeout) } // NewAnthropicOAuthClientWithBaseURL creates a new Anthropic OAuth client using a custom messages endpoint. // This is useful for testing and for deployments that route requests through a proxy. -func NewAnthropicOAuthClientWithBaseURL(accessToken, refreshToken string, expiresAt time.Time, model, baseURL string) *AnthropicOAuthClient { +// timeout is optional - pass 0 to use the default 5 minute timeout +func NewAnthropicOAuthClientWithBaseURL(accessToken, refreshToken string, expiresAt time.Time, model, baseURL string, timeout time.Duration) *AnthropicOAuthClient { if baseURL == "" { baseURL = "https://api.anthropic.com/v1/messages?beta=true" } + if timeout <= 0 { + timeout = 300 * time.Second // Default 5 minutes + } return &AnthropicOAuthClient{ accessToken: accessToken, refreshToken: refreshToken, @@ -337,7 +342,7 @@ func NewAnthropicOAuthClientWithBaseURL(accessToken, refreshToken string, expire model: model, baseURL: baseURL, client: &http.Client{ - Timeout: 300 * time.Second, + Timeout: timeout, }, } } diff --git a/internal/ai/providers/anthropic_oauth_test.go b/internal/ai/providers/anthropic_oauth_test.go index cd426ce23..46c97c1d2 100644 --- a/internal/ai/providers/anthropic_oauth_test.go +++ b/internal/ai/providers/anthropic_oauth_test.go @@ -265,7 +265,7 @@ func TestAnthropicOAuthClient_forceRefreshToken_UpdatesAndCallsCallback(t *testi oauthTokenURL = server.URL oauthHTTPClient = server.Client() - client := NewAnthropicOAuthClient("access_old", "refresh_old", time.Now().Add(-time.Minute), "claude-3") + client := NewAnthropicOAuthClient("access_old", "refresh_old", time.Now().Add(-time.Minute), "claude-3", 0) var cbTokens *OAuthTokens client.SetTokenRefreshCallback(func(tokens *OAuthTokens) { cbTokens = tokens }) @@ -361,6 +361,7 @@ func TestAnthropicOAuthClient_Chat_RefreshesOn401AndRetriesImmediately(t *testin time.Now().Add(10*time.Minute), // valid token, so refresh is driven by 401 "claude-3", server.URL+"/v1/messages?beta=true", + 0, ) client.client = server.Client() @@ -401,6 +402,7 @@ func TestAnthropicOAuthClient_ListModels_UsesConfiguredHost(t *testing.T) { time.Now().Add(10*time.Minute), "claude-3", server.URL+"/v1/messages?beta=true", + 0, ) client.client = server.Client() diff --git a/internal/ai/providers/anthropic_test.go b/internal/ai/providers/anthropic_test.go index 94e63361e..baf0a98c9 100644 --- a/internal/ai/providers/anthropic_test.go +++ b/internal/ai/providers/anthropic_test.go @@ -12,14 +12,14 @@ import ( // Anthropic tests focus on request/response correctness and endpoint behavior. func TestAnthropicClient_Name(t *testing.T) { - client := NewAnthropicClient("test-key", "claude-3-5-sonnet") + client := NewAnthropicClient("test-key", "claude-3-5-sonnet", 0) if client.Name() != "anthropic" { t.Errorf("Expected 'anthropic', got '%s'", client.Name()) } } func TestAnthropicClient_Chat_ContextCanceled(t *testing.T) { - client := NewAnthropicClient("test-key", "claude-3-5-sonnet") + client := NewAnthropicClient("test-key", "claude-3-5-sonnet", 0) ctx, cancel := context.WithCancel(context.Background()) cancel() // Cancel immediately @@ -34,7 +34,7 @@ func TestAnthropicClient_Chat_ContextCanceled(t *testing.T) { } func TestAnthropicClient_Chat_Timeout(t *testing.T) { - client := NewAnthropicClient("test-key", "claude-3-5-sonnet") + client := NewAnthropicClient("test-key", "claude-3-5-sonnet", 0) ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond) defer cancel() @@ -108,7 +108,7 @@ func TestAnthropicClient_Chat_Success_TextAndToolCalls(t *testing.T) { })) defer server.Close() - client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages") + client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages", 0) out, err := client.Chat(context.Background(), ChatRequest{ System: "You are helpful", Messages: []Message{{Role: "user", Content: "Hello"}}, @@ -162,7 +162,7 @@ func TestAnthropicClient_Chat_ToolResultInRequest(t *testing.T) { })) defer server.Close() - client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages") + client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages", 0) _, err := client.Chat(context.Background(), ChatRequest{ Messages: []Message{ { @@ -227,7 +227,7 @@ func TestAnthropicClient_ListModels_UsesConfiguredHost(t *testing.T) { })) defer server.Close() - client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages") + client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages", 0) models, err := client.ListModels(context.Background()) if err != nil { t.Fatalf("ListModels: %v", err) @@ -248,7 +248,7 @@ func TestAnthropicClient_TestConnection_CallsListModels(t *testing.T) { })) defer server.Close() - client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages") + client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages", 0) if err := client.TestConnection(context.Background()); err != nil { t.Fatalf("TestConnection: %v", err) } diff --git a/internal/ai/providers/factory.go b/internal/ai/providers/factory.go index 9bca0a629..c1c068b2e 100644 --- a/internal/ai/providers/factory.go +++ b/internal/ai/providers/factory.go @@ -23,12 +23,15 @@ func NewFromConfig(cfg *config.AIConfig) (Provider, error) { return providerClient, nil } + // Get the configured timeout + timeout := cfg.GetRequestTimeout() + // Fall back to legacy single-provider format switch cfg.Provider { case config.AIProviderAnthropic: // If we have an API key (from direct entry or OAuth-created), use regular client if cfg.APIKey != "" { - return NewAnthropicClient(cfg.APIKey, cfg.GetModel()), nil + return NewAnthropicClient(cfg.APIKey, cfg.GetModel(), timeout), nil } // Pro/Max users without org:create_api_key will use OAuth tokens directly if cfg.IsUsingOAuth() && cfg.OAuthAccessToken != "" { @@ -37,6 +40,7 @@ func NewFromConfig(cfg *config.AIConfig) (Provider, error) { cfg.OAuthRefreshToken, cfg.OAuthExpiresAt, cfg.GetModel(), + timeout, ) return client, nil } @@ -46,23 +50,23 @@ func NewFromConfig(cfg *config.AIConfig) (Provider, error) { if cfg.APIKey == "" { return nil, fmt.Errorf("OpenAI API key is required") } - return NewOpenAIClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL()), nil + return NewOpenAIClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL(), timeout), nil case config.AIProviderOllama: - return NewOllamaClient(cfg.GetModel(), cfg.GetBaseURL()), nil + return NewOllamaClient(cfg.GetModel(), cfg.GetBaseURL(), timeout), nil case config.AIProviderDeepSeek: if cfg.APIKey == "" { return nil, fmt.Errorf("DeepSeek API key is required") } // DeepSeek uses OpenAI-compatible API - return NewOpenAIClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL()), nil + return NewOpenAIClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL(), timeout), nil case config.AIProviderGemini: if cfg.APIKey == "" { return nil, fmt.Errorf("Gemini API key is required") } - return NewGeminiClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL()), nil + return NewGeminiClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL(), timeout), nil default: return nil, fmt.Errorf("unknown provider: %s", cfg.Provider) @@ -75,6 +79,9 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err return nil, fmt.Errorf("AI config is nil") } + // Get the configured timeout + timeout := cfg.GetRequestTimeout() + switch provider { case config.AIProviderAnthropic: // Check for OAuth first @@ -84,6 +91,7 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err cfg.OAuthRefreshToken, cfg.OAuthExpiresAt, model, + timeout, ), nil } // Then check for per-provider API key @@ -91,7 +99,7 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err if apiKey == "" { return nil, fmt.Errorf("Anthropic API key not configured") } - return NewAnthropicClient(apiKey, model), nil + return NewAnthropicClient(apiKey, model, timeout), nil case config.AIProviderOpenAI: apiKey := cfg.GetAPIKeyForProvider(config.AIProviderOpenAI) @@ -99,7 +107,7 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err return nil, fmt.Errorf("OpenAI API key not configured") } baseURL := cfg.GetBaseURLForProvider(config.AIProviderOpenAI) - return NewOpenAIClient(apiKey, model, baseURL), nil + return NewOpenAIClient(apiKey, model, baseURL, timeout), nil case config.AIProviderDeepSeek: apiKey := cfg.GetAPIKeyForProvider(config.AIProviderDeepSeek) @@ -107,11 +115,11 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err return nil, fmt.Errorf("DeepSeek API key not configured") } baseURL := cfg.GetBaseURLForProvider(config.AIProviderDeepSeek) - return NewOpenAIClient(apiKey, model, baseURL), nil + return NewOpenAIClient(apiKey, model, baseURL, timeout), nil case config.AIProviderOllama: baseURL := cfg.GetBaseURLForProvider(config.AIProviderOllama) - return NewOllamaClient(model, baseURL), nil + return NewOllamaClient(model, baseURL, timeout), nil case config.AIProviderGemini: apiKey := cfg.GetAPIKeyForProvider(config.AIProviderGemini) @@ -119,7 +127,7 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err return nil, fmt.Errorf("Gemini API key not configured") } baseURL := cfg.GetBaseURLForProvider(config.AIProviderGemini) - return NewGeminiClient(apiKey, model, baseURL), nil + return NewGeminiClient(apiKey, model, baseURL, timeout), nil default: return nil, fmt.Errorf("unknown provider: %s", provider) diff --git a/internal/ai/providers/gemini.go b/internal/ai/providers/gemini.go index 12cd99920..2e2ebca6c 100644 --- a/internal/ai/providers/gemini.go +++ b/internal/ai/providers/gemini.go @@ -28,7 +28,8 @@ type GeminiClient struct { } // NewGeminiClient creates a new Gemini API client -func NewGeminiClient(apiKey, model, baseURL string) *GeminiClient { +// timeout is optional - pass 0 to use the default 5 minute timeout +func NewGeminiClient(apiKey, model, baseURL string, timeout time.Duration) *GeminiClient { if baseURL == "" { baseURL = geminiAPIURL } @@ -36,13 +37,15 @@ func NewGeminiClient(apiKey, model, baseURL string) *GeminiClient { if strings.HasPrefix(model, "gemini:") { model = strings.TrimPrefix(model, "gemini:") } + if timeout <= 0 { + timeout = 300 * time.Second // Default 5 minutes + } return &GeminiClient{ apiKey: apiKey, model: model, baseURL: baseURL, client: &http.Client{ - // 5 minutes timeout - large models can take a long time - Timeout: 300 * time.Second, + Timeout: timeout, }, } } diff --git a/internal/ai/providers/gemini_test.go b/internal/ai/providers/gemini_test.go index 4ec84c93f..d2696542a 100644 --- a/internal/ai/providers/gemini_test.go +++ b/internal/ai/providers/gemini_test.go @@ -13,7 +13,7 @@ import ( func TestNewGeminiClient(t *testing.T) { t.Parallel() - client := NewGeminiClient("test-api-key", "gemini-pro", "") + client := NewGeminiClient("test-api-key", "gemini-pro", "", 0) if client == nil { t.Fatal("expected non-nil client") } @@ -31,7 +31,7 @@ func TestNewGeminiClient(t *testing.T) { func TestNewGeminiClient_StripPrefix(t *testing.T) { t.Parallel() - client := NewGeminiClient("api-key", "gemini:gemini-1.5-pro", "") + client := NewGeminiClient("api-key", "gemini:gemini-1.5-pro", "", 0) if client.model != "gemini-1.5-pro" { t.Errorf("expected model with prefix stripped, got %q", client.model) } @@ -40,7 +40,7 @@ func TestNewGeminiClient_StripPrefix(t *testing.T) { func TestNewGeminiClient_CustomBaseURL(t *testing.T) { t.Parallel() - client := NewGeminiClient("api-key", "gemini-pro", "https://custom.api.example.com") + client := NewGeminiClient("api-key", "gemini-pro", "https://custom.api.example.com", 0) if client.baseURL != "https://custom.api.example.com" { t.Errorf("expected custom baseURL, got %q", client.baseURL) } @@ -49,7 +49,7 @@ func TestNewGeminiClient_CustomBaseURL(t *testing.T) { func TestGeminiClient_Name(t *testing.T) { t.Parallel() - client := NewGeminiClient("key", "model", "") + client := NewGeminiClient("key", "model", "", 0) if client.Name() != "gemini" { t.Errorf("expected Name() to return 'gemini', got %q", client.Name()) } @@ -87,7 +87,7 @@ func TestGeminiClient_Chat_Success(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() resp, err := client.Chat(ctx, ChatRequest{ @@ -135,7 +135,7 @@ func TestGeminiClient_Chat_WithSystemPrompt(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -178,7 +178,7 @@ func TestGeminiClient_Chat_WithMaxTokens(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -221,7 +221,7 @@ func TestGeminiClient_Chat_ToolCalls(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() resp, err := client.Chat(ctx, ChatRequest{ @@ -269,7 +269,7 @@ func TestGeminiClient_Chat_APIError(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("invalid-key", "gemini-pro", server.URL) + client := NewGeminiClient("invalid-key", "gemini-pro", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -293,7 +293,7 @@ func TestGeminiClient_Chat_NoCandidates(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -325,7 +325,7 @@ func TestGeminiClient_Chat_SafetyBlocked(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -351,7 +351,7 @@ func TestGeminiClient_Chat_PromptBlocked(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -413,7 +413,7 @@ func TestGeminiClient_ListModels_Success(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() models, err := client.ListModels(ctx) @@ -435,7 +435,7 @@ func TestGeminiClient_ListModels_Error(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() _, err := client.ListModels(ctx) @@ -454,7 +454,7 @@ func TestGeminiClient_TestConnection(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() err := client.TestConnection(ctx) @@ -465,7 +465,7 @@ func TestGeminiClient_TestConnection(t *testing.T) { } func TestGeminiClient_Chat_NetworkError(t *testing.T) { - client := NewGeminiClient("test-key", "gemini-pro", "http://localhost:99999") + client := NewGeminiClient("test-key", "gemini-pro", "http://localhost:99999", 0) ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() @@ -497,7 +497,7 @@ func TestGeminiClient_Chat_RoleConversion(t *testing.T) { })) defer server.Close() - client := NewGeminiClient("test-key", "gemini-pro", server.URL) + client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ diff --git a/internal/ai/providers/integration_test.go b/internal/ai/providers/integration_test.go index e5f1be8b2..667082fb1 100644 --- a/internal/ai/providers/integration_test.go +++ b/internal/ai/providers/integration_test.go @@ -27,7 +27,7 @@ func getOllamaURL() string { } func TestIntegration_Ollama_TestConnection(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() @@ -40,7 +40,7 @@ func TestIntegration_Ollama_TestConnection(t *testing.T) { } func TestIntegration_Ollama_ListModels(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() @@ -61,7 +61,7 @@ func TestIntegration_Ollama_ListModels(t *testing.T) { } func TestIntegration_Ollama_SimpleChat(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() @@ -86,7 +86,7 @@ func TestIntegration_Ollama_SimpleChat(t *testing.T) { } func TestIntegration_Ollama_SystemPrompt(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() @@ -116,7 +116,7 @@ func TestIntegration_Ollama_SystemPrompt(t *testing.T) { } func TestIntegration_Ollama_MultiTurnConversation(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) defer cancel() @@ -155,7 +155,7 @@ func TestIntegration_Ollama_MultiTurnConversation(t *testing.T) { } func TestIntegration_Ollama_TokenCounting(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() @@ -183,7 +183,7 @@ func TestIntegration_Ollama_TokenCounting(t *testing.T) { } func TestIntegration_Ollama_ErrorHandling_BadModel(t *testing.T) { - client := providers.NewOllamaClient("nonexistent-model-12345", getOllamaURL()) + client := providers.NewOllamaClient("nonexistent-model-12345", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() @@ -202,7 +202,7 @@ func TestIntegration_Ollama_ErrorHandling_BadModel(t *testing.T) { } func TestIntegration_Ollama_Timeout(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) // Very short timeout - should fail ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond) @@ -225,7 +225,7 @@ func TestIntegration_Ollama_Timeout(t *testing.T) { // --- More useful tests below --- func TestIntegration_Ollama_JSONOutput(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() @@ -267,7 +267,7 @@ func TestIntegration_Ollama_JSONOutput(t *testing.T) { } func TestIntegration_Ollama_LongResponse(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) defer cancel() @@ -292,7 +292,7 @@ func TestIntegration_Ollama_LongResponse(t *testing.T) { } func TestIntegration_Ollama_EmptyMessage(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() @@ -313,7 +313,7 @@ func TestIntegration_Ollama_EmptyMessage(t *testing.T) { } func TestIntegration_Ollama_SpecialCharacters(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() @@ -336,7 +336,7 @@ func TestIntegration_Ollama_SpecialCharacters(t *testing.T) { } func TestIntegration_Ollama_ConcurrentRequests(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) const numRequests = 3 results := make(chan error, numRequests) @@ -372,7 +372,7 @@ func TestIntegration_Ollama_ConcurrentRequests(t *testing.T) { func TestIntegration_Ollama_InfrastructureAnalysis(t *testing.T) { // This simulates what Pulse actually does - send infrastructure context - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) defer cancel() @@ -416,7 +416,7 @@ What should I investigate first?` } func TestIntegration_Ollama_ModelName_Preserved(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() @@ -441,7 +441,7 @@ func TestIntegration_Ollama_ModelName_Preserved(t *testing.T) { } func TestIntegration_Ollama_StopReason(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() @@ -461,7 +461,7 @@ func TestIntegration_Ollama_StopReason(t *testing.T) { } func TestIntegration_Ollama_VeryLongInput(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) defer cancel() @@ -486,7 +486,7 @@ func TestIntegration_Ollama_VeryLongInput(t *testing.T) { } func TestIntegration_Ollama_RapidFireRequests(t *testing.T) { - client := providers.NewOllamaClient("tinyllama", getOllamaURL()) + client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0)) // Send 5 requests in rapid succession for i := 0; i < 5; i++ { diff --git a/internal/ai/providers/ollama.go b/internal/ai/providers/ollama.go index 837c45c2d..d2f958d51 100644 --- a/internal/ai/providers/ollama.go +++ b/internal/ai/providers/ollama.go @@ -19,7 +19,8 @@ type OllamaClient struct { } // NewOllamaClient creates a new Ollama API client -func NewOllamaClient(model, baseURL string) *OllamaClient { +// timeout is optional - pass 0 to use the default 5 minute timeout +func NewOllamaClient(model, baseURL string, timeout time.Duration) *OllamaClient { if baseURL == "" { baseURL = "http://localhost:11434" } @@ -28,11 +29,14 @@ func NewOllamaClient(model, baseURL string) *OllamaClient { baseURL = strings.TrimSuffix(baseURL, "/") baseURL = strings.TrimSuffix(baseURL, "/api") baseURL = strings.TrimSuffix(baseURL, "/") // In case it was /api/ + if timeout <= 0 { + timeout = 300 * time.Second // Default 5 minutes + } return &OllamaClient{ model: model, baseURL: baseURL, client: &http.Client{ - Timeout: 300 * time.Second, // Local models can be slow + Timeout: timeout, }, } } diff --git a/internal/ai/providers/ollama_test.go b/internal/ai/providers/ollama_test.go index 2e87108a9..4313a007b 100644 --- a/internal/ai/providers/ollama_test.go +++ b/internal/ai/providers/ollama_test.go @@ -47,7 +47,7 @@ func TestOllamaClient_Chat_Success(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("llama2", server.URL) + client := NewOllamaClient("llama2", server.URL, 0) ctx := context.Background() resp, err := client.Chat(ctx, ChatRequest{ @@ -92,7 +92,7 @@ func TestOllamaClient_Chat_WithSystemPrompt(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("llama2", server.URL) + client := NewOllamaClient("llama2", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -133,7 +133,7 @@ func TestOllamaClient_Chat_WithOptions(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("llama2", server.URL) + client := NewOllamaClient("llama2", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -166,7 +166,7 @@ func TestOllamaClient_Chat_APIError(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("nonexistent", server.URL) + client := NewOllamaClient("nonexistent", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -179,7 +179,7 @@ func TestOllamaClient_Chat_APIError(t *testing.T) { } func TestOllamaClient_Chat_NetworkError(t *testing.T) { - client := NewOllamaClient("llama2", "http://localhost:99999") + client := NewOllamaClient("llama2", "http://localhost:99999", 0) ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() @@ -211,7 +211,7 @@ func TestOllamaClient_Chat_ModelFallback(t *testing.T) { defer server.Close() // Client with no default model - client := NewOllamaClient("", server.URL) + client := NewOllamaClient("", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -246,7 +246,7 @@ func TestOllamaClient_Chat_StripModelPrefix(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("default", server.URL) + client := NewOllamaClient("default", server.URL, 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -274,7 +274,7 @@ func TestOllamaClient_TestConnection_Success(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("llama2", server.URL) + client := NewOllamaClient("llama2", server.URL, 0) ctx := context.Background() err := client.TestConnection(ctx) @@ -290,7 +290,7 @@ func TestOllamaClient_TestConnection_Failure(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("llama2", server.URL) + client := NewOllamaClient("llama2", server.URL, 0) ctx := context.Background() err := client.TestConnection(ctx) @@ -327,7 +327,7 @@ func TestOllamaClient_ListModels_Success(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("llama2", server.URL) + client := NewOllamaClient("llama2", server.URL, 0) ctx := context.Background() models, err := client.ListModels(ctx) @@ -352,7 +352,7 @@ func TestOllamaClient_ListModels_Failure(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("llama2", server.URL) + client := NewOllamaClient("llama2", server.URL, 0) ctx := context.Background() _, err := client.ListModels(ctx) @@ -376,7 +376,7 @@ func TestNewOllamaClient_NormalizesBaseURL(t *testing.T) { for _, tc := range tests { t.Run(tc.in, func(t *testing.T) { - client := NewOllamaClient("llama3", tc.in) + client := NewOllamaClient("llama3", tc.in, 0) if client.baseURL != tc.expected { t.Fatalf("baseURL = %q, want %q", client.baseURL, tc.expected) } @@ -408,7 +408,7 @@ func TestOllamaClient_Chat_ToolCallsResponse(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("llama3", server.URL) + client := NewOllamaClient("llama3", server.URL, 0) out, err := client.Chat(context.Background(), ChatRequest{ Messages: []Message{{Role: "user", Content: "What time is it?"}}, }) @@ -442,7 +442,7 @@ func TestOllamaClient_Chat_ToolCallsAndToolResultsInRequest(t *testing.T) { })) defer server.Close() - client := NewOllamaClient("llama3", server.URL) + client := NewOllamaClient("llama3", server.URL, 0) _, err := client.Chat(context.Background(), ChatRequest{ System: "system prompt", Messages: []Message{ diff --git a/internal/ai/providers/openai.go b/internal/ai/providers/openai.go index 5cbb8b255..3a841ac75 100644 --- a/internal/ai/providers/openai.go +++ b/internal/ai/providers/openai.go @@ -30,7 +30,8 @@ type OpenAIClient struct { } // NewOpenAIClient creates a new OpenAI API client -func NewOpenAIClient(apiKey, model, baseURL string) *OpenAIClient { +// timeout is optional - pass 0 to use the default 5 minute timeout +func NewOpenAIClient(apiKey, model, baseURL string, timeout time.Duration) *OpenAIClient { if baseURL == "" { baseURL = openaiAPIURL } @@ -40,13 +41,15 @@ func NewOpenAIClient(apiKey, model, baseURL string) *OpenAIClient { } else if strings.HasPrefix(model, "deepseek:") { model = strings.TrimPrefix(model, "deepseek:") } + if timeout <= 0 { + timeout = 300 * time.Second // Default 5 minutes + } return &OpenAIClient{ apiKey: apiKey, model: model, baseURL: baseURL, client: &http.Client{ - // 5 minutes timeout - DeepSeek reasoning models can take a long time - Timeout: 300 * time.Second, + Timeout: timeout, }, } } diff --git a/internal/ai/providers/openai_test.go b/internal/ai/providers/openai_test.go index 6c120b231..2f71c0d76 100644 --- a/internal/ai/providers/openai_test.go +++ b/internal/ai/providers/openai_test.go @@ -115,7 +115,7 @@ func TestOpenAIClient_Chat_Success(t *testing.T) { defer server.Close() // Create client with mock server URL - client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions") + client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions", 0) // Execute chat request ctx := context.Background() @@ -158,7 +158,7 @@ func TestOpenAIClient_Chat_APIError(t *testing.T) { })) defer server.Close() - client := NewOpenAIClient("invalid-key", "gpt-4o", server.URL+"/v1/chat/completions") + client := NewOpenAIClient("invalid-key", "gpt-4o", server.URL+"/v1/chat/completions", 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -172,7 +172,7 @@ func TestOpenAIClient_Chat_APIError(t *testing.T) { func TestOpenAIClient_Chat_NetworkError(t *testing.T) { // Create client pointing to non-existent server - client := NewOpenAIClient("test-key", "gpt-4o", "http://localhost:99999") + client := NewOpenAIClient("test-key", "gpt-4o", "http://localhost:99999", 0) ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() @@ -193,7 +193,7 @@ func TestOpenAIClient_Chat_ContextCanceled(t *testing.T) { })) defer server.Close() - client := NewOpenAIClient("test-key", "gpt-4o", server.URL+"/v1/chat/completions") + client := NewOpenAIClient("test-key", "gpt-4o", server.URL+"/v1/chat/completions", 0) ctx, cancel := context.WithCancel(context.Background()) cancel() // Cancel immediately @@ -236,7 +236,7 @@ func TestOpenAIClient_Chat_WithSystemPrompt(t *testing.T) { })) defer server.Close() - client := NewOpenAIClient("test-key", "gpt-4o", server.URL+"/v1/chat/completions") + client := NewOpenAIClient("test-key", "gpt-4o", server.URL+"/v1/chat/completions", 0) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ @@ -284,7 +284,7 @@ func TestOpenAIClient_ListModels_UsesConfiguredHostAndFilters(t *testing.T) { })) defer server.Close() - client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions") + client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions", 0) models, err := client.ListModels(context.Background()) if err != nil { t.Fatalf("ListModels: %v", err) @@ -312,7 +312,7 @@ func TestOpenAIClient_TestConnection_CallsListModels(t *testing.T) { })) defer server.Close() - client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions") + client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions", 0) if err := client.TestConnection(context.Background()); err != nil { t.Fatalf("TestConnection: %v", err) } @@ -345,7 +345,7 @@ func TestOpenAIClient_Chat_UsesMaxCompletionTokensForOpenAI(t *testing.T) { })) defer server.Close() - client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions") + client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions", 0) _, err := client.Chat(context.Background(), ChatRequest{ Messages: []Message{{Role: "user", Content: "Hello"}}, MaxTokens: 123, @@ -383,7 +383,7 @@ func TestOpenAIClient_Chat_GPT52NonChat_UsesCompletionsEndpointAndPrompt(t *test t.Fatalf("parse server url: %v", err) } - client := NewOpenAIClient("test-api-key", "gpt-5.2-pro", "https://api.openai.com/v1/chat/completions") + client := NewOpenAIClient("test-api-key", "gpt-5.2-pro", "https://api.openai.com/v1/chat/completions", 0) client.client.Transport = rewriteToServerTransport{serverBase: serverURL, rt: http.DefaultTransport} _, err = client.Chat(context.Background(), ChatRequest{ @@ -426,7 +426,7 @@ func TestOpenAIClient_ListModels_DeepSeekUsesModelsEndpoint(t *testing.T) { t.Fatalf("parse server url: %v", err) } - client := NewOpenAIClient("test-api-key", "deepseek-chat", "https://api.deepseek.com/v1/chat/completions") + client := NewOpenAIClient("test-api-key", "deepseek-chat", "https://api.deepseek.com/v1/chat/completions", 0) client.client.Transport = rewriteToServerTransport{serverBase: serverURL, rt: http.DefaultTransport} _, err = client.ListModels(context.Background()) @@ -462,7 +462,7 @@ func TestOpenAIClient_Chat_O1OmitsTemperature(t *testing.T) { })) defer server.Close() - client := NewOpenAIClient("test-api-key", "o1-mini", server.URL+"/v1/chat/completions") + client := NewOpenAIClient("test-api-key", "o1-mini", server.URL+"/v1/chat/completions", 0) _, err := client.Chat(context.Background(), ChatRequest{ Messages: []Message{{Role: "user", Content: "Hello"}}, Model: "o1-mini", diff --git a/internal/config/ai.go b/internal/config/ai.go index 2210c6bf9..b1ac165c6 100644 --- a/internal/config/ai.go +++ b/internal/config/ai.go @@ -53,6 +53,10 @@ type AIConfig struct { // Alert-triggered AI analysis - analyze specific resources when alerts fire AlertTriggeredAnalysis bool `json:"alert_triggered_analysis"` // Enable AI analysis when alerts fire (token-efficient) + // Request timeout - how long to wait for AI responses (default: 300s / 5 min) + // Increase this for slow hardware running local models (e.g., Ollama on low-power devices) + RequestTimeoutSeconds int `json:"request_timeout_seconds,omitempty"` + // AI cost controls // Budget is expressed as an estimated USD amount over a 30-day window (pro-rated in UI for other ranges). CostBudgetUSD30d float64 `json:"cost_budget_usd_30d,omitempty"` @@ -454,3 +458,12 @@ func (c *AIConfig) IsPatrolEnabled() bool { func (c *AIConfig) IsAlertTriggeredAnalysisEnabled() bool { return c.AlertTriggeredAnalysis } + +// GetRequestTimeout returns the timeout duration for AI requests +// Default is 5 minutes (300 seconds) if not configured +func (c *AIConfig) GetRequestTimeout() time.Duration { + if c.RequestTimeoutSeconds > 0 { + return time.Duration(c.RequestTimeoutSeconds) * time.Second + } + return 300 * time.Second // 5 minutes default +}