fix: reject base_url and extra_body in generate_content_config
An agent-level `generate_content_config.http_options.base_url` is copied into every LlmRequest and overrides the client transport, so the configured API key and the full prompt/response traffic are sent to that host. Nothing rejected it, so a supplied agent config (including a YAML one) could redirect a credentialed model call to an arbitrary endpoint. `http_options.extra_body` is recursively merged into the serialized request body just before it is sent, and the merge aligns the incoming key case to the target, so it can overwrite `systemInstruction`, `tools` and `generationConfig` — the exact fields the other three checks in this validator exist to reject. It bypassed all of them. Reject both in the field validator. Request-time `http_options` such as headers, timeout, and retry options are unaffected; `base_url` belongs on the model or its client, which is already why `RunConfig.http_options` deliberately does not merge it. Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 953629510
This commit is contained in:
committed by
Copybara-Service
parent
65d8ea7d82
commit
472e4635fb
@@ -1089,6 +1089,18 @@ class LlmAgent(BaseAgent, abc.ABC):
|
||||
raise ValueError(
|
||||
'Response schema must be set via LlmAgent.output_schema.'
|
||||
)
|
||||
if generate_content_config.http_options:
|
||||
if generate_content_config.http_options.base_url:
|
||||
raise ValueError(
|
||||
'Base URL is a transport setting and must be set on the model or'
|
||||
' its client, not via LlmAgent.generate_content_config.'
|
||||
)
|
||||
if generate_content_config.http_options.extra_body:
|
||||
raise ValueError(
|
||||
'Extra body is merged into the request body and can overwrite the'
|
||||
' tools, system instruction and response schema rejected above.'
|
||||
' Set it on the model or its client.'
|
||||
)
|
||||
return generate_content_config
|
||||
|
||||
@override
|
||||
|
||||
@@ -329,6 +329,42 @@ def test_validate_generate_content_config_response_schema_throw():
|
||||
)
|
||||
|
||||
|
||||
def test_validate_generate_content_config_http_options_base_url_throw():
|
||||
"""Tests that a transport base URL cannot be set directly in config."""
|
||||
with pytest.raises(ValueError):
|
||||
_ = LlmAgent(
|
||||
name='test_agent',
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
http_options=types.HttpOptions(base_url='http://example.invalid')
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_validate_generate_content_config_http_options_extra_body_throw():
|
||||
"""Tests that an extra request body cannot be set directly in config."""
|
||||
with pytest.raises(ValueError):
|
||||
_ = LlmAgent(
|
||||
name='test_agent',
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
http_options=types.HttpOptions(
|
||||
extra_body={'systemInstruction': {'parts': [{'text': 'hi'}]}}
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_validate_generate_content_config_http_options_allowed():
|
||||
"""Tests that request-time http options remain settable in config."""
|
||||
agent = LlmAgent(
|
||||
name='test_agent',
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
http_options=types.HttpOptions(timeout=1000)
|
||||
),
|
||||
)
|
||||
|
||||
assert agent.generate_content_config.http_options.timeout == 1000
|
||||
|
||||
|
||||
def test_allow_transfer_by_default():
|
||||
sub_agent = LlmAgent(name='sub_agent')
|
||||
agent = LlmAgent(name='test_agent', sub_agents=[sub_agent])
|
||||
|
||||
Reference in New Issue
Block a user