diff --git a/api/core/mcp/server/streamable_http.py b/api/core/mcp/server/streamable_http.py index 884610ca82..8869bce8b1 100644 --- a/api/core/mcp/server/streamable_http.py +++ b/api/core/mcp/server/streamable_http.py @@ -215,16 +215,22 @@ def extract_answer_from_response(app: App, response: Any) -> str: def process_streaming_response(response: RateLimitGenerator) -> str: """Process streaming response for agent chat mode""" answer = "" + last_thought = "" for item in response.generator: if isinstance(item, str) and item.startswith("data: "): try: json_str = item[6:].strip() parsed_data = json.loads(json_str) - if parsed_data.get("event") == "agent_thought": - answer += parsed_data.get("thought", "") + event = parsed_data.get("event") + if event in ("message", "agent_message"): + answer += parsed_data.get("answer", "") + elif event == "agent_thought": + thought = parsed_data.get("thought", "") + if thought: + last_thought = thought except json.JSONDecodeError: continue - return answer + return answer or last_thought def process_mapping_response(app: App, response: Mapping) -> str: diff --git a/api/tests/unit_tests/core/mcp/server/test_streamable_http.py b/api/tests/unit_tests/core/mcp/server/test_streamable_http.py index 57456085c3..5bf48ebf87 100644 --- a/api/tests/unit_tests/core/mcp/server/test_streamable_http.py +++ b/api/tests/unit_tests/core/mcp/server/test_streamable_http.py @@ -357,21 +357,49 @@ class TestUtilityFunctions: assert result == expected def test_extract_answer_from_streaming_response(self): - """Test extracting answer from streaming response""" + """Test extracting answer from streaming response with agent_message events""" app = Mock(spec=App) - # Mock RateLimitGenerator mock_generator = Mock(spec=RateLimitGenerator) mock_generator.generator = [ 'data: {"event": "agent_thought", "thought": "thinking..."}', - 'data: {"event": "agent_thought", "thought": "more thinking"}', - 'data: {"event": "other", "content": "ignore this"}', + 'data: {"event": "agent_message", "answer": "Hello "}', + 'data: {"event": "agent_message", "answer": "World"}', + 'data: {"event": "message_end", "metadata": {}}', "not data format", ] result = extract_answer_from_response(app, mock_generator) - assert result == "thinking...more thinking" + assert result == "Hello World" + + def test_extract_answer_from_streaming_response_message_event(self): + """Test extracting answer from streaming response with message event""" + app = Mock(spec=App) + + mock_generator = Mock(spec=RateLimitGenerator) + mock_generator.generator = [ + 'data: {"event": "message", "answer": "Hello from chat"}', + ] + + result = extract_answer_from_response(app, mock_generator) + + assert result == "Hello from chat" + + def test_extract_answer_from_streaming_response_fallback_to_thought(self): + """Test extracting answer falls back to thought when no message events""" + app = Mock(spec=App) + + mock_generator = Mock(spec=RateLimitGenerator) + mock_generator.generator = [ + 'data: {"event": "agent_thought", "thought": "thinking..."}', + 'data: {"event": "agent_thought", "thought": "more thinking"}', + 'data: {"event": "other", "content": "ignore this"}', + ] + + result = extract_answer_from_response(app, mock_generator) + + assert result == "more thinking" def test_process_mapping_response_invalid_mode(self): """Test processing mapping response with invalid app mode"""