Eino ADK: ChatModelAgentMiddleware

type ChatModelAgentMiddleware interface {
    // BeforeAgent is called before each agent run, allows modifying instruction and tools configuration
    BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error)

    // BeforeModelRewriteState is called before each model call
    // The returned state will be persisted to the agent's internal state and passed to the model
    // The returned context will be propagated to the model call and subsequent handlers
    BeforeModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error)

    // AfterModelRewriteState is called after each model call
    // The input state contains the model response as the last message
    AfterModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error)

    // WrapInvokableToolCall wraps the synchronous execution of a tool with custom behavior
    // If no wrapping is needed, return the original endpoint and nil error
    // Only called for tools that implement InvokableTool
    WrapInvokableToolCall(ctx context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) (InvokableToolCallEndpoint, error)

    // WrapStreamableToolCall wraps the streaming execution of a tool with custom behavior
    // If no wrapping is needed, return the original endpoint and nil error
    // Only called for tools that implement StreamableTool
    WrapStreamableToolCall(ctx context.Context, endpoint StreamableToolCallEndpoint, tCtx *ToolContext) (StreamableToolCallEndpoint, error)

    // WrapEnhancedInvokableToolCall wraps the synchronous execution of an enhanced tool with custom behavior
    WrapEnhancedInvokableToolCall(ctx context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) (EnhancedInvokableToolCallEndpoint, error)

    // WrapEnhancedStreamableToolCall wraps the streaming execution of an enhanced tool with custom behavior
    WrapEnhancedStreamableToolCall(ctx context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) (EnhancedStreamableToolCallEndpoint, error)

    // WrapModel wraps the chat model with custom behavior
    // If no wrapping is needed, return the original model and nil error
    // Called at request time, executed before each model call
    WrapModel(ctx context.Context, m model.BaseChatModel, mc *ModelContext) (model.BaseChatModel, error)
}

type MyHandler struct {
    *adk.BaseChatModelAgentMiddleware
}

func (h *MyHandler) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, mc *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) {
    return ctx, state, nil
}

// InvokableToolCallEndpoint is the function signature for synchronous tool calls
type InvokableToolCallEndpoint func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error)

// StreamableToolCallEndpoint is the function signature for streaming tool calls
type StreamableToolCallEndpoint func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error)

// EnhancedInvokableToolCallEndpoint is the function signature for enhanced synchronous tool calls
type EnhancedInvokableToolCallEndpoint func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error)

// EnhancedStreamableToolCallEndpoint is the function signature for enhanced streaming tool calls
type EnhancedStreamableToolCallEndpoint func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error)

type ChatModelAgentContext struct {
    // Instruction is the instruction for the current Agent execution
    // Includes agent-configured instructions, framework and AgentMiddleware appended extra instructions,
    // and modifications applied by previous BeforeAgent handlers
    Instruction string

    // Tools are the original tools (without any wrappers or tool middleware) currently configured for Agent execution
    // Includes tools passed in AgentConfig, tools implicitly added by the framework (like transfer/exit tools),
    // and other tools added by middleware
    Tools []tool.BaseTool

    // ReturnDirectly is the set of tool names currently configured to make the Agent return directly
    ReturnDirectly map[string]bool
}

type ChatModelAgentState struct {
    // Messages contains all messages in the current conversation session
    Messages []Message
}

type ToolContext struct {
    // Name is the tool name
    Name string

    // CallID is the unique identifier for this specific tool call
    CallID string
}

func (h *MyHandler) WrapInvokableToolCall(ctx context.Context, endpoint adk.InvokableToolCallEndpoint, tCtx *adk.ToolContext) (adk.InvokableToolCallEndpoint, error) {
    return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
        log.Printf("Tool %s (call %s) starting with args: %s", tCtx.Name, tCtx.CallID, argumentsInJSON)
        
        result, err := endpoint(ctx, argumentsInJSON, opts...)
        
        if err != nil {
            log.Printf("Tool %s failed: %v", tCtx.Name, err)
            return "", err
        }
        
        log.Printf("Tool %s completed with result: %s", tCtx.Name, result)
        return result, nil
    }, nil
}

type ModelContext struct {
    // Tools is the list of tools currently configured for the agent
    // Populated at request time, contains the tools that will be sent to the model
    Tools []*schema.ToolInfo

    // ModelRetryConfig contains the retry configuration for the model
    // Populated at request time from the agent's ModelRetryConfig
    // Used by EventSenderModelWrapper to appropriately wrap stream errors
    ModelRetryConfig *ModelRetryConfig
}

func (h *MyHandler) WrapModel(ctx context.Context, m model.BaseChatModel, mc *adk.ModelContext) (model.BaseChatModel, error) {
    return &myModelWrapper{
        inner: m,
        tools: mc.Tools,
    }, nil
}

type myModelWrapper struct {
    inner model.BaseChatModel
    tools []*schema.ToolInfo
}

func (w *myModelWrapper) Generate(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
    log.Printf("Model called with %d tools", len(w.tools))
    return w.inner.Generate(ctx, msgs, opts...)
}

func (w *myModelWrapper) Stream(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
    return w.inner.Stream(ctx, msgs, opts...)
}

// SetRunLocalValue sets a key-value pair that persists during the current agent Run() call
// The value is scoped to this specific execution and is not shared between different Run() calls or agent instances
//
// Values stored here are compatible with interrupt/resume cycles - they are serialized and restored when the agent resumes
// For custom types, they must be registered in init() using schema.RegisterName[T]() to ensure proper serialization
//
// This function can only be called from within a ChatModelAgentMiddleware during agent execution
// Returns an error if called outside of agent execution context
func SetRunLocalValue(ctx context.Context, key string, value any) error

// GetRunLocalValue retrieves a value set during the current agent Run() call
// The value is scoped to this specific execution and is not shared between different Run() calls or agent instances
//
// Values stored via SetRunLocalValue are compatible with interrupt/resume cycles - they are serialized and restored when the agent resumes
// For custom types, they must be registered in init() using schema.RegisterName[T]() to ensure proper serialization
//
// This function can only be called from within a ChatModelAgentMiddleware during agent execution
// Returns (value, true, nil) if found, (nil, false, nil) if not found,
// returns error if called outside of agent execution context
func GetRunLocalValue(ctx context.Context, key string) (any, bool, error)

// DeleteRunLocalValue deletes a value set during the current agent Run() call
//
// This function can only be called from within a ChatModelAgentMiddleware during agent execution
// Returns an error if called outside of agent execution context
func DeleteRunLocalValue(ctx context.Context, key string) error

func init() {
    schema.RegisterName[*MyCustomData]("my_package.MyCustomData")
}

type MyCustomData struct {
    Count int
    Name  string
}

type MyHandler struct {
    *adk.BaseChatModelAgentMiddleware
}

func (h *MyHandler) WrapInvokableToolCall(ctx context.Context, endpoint adk.InvokableToolCallEndpoint, tCtx *adk.ToolContext) (adk.InvokableToolCallEndpoint, error) {
    return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
        result, err := endpoint(ctx, argumentsInJSON, opts...)
        
        data := &MyCustomData{Count: 1, Name: tCtx.Name}
        if err := adk.SetRunLocalValue(ctx, "my_handler.last_tool", data); err != nil {
            log.Printf("Failed to set run local value: %v", err)
        }
        
        return result, err
    }, nil
}

func (h *MyHandler) AfterModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, mc *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) {
    if val, found, err := adk.GetRunLocalValue(ctx, "my_handler.last_tool"); err == nil && found {
        if data, ok := val.(*MyCustomData); ok {
            log.Printf("Last tool was: %s (count: %d)", data.Name, data.Count)
        }
    }
    return ctx, state, nil
}

// SendEvent sends a custom AgentEvent to the event stream during agent execution
// Allows ChatModelAgentMiddleware implementations to emit custom events,
// which will be received by callers iterating over the agent's event stream
//
// This function can only be called from within a ChatModelAgentMiddleware during agent execution
// Returns an error if called outside of agent execution context
func SendEvent(ctx context.Context, event *AgentEvent) error

type State struct {
    Messages []Message
    extra    map[string]any  // unexported, access via SetRunLocalValue/GetRunLocalValue

    // The following are internal fields - do not access directly
    // Kept exported for backward compatibility with existing checkpoints
    ReturnDirectlyToolCallID string
    ToolGenActions           map[string]*AgentAction
    AgentName                string
    RemainingIterations      int

    internals map[string]any
}

middleware := adk.AgentMiddleware{
    BeforeChatModel: func(ctx context.Context, state *adk.ChatModelAgentState) error {
        return nil
    },
}

type MyHandler struct {
    *adk.BaseChatModelAgentMiddleware
}

func (h *MyHandler) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, mc *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) {
    newCtx := context.WithValue(ctx, myKey, myValue)
    return newCtx, state, nil
}

compose.ProcessState(ctx, func(_ context.Context, st *adk.State) error {
    st.Extra["myKey"] = myValue
    return nil
})

if err := adk.SetRunLocalValue(ctx, "myKey", myValue); err != nil {
    return ctx, state, err
}

if val, found, err := adk.GetRunLocalValue(ctx, "myKey"); err == nil && found {
}