diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py index 2b37c0c4906..04541d8863e 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py @@ -38,6 +38,7 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): params=params, litellm_params=litellm_params, agent_extra_headers=kwargs.get("agent_extra_headers"), + timeout=kwargs.get("timeout"), ) async def handle_streaming( @@ -58,5 +59,6 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): params=params, litellm_params=litellm_params, agent_extra_headers=kwargs.get("agent_extra_headers"), + timeout=kwargs.get("timeout"), ): yield chunk diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index db57072ca38..64979848110 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -31,6 +31,7 @@ class BedrockAgentCoreA2AHandler: params: dict[str, Any], litellm_params: dict[str, Any], agent_extra_headers: dict[str, str] | None = None, + timeout: float | None = None, ) -> dict[str, Any]: """ Handle non-streaming A2A request to AgentCore. @@ -62,6 +63,7 @@ class BedrockAgentCoreA2AHandler: url, headers=headers, data=body, + timeout=timeout, ) response.raise_for_status() response_data: Final = response.json() @@ -77,6 +79,7 @@ class BedrockAgentCoreA2AHandler: params: dict[str, Any], litellm_params: dict[str, Any], agent_extra_headers: dict[str, str] | None = None, + timeout: float | None = None, ) -> AsyncIterator[dict[str, Any]]: """ Handle streaming A2A request to AgentCore. @@ -110,6 +113,7 @@ class BedrockAgentCoreA2AHandler: headers=headers, data=body, stream=True, + timeout=timeout, ) response.raise_for_status() diff --git a/litellm/a2a_protocol/providers/langflow/config.py b/litellm/a2a_protocol/providers/langflow/config.py index 54d403f88c0..eec32d4e274 100644 --- a/litellm/a2a_protocol/providers/langflow/config.py +++ b/litellm/a2a_protocol/providers/langflow/config.py @@ -6,6 +6,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2ACompletionBridgeHandler, ) from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig +from litellm.interactions.agents.utils import merge_agent_headers from litellm.llms.langflow.a2a import merge_a2a_session_into_litellm_params @@ -28,11 +29,16 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): litellm_params = merge_a2a_session_into_litellm_params( litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) ) + forwarded_headers = merge_agent_headers( + dynamic_headers=kwargs.get("agent_extra_headers"), + static_headers=kwargs.get("agent_static_headers"), + ) return await A2ACompletionBridgeHandler.handle_non_streaming( request_id=request_id, params=params, litellm_params=litellm_params, api_base=api_base, + agent_extra_headers=forwarded_headers, _skip_a2a_provider_routing=True, ) @@ -51,11 +57,16 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): litellm_params = merge_a2a_session_into_litellm_params( litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) ) + forwarded_headers = merge_agent_headers( + dynamic_headers=kwargs.get("agent_extra_headers"), + static_headers=kwargs.get("agent_static_headers"), + ) async for chunk in A2ACompletionBridgeHandler.handle_streaming( request_id=request_id, params=params, litellm_params=litellm_params, api_base=api_base, + agent_extra_headers=forwarded_headers, _skip_a2a_provider_routing=True, ): yield chunk diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index fe9d1ac3c15..0af22578ccc 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -210,12 +210,25 @@ class WatsonxOrchestrateHandler: client: AsyncHTTPHandler, max_attempts: int = _MAX_POLL_ATTEMPTS, interval_s: float = _POLL_INTERVAL_S, + timeout: float | None = None, ) -> _WXORun: url: Final = f"{base_url}/v1/orchestrate/runs/{run_id}" + deadline: float | None = time.monotonic() + max(timeout, 0) if timeout is not None else None for attempt in range(max_attempts): - await asyncio.sleep(interval_s) - response = await client.get(url, headers=auth_headers) + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise asyncio.TimeoutError(f"WXO run '{run_id}' exceeded timeout of {timeout}s") + if interval_s > 0: + await asyncio.sleep(min(interval_s, remaining)) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise asyncio.TimeoutError(f"WXO run '{run_id}' exceeded timeout of {timeout}s") + response = await asyncio.wait_for(client.get(url, headers=auth_headers), timeout=remaining) + else: + await asyncio.sleep(interval_s) + response = await client.get(url, headers=auth_headers) response.raise_for_status() result = WatsonxOrchestrateHandler._run_body(response) status = result.get("status", "") @@ -233,6 +246,7 @@ class WatsonxOrchestrateHandler: base_url: str, auth_headers: dict[str, str], client: AsyncHTTPHandler, + timeout: float | None = None, ) -> _WXORun: status = run_data.get("status", "") if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES: @@ -244,6 +258,7 @@ class WatsonxOrchestrateHandler: run_id=run_id, auth_headers=auth_headers, client=client, + timeout=timeout, ) status = run_data.get("status", "") @@ -341,6 +356,7 @@ class WatsonxOrchestrateHandler: base_url=base_url, auth_headers=auth_headers, client=client, + timeout=timeout, ) response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(run_data) @@ -418,6 +434,7 @@ class WatsonxOrchestrateHandler: base_url=base_url, auth_headers=auth_headers, client=client, + timeout=timeout, ) accumulated_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result) else: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 1e416a21e84..7cfaddc0b0e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1180,6 +1180,15 @@ class CustomStreamWrapper: completion_obj: dict[str, Any], ) -> _ProviderChunkResult: response_obj: dict[str, Any] = {} + if isinstance(chunk, ModelResponseStream) and self.custom_llm_provider == "a2a": + model_response = chunk + model_response.model = self.model + finish_reasons = [getattr(choice, "finish_reason", None) for choice in chunk.choices] + if finish_reasons and all(isinstance(reason, str) and reason for reason in finish_reasons): + self.received_finish_reason = finish_reasons[0] + self.sent_last_chunk = True + return _ProviderChunkEarlyReturn(model_response) + if ( isinstance(chunk, ModelResponseStream) and self.custom_llm_provider is not None @@ -1219,6 +1228,9 @@ class CustomStreamWrapper: raise StopIteration anthropic_response_obj: Final[GChunk] = cast(GChunk, chunk) completion_obj["content"] = anthropic_response_obj["text"] + chunk_index = anthropic_response_obj.get("index") + if isinstance(chunk_index, int): + model_response.choices[0].index = chunk_index if anthropic_response_obj["is_finished"]: self.received_finish_reason = anthropic_response_obj["finish_reason"] diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index ef36296fff5..c2b3750bada 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -850,17 +850,20 @@ async def invoke_agent_a2a( agent_extra_headers=agent_extra_headers, ) + post_call_succeeded = False try: response = await proxy_logging_obj.post_call_success_hook( user_api_key_dict=user_api_key_dict, data=data, response=response, ) + post_call_succeeded = True finally: _enqueue_fn: Final = getattr(logging_obj, "_enqueue_deferred_logging", None) if _enqueue_fn is not None: logging_obj._enqueue_deferred_logging = None - _enqueue_fn(response) + if post_call_succeeded: + _enqueue_fn(response) response_dict: Final[dict[str, Any]] = ( response.model_dump(mode="json", exclude_none=True) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 36557d8bfeb..3e14d31f008 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -459,14 +459,17 @@ def _get_agent_dynamic_headers( return dynamic_headers -def _get_agent_identity_headers(user_api_key_dict: UserAPIKeyAuth | None) -> dict[str, str]: - if user_api_key_dict is None: - return {} +def _get_agent_identity_headers( + user_api_key_dict: UserAPIKeyAuth | None, + trace_id: object | None = None, +) -> dict[str, str]: headers: dict[str, str] = {} - if user_api_key_dict.user_id: + if user_api_key_dict is not None and user_api_key_dict.user_id: headers["X-LiteLLM-User-Id"] = user_api_key_dict.user_id - if user_api_key_dict.team_id: + if user_api_key_dict is not None and user_api_key_dict.team_id: headers["X-LiteLLM-Team-Id"] = user_api_key_dict.team_id + if trace_id: + headers["X-LiteLLM-Trace-Id"] = str(trace_id) return headers @@ -585,7 +588,7 @@ async def route_a2a_agent_request( ) registered_static_headers = merge_agent_headers( dynamic_headers=registered_static_headers, - static_headers=_get_agent_identity_headers(user_api_key_dict), + static_headers=_get_agent_identity_headers(user_api_key_dict, data.get("litellm_trace_id")), ) if ( registered_provider diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 035bc770aed..c6f5cfbe5db 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1890,6 +1890,13 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, trust_client_model_info=False, ) + if isinstance(self.data.get("model"), str) and self.data["model"].startswith("a2a/"): + self.data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=self.data, + call_type=route_type, + guardrails_only=True, + ) # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in