Skip to content

fix(agent): honor AddTools/RemoveTools mid-turn at next LLM step (#76)#77

Merged
ezynda3 merged 2 commits into
masterfrom
fix/76-addtools-midturn
Jun 25, 2026
Merged

fix(agent): honor AddTools/RemoveTools mid-turn at next LLM step (#76)#77
ezynda3 merged 2 commits into
masterfrom
fix/76-addtools-midturn

Conversation

@ezynda3

@ezynda3 ezynda3 commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

Description

Kit.AddTools / Kit.RemoveTools document 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 next PromptResult. 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 fantasy Stream call, which captures the tool snapshot taken when Stream began. Swapping the rebuilt fantasy agent mid-turn just reassigned a pointer the running stream never re-read. Kit's PrepareStep callback set result.Messages but never result.Tools — the only per-step tool-override mechanism fantasy exposes.

This PR wires PrepareStep to populate PrepareStepResult.Tools from 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-Stream execution model while honoring the documented "next LLM step" contract.

Fixes #76

Type of Change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Documentation update

Checklist

  • My code follows the style guidelines of this project
  • I have performed a self-review of my own code
  • I have added tests that prove my fix is effective
  • New and existing unit tests pass locally with my changes (go test ./... -race, golangci-lint run)
  • I have updated the documentation where necessary

Additional Information

  • internal/agent/agent.go: Always wire PrepareStep and set result.Tools = composeAllTools() each step. Extracted a composeAllTools() helper (shared with rebuildFantasyAgent) so the per-step set matches the baked-in composition — the steady state is unchanged. Added a sync.RWMutex guarding extraTools so the in-flight Stream can read the live set while a concurrent SetExtraTools mutates it (verified under -race).
  • pkg/kit/kit.go: Fixed the contradictory AddTools doc 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

  • Bug Fixes
    • Improved tool availability during streaming turns so changes to the tool set are reflected on the next step within the same turn.
    • Made tool snapshots safer and more consistent when tools are added or removed while a response is in progress.
    • Ensured steering messages and callback handling are applied in the correct order during step preparation.
  • Documentation
    • Clarified when newly added tools become available during an active turn.

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
@mark-iii-labs-huly

Copy link
Copy Markdown

Connected to Huly®: KIT-78

@coderabbitai

coderabbitai Bot commented Jun 25, 2026

Copy link
Copy Markdown

Review Change Stack

Warning

Review limit reached

@ezynda3, we couldn't start this review because you've reached your PR review rate limit.

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 @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

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.

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 44c1767f-45a6-4a2c-8d0f-513d71fca0ee

📥 Commits

Reviewing files that changed from the base of the PR and between 77ffb63 and a331588.

📒 Files selected for processing (1)
  • internal/agent/agent.go
📝 Walkthrough

Walkthrough

The 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 AddTools comment now describes step-level visibility within an active turn.

Changes

Mid-turn tool recomposition

Layer / File(s) Summary
Live tool snapshot plumbing
internal/agent/agent.go
Adds mutex-protected access to extraTools, centralizes tool composition in composeAllTools, and routes rebuilds/getters/setters through the shared snapshot path.
Per-step tool recomposition
internal/agent/agent.go
Always wires PrepareStep, drains steering messages, calls the callback hook, applies cache control, and fills step tools from the current composed tool list.
Mid-turn tool visibility tests and docs
internal/agent/agent_midturn_tools_test.go, pkg/kit/kit.go
Adds regression tests for tool addition and removal between streaming steps and updates the AddTools comment to describe step-level visibility.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related issues

  • mark3labs/kit issue 74: The agent now re-reads and synchronizes extra tools during streaming steps, matching the runtime tool mutation objective.

Possibly related PRs

  • mark3labs/kit#75: Covers the same live tool mutation flow, including the agent-side composeAllTools and PrepareStep changes.

Poem

A rabbit went hopping through turn after turn,
While tools in the burrow began to return.
With mutex and step by step, bright and new,
I nibbled a carrot and called foo too. 🐰

🚥 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 The title clearly summarizes the main change: honoring mid-turn AddTools/RemoveTools on the next LLM step.
Linked Issues check ✅ Passed The change re-reads the live tool set each step, updates docs, and adds regression tests for mid-turn tool addition/removal.
Out of Scope Changes check ✅ Passed All changes are directly related to tool recomposition, concurrency safety, docs, or tests for the linked mid-turn behavior.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/76-addtools-midturn

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.

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

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

📥 Commits

Reviewing files that changed from the base of the PR and between 821ef11 and 77ffb63.

📒 Files selected for processing (3)
  • internal/agent/agent.go
  • internal/agent/agent_midturn_tools_test.go
  • pkg/kit/kit.go

Comment thread internal/agent/agent.go
Comment thread internal/agent/agent.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
@ezynda3 ezynda3 merged commit d99b76c into master Jun 25, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

AddTools/RemoveTools do not take effect mid-turn (only on next PromptResult), contradicting docs

1 participant