From cce29e724064331e00382a73a239efe8dd44bcfe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97SO?= <142557582+Linxiushen@users.noreply.github.com> Date: Sun, 9 Aug 2026 14:44:50 +0800 Subject: [PATCH] fix(agentlang): scope stream fallback correlations --- .../agentlang/event/correlation_id_manager.py | 33 +++- .../llms/processors/processor_manager.py | 4 +- .../llms/processors/regular_call_processor.py | 2 +- backend/super-magic/app/magic/agent.py | 9 + .../test_correlation_id_manager_scope.py | 163 ++++++++++++++++++ .../core/test_agent_correlation_cleanup.py | 58 +++++++ 6 files changed, 258 insertions(+), 11 deletions(-) create mode 100644 backend/super-magic/tests/core/test_agent_correlation_cleanup.py diff --git a/backend/super-magic/agentlang/agentlang/event/correlation_id_manager.py b/backend/super-magic/agentlang/agentlang/event/correlation_id_manager.py index fb85181258..17d8491c9d 100644 --- a/backend/super-magic/agentlang/agentlang/event/correlation_id_manager.py +++ b/backend/super-magic/agentlang/agentlang/event/correlation_id_manager.py @@ -109,10 +109,11 @@ def __init__(self): # 同一个全局 manager 内，不同 AgentContext 的事件互不消费。 # 注意：由于 asyncio 是单线程模型，不需要线程锁 self._active_correlations: Dict[Tuple[str, EventPairType], str] = {} - # V2 流式中断降级时保存的 correlation_id： + # V2 流式中断降级时按 AgentContext scope 保存 correlation_id。 + # 同一 AgentContext 的 Agent loop 串行发起 LLM 请求，不同 context 可并发运行。 # V2 流式 chunk 以 request_id 作为 correlation_id 推送前端， # 流式中断降级非流式时需沿用，否则前端会将 chunk 和最终消息视为两条消息。 - self._stream_fallback_cid: Optional[str] = None + self._stream_fallback_cids: Dict[str, str] = {} def generate_for_before_event( self, @@ -231,23 +232,39 @@ def _normalize_scope_id(scope_id: Optional[str]) -> str: def _active_key(event_pair_type: EventPairType, scope_id: str) -> Tuple[str, EventPairType]: return scope_id, event_pair_type - def set_stream_fallback_cid(self, cid: Optional[str]) -> None: + def set_stream_fallback_cid( + self, + cid: Optional[str], + scope_id: Optional[str] = None, + ) -> None: """保存 V2 流式中断后的降级 correlation_id Args: cid: 流式请求的 request_id（将作为降级后非流式消息的 correlation_id），传 None 表示清除 + scope_id: AgentContext 唯一标识；未传时使用全局 scope，兼容无 AgentContext 场景。 """ - self._stream_fallback_cid = cid + if cid is None: + self.clear_stream_fallback_cid(scope_id) + return + normalized_scope_id = self._normalize_scope_id(scope_id) + self._stream_fallback_cids[normalized_scope_id] = cid + + def clear_stream_fallback_cid(self, scope_id: Optional[str] = None) -> None: + """清除指定 scope 未消费的 V2 流式降级 correlation_id。""" + normalized_scope_id = self._normalize_scope_id(scope_id) + self._stream_fallback_cids.pop(normalized_scope_id, None) - def pop_stream_fallback_cid(self) -> Optional[str]: + def pop_stream_fallback_cid(self, scope_id: Optional[str] = None) -> Optional[str]: """取出并清除 V2 流式降级 correlation_id（仅使用一次） + Args: + scope_id: AgentContext 唯一标识；未传时使用全局 scope，兼容无 AgentContext 场景。 + Returns: Optional[str]: 已保存的 correlation_id，如果没有则返回 None """ - cid = self._stream_fallback_cid - self._stream_fallback_cid = None - return cid + normalized_scope_id = self._normalize_scope_id(scope_id) + return self._stream_fallback_cids.pop(normalized_scope_id, None) # 全局单例实例 diff --git a/backend/super-magic/agentlang/agentlang/llms/processors/processor_manager.py b/backend/super-magic/agentlang/agentlang/llms/processors/processor_manager.py index 5bc3b863ad..f0200a96e9 100644 --- a/backend/super-magic/agentlang/agentlang/llms/processors/processor_manager.py +++ b/backend/super-magic/agentlang/agentlang/llms/processors/processor_manager.py @@ -74,7 +74,7 @@ async def execute_llm_call( # 保存流式 request_id 到 CorrelationIdManager，供降级非流式时恢复 correlation_id 一致性： # V2 流式 chunk 以 request_id 作为 correlation_id 推送给前端， # 若流式中断降级非流式，非流式的最终消息也应携带相同的 correlation_id - correlation_manager.set_stream_fallback_cid(request_id) + correlation_manager.set_stream_fallback_cid(request_id, correlation_scope_id) # 流式调用，异常直接抛出，不做任何重试/fallback 决策 response = await StreamingCallProcessor.call_with_stream( @@ -89,7 +89,7 @@ async def execute_llm_call( enable_llm_response_events=enable_llm_response_events, ) # 流式成功，清除 fallback 标记，避免影响后续独立的非流式调用 - correlation_manager.set_stream_fallback_cid(None) + correlation_manager.clear_stream_fallback_cid(correlation_scope_id) return response else: # 非流式调用 diff --git a/backend/super-magic/agentlang/agentlang/llms/processors/regular_call_processor.py b/backend/super-magic/agentlang/agentlang/llms/processors/regular_call_processor.py index cc33b0e1a4..2dfdc9b4f0 100644 --- a/backend/super-magic/agentlang/agentlang/llms/processors/regular_call_processor.py +++ b/backend/super-magic/agentlang/agentlang/llms/processors/regular_call_processor.py @@ -154,7 +154,7 @@ async def _trigger_response_events( stream_fallback_cid: Optional[str] = None correlation_scope_id = getattr(agent_context, "context_id", None) if not cm.get_active_correlation_id(EventPairType.AGENT_REPLY, correlation_scope_id): - stream_fallback_cid = cm.pop_stream_fallback_cid() + stream_fallback_cid = cm.pop_stream_fallback_cid(correlation_scope_id) if stream_fallback_cid: logger.info( f"[{request_id}] 检测到 V2 流式降级场景，" diff --git a/backend/super-magic/app/magic/agent.py b/backend/super-magic/app/magic/agent.py index 8bfa101431..b743c14442 100644 --- a/backend/super-magic/app/magic/agent.py +++ b/backend/super-magic/app/magic/agent.py @@ -321,12 +321,19 @@ def close(self) -> None: return self._closed = True + self._clear_stream_fallback_correlation() + if self._context_registered: from app.core.context.agent_context_registry import AgentContextRegistry AgentContextRegistry.get_instance().unregister(self.agent_context) self._context_registered = False logger.info(f"Agent context 已注销: {self.agent_context.get_agent_session_label()}") + def _clear_stream_fallback_correlation(self) -> None: + """清理当前 AgentContext 未消费的流式降级标记。""" + from agentlang.event import get_correlation_manager + get_correlation_manager().clear_stream_fallback_cid(self.agent_context.context_id) + def dispose(self) -> None: """兼容性别名，语义等同于 close()。""" self.close() @@ -809,6 +816,8 @@ async def run_main_agent(self, query: str): final_task_state=final_task_state, )) finally: + # 确定性错误或用户取消可能跳过非流式 fallback，避免标记污染后续 run。 + self._clear_stream_fallback_correlation() # 每次请求结束后回收内存碎片：强制 GC 并归还空闲 arena 给 OS self._reclaim_memory() diff --git a/backend/super-magic/tests/agentlang/test_correlation_id_manager_scope.py b/backend/super-magic/tests/agentlang/test_correlation_id_manager_scope.py index ea5e99d68a..a8323315ea 100644 --- a/backend/super-magic/tests/agentlang/test_correlation_id_manager_scope.py +++ b/backend/super-magic/tests/agentlang/test_correlation_id_manager_scope.py @@ -1,4 +1,53 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + from agentlang.event.correlation_id_manager import CorrelationIdManager, EventPairType +from agentlang.event.reply_event_manager import ReplyEventManager +from agentlang.event.think_event_manager import ThinkEventManager +from agentlang.llms.processors.processor_config import ProcessorConfig +from agentlang.llms.processors.processor_manager import ProcessorManager +from agentlang.llms.processors.streaming_call_processor import StreamingCallProcessor + + +class _AgentContext: + def __init__(self, context_id: str): + self.context_id = context_id + self.metadata = {} + + def set_metadata(self, key, value): + self.metadata[key] = value + + +async def _execute_stream(context_id, request_id): + return await ProcessorManager.execute_llm_call( + client=None, + llm_config=None, + request_params={}, + model_id="mock-model", + processor_config=ProcessorConfig(use_stream_mode=True), + agent_context=_AgentContext(context_id), + request_id=request_id, + ) + + +async def _execute_regular(context_id, request_id, response): + client = SimpleNamespace( + chat=SimpleNamespace( + completions=SimpleNamespace(create=AsyncMock(return_value=response)), + ) + ) + return await ProcessorManager.execute_llm_call( + client=client, + llm_config=SimpleNamespace(name="Mock Model"), + request_params={}, + model_id="mock-model", + processor_config=ProcessorConfig(use_stream_mode=False), + agent_context=_AgentContext(context_id), + request_id=request_id, + ) def test_tool_call_correlation_is_isolated_by_scope(): @@ -26,3 +75,117 @@ def test_correlation_manager_keeps_global_scope_compatibility(): assert retried_correlation_id == first_correlation_id assert manager.consume_for_after_event(EventPairType.TOOL_CALL) == first_correlation_id + + +def test_stream_fallback_correlation_is_isolated_by_scope(): + manager = CorrelationIdManager() + + manager.set_stream_fallback_cid("parent-request", "mock-parent-context") + manager.set_stream_fallback_cid("child-request", "mock-child-context") + + assert manager.pop_stream_fallback_cid("mock-parent-context") == "parent-request" + assert manager.pop_stream_fallback_cid("mock-child-context") == "child-request" + + +def test_clearing_stream_fallback_only_affects_target_scope(): + manager = CorrelationIdManager() + + manager.set_stream_fallback_cid("parent-request", "mock-parent-context") + manager.set_stream_fallback_cid("child-request", "mock-child-context") + manager.clear_stream_fallback_cid("mock-child-context") + + assert manager.pop_stream_fallback_cid("mock-child-context") is None + assert manager.pop_stream_fallback_cid("mock-parent-context") == "parent-request" + + +def test_stream_fallback_correlation_keeps_global_scope_compatibility(): + manager = CorrelationIdManager() + + manager.set_stream_fallback_cid("stale-request") + manager.set_stream_fallback_cid(None) + assert manager.pop_stream_fallback_cid() is None + + manager.set_stream_fallback_cid("global-request") + + assert manager.pop_stream_fallback_cid() == "global-request" + assert manager.pop_stream_fallback_cid() is None + + +@pytest.mark.asyncio +async def test_concurrent_stream_failures_keep_each_context_fallback(monkeypatch): + manager = CorrelationIdManager() + entered_requests = set() + both_streams_entered = asyncio.Event() + + async def fail_after_both_streams_enter(**kwargs): + entered_requests.add(kwargs["request_id"]) + if len(entered_requests) == 2: + both_streams_entered.set() + await asyncio.wait_for(both_streams_entered.wait(), timeout=5) + raise RuntimeError("mock stream failure") + + monkeypatch.setattr("agentlang.event.get_correlation_manager", lambda: manager) + monkeypatch.setattr( + StreamingCallProcessor, + "call_with_stream", + fail_after_both_streams_enter, + ) + + results = await asyncio.gather( + _execute_stream("mock-parent-context", "parent-request"), + _execute_stream("mock-child-context", "child-request"), + return_exceptions=True, + ) + + assert all(isinstance(result, RuntimeError) for result in results) + after_reply = AsyncMock() + before_reply = AsyncMock() + monkeypatch.setattr(ThinkEventManager, "trigger_after_think", AsyncMock(return_value=True)) + monkeypatch.setattr(ReplyEventManager, "trigger_before_reply", before_reply) + monkeypatch.setattr(ReplyEventManager, "trigger_after_reply", after_reply) + + response = SimpleNamespace( + choices=[SimpleNamespace( + message=SimpleNamespace( + reasoning_content=None, + content="mock final response", + tool_calls=None, + ) + )] + ) + + await _execute_regular("mock-parent-context", "parent-fallback-request", response) + await _execute_regular("mock-child-context", "child-fallback-request", response) + + before_reply.assert_not_awaited() + fallback_correlations = [ + call.kwargs["correlation_id"] + for call in after_reply.await_args_list + ] + assert fallback_correlations == ["parent-request", "child-request"] + assert manager.pop_stream_fallback_cid("mock-parent-context") is None + assert manager.pop_stream_fallback_cid("mock-child-context") is None + + +@pytest.mark.asyncio +async def test_successful_stream_clears_only_its_own_context(monkeypatch): + manager = CorrelationIdManager() + + async def fail_parent_and_complete_child(**kwargs): + if kwargs["request_id"] == "parent-request": + raise RuntimeError("mock parent stream failure") + return "mock child response" + + monkeypatch.setattr("agentlang.event.get_correlation_manager", lambda: manager) + monkeypatch.setattr( + StreamingCallProcessor, + "call_with_stream", + fail_parent_and_complete_child, + ) + + with pytest.raises(RuntimeError, match="mock parent stream failure"): + await _execute_stream("mock-parent-context", "parent-request") + assert await _execute_stream("mock-child-context", "child-request") == "mock child response" + + assert manager.pop_stream_fallback_cid("mock-parent-context") == "parent-request" + assert manager.pop_stream_fallback_cid("mock-child-context") is None diff --git a/backend/super-magic/tests/core/test_agent_correlation_cleanup.py b/backend/super-magic/tests/core/test_agent_correlation_cleanup.py new file mode 100644 index 0000000000..af643f5801 --- /dev/null +++ b/backend/super-magic/tests/core/test_agent_correlation_cleanup.py @@ -0,0 +1,58 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +import app.service # noqa: F401 # Ensure service package finishes initialization before importing Agent. +from agentlang.event.correlation_id_manager import CorrelationIdManager +from app.magic.agent import Agent + + +def _create_uninitialized_agent(context_id: str) -> Agent: + agent = Agent.__new__(Agent) + agent.agent_context = SimpleNamespace( + context_id=context_id, + dispatch_event=AsyncMock(), + ) + agent.agent_name = "mock-agent" + agent._closed = False + agent._context_registered = False + return agent + + +def test_agent_close_clears_only_its_stream_fallback_scope(monkeypatch): + manager = CorrelationIdManager() + manager.set_stream_fallback_cid("parent-request", "parent-context") + manager.set_stream_fallback_cid("child-request", "child-context") + monkeypatch.setattr("agentlang.event.get_correlation_manager", lambda: manager) + + agent = _create_uninitialized_agent("child-context") + agent.close() + agent.close() + + assert manager.pop_stream_fallback_cid("child-context") is None + assert manager.pop_stream_fallback_cid("parent-context") == "parent-request" + + +@pytest.mark.asyncio +async def test_run_main_agent_cancellation_clears_only_its_stream_fallback_scope(monkeypatch): + manager = CorrelationIdManager() + manager.set_stream_fallback_cid("parent-request", "parent-context") + manager.set_stream_fallback_cid("child-request", "child-context") + monkeypatch.setattr("agentlang.event.get_correlation_manager", lambda: manager) + monkeypatch.setattr( + "app.magic.agent.BeforeMainAgentRunEventData", + lambda **kwargs: SimpleNamespace(**kwargs), + ) + + agent = _create_uninitialized_agent("child-context") + agent.run = AsyncMock(side_effect=asyncio.CancelledError()) + agent._reclaim_memory = Mock() + + with pytest.raises(asyncio.CancelledError): + await agent.run_main_agent("mock query") + + agent._reclaim_memory.assert_called_once_with() + assert manager.pop_stream_fallback_cid("child-context") is None + assert manager.pop_stream_fallback_cid("parent-context") == "parent-request"