Fix cancelling in-flight LLM operations#175
Conversation
|
Thanks for the pull request. A maintainer will review it when available. Please keep the PR focused, explain the why in the description, and make sure local checks pass before requesting review. Contribution guide: https://github.com/AI-Shell-Team/aish/blob/main/CONTRIBUTING.md |
|
Template check passed. Thanks for updating the pull request description. |
|
Warning Rate limit exceeded
You’ve run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds early cancellation-token checks to LLMSession timeout handlers so actual cancellation emits cancelled events and raises AnyIO cancelled; adds tests verifying TimeoutError still reports status="timeout" and is not marked as cancelled. ChangesLLM timeout & cancellation behavior
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/aish/shell/runtime/ai.py`:
- Around line 169-175: The _cancel_running_task() helper currently calls
loop.call_soon_threadsafe(task.cancel) but may return before cancellation is
delivered and can raise RuntimeError if the loop closes; update
_cancel_running_task() and the _run_async_in_thread cancellation path so you:
(1) guard call_soon_threadsafe with loop.is_running() and catch RuntimeError,
(2) use asyncio.run_coroutine_threadsafe or schedule a small coroutine on the
target loop that cancels the task and awaits task completion, and (3) block the
caller (with a bounded timeout) until task.done() or the awaitable from
run_coroutine_threadsafe finishes, then only return/raise KeyboardInterrupt
after confirming task.cancelled() or task.done(); apply this logic to the
cancellation flows in _cancel_running_task() and the cancellation handling
inside _run_async_in_thread to avoid leaving the coroutine running in
background.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: fb79ec45-81d3-4e6d-935f-0275bfdaddf4
📒 Files selected for processing (3)
src/aish/llm/session.pysrc/aish/shell/runtime/ai.pytests/shell/runtime/test_shell_pty_core.py
| def _cancel_running_task() -> None: | ||
| loop = loop_box[0] | ||
| task = task_box[0] | ||
| if loop is None or task is None or task.done() or loop.is_closed(): | ||
| return | ||
| loop.call_soon_threadsafe(task.cancel) | ||
|
|
There was a problem hiding this comment.
Ensure cancellation is actually delivered before returning cancellation to caller.
On Line 174/Line 186–189, _run_async_in_thread can raise KeyboardInterrupt even if task cancellation was never delivered (loop not ready / future still running), and call_soon_threadsafe can race with loop close and raise RuntimeError. That can leave the coroutine running in the background after the caller thinks it was cancelled.
Suggested hardening
- def _cancel_running_task() -> None:
+ def _cancel_running_task() -> bool:
loop = loop_box[0]
task = task_box[0]
if loop is None or task is None or task.done() or loop.is_closed():
- return
- loop.call_soon_threadsafe(task.cancel)
+ return False
+ try:
+ loop.call_soon_threadsafe(task.cancel)
+ return True
+ except RuntimeError:
+ # loop closed between check and scheduling
+ return False
...
if cancellation_token and cancellation_token.is_cancelled():
- if loop_ready.wait(timeout=0.2):
- _cancel_running_task()
+ cancel_sent = False
+ if loop_ready.wait(timeout=0.2):
+ cancel_sent = _cancel_running_task()
try:
future.result(timeout=2.0)
except FutureTimeoutError:
- pass
+ if cancel_sent and not future.done():
+ continue
raise KeyboardInterrupt("AI operation cancelled by user")Also applies to: 183-189
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aish/shell/runtime/ai.py` around lines 169 - 175, The
_cancel_running_task() helper currently calls
loop.call_soon_threadsafe(task.cancel) but may return before cancellation is
delivered and can raise RuntimeError if the loop closes; update
_cancel_running_task() and the _run_async_in_thread cancellation path so you:
(1) guard call_soon_threadsafe with loop.is_running() and catch RuntimeError,
(2) use asyncio.run_coroutine_threadsafe or schedule a small coroutine on the
target loop that cancels the task and awaits task completion, and (3) block the
caller (with a bounded timeout) until task.done() or the awaitable from
run_coroutine_threadsafe finishes, then only return/raise KeyboardInterrupt
after confirming task.cancelled() or task.done(); apply this logic to the
cancellation flows in _cancel_running_task() and the cancellation handling
inside _run_async_in_thread to avoid leaving the coroutine running in
background.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/aish/llm/session.py`:
- Around line 1777-1784: The except TimeoutError handler that checks
self.cancellation_token and then raises anyio.get_cancelled_exc_class() loses
the original TimeoutError context; capture the caught exception (e.g., "except
TimeoutError as err:") and re-raise the cancellation exception using "raise
anyio.get_cancelled_exc_class() from err" so exception chaining is preserved;
apply the same change to the other identical spot that re-raises the
cancellation exception (the code that calls events.emit_cancelled /
events.emit_generation_end and raises the anyio cancellation class).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 7d517254-0896-4b5b-9cd9-cf4c64582bed
📒 Files selected for processing (2)
src/aish/llm/session.pytests/llm/test_llm_events.py
1 participant
Summary
Ctrl+C取消 AI 请求后，底层 LLM 请求有时仍会在后台继续运行，之后超时时又打印LLM request timed out，容易让用户误以为是网络问题或取消没有生效。Change Type
Scope
User-visible Changes
Ctrl+C取消 AI 请求后，已取消的旧请求不会再延迟打印LLM request timed out。Compatibility
Testing
uv --native-tls run --group dev python -m pytest tests/shell/runtime/test_shell_pty_core.py -q69 passeduv --native-tls run --group dev ruff check src/aish/shell/runtime/ai.py src/aish/llm/session.py tests/shell/runtime/test_shell_pty_core.pypython3 -m compileall -q src/aish/shell/runtime/ai.py src/aish/llm/session.py tests/shell/runtime/test_shell_pty_core.pygit diff --checkAdditional related test run:
uv --native-tls run --group dev python -m pytest tests/llm tests/tools/test_final_answer.py tests/test_ask_user_tool.py -q74 passed, 1 skipped, 1 failedtests/llm/test_litellm_langfuse.py::testshell_integrationChecklist
Summary by CodeRabbit
Bug Fixes
Tests