| 1 | //! Persistent RLM session state for the v0.8.33 head/hands tool surface. |
| 2 | |
| 3 | use std::collections::HashMap; |
| 4 | use std::path::{Path, PathBuf}; |
| 5 | use std::sync::Arc; |
| 6 | use std::time::{Duration, Instant}; |
| 7 | |
| 8 | use serde::{Deserialize, Serialize}; |
| 9 | use serde_json::{Value, json}; |
| 10 | use tokio::sync::Mutex; |
| 11 | use uuid::Uuid; |
| 12 | |
| 13 | use crate::repl::PythonRuntime; |
| 14 | use codewhale_models::{ContentBlock, Message, SystemPrompt}; |
| 15 | |
| 16 | pub type SharedRlmSessionStore = Arc<Mutex<HashMap<String, Arc<Mutex<RlmSession>>>>>; |
| 17 | |
| 18 | #[must_use] |
| 19 | pub fn new_shared_rlm_session_store() -> SharedRlmSessionStore { |
| 20 | Arc::new(Mutex::new(HashMap::new())) |
| 21 | } |
| 22 | |
| 23 | #[derive(Debug)] |
| 24 | pub struct RlmSession { |
| 25 | pub name: String, |
| 26 | pub id: String, |
| 27 | pub kernel: Option<PythonRuntime>, |
| 28 | pub context_meta: ContextMeta, |
| 29 | pub config: RlmSessionConfig, |
| 30 | pub rpc_count: u32, |
| 31 | pub total_duration: Duration, |
| 32 | pub peak_var_count: usize, |
| 33 | pub final_count: usize, |
| 34 | pub created_at: Instant, |
| 35 | pub last_used_at: Instant, |
| 36 | pub context_path: PathBuf, |
| 37 | } |
| 38 | |
| 39 | impl RlmSession { |
| 40 | #[must_use] |
| 41 | pub fn new( |
| 42 | name: String, |
| 43 | kernel: PythonRuntime, |
| 44 | context_meta: ContextMeta, |
| 45 | context_path: PathBuf, |
| 46 | ) -> Self { |
| 47 | let now = Instant::now(); |
| 48 | Self { |
| 49 | name, |
| 50 | id: format!("rlm:{}", Uuid::new_v4().simple()), |
| 51 | kernel: Some(kernel), |
| 52 | context_meta, |
| 53 | config: RlmSessionConfig::default(), |
| 54 | rpc_count: 0, |
| 55 | total_duration: Duration::ZERO, |
| 56 | peak_var_count: 0, |
| 57 | final_count: 0, |
| 58 | created_at: now, |
| 59 | last_used_at: now, |
| 60 | context_path, |
| 61 | } |
| 62 | } |
| 63 | } |
| 64 | |
| 65 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 66 | pub struct ContextMeta { |
| 67 | pub length: usize, |
| 68 | #[serde(rename = "type")] |
| 69 | pub type_name: String, |
| 70 | pub preview_500: String, |
| 71 | pub sha256: String, |
| 72 | } |
| 73 | |
| 74 | impl ContextMeta { |
| 75 | #[must_use] |
| 76 | pub fn from_body(body: &str, type_name: impl Into<String>) -> Self { |
| 77 | Self { |
| 78 | length: body.chars().count(), |
| 79 | type_name: type_name.into(), |
| 80 | preview_500: body.chars().take(500).collect(), |
| 81 | sha256: sha256_hex(body.as_bytes()), |
| 82 | } |
| 83 | } |
| 84 | } |
| 85 | |
| 86 | #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] |
| 87 | #[serde(rename_all = "snake_case")] |
| 88 | pub enum OutputFeedback { |
| 89 | Full, |
| 90 | Metadata, |
| 91 | } |
| 92 | |
| 93 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 94 | pub struct RlmSessionConfig { |
| 95 | pub output_feedback: OutputFeedback, |
| 96 | pub sub_query_timeout_secs: u64, |
| 97 | pub sub_rlm_max_depth: u32, |
| 98 | pub share_session: bool, |
| 99 | } |
| 100 | |
| 101 | impl Default for RlmSessionConfig { |
| 102 | fn default() -> Self { |
| 103 | Self { |
| 104 | output_feedback: OutputFeedback::Full, |
| 105 | sub_query_timeout_secs: 120, |
| 106 | sub_rlm_max_depth: 1, |
| 107 | share_session: false, |
| 108 | } |
| 109 | } |
| 110 | } |
| 111 | |
| 112 | pub fn write_context_file(body: &str) -> std::io::Result<PathBuf> { |
| 113 | let dir = std::env::temp_dir().join("deepseek_rlm_ctx"); |
| 114 | std::fs::create_dir_all(&dir)?; |
| 115 | let path = dir.join(format!( |
| 116 | "session_{}_{}.txt", |
| 117 | std::process::id(), |
| 118 | Uuid::new_v4().simple() |
| 119 | )); |
| 120 | std::fs::write(&path, body)?; |
| 121 | Ok(path) |
| 122 | } |
| 123 | |
| 124 | #[derive(Debug, Clone)] |
| 125 | pub struct SessionObjectSnapshot { |
| 126 | pub session_id: String, |
| 127 | pub model: String, |
| 128 | pub workspace: PathBuf, |
| 129 | pub system_prompt: Option<SystemPrompt>, |
| 130 | pub messages: Vec<Message>, |
| 131 | } |
| 132 | |
| 133 | impl SessionObjectSnapshot { |
| 134 | #[must_use] |
| 135 | pub fn new( |
| 136 | session_id: String, |
| 137 | model: String, |
| 138 | workspace: PathBuf, |
| 139 | system_prompt: Option<SystemPrompt>, |
| 140 | messages: Vec<Message>, |
| 141 | ) -> Self { |
| 142 | Self { |
| 143 | session_id, |
| 144 | model, |
| 145 | workspace, |
| 146 | system_prompt, |
| 147 | messages, |
| 148 | } |
| 149 | } |
| 150 | |
| 151 | #[must_use] |
| 152 | pub fn object_cards(&self) -> Vec<SessionObjectCard> { |
| 153 | let mut cards = Vec::new(); |
| 154 | for object in self.base_objects() { |
| 155 | cards.push(SessionObjectCard::from_resolved(&object)); |
| 156 | } |
| 157 | for index in 0..self.messages.len() { |
| 158 | if let Some(object) = self.resolve(&format!("session://active/messages/{index}")) { |
| 159 | cards.push(SessionObjectCard::from_resolved(&object)); |
| 160 | } |
| 161 | } |
| 162 | cards |
| 163 | } |
| 164 | |
| 165 | #[must_use] |
| 166 | pub fn resolve(&self, object_ref: &str) -> Option<ResolvedSessionObject> { |
| 167 | let normalized = normalize_session_object_ref(object_ref); |
| 168 | match normalized.as_str() { |
| 169 | "session://active/session" => Some(self.session_metadata_object()), |
| 170 | "session://active/system_prompt" => self.system_prompt_object(), |
| 171 | "session://active/transcript" => Some(self.transcript_object()), |
| 172 | "session://active/latest_user" => self.latest_user_object(), |
| 173 | _ => self.message_object(&normalized), |
| 174 | } |
| 175 | } |
| 176 | |
| 177 | fn base_objects(&self) -> Vec<ResolvedSessionObject> { |
| 178 | let mut objects = vec![self.session_metadata_object()]; |
| 179 | if let Some(object) = self.system_prompt_object() { |
| 180 | objects.push(object); |
| 181 | } |
| 182 | objects.push(self.transcript_object()); |
| 183 | if let Some(object) = self.latest_user_object() { |
| 184 | objects.push(object); |
| 185 | } |
| 186 | objects |
| 187 | } |
| 188 | |
| 189 | fn session_metadata_object(&self) -> ResolvedSessionObject { |
| 190 | let body = json!({ |
| 191 | "session_id": self.session_id, |
| 192 | "model": self.model, |
| 193 | "workspace": self.workspace.display().to_string(), |
| 194 | "message_count": self.messages.len(), |
| 195 | "object_refs": { |
| 196 | "system_prompt": "session://active/system_prompt", |
| 197 | "transcript": "session://active/transcript", |
| 198 | "latest_user": "session://active/latest_user", |
| 199 | "message_prefix": "session://active/messages/" |
| 200 | } |
| 201 | }) |
| 202 | .to_string(); |
| 203 | ResolvedSessionObject::new( |
| 204 | "session://active/session", |
| 205 | "session_metadata", |
| 206 | "Active session metadata", |
| 207 | body, |
| 208 | ) |
| 209 | } |
| 210 | |
| 211 | fn system_prompt_object(&self) -> Option<ResolvedSessionObject> { |
| 212 | let prompt = self.system_prompt.as_ref()?; |
| 213 | Some(ResolvedSessionObject::new( |
| 214 | "session://active/system_prompt", |
| 215 | "system_prompt", |
| 216 | "Active system prompt", |
| 217 | render_system_prompt(prompt), |
| 218 | )) |
| 219 | } |
| 220 | |
| 221 | fn transcript_object(&self) -> ResolvedSessionObject { |
| 222 | let body = self |
| 223 | .messages |
| 224 | .iter() |
| 225 | .enumerate() |
| 226 | .map(|(index, message)| compact_message_json(index, message).to_string()) |
| 227 | .collect::<Vec<_>>() |
| 228 | .join("\n"); |
| 229 | ResolvedSessionObject::new( |
| 230 | "session://active/transcript", |
| 231 | "transcript", |
| 232 | "Active transcript as JSONL", |
| 233 | body, |
| 234 | ) |
| 235 | } |
| 236 | |
| 237 | fn latest_user_object(&self) -> Option<ResolvedSessionObject> { |
| 238 | self.messages |
| 239 | .iter() |
| 240 | .enumerate() |
| 241 | .rev() |
| 242 | .find(|(_, message)| message.role == "user") |
| 243 | .map(|(index, message)| message_resolved_object(index, message, "Latest user message")) |
| 244 | } |
| 245 | |
| 246 | fn message_object(&self, normalized: &str) -> Option<ResolvedSessionObject> { |
| 247 | let index = normalized |
| 248 | .strip_prefix("session://active/messages/")? |
| 249 | .parse::<usize>() |
| 250 | .ok()?; |
| 251 | self.messages |
| 252 | .get(index) |
| 253 | .map(|message| message_resolved_object(index, message, "Transcript message")) |
| 254 | } |
| 255 | } |
| 256 | |
| 257 | #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] |
| 258 | pub struct SessionObjectCard { |
| 259 | pub id: String, |
| 260 | pub kind: String, |
| 261 | pub title: String, |
| 262 | pub length: usize, |
| 263 | pub preview_500: String, |
| 264 | pub sha256: String, |
| 265 | } |
| 266 | |
| 267 | impl SessionObjectCard { |
| 268 | #[must_use] |
| 269 | pub fn from_resolved(object: &ResolvedSessionObject) -> Self { |
| 270 | Self { |
| 271 | id: object.id.clone(), |
| 272 | kind: object.kind.clone(), |
| 273 | title: object.title.clone(), |
| 274 | length: object.body.chars().count(), |
| 275 | preview_500: object.body.chars().take(500).collect(), |
| 276 | sha256: sha256_hex(object.body.as_bytes()), |
| 277 | } |
| 278 | } |
| 279 | } |
| 280 | |
| 281 | #[derive(Debug, Clone)] |
| 282 | pub struct ResolvedSessionObject { |
| 283 | pub id: String, |
| 284 | pub kind: String, |
| 285 | pub title: String, |
| 286 | pub body: String, |
| 287 | } |
| 288 | |
| 289 | impl ResolvedSessionObject { |
| 290 | fn new( |
| 291 | id: impl Into<String>, |
| 292 | kind: impl Into<String>, |
| 293 | title: impl Into<String>, |
| 294 | body: impl Into<String>, |
| 295 | ) -> Self { |
| 296 | Self { |
| 297 | id: id.into(), |
| 298 | kind: kind.into(), |
| 299 | title: title.into(), |
| 300 | body: body.into(), |
| 301 | } |
| 302 | } |
| 303 | } |
| 304 | |
| 305 | fn normalize_session_object_ref(object_ref: &str) -> String { |
| 306 | let trimmed = object_ref.trim(); |
| 307 | if trimmed.starts_with("session://") { |
| 308 | trimmed.to_string() |
| 309 | } else { |
| 310 | format!("session://active/{}", trimmed.trim_start_matches('/')) |
| 311 | } |
| 312 | } |
| 313 | |
| 314 | fn render_system_prompt(prompt: &SystemPrompt) -> String { |
| 315 | match prompt { |
| 316 | SystemPrompt::Text(text) => text.clone(), |
| 317 | SystemPrompt::Blocks(blocks) => blocks |
| 318 | .iter() |
| 319 | .map(|block| block.text.as_str()) |
| 320 | .collect::<Vec<_>>() |
| 321 | .join("\n\n"), |
| 322 | } |
| 323 | } |
| 324 | |
| 325 | fn message_resolved_object(index: usize, message: &Message, title: &str) -> ResolvedSessionObject { |
| 326 | ResolvedSessionObject::new( |
| 327 | format!("session://active/messages/{index}"), |
| 328 | "message", |
| 329 | format!("{title} {index} ({})", message.role), |
| 330 | compact_message_json(index, message).to_string(), |
| 331 | ) |
| 332 | } |
| 333 | |
| 334 | fn compact_message_json(index: usize, message: &Message) -> Value { |
| 335 | json!({ |
| 336 | "index": index, |
| 337 | "role": message.role, |
| 338 | "content": message.content.iter().map(compact_content_block).collect::<Vec<_>>(), |
| 339 | }) |
| 340 | } |
| 341 | |
| 342 | fn compact_content_block(block: &ContentBlock) -> Value { |
| 343 | match block { |
| 344 | ContentBlock::Text { text, .. } => json!({ |
| 345 | "type": "text", |
| 346 | "text": text, |
| 347 | }), |
| 348 | ContentBlock::Thinking { thinking, .. } => json!({ |
| 349 | "type": "thinking", |
| 350 | "redacted": true, |
| 351 | "chars": thinking.chars().count(), |
| 352 | "sha256": sha256_hex(thinking.as_bytes()), |
| 353 | "preview_240": truncate_chars(thinking, 240), |
| 354 | }), |
| 355 | ContentBlock::ToolUse { |
| 356 | id, |
| 357 | name, |
| 358 | input, |
| 359 | caller, |
| 360 | .. |
| 361 | } => json!({ |
| 362 | "type": "tool_use", |
| 363 | "id": id, |
| 364 | "name": name, |
| 365 | "input": input, |
| 366 | "caller": caller, |
| 367 | }), |
| 368 | ContentBlock::ToolResult { |
| 369 | tool_use_id, |
| 370 | content, |
| 371 | is_error, |
| 372 | content_blocks, |
| 373 | .. |
| 374 | } => { |
| 375 | let chars = content.chars().count(); |
| 376 | let large = chars > 2_000; |
| 377 | json!({ |
| 378 | "type": "tool_result", |
| 379 | "tool_use_id": tool_use_id, |
| 380 | "is_error": is_error, |
| 381 | "content": if large { Value::Null } else { Value::String(content.clone()) }, |
| 382 | "content_preview": truncate_chars(content, 500), |
| 383 | "content_chars": chars, |
| 384 | "content_sha256": sha256_hex(content.as_bytes()), |
| 385 | "content_redacted": large, |
| 386 | "content_blocks": crate::image_attach::safe_tool_result_content_blocks( |
| 387 | content_blocks.as_deref(), |
| 388 | ), |
| 389 | }) |
| 390 | } |
| 391 | ContentBlock::ServerToolUse { id, name, input } => json!({ |
| 392 | "type": "server_tool_use", |
| 393 | "id": id, |
| 394 | "name": name, |
| 395 | "input": input, |
| 396 | }), |
| 397 | ContentBlock::ToolSearchToolResult { |
| 398 | tool_use_id, |
| 399 | content, |
| 400 | } => json!({ |
| 401 | "type": "tool_search_tool_result", |
| 402 | "tool_use_id": tool_use_id, |
| 403 | "content": content, |
| 404 | }), |
| 405 | ContentBlock::CodeExecutionToolResult { |
| 406 | tool_use_id, |
| 407 | content, |
| 408 | } => json!({ |
| 409 | "type": "code_execution_tool_result", |
| 410 | "tool_use_id": tool_use_id, |
| 411 | "content": content, |
| 412 | }), |
| 413 | ContentBlock::ImageUrl { .. } => serde_json::Value::Null, |
| 414 | } |
| 415 | } |
| 416 | |
| 417 | fn truncate_chars(text: &str, max_chars: usize) -> String { |
| 418 | if text.chars().count() <= max_chars { |
| 419 | return text.to_string(); |
| 420 | } |
| 421 | let take = max_chars.saturating_sub(3); |
| 422 | let mut out: String = text.chars().take(take).collect(); |
| 423 | out.push_str("..."); |
| 424 | out |
| 425 | } |
| 426 | |
| 427 | #[must_use] |
| 428 | pub fn derive_session_name(source_hint: Option<&str>) -> String { |
| 429 | let hint = source_hint |
| 430 | .and_then(|raw| { |
| 431 | Path::new(raw) |
| 432 | .file_name() |
| 433 | .and_then(|name| name.to_str()) |
| 434 | .or(Some(raw)) |
| 435 | }) |
| 436 | .unwrap_or("context"); |
| 437 | let mut out = String::new(); |
| 438 | for ch in hint.chars() { |
| 439 | if ch.is_ascii_alphanumeric() { |
| 440 | out.push(ch.to_ascii_lowercase()); |
| 441 | } else if !out.ends_with('_') { |
| 442 | out.push('_'); |
| 443 | } |
| 444 | if out.len() >= 48 { |
| 445 | break; |
| 446 | } |
| 447 | } |
| 448 | let out = out.trim_matches('_'); |
| 449 | if out.is_empty() { |
| 450 | "context".to_string() |
| 451 | } else { |
| 452 | out.to_string() |
| 453 | } |
| 454 | } |
| 455 | |
| 456 | fn sha256_hex(bytes: &[u8]) -> String { |
| 457 | crate::hashing::sha256_hex(bytes) |
| 458 | } |
| 459 | |
| 460 | #[cfg(test)] |
| 461 | mod tests { |
| 462 | use super::*; |
| 463 | use codewhale_models::Role; |
| 464 | |
| 465 | #[test] |
| 466 | fn derive_session_name_slugifies_path() { |
| 467 | assert_eq!( |
| 468 | derive_session_name(Some("src/Big File.rs")), |
| 469 | "big_file_rs".to_string() |
| 470 | ); |
| 471 | } |
| 472 | |
| 473 | #[test] |
| 474 | fn context_meta_hashes_and_previews_body() { |
| 475 | let meta = ContextMeta::from_body("abcdef", "text"); |
| 476 | assert_eq!(meta.length, 6); |
| 477 | assert_eq!(meta.preview_500, "abcdef"); |
| 478 | assert_eq!( |
| 479 | meta.sha256, |
| 480 | "bef57ec7f53a6d40beb640a780a639c83bc29ac8a9816f1fc6c5c6dcd93c4721" |
| 481 | ); |
| 482 | } |
| 483 | |
| 484 | #[test] |
| 485 | fn session_objects_expose_prompt_and_transcript_cards() { |
| 486 | let snapshot = SessionObjectSnapshot::new( |
| 487 | "session-1".to_string(), |
| 488 | "deepseek-v4-pro".to_string(), |
| 489 | PathBuf::from("/tmp/work"), |
| 490 | Some(SystemPrompt::Text("system body".to_string())), |
| 491 | vec![Message { |
| 492 | role: Role::User, |
| 493 | content: vec![ContentBlock::Text { |
| 494 | text: "hello RLM".to_string(), |
| 495 | cache_control: None, |
| 496 | }], |
| 497 | }], |
| 498 | ); |
| 499 | |
| 500 | let cards = snapshot.object_cards(); |
| 501 | assert!( |
| 502 | cards |
| 503 | .iter() |
| 504 | .any(|card| card.id == "session://active/system_prompt") |
| 505 | ); |
| 506 | assert!( |
| 507 | cards |
| 508 | .iter() |
| 509 | .any(|card| card.id == "session://active/messages/0") |
| 510 | ); |
| 511 | |
| 512 | let transcript = snapshot |
| 513 | .resolve("session://active/transcript") |
| 514 | .expect("transcript object"); |
| 515 | assert!(transcript.body.contains("hello RLM")); |
| 516 | } |
| 517 | |
| 518 | #[test] |
| 519 | fn session_object_transcript_keeps_large_tool_results_compact() { |
| 520 | let large = "tool output\n".repeat(400); |
| 521 | let snapshot = SessionObjectSnapshot::new( |
| 522 | "session-1".to_string(), |
| 523 | "deepseek-v4-pro".to_string(), |
| 524 | PathBuf::from("/tmp/work"), |
| 525 | None, |
| 526 | vec![Message { |
| 527 | role: Role::User, |
| 528 | content: vec![ContentBlock::ToolResult { |
| 529 | execution_id: None, |
| 530 | tool_use_id: "call_1".to_string(), |
| 531 | content: large.clone(), |
| 532 | is_error: None, |
| 533 | content_blocks: None, |
| 534 | }], |
| 535 | }], |
| 536 | ); |
| 537 | |
| 538 | let object = snapshot |
| 539 | .resolve("session://active/messages/0") |
| 540 | .expect("message object"); |
| 541 | assert!(object.body.contains("\"content_redacted\":true")); |
| 542 | assert!(object.body.len() < large.len()); |
| 543 | } |
| 544 | } |
| 545 |