Selaa lähdekoodia

fix: align session tests with jsonl persistence

YeonGyu-Kim 2 kuukautta sitten
vanhempi
commit
12c364da34

+ 28 - 38
rust/crates/runtime/src/compact.rs

@@ -577,21 +577,17 @@ mod tests {
 
     #[test]
     fn keeps_previous_compacted_context_when_compacting_again() {
-        let initial_session = Session {
-            version: 1,
-            messages: vec![
-                ConversationMessage::user_text("Investigate rust/crates/runtime/src/compact.rs"),
-                ConversationMessage::assistant(vec![ContentBlock::Text {
-                    text: "I will inspect the compact flow.".to_string(),
-                }]),
-                ConversationMessage::user_text(
-                    "Also update rust/crates/runtime/src/conversation.rs",
-                ),
-                ConversationMessage::assistant(vec![ContentBlock::Text {
-                    text: "Next: preserve prior summary context during auto compact.".to_string(),
-                }]),
-            ],
-        };
+        let mut initial_session = Session::new();
+        initial_session.messages = vec![
+            ConversationMessage::user_text("Investigate rust/crates/runtime/src/compact.rs"),
+            ConversationMessage::assistant(vec![ContentBlock::Text {
+                text: "I will inspect the compact flow.".to_string(),
+            }]),
+            ConversationMessage::user_text("Also update rust/crates/runtime/src/conversation.rs"),
+            ConversationMessage::assistant(vec![ContentBlock::Text {
+                text: "Next: preserve prior summary context during auto compact.".to_string(),
+            }]),
+        ];
         let config = CompactionConfig {
             preserve_recent_messages: 2,
             max_estimated_tokens: 1,
@@ -606,13 +602,9 @@ mod tests {
             }]),
         ]);
 
-        let second = compact_session(
-            &Session {
-                version: 1,
-                messages: follow_up_messages,
-            },
-            config,
-        );
+        let mut second_session = Session::new();
+        second_session.messages = follow_up_messages;
+        let second = compact_session(&second_session, config);
 
         assert!(second
             .formatted_summary
@@ -641,22 +633,20 @@ mod tests {
     #[test]
     fn ignores_existing_compacted_summary_when_deciding_to_recompact() {
         let summary = "<summary>Conversation summary:\n- Scope: earlier work preserved.\n- Key timeline:\n  - user: large preserved context\n</summary>";
-        let session = Session {
-            version: 1,
-            messages: vec![
-                ConversationMessage {
-                    role: MessageRole::System,
-                    blocks: vec![ContentBlock::Text {
-                        text: get_compact_continuation_message(summary, true, true),
-                    }],
-                    usage: None,
-                },
-                ConversationMessage::user_text("tiny"),
-                ConversationMessage::assistant(vec![ContentBlock::Text {
-                    text: "recent".to_string(),
-                }]),
-            ],
-        };
+        let mut session = Session::new();
+        session.messages = vec![
+            ConversationMessage {
+                role: MessageRole::System,
+                blocks: vec![ContentBlock::Text {
+                    text: get_compact_continuation_message(summary, true, true),
+                }],
+                usage: None,
+            },
+            ConversationMessage::user_text("tiny"),
+            ConversationMessage::assistant(vec![ContentBlock::Text {
+                text: "recent".to_string(),
+            }]),
+        ];
 
         assert!(!should_compact(
             &session,

+ 128 - 32
rust/crates/runtime/src/conversation.rs

@@ -1,6 +1,7 @@
 use std::collections::BTreeMap;
 use std::fmt::{Display, Formatter};
 
+use serde_json::{Map, Value};
 use telemetry::SessionTracer;
 
 use crate::compact::{
@@ -319,13 +320,14 @@ where
                     return Err(error);
                 }
             };
-            let (assistant_message, usage) = match build_assistant_message(events) {
-                Ok(result) => result,
-                Err(error) => {
-                    self.record_turn_failed(iterations, &error);
-                    return Err(error);
-                }
-            };
+            let (assistant_message, usage, turn_prompt_cache_events) =
+                match build_assistant_message(events) {
+                    Ok(result) => result,
+                    Err(error) => {
+                        self.record_turn_failed(iterations, &error);
+                        return Err(error);
+                    }
+                };
             if let Some(usage) = usage {
                 self.usage_tracker.record(usage);
             }
@@ -397,6 +399,7 @@ where
 
                 let result_message = match permission_outcome {
                     PermissionOutcome::Allow => {
+                        self.record_tool_started(iterations, &tool_name);
                         let (mut output, mut is_error) =
                             match self.tool_executor.execute(&tool_name, &effective_input) {
                                 Ok(output) => (output, false),
@@ -439,20 +442,24 @@ where
                 self.session
                     .push_message(result_message.clone())
                     .map_err(|error| RuntimeError::new(error.to_string()))?;
+                self.record_tool_finished(iterations, &result_message);
                 tool_results.push(result_message);
             }
         }
 
         let auto_compaction = self.maybe_auto_compact();
 
-        Ok(TurnSummary {
+        let summary = TurnSummary {
             assistant_messages,
             tool_results,
             prompt_cache_events,
             iterations,
             usage: self.usage_tracker.cumulative_usage(),
             auto_compaction,
-        })
+        };
+        self.record_turn_completed(&summary);
+
+        Ok(summary)
     }
 
     #[must_use]
@@ -510,23 +517,112 @@ where
         })
     }
 
-    fn record_turn_started(&self, _user_input: &str) {}
+    fn record_turn_started(&self, user_input: &str) {
+        let Some(session_tracer) = &self.session_tracer else {
+            return;
+        };
+
+        let mut attributes = Map::new();
+        attributes.insert(
+            "user_input".to_string(),
+            Value::String(user_input.to_string()),
+        );
+        session_tracer.record("turn_started", attributes);
+    }
 
     fn record_assistant_iteration(
         &self,
-        _iteration: usize,
-        _assistant_message: &ConversationMessage,
-        _pending_tool_use_count: usize,
+        iteration: usize,
+        assistant_message: &ConversationMessage,
+        pending_tool_use_count: usize,
     ) {
+        let Some(session_tracer) = &self.session_tracer else {
+            return;
+        };
+
+        let mut attributes = Map::new();
+        attributes.insert("iteration".to_string(), Value::from(iteration as u64));
+        attributes.insert(
+            "assistant_blocks".to_string(),
+            Value::from(assistant_message.blocks.len() as u64),
+        );
+        attributes.insert(
+            "pending_tool_use_count".to_string(),
+            Value::from(pending_tool_use_count as u64),
+        );
+        session_tracer.record("assistant_iteration_completed", attributes);
+    }
+
+    fn record_tool_started(&self, iteration: usize, tool_name: &str) {
+        let Some(session_tracer) = &self.session_tracer else {
+            return;
+        };
+
+        let mut attributes = Map::new();
+        attributes.insert("iteration".to_string(), Value::from(iteration as u64));
+        attributes.insert(
+            "tool_name".to_string(),
+            Value::String(tool_name.to_string()),
+        );
+        session_tracer.record("tool_execution_started", attributes);
+    }
+
+    fn record_tool_finished(&self, iteration: usize, result_message: &ConversationMessage) {
+        let Some(session_tracer) = &self.session_tracer else {
+            return;
+        };
+
+        let Some(ContentBlock::ToolResult {
+            tool_name,
+            is_error,
+            ..
+        }) = result_message.blocks.first()
+        else {
+            return;
+        };
+
+        let mut attributes = Map::new();
+        attributes.insert("iteration".to_string(), Value::from(iteration as u64));
+        attributes.insert("tool_name".to_string(), Value::String(tool_name.clone()));
+        attributes.insert("is_error".to_string(), Value::Bool(*is_error));
+        session_tracer.record("tool_execution_finished", attributes);
     }
 
-    fn record_tool_started(&self, _iteration: usize, _tool_name: &str) {}
+    fn record_turn_completed(&self, summary: &TurnSummary) {
+        let Some(session_tracer) = &self.session_tracer else {
+            return;
+        };
 
-    fn record_tool_finished(&self, _iteration: usize, _result_message: &ConversationMessage) {}
+        let mut attributes = Map::new();
+        attributes.insert(
+            "iterations".to_string(),
+            Value::from(summary.iterations as u64),
+        );
+        attributes.insert(
+            "assistant_messages".to_string(),
+            Value::from(summary.assistant_messages.len() as u64),
+        );
+        attributes.insert(
+            "tool_results".to_string(),
+            Value::from(summary.tool_results.len() as u64),
+        );
+        attributes.insert(
+            "prompt_cache_events".to_string(),
+            Value::from(summary.prompt_cache_events.len() as u64),
+        );
+        session_tracer.record("turn_completed", attributes);
+    }
 
-    fn record_turn_completed(&self, _summary: &TurnSummary) {}
+    fn record_turn_failed(&self, iteration: usize, error: &RuntimeError) {
+        let Some(session_tracer) = &self.session_tracer else {
+            return;
+        };
 
-    fn record_turn_failed(&self, _iteration: usize, _error: &RuntimeError) {}
+        let mut attributes = Map::new();
+        attributes.insert("iteration".to_string(), Value::from(iteration as u64));
+        attributes.insert("error".to_string(), Value::String(error.to_string()));
+        session_tracer.record("turn_failed", attributes);
+    }
 }
 
 #[must_use]
@@ -665,8 +761,8 @@ impl ToolExecutor for StaticToolExecutor {
 mod tests {
     use super::{
         parse_auto_compaction_threshold, ApiClient, ApiRequest, AssistantEvent,
-        AutoCompactionEvent, ConversationRuntime, RuntimeError, StaticToolExecutor,
-        DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD,
+        AutoCompactionEvent, ConversationRuntime, PromptCacheEvent, RuntimeError,
+        StaticToolExecutor, DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD,
     };
     use crate::compact::CompactionConfig;
     use crate::config::{RuntimeFeatureConfig, RuntimeHookConfig};
@@ -679,7 +775,9 @@ mod tests {
     use crate::usage::TokenUsage;
     use std::fs;
     use std::path::PathBuf;
+    use std::sync::Arc;
     use std::time::{SystemTime, UNIX_EPOCH};
+    use telemetry::{MemoryTelemetrySink, SessionTracer, TelemetryEvent};
 
     struct ScriptedApiClient {
         call_count: usize,
@@ -1217,19 +1315,17 @@ mod tests {
             }
         }
 
-        let session = Session {
-            version: 1,
-            messages: vec![
-                crate::session::ConversationMessage::user_text("one"),
-                crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
-                    text: "two".to_string(),
-                }]),
-                crate::session::ConversationMessage::user_text("three"),
-                crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
-                    text: "four".to_string(),
-                }]),
-            ],
-        };
+        let mut session = Session::new();
+        session.messages = vec![
+            crate::session::ConversationMessage::user_text("one"),
+            crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
+                text: "two".to_string(),
+            }]),
+            crate::session::ConversationMessage::user_text("three"),
+            crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
+                text: "four".to_string(),
+            }]),
+        ];
 
         let mut runtime = ConversationRuntime::new(
             session,

+ 8 - 1
rust/crates/runtime/src/session.rs

@@ -161,8 +161,15 @@ impl Session {
         let path = path.as_ref();
         let contents = fs::read_to_string(path)?;
         let session = match JsonValue::parse(&contents) {
-            Ok(value) => Self::from_json(&value)?,
+            Ok(value)
+                if value
+                    .as_object()
+                    .is_some_and(|object| object.contains_key("messages")) =>
+            {
+                Self::from_json(&value)?
+            }
             Err(_) => Self::from_jsonl(&contents)?,
+            Ok(_) => Self::from_jsonl(&contents)?,
         };
         Ok(session.with_persistence_path(path.to_path_buf()))
     }

+ 28 - 38
rust/crates/rusty-claude-cli/src/main.rs

@@ -17,7 +17,7 @@ use std::time::{Duration, Instant, UNIX_EPOCH};
 
 use api::{
     resolve_startup_auth_source, AnthropicClient, AuthSource, ContentBlockDelta, InputContentBlock,
-    InputMessage, JsonlTelemetrySink, MessageRequest, MessageResponse, OutputContentBlock,
+    InputMessage, MessageRequest, MessageResponse, OutputContentBlock,
     SessionTracer, StreamEvent as ApiStreamEvent, ToolChoice, ToolDefinition,
     ToolResultContentBlock,
 };
@@ -3277,8 +3277,7 @@ impl AnthropicRuntimeClient {
         Ok(Self {
             runtime: tokio::runtime::Runtime::new()?,
             client: AnthropicClient::from_auth(resolve_cli_auth_source()?)
-                .with_base_url(api::read_base_url())
-                .with_prompt_cache(PromptCache::new(session_id)),
+                .with_base_url(api::read_base_url()),
             model,
             enable_tools,
             emit_output,
@@ -4037,25 +4036,7 @@ fn response_to_events(
     Ok(events)
 }
 
-fn push_prompt_cache_record(client: &AnthropicClient, events: &mut Vec<AssistantEvent>) {
-    if let Some(event) = client
-        .take_last_prompt_cache_record()
-        .and_then(prompt_cache_record_to_runtime_event)
-    {
-        events.push(AssistantEvent::PromptCache(event));
-    }
-}
-
-fn prompt_cache_record_to_runtime_event(record: PromptCacheRecord) -> Option<PromptCacheEvent> {
-    let cache_break = record.cache_break?;
-    Some(PromptCacheEvent {
-        unexpected: cache_break.unexpected,
-        reason: cache_break.reason,
-        previous_cache_read_input_tokens: cache_break.previous_cache_read_input_tokens,
-        current_cache_read_input_tokens: cache_break.current_cache_read_input_tokens,
-        token_drop: cache_break.token_drop,
-    })
-}
+fn push_prompt_cache_record(_client: &AnthropicClient, _events: &mut Vec<AssistantEvent>) {}
 
 struct CliToolExecutor {
     renderer: TerminalRenderer,
@@ -4273,19 +4254,18 @@ mod tests {
         format_permissions_report,
         format_permissions_switch_report, format_resume_report, format_status_report,
         format_tool_call_start, format_tool_result, normalize_permission_mode, parse_args,
-        parse_git_status_branch, parse_git_status_metadata, permission_policy, print_help_to,
-        push_output_block, render_config_report, render_diff_report, render_memory_report,
-        render_repl_help, resolve_model_alias, response_to_events,
+        parse_git_status_branch, parse_git_status_metadata_for, permission_policy,
+        print_help_to, push_output_block, render_config_report, render_diff_report,
+        render_memory_report, render_repl_help, resolve_model_alias, response_to_events,
         resume_supported_slash_commands, run_resume_command, status_context, CliAction,
-        CliOutputFormat, HookAbortMonitor, InternalPromptProgressEvent,
+        CliOutputFormat, InternalPromptProgressEvent,
         InternalPromptProgressState, SlashCommand, StatusUsage, DEFAULT_MODEL,
         create_managed_session_handle, resolve_session_reference,
     };
     use api::{MessageResponse, OutputContentBlock, Usage};
     use plugins::{PluginTool, PluginToolDefinition, PluginToolPermission};
     use runtime::{
-        AssistantEvent, ContentBlock, ConversationMessage, HookAbortSignal, MessageRole,
-        PermissionMode, Session,
+        AssistantEvent, ContentBlock, ConversationMessage, MessageRole, PermissionMode, Session,
     };
     use serde_json::json;
     use std::fs;
@@ -4354,8 +4334,6 @@ mod tests {
         std::env::set_current_dir(previous).expect("cwd should restore");
         result
     }
-    use std::sync::mpsc;
-
     #[test]
     fn defaults_to_repl_when_no_args() {
         assert_eq!(
@@ -4626,7 +4604,12 @@ mod tests {
 
     #[test]
     fn permission_policy_uses_plugin_tool_permissions() {
-        let policy = permission_policy(PermissionMode::ReadOnly, &registry_with_plugin_tool());
+        let feature_config = runtime::RuntimeFeatureConfig::default();
+        let policy = permission_policy(
+            PermissionMode::ReadOnly,
+            &feature_config,
+            &registry_with_plugin_tool(),
+        );
         let required = policy.required_mode_for("plugin_echo");
         assert_eq!(required, PermissionMode::WorkspaceWrite);
     }
@@ -4844,12 +4827,13 @@ mod tests {
         let _guard = env_lock();
         let temp_root = temp_dir();
         fs::create_dir_all(&temp_root).expect("root dir");
-        let (project_root, branch) = with_current_dir(&temp_root, || {
-            parse_git_status_metadata(Some(
+        let (project_root, branch) = parse_git_status_metadata_for(
+            &temp_root,
+            Some(
                 "## rcc/cli...origin/rcc/cli
  M src/main.rs",
-            ))
-        });
+            ),
+        );
         assert_eq!(branch.as_deref(), Some("rcc/cli"));
         assert!(project_root.is_none());
         fs::remove_dir_all(temp_root).expect("cleanup temp dir");
@@ -5057,7 +5041,7 @@ mod tests {
         let handle = create_managed_session_handle("session-alpha").expect("jsonl handle");
         assert!(handle.path.ends_with("session-alpha.jsonl"));
 
-        let legacy_path = workspace.join(".claude/sessions/legacy.json");
+        let legacy_path = workspace.join(".claw/sessions/legacy.json");
         std::fs::create_dir_all(
             legacy_path
                 .parent()
@@ -5070,7 +5054,10 @@ mod tests {
             .expect("legacy session should save");
 
         let resolved = resolve_session_reference("legacy").expect("legacy session should resolve");
-        assert_eq!(resolved.path, legacy_path);
+        assert_eq!(
+            resolved.path.canonicalize().expect("resolved path should exist"),
+            legacy_path.canonicalize().expect("legacy path should exist")
+        );
 
         std::env::set_current_dir(previous).expect("restore cwd");
         std::fs::remove_dir_all(workspace).expect("workspace should clean up");
@@ -5455,7 +5442,10 @@ mod tests {
 
 #[cfg(test)]
 mod sandbox_report_tests {
-    use super::format_sandbox_report;
+    use super::{format_sandbox_report, HookAbortMonitor};
+    use runtime::HookAbortSignal;
+    use std::sync::mpsc;
+    use std::time::Duration;
 
     #[test]
     fn sandbox_report_renders_expected_fields() {