1use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6use std::time::{Duration, Instant};
7
8use serde::{Deserialize, Serialize};
9use serde_json::{Value, json};
10use tokio::sync::Mutex;
11use uuid::Uuid;
12
13use crate::models::{ContentBlock, Message, SystemPrompt};
14use crate::repl::PythonRuntime;
15
16pub type SharedRlmSessionStore = Arc<Mutex<HashMap<String, Arc<Mutex<RlmSession>>>>>;
17
18#[must_use]
19pub fn new_shared_rlm_session_store() -> SharedRlmSessionStore {
20 Arc::new(Mutex::new(HashMap::new()))
21}
22
23#[derive(Debug)]
24pub 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
39impl 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)]
66pub 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
74impl 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")]
88pub enum OutputFeedback {
89 Full,
90 Metadata,
91}
92
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub 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
101impl 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
112pub 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)]
125pub 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
133impl 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)]
258pub 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
267impl 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)]
282pub struct ResolvedSessionObject {
283 pub id: String,
284 pub kind: String,
285 pub title: String,
286 pub body: String,
287}
288
289impl 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
305fn 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
314fn 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
325fn 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
334fn 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
342fn 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 } => json!({
361 "type": "tool_use",
362 "id": id,
363 "name": name,
364 "input": input,
365 "caller": caller,
366 }),
367 ContentBlock::ToolResult {
368 tool_use_id,
369 content,
370 is_error,
371 content_blocks,
372 } => {
373 let chars = content.chars().count();
374 let large = chars > 2_000;
375 json!({
376 "type": "tool_result",
377 "tool_use_id": tool_use_id,
378 "is_error": is_error,
379 "content": if large { Value::Null } else { Value::String(content.clone()) },
380 "content_preview": truncate_chars(content, 500),
381 "content_chars": chars,
382 "content_sha256": sha256_hex(content.as_bytes()),
383 "content_redacted": large,
384 "content_blocks": crate::image_attach::safe_tool_result_content_blocks(
385 content_blocks.as_deref(),
386 ),
387 })
388 }
389 ContentBlock::ServerToolUse { id, name, input } => json!({
390 "type": "server_tool_use",
391 "id": id,
392 "name": name,
393 "input": input,
394 }),
395 ContentBlock::ToolSearchToolResult {
396 tool_use_id,
397 content,
398 } => json!({
399 "type": "tool_search_tool_result",
400 "tool_use_id": tool_use_id,
401 "content": content,
402 }),
403 ContentBlock::CodeExecutionToolResult {
404 tool_use_id,
405 content,
406 } => json!({
407 "type": "code_execution_tool_result",
408 "tool_use_id": tool_use_id,
409 "content": content,
410 }),
411 ContentBlock::ImageUrl { .. } => serde_json::Value::Null,
412 }
413}
414
415fn truncate_chars(text: &str, max_chars: usize) -> String {
416 if text.chars().count() <= max_chars {
417 return text.to_string();
418 }
419 let take = max_chars.saturating_sub(3);
420 let mut out: String = text.chars().take(take).collect();
421 out.push_str("...");
422 out
423}
424
425#[must_use]
426pub fn derive_session_name(source_hint: Option<&str>) -> String {
427 let hint = source_hint
428 .and_then(|raw| {
429 Path::new(raw)
430 .file_name()
431 .and_then(|name| name.to_str())
432 .or(Some(raw))
433 })
434 .unwrap_or("context");
435 let mut out = String::new();
436 for ch in hint.chars() {
437 if ch.is_ascii_alphanumeric() {
438 out.push(ch.to_ascii_lowercase());
439 } else if !out.ends_with('_') {
440 out.push('_');
441 }
442 if out.len() >= 48 {
443 break;
444 }
445 }
446 let out = out.trim_matches('_');
447 if out.is_empty() {
448 "context".to_string()
449 } else {
450 out.to_string()
451 }
452}
453
454fn sha256_hex(bytes: &[u8]) -> String {
455 crate::hashing::sha256_hex(bytes)
456}
457
458#[cfg(test)]
459mod tests {
460 use super::*;
461
462 #[test]
463 fn derive_session_name_slugifies_path() {
464 assert_eq!(
465 derive_session_name(Some("src/Big File.rs")),
466 "big_file_rs".to_string()
467 );
468 }
469
470 #[test]
471 fn context_meta_hashes_and_previews_body() {
472 let meta = ContextMeta::from_body("abcdef", "text");
473 assert_eq!(meta.length, 6);
474 assert_eq!(meta.preview_500, "abcdef");
475 assert_eq!(
476 meta.sha256,
477 "bef57ec7f53a6d40beb640a780a639c83bc29ac8a9816f1fc6c5c6dcd93c4721"
478 );
479 }
480
481 #[test]
482 fn session_objects_expose_prompt_and_transcript_cards() {
483 let snapshot = SessionObjectSnapshot::new(
484 "session-1".to_string(),
485 "deepseek-v4-pro".to_string(),
486 PathBuf::from("/tmp/work"),
487 Some(SystemPrompt::Text("system body".to_string())),
488 vec![Message {
489 role: "user".to_string(),
490 content: vec![ContentBlock::Text {
491 text: "hello RLM".to_string(),
492 cache_control: None,
493 }],
494 }],
495 );
496
497 let cards = snapshot.object_cards();
498 assert!(
499 cards
500 .iter()
501 .any(|card| card.id == "session://active/system_prompt")
502 );
503 assert!(
504 cards
505 .iter()
506 .any(|card| card.id == "session://active/messages/0")
507 );
508
509 let transcript = snapshot
510 .resolve("session://active/transcript")
511 .expect("transcript object");
512 assert!(transcript.body.contains("hello RLM"));
513 }
514
515 #[test]
516 fn session_object_transcript_keeps_large_tool_results_compact() {
517 let large = "tool output\n".repeat(400);
518 let snapshot = SessionObjectSnapshot::new(
519 "session-1".to_string(),
520 "deepseek-v4-pro".to_string(),
521 PathBuf::from("/tmp/work"),
522 None,
523 vec![Message {
524 role: "user".to_string(),
525 content: vec![ContentBlock::ToolResult {
526 tool_use_id: "call_1".to_string(),
527 content: large.clone(),
528 is_error: None,
529 content_blocks: None,
530 }],
531 }],
532 );
533
534 let object = snapshot
535 .resolve("session://active/messages/0")
536 .expect("message object");
537 assert!(object.body.contains("\"content_redacted\":true"));
538 assert!(object.body.len() < large.len());
539 }
540}