feat(agent sessions): attribute stored KV cache blocks to sessions (#37482)

Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
This commit is contained in:
ishandhanani
2026-09-14 21:34:49 -07:00
committed by GitHub
parent c9fbe5f655
commit 3f871a246c
22 changed files with 905 additions and 451 deletions
@@ -27,9 +27,9 @@ pub async fn make_pub_bound() -> (PubSocket, u16) {
(sock, port)
}
/// Encode a single `BlockStored` event in the wire format msgspec
/// emits. Layout: `["BlockStored", block_hashes, parent, token_ids,
/// block_size, lora_id, medium]`.
/// Encode a single `BlockStored` event in the wire format msgspec emits: a
/// tagged map `{"type": "BlockStored", "block_hashes", "parent_block_hash",
/// "token_ids", "block_size", "lora_id", "medium"}`.
pub fn encode_block_stored_event(
block_hashes: &[i64],
parent: Option<i64>,
@@ -37,24 +37,31 @@ pub fn encode_block_stored_event(
block_size: u32,
) -> Vec<u8> {
let mut buf = Vec::new();
mp::write_array_len(&mut buf, 7).unwrap();
mp::write_map_len(&mut buf, 7).unwrap();
mp::write_str(&mut buf, "type").unwrap();
mp::write_str(&mut buf, "BlockStored").unwrap();
mp::write_str(&mut buf, "block_hashes").unwrap();
mp::write_array_len(&mut buf, block_hashes.len() as u32).unwrap();
for v in block_hashes {
mp::write_sint(&mut buf, *v).unwrap();
}
mp::write_str(&mut buf, "parent_block_hash").unwrap();
match parent {
Some(v) => {
mp::write_sint(&mut buf, v).unwrap();
}
None => mp::write_nil(&mut buf).unwrap(),
}
mp::write_str(&mut buf, "token_ids").unwrap();
mp::write_array_len(&mut buf, token_ids.len() as u32).unwrap();
for v in token_ids {
mp::write_uint(&mut buf, *v as u64).unwrap();
}
mp::write_str(&mut buf, "block_size").unwrap();
mp::write_uint(&mut buf, block_size as u64).unwrap();
mp::write_nil(&mut buf).unwrap(); // lora_id
mp::write_str(&mut buf, "lora_id").unwrap();
mp::write_nil(&mut buf).unwrap();
mp::write_str(&mut buf, "medium").unwrap();
mp::write_str(&mut buf, "GPU").unwrap();
buf
}