feat(llm): sub-agent tool loop and spawn test infrastructure (#330) - #339
Conversation
Extract a reusable native tool calling loop with cancel cascade and FIFO mock LLM responses so spawn paths can be integration-tested without network or shell E2E. Closes AI-Shell-Team#330.
📝 WalkthroughWalkthroughAdds a new ChangesSub-agent Spawn and Tool Loop
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Parent as LlmSession (parent)
participant spawn
participant Sub as LlmSession (sub)
participant run_tool_loop_until_done
participant Tool
spawn->>Sub: create_subsession + configure
spawn->>run_tool_loop_until_done: run(ToolLoopConfig, prompt)
run_tool_loop_until_done->>Sub: chat_completion_raw(messages, tool_specs)
Sub-->>run_tool_loop_until_done: assistant text + tool_calls
run_tool_loop_until_done->>Tool: execute(tool_call args)
Tool-->>run_tool_loop_until_done: tool result message
par parent cancellation watch
spawn->>Parent: poll cancellation token
Parent-->>spawn: cancelled
spawn->>Sub: cancel sub token
end
run_tool_loop_until_done-->>spawn: LoopOutcome(status, text, new_messages)
spawn-->>Parent: SpawnResult(text, status)
Possibly related issues
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 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 |
|
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 |
|
This pull request description looks incomplete. Please update the missing sections below before review. Missing items:
|
There was a problem hiding this comment.
🧹 Nitpick comments (4)
crates/aish-llm/src/agents/mock_llm.rs (1)
6-42: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider gating these behind
#[cfg(test)].
mock_text_response/mock_tool_call_responseare fullypub(unliketest_chat_responses/set_test_chat_responsesinsession.rs, which are#[cfg(test)]-gated). Per the stack outline,agents/mod.rsre-exports these from the crate root, meaning test-only scripted-response builders ship in release builds and become part of the public API surface. Since they're only consumed from#[cfg(test)] mod testsblocks intool_loop.rs/spawn.rs, gating them with#[cfg(test)](or atest-utilsfeature) would keep the release API free of test scaffolding while still working, sincecfg(test)items are visible crate-wide duringcargo test.♻️ Suggested gating
+#[cfg(test)] pub fn mock_text_response(text: &str) -> LlmResponse { ... } +#[cfg(test)] pub fn mock_tool_call_response(calls: &[(&str, &str, &str)]) -> LlmResponse { ... }🤖 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 `@crates/aish-llm/src/agents/mock_llm.rs` around lines 6 - 42, Gate the test-only helpers in mock_llm.rs behind #[cfg(test)] so they don’t ship in the release API surface: mock_text_response and mock_tool_call_response are currently fully public and re-exported through agents/mod.rs, but they’re only used from test modules like tool_loop.rs and spawn.rs. Update the definitions (and any re-exports if needed) so these scripted-response builders remain available to cargo test while staying out of non-test builds, consistent with session.rs test-only helpers.crates/aish-llm/src/lib.rs (1)
38-41: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueGeneric top-level
spawnre-export.Re-exporting a bare
spawnfunction name at the crate root risks ambiguity for consumers who glob-import alongsidetokio::spawn/std::thread::spawn. Consider re-exporting only viaaish_llm::agents::spawn(already available) and dropping it from the crate-root re-export list, or keep both but be aware of the naming collision risk.🤖 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 `@crates/aish-llm/src/lib.rs` around lines 38 - 41, The crate-root re-export list in aish_llm::lib currently exposes a bare spawn symbol, which can collide with common imports like tokio::spawn or std::thread::spawn. Update the top-level pub use in the agents re-export block to remove spawn from the crate root while keeping it available through aish_llm::agents::spawn, and leave the other re-exports unchanged.crates/aish-llm/src/agents/mod.rs (1)
5-9: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider gating mock LLM test helpers behind
#[cfg(test)].
mock_text_response/mock_tool_call_responseare test-seam helpers per the issue (mock LLM injection for spawn integration tests), but re-exporting them unconditionally makes them part of the crate's always-compiled public API. Ifmock_llm.rsdoesn't itself gate these with#[cfg(test)]or atest-utilsfeature, they'll ship in release builds and become a de facto stable API surface that downstream consumers could depend on.♻️ Suggested gating
-mod mock_llm; +#[cfg(test)] +mod mock_llm; mod spawn; mod tool_loop; -pub use mock_llm::{mock_text_response, mock_tool_call_response}; +#[cfg(test)] +pub use mock_llm::{mock_text_response, mock_tool_call_response};Please confirm mock_llm.rs's visibility/attributes to verify whether this gating is already applied there.
🤖 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 `@crates/aish-llm/src/agents/mod.rs` around lines 5 - 9, The `mock_text_response` and `mock_tool_call_response` helpers are being re-exported from `agents::mod` unconditionally, which exposes test-only mock LLM APIs in normal builds. Verify how `mock_llm.rs` is annotated, and if it is not already gated, add `#[cfg(test)]` (or an equivalent test-utils feature gate) to the `mock_llm` module and its re-export so these helpers are only available to the spawn integration tests. Keep the public API in `agents::mod` limited to non-test symbols like `spawn` and `tool_loop`.crates/aish-llm/src/agents/spawn.rs (1)
24-67: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
SpawnResultdiscards fatal error and produced messages.On a
Fataloutcome,LoopOutcome::error(anAishError) is dropped andtextis empty, so a spawn caller cannot tell why the sub-agent failed.new_messagesis likewise lost. Consider surfacing at least the error so callers can distinguishFatalfrom an empty completion.♻️ Proposed change
pub struct SpawnResult { pub text: String, pub status: LoopStatus, + pub error: Option<aish_core::AishError>, }SpawnResult { text: outcome.text, status: outcome.status, + error: outcome.error, }🤖 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 `@crates/aish-llm/src/agents/spawn.rs` around lines 24 - 67, `spawn` currently turns a `Fatal` LoopOutcome from `run_tool_loop_until_done` into a `SpawnResult` that loses the underlying `AishError` and any `new_messages`, making failures indistinguishable from empty success. Update `SpawnResult` and the `spawn` flow so the fatal error is preserved and returned or exposed to callers, and make sure the `LoopOutcome` handling in `spawn`/`SpawnResult` distinguishes `Fatal` from normal completion while still carrying the generated messages where needed.
🤖 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.
Nitpick comments:
In `@crates/aish-llm/src/agents/mock_llm.rs`:
- Around line 6-42: Gate the test-only helpers in mock_llm.rs behind
#[cfg(test)] so they don’t ship in the release API surface: mock_text_response
and mock_tool_call_response are currently fully public and re-exported through
agents/mod.rs, but they’re only used from test modules like tool_loop.rs and
spawn.rs. Update the definitions (and any re-exports if needed) so these
scripted-response builders remain available to cargo test while staying out of
non-test builds, consistent with session.rs test-only helpers.
In `@crates/aish-llm/src/agents/mod.rs`:
- Around line 5-9: The `mock_text_response` and `mock_tool_call_response`
helpers are being re-exported from `agents::mod` unconditionally, which exposes
test-only mock LLM APIs in normal builds. Verify how `mock_llm.rs` is annotated,
and if it is not already gated, add `#[cfg(test)]` (or an equivalent test-utils
feature gate) to the `mock_llm` module and its re-export so these helpers are
only available to the spawn integration tests. Keep the public API in
`agents::mod` limited to non-test symbols like `spawn` and `tool_loop`.
In `@crates/aish-llm/src/agents/spawn.rs`:
- Around line 24-67: `spawn` currently turns a `Fatal` LoopOutcome from
`run_tool_loop_until_done` into a `SpawnResult` that loses the underlying
`AishError` and any `new_messages`, making failures indistinguishable from empty
success. Update `SpawnResult` and the `spawn` flow so the fatal error is
preserved and returned or exposed to callers, and make sure the `LoopOutcome`
handling in `spawn`/`SpawnResult` distinguishes `Fatal` from normal completion
while still carrying the generated messages where needed.
In `@crates/aish-llm/src/lib.rs`:
- Around line 38-41: The crate-root re-export list in aish_llm::lib currently
exposes a bare spawn symbol, which can collide with common imports like
tokio::spawn or std::thread::spawn. Update the top-level pub use in the agents
re-export block to remove spawn from the crate root while keeping it available
through aish_llm::agents::spawn, and leave the other re-exports unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: f11c40b0-9059-40d9-8f4d-e9f404120a57
📒 Files selected for processing (6)
crates/aish-llm/src/agents/mock_llm.rscrates/aish-llm/src/agents/mod.rscrates/aish-llm/src/agents/spawn.rscrates/aish-llm/src/agents/tool_loop.rscrates/aish-llm/src/lib.rscrates/aish-llm/src/session.rs
Summary
aish-llm::agentswithrun_tool_loop_until_done(Complete / Incomplete / Cancelled / Fatal) for reusable native tool calling loops.spawn()with parent→child cancel cascade and FIFO mock LLM response injection for integration tests without network.Agenttool yet;/diagnoseandsystem_diagnose_agentunchanged.Closes #330.
Test plan
cargo test -p aish-llm(205 tests)make format-check && make linttest_spawn_cascades_parent_cancel, cancel while runningSummary by CodeRabbit
New Features
Bug Fixes
Tests