feat(adk): add AfterAgent method to ChatModelAgentMiddleware

[shentongmartin]

Problem

Middleware authors need to perform cleanup operations after the ChatModelAgent completes execution. For example, during execution, a middleware might store temporary file paths in RunLocalValues. When the agent finishes, the middleware needs to delete these temporary files.

The challenge: GetRunLocalValue only works while the graph is executing. Once the graph exits, the state context is lost and RunLocalValues become inaccessible.

Solution

Add AfterAgent method to ChatModelAgentMiddleware interface:

AfterAgent(ctx context.Context, state *ChatModelAgentState) error

This method is called before the graph exits, so GetRunLocalValue still works. The state parameter contains all messages including the final one.

Graph topology for react path:

Before:
  ChatModel --[no tools]--> END
  ToolNodeToEndConverter --> END

After:
  ChatModel --[no tools]--> AfterAgentNode --> END
  ToolNodeToEndConverter --> AfterAgentNode --> END

Key Insight

The stateModelWrapper already appends the model response to State.Messages and writes it back via ProcessState before any post-processing. This means in AfterAgent, we can simply read State.Messages directly - no need to copy/concatenate the stream ourselves.

Initial implementation copied and concatenated the stream to get the final message:

streams := msgStream.Copy(2)
finalMsg, err := schema.ConcatMessageStream(streams[0])
state := &ChatModelAgentState{Messages: append(stateMessages, finalMsg)}

But this was redundant and would have duplicated the final message. The simplified implementation just reads from state:

var stateMessages []Message
compose.ProcessState(ctx, func(_ context.Context, st *State) error {
    stateMessages = st.Messages
    return nil
})
state := &ChatModelAgentState{Messages: stateMessages}

Summary

Problem	Solution
Middleware cannot clean up RunLocalValues after agent execution because state is lost when graph exits	Add AfterAgent method that is called before graph exits, while GetRunLocalValue still works

[codecov]

Codecov Report

❌ Patch coverage is 92.59259% with 4 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (alpha/08@6565ff4). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
adk/handler.go	33.33%	2 Missing ⚠️
adk/react.go	92.30%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             alpha/08     #735   +/-   ##
===========================================
  Coverage            ?   80.75%           
===========================================
  Files               ?      132           
  Lines               ?    13638           
  Branches            ?        0           
===========================================
  Hits                ?    11013           
  Misses              ?     1812           
  Partials            ?      813           

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