Agent loop: subagents for heavy work + safe-point message queue - #2335
Agent loop: subagents for heavy work + safe-point message queue#2335jaylfc wants to merge 1 commit into
Conversation
|
ⓘ Qodo reviews are paused because the subscription is no longer active. Ask your workspace admin to reactivate the subscription to resume reviews. Manage billing |
📝 WalkthroughWalkthroughAdds ChangesAgentLoop orchestration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant MainChatLoop
participant AgentLoop
participant SubagentWorker
participant ProgressSink
MainChatLoop->>AgentLoop: handle_message(content)
AgentLoop->>SubagentWorker: spawn worker
SubagentWorker->>ProgressSink: forward progress
MainChatLoop->>AgentLoop: reach_safe_point()
AgentLoop->>SubagentWorker: cancel if redirect is queued
AgentLoop-->>MainChatLoop: return delivered messages and status
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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 |
| if any(m.is_redirect for m in queued): | ||
| await self.cancel_subagents(reason="redirect at safe point") | ||
|
|
||
| async with self._lock: |
There was a problem hiding this comment.
[WARNING]: reach_safe_point releases the lock between SAFE_POINT and IDLE transitions. During this window, handle_message can acquire the lock and start a new turn (treating SAFE_POINT as idle), potentially corrupting state or losing queued messages.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| entry.task_obj.cancel() | ||
| # Wait outside the lock so the loop isn't held while workers unwind. | ||
| if to_cancel: | ||
| _done, pending = await asyncio.wait(to_cancel, timeout=10.0) |
There was a problem hiding this comment.
[WARNING]: cancel_subagents leaves tasks pending after the 10s timeout without forcefully terminating them. A worker that swallows CancelledError will run forever, leaking resources.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| if e.task_obj is not None and not e.task_obj.done() | ||
| ] | ||
| if tasks: | ||
| await asyncio.wait_for( |
There was a problem hiding this comment.
[WARNING]: await_all_subagents does not cancel individual tasks when asyncio.wait_for times out. The gather coroutine is cancelled, but the underlying tasks continue running indefinitely.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
|
||
| Safe to call from any context (no I/O, no await). | ||
| """ | ||
| subagents: list[dict[str, Any]] = [] |
There was a problem hiding this comment.
[SUGGESTION]: status() reads shared state (_state, _subagents, _message_queue, _current_turn_id) without acquiring self._lock, which could return an inconsistent snapshot during concurrent mutations.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 4 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
SUGGESTION
Files Reviewed (1 file)
Fix these issues in Kilo Cloud Reviewed by step-3.7-flash · Input: 56.5K · Output: 20.2K · Cached: 301.2K |
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
docs/design/agent-loop-subagents.md (1)
64-73: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd a language to the fenced code block.
markdownlint reports MD040 for this block. Use
textfor the state diagram.-``` +```text IDLE🤖 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 `@docs/design/agent-loop-subagents.md` around lines 64 - 73, Update the fenced code block containing the IDLE/WORKING/SAFE_POINT state diagram to specify the text language, changing the opening fence to ```text while preserving the diagram content.Source: Linters/SAST tools
tests/test_agent_loop.py (2)
312-360: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover the async-sink and timeout paths.
Two behaviours have no coverage:
_progressschedules a task when the sink returns a coroutine (tinyagentos/agent_loop.pylines 247-252)._RecordingSink.__call__is sync, so this branch never runs.await_subagent(sub_id, timeout=...)andawait_all_subagents(timeout=...)are never called with a timeout. Both currently cancel the subagents on timeout; see the comments on those methods.Add a coroutine sink and a short-timeout test to lock in the intended behaviour.
🤖 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 `@tests/test_agent_loop.py` around lines 312 - 360, Extend the tests around AgentLoop with an async sink whose __call__ returns a coroutine, then verify _progress schedules and delivers progress through that sink. Add short-timeout coverage for await_subagent and await_all_subagents using workers that do not finish promptly, asserting the timeout path cancels the relevant subagents and preserves the documented state.
108-211: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a regression test for a message that arrives during
SAFE_POINT.The suite covers
IDLEandWORKINGintake, but not the window in whichreach_safe_pointhas setSAFE_POINTand is awaitingcancel_subagents. That window is where the queue clear and the turn-id clobber occur; see the comment ontinyagentos/agent_loop.pylines 202-210. A test that sends a message while a slow-unwinding subagent is being cancelled would pin the fixed behaviour.🤖 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 `@tests/test_agent_loop.py` around lines 108 - 211, Add an async regression test covering message intake while AgentLoop.reach_safe_point() is in SAFE_POINT and awaiting cancel_subagents(). Use a slow-unwinding subagent whose cancellation can be observed, send a message during that window, then await the safe point and assert the message is delivered with its original turn association while the queue is not cleared or the turn ID clobbered.
🤖 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 `@tinyagentos/agent_loop.py`:
- Around line 404-416: The await API at tinyagentos/agent_loop.py lines 404-416
must stop cancelling live subagent tasks on timeout: resolve entry.task_obj to a
local asyncio.Task, use asyncio.wait({task_obj}, timeout=timeout), and raise
asyncio.TimeoutError when it is not done, preserving cancellation handling and
result retrieval. Apply the same non-cancelling asyncio.wait approach to lines
433-449, replacing the wait_for/gather pair and raising asyncio.TimeoutError
when pending is non-empty.
- Around line 202-210: Update the message-handling branch around the state
transition to treat SAFE_POINT like WORKING, appending the message and returning
QUEUED instead of starting a turn. Remove the unconditional
_message_queue.clear() from this path, and make reach_safe_point’s final
IDLE/None transition conditional so it does not overwrite a turn started after
queue draining.
- Around line 245-254: Update the progress sink handling around self._sink(msg)
to retain each scheduled coroutine task in a persistent set and remove it via a
completion callback after it finishes. Ensure the set is owned by the agent loop
instance and initialized appropriately, while preserving the existing
best-effort scheduling and exception logging behavior.
---
Nitpick comments:
In `@docs/design/agent-loop-subagents.md`:
- Around line 64-73: Update the fenced code block containing the
IDLE/WORKING/SAFE_POINT state diagram to specify the text language, changing the
opening fence to ```text while preserving the diagram content.
In `@tests/test_agent_loop.py`:
- Around line 312-360: Extend the tests around AgentLoop with an async sink
whose __call__ returns a coroutine, then verify _progress schedules and delivers
progress through that sink. Add short-timeout coverage for await_subagent and
await_all_subagents using workers that do not finish promptly, asserting the
timeout path cancels the relevant subagents and preserves the documented state.
- Around line 108-211: Add an async regression test covering message intake
while AgentLoop.reach_safe_point() is in SAFE_POINT and awaiting
cancel_subagents(). Use a slow-unwinding subagent whose cancellation can be
observed, send a message during that window, then await the safe point and
assert the message is delivered with its original turn association while the
queue is not cleared or the turn ID clobbered.
🪄 Autofix
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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 1b2465c2-5068-449a-86e8-a5b204e26616
📒 Files selected for processing (4)
changelog.d/tsk-rl2lfb-agent-loop.mddocs/design/agent-loop-subagents.mdtests/test_agent_loop.pytinyagentos/agent_loop.py
| async with self._lock: | ||
| if self._state == LoopState.WORKING: | ||
| self._message_queue.append(msg) | ||
| return LoopAction.QUEUED | ||
| # IDLE (or briefly SAFE_POINT) -> start a new turn. | ||
| self._state = LoopState.WORKING | ||
| self._current_turn_id = msg_id | ||
| self._message_queue.clear() | ||
| return LoopAction.IMMEDIATE |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
Fix the state-clobber race and the destructive queue clear.
reach_safe_point releases self._lock between line 332 and line 339 while cancel_subagents runs, which can take up to 10 seconds. A message that arrives in that window sees SAFE_POINT, falls into this branch, returns IMMEDIATE, and sets _state = WORKING plus _current_turn_id = msg_id. reach_safe_point then reacquires the lock and overwrites both with IDLE/None. The caller drives a turn that the loop no longer tracks.
Line 209 also clears the queue unconditionally. Any message appended in that same window is discarded, which breaks the "never dropped" guarantee stated in the module docstring and in changelog.d/tsk-rl2lfb-agent-loop.md.
Two changes are needed: treat SAFE_POINT as busy (queue the message) and remove the blanket clear(), since reach_safe_point already drains the queue.
🐛 Proposed fix
async with self._lock:
- if self._state == LoopState.WORKING:
+ if self._state in (LoopState.WORKING, LoopState.SAFE_POINT):
self._message_queue.append(msg)
return LoopAction.QUEUED
- # IDLE (or briefly SAFE_POINT) -> start a new turn.
+ # IDLE -> start a new turn. The queue is already drained by
+ # reach_safe_point, so nothing is discarded here.
self._state = LoopState.WORKING
self._current_turn_id = msg_id
- self._message_queue.clear()
return LoopAction.IMMEDIATEreach_safe_point must then avoid overwriting a turn that a queued message starts; keep the final transition conditional:
async with self._lock:
- self._state = LoopState.IDLE
- self._current_turn_id = None
+ if self._state == LoopState.SAFE_POINT:
+ self._state = LoopState.IDLE
+ self._current_turn_id = None
self._delivered.extend(queued)🤖 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 `@tinyagentos/agent_loop.py` around lines 202 - 210, Update the
message-handling branch around the state transition to treat SAFE_POINT like
WORKING, appending the message and returning QUEUED instead of starting a turn.
Remove the unconditional _message_queue.clear() from this path, and make
reach_safe_point’s final IDLE/None transition conditional so it does not
overwrite a turn started after queue draining.
| try: | ||
| res = self._sink(msg) | ||
| if asyncio.iscoroutine(res): | ||
| # Best-effort schedule; the sink owns its lifecycle. | ||
| try: | ||
| asyncio.get_running_loop().create_task(res) | ||
| except RuntimeError: | ||
| pass | ||
| except Exception: | ||
| logger.exception("subagent %s: progress sink raised", sub_id) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
asyncio create_task strong reference required task garbage collected
💡 Result:
Yes, it is necessary to maintain a strong reference to a Task created by asyncio.create_task to prevent it from being garbage collected prematurely [1][2][3]. The asyncio event loop maintains only weak references to scheduled tasks [1][4]. If a task is not referenced elsewhere in your code—meaning it has no strong references—the Python garbage collector may destroy the task object at any time, even while it is still pending or suspended [1][5]. If this happens, you may encounter an error message such as "Task was destroyed but it is pending!" [2][6][7]. To reliably manage background tasks (often called "fire-and-forget" tasks), you should store them in a collection, such as a set, to keep a strong reference [1][4]. It is also standard practice to remove the task from that collection once it has finished to avoid memory leaks [1][4]. Example of recommended practice: background_tasks = set task = asyncio.create_task(some_coroutine) # Keep a strong reference background_tasks.add(task) # Remove the reference automatically when the task finishes task.add_done_callback(background_tasks.discard) For more structured concurrency in modern Python (3.11+), it is highly recommended to use asyncio.TaskGroup instead, which automatically manages these references and handles task lifetimes safely [1].
Citations:
- 1: https://docs.python.org/3/library/asyncio-task.html
- 2: asyncio.create_task() documentation should mention user needs to keep reference to the task python/cpython#88831
- 3: asyncio: Use strong references for free-flying tasks python/cpython#91887
- 4: https://docs.python.org/release/3.12.1/library/asyncio-task.html
- 5: https://docs.python.org/3.10/library/asyncio-task.html
- 6: https://stackoverflow.com/questions/71938799/python-asyncio-create-task-really-need-to-keep-a-reference
- 7: https://bugs.python.org/issue44665
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate and inspect the relevant source and class structure without running repository code.
if [ -f tinyagentos/agent_loop.py ]; then
echo "---- agent_loop.py outline ----"
ast-grep outline tinyagentos/agent_loop.py --view compact || true
echo
echo "---- relevant source lines ----"
sed -n '1,120p;210,285p' tinyagentos/agent_loop.py | nl -ba
else
echo "tinyagentos/agent_loop.py not found"
fd -a agent_loop.py .
fi
echo
echo "---- sink references in agent_loop.py ----"
rg -n "_sink|_progress|sink_tasks|TaskGroup|run_in_executor|create_task|_runner|cancel_subagents" tinyagentos/agent_loop.py || true
echo
echo "---- Python asyncio task GC behavior probe ----"
python3 - <<'PY'
import asyncio, gc, weakref, sys
def warn_handler(*args, **kw):
print("WARNING:", args[0], file=sys.stderr)
def run_sink_with_progress_loop(sink_name):
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
sink_called = False
async def sink():
nonlocal sink_called
sink_called = True
async def runner():
loop.create_task(sink())
asyncio.get_running_loop().stop()
loop.run_coroutine_threadsafe(runner(), loop).result(timeout=5)
finally:
loop.close()
return sink_called
def run_sink_without_progress_loop(sink_name):
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
sink_called = False
async def sink():
nonlocal sink_called
sink_called = True
async def runner():
# Drop reference to task immediately; force GC while message is in sink queue.
t = asyncio.get_running_loop().create_task(sink())
del t
gc.collect()
await asyncio.sleep(0.01)
asyncio.get_running_loop().stop()
loop.run_coroutine_threadsafe(runner(), loop).result(timeout=5)
finally:
loop.close()
return sink_called
import warnings
warnings.filterwarnings("always", category=PendingDeprecationWarning)
warnings.filterwarnings("always", category=RuntimeWarning)
print("task with progress loop referenced:", run_sink_with_progress_loop("referenced"))
print("task without progress loop unreferenced:", run_sink_without_progress_loop("unreferenced"))
PYRepository: jaylfc/taOS
Length of output: 406
🏁 Script executed:
#!/bin/bash
set -euo pipefail
if [ -f tinyagentos/agent_loop.py ]; then
echo "---- agent_loop.py class/method identifiers ----"
ast-grep outline tinyagentos/agent_loop.py || true
echo
echo "---- relevant source lines ----"
sed -n '1,140p;210,285p' tinyagentos/agent_loop.py
else
echo "tinyagentos/agent_loop.py not found"
fd -a agent_loop.py .
fi
echo
echo "---- sink references in agent_loop.py ----"
rg -n "_sink|_progress|sink_tasks|TaskGroup|run_in_executor|create_task|_runner|cancel_subagents" tinyagentos/agent_loop.py || true
echo
echo "---- Python asyncio task GC behavior probe ----"
python3 - <<'PY'
import asyncio, gc, sys
def run_sink_with_progress_loop():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
called = False
async def sink():
nonlocal called
called = True
async def runner():
loop.create_task(sink())
asyncio.get_running_loop().stop()
try:
asyncio.run_coroutine_threadsafe(runner(), loop).result(timeout=5)
finally:
loop.close()
return called
def run_sink_without_progress_loop():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
called = False
async def sink():
nonlocal called
called = True
async def runner():
t = asyncio.get_running_loop().create_task(sink())
ref_count_before = sys.getrefcount(t)
del t
gc.collect()
ref_count_after = sys.getrefcount(t)
await asyncio.sleep(0.01)
asyncio.get_running_loop().stop()
try:
asyncio.run_coroutine_threadsafe(runner(), loop).result(timeout=5)
finally:
loop.close()
return called
print("task with progress loop referenced:", run_sink_with_progress_loop())
print("task without progress loop unreferenced:", run_sink_without_progress_loop())
PYRepository: jaylfc/taOS
Length of output: 10403
🌐 Web query:
Python 3.9+ Task __del__ warning Pending tasks no garbage collected asyncio garbage collection
💡 Result:
The warning "Task was destroyed but it is pending!" occurs because asyncio maintains only weak references to scheduled tasks [1][2]. If you create a task (e.g., using asyncio.create_task) but do not keep a strong reference to it, the Python garbage collector may destroy the task object while it is still pending or suspended [2][3][4]. This typically happens when the task is waiting on a future or coroutine that is not referenced elsewhere, causing the task to become unreachable to the garbage collector [5][6][3]. To prevent this issue and the associated warning, you must maintain a strong reference to the task until it completes [2][4]. Recommended approaches include: 1. Use asyncio.TaskGroup (Recommended for Python 3.11+): TaskGroup automatically manages the lifecycle of tasks, ensuring they are properly awaited and exceptions are propagated [2][3]. async with asyncio.TaskGroup as tg: tg.create_task(some_coro) 2. Maintain a collection: If you need to manage fire-and-forget tasks, store them in a collection (like a set) and remove them using a done callback to avoid memory leaks [2][7]. background_tasks = set task = asyncio.create_task(some_coro) background_tasks.add(task) task.add_done_callback(background_tasks.discard) This behavior is fundamental to how asyncio's task scheduling interacts with Python's memory management [3][4]. While internal changes have been proposed to make this more robust, keeping a strong reference remains the standard, supported practice [8][9].
Citations:
- 1: https://bugs.python.org/issue44665
- 2: https://docs.python.org/3/library/asyncio-task.html
- 3: https://stackoverflow.com/questions/71938799/python-asyncio-create-task-really-need-to-keep-a-reference
- 4: https://superfastpython.com/asyncio-disappearing-task-bug/
- 5: https://stackoverflow.com/questions/44304603/python-asyncio-unreferenced-tasks-are-destroyed-by-garbage-collector
- 6: asyncio doesn't warn if a task is destroyed during its execution python/cpython#65362
- 7: https://docs.python.org/3.10/library/asyncio-task.html
- 8: asyncio: Use strong references for free-flying tasks python/cpython#91887
- 9: asyncio.create_task() documentation should mention user needs to keep reference to the task python/cpython#88831
Keep the sink task alive until it finishes.
asyncio.create_task stores the task with only a weak reference. The task created at line 250 is not referenced elsewhere, so the garbage collector can drop it while it is pending; the sink coroutine then may not run and the progress write is lost. Keep the task in a set and discard it when it completes.
🤖 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 `@tinyagentos/agent_loop.py` around lines 245 - 254, Update the progress sink
handling around self._sink(msg) to retain each scheduled coroutine task in a
persistent set and remove it via a completion callback after it finishes. Ensure
the set is owned by the agent loop instance and initialized appropriately, while
preserving the existing best-effort scheduling and exception logging behavior.
| async with self._lock: | ||
| entry = self._subagents.get(sub_id) | ||
| if entry is None: | ||
| raise KeyError(sub_id) | ||
| try: | ||
| await asyncio.wait_for(entry.task_obj, timeout=timeout) | ||
| except asyncio.CancelledError: | ||
| # Distinguish "the subagent itself was cancelled" from "this | ||
| # caller was cancelled". Only swallow the former. | ||
| if entry.handle.state == "cancelled": | ||
| return entry.handle.result | ||
| raise | ||
| return entry.handle.result |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Timeouts cancel the subagents in both await APIs. asyncio.wait_for cancels the future it wraps when the timeout expires. Both methods pass live subagent tasks to it, so a timed-out observation destroys the work instead of only ending the wait. This contradicts the docstring of each method.
tinyagentos/agent_loop.py#L404-L416: replaceasyncio.wait_for(entry.task_obj, timeout=timeout)withasyncio.wait({task_obj}, timeout=timeout)and raiseasyncio.TimeoutErrorwhen the task is not in the done set. Also resolveentry.task_objinto a local first, because its type isasyncio.Task | None.tinyagentos/agent_loop.py#L433-L449: replace thewait_for/gatherpair withasyncio.wait(tasks, timeout=timeout)and raiseasyncio.TimeoutErrorwhenpendingis non-empty.
📍 Affects 1 file
tinyagentos/agent_loop.py#L404-L416(this comment)tinyagentos/agent_loop.py#L433-L449
🤖 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 `@tinyagentos/agent_loop.py` around lines 404 - 416, The await API at
tinyagentos/agent_loop.py lines 404-416 must stop cancelling live subagent tasks
on timeout: resolve entry.task_obj to a local asyncio.Task, use
asyncio.wait({task_obj}, timeout=timeout), and raise asyncio.TimeoutError when
it is not done, preserving cancellation handling and result retrieval. Apply the
same non-cancelling asyncio.wait approach to lines 433-449, replacing the
wait_for/gather pair and raising asyncio.TimeoutError when pending is non-empty.
|
BLOCKING, two independent blockers. The library core is genuinely good — real un-mocked tests, no hallucinated APIs, clean cancel-propagation — which is why this is worth fixing forward rather than closing.
Test gaps sit exactly on the bugs: no test for a message arriving during the safe-point/cancel window, none for Bot adjudication: CodeRabbit's Critical (the race) and Majors (wait_for, GC'd task) are all correct, verified independently. Kilo's force-terminate and status()-lock findings are noise. Fix-forward card follows; do not open a new PR. |
Fix-forward #2335: safe-point race + wait_for cancellation + task_utils reuse + wiring design note
|
Superseded by #2337 (merged): same four files rebuilt with the safe-point race fixed (red-proven against this PR's code), non-destructive await semantics, task_utils reuse, and an honest changelog. Wiring is deliberately deferred to the AgentLoop/AgentChatRouter design decision. |
CARD TITLE (intent, not commit subject): Agent loop: subagents for heavy work + safe-point message queue
Autonomous build of board card tsk-rl2lfb.
Files:
changelog.d/tsk-rl2lfb-agent-loop.md | 10 +
docs/design/agent-loop-subagents.md | 95 ++++++++
tests/test_agent_loop.py | 456 +++++++++++++++++++++++++++++++++++
tinyagentos/agent_loop.py | 449 ++++++++++++++++++++++++++++++++++
4 files changed, 1010 insertions(+)
Summary by CodeRabbit
New Features
Documentation
Tests