callback_context.go

Overview

The callback_context.go file implements a mutable callback context used during agent invocation callbacks, extending a read-only invocation context with state mutation and artifact saving capabilities. It is part of the broader Agent Invocation Context framework (Agent Invocation Context) that encapsulates all relevant information and state management during an agent's invocation lifecycle.

This file defines the callbackContext struct, which provides access to mutable session state, artifact saving with version tracking, and references to the underlying invocation context. It supports tracking state deltas and artifact version changes during an invocation, enabling fine-grained updates to session state and artifact metadata that are committed at the end of the invocation. The callbackContext also wraps artifact operations to record changes into an EventActions structure, facilitating event-driven workflows and version consistency.

Entities and Key Structures

internalArtifacts

• Description:
  An internal wrapper around the agent.Artifacts interface that intercepts Save calls to track artifact version changes within the current invocation's event actions.

• Fields:

  • Artifacts agent.Artifacts — The underlying artifact service used to persist artifacts.

  • eventActions *session.EventActions — Tracks state and artifact deltas during the invocation.

• Methods:

  • Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error)
    Saves an artifact and updates the ArtifactDelta in eventActions with the new version number for the saved artifact.

• Usage:
  Used internally by callbackContext to provide artifact saving functionality while recording artifact version changes for event processing.

callbackContext

• Description:
  Implements the agent.CallbackContext interface by extending a ReadonlyContext with mutable state and artifact saving support. It maintains a reference to the underlying agent.InvocationContext and an EventActions structure that records state and artifact changes.

• Fields:

  • Embeds agent.ReadonlyContext for read-only accessors like agent name and user content.

  • artifacts *internalArtifacts — Wrapped artifact service for saving with delta tracking.

  • invocationCtx agent.InvocationContext — The underlying invocation context.

  • eventActions *session.EventActions — Tracks state and artifact changes during the invocation.

• Methods:

  • Artifacts() agent.Artifacts
    Returns the wrapped artifact service with version tracking.

  • AgentName() string
    Returns the name of the agent from the invocation context.

  • ReadonlyState() session.ReadonlyState
    Returns a read-only view of the session state.

  • State() session.State
    Returns a mutable state interface backed by state deltas in eventActions.

  • InvocationID() string
    Returns the unique invocation identifier.

  • UserContent() *genai.Content
    Returns user message content triggering the invocation.

• Construction:

  • Use NewCallbackContext(ctx agent.InvocationContext) to create a context with an empty state delta.

  • Use NewCallbackContextWithDelta(ctx agent.InvocationContext, stateDelta map[string]any) to create a context initialized with a given state delta.

• Internal helper:

  • newCallbackContext(ctx agent.InvocationContext, stateDelta map[string]any) *callbackContext initializes the struct, wrapping the read-only context and setting up EventActions.

• Usage:
  Used during agent lifecycle callbacks to enable mutation of state and saving of artifacts with consistent tracking of changes.

callbackContextState

• Description:
  Implements the session.State interface for mutable session state that considers both the session's persisted state and the in-memory deltas tracked in EventActions.

• Fields:

  • ctx *callbackContext — Reference to the parent callback context.

• Methods:

  • Get(key string) (any, error)
    Returns the value for a key, checking deltas first before falling back to the session state.

  • Set(key string, val any) error
    Sets a value in the state delta map and also updates the underlying session state.

  • All() iter.Seq2[string, any]
    Returns an iterator over all key-value pairs in the session state.

• Usage:
  Provides a view of mutable session state that merges persisted state with uncommitted changes during invocation.

Important Implementation Details

• Artifact Version Tracking:
  When saving artifacts via internalArtifacts.Save, the version returned by the save operation is recorded in eventActions.ArtifactDelta keyed by artifact name. This allows the system to track what versions have been updated during the invocation lifecycle, essential for event generation and potential conflict resolution.

• State Delta Management:
  The callbackContext maintains a mutable state delta map inside session.EventActions.StateDelta. Reads first check this delta before falling back to the session's persisted state. Writes update both the delta and the underlying session state, ensuring changes are staged and can be committed or reverted as part of the lifecycle.

• Concurrency Considerations:
  The Save method includes a TODO comment noting the need for a read-write lock or concurrency control to ensure that only newer artifact versions replace older ones, especially when multiple tools might save the same artifact concurrently.

• ReadOnlyContext Wrapping:
  The callback context embeds a read-only context created from the invocation context (NewReadonlyContext(ctx)) to provide safe access to immutable data like agent name and user content, while layering on mutation capabilities.

Interaction With Other System Components

• InvocationContext (from Agent Invocation Context):
  The callback context wraps an existing invocation context, extending it to support state mutation and artifact saving with version tracking.

• Session Management (Session Management):
  Uses the session interface to get and set session state, augmented with mutable deltas for the current invocation.

• Artifact Management (Artifact Management):
  Delegates artifact saving to the underlying artifact service, adding additional tracking of artifact version deltas.

• Agent Lifecycle and Callbacks (Agent Lifecycle and Callbacks):
  Designed to be used during callback phases in the agent lifecycle where mutable state and artifact changes are needed.

Usage Examples

Creating a New Callback Context

// Create a callback context from an existing invocation context with an empty state delta.
cbCtx := NewCallbackContext(invocationCtx)

// Or create with an existing state delta.
stateDelta := map[string]any{"lastResult": "success"}
cbCtx := NewCallbackContextWithDelta(invocationCtx, stateDelta)

Reading and Writing State

// Get a session state value, checking both delta and persisted state.
val, err := cbCtx.State().Get("someKey")

// Set a session state value, updating the delta and session.
err := cbCtx.State().Set("someKey", "newValue")

Saving an Artifact and Tracking Version

resp, err := cbCtx.Artifacts().Save(ctx, "config.json", contentPart)
if err != nil {
    // Handle error
}
// The artifact version is tracked internally for event processing.

Structure Diagram

classDiagram
class internalArtifacts {
-Artifacts
-eventActions
+Save()
}
class callbackContext {
-ReadonlyContext
-artifacts
-invocationCtx
-eventActions
+Artifacts()
+AgentName()
+ReadonlyState()
+State()
+InvocationID()
+UserContent()
}
class callbackContextState {
-ctx
+Get()
+Set()
+All()
}
callbackContext o-- internalArtifacts : uses
callbackContext "1" *-- callbackContextState : provides

This diagram depicts the relationship between the main entities:

• callbackContext composes a ReadonlyContext and an internalArtifacts wrapper.

• internalArtifacts wraps the underlying artifact service adding version tracking.

• callbackContextState provides mutable state access layered over the callback context.

References

• Agent Invocation Context — for overall context interfaces and lifecycle.

• Session Management — for session state interfaces.

• Artifact Management — for artifact service interfaces.

• Agent Lifecycle and Callbacks — for callback integration and event handling.