Agent Lifecycle and Callbacks

Purpose

Within the broader AI Agent Framework that manages agent creation, execution, and composition, the Agent Lifecycle and Callbacks subtopic addresses the need for fine-grained control over the agent's execution phases. Specifically, it enables developers to inject custom logic before an agent begins its main run and after an agent completes execution. This capability is crucial for customizing agent behavior dynamically, handling pre-processing or post-processing tasks, modifying agent state, or short-circuiting the normal run flow based on certain conditions.

This subtopic solves the challenge of extending agent execution without modifying core agent logic or the main run function. It provides structured hooks to:

Perform validation, initialization, or early returns before the agent runs.
Modify outputs, update state, or trigger side effects after the agent finishes.
Abort execution early based on callback results.

These lifecycle callbacks integrate seamlessly with the agent invocation context, allowing rich access to session state, artifacts, user inputs, and invocation metadata.

Functionality

Before and After Agent Callbacks

Agents created via agent.New accept two callback lists in their configuration:

BeforeAgentCallbacks []BeforeAgentCallback
AfterAgentCallbacks []AfterAgentCallback

Each callback is a function receiving a CallbackContext and returns either:

A non-nil genai.Content to produce an immediate event and potentially skip the main agent run (for before callbacks).
An error to terminate the lifecycle with failure.

The lifecycle within the agent's Run method executes as follows:

Before Callbacks Execution
- Run sequentially.
- If any callback returns non-nil content or error, the main agent run is skipped.
- A new session event is created from the callback's content or error.
- The invocation is marked as ended (ctx.EndInvocation()).
- If callbacks only update state delta without returning content, an event with the state delta is emitted.
Main Agent Run
- The agent's primary run function executes only if no before callback short-circuits.
- Events yielded from the run function are emitted, with automatic author assignment if missing.
After Callbacks Execution
- Run sequentially if the invocation has not been ended prematurely.
- If any callback returns content or error, a new event is created and emitted.
- State delta modifications produce events even if no content is returned.

Invocation Context and CallbackContext

Callbacks receive a rich context interface (CallbackContext), which provides access to:

The current agent (AgentName()).
Session state (read-only and mutable).
Artifacts associated with the session.
Invocation metadata such as invocation ID, branch, user content.
Session identifiers and app information.

This context allows callbacks to inspect and modify the agent's environment safely and effectively.

State Delta Mechanism

Callbacks can mutate session state by setting key-value pairs on the callback context's state object. These changes are aggregated into a state delta, which is included in the resulting session event's actions (EventActions.StateDelta).

This mechanism supports stateful side effects from callbacks without requiring explicit event creation code in the callback itself.

Early Termination of Invocation

Callbacks can explicitly stop further agent execution by marking the invocation ended (ctx.EndInvocation()), either by returning content or by setting this flag directly. This feature is useful for implementing guard clauses or conditional shortcuts in agent behavior.

Code Snippet Illustrating Core Lifecycle Flow

func (a *agent) Run(ctx InvocationContext) iter.Seq2[*session.Event, error] {
    return func(yield func(*session.Event, error) bool) {
        // Run before callbacks, potentially skip main run
        event, err := runBeforeAgentCallbacks(ctx)
        if event != nil || err != nil {
            if !yield(event, err) {
                return
            }
        }
        if ctx.Ended() {
            return
        }

        // Run main agent logic
        for event, err := range a.run(ctx) {
            if event != nil && event.Author == "" {
                event.Author = ctx.Agent().Name()
            }
            if !yield(event, err) {
                return
            }
        }
        if ctx.Ended() {
            return
        }

        // Run after callbacks, emit additional events if any
        event, err = runAfterAgentCallbacks(ctx)
        if event != nil || err != nil {
            yield(event, err)
        }
    }
}

Integration

The agent lifecycle and callback mechanism is a foundational part of the overall AI Agent Framework. It enhances the core agent interface by allowing modular extension points in the agent invocation process without changing the main run logic or the agent tree structure.

Complementing Agent Invocation Context

This subtopic builds upon the Agent Invocation Context subtopic by providing callbacks direct access to the invocation context, session, and artifacts. This enables callbacks to interact with the same data and services that the agent uses during its run.

Interaction with Session Management

Callbacks can mutate session state and append new events, thereby integrating closely with Session Management. State deltas from callbacks propagate to the session, affecting subsequent agent invocations and user interactions.

Coordination with Agent Execution Runner

While the Agent Execution Runner orchestrates overall agent execution within sessions, the lifecycle callbacks provide hooks inside the agent's run process, enabling pre- and post-processing of events emitted by the runner.

Enhancing Tool and Workflow Agents

Callback hooks can be used to customize behavior in specialized agents such as those in Agent Workflow Management or LLM Integration and Agents, enabling use cases like logging, telemetry, dynamic context injection, or conditional execution flows.

Diagram

sequenceDiagram
participant Caller as Runner/User
participant Agent as Agent Instance
participant BeforeCB as BeforeAgentCallbacks
participant MainRun as Agent Run Function
participant AfterCB as AfterAgentCallbacks
participant Session as Session Store
Caller->>Agent: Run(ctx)
Agent->>BeforeCB: Execute callbacks sequentially
alt Before callback returns content/error
BeforeCB-->>Agent: Event + EndInvocation
Agent-->>Caller: Yield event (Skip main run)
else No before callback interruption
BeforeCB-->>Agent: nil
Agent->>MainRun: Execute main run logic
MainRun-->>Agent: Yield events
Agent->>AfterCB: Execute callbacks sequentially
alt After callback returns content/error
AfterCB-->>Agent: Event
Agent-->>Caller: Yield event
else No after callback event
AfterCB-->>Agent: nil
end
end

This sequence diagram highlights the flow of control and event generation through before callbacks, main agent run, and after callbacks, emphasizing how callbacks can short-circuit or extend the agent lifecycle.