codewhale_tui/rlm/
session.rs

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

codewhale_tui/rlm/
session.rs
