[FEATURE] Pass invocation_state to all hook events

[mkmeral]

Problem Statement

Currently, invocation_state is successfully passed to tool-related hook events (BeforeToolInvocationEvent and AfterToolInvocationEvent), which has proven valuable for tool configuration and coordination. With the recent addition of multi-agent patterns (Graph and Swarm) that support passing invocation_state to propagate shared configuration across agents, we have an opportunity to extend this powerful pattern to enhance hook functionality even further.

By making invocation_state available to additional hook events, we could enable richer use cases and more sophisticated cross-cutting concerns:

• BeforeInvocationEvent - Enable request tracking and dynamic configuration
• AfterInvocationEvent - Support result processing with context awareness
• BeforeModelInvocationEvent - Allow model configuration based on shared state
• AfterModelInvocationEvent - Enable context-aware response processing
• MessageAddedEvent - Support message enrichment with workflow metadata

This enhancement would particularly benefit multi-agent scenarios where the invocation state carries important context across the entire system, enabling developers to build more powerful observability, configuration, and coordination patterns.

Proposed Solution

Add invocation_state as an attribute to all hook event classes, not just tool-related ones. This would involve:

1. Updating hook event classes to include an invocation_state attribute:

class BeforeInvocationEvent:
    agent: Agent
    invocation_state: dict[str, Any]  # Add this
    # ... other attributes

class AfterInvocationEvent:
    agent: Agent
    result: Optional[AgentResult]
    exception: Optional[Exception]
    invocation_state: dict[str, Any]  # Add this
    # ... other attributes

class BeforeModelInvocationEvent:
    agent: Agent
    invocation_state: dict[str, Any]  # Add this
    # ... other attributes

# Similar for other events

2. Ensuring the event loop passes invocation_state when creating these events

Use Case

1. Multi-Agent Request Tracking

In a Graph or Swarm, track and log requests across all agents using shared session information:

class MultiAgentTracker(HookProvider):
    def register_hooks(self, registry: HookRegistry) -> None:
        registry.add_callback(BeforeInvocationEvent, self.track_start)
        registry.add_callback(AfterInvocationEvent, self.track_end)
    
    def track_start(self, event: BeforeInvocationEvent) -> None:
        session_id = event.invocation_state.get("session_id")
        request_id = event.invocation_state.get("request_id")
        logger.info(f"Agent {event.agent.name} starting in session {session_id}, request {request_id}")
    
    def track_end(self, event: AfterInvocationEvent) -> None:
        session_id = event.invocation_state.get("session_id")
        metrics = event.invocation_state.get("shared_metrics", {})
        # Update shared metrics for the multi-agent execution

2. Dynamic Agent Configuration

Modify agent behavior based on multi-agent context:

class DynamicAgentConfig(HookProvider):
    def register_hooks(self, registry: HookRegistry) -> None:
        registry.add_callback(BeforeModelInvocationEvent, self.configure_model)
    
    def configure_model(self, event: BeforeModelInvocationEvent) -> None:
        # Access multi-agent context
        agent_role = event.invocation_state.get("agent_role")
        environment = event.invocation_state.get("environment")
        
        # Dynamically adjust model parameters based on context
        if agent_role == "critic" and environment == "production":
            # Make critic more conservative in production
            event.agent.model.params["temperature"] = 0.2

3. Message Enrichment

Add metadata to messages based on shared state:

class MessageEnricher(HookProvider):
    def register_hooks(self, registry: HookRegistry) -> None:
        registry.add_callback(MessageAddedEvent, self.enrich_message)
    
    def enrich_message(self, event: MessageAddedEvent) -> None:
        # Add metadata from invocation state
        workflow_id = event.invocation_state.get("workflow_id")
        stage = event.invocation_state.get("current_stage")
        
        # Tag messages with workflow context (for debugging/tracing)
        event.message.metadata = {
            "workflow_id": workflow_id,
            "stage": stage,
            "agent": event.agent.name
        }

4. Conditional Hook Logic

Enable/disable hook behavior based on invocation context:

class ConditionalLogger(HookProvider):
    def register_hooks(self, registry: HookRegistry) -> None:
        registry.add_callback(BeforeInvocationEvent, self.maybe_log)
    
    def maybe_log(self, event: BeforeInvocationEvent) -> None:
        # Only log in debug mode
        if event.invocation_state.get("debug_mode"):
            print(f"Debug: Agent {event.agent.name} invoked")
        
        # Different behavior for different environments
        if event.invocation_state.get("environment") == "staging":
            # Additional staging-specific logic
            pass

Alternatives Solutions

No response

Additional Context

No response

[zastrowm]

It might be simpler to start with these 4 hooks:

BeforeInvocationEvent - Enable request tracking and dynamic configuration
AfterInvocationEvent - Support result processing with context awareness
BeforeModelInvocationEvent - Allow model configuration based on shared state
AfterModelInvocationEvent - Enable context-aware response processing

And pick up this one as a follow-up

MessageAddedEvent - Support message enrichment with workflow metadata

Only because MessageAddedEvent is spread throughout more places, and I think the first 4 benefit the most from these