diff --git a/rust/crates/runtime/src/session.rs b/rust/crates/runtime/src/session.rs index 6d6ff8fb..79a3d684 100644 --- a/rust/crates/runtime/src/session.rs +++ b/rust/crates/runtime/src/session.rs @@ -13,6 +13,9 @@ use serde::{Deserialize, Serialize}; const SESSION_VERSION: u32 = 1; const ROTATE_AFTER_BYTES: u64 = 256 * 1024; const MAX_ROTATED_FILES: usize = 3; +const MAX_JSONL_FIELD_CHARS: usize = 16 * 1024; +const JSONL_TRUNCATION_MARKER: &str = "… [truncated for session JSONL]"; +const JSONL_REDACTION_MARKER: &str = "[redacted]"; static SESSION_ID_COUNTER: AtomicU64 = AtomicU64::new(0); static LAST_TIMESTAMP_MS: AtomicU64 = AtomicU64::new(0); @@ -920,7 +923,7 @@ impl SessionCompaction { ); object.insert( "summary".to_string(), - JsonValue::String(self.summary.clone()), + JsonValue::String(sanitize_jsonl_field(&self.summary)), ); Ok(JsonValue::Object(object)) } @@ -980,7 +983,10 @@ impl SessionPromptEntry { "timestamp_ms".to_string(), JsonValue::Number(i64::try_from(self.timestamp_ms).unwrap_or(i64::MAX)), ); - object.insert("text".to_string(), JsonValue::String(self.text.clone())); + object.insert( + "text".to_string(), + JsonValue::String(sanitize_jsonl_field(&self.text)), + ); JsonValue::Object(object) } @@ -998,10 +1004,156 @@ impl SessionPromptEntry { fn message_record(message: &ConversationMessage) -> JsonValue { let mut object = BTreeMap::new(); object.insert("type".to_string(), JsonValue::String("message".to_string())); - object.insert("message".to_string(), message.to_json()); + object.insert("message".to_string(), persisted_message_json(message)); JsonValue::Object(object) } +fn persisted_message_json(message: &ConversationMessage) -> JsonValue { + let mut object = BTreeMap::new(); + object.insert( + "role".to_string(), + JsonValue::String( + match message.role { + MessageRole::System => "system", + MessageRole::User => "user", + MessageRole::Assistant => "assistant", + MessageRole::Tool => "tool", + } + .to_string(), + ), + ); + object.insert( + "blocks".to_string(), + JsonValue::Array(message.blocks.iter().map(persisted_block_json).collect()), + ); + if let Some(usage) = message.usage { + object.insert("usage".to_string(), usage_to_json(usage)); + } + JsonValue::Object(object) +} + +fn persisted_block_json(block: &ContentBlock) -> JsonValue { + let mut object = BTreeMap::new(); + match block { + ContentBlock::Text { text } => { + object.insert("type".to_string(), JsonValue::String("text".to_string())); + object.insert("text".to_string(), JsonValue::String(sanitize_jsonl_field(text))); + } + ContentBlock::Thinking { + thinking, + signature, + } => { + object.insert( + "type".to_string(), + JsonValue::String("thinking".to_string()), + ); + object.insert( + "thinking".to_string(), + JsonValue::String(sanitize_jsonl_field(thinking)), + ); + if let Some(signature) = signature { + object.insert( + "signature".to_string(), + JsonValue::String(sanitize_jsonl_field(signature)), + ); + } + } + ContentBlock::ToolUse { id, name, input } => { + object.insert( + "type".to_string(), + JsonValue::String("tool_use".to_string()), + ); + object.insert("id".to_string(), JsonValue::String(sanitize_jsonl_field(id))); + object.insert("name".to_string(), JsonValue::String(name.clone())); + object.insert( + "input".to_string(), + JsonValue::String(sanitize_jsonl_field(input)), + ); + } + ContentBlock::ToolResult { + tool_use_id, + tool_name, + output, + is_error, + } => { + object.insert( + "type".to_string(), + JsonValue::String("tool_result".to_string()), + ); + object.insert( + "tool_use_id".to_string(), + JsonValue::String(sanitize_jsonl_field(tool_use_id)), + ); + object.insert("tool_name".to_string(), JsonValue::String(tool_name.clone())); + object.insert( + "output".to_string(), + JsonValue::String(sanitize_jsonl_field(output)), + ); + object.insert("is_error".to_string(), JsonValue::Bool(*is_error)); + } + } + JsonValue::Object(object) +} + +fn sanitize_jsonl_field(value: &str) -> String { + truncate_jsonl_field(&redact_jsonl_secrets(value)) +} + +fn truncate_jsonl_field(value: &str) -> String { + let char_count = value.chars().count(); + if char_count <= MAX_JSONL_FIELD_CHARS { + return value.to_string(); + } + + let keep = MAX_JSONL_FIELD_CHARS.saturating_sub(JSONL_TRUNCATION_MARKER.chars().count()); + let mut truncated = value.chars().take(keep).collect::(); + truncated.push_str(JSONL_TRUNCATION_MARKER); + truncated +} + +fn redact_jsonl_secrets(value: &str) -> String { + let mut redacted = value.to_string(); + for marker in [ + "ANTHROPIC_API_KEY=", + "ANTHROPIC_AUTH_TOKEN=", + "OPENAI_API_KEY=", + "DASHSCOPE_API_KEY=", + "XAI_API_KEY=", + "Authorization: Bearer ", + "authorization: Bearer ", + "Bearer sk-", + "sk-ant-", + ] { + redacted = redact_after_marker(&redacted, marker); + } + redacted +} + +fn redact_after_marker(value: &str, marker: &str) -> String { + let mut output = String::with_capacity(value.len()); + let mut rest = value; + + while let Some(index) = rest.find(marker) { + let (before, after_before) = rest.split_at(index); + output.push_str(before); + output.push_str(marker); + output.push_str(JSONL_REDACTION_MARKER); + + let secret_start = marker.len(); + let after_marker = &after_before[secret_start..]; + let secret_end = after_marker + .char_indices() + .find_map(|(idx, ch)| { + (ch.is_whitespace() || matches!(ch, '\'' | '"' | ',' | '}' | ']')).then_some(idx) + }) + .unwrap_or(after_marker.len()); + rest = &after_marker[secret_end..]; + } + + output.push_str(rest); + output +} + fn usage_to_json(usage: TokenUsage) -> JsonValue { let mut object = BTreeMap::new(); object.insert(