fix(agent): honor AddTools/RemoveTools mid-turn at next LLM step (#76)#77
Conversation
Runtime AddTools/RemoveTools changes were not visible to an in-flight
turn — they only took effect on the next PromptResult, contradicting the
documented "next LLM step" contract. The whole agentic loop runs inside a
single fantasy Stream call that captures the tool snapshot taken when
Stream began, so swapping the rebuilt fantasy agent mid-turn had no
effect on the running stream.
- Always wire PrepareStep and set PrepareStepResult.Tools from the live
tool set each step, so fantasy re-reads tools per step
- Extract composeAllTools() shared with rebuildFantasyAgent so the
per-step set matches the baked-in composition (steady state unchanged)
- Guard extraTools with a RWMutex so the in-flight Stream can read the
live set while a concurrent SetExtraTools mutates it
- Fix contradictory AddTools doc ("next turn" -> "next LLM step")
- Add regression tests for mid-turn tool addition and removal
Fixes #76
|
Connected to Huly®: KIT-78 |
|
Warning Review limit reached
More reviews will be available in 46 minutes and 6 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits. 🚦 How do rate limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please see our Fair Usage Limits Policy for further information. 📝 WalkthroughWalkthroughThe agent now rebuilds tool lists from shared live state, re-reading extra tools on each streaming step. Mid-turn tool additions and removals are covered by new tests, and the ChangesMid-turn tool recomposition
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@internal/agent/agent.go`:
- Around line 1159-1163: `SetExtraTools` is rebuilding `a.fantasyAgent` without
the same serialization used for other runtime mutations, which can race with
prompt/tool updates and leave a stale snapshot. Update `Agent.SetExtraTools` to
take the existing rebuild protection used around `rebuildFantasyAgent()` and
ensure the rebuild happens under that shared lock, alongside the
`promptMu`-guarded paths, so all `fantasyAgent` rebuilds are serialized
consistently.
- Around line 1160-1162: `SetExtraTools` is storing the caller-owned slice
directly in `a.extraTools`, which aliases external backing storage and can be
mutated outside `toolsMu` protection. Update the `SetExtraTools` method in
`agent.go` to clone the incoming `extraTools` slice before assigning it to
`a.extraTools`, so `composeAllTools()` only sees changes made under the agent’s
lock.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 0f99e150-bea3-4b73-8c15-ba3b1da1120f
📒 Files selected for processing (3)
internal/agent/agent.gointernal/agent/agent_midturn_tools_test.gopkg/kit/kit.go
Address CodeRabbit review on #77: - Clone the incoming slice before storing it so later caller mutations cannot bypass toolsMu and race composeAllTools - Serialize the fantasyAgent rebuild under promptMu, consistent with the other rebuild paths (SetSystemPrompt, SetModel), so concurrent rebuilds can't last-writer-win a stale snapshot
Description
Kit.AddTools/Kit.RemoveToolsdocument that, when a turn is in progress, the change "takes effect starting from the next LLM step." In practice the added/removed tools were not visible to the in-flight turn at all — they only took effect on the nextPromptResult. This broke progressive-disclosure tool loading, where an agent calls a control tool mid-turn and is expected to call the freshly-loaded tools on its next step within the same turn.The root cause is that the entire multi-step agentic loop runs inside a single
fantasyStreamcall, which captures the tool snapshot taken whenStreambegan. Swapping the rebuilt fantasy agent mid-turn just reassigned a pointer the running stream never re-read. Kit'sPrepareStepcallback setresult.Messagesbut neverresult.Tools— the only per-step tool-override mechanism fantasy exposes.This PR wires
PrepareStepto populatePrepareStepResult.Toolsfrom the agent's live composed tool set on every step, so each step re-reads the current tools instead of the start-of-turn snapshot. This keeps the single-Streamexecution model while honoring the documented "next LLM step" contract.Fixes #76
Type of Change
Checklist
go test ./... -race,golangci-lint run)Additional Information
internal/agent/agent.go: Always wirePrepareStepand setresult.Tools = composeAllTools()each step. Extracted acomposeAllTools()helper (shared withrebuildFantasyAgent) so the per-step set matches the baked-in composition — the steady state is unchanged. Added async.RWMutexguardingextraToolsso the in-flightStreamcan read the live set while a concurrentSetExtraToolsmutates it (verified under-race).pkg/kit/kit.go: Fixed the contradictoryAddToolsdoc comment ("next turn" → "next LLM step of the current turn").internal/agent/agent_midturn_tools_test.go(new): Regression tests for mid-turn tool addition and removal. Both fail before the fix and pass after.Backward compatibility: No API or behavior changes for existing callers. In the steady state (no mid-turn tool mutation), the per-step tool list is identical to the previous start-of-turn snapshot.
Summary by CodeRabbit