Coverage for merco/core/interrupt.py: 100%

1"""Agent 中断清理管线，处理中断时的资源清理。""" 

2 

3import logging 

4from abc import ABC, abstractmethod 

5from dataclasses import dataclass, field 

6from typing import Any 

7 

8logger = logging.getLogger("merco.core.interrupt") 

9 

10 

11@dataclass 

12class CleanupContext: 

13 """中断清理上下文。""" 

14 agent: Any # Agent 类型，避免循环导入 

15 cancelled_tool_calls: list[dict] 

16 session_id: str 

17 results: dict[str, Any] = field(default_factory=dict) 

18 

19 

20class CleanupProcessor(ABC): 

21 """清理处理器基类。""" 

22 name: str = "" 

23 

24 @abstractmethod 

25 async def process(self, ctx: CleanupContext) -> bool: 

26 """返回 True 表示已处理，停止管线。""" 

27 ... 

28 

29 

30class InjectCancelMessages(CleanupProcessor): 

31 """为孤儿 tool_calls 注入取消消息。""" 

32 name = "inject_cancel" 

33 

34 async def process(self, ctx: CleanupContext) -> bool: 

35 completed_ids = set() 

36 for msg in ctx.agent.context.messages: 

37 if msg.get("tool_call_id"): 

38 completed_ids.add(msg["tool_call_id"]) 

39 

40 for msg in reversed(ctx.agent.context.messages): 

41 if msg.get("role") != "assistant": 

42 continue 

43 for tc in (msg.get("tool_calls") or []): 

44 tc_id = tc.get("id") if isinstance(tc, dict) else None 

45 if tc_id and tc_id not in completed_ids: 

46 tool_msg = { 

47 "role": "tool", 

48 "tool_call_id": tc_id, 

49 "content": "取消 (Ctrl+C)" 

50 } 

51 ctx.agent.context.add(tool_msg) 

52 ctx.agent.session.add_message("tool", "取消 (Ctrl+C)", tool_call_id=tc_id) 

53 return False 

54 

55 

56class TerminateSubprocesses(CleanupProcessor): 

57 """kill 所有运行中的子进程。""" 

58 name = "kill_subprocesses" 

59 

60 async def process(self, ctx: CleanupContext) -> bool: 

61 bash_tool = ctx.agent.tool_registry.get("bash") if ctx.agent.tool_registry else None 

62 if bash_tool and hasattr(bash_tool, '_active_processes'): 

63 for proc in list(bash_tool._active_processes): 

64 proc.kill() 

65 bash_tool._active_processes.clear() 

66 return False 

67 

68 

69class CloseMCPConnections(CleanupProcessor): 

70 """关闭 MCP 连接。""" 

71 name = "close_mcp" 

72 

73 async def process(self, ctx: CleanupContext) -> bool: 

74 if ctx.agent.mcp_manager: 

75 await ctx.agent.mcp_manager.shutdown() 

76 return False 

77 

78 

79class EmitInterruptHooks(CleanupProcessor): 

80 """发射中断钩子。""" 

81 name = "emit_hooks" 

82 

83 async def process(self, ctx: CleanupContext) -> bool: 

84 await ctx.agent.hooks.emit( 

85 "agent.interrupted", 

86 interrupted_tools=len(ctx.cancelled_tool_calls), 

87 session_id=ctx.session_id, 

88 ) 

89 return False 

90 

91 

92class SavePartialState(CleanupProcessor): 

93 """保存 session + observer 快照。""" 

94 name = "save_state" 

95 

96 async def process(self, ctx: CleanupContext) -> bool: 

97 ctx.agent.observer.save() 

98 ctx.agent.session.metadata["observer"] = ctx.agent.observer.snapshot() 

99 ctx.agent.session.save() 

100 ctx.agent._session_store.save_metadata(ctx.agent.session.id, ctx.agent.session.metadata) 

101 return False 

102 

103 

104class InterruptCleanupPipeline: 

105 """中断清理管线。""" 

106 

107 def __init__(self): 

108 self._processors: list[CleanupProcessor] = [] 

109 

110 def use(self, processor: CleanupProcessor) -> "InterruptCleanupPipeline": 

111 self._processors.append(processor) 

112 return self 

113 

114 async def process(self, ctx: CleanupContext) -> None: 

115 for processor in self._processors: 

116 try: 

117 if await processor.process(ctx): 

118 return 

119 except Exception: 

120 logger.warning("清理处理器 '%s' 异常", processor.name, exc_info=True)