diff --git a/go/adk/pkg/agent/agent.go b/go/adk/pkg/agent/agent.go index 6d6bab4aa..ecc620b6f 100644 --- a/go/adk/pkg/agent/agent.go +++ b/go/adk/pkg/agent/agent.go @@ -305,9 +305,14 @@ func CreateLLM(ctx context.Context, m adk.Model, log logr.Logger) (adkmodel.LLM, if modelName == "" { return nil, fmt.Errorf("bedrock requires a model name (e.g. anthropic.claude-3-sonnet-20240229-v1:0)") } - // Use Bedrock Converse API for ALL models (including Anthropic) + // Use Bedrock Converse API for ALL models (including Anthropic). + // ReadTimeout maps to the overall HTTP client timeout (whole Converse + // request) and ConnectTimeout to the dialer, mirroring the Python ADK's + // botocore read/connect timeouts so the config is honored on both runtimes. + tc := transportConfigFromBase(m.BaseModel, m.ReadTimeout) + tc.ConnectTimeout = m.ConnectTimeout cfg := &models.BedrockConfig{ - TransportConfig: transportConfigFromBase(m.BaseModel, nil), + TransportConfig: tc, Model: modelName, Region: region, AdditionalModelRequestFields: m.AdditionalModelRequestFields, diff --git a/go/adk/pkg/agent/agent_test.go b/go/adk/pkg/agent/agent_test.go index 4f4c5fb45..fdaa64bf6 100644 --- a/go/adk/pkg/agent/agent_test.go +++ b/go/adk/pkg/agent/agent_test.go @@ -1,6 +1,7 @@ package agent import ( + "context" "encoding/json" "os" "path/filepath" @@ -286,6 +287,59 @@ func TestConfigDeserialization_Bedrock(t *testing.T) { } } +// TestCreateLLM_BedrockTimeouts verifies that the Bedrock read/connect timeouts +// from the model config are wired into the Go ADK's HTTP transport, so the same +// ModelConfig fields are honored by both the Python and Go runtimes. +func TestCreateLLM_BedrockTimeouts(t *testing.T) { + read := 1800 + connect := 10 + m := &adk.Bedrock{ + BaseModel: adk.BaseModel{Type: "bedrock", Model: "anthropic.claude-3-sonnet-20240229-v1:0"}, + Region: "us-east-1", + ReadTimeout: &read, + ConnectTimeout: &connect, + } + + llm, err := CreateLLM(context.Background(), m, logr.Discard()) + if err != nil { + t.Fatalf("CreateLLM: %v", err) + } + bm, ok := llm.(*models.BedrockModel) + if !ok { + t.Fatalf("model is %T, want *models.BedrockModel", llm) + } + if bm.Config.Timeout == nil || *bm.Config.Timeout != read { + t.Errorf("Config.Timeout = %v, want %d", bm.Config.Timeout, read) + } + if bm.Config.ConnectTimeout == nil || *bm.Config.ConnectTimeout != connect { + t.Errorf("Config.ConnectTimeout = %v, want %d", bm.Config.ConnectTimeout, connect) + } +} + +// TestCreateLLM_BedrockTimeoutsUnset verifies that omitting the timeouts leaves +// them nil so each runtime keeps its own default behavior. +func TestCreateLLM_BedrockTimeoutsUnset(t *testing.T) { + m := &adk.Bedrock{ + BaseModel: adk.BaseModel{Type: "bedrock", Model: "anthropic.claude-3-sonnet-20240229-v1:0"}, + Region: "us-east-1", + } + + llm, err := CreateLLM(context.Background(), m, logr.Discard()) + if err != nil { + t.Fatalf("CreateLLM: %v", err) + } + bm, ok := llm.(*models.BedrockModel) + if !ok { + t.Fatalf("model is %T, want *models.BedrockModel", llm) + } + if bm.Config.Timeout != nil { + t.Errorf("Config.Timeout = %v, want nil", *bm.Config.Timeout) + } + if bm.Config.ConnectTimeout != nil { + t.Errorf("Config.ConnectTimeout = %v, want nil", *bm.Config.ConnectTimeout) + } +} + func TestBuildAgentTools_WiresSkillsToolsFromEnv(t *testing.T) { skillsDir := t.TempDir() skillDir := filepath.Join(skillsDir, "csv-to-json") diff --git a/go/adk/pkg/models/base.go b/go/adk/pkg/models/base.go index 41294b42a..3e23a343e 100644 --- a/go/adk/pkg/models/base.go +++ b/go/adk/pkg/models/base.go @@ -2,6 +2,8 @@ package models import ( "encoding/json" + "fmt" + "net" "net/http" "strings" "time" @@ -19,7 +21,8 @@ type TransportConfig struct { TLSCACertPath *string TLSDisableSystemCAs *bool APIKeyPassthrough bool - Timeout *int // seconds; nil = defaultTimeout + Timeout *int // seconds; nil = defaultTimeout. Bounds the whole request (connect + write + read). + ConnectTimeout *int // seconds; nil = Go transport default. Bounds connection establishment only. } // BuildHTTPClient creates an http.Client with the full transport stack: @@ -35,6 +38,16 @@ func BuildHTTPClient(tc TransportConfig) (*http.Client, error) { return nil, err } + // Apply the connect timeout to the transport's dialer. Clone first so we + // never mutate the shared http.DefaultTransport that BuildTLSTransport + // returns when no TLS override is set. + if tc.ConnectTimeout != nil { + transport, err = withConnectTimeout(transport, time.Duration(*tc.ConnectTimeout)*time.Second) + if err != nil { + return nil, err + } + } + if len(tc.Headers) > 0 { transport = &headerTransport{base: transport, headers: tc.Headers} } @@ -47,6 +60,21 @@ func BuildHTTPClient(tc TransportConfig) (*http.Client, error) { return &http.Client{Timeout: timeout, Transport: transport}, nil } +// withConnectTimeout returns a copy of base (which must be an *http.Transport) +// whose DialContext enforces the given connection-establishment timeout. The +// base transport is cloned so callers can safely pass http.DefaultTransport +// without mutating the shared global. +func withConnectTimeout(base http.RoundTripper, connectTimeout time.Duration) (http.RoundTripper, error) { + t, ok := base.(*http.Transport) + if !ok { + return nil, fmt.Errorf("BuildHTTPClient: connect timeout requires an *http.Transport base, got %T", base) + } + cloned := t.Clone() + dialer := &net.Dialer{Timeout: connectTimeout, KeepAlive: 30 * time.Second} + cloned.DialContext = dialer.DialContext + return cloned, nil +} + // BearerTokenKey is the context key for storing the bearer token for API key passthrough var BearerTokenKey = &contextKey{} diff --git a/go/adk/pkg/models/tls_test.go b/go/adk/pkg/models/tls_test.go index 0189d419d..41998af56 100644 --- a/go/adk/pkg/models/tls_test.go +++ b/go/adk/pkg/models/tls_test.go @@ -2,6 +2,7 @@ package models import ( "encoding/pem" + "fmt" "net/http" "net/http/httptest" "os" @@ -134,6 +135,35 @@ func TestBuildHTTPClient_Timeout(t *testing.T) { } } +// Should set a connect timeout on a cloned transport without mutating the +// shared http.DefaultTransport. +func TestBuildHTTPClient_ConnectTimeout(t *testing.T) { + def, ok := http.DefaultTransport.(*http.Transport) + if !ok { + t.Fatalf("http.DefaultTransport is %T, want *http.Transport", http.DefaultTransport) + } + origDial := fmt.Sprintf("%p", def.DialContext) + + seconds := 7 + client, err := BuildHTTPClient(TransportConfig{ConnectTimeout: &seconds}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + tr, ok := client.Transport.(*http.Transport) + if !ok { + t.Fatalf("expected *http.Transport, got %T", client.Transport) + } + if tr == def { + t.Fatal("expected a cloned transport, got shared http.DefaultTransport") + } + if tr.DialContext == nil { + t.Fatal("expected DialContext to be set for connect timeout") + } + if newDial := fmt.Sprintf("%p", def.DialContext); newDial != origDial { + t.Error("http.DefaultTransport.DialContext was mutated by BuildHTTPClient") + } +} + // Should inject headers if specified func TestBuildHTTPClient_HeadersInjected(t *testing.T) { var got string diff --git a/go/api/adk/types.go b/go/api/adk/types.go index bcc587cf4..67bbd5413 100644 --- a/go/api/adk/types.go +++ b/go/api/adk/types.go @@ -263,6 +263,15 @@ type Bedrock struct { // "5m" (default) or "1h". See the v1alpha2.BedrockConfig CRD doc for the // cost/compatibility trade-offs of "1h". CacheTTL string `json:"cache_ttl,omitempty"` + // ReadTimeout is the Bedrock HTTP client read timeout in seconds. Nil keeps + // each runtime's default. Python ADK: overrides botocore's ~60s read timeout, + // which otherwise aborts long completions with a ReadTimeoutError. Go ADK: + // bounds the whole Converse request (overall HTTP client timeout, default 30m). + ReadTimeout *int `json:"read_timeout,omitempty"` + // ConnectTimeout is the Bedrock HTTP client connection-establishment timeout + // in seconds. Nil keeps each runtime's default (Python ADK: botocore; Go ADK: + // net dialer). Bounds connection setup only, not the response read. + ConnectTimeout *int `json:"connect_timeout,omitempty"` } func (b *Bedrock) MarshalJSON() ([]byte, error) { diff --git a/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml b/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml index c7a4163d6..d2e7d2a34 100644 --- a/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml +++ b/go/api/config/crd/bases/kagent.dev_modelconfigs.yaml @@ -503,6 +503,14 @@ spec: - 5m - 1h type: string + connectTimeout: + description: |- + ConnectTimeout is the Bedrock HTTP client connection-establishment timeout + in seconds, applied by both the Python and Go ADK runtimes. It bounds + connection setup only, not the response read. When unset, each runtime's + default is used (Python ADK: botocore; Go ADK: net dialer). + minimum: 1 + type: integer promptCaching: default: false description: |- @@ -521,6 +529,17 @@ spec: See https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html for the current list of supported models and minimum prefix sizes. type: boolean + readTimeout: + description: |- + ReadTimeout is the Bedrock HTTP client read timeout in seconds, applied by + both the Python and Go ADK runtimes. Raise this for agents that make long + Converse calls (large tool-augmented turns, extended reasoning). On the + Python ADK it overrides botocore's ~60s read timeout, which otherwise + aborts long completions with a ReadTimeoutError; on the Go ADK it bounds + the whole Converse request (default 30m). When unset, each runtime's + default is used. + minimum: 1 + type: integer region: description: AWS region where the Bedrock model is available (e.g., us-east-1, us-west-2) diff --git a/go/api/v1alpha2/modelconfig_types.go b/go/api/v1alpha2/modelconfig_types.go index b46d07c81..bc15acb1c 100644 --- a/go/api/v1alpha2/modelconfig_types.go +++ b/go/api/v1alpha2/modelconfig_types.go @@ -294,6 +294,25 @@ type BedrockConfig struct { // +kubebuilder:validation:Enum="5m";"1h" // +kubebuilder:default="5m" CacheTTL string `json:"cacheTTL,omitempty"` + + // ReadTimeout is the Bedrock HTTP client read timeout in seconds, applied by + // both the Python and Go ADK runtimes. Raise this for agents that make long + // Converse calls (large tool-augmented turns, extended reasoning). On the + // Python ADK it overrides botocore's ~60s read timeout, which otherwise + // aborts long completions with a ReadTimeoutError; on the Go ADK it bounds + // the whole Converse request (default 30m). When unset, each runtime's + // default is used. + // +optional + // +kubebuilder:validation:Minimum=1 + ReadTimeout *int `json:"readTimeout,omitempty"` + + // ConnectTimeout is the Bedrock HTTP client connection-establishment timeout + // in seconds, applied by both the Python and Go ADK runtimes. It bounds + // connection setup only, not the response read. When unset, each runtime's + // default is used (Python ADK: botocore; Go ADK: net dialer). + // +optional + // +kubebuilder:validation:Minimum=1 + ConnectTimeout *int `json:"connectTimeout,omitempty"` } // SAPAICoreConfig contains SAP AI Core-specific configuration options. diff --git a/go/api/v1alpha2/zz_generated.deepcopy.go b/go/api/v1alpha2/zz_generated.deepcopy.go index 98d89ec39..15a9de59e 100644 --- a/go/api/v1alpha2/zz_generated.deepcopy.go +++ b/go/api/v1alpha2/zz_generated.deepcopy.go @@ -687,6 +687,16 @@ func (in *BedrockConfig) DeepCopyInto(out *BedrockConfig) { *out = new(apiextensionsv1.JSON) (*in).DeepCopyInto(*out) } + if in.ReadTimeout != nil { + in, out := &in.ReadTimeout, &out.ReadTimeout + *out = new(int) + **out = **in + } + if in.ConnectTimeout != nil { + in, out := &in.ConnectTimeout, &out.ConnectTimeout + *out = new(int) + **out = **in + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new BedrockConfig. diff --git a/go/core/internal/controller/translator/agent/adk_api_translator.go b/go/core/internal/controller/translator/agent/adk_api_translator.go index 50b680756..c76228d31 100644 --- a/go/core/internal/controller/translator/agent/adk_api_translator.go +++ b/go/core/internal/controller/translator/agent/adk_api_translator.go @@ -831,6 +831,8 @@ func (a *adkApiTranslator) translateModel(ctx context.Context, namespace, modelC AdditionalModelRequestFields: additionalFields, PromptCaching: model.Spec.Bedrock.PromptCaching, CacheTTL: model.Spec.Bedrock.CacheTTL, + ReadTimeout: model.Spec.Bedrock.ReadTimeout, + ConnectTimeout: model.Spec.Bedrock.ConnectTimeout, } // Populate TLS fields in BaseModel diff --git a/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml b/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml index c7a4163d6..d2e7d2a34 100644 --- a/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml +++ b/helm/kagent-crds/templates/kagent.dev_modelconfigs.yaml @@ -503,6 +503,14 @@ spec: - 5m - 1h type: string + connectTimeout: + description: |- + ConnectTimeout is the Bedrock HTTP client connection-establishment timeout + in seconds, applied by both the Python and Go ADK runtimes. It bounds + connection setup only, not the response read. When unset, each runtime's + default is used (Python ADK: botocore; Go ADK: net dialer). + minimum: 1 + type: integer promptCaching: default: false description: |- @@ -521,6 +529,17 @@ spec: See https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html for the current list of supported models and minimum prefix sizes. type: boolean + readTimeout: + description: |- + ReadTimeout is the Bedrock HTTP client read timeout in seconds, applied by + both the Python and Go ADK runtimes. Raise this for agents that make long + Converse calls (large tool-augmented turns, extended reasoning). On the + Python ADK it overrides botocore's ~60s read timeout, which otherwise + aborts long completions with a ReadTimeoutError; on the Go ADK it bounds + the whole Converse request (default 30m). When unset, each runtime's + default is used. + minimum: 1 + type: integer region: description: AWS region where the Bedrock model is available (e.g., us-east-1, us-west-2) diff --git a/python/packages/kagent-adk/src/kagent/adk/models/_bedrock.py b/python/packages/kagent-adk/src/kagent/adk/models/_bedrock.py index 5323a7ada..8d66c6edc 100644 --- a/python/packages/kagent-adk/src/kagent/adk/models/_bedrock.py +++ b/python/packages/kagent-adk/src/kagent/adk/models/_bedrock.py @@ -16,7 +16,9 @@ from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional import boto3 +from botocore.config import Config as BotocoreConfig from google.adk.models import BaseLlm +from pydantic import Field from google.adk.models.llm_response import LlmResponse from google.genai import types @@ -85,10 +87,24 @@ def _get_bedrock_client( tls_disable_verify: Optional[bool] = None, tls_ca_cert_path: Optional[str] = None, tls_disable_system_cas: Optional[bool] = None, + read_timeout: Optional[int] = None, + connect_timeout: Optional[int] = None, ): region = os.environ.get("AWS_DEFAULT_REGION") or os.environ.get("AWS_REGION") or "us-east-1" kwargs: dict[str, Any] = {"region_name": region} + # botocore defaults to a 60s read timeout, which aborts long Bedrock + # completions (large tool-augmented turns, extended reasoning) with a + # ReadTimeoutError. Allow callers to raise it. Only build a Config when an + # override is set so we keep botocore's defaults otherwise. + timeout_config: dict[str, Any] = {} + if read_timeout is not None: + timeout_config["read_timeout"] = read_timeout + if connect_timeout is not None: + timeout_config["connect_timeout"] = connect_timeout + if timeout_config: + kwargs["config"] = BotocoreConfig(**timeout_config) + if extra_headers: # boto3 doesn't support custom headers natively; log and ignore logger.warning("extra_headers are not supported for Bedrock models and will be ignored.") @@ -279,6 +295,12 @@ class KAgentBedrockLlm(KAgentTLSMixin, BaseLlm): # "5m"/None (Bedrock's default 5-minute cache) or "1h" (extended-TTL, # narrower model support and higher cache-write cost). See _cache_point_block. cache_ttl: Optional[str] = None + # Bedrock HTTP client timeouts in seconds. read_timeout overrides botocore's + # ~60s default, which otherwise aborts long completions with a + # ReadTimeoutError. None keeps botocore's defaults. Constrained to >= 1 so + # invalid values can't reach the boto3 client (matches the CRD contract). + read_timeout: Optional[int] = Field(default=None, ge=1) + connect_timeout: Optional[int] = Field(default=None, ge=1) model_config = {"arbitrary_types_allowed": True} @@ -289,6 +311,8 @@ def _client(self): tls_disable_verify=self.tls_disable_verify, tls_ca_cert_path=self.tls_ca_cert_path, tls_disable_system_cas=self.tls_disable_system_cas, + read_timeout=self.read_timeout, + connect_timeout=self.connect_timeout, ) @classmethod diff --git a/python/packages/kagent-adk/src/kagent/adk/types.py b/python/packages/kagent-adk/src/kagent/adk/types.py index 867ac317d..ddddd5b2c 100644 --- a/python/packages/kagent-adk/src/kagent/adk/types.py +++ b/python/packages/kagent-adk/src/kagent/adk/types.py @@ -327,6 +327,12 @@ class Bedrock(BaseLLM): # extended-TTL caching. "1h" is supported on fewer models and billed at a # higher cache-write rate, so it is not strictly better than "5m". cache_ttl: Literal["5m", "1h"] | None = None + # Bedrock HTTP client timeouts in seconds. read_timeout overrides botocore's + # ~60s default, which otherwise aborts long completions with a + # ReadTimeoutError. None keeps botocore's defaults. Constrained to >= 1 to + # match the ModelConfig CRD (readTimeout/connectTimeout minimum: 1). + read_timeout: int | None = Field(default=None, ge=1) + connect_timeout: int | None = Field(default=None, ge=1) type: Literal["bedrock"] @@ -709,6 +715,8 @@ def _create_llm_from_model_config(model_config: ModelUnion): additional_model_request_fields=model_config.additional_model_request_fields, prompt_caching=model_config.prompt_caching, cache_ttl=model_config.cache_ttl, + read_timeout=model_config.read_timeout, + connect_timeout=model_config.connect_timeout, **_transport_kwargs(model_config), ) if model_config.type == "sap_ai_core": diff --git a/python/packages/kagent-adk/tests/unittests/models/test_bedrock.py b/python/packages/kagent-adk/tests/unittests/models/test_bedrock.py index 08cbfad5b..37c82ec57 100644 --- a/python/packages/kagent-adk/tests/unittests/models/test_bedrock.py +++ b/python/packages/kagent-adk/tests/unittests/models/test_bedrock.py @@ -172,6 +172,29 @@ def test_defaults_to_us_east_1(self): _get_bedrock_client() assert mock_boto.call_args.kwargs["region_name"] == "us-east-1" + def test_no_config_when_timeouts_unset(self): + with mock.patch("kagent.adk.models._bedrock.boto3.client") as mock_boto: + _get_bedrock_client() + assert "config" not in mock_boto.call_args.kwargs + + def test_read_timeout_sets_botocore_config(self): + from botocore.config import Config as BotocoreConfig + + with mock.patch("kagent.adk.models._bedrock.boto3.client") as mock_boto: + _get_bedrock_client(read_timeout=1800) + config = mock_boto.call_args.kwargs["config"] + assert config.read_timeout == 1800 + # connect_timeout untouched -> botocore keeps its own default, + # whatever that is for the installed version (don't hard-code it). + assert config.connect_timeout == BotocoreConfig().connect_timeout + + def test_read_and_connect_timeout(self): + with mock.patch("kagent.adk.models._bedrock.boto3.client") as mock_boto: + _get_bedrock_client(read_timeout=900, connect_timeout=10) + config = mock_boto.call_args.kwargs["config"] + assert config.read_timeout == 900 + assert config.connect_timeout == 10 + class TestKAgentBedrockLlm: def test_default_construction(self): @@ -182,6 +205,19 @@ def test_non_anthropic_model_accepted(self): llm = KAgentBedrockLlm(model="meta.llama3-8b-instruct-v1:0") assert llm.model == "meta.llama3-8b-instruct-v1:0" + def test_positive_timeouts_accepted(self): + llm = KAgentBedrockLlm(model="meta.llama3-8b-instruct-v1:0", read_timeout=1800, connect_timeout=10) + assert llm.read_timeout == 1800 + assert llm.connect_timeout == 10 + + @pytest.mark.parametrize("field", ["read_timeout", "connect_timeout"]) + @pytest.mark.parametrize("value", [0, -5]) + def test_non_positive_timeouts_rejected(self, field, value): + from pydantic import ValidationError + + with pytest.raises(ValidationError): + KAgentBedrockLlm(model="meta.llama3-8b-instruct-v1:0", **{field: value}) + @pytest.mark.asyncio async def test_generate_calls_converse(self): llm = KAgentBedrockLlm(model="us.anthropic.claude-sonnet-4-20250514-v1:0")