YiemAgent/learning/05-SESSION_MANAGEMENT.md

AgentCore.jl - Session Management Deep Dive

Session Architecture with Data Flow

┌─────────────────────────────────────────────────────────────────────────────┐
│                        Session Layer                                        │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│  Session = Tree of Entries                                                  │
│                                                                             │
│  Each entry represents a change in conversation state                       │
│                                                                             │
│  Branch Navigation:                                                         │
│  ┌─────┐     ┌─────┐      ┌─────┐      ┌─────┐     ┌─────┐                  │
│  │ E1  │────▶│ E2  │────▶│ E3  │────▶│ E4  │────▶│ E5  │ (current leaf)   │
│  └─────┘     └─────┘      └─────┘      └─────┘     └─────┘                  │
│      │           │           │           │           │                      │
│      ▼           ▼           ▼           ▼           ▼                      │
│   Message    Message   Compaction   Message    BranchSummary                │
│                                                                             │
│  Data Flow:                                                                 │
│  AgentMessage[] (AgentState.messages)                                       │
│      │                                                                      │
│      └─► appendMessage() → MessageEntry                                     │
│           └─► storage.appendEntry() → JSONL file                            │
│                                                                             │
│  To navigate to E2 (fork point):                                            │
│  Session.moveTo(E2)                                                         │
│  ┌─────┐      ┌─────┐     ┌─────┐      ┌─────┐     ┌─────┐                  │
│  │ E1  │────▶│ E2  │────▶│ E3' │────▶│ E4' │────▶│ E5' │ (new branch)     │
│  └─────┘      └─────┘     └─────┘      └─────┘     └─────┘                  │
│      │           │                                                          │
│      │           ▼ create BranchSummary                                     │
│      │      ┌─────┐                                                         │
│      └──────│ E6  │ (branch summary)                                        │
│             └─────┘                                                         │
└─────────────────────────────────────────────────────────────────────────────┘

Data Flow: AgentMessage → SessionTreeEntry

AgentState.messages::Vector{AgentMessage}
    │
    ├─► For each message in messages:
    │       │
    │       ▼
    │   ┌──────────────────────────────────────────────────────────────┐
    │   │ appendMessage(session, AgentMessage)                         │
    │   │   Input: message::AgentMessage                               │
    │   │   Output: entry_id::String                                   │
    │   │                                                              │
    │   │   Steps:                                                     │
    │   │   1. Create MessageEntry:                                    │
    │   │      - type: "message"                                       │
    │   │      - id: createEntryId(storage)                            │
    │   │      - parent_id: getLeafId(storage)                         │
    │   │      - timestamp: create_timestamp()                         │
    │   │      - message: copy(message)                                │
    │   │   2. storage.appendEntry(entry)                              │
    │   │      - Write JSONL line to file                              │
    │   │      - Update leaf_id                                        │
    │   │   3. Return entry.id                                         │
    │   └──────────────────────────────────────────────────────────────┘
    │
    └─► Entry stored in JSONL:
        {"type":"message","id":"msg_1","parent_id":null,"timestamp":"...","message":{...}}

Entry Types

Entry Types

abstract type SessionTreeEntry end

1. MessageEntry

struct MessageEntry <: SessionTreeEntry
    type::String           # "message"
    id::String             # Unique entry ID
    parent_id::Union{String, Nothing}
    timestamp::String      # ISO 8601 timestamp
    message::AgentMessage  # The actual message
end

Represents: A user, assistant, or tool message

2. ThinkingLevelChangeEntry

struct ThinkingLevelChangeEntry <: SessionTreeEntry
    type::String           # "thinking_level_change"
    id::String
    parent_id::Union{String, Nothing}
    timestamp::String
    thinking_level::String # "off", "minimal", "low", "medium", etc.
end

Represents: Change in model thinking level

3. ModelChangeEntry

struct ModelChangeEntry <: SessionTreeEntry
    type::String           # "model_change"
    id::String
    parent_id::Union{String, Nothing}
    timestamp::String
    provider::String       # "openai", "anthropic", etc.
    model_id::String       # Model identifier
end

Represents: Change in model

4. ActiveToolsChangeEntry

struct ActiveToolsChangeEntry <: SessionTreeEntry
    type::String           # "active_tools_change"
    id::String
    parent_id::Union{String, Nothing}
    timestamp::String
    active_tool_names::Vector{String}
end

Represents: Change in active tools

5. CompactionEntry

struct CompactionEntry <: SessionTreeEntry
    type::String           # "compaction"
    id::String
    parent_id::Union{String, Nothing}
    timestamp::String
    summary::String        # Summary of compacted history
    first_kept_entry_id::Union{String, Nothing}
    tokens_before::Int64   # Context size before compaction
    retained_tail::Union{Vector{AgentMessage}, Nothing}
    details::Union{Any, Nothing}
    usage::Union{Usage, Nothing}
    from_hook::Bool        # Whether triggered by hook
end

Represents: Context window compression

Key fields:

• summary: Summary of removed messages
• first_kept_entry_id: First entry that was kept
• tokens_before: Context size before compaction
• retained_tail: Messages kept after compaction point

6. BranchSummaryEntry

struct BranchSummaryEntry <: SessionTreeEntry
    type::String           # "branch_summary"
    id::String
    parent_id::Union{String, Nothing}
    timestamp::String
    from_id::String        # Branch point entry ID
    summary::String        # Summary of branch history
    details::Union{Any, Nothing}
    usage::Union{Usage, Nothing}
    from_hook::Bool
end

Represents: Branch point with summary

7. CustomEntry

struct CustomEntry <: SessionTreeEntry
    type::String           # Custom type
    id::String
    parent_id::Union{String, Nothing}
    timestamp::String
    custom_type::String
    data::Union{Any, Nothing}
end

Represents: Custom application-specific data

8. CustomMessageEntry

struct CustomMessageEntry <: SessionTreeEntry
    type::String
    id::String
    parent_id::Union{String, Nothing}
    timestamp::String
    custom_type::String
    content::String
    details::Union{Any, Nothing}
    display::Bool
end

Represents: Custom message to display to user

9. LabelEntry

struct LabelEntry <: SessionTreeEntry
    type::String
    id::String
    parent_id::Union{String, Nothing}
    timestamp::String
    target_id::String      # Entry being labeled
    label::Union{String, Nothing}
end

Represents: Label/note on an entry

10. SessionInfoEntry

struct SessionInfoEntry <: SessionTreeEntry
    type::String
    id::String
    parent_id::Union{String, Nothing}
    timestamp::String
    name::Union{String, Nothing}
end

Represents: Session metadata (name, etc.)

11. LeafEntry

struct LeafEntry <: SessionTreeEntry
    type::String
    id::String
    parent_id::Union{String, Nothing}
    timestamp::String
    target_id::Union{String, Nothing}
end

Represents: Change in current leaf (branch pointer)

Session Storage Interface

abstract type SessionStorage{T<:SessionMetadata} end

Storage Methods

# Metadata
getMetadata(storage::SessionStorage)::Promise{T}

# Leaf management
getLeafId(storage::SessionStorage)::Promise{Union{String, Nothing}}
setLeafId(storage::SessionStorage, leaf_id::String)::Promise{Nothing}

# Entry management
createEntryId(storage::SessionStorage)::Promise{String}
appendEntry(storage::SessionStorage, entry::SessionTreeEntry)::Promise{Nothing}
getEntry(storage::SessionStorage, id::String)::Promise{Union{SessionTreeEntry, Nothing}}

# Query
findEntries(storage::SessionStorage, type::String)::Promise{Vector{SessionTreeEntry}}
getLabel(storage::SessionStorage, id::String)::Promise{Union{String, Nothing}}
getSessionName(storage::SessionStorage)::Promise{Union{String, Nothing}}

# Branch navigation
getPathToRootOrCompaction(
    storage::SessionStorage,
    leaf_id::String,
)::Promise{Vector{SessionTreeEntry}}

getEntries(storage::SessionStorage, options::Dict{String, Any})::Promise{Vector{SessionTreeEntry}}

# Stats
getSessionStats(storage::SessionStorage)::Promise{SessionStats}

JsonlSessionStorage

┌─────────────────────────────────────────────────────────────────────────────┐
│                        JSONL Storage Format                                 │
└─────────────────────────────────────────────────────────────────────────────┘

File: session.jsonl

Entry 1 (Metadata):
{"type":"session","id":"meta_1","created_at":"2024-01-01T00:00:00Z","cwd":"/path","path":"/path/session.jsonl"}

Entry 2 (Message):
{"type":"message","id":"msg_1","parent_id":null,"timestamp":"2024-01-01T00:00:01Z","message":{"role":"user","content":[{"type":"text","text":"Hello"}]}}

Entry 3 (Thinking Level):
{"type":"thinking_level_change","id":"tl_1","parent_id":"msg_1","timestamp":"2024-01-01T00:00:02Z","thinking_level":"medium"}

Entry 4 (Model Change):
{"type":"model_change","id":"mc_1","parent_id":"tl_1","timestamp":"2024-01-01T00:00:03Z","provider":"openai","model_id":"gpt-4"}

Entry 5 (Compaction):
{"type":"compaction","id":"comp_1","parent_id":"mc_1","timestamp":"2024-01-01T00:00:04Z","summary":"Previous messages summarized...","first_kept_entry_id":"msg_3","tokens_before":100000,"tokens_after":50000}

Entry 6 (Branch Summary):
{"type":"branch_summary","id":"branch_1","parent_id":"comp_1","timestamp":"2024-01-01T00:00:05Z","from_id":"msg_3","summary":"Branch from message 3"}

Entry 7 (Active Tools):
{"type":"active_tools_change","id":"tools_1","parent_id":"branch_1","timestamp":"2024-01-01T00:00:06Z","active_tool_names":["bash","read"]}

Entry 8 (Leaf):
{"type":"leaf","id":"leaf_1","parent_id":"tools_1","timestamp":"2024-01-01T00:00:07Z","target_id":"msg_5"}

Notes:
- Each line is a JSON object (JSONL format)
- parent_id references previous entry (linked list structure)
- Leaf entry points to current position in tree
- To fork, create new branch from any entry

InMemorySessionStorage

mutable struct InMemorySessionStorage
    metadata::SessionMetadata
    leaf_id::Union{String, Nothing}
    entries::Dict{String, SessionTreeEntry}
    labels::Dict{String, String}
end

Purpose: Testing and temporary sessions

Advantages:

• Fast (no I/O)
• Easy to inspect
• Perfect for tests

Session Class

mutable struct Session{T<:SessionMetadata}
    storage::SessionStorage{T}
    context_build_options::SessionContextBuildOptions
end

Session Methods

appendMessage()

function appendMessage(session::Session, message::AgentMessage)::String
    entry = MessageEntry(
        "message",
        createEntryId(session.storage),
        getLeafId(session.storage),
        create_timestamp(),
        message,
    )
    return appendTypedEntry(session, entry)
end

Usage:

session = Session(storage)

# Add user message
user_id = appendMessage(session, UserMessage("user", [TextContent("Hello")], timestamp))

# Add assistant message
assistant_id = appendMessage(session, AssistantMessage(...))

# Add tool result
tool_id = appendMessage(session, ToolResultMessage(...))

appendThinkingLevelChange()

function appendThinkingLevelChange(
    session::Session,
    thinking_level::String,
)::String
    entry = ThinkingLevelChangeEntry(
        "thinking_level_change",
        createEntryId(session.storage),
        getLeafId(session.storage),
        create_timestamp(),
        thinking_level,
    )
    return appendTypedEntry(session, entry)
end

appendCompaction()

function appendCompaction(
    session::Session,
    summary::String,
    first_kept_entry_id::Union{String, Nothing},
    tokens_before::Int64,
    details::Union{Any, Nothing}=nothing,
    from_hook::Bool=false,
    usage::Union{Usage, Nothing}=nothing,
    retained_tail::Union{Vector{AgentMessage}, Nothing}=nothing,
)::String
    entry = CompactionEntry(
        "compaction",
        createEntryId(session.storage),
        getLeafId(session.storage),
        create_timestamp(),
        summary,
        first_kept_entry_id,
        tokens_before,
        retained_tail,
        details,
        usage,
        from_hook,
    )
    return appendTypedEntry(session, entry)
end

moveTo()

function moveTo(
    session::Session,
    entry_id::Union{String, Nothing},
    summary::Union{Dict{String, Any}, Nothing}=nothing,
)::Union{String, Nothing}
    # Set new leaf
    setLeafId(session.storage, entry_id)
    
    # Optionally create branch summary
    if !isnothing(summary)
        return appendTypedEntry(session, BranchSummaryEntry(
            "branch_summary",
            createEntryId(session.storage),
            entry_id,
            create_timestamp(),
            entry_id,
            summary["summary"],
            get(summary, "details", nothing),
            get(summary, "usage", nothing),
            get(summary, "from_hook", false),
        ))
    end
    
    return nothing
end

Usage:

# Fork from a specific point
session.moveTo(msg_3_id)

# Branch with summary
session.moveTo(
    msg_3_id,
    Dict(
        "summary" => "User wanted to focus on file operations",
        "details" => Dict("focus" => "files"),
    )
)

Build Session Context

function buildSessionContext(
    path_entries::Vector{SessionTreeEntry},
    options::SessionContextBuildOptions=SessionContextBuildOptions(),
)::SessionContext
    state = deriveSessionContextState(path_entries)
    context_entries = buildContextEntries(path_entries, options)
    messages = SessionTreeEntry[]
    for (i, entry) in enumerate(context_entries)
        append!(messages, sessionEntryToContextMessages(entry, i, context_entries, options))
    end
    return SessionContext(messages, state.thinking_level, state.model, state.active_tool_names)
end

Context Entry Transform

function defaultContextEntryTransform(
    path_entries::Vector{SessionTreeEntry},
)::Vector{SessionTreeEntry}
    compaction = nothing
    for entry in path_entries
        if entry isa CompactionEntry
            compaction = entry
            break
        end
    end
    
    if isnothing(compaction)
        return copy(path_entries)
    end
    
    # Include compaction entry
    entries = [compaction]
    
    # Include retained tail if present
    if !isnothing(compaction.retained_tail)
        compaction_idx = findfirst(e -> e.id == compaction.id, path_entries)
        append!(entries, path_entries[compaction_idx+1:end])
        return entries
    end
    
    # Otherwise include entries after first_kept_entry_id
    if !isnothing(compaction.first_kept_entry_id)
        found_first_kept = false
        compaction_idx = findfirst(e -> e.id == compaction.id, path_entries)
        for i in 1:compaction_idx-1
            entry = path_entries[i]
            if entry.id == compaction.first_kept_entry_id
                found_first_kept = true
            end
            if found_first_kept
                push!(entries, entry)
            end
        end
    end
    
    # Include entries after compaction
    compaction_idx = findfirst(e -> e.id == compaction.id, path_entries)
    append!(entries, path_entries[compaction_idx+1:end])
    
    return entries
end

Session Entry to Context Messages

function sessionEntryToContextMessages(
    entry::SessionTreeEntry,
    index::Int64,
    entries::Vector{SessionTreeEntry},
    options::SessionContextBuildOptions=SessionContextBuildOptions(),
)::Vector{AgentMessage}
    if entry isa MessageEntry
        return [entry.message]
    end
    
    if entry isa CustomMessageEntry
        return [createCustomMessage(
            entry.custom_type,
            entry.content,
            entry.display,
            entry.details,
            entry.timestamp,
        )]
    end
    
    if entry isa CompactionEntry
        messages = [createCompactionSummaryMessage(
            entry.summary,
            entry.tokens_before,
            entry.timestamp,
        )]
        if !isnothing(entry.retained_tail)
            append!(messages, entry.retained_tail)
        end
        return messages
    end
    
    if entry isa BranchSummaryEntry
        return [createBranchSummaryMessage(
            entry.summary,
            entry.from_id,
            entry.timestamp,
        )]
    end
    
    if entry isa CustomEntry
        # Custom projectors can transform custom entries
        if !isnothing(options.entry_projectors) && haskey(options.entry_projectors, entry.custom_type)
            projector = options.entry_projectors[entry.custom_type]
            return projector(entry, index, entries)
        end
        return AgentMessage[]
    end
    
    return AgentMessage[]
end

Branch Navigation

Scenario: User wants to explore a different path

Initial Branch (current path):
┌─────┐     ┌─────┐     ┌─────┐     ┌─────┐
│ E1  │────▶│ E2  │────▶│ E3  │────▶│ E4  │ (leaf)
└─────┘     └─────┘     └─────┘     └─────┘
  │           │           │           │
Message     Message   Compaction   Message

Step 1: Fork from E2
┌─────┐     ┌─────┐     ┌─────┐
│ E1  │────▶│ E2  │─────────────────┐
└─────┘     └─────┘                 │
  │           │                     │
  │           ▼ create BranchSummary│
  │      ┌─────┐                     │
  │      │ E5  │ (branch summary)   │
  │      └─────┘                     │
  └──────────────────────────────────┘
         (new branch from E2)

Step 2: Continue on new branch
┌─────┐     ┌─────┐     ┌─────┐     ┌─────┐     ┌─────┐
│ E1  │────▶│ E2  │────▶│ E3' │────▶│ E4' │────▶│ E5' │ (new leaf)
└─────┘     └─────┘     └─────┘     └─────┘     └─────┘

Current branch now is:
[ E1, E2, E3', E4', E5' ]

Original branch is:
[ E1, E2, E5 ]  (E3, E4 are now separate branch)

Key Points:
- Shared entries: E1, E2
- Branch point: E2
- Branch summary: E5 (points to E2)
- Each branch has independent tail

Compaction Strategy

Why Compaction?

LLM context windows have limits:

• GPT-4: 128K tokens
• Claude 2: 100K tokens
• Llama 2: 4K tokens

Problem: Conversations grow unbounded Solution: Compaction - summarize old messages

Compaction Process

# 1. Identify messages to compact
# - Keep recent N messages (e.g., last 2 turns)
# - Summarize everything before

# 2. Generate summary
# - Use LLM to summarize
# - Include key facts, decisions, user preferences

# 3. Create CompactionEntry
# - summary: The summary text
# - first_kept_entry_id: First entry that was NOT compacted
# - tokens_before: Context size before compaction
# - retained_tail: Messages kept after compaction point

# 4. Update storage
# - Append CompactionEntry
# - Update leaf to CompactionEntry

Compaction Example

# Before compaction (100K tokens):
[
  msg_1,  # User: "I need to set up a project"
  msg_2,  # Assistant: "Sure, what language?"
  msg_3,  # User: "Python"
  msg_4,  # Assistant: "I'll create a Python project"
  msg_5,  # User: "With FastAPI"
  msg_6,  # Assistant: "Creating FastAPI project..."
  msg_7,  # Tool: bash("mkdir myapp")
  msg_8,  # Tool: write("myapp/main.py", ...)
  msg_9,  # Assistant: "Project created!"
  msg_10, # User: "Can you add auth?"
  msg_11, # Assistant: "Adding auth..."
  msg_12, # User: "Use JWT"
  msg_13, # Assistant: "Implementing JWT..."
  msg_14, # Tool: bash("pip install jwt")
  msg_15, # Tool: write("myapp/auth.py", ...)
  msg_16, # Assistant: "Auth implemented!"
]

# After compaction (20K tokens):
[
  compaction_entry,  # Summary of msg_1 to msg_10
  msg_11,            # Keep recent messages
  msg_12,
  msg_13,
  msg_14,
  msg_15,
  msg_16,
]

# Compaction summary:
"""
Previous conversation summary:
- User wanted to create a Python project
- Chose FastAPI framework
- Assistant created project structure in myapp/
- User requested authentication
- Chose JWT for auth
- Assistant implemented JWT auth in myapp/auth.py
"""

Complete Session Example

using AgentCore

# 1. Create storage
storage = JsonlSessionStorage(
    SessionMetadata("session_1", "2024-01-01T00:00:00Z"),
    "/path/to/session.jsonl",
)

# 2. Create session
session = Session(storage)

# 3. Add messages
msg1_id = appendMessage(session, UserMessage("user", [TextContent("Hello")], timestamp))
msg2_id = appendMessage(session, AssistantMessage("assistant", [TextContent("Hi!")], ...))

# 4. Change thinking level
tl_id = appendThinkingLevelChange(session, "medium")

# 5. Change model
mc_id = appendModelChange(session, "openai", "gpt-4")

# 6. Add more messages
msg3_id = appendMessage(session, UserMessage("user", [TextContent("What can you do?")], timestamp))
msg4_id = appendMessage(session, AssistantMessage("assistant", [TextContent("I can...")], ...))

# 7. Compact context (100K tokens → 20K)
compact_id = appendCompaction(
    session,
    "User asked about capabilities and assistant explained",
    msg2_id,
    100000,
    Dict("summary_length" => 50),
    false,
    usage,
    [msg3, msg4],  # Retained tail
)

# 8. Fork and branch
session.moveTo(msg2_id)  # Go back to msg2

# 9. Create new branch
branch_id = appendBranchSummary(
    session,
    "User changed direction to focus on file operations",
    msg2_id,
    Dict("focus" => "files"),
)

# 10. Continue on new branch
msg5_id = appendMessage(session, UserMessage("user", [TextContent("Let's work with files")], timestamp))

# 11. Query session context
context = buildSessionContext(session)

# 12. Get stats
stats = getSessionStats(session)
println("Messages: $(stats.message_count)")
println("Total tokens: $(stats.total_tokens)")
println("Cost: $$(stats.cost_total)")

Best Practices

1. Use compaction for long conversations to stay within context limits
2. Create branch summaries when forking to document divergent paths
3. Retain tail messages after compaction for context
4. Track token usage to optimize compaction timing
5. Use InMemorySessionStorage for testing