# Session Module
The session module manages distributed block transfer sessions between
instances. Sessions coordinate the search, staging, and RDMA transfer of
KV cache blocks between a requesting instance (Prefill) and a serving
instance (Decode).
## Protocol Overview
### Onboard Protocol (InitiatorSession ↔ ResponderSession)
Multi-peer search and staging using `OnboardMessage`:
```text
Initiator (Prefill) Responder (Decode)
│ │
│──── CreateSession ────────────▶│
│ │ search local G2/G3
│◀─── G2Results ────────────────│
│◀─── G3Results ────────────────│ (if G3 blocks found)
│◀─── SearchComplete ───────────│
│ │
│──── HoldBlocks ───────────────▶│
│◀─── Acknowledged ─────────────│
│ │
│──── StageBlocks ──────────────▶│ G3→G2 staging (optional)
│◀─── BlocksReady ──────────────│
│ │
│ RDMA pull (remote G2→local G2)
│ │
│──── CloseSession ─────────────▶│
```
When G4 (object storage) is configured, the initiator also runs a parallel
G4 search via internal `G4Results`/`G4LoadComplete` messages (not sent over
the network).
### Unified Session Protocol (SessionHandle ↔ ServerSession)
Point-to-point session using `SessionMessage`:
```text
Controller (Prefill) ServerSession (Decode)
│ │
│──── Attach ───────────────────▶│
│◀─── StateResponse ────────────│ (current state snapshot)
│ │
│──── TriggerStaging ───────────▶│ (if G3 blocks pending)
│◀─── BlocksStaged ────────────│ (newly staged blocks)
│ │
│ RDMA pull (remote G2→local G2)
│ │
│──── BlocksPulled ─────────────▶│ (release pulled blocks)
│──── Detach ───────────────────▶│
```
Control can be transferred bidirectionally via `YieldControl`/`AcquireControl`.
For layerwise transfer, `BlocksStaged` includes an optional `layer_range`.
## Session Types
| **ServerSession** | Holds blocks, exposes for pull | SessionMessage | Merged from former EndpointSession + ControllableSession |
| **SessionHandle** | Client-side control | SessionMessage | Attach/detach, state queries, RDMA pulls |
| **InitiatorSession** | Multi-peer search orchestrator | OnboardMessage | Created by InstanceLeader for distributed search |
| **ResponderSession** | Responds to search requests | OnboardMessage | Searches local G2/G3, holds blocks, stages on request |
### ServerSession
Server-side session that holds blocks and exposes them for remote RDMA pull.
Supports two modes:
- **G2-only**: Blocks are already in G2 with pre-assigned layout handles
(`BlockMetadataMap::Direct`). `TriggerStaging` is a no-op. Created via
`ServerSession::new_g2_only()` or the `create_server_session()` factory.
- **Staging**: G3 blocks need to be staged to G2. Layout handles are assigned
round-robin across workers (`BlockMetadataMap::RoundRobin`). Supports
`auto_stage` option via `ServerSessionOptions`. Created via
`ServerSession::new_with_staging()`.
Created from `InstanceLeader` via:
- `create_endpoint_session()` / `create_endpoint_session_for_blocks()` — G2-only
- `create_controllable_session()` / `create_controllable_session_with_options()` — with staging
**ServerSessionHandle** provides local control: `notify_layers_ready()` for
layerwise transfer notifications, and `close()` for graceful shutdown.
### InitiatorSession
The requesting side. Sends `CreateSession` to one or more remote instances,
collects results, applies first-responder-wins deduplication, and orchestrates
staging and RDMA pulls. Supports three staging modes:
- **Hold**: Find and hold blocks (G2+G3), no staging
- **Prepare**: Stage G3→G2 everywhere, keep session alive
- **Full**: Stage G3→G2 + RDMA pull remote G2→local G2, session completes
Created by `InstanceLeader::find_matches_with_options()` when
`search_remote == true`.
### ResponderSession
The serving side. Receives `CreateSession`, searches local block managers
(G2 then G3 for remaining), holds matched blocks via `BlockHolder`, and
responds with match results. Handles staging requests and keeps blocks
alive until the session ends.
## Core Building Blocks
### BlockHolder
RAII container for holding blocks during sessions. Tier-agnostic (`BlockHolder<G2>`,
`BlockHolder<G3>`). Blocks are automatically released when the holder is dropped,
preventing leaks even if session handling panics. Key operations: `retain()`,
`release()`, `extend()`, `take_all()`.
### SessionEndpoint
Point-to-point session primitive with a state machine. Encapsulates:
- Identity (`session_id`, `instance_id`)
- State machine (`ControlRole` + `AttachmentState` + `SessionPhase`)
- Message receive channel (`mpsc::Receiver<SessionMessage>`)
- State publication via watch channel
- Transport for sending messages to peer
Used internally by `ServerSession`. Does NOT handle block holding or staging
logic — those are the caller's responsibility.
### SessionHandle
Handle for controlling a remote session. Supports:
- State observation: `current_state()`, `wait_for_ready()`, `wait_for_complete()`
- Control commands: `trigger_staging()`, `mark_blocks_pulled()`, `detach()`
- Bidirectional control: `yield_control()`, `acquire_control()`
- RDMA transfers: `pull_blocks_rdma()`, `pull_blocks_rdma_with_options()`
Used by the controller side (Prefill) to drive a remote `ServerSession` (Decode).
### SessionHandleStateTx
Sender side of the state observation channel. Used by the session receiver
task to forward `StateResponse` and `BlocksStaged` messages into the
watch channel that `SessionHandle` observes.
### Staging
Shared G3→G2 staging logic extracted into `staging::stage_g3_to_g2()`.
Core kernel: allocate G2 destinations → execute local transfer (G3→G2) →
register new blocks with source sequence hashes. Used by `InitiatorSession`,
`ResponderSession`, and `ServerSession` to avoid code duplication.
## Transport Layer
`MessageTransport` is an enum with two variants:
- **`VeloTransport`**: Uses Velo active messages for distributed
communication between instances.
- **`LocalTransport`**: Direct channel dispatch for in-process testing
without network overhead.
Methods:
- `send()` — Send an `OnboardMessage` to a target instance
- `send_session()` — Send a `SessionMessage` to a target instance
- `request_metadata()` — RPC call to get remote worker layout metadata for RDMA
## Message Types
| `OnboardMessage` | Initiator ↔ Responder | Block search, hold/drop, staging requests |
| `SessionMessage` | Controller ↔ ServerSession | Attach/detach, control transfer, block ops, state sync |
### OnboardMessage Variants
| `CreateSession` | Initiator | Start new session with sequence hashes |
| `G2Results` | Responder | G2 search matches (sequence hashes + block IDs) |
| `G3Results` | Responder | G3 search matches (sequence hashes only) |
| `SearchComplete` | Responder | All local searching done |
| `HoldBlocks` | Initiator | Which blocks to hold vs drop |
| `Acknowledged` | Responder | Hold/drop processed |
| `StageBlocks` | Initiator | G3 hashes to stage to G2 |
| `BlocksReady` | Responder | Newly staged G2 blocks ready |
| `ReleaseBlocks` | Initiator | Release specific blocks |
| `CloseSession` | Initiator | Session complete, clean up |
| `G4Results` | Internal | Object storage search results (not sent over network) |
| `G4LoadComplete` | Internal | Object storage load results (not sent over network) |
### SessionMessage Variants
| `Attach` | Connection | Peer attaches with a control role |
| `Detach` | Connection | Peer detaches gracefully |
| `YieldControl` | Control | Sender yields controller role |
| `AcquireControl` | Control | Sender acquires controller role |
| `TriggerStaging` | Block ops | Request G3→G2 staging |
| `HoldBlocks` | Block ops | Request blocks be held |
| `ReleaseBlocks` | Block ops | Release specific blocks |
| `BlocksPulled` | Block ops | Notify blocks were pulled via RDMA |
| `StateResponse` | State sync | Full state snapshot (phase, role, blocks) |
| `BlocksStaged` | State sync | Newly staged blocks (with optional layer range) |
| `Close` | Lifecycle | Graceful session close |
| `Error` | Lifecycle | Report error |
## State Machine
### SessionPhase
Lifecycle of block operations. Staging is optional — blocks already in the
target tier (G2) skip it:
```text
Searching → Holding ──────────────────── Ready → Complete
└── Staging ────────┘
└── Complete (direct pull, no staging needed)
└── Failed
```
### ControlRole
Dynamic role in session relationship:
- `Neutral` — Initial state, can transition either way
- `Controller` — Issues commands to peer
- `Controllee` — Executes commands from peer
Supports bidirectional transfer via `YieldControl`/`AcquireControl`.
### AttachmentState
Peer connection state: `Unattached` (waiting) or `Attached { peer }` (connected).
## Dispatch Functions
- **`dispatch_onboard_message`**: Routes `OnboardMessage` to per-session task
channels by session ID. Used by the Velo onboard handler.
- **`dispatch_session_message`**: Routes `SessionMessage` to per-session task
channels by session ID. Used by the Velo session handler.
## File Structure
```text
session/
├── mod.rs # Module declarations, dispatch functions, re-exports
├── blocks.rs # BlockHolder<T> — RAII block container
├── endpoint.rs # SessionEndpoint — state machine primitive
├── handle.rs # SessionHandle + SessionHandleStateTx
├── server_session.rs # ServerSession + ServerSessionHandle + BlockMetadataMap
├── staging.rs # Shared stage_g3_to_g2() function
├── state.rs # SessionPhase, ControlRole, AttachmentState
├── messages.rs # OnboardMessage, SessionMessage, BlockInfo, etc.
├── initiator.rs # InitiatorSession (multi-peer orchestrator)
├── responder.rs # ResponderSession (search + hold + stage)
└── transport.rs # MessageTransport (Velo + Local)
```