refactor(adk): replace AfterToolCallsRewriteState with per-run WithAfterToolCallsHook option (#985)

Author: shentongmartin

adk/chatmodel.go

@@ -54,6 +54,8 @@ type typedChatModelAgentExecCtx[M messageType] struct {
 	// Invariant: any code path that emits model output events MUST check this flag.
 	suppressEventSend  bool
 	retryVerdictSignal *retryVerdictSignal
+
+	afterToolCallsHook func(ctx context.Context) error
 }
 
 func (e *typedChatModelAgentExecCtx[M]) send(event *TypedAgentEvent[M]) {
@@ -87,6 +89,8 @@ type chatModelAgentRunOptions struct {
 	agentToolOptions map[string][]AgentRunOption
 
 	historyModifier func(context.Context, []Message) []Message
+
+	afterToolCallsHook func(ctx context.Context) error
 }
 
 // WithChatModelOptions sets options for the underlying chat model.
@@ -118,6 +122,17 @@ func WithHistoryModifier(f func(context.Context, []Message) []Message) AgentRunO
 	})
 }
 
+// WithAfterToolCallsHook registers a per-run hook that fires synchronously after
+// all tool calls in a react iteration complete, before the next ChatModel call.
+//
+// This is suitable for TurnLoop Push+Preempt patterns where the pushed item
+// must be visible to the next turn's GenInput.
+func WithAfterToolCallsHook(fn func(ctx context.Context) error) AgentRunOption {
+	return WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) {
+		t.afterToolCallsHook = fn
+	})
+}
+
 type ToolsConfig struct {
 	compose.ToolsNodeConfig
 
@@ -461,6 +476,8 @@ type typedRunParams[M messageType] struct {
 	cancelCtx      *cancelContext
 	cancelCtxOwned bool
 	composeOpts    []compose.Option
+
+	afterToolCallsHook func(ctx context.Context) error
 }
 
 type typedRunFunc[M messageType] func(ctx context.Context, p *typedRunParams[M])
@@ -1130,6 +1147,7 @@ func (a *TypedChatModelAgent[M]) buildMessageReActRunFunc(ctx context.Context, b
 			generator:                mp.generator,
 			cancelCtx:                cancelCtx,
 			failoverLastSuccessModel: msgModel,
+			afterToolCallsHook:       mp.afterToolCallsHook,
 		})
 
 		// Pre-execution cancel check
@@ -1406,6 +1424,7 @@ func (a *TypedChatModelAgent[M]) Run(ctx context.Context, input *TypedAgentInput
 
 	co := getComposeOptions(opts)
 	co = append(co, compose.WithCheckPointID(bridgeCheckpointID))
+	runOps := GetImplSpecificOptions[chatModelAgentRunOptions](nil, opts...)
 
 	if bc != nil {
 		co = append(co, compose.WithChatModelOption(model.WithTools(bc.toolInfos)))
@@ -1439,14 +1458,15 @@ func (a *TypedChatModelAgent[M]) Run(ctx context.Context, input *TypedAgentInput
 		}
 
 		run(ctx, &typedRunParams[M]{
-			input:          input,
-			generator:      generator,
-			store:          newBridgeStore(),
-			instruction:    instruction,
-			returnDirectly: returnDirectly,
-			cancelCtx:      cancelCtx,
-			cancelCtxOwned: cancelCtxOwned,
-			composeOpts:    co,
+			input:              input,
+			generator:          generator,
+			store:              newBridgeStore(),
+			instruction:        instruction,
+			returnDirectly:     returnDirectly,
+			cancelCtx:          cancelCtx,
+			cancelCtxOwned:     cancelCtxOwned,
+			composeOpts:        co,
+			afterToolCallsHook: runOps.afterToolCallsHook,
 		})
 	}()
 
@@ -1480,6 +1500,7 @@ func (a *TypedChatModelAgent[M]) Resume(ctx context.Context, info *ResumeInfo, o
 
 	co := getComposeOptions(opts)
 	co = append(co, compose.WithCheckPointID(bridgeCheckpointID))
+	resumeRunOps := GetImplSpecificOptions[chatModelAgentRunOptions](nil, opts...)
 
 	if bc != nil {
 		co = append(co, compose.WithChatModelOption(model.WithTools(bc.toolInfos)))
@@ -1563,14 +1584,15 @@ func (a *TypedChatModelAgent[M]) Resume(ctx context.Context, info *ResumeInfo, o
 		}
 
 		run(ctx, &typedRunParams[M]{
-			input:          &TypedAgentInput[M]{EnableStreaming: info.EnableStreaming},
-			generator:      generator,
-			store:          newResumeBridgeStore(bridgeCheckpointID, stateByte),
-			instruction:    instruction,
-			returnDirectly: returnDirectly,
-			cancelCtx:      cancelCtx,
-			cancelCtxOwned: cancelCtxOwned,
-			composeOpts:    co,
+			input:              &TypedAgentInput[M]{EnableStreaming: info.EnableStreaming},
+			generator:          generator,
+			store:              newResumeBridgeStore(bridgeCheckpointID, stateByte),
+			instruction:        instruction,
+			returnDirectly:     returnDirectly,
+			cancelCtx:          cancelCtx,
+			cancelCtxOwned:     cancelCtxOwned,
+			composeOpts:        co,
+			afterToolCallsHook: resumeRunOps.afterToolCallsHook,
 		})
 	}()
 

adk/handler.go

@@ -161,14 +161,6 @@ type TypedChatModelAgentMiddleware[M messageType] interface {
 	//   - DeferredToolInfos: tools for server-side search (nil if unused)
 	AfterModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *ModelContext) (context.Context, *TypedChatModelAgentState[M], error)
 
-	// AfterToolCallsRewriteState is called after all concurrent tool calls in an iteration complete.
-	// The input state includes all messages up to and including the tool call results.
-	// The returned state is persisted to the agent's internal state.
-	//
-	// The ToolCallsContext provides metadata about the tool calls that just completed,
-	// derived from the assistant message's ToolCalls field.
-	AfterToolCallsRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], tc *ToolCallsContext) (context.Context, *TypedChatModelAgentState[M], error)
-
 	// WrapInvokableToolCall wraps a tool's synchronous execution with custom behavior.
 	// Return the input endpoint unchanged and nil error if no wrapping is needed.
 	//
@@ -297,10 +289,6 @@ func (b *TypedBaseChatModelAgentMiddleware[M]) AfterModelRewriteState(ctx contex
 	return ctx, state, nil
 }
 
-func (b *TypedBaseChatModelAgentMiddleware[M]) AfterToolCallsRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], tc *ToolCallsContext) (context.Context, *TypedChatModelAgentState[M], error) {
-	return ctx, state, nil
-}
-
 func processTypedState(ctx context.Context, fn func(extra map[string]any) map[string]any) error {
 	runCtx := getRunCtx(ctx)
 	if runCtx != nil && runCtx.AgenticRootInput != nil {