30 KiB
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: session::Session, message::AgentMessage │
│ │ Output: entry_id::String │
│ │ │
│ │ Steps: │
│ │ 1. Create MessageEntry: │
│ │ - base: SessionTreeEntryBase(type, id, leaf_id, time) │
│ │ - message: the AgentMessage │
│ │ 2. storage.appendEntry(entry) │
│ │ - In-memory: push to entries vector, update by_id dict │
│ │ - JSONL: would append to file (TODO) │
│ │ 3. Return entry.id │
│ └──────────────────────────────────────────────────────────────┘
│
└─► Entry stored in JSONL (conceptual):
{"type":"message","id":"msg_1","parent_id":null,"timestamp":"...","message":{...}}
Entry Types
All entry types extend abstract type SessionTreeEntry end and embed a
base::SessionTreeEntryBase struct containing type, id, parent_id, and timestamp.
abstract type SessionTreeEntry end
struct SessionTreeEntryBase
type::String
id::String
parent_id::Union{String, Nothing}
timestamp::String
end
1. MessageEntry
struct MessageEntry <: SessionTreeEntry
base::SessionTreeEntryBase
message::AgentMessage
end
Represents: A user, assistant, or tool message
2. ThinkingLevelChangeEntry
struct ThinkingLevelChangeEntry <: SessionTreeEntry
base::SessionTreeEntryBase
thinking_level::String
end
Represents: Change in model thinking level
3. ModelChangeEntry
struct ModelChangeEntry <: SessionTreeEntry
base::SessionTreeEntryBase
provider::String
model_id::String
end
Represents: Change in model
4. ActiveToolsChangeEntry
struct ActiveToolsChangeEntry <: SessionTreeEntry
base::SessionTreeEntryBase
active_tool_names::Vector{String}
end
Represents: Change in active tools
5. CompactionEntry
struct CompactionEntry{T} <: SessionTreeEntry
base::SessionTreeEntryBase
summary::String
first_kept_entry_id::Union{String, Nothing}
tokens_before::Int64
retained_tail::Union{Vector{AgentMessage}, Nothing}
details::Union{T, Nothing}
usage::Union{Usage, Nothing}
from_hook::Bool
end
Represents: Context window compression
Key fields:
summary: Summary of removed messagesfirst_kept_entry_id: First entry that was kepttokens_before: Context size before compactionretained_tail: Messages kept after compaction point
6. BranchSummaryEntry
struct BranchSummaryEntry{T} <: SessionTreeEntry
base::SessionTreeEntryBase
from_id::String
summary::String
details::Union{T, Nothing}
usage::Union{Usage, Nothing}
from_hook::Bool
end
Represents: Branch point with summary
7. CustomEntry
struct CustomEntry{T} <: SessionTreeEntry
base::SessionTreeEntryBase
custom_type::String
data::Union{T, Nothing}
end
Represents: Custom application-specific data
8. CustomMessageEntry
struct CustomMessageEntry{T} <: SessionTreeEntry
base::SessionTreeEntryBase
custom_type::String
content::String
details::Union{T, Nothing}
display::Bool
end
Represents: Custom message to display to user
9. LabelEntry
struct LabelEntry <: SessionTreeEntry
base::SessionTreeEntryBase
target_id::String
label::Union{String, Nothing}
end
Represents: Label/note on an entry
10. SessionInfoEntry
struct SessionInfoEntry <: SessionTreeEntry
base::SessionTreeEntryBase
name::Union{String, Nothing}
end
Represents: Session metadata (name, etc.)
11. LeafEntry
struct LeafEntry <: SessionTreeEntry
base::SessionTreeEntryBase
target_id::Union{String, Nothing}
end
Represents: Change in current leaf (branch pointer)
Session Storage Interface
abstract type SessionStorage{T<:SessionMetadata} end
Storage Methods (actual implementation signatures)
# Metadata
getMetadata(storage::SessionStorage)::T
# Leaf management
getLeafId(storage::SessionStorage)::Union{String, Nothing}
setLeafId(storage::SessionStorage, leaf_id::Union{String, Nothing})::Nothing
# Entry management
createEntryId(storage::SessionStorage)::String
appendEntry(storage::SessionStorage, entry::SessionTreeEntry)::Nothing
getEntry(storage::SessionStorage, id::String)::Union{SessionTreeEntry, Nothing}
# Query
findEntries(storage::SessionStorage, type::String)::Vector{SessionTreeEntry}
getLabel(storage::SessionStorage, id::String)::Union{String, Nothing}
getSessionName(storage::SessionStorage)::Union{String, Nothing}
# Branch navigation
getPathToRootOrCompaction(storage::SessionStorage, leaf_id::Union{String, Nothing})::Vector{SessionTreeEntry}
getEntries(storage::SessionStorage, options::Dict{String, Any})::Vector{SessionTreeEntry}
# Stats
getSessionStats(storage::SessionStorage)::SessionStats
JsonlSessionStorage
mutable struct JsonlSessionStorage{T<:SessionMetadata} <: SessionStorage{T}
file_path::String
metadata::T
entries::Vector{SessionTreeEntry} # ordered list
by_id::Dict{String, SessionTreeEntry} # fast lookup by id
labels_by_id::Dict{String, String} # label cache
current_leaf_id::Union{String, Nothing} # current branch tip
end
┌─────────────────────────────────────────────────────────────────────────────┐
│ JSONL Storage Format │
└─────────────────────────────────────────────────────────────────────────────┘
File: session.jsonl (conceptual - not yet implemented)
Entry 1 (Metadata via SessionHeader):
{"type":"session","version":3,"id":"meta_1","timestamp":"...","cwd":"/path","parent_session":null,"metadata":{}}
Entry 2 (Message):
{"type":"message","id":"msg_1","parent_id":null,"timestamp":"...","message":{"role":"user",...}}
Entry 3 (Thinking Level):
{"type":"thinking_level_change","id":"tl_1","parent_id":"msg_1","timestamp":"...","thinking_level":"medium"}
Entry 4 (Model Change):
{"type":"model_change","id":"mc_1","parent_id":"tl_1","timestamp":"...","provider":"openai","model_id":"gpt-4"}
Entry 5 (Compaction):
{"type":"compaction","id":"comp_1","parent_id":"mc_1","timestamp":"...","summary":"...","first_kept_entry_id":"msg_3","tokens_before":100000}
Entry 6 (Branch Summary):
{"type":"branch_summary","id":"branch_1","parent_id":"comp_1","timestamp":"...","from_id":"msg_3","summary":"..."}
Entry 7 (Active Tools):
{"type":"active_tools_change","id":"tools_1","parent_id":"branch_1","timestamp":"...","active_tool_names":["bash","read"]}
Entry 8 (Leaf):
{"type":"leaf","id":"leaf_1","parent_id":"tools_1","timestamp":"...","target_id":"msg_5"}
Notes:
- Each line is a JSON object (JSONL format) - TODO: file I/O not yet implemented
- parent_id references previous entry (linked list structure)
- Leaf entry points to current position in tree
- To fork, create new branch from any entry
- In-memory mode uses Vector + Dict by_id for fast access
InMemorySessionStorage
mutable struct InMemorySessionStorage{T<:SessionMetadata} <: SessionStorage{T}
metadata::T
entries::Vector{SessionTreeEntry}
by_id::Dict{String, SessionTreeEntry}
labels_by_id::Dict{String, String}
leaf_id::Union{String, Nothing}
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
function Session(storage::SessionStorage, context_build_options=SessionContextBuildOptions(nothing, nothing))
new{typeof(storage.metadata)}(storage, context_build_options)
end
end
SessionContextBuildOptions
mutable struct SessionContextBuildOptions
entry_transforms::Union{Vector{Function}, Nothing}
entry_projectors::Union{Dict{String, Function}, Nothing}
end
Session Methods
appendMessage()
function appendMessage(session::Session, message::AgentMessage)::String
return appendTypedEntry(session, MessageEntry(
SessionTreeEntryBase("message", createEntryId(session.storage), getLeafId(session.storage), create_timestamp()),
message,
))
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
return appendTypedEntry(session, ThinkingLevelChangeEntry(
SessionTreeEntryBase("thinking_level_change", createEntryId(session.storage), getLeafId(session.storage), create_timestamp()),
thinking_level,
))
end
appendModelChange()
function appendModelChange(session::Session, provider::String, model_id::String)::String
return appendTypedEntry(session, ModelChangeEntry(
SessionTreeEntryBase("model_change", createEntryId(session.storage), getLeafId(session.storage), create_timestamp()),
provider,
model_id,
))
end
appendActiveToolsChange()
function appendActiveToolsChange(session::Session, active_tool_names::Vector{String})::String
return appendTypedEntry(session, ActiveToolsChangeEntry(
SessionTreeEntryBase("active_tools_change", createEntryId(session.storage), getLeafId(session.storage), create_timestamp()),
active_tool_names,
))
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
return appendTypedEntry(session, CompactionEntry(
SessionTreeEntryBase("compaction", createEntryId(session.storage), getLeafId(session.storage), create_timestamp()),
summary,
first_kept_entry_id,
tokens_before,
retained_tail,
details,
usage,
from_hook,
))
end
moveTo()
function moveTo(
session::Session,
entry_id::Union{String, Nothing},
summary::Union{Dict{String, Any}, Nothing}=nothing,
)::Union{String, Nothing}
# Validate entry exists
if !isnothing(entry_id) && isnothing(getEntry(session, entry_id))
throw(SessionError("not_found", "Entry $(entry_id) not found"))
end
# Set new leaf (creates a LeafEntry)
setLeafId(session.storage, entry_id)
# Optionally create branch summary
if isnothing(summary)
return nothing
end
return appendTypedEntry(session, BranchSummaryEntry(
SessionTreeEntryBase("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
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"),
)
)
How it works:
- Validates the target entry exists
- Calls
setLeafId()which creates aLeafEntrywithtarget_id = entry_id - If
summaryis provided, creates aBranchSummaryEntryas a child of the target entry - The new leaf now points to
entry_id, making it the root of a new branch
Build Session Context
function buildSessionContext(
path_entries::Vector{SessionTreeEntry},
options::SessionContextBuildOptions=SessionContextBuildOptions(nothing, nothing),
)::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
function deriveSessionContextState(path_entries::Vector{SessionTreeEntry})::Dict{String, Any}
thinking_level = "off"
model = nothing
active_tool_names = nothing
for entry in path_entries
if entry isa ThinkingLevelChangeEntry
thinking_level = entry.thinking_level
elseif entry isa ModelChangeEntry
model = Dict("provider" => entry.provider, "modelId" => entry.model_id)
elseif entry isa MessageEntry && entry.message.role == "assistant"
model = Dict("provider" => entry.message.provider, "modelId" => entry.message.model)
elseif entry isa ActiveToolsChangeEntry
active_tool_names = copy(entry.active_tool_names)
end
end
return Dict(
"thinking_level" => thinking_level,
"model" => model,
"active_tool_names" => 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
entries::Vector{SessionTreeEntry} = [compaction]
compaction_idx = findfirst(
(entry) -> entry isa CompactionEntry && entry.id == compaction.id,
path_entries,
)
if !isnothing(compaction.retained_tail)
for i in compaction_idx+1:length(path_entries)
push!(entries, path_entries[i])
end
return entries
end
if !isnothing(compaction.first_kept_entry_id)
found_first_kept = false
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
for i in compaction_idx+1:length(path_entries)
push!(entries, path_entries[i])
end
return entries
end
function buildContextEntries(
path_entries::Vector{SessionTreeEntry},
options::SessionContextBuildOptions=SessionContextBuildOptions(nothing, nothing),
)::Vector{SessionTreeEntry}
entries = defaultContextEntryTransform(path_entries)
if !isnothing(options.entry_transforms)
for transform in options.entry_transforms
entries = transform(entries)
end
end
return entries
end
Session Entry to Context Messages
function sessionEntryToContextMessages(
entry::SessionTreeEntry,
index::Int64,
entries::Vector{SessionTreeEntry},
options::SessionContextBuildOptions=SessionContextBuildOptions(nothing, nothing),
)::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
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
getPathToRootOrCompaction
Walks from a leaf back to the root, handling compaction entries:
# When encountering a CompactionEntry:
# - If retained_tail is set: stop (compaction covers the tail)
# - Otherwise: skip to first_kept_entry_id and continue walking
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
# - Leaf automatically points to CompactionEntry (leafIdAfterEntry)
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(
"/path/to/session.jsonl",
SessionHeader("session", 3, "session_1", created_at, "/path", nothing, nothing),
SessionTreeEntry[],
nothing,
)
# 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. Continue on new branch (moveTo creates branch summary when summary is provided)
branch_id = moveTo(
session,
msg2_id,
Dict("summary" => "User changed direction", "details" => 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 = buildContext(session)
# 12. Get stats
stats = getSessionStats(session)
println("Messages: $(stats.message_count)")
println("Total tokens: $(stats.total_tokens)")
println("Cost: \$(stats.cost_total)")
Session Repo Interface
Session Repository Methods
# Create a new session
create(repo::SessionRepo, options::TCreateOptions)::Session
# Open an existing session
open(repo::SessionRepo, metadata::TMetadata)::Session
# List sessions
list(repo::SessionRepo, options::TListOptions)::Vector{TMetadata}
# Delete a session
delete(repo::SessionRepo, metadata::TMetadata)::Nothing
# Fork a session (copy branch from entry)
fork(repo::SessionRepo, source::TMetadata, options::Dict{String, Any})::Session
JSONL vs In-Memory Repos
| Feature | JsonlSessionRepo | InMemorySessionRepo |
|---|---|---|
| Persistence | File-based (TODO) | In-memory only |
| Use case | Production | Testing |
| Fork | Not implemented | Uses getEntriesToFork |
| Metadata | JsonlSessionMetadata | SessionMetadata |
Fork Behavior (getEntriesToFork)
function getEntriesToFork(storage, options)::Vector{SessionTreeEntry}
# If no entryId specified, fork from current leaf (full copy)
if !haskey(options, :entryId) || isnothing(options[:entryId])
return getEntries(storage, Dict{String, Any}())
end
target = getEntry(storage, options[:entryId])
position = get(options, "position", "before")
if position == "at"
# Fork includes the target entry
effective_leaf_id = target.id
else
# Fork before the target (parent)
# Target must be a user message
if target isa MessageEntry && target.message.role != "user"
throw(SessionError("invalid_fork_target", "Not a user message"))
end
effective_leaf_id = target.parent_id
end
return getPathToRootOrCompaction(storage, effective_leaf_id)
end
Best Practices
- Use compaction for long conversations to stay within context limits
- Create branch summaries when forking to document divergent paths (via
moveTo()with summary) - Retain tail messages after compaction for context (
retained_tailfield) - Track token usage to optimize compaction timing
- Use InMemorySessionStorage for testing
- Use
getBranch(session)to get the current path from leaf to root/compaction - Use
buildContext(session)as the convenient Session method for building context - Use
mergeContextBuildOptions(session, options)to combine session-level and call-level transforms/projectors