mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/cli: fix Ctrl+C after mid-turn prompts (SIGINT cancel)
asyncio signal handlers freeze ContextVars at install time; reinstalling while allow_task_cancel was still False left cancel permanently disabled. Use a process-level prompt depth counter and reinstall only after depth=0.
This commit is contained in:
@@ -1,16 +1,22 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import signal
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
|
||||
from plyngent.cli.interrupt import (
|
||||
allow_task_cancel,
|
||||
pause_task_cancel_for_prompt,
|
||||
run_in_prompt_thread,
|
||||
set_sigint_reinstall,
|
||||
)
|
||||
from plyngent.cli.limits import prompt_continue_limit, prompt_continue_limit_async
|
||||
from plyngent.cli.retry import run_cancellable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pytest
|
||||
pass
|
||||
|
||||
|
||||
def test_pause_task_cancel_for_prompt() -> None:
|
||||
@@ -20,6 +26,16 @@ def test_pause_task_cancel_for_prompt() -> None:
|
||||
assert allow_task_cancel() is True
|
||||
|
||||
|
||||
def test_nested_pause_depth() -> None:
|
||||
assert allow_task_cancel() is True
|
||||
with pause_task_cancel_for_prompt():
|
||||
assert allow_task_cancel() is False
|
||||
with pause_task_cancel_for_prompt():
|
||||
assert allow_task_cancel() is False
|
||||
assert allow_task_cancel() is False
|
||||
assert allow_task_cancel() is True
|
||||
|
||||
|
||||
def test_prompt_continue_limit_under_pause(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def _confirm(*_a: object, **_k: object) -> bool:
|
||||
assert allow_task_cancel() is False
|
||||
@@ -36,7 +52,6 @@ async def test_run_in_prompt_thread_pauses_cancel() -> None:
|
||||
def work() -> str:
|
||||
return "ok"
|
||||
|
||||
# ContextVar may not propagate to worker threads; assert pause around the call.
|
||||
result = await run_in_prompt_thread(work)
|
||||
assert result == "ok"
|
||||
assert allow_task_cancel() is True
|
||||
@@ -48,3 +63,53 @@ async def test_prompt_continue_limit_async(monkeypatch: pytest.MonkeyPatch) -> N
|
||||
|
||||
monkeypatch.setattr("click.confirm", _confirm)
|
||||
assert await prompt_continue_limit_async("too many rounds") is True
|
||||
|
||||
|
||||
async def test_sigint_cancels_after_prompt_pause() -> None:
|
||||
"""Regression: after a mid-turn prompt, SIGINT must still cancel the turn.
|
||||
|
||||
Previously, reinstalling the asyncio SIGINT handler while
|
||||
allow_task_cancel was still False froze that value into the callback
|
||||
forever (ContextVar snapshot at add_signal_handler).
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
loop.add_signal_handler(signal.SIGINT, lambda: None)
|
||||
loop.remove_signal_handler(signal.SIGINT)
|
||||
except NotImplementedError, RuntimeError, ValueError:
|
||||
pytest.skip("asyncio signal handlers not available on this platform")
|
||||
|
||||
started = asyncio.Event()
|
||||
cancelled = asyncio.Event()
|
||||
|
||||
async def hang() -> None:
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.sleep(3600)
|
||||
except asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
raise
|
||||
|
||||
async def turn() -> None:
|
||||
# Simulate soft-confirm pause mid-turn, then keep streaming.
|
||||
with pause_task_cancel_for_prompt():
|
||||
assert allow_task_cancel() is False
|
||||
assert allow_task_cancel() is True
|
||||
await hang()
|
||||
|
||||
task = asyncio.create_task(run_cancellable(turn()))
|
||||
await started.wait()
|
||||
# Give run_cancellable a tick to install the handler after the pause reinstall.
|
||||
await asyncio.sleep(0)
|
||||
assert allow_task_cancel() is True
|
||||
# Deliver SIGINT the same way Ctrl+C does under asyncio.
|
||||
loop.call_soon(lambda: None) # ensure loop is processing
|
||||
# Invoke the process signal: raise SIGINT to this process.
|
||||
import os
|
||||
|
||||
os.kill(os.getpid(), signal.SIGINT)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert cancelled.is_set()
|
||||
# Cleanup any leftover reinstall hook from run_cancellable
|
||||
set_sigint_reinstall(None)
|
||||
|
||||
Reference in New Issue
Block a user