From a9f5f633b2b67cd595c4efc0ee46aac9a1c4231f Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 18 Feb 2026 12:45:19 +0000 Subject: [PATCH] feat: memory usage metrics (#12120) --- codex-rs/core/src/memories/mod.rs | 1 + codex-rs/core/src/memories/usage.rs | 122 ++++++++++++++++++ codex-rs/core/src/tools/handlers/mod.rs | 2 +- .../core/src/tools/handlers/unified_exec.rs | 6 +- codex-rs/core/src/tools/registry.rs | 2 + 5 files changed, 129 insertions(+), 4 deletions(-) create mode 100644 codex-rs/core/src/memories/usage.rs diff --git a/codex-rs/core/src/memories/mod.rs b/codex-rs/core/src/memories/mod.rs index ddf9878d7..fbc67781b 100644 --- a/codex-rs/core/src/memories/mod.rs +++ b/codex-rs/core/src/memories/mod.rs @@ -11,6 +11,7 @@ mod start; mod storage; #[cfg(test)] mod tests; +pub(crate) mod usage; /// Starts the memory startup pipeline for eligible root sessions. /// This is the single entrypoint that `codex` uses to trigger memory startup. diff --git a/codex-rs/core/src/memories/usage.rs b/codex-rs/core/src/memories/usage.rs new file mode 100644 index 000000000..dbdcd2573 --- /dev/null +++ b/codex-rs/core/src/memories/usage.rs @@ -0,0 +1,122 @@ +use crate::is_safe_command::is_known_safe_command; +use crate::parse_command::parse_command; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; +use crate::tools::handlers::unified_exec::ExecCommandArgs; +use codex_protocol::models::ShellCommandToolCallParams; +use codex_protocol::models::ShellToolCallParams; +use codex_protocol::parse_command::ParsedCommand; +use std::path::PathBuf; + +const MEMORIES_USAGE_METRIC: &str = "codex.memories.usage"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +enum MemoriesUsageKind { + MemoryMd, + MemorySummary, + RawMemories, + RolloutSummaries, + Skills, +} + +impl MemoriesUsageKind { + fn as_tag(self) -> &'static str { + match self { + Self::MemoryMd => "memory_md", + Self::MemorySummary => "memory_summary", + Self::RawMemories => "raw_memories", + Self::RolloutSummaries => "rollout_summaries", + Self::Skills => "skills", + } + } +} + +pub(crate) async fn emit_metric_for_tool_read(invocation: &ToolInvocation, success: bool) { + let kinds = memories_usage_kinds_from_invocation(invocation).await; + if kinds.is_empty() { + return; + } + + let success = if success { "true" } else { "false" }; + for kind in kinds { + invocation.turn.otel_manager.counter( + MEMORIES_USAGE_METRIC, + 1, + &[ + ("kind", kind.as_tag()), + ("tool", invocation.tool_name.as_str()), + ("success", success), + ], + ); + } +} + +async fn memories_usage_kinds_from_invocation( + invocation: &ToolInvocation, +) -> Vec { + let Some((command, _)) = shell_command_for_invocation(invocation) else { + return Vec::new(); + }; + if !is_known_safe_command(&command) { + return Vec::new(); + } + + let parsed_commands = parse_command(&command); + parsed_commands + .into_iter() + .filter_map(|command| match command { + ParsedCommand::Read { path, .. } => get_memory_kind(path.display().to_string()), + ParsedCommand::Search { path, .. } => path.and_then(get_memory_kind), + ParsedCommand::ListFiles { .. } | ParsedCommand::Unknown { .. } => None, + }) + .collect() +} + +fn shell_command_for_invocation(invocation: &ToolInvocation) -> Option<(Vec, PathBuf)> { + let ToolPayload::Function { arguments } = &invocation.payload else { + return None; + }; + + match invocation.tool_name.as_str() { + "shell" => serde_json::from_str::(arguments) + .ok() + .map(|params| (params.command, invocation.turn.resolve_path(params.workdir))), + "shell_command" => serde_json::from_str::(arguments) + .ok() + .map(|params| { + let command = invocation + .session + .user_shell() + .derive_exec_args(¶ms.command, params.login.unwrap_or(true)); + (command, invocation.turn.resolve_path(params.workdir)) + }), + "exec_command" => serde_json::from_str::(arguments) + .ok() + .map(|params| { + ( + crate::tools::handlers::unified_exec::get_command( + ¶ms, + invocation.session.user_shell(), + ), + invocation.turn.resolve_path(params.workdir), + ) + }), + _ => None, + } +} + +fn get_memory_kind(path: String) -> Option { + if path.contains("memories/MEMORY.md") { + Some(MemoriesUsageKind::MemoryMd) + } else if path.contains("memories/memory_summary.md") { + Some(MemoriesUsageKind::MemorySummary) + } else if path.contains("memories/raw_memories.md") { + Some(MemoriesUsageKind::RawMemories) + } else if path.contains("memories/rollout_summaries/") { + Some(MemoriesUsageKind::RolloutSummaries) + } else if path.contains("memories/skills/") { + Some(MemoriesUsageKind::Skills) + } else { + None + } +} diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index b0e1b3c0a..e5229aa1d 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -12,7 +12,7 @@ mod request_user_input; mod search_tool_bm25; mod shell; mod test_sync; -mod unified_exec; +pub(crate) mod unified_exec; mod view_image; pub use plan::PLAN_TOOL; diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index e689e5bba..1bad6fa63 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -26,10 +26,10 @@ use std::sync::Arc; pub struct UnifiedExecHandler; #[derive(Debug, Deserialize)] -struct ExecCommandArgs { +pub(crate) struct ExecCommandArgs { cmd: String, #[serde(default)] - workdir: Option, + pub(crate) workdir: Option, #[serde(default)] shell: Option, #[serde(default = "default_login")] @@ -238,7 +238,7 @@ impl ToolHandler for UnifiedExecHandler { } } -fn get_command(args: &ExecCommandArgs, session_shell: Arc) -> Vec { +pub(crate) fn get_command(args: &ExecCommandArgs, session_shell: Arc) -> Vec { let model_shell = args.shell.as_ref().map(|shell_str| { let mut shell = get_shell_by_model_provided_path(&PathBuf::from(shell_str)); shell.shell_snapshot = crate::shell::empty_shell_snapshot_receiver(); diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 8006c8072..4ffb7ca10 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -6,6 +6,7 @@ use std::time::Instant; use crate::client_common::tools::ToolSpec; use crate::features::Feature; use crate::function_tool::FunctionCallError; +use crate::memories::usage::emit_metric_for_tool_read; use crate::protocol::SandboxPolicy; use crate::sandbox_tags::sandbox_tag; use crate::tools::context::ToolInvocation; @@ -173,6 +174,7 @@ impl ToolRegistry { Ok((preview, success)) => (preview.clone(), *success), Err(err) => (err.to_string(), false), }; + emit_metric_for_tool_read(&invocation, success).await; let hook_abort_error = dispatch_after_tool_use_hook(AfterToolUseHookDispatch { invocation: &invocation, output_preview,