Skip to main content

context.session

Comprehensive agent usage tracking and context management.

This module provides real-time monitoring and rich display capabilities for AI agents:

Usage Tracking:

  • Turn-by-turn token usage (input/output tokens, tool calls, duration)
  • Session-wide cumulative metrics across multiple turns
  • Per-step detailed tracking with tool names and timing
  • Support for both codemode and MCP tool monitoring

Context Management:

  • System prompts and tool definitions analysis
  • Message history processing for context window estimation
  • Token counting with tiktoken integration and fallback estimation
  • Context window utilization calculations

Rich Display:

  • Rich tables for CLI display with turn, step, and session breakdowns
  • Real-time duration tracking with millisecond precision
  • Per-step tool execution details with actual token distributions
  • Context snapshot creation for comprehensive usage analysis

Export & Analytics:

  • CSV export of detailed step-by-step execution data
  • Structured usage data for integration with external analytics
  • Support for both real-time monitoring and post-execution analysis

The core workflow involves:

  1. UsageTracker collects real usage data from agent runs
  2. StepRecord captures individual request/response cycles with timing
  3. ContextSnapshot provides rich display formatting with actual metrics
  4. Rich tables display comprehensive usage breakdowns for CLI interfaces

get_model_context_window​

def get_model_context_window(model: str) -> int

Get context window size for a model.

Arguments:

  • model - Model identifier (e.g., "anthropic:claude-sonnet-4-0", "openai:gpt-4o")

Returns:

Context window size in tokens. Defaults to 128000 if model not found.

count_tokens​

def count_tokens(text: str) -> int

Count tokens in text using tiktoken if available, else estimate.

Uses cl100k_base encoding (GPT-4/Claude compatible) when tiktoken is available. Falls back to ~4 chars per token estimation otherwise.

count_tokens_json​

def count_tokens_json(obj: Any) -> int

Count tokens for a JSON-serializable object.

ToolSnapshot Objects​

@dataclass
class ToolSnapshot()

Snapshot of a tool definition.

ToolDetailSnapshot Objects​

@dataclass
class ToolDetailSnapshot()

Detailed snapshot of a tool including full schema.

source_type​

"function", "mcp", "codemode", etc.

MessageDetailSnapshot Objects​

@dataclass
class MessageDetailSnapshot()

Detailed snapshot of a message with in_context tracking.

role​

"user", "assistant", "system", "tool"

in_context​

Whether this message is in the current context window

tool_name​

For tool calls/results

MessageSnapshot Objects​

@dataclass
class MessageSnapshot()

Snapshot of a single message.

role​

"user", "assistant", "system"

RequestUsageSnapshot Objects​

@dataclass
class RequestUsageSnapshot()

Snapshot of a single model request's usage.

request_num​

1-indexed request number

tool_names​

Tools called in this request

timestamp​

ISO format timestamp

turn_id​

Which turn this step belongs to

duration_ms​

Duration in milliseconds

TurnUsage Objects​

@dataclass
class TurnUsage()

Usage statistics for a single turn (user prompt + agent response).

duration_seconds​

Duration in seconds

SessionUsage Objects​

@dataclass
class SessionUsage()

Cumulative usage statistics across all turns in a session.

duration_seconds​

Total duration in seconds

usage_to_dict​

def usage_to_dict(usage: object) -> dict[str, object]

Convert a usage object (from pydantic-ai) to a dictionary.

UsageTracker Objects​

class UsageTracker()

Tracks usage statistics across a session.

This class encapsulates the logic for tracking token usage, tool calls, and other metrics across multiple turns in a conversation.

Usage: tracker = UsageTracker()

# After each turn:
tracker.record_turn(run_usage, codemode_counts)

# Get current stats:
turn_usage = tracker.get_turn_usage()
session_usage = tracker.get_session_usage()

__init__​

def __init__(codemode: bool = False, session_id: str | None = None)

Initialize the usage tracker.

Arguments:

  • codemode - Whether this is a codemode session (affects tool call tracking)
  • session_id - Optional session identifier for CSV export

start_turn​

def start_turn() -> None

Mark the start of a new turn.

record_turn​

def record_turn(run_usage: object,
                codemode_counts: dict[str, int] | None = None,
                turn_duration: float | None = None) -> dict[str, object]

Record usage from a completed turn.

Arguments:

  • run_usage - Usage object from agent run (run.usage())
  • codemode_counts - Optional dict with codemode/mcp tool call counts
  • turn_duration - Optional turn duration in seconds (auto-calculated if not provided)

Returns:

Dictionary with this turn's usage data

get_turn_data​

def get_turn_data() -> dict[str, object]

Get the current turn's usage data.

get_session_data​

def get_session_data() -> dict[str, object]

Get the session's cumulative usage data.

get_turn_usage​

def get_turn_usage() -> TurnUsage

Build TurnUsage dataclass from current turn data.

get_session_usage​

def get_session_usage() -> SessionUsage

Build SessionUsage dataclass from session data.

ContextSnapshot Objects​

@dataclass
class ContextSnapshot()

Complete snapshot of agent context.

history_tool_call_tokens​

Tool calls from previous turns

history_tool_return_tokens​

Tool returns from previous turns

current_tool_call_tokens​

Tool calls from current turn

current_tool_return_tokens​

Tool returns from current turn

tool_call_tokens​

Total tool call tokens

tool_return_tokens​

Total tool return tokens

history_user_tokens​

User messages from previous turns

history_assistant_tokens​

Assistant messages from previous turns

current_user_tokens​

Current user message

current_assistant_tokens​

Current assistant response

current_message_tokens​

Alias for current_user_tokens

user_message_tokens​

Total user: history_user + current_user

total_tokens​

Our estimate based on tiktoken

model_output_tokens​

Model's reported output tokens

per_request_usage​

Per-request token usage

context_window​

Default context window

get_context_total​

def get_context_total() -> int

Get total context tokens (what's in the context window).

get_context_percentage​

def get_context_percentage() -> float

Get context usage as percentage of context window.

get_history_total​

def get_history_total() -> int

Get total history tokens (including tool usage from previous turns).

to_dict​

def to_dict() -> dict[str, Any]

Convert to dictionary for API response.

to_table​

def to_table(show_context: bool = True) -> "Table"

Generate a Rich Table for CLI display.

Arguments:

  • show_context - Whether to show the CONTEXT section. Defaults to True.

    Returns a 2 or 3-section table:

    • CONTEXT: System prompts, tool definitions, history, current user (with % of context window)
    • THIS TURN: Tool calls, requests, input/output tokens for current turn
    • SESSION: Cumulative totals across all turns

FullContextSnapshot Objects​

@dataclass
class FullContextSnapshot()

Complete detailed snapshot of agent context with all configuration.

This includes everything needed to understand what's in the context window:

  • Model configuration (name, context window, settings)
  • System prompts (full text)
  • Tool definitions (full schemas, source code if available)
  • Message history (complete with in_context flags)
  • Memory blocks (if available)

to_dict​

def to_dict() -> dict[str, Any]

Convert to dictionary for API response.

list_available_tools​

async def list_available_tools(
        toolsets: Sequence[Any]
) -> list[tuple[str, str | None, dict[str, Any]]]

List all available tools from toolsets.

This is the preferred async method for extracting tool definitions. It properly handles CodemodeToolset's registry discovery.

Arguments:

  • toolsets - Sequence of pydantic-ai toolsets.

Returns:

List of (name, description, parameters_schema) tuples.

extract_context_snapshot​

def extract_context_snapshot(
        agent: Any,
        agent_id: str,
        context_window: int = 128000,
        message_history: Sequence[Any] | None = None,
        model_input_tokens: int | None = None,
        model_output_tokens: int | None = None,
        turn_usage: TurnUsage | None = None,
        session_usage: SessionUsage | None = None,
        tool_definitions: list[tuple[str, str | None, dict[str, Any]]]
    | None = None,
        turn_start_time: float | None = None) -> ContextSnapshot

Extract context snapshot from a pydantic-ai agent.

Arguments:

  • agent - The BaseAgent wrapper or pydantic_ai.Agent instance.
  • agent_id - The agent identifier.
  • context_window - The context window size for the model.
  • message_history - Optional message history from the agent run.
  • model_input_tokens - Optional model-reported input tokens (from result.usage).
  • model_output_tokens - Optional model-reported output tokens (from result.usage).
  • turn_usage - Optional turn usage data for current request.
  • session_usage - Optional cumulative session usage data.
  • tool_definitions - Optional pre-fetched tool definitions as list of (name, description, params_schema) tuples. Use this when MCP server tools are cached during the run but cleared after.

Returns:

ContextSnapshot with extracted information.

extract_full_context_snapshot​

def extract_full_context_snapshot(
        agent: Any,
        agent_id: str,
        context_window: int = 128000,
        message_history: Sequence[Any] | None = None) -> FullContextSnapshot

Extract detailed context snapshot from a pydantic-ai agent.

This provides complete introspection including:

  • Model configuration
  • Full system prompts
  • Tool definitions with schemas and source code
  • Complete message history with in_context flags

Arguments:

  • agent - The BaseAgent wrapper or pydantic_ai.Agent instance.
  • agent_id - The agent identifier.
  • context_window - The context window size for the model.
  • message_history - Optional message history from the agent run.

Returns:

FullContextSnapshot with all extracted information.

register_agent​

def register_agent(agent_id: str,
                   agent: Any,
                   info: dict[str, Any] | None = None) -> None

Register an agent for context snapshot lookups.

Arguments:

  • agent_id - The unique identifier for the agent.
  • agent - The agent instance (BaseAgent wrapper or pydantic_ai.Agent).
  • info - Optional additional info about the agent.

unregister_agent​

def unregister_agent(agent_id: str) -> None

Unregister an agent from the registry.

Arguments:

  • agent_id - The unique identifier for the agent.

get_agent_context_snapshot​

def get_agent_context_snapshot(agent_id: str) -> ContextSnapshot | None

Get context snapshot for a registered agent.

Arguments:

  • agent_id - The unique identifier for the agent.

Returns:

ContextSnapshot if agent is found, None otherwise.

get_agent_full_context_snapshot​

def get_agent_full_context_snapshot(
        agent_id: str) -> FullContextSnapshot | None

Get full detailed context snapshot for a registered agent.

This provides complete introspection including:

  • Model configuration
  • Full system prompts
  • Tool definitions with schemas and source code
  • Complete message history with in_context flags

Arguments:

  • agent_id - The unique identifier for the agent.

Returns:

FullContextSnapshot if agent is found, None otherwise.