diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index 922052cac17..1401d4ff10a 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -138,6 +138,33 @@ class WatsonxOrchestrateHandler: f"{max_attempts * interval_s:.0f}s" ) + @staticmethod + async def _get_successful_run_data( + run_data: Dict[str, Any], + base_url: str, + auth_headers: Dict[str, str], + client: AsyncHTTPHandler, + ) -> Dict[str, Any]: + status = run_data.get("status", "") + if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES: + run_id = run_data.get("run_id") or run_data.get("id") or "" + if not run_id: + raise ValueError(f"WXO: No run_id in response: {run_data}") + run_data = await WatsonxOrchestrateHandler._poll_run( + base_url=base_url, + run_id=run_id, + auth_headers=auth_headers, + client=client, + ) + status = run_data.get("status", "") + + if status not in WatsonxOrchestrateTransformation.SUCCESS_STATES: + raise RuntimeError( + f"WXO run ended with non-success status '{status}': {run_data}" + ) + + return run_data + @staticmethod async def _accumulate_wxo_sse_text(response: Any) -> str: accumulated_text = "" @@ -237,23 +264,12 @@ class WatsonxOrchestrateHandler: run_response.raise_for_status() run_data: Dict[str, Any] = run_response.json() - status = run_data.get("status", "") - if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES: - run_id = run_data.get("run_id") or run_data.get("id") or "" - if not run_id: - raise ValueError(f"WXO: No run_id in response: {run_data}") - run_data = await WatsonxOrchestrateHandler._poll_run( - base_url=base_url, - run_id=run_id, - auth_headers=auth_headers, - client=client, - ) - status = run_data.get("status", "") - - if status not in WatsonxOrchestrateTransformation.SUCCESS_STATES: - raise RuntimeError( - f"WXO run ended with non-success status '{status}': {run_data}" - ) + run_data = await WatsonxOrchestrateHandler._get_successful_run_data( + run_data=run_data, + base_url=base_url, + auth_headers=auth_headers, + client=client, + ) response_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result( run_data @@ -313,6 +329,12 @@ class WatsonxOrchestrateHandler: if "text/event-stream" not in content_type: response_body = await response.aread() result = json.loads(response_body) + result = await WatsonxOrchestrateHandler._get_successful_run_data( + run_data=result, + base_url=base_url, + auth_headers=auth_headers, + client=client, + ) accumulated_text = ( WatsonxOrchestrateTransformation.extract_text_from_wxo_result( result diff --git a/litellm/proxy/public_endpoints/agent_create_fields.json b/litellm/proxy/public_endpoints/agent_create_fields.json index e58bd97cce7..c5c0437bf4f 100644 --- a/litellm/proxy/public_endpoints/agent_create_fields.json +++ b/litellm/proxy/public_endpoints/agent_create_fields.json @@ -242,7 +242,7 @@ "label": "Username (CP4D only)", "placeholder": "admin", "tooltip": "Your CP4D username. Required when auth_mode is 'cp4d'.", - "required": false, + "required": true, "field_type": "text", "default_value": null, "include_in_litellm_params": true diff --git a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index 4f5f063b37c..c7393851dc3 100644 --- a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -1,6 +1,7 @@ import json import os import sys +from pathlib import Path import pytest @@ -69,6 +70,33 @@ class _InvalidJsonStreamClient: raise AssertionError(url) +class _JsonStreamResponse: + headers = {"content-type": "application/json"} + + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + pass + + async def aread(self): + return json.dumps(self.payload).encode() + + +class _JsonStreamClient: + def __init__(self, stream_payload): + self.stream_payload = stream_payload + self.post_urls = [] + + async def post(self, url, **kwargs): + self.post_urls.append(url) + if "identity/token" in url: + return _JsonResponse({"access_token": "token", "expires_in": 3600}) + if url.endswith("/runs/stream"): + return _JsonStreamResponse(self.stream_payload) + raise AssertionError(url) + + class TestWatsonxOrchestrateTransformation: def test_get_api_base_url(self): url = WatsonxOrchestrateTransformation.get_api_base_url( @@ -222,6 +250,91 @@ async def test_short_lived_tokens_are_not_served_from_cache(): assert client.calls == 2 +@pytest.mark.asyncio +async def test_handle_streaming_polls_non_sse_json_until_complete(monkeypatch): + client = _JsonStreamClient({"status": "running", "run_id": "run-1"}) + poll_calls = [] + + async def poll_run(base_url, run_id, auth_headers, client, **kwargs): + poll_calls.append((base_url, run_id, auth_headers, client)) + return {"status": "completed", "results": "polled text"} + + monkeypatch.setattr( + WatsonxOrchestrateHandler, + "_http_client", + lambda timeout=90.0: client, + ) + monkeypatch.setattr(WatsonxOrchestrateHandler, "_poll_run", poll_run) + + params = { + "message": { + "parts": [ + {"kind": "text", "text": "Hello"}, + ], + } + } + litellm_params = { + "cp4d_host": "https://cpd.example.com", + "instance_id": "instance-id", + "wxo_agent_id": "agent-id", + "api_key": "pending-json-stream-cache-key", + "auth_mode": "ibm_cloud", + } + + events = [ + event + async for event in WatsonxOrchestrateHandler.handle_streaming( + request_id="req-1", + params=params, + litellm_params=litellm_params, + delay_ms=0, + ) + ] + artifact_text = "".join( + event["result"]["artifact"]["parts"][0]["text"] + for event in events + if event["result"].get("kind") == "artifact-update" + ) + + assert len(poll_calls) == 1 + assert poll_calls[0][1] == "run-1" + assert artifact_text == "polled text" + + +@pytest.mark.asyncio +async def test_handle_streaming_raises_for_non_sse_json_failure(monkeypatch): + client = _JsonStreamClient({"status": "failed", "run_id": "run-1"}) + monkeypatch.setattr( + WatsonxOrchestrateHandler, + "_http_client", + lambda timeout=90.0: client, + ) + + params = { + "message": { + "parts": [ + {"kind": "text", "text": "Hello"}, + ], + } + } + litellm_params = { + "cp4d_host": "https://cpd.example.com", + "instance_id": "instance-id", + "wxo_agent_id": "agent-id", + "api_key": "failed-json-stream-cache-key", + "auth_mode": "ibm_cloud", + } + + with pytest.raises(RuntimeError, match="non-success status 'failed'"): + async for _ in WatsonxOrchestrateHandler.handle_streaming( + request_id="req-1", + params=params, + litellm_params=litellm_params, + delay_ms=0, + ): + pass + + @pytest.mark.asyncio async def test_handle_streaming_does_not_fallback_on_invalid_json(monkeypatch): client = _InvalidJsonStreamClient() @@ -263,3 +376,18 @@ def test_config_manager_returns_wxo_provider(): ) assert config is not None assert config.__class__.__name__ == "WatsonxOrchestrateA2AConfig" + + +def test_wxo_cp4d_default_requires_username_in_dashboard_fields(): + fields_path = ( + Path(__file__).resolve().parents[5] + / "litellm/proxy/public_endpoints/agent_create_fields.json" + ) + agent_fields = json.loads(fields_path.read_text()) + wxo_agent = next( + agent for agent in agent_fields if agent["agent_type"] == "watsonx_orchestrate" + ) + fields_by_key = {field["key"]: field for field in wxo_agent["credential_fields"]} + + assert fields_by_key["auth_mode"]["default_value"] == "cp4d" + assert fields_by_key["username"]["required"] is True