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 ..
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 let chars = content.chars().count();
375 let large = chars > 2_000;
376 json!({
377 "type": "tool_result",
378 "tool_use_id": tool_use_id,
379 "is_error": is_error,
380 "content": if large { Value::Null } else { Value::String(content.clone()) },
381 "content_preview": truncate_chars(content, 500),
382 "content_chars": chars,
383 "content_sha256": sha256_hex(content.as_bytes()),
384 "content_redacted": large,
385 "content_blocks": crate::image_attach::safe_tool_result_content_blocks(
386 content_blocks.as_deref(),
387 ),
388 })
389 }
390 ContentBlock::ServerToolUse { id, name, input } => json!({
391 "type": "server_tool_use",
392 "id": id,
393 "name": name,
394 "input": input,
395 }),
396 ContentBlock::ToolSearchToolResult {
397 tool_use_id,
398 content,
399 } => json!({
400 "type": "tool_search_tool_result",
401 "tool_use_id": tool_use_id,
402 "content": content,
403 }),
404 ContentBlock::CodeExecutionToolResult {
405 tool_use_id,
406 content,
407 } => json!({
408 "type": "code_execution_tool_result",
409 "tool_use_id": tool_use_id,
410 "content": content,
411 }),
412 ContentBlock::ImageUrl { .. } => serde_json::Value::Null,
413 }
414}
415
416fn truncate_chars(text: &str, max_chars: usize) -> String {
417 if text.chars().count() <= max_chars {
418 return text.to_string();
419 }
420 let take = max_chars.saturating_sub(3);
421 let mut out: String = text.chars().take(take).collect();
422 out.push_str("...");
423 out
424}
425
426#[must_use]
427pub fn derive_session_name(source_hint: Option<&str>) -> String {
428 let hint = source_hint
429 .and_then(|raw| {
430 Path::new(raw)
431 .file_name()
432 .and_then(|name| name.to_str())
433 .or(Some(raw))
434 })
435 .unwrap_or("context");
436 let mut out = String::new();
437 for ch in hint.chars() {
438 if ch.is_ascii_alphanumeric() {
439 out.push(ch.to_ascii_lowercase());
440 } else if !out.ends_with('_') {
441 out.push('_');
442 }
443 if out.len() >= 48 {
444 break;
445 }
446 }
447 let out = out.trim_matches('_');
448 if out.is_empty() {
449 "context".to_string()
450 } else {
451 out.to_string()
452 }
453}
454
455fn sha256_hex(bytes: &[u8]) -> String {
456 crate::hashing::sha256_hex(bytes)
457}
458
459#[cfg(test)]
460mod tests {
461 use super::*;
462
463 #[test]
464 fn derive_session_name_slugifies_path() {
465 assert_eq!(
466 derive_session_name(Some("src/Big File.rs")),
467 "big_file_rs".to_string()
468 );
469 }
470
471 #[test]
472 fn context_meta_hashes_and_previews_body() {
473 let meta = ContextMeta::from_body("abcdef", "text");
474 assert_eq!(meta.length, 6);
475 assert_eq!(meta.preview_500, "abcdef");
476 assert_eq!(
477 meta.sha256,
478 "bef57ec7f53a6d40beb640a780a639c83bc29ac8a9816f1fc6c5c6dcd93c4721"
479 );
480 }
481
482 #[test]
483 fn session_objects_expose_prompt_and_transcript_cards() {
484 let snapshot = SessionObjectSnapshot::new(
485 "session-1".to_string(),
486 "deepseek-v4-pro".to_string(),
487 PathBuf::from("/tmp/work"),
488 Some(SystemPrompt::Text("system body".to_string())),
489 vec![Message {
490 role: "user".to_string(),
491 content: vec![ContentBlock::Text {
492 text: "hello RLM".to_string(),
493 cache_control: None,
494 }],
495 }],
496 );
497
498 let cards = snapshot.object_cards();
499 assert!(
500 cards
501 .iter()
502 .any(|card| card.id == "session://active/system_prompt")
503 );
504 assert!(
505 cards
506 .iter()
507 .any(|card| card.id == "session://active/messages/0")
508 );
509
510 let transcript = snapshot
511 .resolve("session://active/transcript")
512 .expect("transcript object");
513 assert!(transcript.body.contains("hello RLM"));
514 }
515
516 #[test]
517 fn session_object_transcript_keeps_large_tool_results_compact() {
518 let large = "tool output\n".repeat(400);
519 let snapshot = SessionObjectSnapshot::new(
520 "session-1".to_string(),
521 "deepseek-v4-pro".to_string(),
522 PathBuf::from("/tmp/work"),
523 None,
524 vec![Message {
525 role: "user".to_string(),
526 content: vec![ContentBlock::ToolResult {
527 tool_use_id: "call_1".to_string(),
528 content: large.clone(),
529 is_error: None,
530 content_blocks: None,
531 }],
532 }],
533 );
534
535 let object = snapshot
536 .resolve("session://active/messages/0")
537 .expect("message object");
538 assert!(object.body.contains("\"content_redacted\":true"));
539 assert!(object.body.len() < large.len());
540 }
541}