use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use tokio::sync::Mutex;
use uuid::Uuid;
use crate::models::{ContentBlock, Message, SystemPrompt};
use crate::repl::PythonRuntime;
pub type SharedRlmSessionStore = Arc<Mutex<HashMap<String, Arc<Mutex<RlmSession>>>>>;
#[must_use]
pub fn new_shared_rlm_session_store() -> SharedRlmSessionStore {
Arc::new(Mutex::new(HashMap::new()))
}
#[derive(Debug)]
pub struct RlmSession {
pub name: String,
pub id: String,
pub kernel: Option<PythonRuntime>,
pub context_meta: ContextMeta,
pub config: RlmSessionConfig,
pub rpc_count: u32,
pub total_duration: Duration,
pub peak_var_count: usize,
pub final_count: usize,
pub created_at: Instant,
pub last_used_at: Instant,
pub context_path: PathBuf,
}
impl RlmSession {
#[must_use]
pub fn new(
name: String,
kernel: PythonRuntime,
context_meta: ContextMeta,
context_path: PathBuf,
) -> Self {
let now = Instant::now();
Self {
name,
id: format!("rlm:{}", Uuid::new_v4().simple()),
kernel: Some(kernel),
context_meta,
config: RlmSessionConfig::default(),
rpc_count: 0,
total_duration: Duration::ZERO,
peak_var_count: 0,
final_count: 0,
created_at: now,
last_used_at: now,
context_path,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextMeta {
pub length: usize,
#[serde(rename = "type")]
pub type_name: String,
pub preview_500: String,
pub sha256: String,
}
impl ContextMeta {
#[must_use]
pub fn from_body(body: &str, type_name: impl Into<String>) -> Self {
Self {
length: body.chars().count(),
type_name: type_name.into(),
preview_500: body.chars().take(500).collect(),
sha256: sha256_hex(body.as_bytes()),
}
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum OutputFeedback {
Full,
Metadata,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RlmSessionConfig {
pub output_feedback: OutputFeedback,
pub sub_query_timeout_secs: u64,
pub sub_rlm_max_depth: u32,
pub share_session: bool,
}
impl Default for RlmSessionConfig {
fn default() -> Self {
Self {
output_feedback: OutputFeedback::Full,
sub_query_timeout_secs: 120,
sub_rlm_max_depth: 1,
share_session: false,
}
}
}
pub fn write_context_file(body: &str) -> std::io::Result<PathBuf> {
let dir = std::env::temp_dir().join("deepseek_rlm_ctx");
std::fs::create_dir_all(&dir)?;
let path = dir.join(format!(
"session_{}_{}.txt",
std::process::id(),
Uuid::new_v4().simple()
));
std::fs::write(&path, body)?;
Ok(path)
}
#[derive(Debug, Clone)]
pub struct SessionObjectSnapshot {
pub session_id: String,
pub model: String,
pub workspace: PathBuf,
pub system_prompt: Option<SystemPrompt>,
pub messages: Vec<Message>,
}
impl SessionObjectSnapshot {
#[must_use]
pub fn new(
session_id: String,
model: String,
workspace: PathBuf,
system_prompt: Option<SystemPrompt>,
messages: Vec<Message>,
) -> Self {
Self {
session_id,
model,
workspace,
system_prompt,
messages,
}
}
#[must_use]
pub fn object_cards(&self) -> Vec<SessionObjectCard> {
let mut cards = Vec::new();
for object in self.base_objects() {
cards.push(SessionObjectCard::from_resolved(&object));
}
for index in 0..self.messages.len() {
if let Some(object) = self.resolve(&format!("session://active/messages/{index}")) {
cards.push(SessionObjectCard::from_resolved(&object));
}
}
cards
}
#[must_use]
pub fn resolve(&self, object_ref: &str) -> Option<ResolvedSessionObject> {
let normalized = normalize_session_object_ref(object_ref);
match normalized.as_str() {
"session://active/session" => Some(self.session_metadata_object()),
"session://active/system_prompt" => self.system_prompt_object(),
"session://active/transcript" => Some(self.transcript_object()),
"session://active/latest_user" => self.latest_user_object(),
_ => self.message_object(&normalized),
}
}
fn base_objects(&self) -> Vec<ResolvedSessionObject> {
let mut objects = vec![self.session_metadata_object()];
if let Some(object) = self.system_prompt_object() {
objects.push(object);
}
objects.push(self.transcript_object());
if let Some(object) = self.latest_user_object() {
objects.push(object);
}
objects
}
fn session_metadata_object(&self) -> ResolvedSessionObject {
let body = json!({
"session_id": self.session_id,
"model": self.model,
"workspace": self.workspace.display().to_string(),
"message_count": self.messages.len(),
"object_refs": {
"system_prompt": "session://active/system_prompt",
"transcript": "session://active/transcript",
"latest_user": "session://active/latest_user",
"message_prefix": "session://active/messages/"
}
})
.to_string();
ResolvedSessionObject::new(
"session://active/session",
"session_metadata",
"Active session metadata",
body,
)
}
fn system_prompt_object(&self) -> Option<ResolvedSessionObject> {
let prompt = self.system_prompt.as_ref()?;
Some(ResolvedSessionObject::new(
"session://active/system_prompt",
"system_prompt",
"Active system prompt",
render_system_prompt(prompt),
))
}
fn transcript_object(&self) -> ResolvedSessionObject {
let body = self
.messages
.iter()
.enumerate()
.map(|(index, message)| compact_message_json(index, message).to_string())
.collect::<Vec<_>>()
.join("\n");
ResolvedSessionObject::new(
"session://active/transcript",
"transcript",
"Active transcript as JSONL",
body,
)
}
fn latest_user_object(&self) -> Option<ResolvedSessionObject> {
self.messages
.iter()
.enumerate()
.rev()
.find(|(_, message)| message.role == "user")
.map(|(index, message)| message_resolved_object(index, message, "Latest user message"))
}
fn message_object(&self, normalized: &str) -> Option<ResolvedSessionObject> {
let index = normalized
.strip_prefix("session://active/messages/")?
.parse::<usize>()
.ok()?;
self.messages
.get(index)
.map(|message| message_resolved_object(index, message, "Transcript message"))
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct SessionObjectCard {
pub id: String,
pub kind: String,
pub title: String,
pub length: usize,
pub preview_500: String,
pub sha256: String,
}
impl SessionObjectCard {
#[must_use]
pub fn from_resolved(object: &ResolvedSessionObject) -> Self {
Self {
id: object.id.clone(),
kind: object.kind.clone(),
title: object.title.clone(),
length: object.body.chars().count(),
preview_500: object.body.chars().take(500).collect(),
sha256: sha256_hex(object.body.as_bytes()),
}
}
}
#[derive(Debug, Clone)]
pub struct ResolvedSessionObject {
pub id: String,
pub kind: String,
pub title: String,
pub body: String,
}
impl ResolvedSessionObject {
fn new(
id: impl Into<String>,
kind: impl Into<String>,
title: impl Into<String>,
body: impl Into<String>,
) -> Self {
Self {
id: id.into(),
kind: kind.into(),
title: title.into(),
body: body.into(),
}
}
}
fn normalize_session_object_ref(object_ref: &str) -> String {
let trimmed = object_ref.trim();
if trimmed.starts_with("session://") {
trimmed.to_string()
} else {
format!("session://active/{}", trimmed.trim_start_matches('/'))
}
}
fn render_system_prompt(prompt: &SystemPrompt) -> String {
match prompt {
SystemPrompt::Text(text) => text.clone(),
SystemPrompt::Blocks(blocks) => blocks
.iter()
.map(|block| block.text.as_str())
.collect::<Vec<_>>()
.join("\n\n"),
}
}
fn message_resolved_object(index: usize, message: &Message, title: &str) -> ResolvedSessionObject {
ResolvedSessionObject::new(
format!("session://active/messages/{index}"),
"message",
format!("{title} {index} ({})", message.role),
compact_message_json(index, message).to_string(),
)
}
fn compact_message_json(index: usize, message: &Message) -> Value {
json!({
"index": index,
"role": message.role,
"content": message.content.iter().map(compact_content_block).collect::<Vec<_>>(),
})
}
fn compact_content_block(block: &ContentBlock) -> Value {
match block {
ContentBlock::Text { text, .. } => json!({
"type": "text",
"text": text,
}),
ContentBlock::Thinking { thinking } => json!({
"type": "thinking",
"redacted": true,
"chars": thinking.chars().count(),
"sha256": sha256_hex(thinking.as_bytes()),
"preview_240": truncate_chars(thinking, 240),
}),
ContentBlock::ToolUse {
id,
name,
input,
caller,
} => json!({
"type": "tool_use",
"id": id,
"name": name,
"input": input,
"caller": caller,
}),
ContentBlock::ToolResult {
tool_use_id,
content,
is_error,
content_blocks,
} => {
let chars = content.chars().count();
let large = chars > 2_000;
json!({
"type": "tool_result",
"tool_use_id": tool_use_id,
"is_error": is_error,
"content": if large { Value::Null } else { Value::String(content.clone()) },
"content_preview": truncate_chars(content, 500),
"content_chars": chars,
"content_sha256": sha256_hex(content.as_bytes()),
"content_redacted": large,
"content_blocks": content_blocks,
})
}
ContentBlock::ServerToolUse { id, name, input } => json!({
"type": "server_tool_use",
"id": id,
"name": name,
"input": input,
}),
ContentBlock::ToolSearchToolResult {
tool_use_id,
content,
} => json!({
"type": "tool_search_tool_result",
"tool_use_id": tool_use_id,
"content": content,
}),
ContentBlock::CodeExecutionToolResult {
tool_use_id,
content,
} => json!({
"type": "code_execution_tool_result",
"tool_use_id": tool_use_id,
"content": content,
}),
}
}
fn truncate_chars(text: &str, max_chars: usize) -> String {
if text.chars().count() <= max_chars {
return text.to_string();
}
let take = max_chars.saturating_sub(3);
let mut out: String = text.chars().take(take).collect();
out.push_str("...");
out
}
#[must_use]
pub fn derive_session_name(source_hint: Option<&str>) -> String {
let hint = source_hint
.and_then(|raw| {
Path::new(raw)
.file_name()
.and_then(|name| name.to_str())
.or(Some(raw))
})
.unwrap_or("context");
let mut out = String::new();
for ch in hint.chars() {
if ch.is_ascii_alphanumeric() {
out.push(ch.to_ascii_lowercase());
} else if !out.ends_with('_') {
out.push('_');
}
if out.len() >= 48 {
break;
}
}
let out = out.trim_matches('_');
if out.is_empty() {
"context".to_string()
} else {
out.to_string()
}
}
fn sha256_hex(bytes: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(bytes);
format!("{:x}", hasher.finalize())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn derive_session_name_slugifies_path() {
assert_eq!(
derive_session_name(Some("src/Big File.rs")),
"big_file_rs".to_string()
);
}
#[test]
fn context_meta_hashes_and_previews_body() {
let meta = ContextMeta::from_body("abcdef", "text");
assert_eq!(meta.length, 6);
assert_eq!(meta.preview_500, "abcdef");
assert_eq!(
meta.sha256,
"bef57ec7f53a6d40beb640a780a639c83bc29ac8a9816f1fc6c5c6dcd93c4721"
);
}
#[test]
fn session_objects_expose_prompt_and_transcript_cards() {
let snapshot = SessionObjectSnapshot::new(
"session-1".to_string(),
"deepseek-v4-pro".to_string(),
PathBuf::from("/tmp/work"),
Some(SystemPrompt::Text("system body".to_string())),
vec![Message {
role: "user".to_string(),
content: vec![ContentBlock::Text {
text: "hello RLM".to_string(),
cache_control: None,
}],
}],
);
let cards = snapshot.object_cards();
assert!(
cards
.iter()
.any(|card| card.id == "session://active/system_prompt")
);
assert!(
cards
.iter()
.any(|card| card.id == "session://active/messages/0")
);
let transcript = snapshot
.resolve("session://active/transcript")
.expect("transcript object");
assert!(transcript.body.contains("hello RLM"));
}
#[test]
fn session_object_transcript_keeps_large_tool_results_compact() {
let large = "tool output\n".repeat(400);
let snapshot = SessionObjectSnapshot::new(
"session-1".to_string(),
"deepseek-v4-pro".to_string(),
PathBuf::from("/tmp/work"),
None,
vec![Message {
role: "user".to_string(),
content: vec![ContentBlock::ToolResult {
tool_use_id: "call_1".to_string(),
content: large.clone(),
is_error: None,
content_blocks: None,
}],
}],
);
let object = snapshot
.resolve("session://active/messages/0")
.expect("message object");
assert!(object.body.contains("\"content_redacted\":true"));
assert!(object.body.len() < large.len());
}
}