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

internal/agent/agent.go

@@ -163,10 +163,21 @@ type ErrorHandler func(err error)
 // RetryHandler is called when the LLM request is retried.
 type RetryHandler func(attempt int, err error)
 
-// PrepareStepHandler is called between steps to allow message modification.
-// It receives the step number and current messages, and returns replacement
-// messages (or nil to keep unchanged).
-type PrepareStepHandler func(stepNumber int, messages []fantasy.Message) []fantasy.Message
+// PrepareStepUpdate carries per-step overrides returned by a
+// PrepareStepHandler. Nil fields leave the corresponding aspect of the step
+// unchanged.
+type PrepareStepUpdate struct {
+	// Messages, when non-nil, replaces the context window for this step.
+	Messages []fantasy.Message
+	// ToolChoice, when non-nil, overrides the tool-choice mode for this step.
+	ToolChoice *fantasy.ToolChoice
+}
+
+// PrepareStepHandler is called between steps to allow message modification
+// and per-step tool-choice control. It receives the step number and current
+// messages, and returns a PrepareStepUpdate (or nil to keep everything
+// unchanged).
+type PrepareStepHandler func(stepNumber int, messages []fantasy.Message) *PrepareStepUpdate
 
 // GenerateCallbacks consolidates all callback functions for
 // GenerateWithCallbacks into a single struct, replacing what was previously
@@ -923,8 +934,13 @@ func (a *Agent) GenerateWithCallbacks(ctx context.Context, messages []fantasy.Me
 
 			// Phase 2: Run OnPrepareStep hook (if registered).
 			if hasPrepareStepHook {
-				if replacement := cb.OnPrepareStep(opts.StepNumber, result.Messages); replacement != nil {
-					result.Messages = replacement
+				if update := cb.OnPrepareStep(opts.StepNumber, result.Messages); update != nil {
+					if update.Messages != nil {
+						result.Messages = update.Messages
+					}
+					if update.ToolChoice != nil {
+						result.ToolChoice = update.ToolChoice
+					}
 				}
 			}
 

internal/agent/agent_toolchoice_test.go

@@ -0,0 +1,103 @@
+package agent
+
+import (
+	"context"
+	"testing"
+
+	"charm.land/fantasy"
+)
+
+// toolChoiceAgent is a fake fantasy.Agent that simulates a two-step agentic
+// loop and records the ToolChoice each step's PrepareStep callback yields, so
+// tests can assert per-step tool-choice overrides are threaded through.
+type toolChoiceAgent struct {
+	stepChoices []*fantasy.ToolChoice
+}
+
+func (f *toolChoiceAgent) Generate(_ context.Context, _ fantasy.AgentCall) (*fantasy.AgentResult, error) {
+	return &fantasy.AgentResult{}, nil
+}
+
+func (f *toolChoiceAgent) Stream(ctx context.Context, opts fantasy.AgentStreamCall) (*fantasy.AgentResult, error) {
+	for step := range 2 {
+		if opts.PrepareStep != nil {
+			_, prepared, err := opts.PrepareStep(ctx, fantasy.PrepareStepFunctionOptions{
+				StepNumber: step,
+				Model:      nil,
+				Messages:   nil,
+			})
+			if err != nil {
+				return nil, err
+			}
+			f.stepChoices = append(f.stepChoices, prepared.ToolChoice)
+		}
+	}
+	return &fantasy.AgentResult{}, nil
+}
+
+// TestPrepareStepToolChoiceForcing verifies that a PrepareStepHandler can
+// force a specific tool choice on one step and release it on the next. This
+// is the pattern used to guarantee a capture/structured-output tool is
+// invoked: force SpecificToolChoice until the call is observed, then flip
+// back to nil so the turn can end.
+func TestPrepareStepToolChoiceForcing(t *testing.T) {
+	t.Parallel()
+
+	fake := &toolChoiceAgent{}
+	a := &Agent{
+		fantasyAgent:     fake,
+		streamingEnabled: true,
+	}
+
+	forced := fantasy.SpecificToolChoice("record_result")
+	cb := GenerateCallbacks{
+		OnPrepareStep: func(stepNumber int, messages []fantasy.Message) *PrepareStepUpdate {
+			if stepNumber == 0 {
+				return &PrepareStepUpdate{ToolChoice: &forced}
+			}
+			return nil
+		},
+	}
+
+	msgs := []fantasy.Message{fantasy.NewUserMessage("go")}
+	if _, err := a.GenerateWithCallbacks(context.Background(), msgs, cb); err != nil {
+		t.Fatalf("GenerateWithCallbacks returned error: %v", err)
+	}
+
+	if len(fake.stepChoices) != 2 {
+		t.Fatalf("expected 2 prepared steps, got %d", len(fake.stepChoices))
+	}
+	if fake.stepChoices[0] == nil {
+		t.Fatal("step 0: expected forced ToolChoice, got nil")
+	}
+	if *fake.stepChoices[0] != forced {
+		t.Errorf("step 0: expected ToolChoice %q, got %q", forced, *fake.stepChoices[0])
+	}
+	if fake.stepChoices[1] != nil {
+		t.Errorf("step 1: expected nil ToolChoice after release, got %q", *fake.stepChoices[1])
+	}
+}
+
+// TestPrepareStepToolChoiceWithoutHandler verifies that when no
+// PrepareStepHandler is registered, no ToolChoice override is sent — fantasy
+// falls back to its default (auto).
+func TestPrepareStepToolChoiceWithoutHandler(t *testing.T) {
+	t.Parallel()
+
+	fake := &toolChoiceAgent{}
+	a := &Agent{
+		fantasyAgent:     fake,
+		streamingEnabled: true,
+	}
+
+	msgs := []fantasy.Message{fantasy.NewUserMessage("go")}
+	if _, err := a.GenerateWithCallbacks(context.Background(), msgs, GenerateCallbacks{}); err != nil {
+		t.Fatalf("GenerateWithCallbacks returned error: %v", err)
+	}
+
+	for i, tc := range fake.stepChoices {
+		if tc != nil {
+			t.Errorf("step %d: expected nil ToolChoice without handler, got %q", i, *tc)
+		}
+	}
+}

pkg/kit/hooks.go

@@ -138,11 +138,24 @@ type PrepareStepHook struct {
 	Messages []LLMMessage
 }
 
-// PrepareStepResult can replace the context window between steps.
+// PrepareStepResult can replace the context window and override the
+// tool-choice mode between steps.
 type PrepareStepResult struct {
 	// Messages replaces the entire context window for this step. If nil,
 	// the original messages (including any steering) are used unchanged.
 	Messages []LLMMessage
+	// ToolChoice, when non-nil, overrides the tool-choice mode for this step
+	// only. Use [LLMToolChoiceRequired] or [LLMSpecificToolChoice] to force a
+	// tool call, [LLMToolChoiceNone] to forbid tool calls, or
+	// [LLMToolChoiceAuto] to explicitly restore the default. If nil, the
+	// provider default (auto) is used.
+	//
+	// Note that a turn only ends when the model stops emitting tool calls, so
+	// forcing a tool choice on every step prevents the turn from finishing.
+	// Callers forcing a tool call should flip back to nil (or
+	// LLMToolChoiceAuto) on subsequent steps once the desired call has been
+	// observed.
+	ToolChoice *LLMToolChoice
 }
 
 // ---------------------------------------------------------------------------
@@ -275,12 +288,14 @@ func (m *Kit) OnBeforeCompact(p HookPriority, h func(BeforeCompactHook) *BeforeC
 // OnPrepareStep registers a hook that fires between steps within a multi-step
 // agent turn, after steering messages are injected and before the messages are
 // sent to the LLM. Return a non-nil PrepareStepResult with Messages to replace
-// the entire context window for this step. Hooks execute in priority order;
-// the first non-nil result wins. Returns an unregister function.
+// the entire context window for this step, and/or with ToolChoice to override
+// the tool-choice mode for this step. Hooks execute in priority order; the
+// first non-nil result wins. Returns an unregister function.
 //
 // This is the most powerful interception point in the agent lifecycle. It
 // enables patterns like transforming tool results, dynamic tool filtering,
-// and mid-turn context injection.
+// mid-turn context injection, and forcing a specific tool call on a step
+// (e.g. guaranteeing a structured-output capture tool is invoked).
 func (m *Kit) OnPrepareStep(p HookPriority, h func(PrepareStepHook) *PrepareStepResult) func() {
 	return m.prepareStep.register(p, h)
 }

pkg/kit/hooks_test.go

@@ -578,6 +578,45 @@ func TestPrepareStepHookRegistry(t *testing.T) {
 	}
 }
 
+// TestPrepareStepHookToolChoice verifies that a PrepareStep hook can set a
+// per-step tool-choice override without replacing messages.
+func TestPrepareStepHookToolChoice(t *testing.T) {
+	hr := newHookRegistry[PrepareStepHook, PrepareStepResult]()
+
+	forced := LLMSpecificToolChoice("record_result")
+	hr.register(HookPriorityNormal, func(h PrepareStepHook) *PrepareStepResult {
+		if h.StepNumber == 0 {
+			return &PrepareStepResult{ToolChoice: &forced}
+		}
+		return nil
+	})
+
+	input := PrepareStepHook{
+		StepNumber: 0,
+		Messages:   []LLMMessage{NewLLMUserMessage("hello")},
+	}
+	result := hr.run(input)
+	if result == nil {
+		t.Fatal("expected non-nil result for step 0")
+		return
+	}
+	if result.Messages != nil {
+		t.Errorf("expected nil Messages (unchanged), got %d messages", len(result.Messages))
+	}
+	if result.ToolChoice == nil {
+		t.Fatal("expected ToolChoice to be set")
+	}
+	if *result.ToolChoice != forced {
+		t.Errorf("expected ToolChoice %q, got %q", forced, *result.ToolChoice)
+	}
+
+	// Step 1 — no override.
+	input.StepNumber = 1
+	if result := hr.run(input); result != nil {
+		t.Errorf("expected nil result for step 1, got %+v", result)
+	}
+}
+
 // TestPrepareStepHookPriority verifies that PrepareStep hooks respect priority ordering.
 func TestPrepareStepHookPriority(t *testing.T) {
 	hr := newHookRegistry[PrepareStepHook, PrepareStepResult]()

pkg/kit/kit.go

@@ -2695,15 +2695,18 @@ func (m *Kit) generate(ctx context.Context, messages []fantasy.Message) (*agent.
 			if !m.prepareStep.hasHooks() {
 				return nil
 			}
-			return func(stepNumber int, messages []fantasy.Message) []fantasy.Message {
+			return func(stepNumber int, messages []fantasy.Message) *agent.PrepareStepUpdate {
 				hookResult := m.prepareStep.run(PrepareStepHook{
 					StepNumber: stepNumber,
 					Messages:   messages,
 				})
-				if hookResult != nil && hookResult.Messages != nil {
-					return hookResult.Messages
+				if hookResult == nil || (hookResult.Messages == nil && hookResult.ToolChoice == nil) {
+					return nil
+				}
+				return &agent.PrepareStepUpdate{
+					Messages:   hookResult.Messages,
+					ToolChoice: hookResult.ToolChoice,
 				}
-				return nil
 			}
 		}(),
 	})

pkg/kit/types.go

@@ -228,6 +228,26 @@ const (
 // LLMToolInfo describes a tool's name, description, and JSON-Schema parameters.
 type LLMToolInfo = fantasy.ToolInfo
 
+// LLMToolChoice controls which tool (if any) the model is allowed or required
+// to call on a given LLM step. Use the LLMToolChoice* constants for the
+// generic modes, or [LLMSpecificToolChoice] to force a single named tool.
+type LLMToolChoice = fantasy.ToolChoice
+
+// LLMToolChoice mode constants mirror the provider's tool-choice values.
+const (
+	// LLMToolChoiceNone forbids tool calls for the step.
+	LLMToolChoiceNone = fantasy.ToolChoiceNone
+	// LLMToolChoiceAuto lets the model decide whether to call tools. This is
+	// the default behaviour when no tool choice is set.
+	LLMToolChoiceAuto = fantasy.ToolChoiceAuto
+	// LLMToolChoiceRequired forces the model to call some tool on the step.
+	LLMToolChoiceRequired = fantasy.ToolChoiceRequired
+)
+
+// LLMSpecificToolChoice returns a tool choice that forces the model to call
+// the named tool on the step.
+var LLMSpecificToolChoice = fantasy.SpecificToolChoice
+
 // LLMProviderOptions carries provider-specific key/value option maps, keyed
 // by provider name (e.g. "anthropic"). Use this when configuring or
 // inspecting provider-specific tool behaviour.