diff --git a/src/kimi_cli/soul/kimisoul.py b/src/kimi_cli/soul/kimisoul.py index 3f14c2a2f7..dac4a1f29d 100644 --- a/src/kimi_cli/soul/kimisoul.py +++ b/src/kimi_cli/soul/kimisoul.py @@ -1059,19 +1059,16 @@ async def _agent_loop(self) -> TurnOutcome: # --- StopFailure hook --- from kimi_cli.hooks import events as _hook_events - _hook_task = asyncio.create_task( - self._hook_engine.trigger( - "StopFailure", - matcher_value=type(e).__name__, - input_data=_hook_events.stop_failure( - session_id=self._runtime.session.id, - cwd=str(Path.cwd()), - error_type=type(e).__name__, - error_message=str(e), - ), - ) + self._hook_engine.fire_and_forget_trigger( + "StopFailure", + matcher_value=type(e).__name__, + input_data=_hook_events.stop_failure( + session_id=self._runtime.session.id, + cwd=str(Path.cwd()), + error_type=type(e).__name__, + error_message=str(e), + ), ) - _hook_task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None) # break the agent loop raise @@ -1139,22 +1136,19 @@ async def _append_notification(view: NotificationView) -> None: # --- Notification hook --- from kimi_cli.hooks import events - _hook_task = asyncio.create_task( - self._hook_engine.trigger( - "Notification", - matcher_value=view.event.type, - input_data=events.notification( - session_id=self._runtime.session.id, - cwd=str(Path.cwd()), - sink="llm", - notification_type=view.event.type, - title=view.event.title, - body=view.event.body, - severity=view.event.severity, - ), - ) + self._hook_engine.fire_and_forget_trigger( + "Notification", + matcher_value=view.event.type, + input_data=events.notification( + session_id=self._runtime.session.id, + cwd=str(Path.cwd()), + sink="llm", + notification_type=view.event.type, + title=view.event.title, + body=view.event.body, + severity=view.event.severity, + ), ) - _hook_task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None) await self._runtime.notifications.deliver_pending( "llm", @@ -1629,19 +1623,16 @@ async def _compact_with_retry() -> CompactionResult: track_kwargs["trace_id"] = compaction_result.trace_id track("compaction_finished", **track_kwargs) - _hook_task = asyncio.create_task( - self._hook_engine.trigger( - "PostCompact", - matcher_value=trigger_reason, - input_data=events.post_compact( - session_id=self._runtime.session.id, - cwd=str(Path.cwd()), - trigger=trigger_reason, - estimated_token_count=estimated_token_count, - ), - ) + self._hook_engine.fire_and_forget_trigger( + "PostCompact", + matcher_value=trigger_reason, + input_data=events.post_compact( + session_id=self._runtime.session.id, + cwd=str(Path.cwd()), + trigger=trigger_reason, + estimated_token_count=estimated_token_count, + ), ) - _hook_task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None) @staticmethod def _is_retryable_error(exception: BaseException) -> bool: diff --git a/src/kimi_cli/soul/toolset.py b/src/kimi_cli/soul/toolset.py index 5d66344aaa..5504093fef 100644 --- a/src/kimi_cli/soul/toolset.py +++ b/src/kimi_cli/soul/toolset.py @@ -506,22 +506,17 @@ async def _call(): call_id=tool_call.id, ) # --- PostToolUseFailure (fire-and-forget) --- - _hook_task = asyncio.create_task( - self._hook_engine.trigger( - "PostToolUseFailure", - matcher_value=tool_name, - input_data=events.post_tool_use_failure( - session_id=_get_session_id(), - cwd=str(Path.cwd()), - tool_name=tool_name, - tool_input=tool_input_dict, - error=str(e), - tool_call_id=tool_call.id, - ), - ) - ) - _hook_task.add_done_callback( - lambda t: t.exception() if not t.cancelled() else None + self._hook_engine.fire_and_forget_trigger( + "PostToolUseFailure", + matcher_value=tool_name, + input_data=events.post_tool_use_failure( + session_id=_get_session_id(), + cwd=str(Path.cwd()), + tool_name=tool_name, + tool_input=tool_input_dict, + error=str(e), + tool_call_id=tool_call.id, + ), ) from kimi_cli.telemetry import track @@ -574,21 +569,18 @@ async def _call(): ) # --- PostToolUse (fire-and-forget) --- - _hook_task = asyncio.create_task( - self._hook_engine.trigger( - "PostToolUse", - matcher_value=tool_name, - input_data=events.post_tool_use( - session_id=_get_session_id(), - cwd=str(Path.cwd()), - tool_name=tool_name, - tool_input=tool_input_dict, - tool_output=str(ret)[:2000], - tool_call_id=tool_call.id, - ), - ) + self._hook_engine.fire_and_forget_trigger( + "PostToolUse", + matcher_value=tool_name, + input_data=events.post_tool_use( + session_id=_get_session_id(), + cwd=str(Path.cwd()), + tool_name=tool_name, + tool_input=tool_input_dict, + tool_output=str(ret)[:2000], + tool_call_id=tool_call.id, + ), ) - _hook_task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None) return ToolResult(tool_call_id=tool_call.id, return_value=ret) diff --git a/tests/hooks/test_fire_and_forget.py b/tests/hooks/test_fire_and_forget.py new file mode 100644 index 0000000000..b3bd6c5ecf --- /dev/null +++ b/tests/hooks/test_fire_and_forget.py @@ -0,0 +1,103 @@ +"""Background hook triggers must survive garbage collection. + +``asyncio`` only keeps weak references to running tasks, so +``asyncio.create_task(engine.trigger(...))`` followed by discarding the local +variable lets the GC collect a still-pending task. The hook subprocess is then +never awaited and the hook silently does not run. +``HookEngine.fire_and_forget_trigger`` exists to hold a strong reference for the +lifetime of the task; these tests pin that behaviour and check that the +fire-and-forget call sites use it instead of rolling their own ``create_task``. +""" + +import ast +import gc +import inspect +from pathlib import Path + +import pytest + +from kimi_cli.hooks.config import HookDef +from kimi_cli.hooks.engine import HookEngine +from kimi_cli.soul import kimisoul, toolset + + +@pytest.mark.asyncio +async def test_pending_task_survives_gc_and_still_runs(tmp_path): + """The task must complete even when the caller keeps no reference to it.""" + marker = tmp_path / "hook-ran" + hooks = [ + HookDef( + event="PostToolUse", + command=f"sleep 0.2 && touch {marker}", + timeout=5, + ) + ] + engine = HookEngine(hooks, cwd=str(tmp_path)) + + def fire_without_keeping_a_reference() -> None: + engine.fire_and_forget_trigger("PostToolUse", input_data={"tool_name": "Shell"}) + + fire_without_keeping_a_reference() + gc.collect() + + pending = set(engine._pending_fire_and_forget) + assert len(pending) == 1, "engine must hold a strong reference to the pending task" + + await next(iter(pending)) + assert marker.exists(), "hook command did not run to completion" + + +@pytest.mark.asyncio +async def test_completed_task_is_released(tmp_path): + """The strong reference must not leak once the task finishes.""" + engine = HookEngine([HookDef(event="Stop", command="exit 0", timeout=5)], cwd=str(tmp_path)) + + task = engine.fire_and_forget_trigger("Stop", input_data={}) + await task + + assert task not in engine._pending_fire_and_forget + + +@pytest.mark.asyncio +async def test_failing_command_does_not_leak_the_reference(tmp_path): + """A hook whose command cannot be run must still release its reference.""" + engine = HookEngine( + [HookDef(event="Stop", command="definitely-not-a-real-command", timeout=5)], + cwd=str(tmp_path), + ) + + task = engine.fire_and_forget_trigger("Stop", input_data={}) + await task + + assert task.exception() is None + assert task not in engine._pending_fire_and_forget + + +def _bare_create_task_hook_triggers(module) -> list[str]: + """Return `Event` names triggered via a bare `create_task` in `module`.""" + source = Path(inspect.getsourcefile(module)).read_text(encoding="utf-8") + found: list[str] = [] + + for node in ast.walk(ast.parse(source)): + if not isinstance(node, ast.Call): + continue + func = node.func + if not (isinstance(func, ast.Attribute) and func.attr == "create_task"): + continue + for inner in ast.walk(node): + if ( + isinstance(inner, ast.Call) + and isinstance(inner.func, ast.Attribute) + and inner.func.attr == "trigger" + and inner.args + and isinstance(inner.args[0], ast.Constant) + ): + found.append(inner.args[0].value) + + return found + + +@pytest.mark.parametrize("module", [toolset, kimisoul], ids=lambda m: m.__name__) +def test_no_bare_create_task_around_hook_triggers(module): + """Fire-and-forget call sites must go through `fire_and_forget_trigger`.""" + assert _bare_create_task_hook_triggers(module) == []