feat(hooks): expose per-step ToolChoice via OnPrepareStep

[giltine]

Motivation

fantasy supports per-step tool-choice control (PrepareStepResult.ToolChoice,
applied in agent.go step loop), but kit's public OnPrepareStep hook only
exposes message replacement — the internal PrepareStep closure builds the
fantasy result with Model/Messages/Tools and never threads a tool choice.

This gap matters for structured-output / capture-tool patterns: an agent whose
turn must end with a specific tool call (e.g. record_result) can currently
only be prompted to call it, and models occasionally finish a turn without
doing so. Per-step forcing closes that gap: force SpecificToolChoice(tool)
while the capture hasn't happened, release once it has.

Per-step (rather than session-level) control is deliberate — a turn only ends
when the model stops emitting tool calls, so a permanently forced choice can
never legally finish a turn. The hook is the right granularity.

Changes

• pkg/kit/types.go — re-export fantasy.ToolChoice as LLMToolChoice
  with the None/Auto/Required constants and an LLMSpecificToolChoice
  constructor var, all following the existing LLM-prefixed re-export convention.
• pkg/kit/hooks.go — add ToolChoice *LLMToolChoice to
  PrepareStepResult (nil = provider default, unchanged behaviour). Doc
  comment warns about the always-forced footgun.
• internal/agent/agent.go — PrepareStepHandler now returns a
  *PrepareStepUpdate struct (Messages + ToolChoice, nil fields = no change)
  instead of bare replacement messages; the PrepareStep closure threads
  ToolChoice into fantasy's PrepareStepResult.
• pkg/kit/kit.go — the hook→handler bridge forwards ToolChoice and
  returns a non-nil update when either field is set.

Compatibility

• Public API: additive only. PrepareStepResult gains a field; all existing
  construction uses named fields. Existing hooks returning Messages-only
  results behave identically.
• Internal: PrepareStepHandler signature changes, but it is module-internal
  with a single producer (kit.go bridge) and single consumer (agent closure),
  both updated.
• Extensions bridge (extensions_bridge.go) untouched — extension-facing
  ToolChoice exposure can follow separately if wanted.

Tests

• internal/agent/agent_toolchoice_test.go — fake fantasy.Agent records the
  ToolChoice each step's PrepareStep yields: forced on step 0, released on
  step 1; nil throughout when no handler is registered. Mirrors the
  agent_midturn_tools_test.go harness.
• pkg/kit/hooks_test.go — registry-level test for a ToolChoice-only result
  (Messages nil passes through unchanged).

Summary by CodeRabbit

• New Features
  • Added per-step control over tool usage during agent runs, including the ability to keep tools optional, require them, or force a specific tool for a step.
  • Hooks can now update both the step’s context and tool-choice behavior at the same time.
• Tests
  • Added coverage for step-level tool-choice overrides and default behavior when no override is provided.
• Documentation
  • Updated hook guidance to explain the new tool-choice options and usage patterns.

[mark-iii-labs-huly]

Connected to Huly®: KIT-83

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: c11d7258-16aa-422e-ab5c-99841ee1e0f0

📥 Commits

Reviewing files that changed from the base of the PR and between ed51910 and 6f82962.

📒 Files selected for processing (6)

• internal/agent/agent.go
• internal/agent/agent_toolchoice_test.go
• pkg/kit/hooks.go
• pkg/kit/hooks_test.go
• pkg/kit/kit.go
• pkg/kit/types.go

---

📝 Walkthrough

Walkthrough

This PR adds a per-step tool-choice override capability. PrepareStepHandler now returns a *PrepareStepUpdate struct (Messages and ToolChoice, both optional) instead of a message slice. PrepareStepResult in pkg/kit gains a ToolChoice field, wired through Kit.generate. New public types LLMToolChoice, mode constants, and LLMSpecificToolChoice are exposed, with accompanying tests.

Changes

Per-step ToolChoice override

Layer / File(s)	Summary
PrepareStepHandler contract and streaming loop internal/agent/agent.go	PrepareStepHandler returns *PrepareStepUpdate (optional Messages/ToolChoice); streaming loop conditionally applies each field instead of replacing messages outright.
Agent-level tests internal/agent/agent_toolchoice_test.go	Fake fantasy.Agent records per-step ToolChoice; tests verify forced choice on step 0, release on step 1, and no override without a handler.
Kit hooks contract and wiring pkg/kit/hooks.go, pkg/kit/kit.go	PrepareStepResult gains ToolChoice *LLMToolChoice with updated docs; Kit.generate's OnPrepareStep builds an agent.PrepareStepUpdate from hook results, returning nil when nothing changed.
Public ToolChoice API and hook tests pkg/kit/types.go, pkg/kit/hooks_test.go	Adds LLMToolChoice alias, LLMToolChoiceNone/Auto/Required constants, LLMSpecificToolChoice variable, and a test covering per-step ToolChoice hook behavior.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• mark3labs/kit#77: Updates the same PrepareStep/hook wiring in internal/agent/agent.go and pkg/kit/kit.go to apply per-step Tools overrides, analogous to this PR's ToolChoice overrides.

🚥 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 accurately summarizes the main change: exposing per-step ToolChoice control through OnPrepareStep.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ 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

• X
• Mastodon
• Reddit
• LinkedIn

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