feat(adk): add AfterAgent hook to TypedChatModelAgentMiddleware

[shentongmartin]

Summary

Add an AfterAgent lifecycle hook to TypedChatModelAgentMiddleware that fires when the agent reaches a successful terminal state — either a final answer (model response with no tool calls) or a return-directly tool result.

Motivation

Users need a way to perform post-run logic (e.g., persisting conversation history, emitting metrics, cleaning up resources) after the agent completes successfully. Existing middleware hooks only cover pre-run and per-iteration stages; there was no hook for the terminal state.

Design Decisions

Graph Node Approach

AfterAgent is implemented as a graph node (afterAgentNode_) inserted before compose.END, rather than being called externally after graph execution. This ensures:

• compose.ProcessState is available, giving access to the complete final state including tool results from return-directly paths.
• Consistent execution context with other graph nodes.

Success-Only Semantics

AfterAgent fires only on successful terminal paths. It is NOT called on errors (e.g., ErrExceedMaxIterations, context cancellation, model errors). This avoids the design conflict of combining fail-fast error propagation with an error-input parameter.

Fail-Fast Error Propagation

Consistent with all other middleware hooks (BeforeAgent, BeforeModelRewriteState, AfterModelRewriteState): if any handler's AfterAgent returns an error, subsequent handlers are skipped and the error propagates to the event stream.

Changes

adk/handler.go

• Added AfterAgent(ctx, *TypedChatModelAgentState[M]) (context.Context, error) to TypedChatModelAgentMiddleware interface
• Added no-op implementation on TypedBaseChatModelAgentMiddleware

adk/chatmodel.go

• Added applyAfterAgent method that reads state via compose.ProcessState and calls handlers in order with fail-fast
• Wired afterAgentFunc closures in all three run-func builders: no-tools chain, message ReAct graph, agentic ReAct graph

adk/react.go

• Added afterAgentFunc field to typedReactConfig
• Added afterAgentNode_ in both newReact and newAgenticReact: conditionally inserted before compose.END, with both terminal paths (final answer branch, return-directly converter) routing through it

adk/handler_test.go

• Added TestAfterAgent with 8 subtests covering: FinalAnswer, ReturnDirectly, NotCalledOnModelError, NotCalledOnMaxIterations, ErrorStopsRun, ContextPropagation, NoToolsPath, FailFast
• Extended existing TestCustomHandler to verify AfterAgent fires during normal agent runs

Interface

// AfterAgent is called after the agent run reaches a successful terminal state.
AfterAgent(ctx context.Context, state *TypedChatModelAgentState[M]) (context.Context, error)

Test Coverage

All 8 test cases pass with -race:

• FinalAnswer: AfterAgent fires with correct state after model returns no tool calls
• ReturnDirectly: AfterAgent fires after a return-directly tool completes
• NotCalledOnModelError: AfterAgent is skipped when the model returns an error
• NotCalledOnMaxIterations: AfterAgent is skipped when max iterations exceeded
• ErrorStopsRun: AfterAgent error propagates to the stream
• ContextPropagation: Context modifications from AfterAgent are preserved
• NoToolsPath: AfterAgent fires correctly in the no-tools chain path
• FailFast: Second handler is skipped when first handler returns error

[codecov]

Codecov Report

❌ Patch coverage is 96.36364% with 2 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (alpha/09@58d7b8d). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
adk/chatmodel.go	93.54%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             alpha/09    #1002   +/-   ##
===========================================
  Coverage            ?   82.92%           
===========================================
  Files               ?      162           
  Lines               ?    22092           
  Branches            ?        0           
===========================================
  Hits                ?    18319           
  Misses              ?     2538           
  Partials            ?     1235           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.