Skip to content

feat(llm): sub-agent tool loop and spawn test infrastructure (#330) - #339

Merged
F16shen merged 1 commit into
AI-Shell-Team:mainfrom
F16shen:feat/subagent-tool-loop-spawn-infra
Jul 7, 2026
Merged

feat(llm): sub-agent tool loop and spawn test infrastructure (#330)#339
F16shen merged 1 commit into
AI-Shell-Team:mainfrom
F16shen:feat/subagent-tool-loop-spawn-infra

Conversation

@F16shen

@F16shen F16shen commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add aish-llm::agents with run_tool_loop_until_done (Complete / Incomplete / Cancelled / Fatal) for reusable native tool calling loops.
  • Add spawn() with parent→child cancel cascade and FIFO mock LLM response injection for integration tests without network.
  • No user-visible Agent tool yet; /diagnose and system_diagnose_agent unchanged.

Closes #330.

Test plan

Summary by CodeRabbit

  • New Features

    • Added support for running sub-agent workflows with configurable turn limits and optional system prompts.
    • Introduced reusable helpers for building scripted assistant responses, including plain text replies and tool-call outputs.
    • Expanded the public API to expose agent spawning and tool-loop utilities.
  • Bug Fixes

    • Improved cancellation handling so child agent runs stop promptly when the parent session is cancelled.
  • Tests

    • Added coverage for tool-call parsing, completion flows, turn-limit behavior, and cancellation scenarios.

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.
@coderabbitai

coderabbitai Bot commented Jul 7, 2026

Copy link
Copy Markdown

Review Change Stack

📝 Walkthrough

Walkthrough

Adds a new agents module to the aish-llm crate implementing a reusable native tool-calling loop (run_tool_loop_until_done) and a spawn function for sub-agent sessions with parent-to-child cancellation cascade. Adds mock LLM response helpers and session-level scripted test response injection.

Changes

Sub-agent Spawn and Tool Loop

Layer / File(s) Summary
Mock LLM response helpers
crates/aish-llm/src/agents/mock_llm.rs
Adds `mock_text_response` and `mock_tool_call_response` for scripted non-streaming LLM JSON payloads, plus a unit test validating parsed tool call output.
Session test response queue and accessors
crates/aish-llm/src/session.rs
Adds a test-only FIFO queue of scripted `chat_completion` results, `set_test_chat_responses`, `record_usage_public`, `loop_temperature`/`loop_max_tokens` accessors, and widens `prepare_messages_for_send` visibility to `pub(crate)`.
Tool loop contract and implementation
crates/aish-llm/src/agents/tool_loop.rs
Defines `ToolLoopConfig`, `LoopStatus`, `LoopOutcome`, and implements `run_tool_loop_until_done` handling message building, tool execution, usage recording, and termination (complete/incomplete/cancelled/fatal), with unit tests.
Spawn function and cancellation cascade
crates/aish-llm/src/agents/spawn.rs
Adds `SpawnConfig`, `SpawnResult`, and `spawn` async function configuring a sub-session and running the tool loop while racing a `forward_cancellation` task to cascade parent cancellation, with tests for completion and cancellation.
Module wiring and crate re-exports
crates/aish-llm/src/agents/mod.rs, crates/aish-llm/src/lib.rs
Declares the `agents` module's submodules and re-exports, and exposes agents' public APIs from the crate root.

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)
Loading

Possibly related issues

Poem

A rabbit taps a looping beat,
Tool calls hopping, turn complete,
Cancel tokens cascade down the burrow,
Sub-sessions spawn without a worry,
Mocked replies, no network chase—
🐇 tests all green, a tidy place!

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed It clearly summarizes the new sub-agent tool loop, spawn infrastructure, and test support.
Linked Issues check ✅ Passed The PR adds the reusable tool loop, parent-cancel cascade, and mock LLM scripting required by #330.
Out of Scope Changes check ✅ Passed The changes stay within #330's spawn/tool-loop and test-seam scope, with no unrelated features added.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

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

@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

This pull request description looks incomplete. Please update the missing sections below before review.

Missing items:

  • User-visible Changes
  • Compatibility
  • Testing
  • Change Type
  • Scope

@F16shen
F16shen marked this pull request as ready for review July 7, 2026 09:12

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (4)
crates/aish-llm/src/agents/mock_llm.rs (1)

6-42: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider gating these behind #[cfg(test)].

mock_text_response/mock_tool_call_response are fully pub (unlike test_chat_responses/set_test_chat_responses in session.rs, which are #[cfg(test)]-gated). Per the stack outline, agents/mod.rs re-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 tests blocks in tool_loop.rs/spawn.rs, gating them with #[cfg(test)] (or a test-utils feature) would keep the release API free of test scaffolding while still working, since cfg(test) items are visible crate-wide during cargo 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 value

Generic top-level spawn re-export.

Re-exporting a bare spawn function name at the crate root risks ambiguity for consumers who glob-import alongside tokio::spawn/std::thread::spawn. Consider re-exporting only via aish_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 win

Consider gating mock LLM test helpers behind #[cfg(test)].

mock_text_response/mock_tool_call_response are 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. If mock_llm.rs doesn't itself gate these with #[cfg(test)] or a test-utils feature, 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

SpawnResult discards fatal error and produced messages.

On a Fatal outcome, LoopOutcome::error (an AishError) is dropped and text is empty, so a spawn caller cannot tell why the sub-agent failed. new_messages is likewise lost. Consider surfacing at least the error so callers can distinguish Fatal from 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

📥 Commits

Reviewing files that changed from the base of the PR and between 0acf441 and a9c6ea7.

📒 Files selected for processing (6)
  • crates/aish-llm/src/agents/mock_llm.rs
  • crates/aish-llm/src/agents/mod.rs
  • crates/aish-llm/src/agents/spawn.rs
  • crates/aish-llm/src/agents/tool_loop.rs
  • crates/aish-llm/src/lib.rs
  • crates/aish-llm/src/session.rs

@F16shen
F16shen merged commit 854697c into AI-Shell-Team:main Jul 7, 2026
9 checks passed
@F16shen
F16shen deleted the feat/subagent-tool-loop-spawn-infra branch July 7, 2026 09:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature]: Sub-agent Phase 1 — 抽取 tool loop 与 spawn 测试基建

1 participant