diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index de11d1c15..164e8a085 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1859,6 +1859,7 @@ dependencies = [ "codex-git-utils", "codex-login", "codex-mcp", + "codex-memories-write", "codex-model-provider", "codex-model-provider-info", "codex-models-manager", @@ -2118,6 +2119,7 @@ dependencies = [ "codex-login", "codex-mcp", "codex-mcp-server", + "codex-memories-write", "codex-models-manager", "codex-protocol", "codex-responses-api-proxy", @@ -2373,7 +2375,6 @@ dependencies = [ "codex-login", "codex-mcp", "codex-memories-read", - "codex-memories-write", "codex-model-provider", "codex-model-provider-info", "codex-models-manager", @@ -2386,7 +2387,6 @@ dependencies = [ "codex-rollout", "codex-rollout-trace", "codex-sandboxing", - "codex-secrets", "codex-shell-command", "codex-shell-escalation", "codex-state", @@ -2957,18 +2957,32 @@ version = "0.0.0" dependencies = [ "anyhow", "chrono", + "codex-config", + "codex-core", + "codex-features", "codex-git-utils", + "codex-login", "codex-models-manager", + "codex-otel", "codex-protocol", + "codex-rollout", + "codex-rollout-trace", + "codex-secrets", "codex-state", + "codex-terminal-detection", "codex-utils-absolute-path", "codex-utils-output-truncation", "codex-utils-template", + "core_test_support", + "futures", "pretty_assertions", + "serde", + "serde_json", "tempfile", "tokio", "tracing", "uuid", + "wiremock", ] [[package]] diff --git a/codex-rs/agent-identity/src/lib.rs b/codex-rs/agent-identity/src/lib.rs index bf139f787..69801f6fc 100644 --- a/codex-rs/agent-identity/src/lib.rs +++ b/codex-rs/agent-identity/src/lib.rs @@ -322,6 +322,7 @@ pub fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials { | SessionSource::Exec | SessionSource::Mcp | SessionSource::Custom(_) + | SessionSource::Internal(_) | SessionSource::SubAgent(_) | SessionSource::Unknown => "codex-cli".to_string(), }, diff --git a/codex-rs/analytics/src/reducer.rs b/codex-rs/analytics/src/reducer.rs index 0445975dd..290dbe28d 100644 --- a/codex-rs/analytics/src/reducer.rs +++ b/codex-rs/analytics/src/reducer.rs @@ -106,6 +106,7 @@ impl ThreadMetadataState { | SessionSource::Exec | SessionSource::Mcp | SessionSource::Custom(_) + | SessionSource::Internal(_) | SessionSource::Unknown => (None, None), }; Self { diff --git a/codex-rs/app-server-protocol/schema/typescript/InternalSessionSource.ts b/codex-rs/app-server-protocol/schema/typescript/InternalSessionSource.ts new file mode 100644 index 000000000..47417c516 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/InternalSessionSource.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type InternalSessionSource = "memory_consolidation"; diff --git a/codex-rs/app-server-protocol/schema/typescript/SessionSource.ts b/codex-rs/app-server-protocol/schema/typescript/SessionSource.ts index a80b013b2..3317c228b 100644 --- a/codex-rs/app-server-protocol/schema/typescript/SessionSource.ts +++ b/codex-rs/app-server-protocol/schema/typescript/SessionSource.ts @@ -1,6 +1,7 @@ // GENERATED CODE! DO NOT MODIFY BY HAND! // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { InternalSessionSource } from "./InternalSessionSource"; import type { SubAgentSource } from "./SubAgentSource"; -export type SessionSource = "cli" | "vscode" | "exec" | "mcp" | { "custom": string } | { "subagent": SubAgentSource } | "unknown"; +export type SessionSource = "cli" | "vscode" | "exec" | "mcp" | { "custom": string } | { "internal": InternalSessionSource } | { "subagent": SubAgentSource } | "unknown"; diff --git a/codex-rs/app-server-protocol/schema/typescript/index.ts b/codex-rs/app-server-protocol/schema/typescript/index.ts index a5d3dbe57..a082e045f 100644 --- a/codex-rs/app-server-protocol/schema/typescript/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/index.ts @@ -37,6 +37,7 @@ export type { InitializeCapabilities } from "./InitializeCapabilities"; export type { InitializeParams } from "./InitializeParams"; export type { InitializeResponse } from "./InitializeResponse"; export type { InputModality } from "./InputModality"; +export type { InternalSessionSource } from "./InternalSessionSource"; export type { LocalShellAction } from "./LocalShellAction"; export type { LocalShellExecAction } from "./LocalShellExecAction"; export type { LocalShellStatus } from "./LocalShellStatus"; diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index d59363df4..64f16d162 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -1996,6 +1996,8 @@ impl From for SessionSource { CoreSessionSource::Exec => SessionSource::Exec, CoreSessionSource::Mcp => SessionSource::AppServer, CoreSessionSource::Custom(source) => SessionSource::Custom(source), + // We do not want to render those at the app-server level. + CoreSessionSource::Internal(_) => SessionSource::Unknown, CoreSessionSource::SubAgent(sub) => SessionSource::SubAgent(sub), CoreSessionSource::Unknown => SessionSource::Unknown, } diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index d78486f38..d1721c99d 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -48,6 +48,7 @@ codex-backend-client = { workspace = true } codex-file-search = { workspace = true } codex-chatgpt = { workspace = true } codex-login = { workspace = true } +codex-memories-write = { workspace = true } codex-mcp = { workspace = true } codex-model-provider = { workspace = true } codex-models-manager = { workspace = true } diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index eb67a2051..519ea50fd 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -241,11 +241,10 @@ use codex_core::ForkSnapshot; use codex_core::NewThread; use codex_core::RolloutRecorder; use codex_core::SessionMeta; -use codex_core::StartThreadWithToolsOptions; +use codex_core::StartThreadOptions; use codex_core::SteerInputError; use codex_core::ThreadConfigSnapshot; use codex_core::ThreadManager; -use codex_core::clear_memory_roots_contents; use codex_core::config::Config; use codex_core::config::ConfigOverrides; use codex_core::config::NetworkProxyAuditMetadata; @@ -315,6 +314,7 @@ use codex_mcp::discover_supported_scopes; use codex_mcp::effective_mcp_servers; use codex_mcp::read_mcp_resource as read_mcp_resource_without_thread; use codex_mcp::resolve_oauth_scopes; +use codex_memories_write::clear_memory_roots_contents; use codex_model_provider::ProviderAccountError; use codex_model_provider::create_model_provider; use codex_models_manager::collaboration_mode_presets::CollaborationModesConfig; @@ -542,6 +542,7 @@ pub(crate) enum ApiVersion { #[derive(Clone)] struct ListenerTaskContext { + auth_manager: Arc, thread_manager: Arc, thread_state_manager: ThreadStateManager, outgoing: Arc, @@ -2414,6 +2415,7 @@ impl CodexMessageProcessor { ); typesafe_overrides.ephemeral = ephemeral; let listener_task_context = ListenerTaskContext { + auth_manager: Arc::clone(&self.auth_manager), thread_manager: Arc::clone(&self.thread_manager), thread_state_manager: self.thread_state_manager.clone(), outgoing: Arc::clone(&self.outgoing), @@ -2613,6 +2615,7 @@ impl CodexMessageProcessor { .collect() }; let core_dynamic_tool_count = core_dynamic_tools.len(); + let memory_config = Arc::new(config.clone()); let NewThread { thread_id, @@ -2621,7 +2624,7 @@ impl CodexMessageProcessor { .. } = listener_task_context .thread_manager - .start_thread_with_tools_and_service_name(StartThreadWithToolsOptions { + .start_thread_with_options(StartThreadOptions { config, initial_history: match session_start_source .unwrap_or(codex_app_server_protocol::ThreadStartSource::Startup) @@ -2633,6 +2636,7 @@ impl CodexMessageProcessor { InitialHistory::Cleared } }, + session_source: None, dynamic_tools: core_dynamic_tools, persist_extended_history, metrics_service_name: service_name, @@ -2665,6 +2669,14 @@ impl CodexMessageProcessor { otel.name = "app_server.thread_start.config_snapshot", )) .await; + codex_memories_write::start_memories_startup_task( + Arc::clone(&listener_task_context.thread_manager), + Arc::clone(&listener_task_context.auth_manager), + thread_id, + Arc::clone(&thread), + Arc::clone(&memory_config), + &config_snapshot.session_source, + ); let mut thread = build_thread_from_snapshot( thread_id, &config_snapshot, @@ -3659,13 +3671,13 @@ impl CodexMessageProcessor { params: ThreadLoadedListParams, ) -> Result { let ThreadLoadedListParams { cursor, limit } = params; - let mut data = self + let mut data: Vec = self .thread_manager .list_thread_ids() .await .into_iter() .map(|thread_id| thread_id.to_string()) - .collect::>(); + .collect(); if data.is_empty() { return Ok(ThreadLoadedListResponse { @@ -4166,6 +4178,7 @@ impl CodexMessageProcessor { let instruction_sources = Self::instruction_sources_from_config(&config).await; let response_history = thread_history.clone(); + let memory_config = Arc::new(config.clone()); match self .thread_manager @@ -4240,6 +4253,14 @@ impl CodexMessageProcessor { /*has_live_in_progress_turn*/ false, ); let config_snapshot = codex_thread.config_snapshot().await; + codex_memories_write::start_memories_startup_task( + Arc::clone(&self.thread_manager), + Arc::clone(&self.auth_manager), + thread_id, + Arc::clone(&codex_thread), + Arc::clone(&memory_config), + &config_snapshot.session_source, + ); let sandbox = thread_response_sandbox_policy( &config_snapshot.permission_profile, config_snapshot.cwd.as_path(), @@ -4745,6 +4766,7 @@ impl CodexMessageProcessor { let fallback_model_provider = config.model_provider_id.clone(); let instruction_sources = Self::instruction_sources_from_config(&config).await; let fork_thread_store = configured_thread_store(&config); + let memory_config = Arc::new(config.clone()); let NewThread { thread_id, @@ -4840,6 +4862,14 @@ impl CodexMessageProcessor { /*has_in_progress_turn*/ false, ); let config_snapshot = forked_thread.config_snapshot().await; + codex_memories_write::start_memories_startup_task( + Arc::clone(&self.thread_manager), + Arc::clone(&self.auth_manager), + thread_id, + Arc::clone(&forked_thread), + Arc::clone(&memory_config), + &config_snapshot.session_source, + ); let sandbox = thread_response_sandbox_policy( &config_snapshot.permission_profile, config_snapshot.cwd.as_path(), @@ -4878,7 +4908,6 @@ impl CodexMessageProcessor { return; } }; - self.analytics_events_client.track_response( request_id.connection_id.0, ClientResponse::ThreadFork { @@ -7159,6 +7188,7 @@ impl CodexMessageProcessor { ) -> Result { Self::ensure_conversation_listener_task( ListenerTaskContext { + auth_manager: Arc::clone(&self.auth_manager), thread_manager: Arc::clone(&self.thread_manager), thread_state_manager: self.thread_state_manager.clone(), outgoing: Arc::clone(&self.outgoing), @@ -7276,6 +7306,7 @@ impl CodexMessageProcessor { ) -> Result<(), JSONRPCErrorError> { Self::ensure_listener_task_running_task( ListenerTaskContext { + auth_manager: Arc::clone(&self.auth_manager), thread_manager: Arc::clone(&self.thread_manager), thread_state_manager: self.thread_state_manager.clone(), outgoing: Arc::clone(&self.outgoing), @@ -7324,6 +7355,7 @@ impl CodexMessageProcessor { thread_state.set_listener(cancel_tx, &conversation) }; let ListenerTaskContext { + auth_manager: _, outgoing, thread_manager, thread_state_manager, diff --git a/codex-rs/cli/Cargo.toml b/codex-rs/cli/Cargo.toml index 5f4f3aee1..cdee241b4 100644 --- a/codex-rs/cli/Cargo.toml +++ b/codex-rs/cli/Cargo.toml @@ -35,6 +35,7 @@ codex-exec-server = { workspace = true } codex-execpolicy = { workspace = true } codex-features = { workspace = true } codex-login = { workspace = true } +codex-memories-write = { workspace = true } codex-mcp = { workspace = true } codex-mcp-server = { workspace = true } codex-models-manager = { workspace = true } diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index c134bc8a2..1c2eabf98 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -52,7 +52,6 @@ use crate::marketplace_cmd::MarketplaceCli; use crate::mcp_cmd::McpCli; use codex_core::build_models_manager; -use codex_core::clear_memory_roots_contents; use codex_core::config::Config; use codex_core::config::ConfigOverrides; use codex_core::config::edit::ConfigEditsBuilder; @@ -61,6 +60,7 @@ use codex_features::FEATURES; use codex_features::Stage; use codex_features::is_known_feature_key; use codex_login::AuthManager; +use codex_memories_write::clear_memory_roots_contents; use codex_models_manager::bundled_models_response; use codex_models_manager::collaboration_mode_presets::CollaborationModesConfig; use codex_models_manager::manager::RefreshStrategy; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index c95b57b71..44c6aacac 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -40,7 +40,6 @@ codex-features = { workspace = true } codex-feedback = { workspace = true } codex-login = { workspace = true } codex-memories-read = { workspace = true } -codex-memories-write = { workspace = true } codex-mcp = { workspace = true } codex-model-provider-info = { workspace = true } codex-models-manager = { workspace = true } @@ -71,7 +70,6 @@ codex-utils-path = { workspace = true } codex-utils-plugins = { workspace = true } codex-utils-pty = { workspace = true } codex-utils-readiness = { workspace = true } -codex-secrets = { workspace = true } codex-utils-string = { workspace = true } codex-utils-stream-parser = { workspace = true } codex-utils-template = { workspace = true } diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 187fc29e9..25fc92206 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -28,7 +28,6 @@ use codex_protocol::protocol::Op; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; -use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; use codex_rollout::state_db; @@ -149,16 +148,8 @@ impl AgentControl { } } - /// Create a control-plane handle over the same thread manager with an independent live-agent - /// registry. - pub(crate) fn detached_registry(&self) -> Self { - Self { - manager: self.manager.clone(), - ..Default::default() - } - } - /// Spawn a new agent thread and submit the initial prompt. + #[cfg(test)] pub(crate) async fn spawn_agent( &self, config: crate::config::Config, @@ -830,16 +821,6 @@ impl AgentControl { Ok(thread.subscribe_status()) } - pub(crate) async fn get_total_token_usage(&self, agent_id: ThreadId) -> Option { - let Ok(state) = self.upgrade() else { - return None; - }; - let Ok(thread) = state.get_thread(agent_id).await else { - return None; - }; - thread.total_token_usage().await - } - pub(crate) async fn format_environment_context_subagents( &self, parent_thread_id: ThreadId, diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 3997018ea..fe64099cc 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -77,6 +77,7 @@ use codex_protocol::config_types::Verbosity as VerbosityConfig; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; +use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::W3cTraceContext; @@ -566,7 +567,7 @@ impl ModelClient { } if matches!( self.state.session_source, - SessionSource::SubAgent(SubAgentSource::MemoryConsolidation) + SessionSource::Internal(InternalSessionSource::MemoryConsolidation) ) { extra_headers.insert( X_OPENAI_MEMGEN_REQUEST_HEADER, @@ -1596,15 +1597,23 @@ fn build_responses_headers( } fn subagent_header_value(session_source: &SessionSource) -> Option { - let SessionSource::SubAgent(subagent_source) = session_source else { - return None; - }; - match subagent_source { - SubAgentSource::Review => Some("review".to_string()), - SubAgentSource::Compact => Some("compact".to_string()), - SubAgentSource::MemoryConsolidation => Some("memory_consolidation".to_string()), - SubAgentSource::ThreadSpawn { .. } => Some("collab_spawn".to_string()), - SubAgentSource::Other(label) => Some(label.clone()), + match session_source { + SessionSource::SubAgent(subagent_source) => match subagent_source { + SubAgentSource::Review => Some("review".to_string()), + SubAgentSource::Compact => Some("compact".to_string()), + SubAgentSource::MemoryConsolidation => Some("memory_consolidation".to_string()), + SubAgentSource::ThreadSpawn { .. } => Some("collab_spawn".to_string()), + SubAgentSource::Other(label) => Some(label.clone()), + }, + SessionSource::Internal(InternalSessionSource::MemoryConsolidation) => { + Some("memory_consolidation".to_string()) + } + SessionSource::Cli + | SessionSource::VSCode + | SessionSource::Exec + | SessionSource::Mcp + | SessionSource::Custom(_) + | SessionSource::Unknown => None, } } @@ -1618,6 +1627,7 @@ fn parent_thread_id_header_value(session_source: &SessionSource) -> Option None, } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index e2934dc1c..f7d263c0e 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -18,6 +18,7 @@ use codex_protocol::ThreadId; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; +use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_rollout_trace::ExecutionStatus; @@ -197,6 +198,18 @@ fn build_subagent_headers_sets_other_subagent_label() { assert_eq!(value, Some("memory_consolidation")); } +#[test] +fn build_subagent_headers_sets_internal_memory_consolidation_label() { + let client = test_model_client(SessionSource::Internal( + InternalSessionSource::MemoryConsolidation, + )); + let headers = client.build_subagent_headers(); + let value = headers + .get(X_OPENAI_SUBAGENT_HEADER) + .and_then(|value| value.to_str().ok()); + assert_eq!(value, Some("memory_consolidation")); +} + #[test] fn build_ws_client_metadata_includes_window_lineage_and_turn_metadata() { let parent_thread_id = ThreadId::new(); diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index fa4bd8067..eb5afe57f 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -27,7 +27,6 @@ use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::Submission; use codex_protocol::protocol::ThreadMemoryMode; -use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::W3cTraceContext; use codex_protocol::user_input::UserInput; @@ -86,6 +85,7 @@ pub struct CodexThreadTurnContextOverrides { pub struct CodexThread { pub(crate) codex: Codex, + pub(crate) session_source: SessionSource, rollout_path: Option, out_of_band_elicitation_count: Mutex, _watch_registration: WatchRegistration, @@ -97,10 +97,12 @@ impl CodexThread { pub(crate) fn new( codex: Codex, rollout_path: Option, + session_source: SessionSource, watch_registration: WatchRegistration, ) -> Self { Self { codex, + session_source, rollout_path, out_of_band_elicitation_count: Mutex::new(0), _watch_registration: watch_registration, @@ -115,6 +117,11 @@ impl CodexThread { self.codex.shutdown_and_wait().await } + /// Wait until the underlying session loop has terminated. + pub async fn wait_until_terminated(&self) { + self.codex.session_loop_termination.clone().await; + } + pub async fn apply_goal_resume_runtime_effects(&self) -> anyhow::Result<()> { self.codex .session @@ -268,10 +275,6 @@ impl CodexThread { self.codex.agent_status.clone() } - pub(crate) async fn total_token_usage(&self) -> Option { - self.codex.session.total_token_usage().await - } - /// Returns the complete token usage snapshot currently cached for this thread. /// /// This accessor is intentionally narrower than direct session access: it lets diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 1150aed7e..f8a079b1f 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -62,7 +62,7 @@ use codex_features::MultiAgentV2ConfigToml; use codex_git_utils::resolve_root_git_project_for_trust; use codex_login::AuthManagerConfig; use codex_mcp::McpConfig; -use codex_memories_write::memory_root; +use codex_memories_read::memory_root; use codex_model_provider_info::LEGACY_OLLAMA_CHAT_PROVIDER_ID; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::OLLAMA_CHAT_PROVIDER_REMOVED_ERROR; diff --git a/codex-rs/core/src/context/contextual_user_message.rs b/codex-rs/core/src/context/contextual_user_message.rs index aeb54d61d..cd7788afe 100644 --- a/codex-rs/core/src/context/contextual_user_message.rs +++ b/codex-rs/core/src/context/contextual_user_message.rs @@ -33,34 +33,12 @@ static CONTEXTUAL_USER_FRAGMENTS: &[&dyn FragmentRegistration] = &[ &SUBAGENT_NOTIFICATION_REGISTRATION, ]; -static MEMORY_EXCLUDED_CONTEXTUAL_USER_FRAGMENTS: &[&dyn FragmentRegistration] = &[ - &USER_INSTRUCTIONS_REGISTRATION, - &SKILL_INSTRUCTIONS_REGISTRATION, -]; - fn is_standard_contextual_user_text(text: &str) -> bool { CONTEXTUAL_USER_FRAGMENTS .iter() .any(|fragment| fragment.matches_text(text)) } -/// Returns whether a contextual user fragment should be omitted from memory -/// stage-1 inputs. -/// -/// We exclude injected `AGENTS.md` instructions and skill payloads because -/// they are prompt scaffolding rather than conversation content, so they do -/// not improve the resulting memory. We keep environment context and -/// subagent notifications because they can carry useful execution context or -/// subtask outcomes that should remain visible to memory generation. -pub(crate) fn is_memory_excluded_contextual_user_fragment(content_item: &ContentItem) -> bool { - let ContentItem::InputText { text } = content_item else { - return false; - }; - MEMORY_EXCLUDED_CONTEXTUAL_USER_FRAGMENTS - .iter() - .any(|fragment| fragment.matches_text(text)) -} - pub(crate) fn is_contextual_user_fragment(content_item: &ContentItem) -> bool { let ContentItem::InputText { text } = content_item else { return false; diff --git a/codex-rs/core/src/context/contextual_user_message_tests.rs b/codex-rs/core/src/context/contextual_user_message_tests.rs index d52cf27ad..a90b8f280 100644 --- a/codex-rs/core/src/context/contextual_user_message_tests.rs +++ b/codex-rs/core/src/context/contextual_user_message_tests.rs @@ -33,38 +33,6 @@ fn ignores_regular_user_text() { })); } -#[test] -fn classifies_memory_excluded_fragments() { - let cases = [ - ( - "# AGENTS.md instructions for /tmp\n\n\nbody\n", - true, - ), - ( - "\ndemo\nskills/demo/SKILL.md\nbody\n", - true, - ), - ( - "\n/tmp\n", - false, - ), - ( - "{\"agent_id\":\"a\",\"status\":\"completed\"}", - false, - ), - ]; - - for (text, expected) in cases { - assert_eq!( - is_memory_excluded_contextual_user_fragment(&ContentItem::InputText { - text: text.to_string(), - }), - expected, - "{text}", - ); - } -} - #[test] fn detects_hook_prompt_fragment_and_roundtrips_escaping() { let message = build_hook_prompt_message(&[HookPromptFragment::from_single_hook( diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 848a4f7f0..25a6e1f13 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -31,7 +31,6 @@ pub(crate) use available_plugins_instructions::AvailablePluginsInstructions; pub(crate) use available_skills_instructions::AvailableSkillsInstructions; pub(crate) use collaboration_mode_instructions::CollaborationModeInstructions; pub(crate) use contextual_user_message::is_contextual_user_fragment; -pub(crate) use contextual_user_message::is_memory_excluded_contextual_user_fragment; pub(crate) use contextual_user_message::parse_visible_hook_prompt_message; pub(crate) use environment_context::EnvironmentContext; pub use fragment::ContextualUserFragment; diff --git a/codex-rs/core/src/installation_id.rs b/codex-rs/core/src/installation_id.rs index e9e5445c8..a42e6b6d8 100644 --- a/codex-rs/core/src/installation_id.rs +++ b/codex-rs/core/src/installation_id.rs @@ -16,7 +16,7 @@ use uuid::Uuid; pub(crate) const INSTALLATION_ID_FILENAME: &str = "installation_id"; -pub(crate) async fn resolve_installation_id(codex_home: &AbsolutePathBuf) -> Result { +pub async fn resolve_installation_id(codex_home: &AbsolutePathBuf) -> Result { let path = codex_home.join(INSTALLATION_ID_FILENAME); fs::create_dir_all(codex_home).await?; tokio::task::spawn_blocking(move || { diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index c5e0da9b8..05e34efc8 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -56,8 +56,6 @@ mod original_image_detail; pub use codex_mcp::SandboxState; mod mcp_openai_file; mod mcp_tool_call; -mod memories; -pub use codex_memories_write::clear_memory_roots_contents; pub(crate) mod mention_syntax; pub(crate) mod message_history; pub(crate) mod utils; @@ -119,7 +117,7 @@ pub(crate) mod web_search; pub(crate) mod windows_sandbox_read_grants; pub use thread_manager::ForkSnapshot; pub use thread_manager::NewThread; -pub use thread_manager::StartThreadWithToolsOptions; +pub use thread_manager::StartThreadOptions; pub use thread_manager::ThreadManager; pub use thread_manager::build_models_manager; pub use web_search::web_search_action_detail; @@ -195,6 +193,7 @@ pub use exec_policy::check_execpolicy_for_warnings; pub use exec_policy::format_exec_policy_error_with_source; pub use exec_policy::load_exec_policy; pub use file_watcher::FileWatcherEvent; +pub use installation_id::resolve_installation_id; pub use turn_metadata::build_turn_metadata_header; pub mod compact; pub(crate) mod memory_trace; diff --git a/codex-rs/core/src/memories/mod.rs b/codex-rs/core/src/memories/mod.rs deleted file mode 100644 index efa70f4bb..000000000 --- a/codex-rs/core/src/memories/mod.rs +++ /dev/null @@ -1,73 +0,0 @@ -//! Memory startup extraction and consolidation orchestration. -//! -//! The startup memory pipeline is split into two phases: -//! - Phase 1: select rollouts, extract stage-1 raw memories, persist stage-1 outputs, and enqueue consolidation. -//! - Phase 2: claim a global consolidation lock, materialize consolidation inputs, and dispatch one consolidation agent. - -mod phase1; -mod phase2; -mod start; -#[cfg(test)] -mod tests; - -use codex_protocol::openai_models::ReasoningEffort; - -/// Starts the memory startup pipeline for eligible root sessions. -/// This is the single entrypoint that `codex` uses to trigger memory startup. -/// -/// This is the entry point to read and understand this module. -pub(crate) use start::start_memories_startup_task; - -/// Phase 1 (startup extraction). -mod phase_one { - /// Default model used for phase 1. - pub(super) const MODEL: &str = "gpt-5.4-mini"; - /// Default reasoning effort used for phase 1. - pub(super) const REASONING_EFFORT: super::ReasoningEffort = super::ReasoningEffort::Low; - /// Prompt used for phase 1. - pub(super) const PROMPT: &str = codex_memories_write::STAGE_ONE_PROMPT; - /// Concurrency cap for startup memory extraction and consolidation scheduling. - pub(super) const CONCURRENCY_LIMIT: usize = 8; - /// Lease duration (seconds) for phase-1 job ownership. - pub(super) const JOB_LEASE_SECONDS: i64 = 3_600; - /// Backoff delay (seconds) before retrying a failed stage-1 extraction job. - pub(super) const JOB_RETRY_DELAY_SECONDS: i64 = 3_600; - /// Maximum number of threads to scan. - pub(super) const THREAD_SCAN_LIMIT: usize = 5_000; - /// Size of the batches when pruning old thread memories. - pub(super) const PRUNE_BATCH_SIZE: usize = 200; -} - -/// Phase 2 (aka `Consolidation`). -mod phase_two { - /// Default model used for phase 2. - pub(super) const MODEL: &str = "gpt-5.4"; - /// Default reasoning effort used for phase 2. - pub(super) const REASONING_EFFORT: super::ReasoningEffort = super::ReasoningEffort::Medium; - /// Lease duration (seconds) for phase-2 consolidation job ownership. - pub(super) const JOB_LEASE_SECONDS: i64 = 3_600; - /// Backoff delay (seconds) before retrying a failed phase-2 consolidation - /// job. - pub(super) const JOB_RETRY_DELAY_SECONDS: i64 = 3_600; - /// Heartbeat interval (seconds) for phase-2 running jobs. - pub(super) const JOB_HEARTBEAT_SECONDS: u64 = 90; -} - -mod metrics { - /// Number of phase-1 startup jobs grouped by status. - pub(super) const MEMORY_PHASE_ONE_JOBS: &str = "codex.memory.phase1"; - /// End-to-end latency for a single phase-1 startup run. - pub(super) const MEMORY_PHASE_ONE_E2E_MS: &str = "codex.memory.phase1.e2e_ms"; - /// Number of raw memories produced by phase-1 startup extraction. - pub(super) const MEMORY_PHASE_ONE_OUTPUT: &str = "codex.memory.phase1.output"; - /// Histogram for aggregate token usage across one phase-1 startup run. - pub(super) const MEMORY_PHASE_ONE_TOKEN_USAGE: &str = "codex.memory.phase1.token_usage"; - /// Number of phase-2 startup jobs grouped by status. - pub(super) const MEMORY_PHASE_TWO_JOBS: &str = "codex.memory.phase2"; - /// End-to-end latency for a single phase-2 consolidation run. - pub(super) const MEMORY_PHASE_TWO_E2E_MS: &str = "codex.memory.phase2.e2e_ms"; - /// Number of stage-1 memories included in each phase-2 consolidation step. - pub(super) const MEMORY_PHASE_TWO_INPUT: &str = "codex.memory.phase2.input"; - /// Histogram for aggregate token usage across one phase-2 consolidation run. - pub(super) const MEMORY_PHASE_TWO_TOKEN_USAGE: &str = "codex.memory.phase2.token_usage"; -} diff --git a/codex-rs/core/src/memories/phase1.rs b/codex-rs/core/src/memories/phase1.rs deleted file mode 100644 index 40f65c6db..000000000 --- a/codex-rs/core/src/memories/phase1.rs +++ /dev/null @@ -1,620 +0,0 @@ -use crate::Prompt; -use crate::RolloutRecorder; -use crate::config::Config; -use crate::context::is_memory_excluded_contextual_user_fragment; -use crate::memories::metrics; -use crate::memories::phase_one; -use crate::memories::phase_one::PRUNE_BATCH_SIZE; -use crate::rollout::INTERACTIVE_SESSION_SOURCES; -use crate::rollout::policy::should_persist_response_item_for_memories; -use crate::session::session::Session; -use crate::session::turn_context::TurnContext; -use codex_api::ResponseEvent; -use codex_config::types::MemoriesConfig; -use codex_memories_write::build_stage_one_input_message; -use codex_otel::SessionTelemetry; -use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; -use codex_protocol::config_types::ServiceTier; -use codex_protocol::error::CodexErr; -use codex_protocol::models::BaseInstructions; -use codex_protocol::models::ContentItem; -use codex_protocol::models::ResponseItem; -use codex_protocol::openai_models::ModelInfo; -use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; -use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::TokenUsage; -use codex_rollout_trace::InferenceTraceContext; -use codex_secrets::redact_secrets; -use futures::StreamExt; -use serde::Deserialize; -use serde_json::Value; -use serde_json::json; -use std::path::Path; -use std::sync::Arc; -use tracing::info; -use tracing::warn; - -#[derive(Clone, Debug)] -pub(in crate::memories) struct RequestContext { - pub(in crate::memories) model_info: ModelInfo, - pub(in crate::memories) session_telemetry: SessionTelemetry, - pub(in crate::memories) reasoning_effort: Option, - pub(in crate::memories) reasoning_summary: ReasoningSummaryConfig, - pub(in crate::memories) service_tier: Option, - pub(in crate::memories) turn_metadata_header: Option, -} - -struct JobResult { - outcome: JobOutcome, - token_usage: Option, -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum JobOutcome { - SucceededWithOutput, - SucceededNoOutput, - Failed, -} - -struct Stats { - claimed: usize, - succeeded_with_output: usize, - succeeded_no_output: usize, - failed: usize, - total_token_usage: Option, -} - -/// Phase 1 model output payload. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -struct StageOneOutput { - /// Detailed markdown raw memory for a single rollout. - #[serde(rename = "raw_memory")] - pub(crate) raw_memory: String, - /// Compact summary line used for routing and indexing. - #[serde(rename = "rollout_summary")] - pub(crate) rollout_summary: String, - /// Optional slug used to derive rollout summary artifact filenames. - #[serde(default, rename = "rollout_slug")] - pub(crate) rollout_slug: Option, -} - -/// Runs memory phase 1 in strict step order: -/// 1) claim eligible rollout jobs -/// 2) build one stage-1 request context -/// 3) run stage-1 extraction jobs in parallel -/// 4) emit metrics and logs -pub(in crate::memories) async fn run(session: &Arc, config: &Config) { - let _phase_one_e2e_timer = session - .services - .session_telemetry - .start_timer(metrics::MEMORY_PHASE_ONE_E2E_MS, &[]) - .ok(); - - // 1. Claim startup job. - let Some(claimed_candidates) = claim_startup_jobs(session, &config.memories).await else { - return; - }; - if claimed_candidates.is_empty() { - session.services.session_telemetry.counter( - metrics::MEMORY_PHASE_ONE_JOBS, - /*inc*/ 1, - &[("status", "skipped_no_candidates")], - ); - return; - } - - // 2. Build request. - let stage_one_context = build_request_context(session, config).await; - - // 3. Run the parallel sampling. - let outcomes = run_jobs(session, claimed_candidates, stage_one_context).await; - - // 4. Metrics and logs. - let counts = aggregate_stats(outcomes); - emit_metrics(session, &counts); - info!( - "memory stage-1 extraction complete: {} job(s) claimed, {} succeeded ({} with output, {} no output), {} failed", - counts.claimed, - counts.succeeded_with_output + counts.succeeded_no_output, - counts.succeeded_with_output, - counts.succeeded_no_output, - counts.failed - ); -} - -/// Prune old un-used "dead" raw memories. -pub(in crate::memories) async fn prune(session: &Arc, config: &Config) { - if let Some(db) = session.services.state_db.as_deref() { - let max_unused_days = config.memories.max_unused_days; - match db - .prune_stage1_outputs_for_retention(max_unused_days, PRUNE_BATCH_SIZE) - .await - { - Ok(pruned) => { - if pruned > 0 { - info!( - "memory startup pruned {pruned} stale stage-1 output row(s) older than {max_unused_days} days" - ); - } - } - Err(err) => { - warn!( - "state db prune_stage1_outputs_for_retention failed during memories startup: {err}" - ); - } - } - } -} - -/// JSON schema used to constrain phase-1 model output. -pub fn output_schema() -> Value { - json!({ - "type": "object", - "properties": { - "rollout_summary": { "type": "string" }, - "rollout_slug": { "type": ["string", "null"] }, - "raw_memory": { "type": "string" } - }, - "required": ["rollout_summary", "rollout_slug", "raw_memory"], - "additionalProperties": false - }) -} - -impl RequestContext { - pub(in crate::memories) fn from_turn_context( - turn_context: &TurnContext, - turn_metadata_header: Option, - model_info: ModelInfo, - ) -> Self { - Self { - model_info, - turn_metadata_header, - session_telemetry: turn_context.session_telemetry.clone(), - reasoning_effort: Some(phase_one::REASONING_EFFORT), - reasoning_summary: turn_context.reasoning_summary, - service_tier: turn_context.config.service_tier, - } - } -} - -async fn claim_startup_jobs( - session: &Arc, - memories_config: &MemoriesConfig, -) -> Option> { - let Some(state_db) = session.services.state_db.as_deref() else { - // This should not happen. - warn!("state db unavailable while claiming phase-1 startup jobs; skipping"); - return None; - }; - - let allowed_sources = INTERACTIVE_SESSION_SOURCES - .iter() - .map(ToString::to_string) - .collect::>(); - - match state_db - .claim_stage1_jobs_for_startup( - session.conversation_id, - codex_state::Stage1StartupClaimParams { - scan_limit: phase_one::THREAD_SCAN_LIMIT, - max_claimed: memories_config.max_rollouts_per_startup, - max_age_days: memories_config.max_rollout_age_days, - min_rollout_idle_hours: memories_config.min_rollout_idle_hours, - allowed_sources: allowed_sources.as_slice(), - lease_seconds: phase_one::JOB_LEASE_SECONDS, - }, - ) - .await - { - Ok(claims) => Some(claims), - Err(err) => { - warn!("state db claim_stage1_jobs_for_startup failed during memories startup: {err}"); - session.services.session_telemetry.counter( - metrics::MEMORY_PHASE_ONE_JOBS, - /*inc*/ 1, - &[("status", "failed_claim")], - ); - None - } - } -} - -async fn build_request_context(session: &Arc, config: &Config) -> RequestContext { - let model_name = config - .memories - .extract_model - .clone() - .unwrap_or(phase_one::MODEL.to_string()); - let model = session - .services - .models_manager - .get_model_info(&model_name, &config.to_models_manager_config()) - .await; - let turn_context = session.new_default_turn().await; - RequestContext::from_turn_context( - turn_context.as_ref(), - turn_context.turn_metadata_state.current_header_value(), - model, - ) -} - -async fn run_jobs( - session: &Arc, - claimed_candidates: Vec, - stage_one_context: RequestContext, -) -> Vec { - futures::stream::iter(claimed_candidates.into_iter()) - .map(|claim| { - let session = Arc::clone(session); - let stage_one_context = stage_one_context.clone(); - async move { job::run(session.as_ref(), claim, &stage_one_context).await } - }) - .buffer_unordered(phase_one::CONCURRENCY_LIMIT) - .collect::>() - .await -} - -mod job { - use super::*; - - pub(in crate::memories) async fn run( - session: &Session, - claim: codex_state::Stage1JobClaim, - stage_one_context: &RequestContext, - ) -> JobResult { - let thread = claim.thread; - let (stage_one_output, token_usage) = match sample( - session, - &thread.rollout_path, - &thread.cwd, - stage_one_context, - ) - .await - { - Ok(output) => output, - Err(reason) => { - result::failed( - session, - thread.id, - &claim.ownership_token, - &reason.to_string(), - ) - .await; - return JobResult { - outcome: JobOutcome::Failed, - token_usage: None, - }; - } - }; - - if stage_one_output.raw_memory.is_empty() || stage_one_output.rollout_summary.is_empty() { - return JobResult { - outcome: result::no_output(session, thread.id, &claim.ownership_token).await, - token_usage, - }; - } - - JobResult { - outcome: result::success( - session, - thread.id, - &claim.ownership_token, - thread.updated_at.timestamp(), - &stage_one_output.raw_memory, - &stage_one_output.rollout_summary, - stage_one_output.rollout_slug.as_deref(), - ) - .await, - token_usage, - } - } - - /// Extract the rollout and perform the actual sampling. - async fn sample( - session: &Session, - rollout_path: &Path, - rollout_cwd: &Path, - stage_one_context: &RequestContext, - ) -> anyhow::Result<(StageOneOutput, Option)> { - let (rollout_items, _, _) = RolloutRecorder::load_rollout_items(rollout_path).await?; - let rollout_contents = serialize_filtered_rollout_response_items(&rollout_items)?; - - let prompt = Prompt { - input: vec![ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { - text: build_stage_one_input_message( - &stage_one_context.model_info, - rollout_path, - rollout_cwd, - &rollout_contents, - )?, - }], - phase: None, - }], - tools: Vec::new(), - parallel_tool_calls: false, - base_instructions: BaseInstructions { - text: phase_one::PROMPT.to_string(), - }, - personality: None, - output_schema: Some(output_schema()), - output_schema_strict: true, - }; - - let mut client_session = session.services.model_client.new_session(); - let mut stream = client_session - .stream( - &prompt, - &stage_one_context.model_info, - &stage_one_context.session_telemetry, - stage_one_context.reasoning_effort, - stage_one_context.reasoning_summary, - stage_one_context.service_tier, - stage_one_context.turn_metadata_header.as_deref(), - &InferenceTraceContext::disabled(), - ) - .await?; - - // TODO(jif) we should have a shared helper somewhere for this. - // Unwrap the stream. - let mut result = String::new(); - let mut token_usage = None; - while let Some(message) = stream.next().await.transpose()? { - match message { - ResponseEvent::OutputTextDelta(delta) => result.push_str(&delta), - ResponseEvent::OutputItemDone(item) => { - if result.is_empty() - && let ResponseItem::Message { content, .. } = item - && let Some(text) = crate::compact::content_items_to_text(&content) - { - result.push_str(&text); - } - } - ResponseEvent::Completed { - token_usage: usage, .. - } => { - token_usage = usage; - break; - } - _ => {} - } - } - - let mut output: StageOneOutput = serde_json::from_str(&result)?; - output.raw_memory = redact_secrets(output.raw_memory); - output.rollout_summary = redact_secrets(output.rollout_summary); - output.rollout_slug = output.rollout_slug.map(redact_secrets); - - Ok((output, token_usage)) - } - - mod result { - use super::*; - - pub(in crate::memories) async fn failed( - session: &Session, - thread_id: codex_protocol::ThreadId, - ownership_token: &str, - reason: &str, - ) { - tracing::warn!("Phase 1 job failed for thread {thread_id}: {reason}"); - if let Some(state_db) = session.services.state_db.as_deref() { - let _ = state_db - .mark_stage1_job_failed( - thread_id, - ownership_token, - reason, - phase_one::JOB_RETRY_DELAY_SECONDS, - ) - .await; - } - } - - pub(in crate::memories) async fn no_output( - session: &Session, - thread_id: codex_protocol::ThreadId, - ownership_token: &str, - ) -> JobOutcome { - let Some(state_db) = session.services.state_db.as_deref() else { - return JobOutcome::Failed; - }; - - if state_db - .mark_stage1_job_succeeded_no_output(thread_id, ownership_token) - .await - .unwrap_or(false) - { - JobOutcome::SucceededNoOutput - } else { - JobOutcome::Failed - } - } - - pub(in crate::memories) async fn success( - session: &Session, - thread_id: codex_protocol::ThreadId, - ownership_token: &str, - source_updated_at: i64, - raw_memory: &str, - rollout_summary: &str, - rollout_slug: Option<&str>, - ) -> JobOutcome { - let Some(state_db) = session.services.state_db.as_deref() else { - return JobOutcome::Failed; - }; - - if state_db - .mark_stage1_job_succeeded( - thread_id, - ownership_token, - source_updated_at, - raw_memory, - rollout_summary, - rollout_slug, - ) - .await - .unwrap_or(false) - { - JobOutcome::SucceededWithOutput - } else { - JobOutcome::Failed - } - } - } - - /// Serializes filtered stage-1 memory items for prompt inclusion. - pub(super) fn serialize_filtered_rollout_response_items( - items: &[RolloutItem], - ) -> codex_protocol::error::Result { - let filtered = items - .iter() - .filter_map(|item| { - if let RolloutItem::ResponseItem(item) = item { - sanitize_response_item_for_memories(item) - } else { - None - } - }) - .collect::>(); - let serialized = serde_json::to_string(&filtered).map_err(|err| { - CodexErr::InvalidRequest(format!("failed to serialize rollout memory: {err}")) - })?; - Ok(redact_secrets(serialized)) - } - - fn sanitize_response_item_for_memories(item: &ResponseItem) -> Option { - let ResponseItem::Message { - id, - role, - content, - phase, - } = item - else { - return should_persist_response_item_for_memories(item).then(|| item.clone()); - }; - - if role == "developer" { - return None; - } - - if role != "user" { - return Some(item.clone()); - } - - let content = content - .iter() - .filter(|content_item| !is_memory_excluded_contextual_user_fragment(content_item)) - .cloned() - .collect::>(); - if content.is_empty() { - return None; - } - - Some(ResponseItem::Message { - id: id.clone(), - role: role.clone(), - content, - phase: phase.clone(), - }) - } -} - -fn aggregate_stats(outcomes: Vec) -> Stats { - let claimed = outcomes.len(); - let mut succeeded_with_output = 0; - let mut succeeded_no_output = 0; - let mut failed = 0; - let mut total_token_usage = TokenUsage::default(); - let mut has_token_usage = false; - - for outcome in outcomes { - match outcome.outcome { - JobOutcome::SucceededWithOutput => succeeded_with_output += 1, - JobOutcome::SucceededNoOutput => succeeded_no_output += 1, - JobOutcome::Failed => failed += 1, - } - - if let Some(token_usage) = outcome.token_usage { - total_token_usage.add_assign(&token_usage); - has_token_usage = true; - } - } - - Stats { - claimed, - succeeded_with_output, - succeeded_no_output, - failed, - total_token_usage: has_token_usage.then_some(total_token_usage), - } -} - -fn emit_metrics(session: &Session, counts: &Stats) { - if counts.claimed > 0 { - session.services.session_telemetry.counter( - metrics::MEMORY_PHASE_ONE_JOBS, - counts.claimed as i64, - &[("status", "claimed")], - ); - } - if counts.succeeded_with_output > 0 { - session.services.session_telemetry.counter( - metrics::MEMORY_PHASE_ONE_JOBS, - counts.succeeded_with_output as i64, - &[("status", "succeeded")], - ); - session.services.session_telemetry.counter( - metrics::MEMORY_PHASE_ONE_OUTPUT, - counts.succeeded_with_output as i64, - &[], - ); - } - if counts.succeeded_no_output > 0 { - session.services.session_telemetry.counter( - metrics::MEMORY_PHASE_ONE_JOBS, - counts.succeeded_no_output as i64, - &[("status", "succeeded_no_output")], - ); - } - if counts.failed > 0 { - session.services.session_telemetry.counter( - metrics::MEMORY_PHASE_ONE_JOBS, - counts.failed as i64, - &[("status", "failed")], - ); - } - if let Some(token_usage) = counts.total_token_usage.as_ref() { - session.services.session_telemetry.histogram( - metrics::MEMORY_PHASE_ONE_TOKEN_USAGE, - token_usage.total_tokens.max(0), - &[("token_type", "total")], - ); - session.services.session_telemetry.histogram( - metrics::MEMORY_PHASE_ONE_TOKEN_USAGE, - token_usage.input_tokens.max(0), - &[("token_type", "input")], - ); - session.services.session_telemetry.histogram( - metrics::MEMORY_PHASE_ONE_TOKEN_USAGE, - token_usage.cached_input(), - &[("token_type", "cached_input")], - ); - session.services.session_telemetry.histogram( - metrics::MEMORY_PHASE_ONE_TOKEN_USAGE, - token_usage.output_tokens.max(0), - &[("token_type", "output")], - ); - session.services.session_telemetry.histogram( - metrics::MEMORY_PHASE_ONE_TOKEN_USAGE, - token_usage.reasoning_output_tokens.max(0), - &[("token_type", "reasoning_output")], - ); - } -} - -#[cfg(test)] -#[path = "phase1_tests.rs"] -mod tests; diff --git a/codex-rs/core/src/memories/phase1_tests.rs b/codex-rs/core/src/memories/phase1_tests.rs deleted file mode 100644 index 18a46d174..000000000 --- a/codex-rs/core/src/memories/phase1_tests.rs +++ /dev/null @@ -1,152 +0,0 @@ -use super::JobOutcome; -use super::JobResult; -use super::aggregate_stats; -use super::job::serialize_filtered_rollout_response_items; -use codex_protocol::models::ContentItem; -use codex_protocol::models::FunctionCallOutputBody; -use codex_protocol::models::FunctionCallOutputPayload; -use codex_protocol::models::ResponseItem; -use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::TokenUsage; -use pretty_assertions::assert_eq; - -#[test] -fn serializes_memory_rollout_with_agents_removed_but_environment_kept() { - let mixed_contextual_message = ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ - ContentItem::InputText { - text: "# AGENTS.md instructions for /tmp\n\n\nbody\n" - .to_string(), - }, - ContentItem::InputText { - text: "\n/tmp\n".to_string(), - }, - ], - phase: None, - }; - let skill_message = ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { - text: "\ndemo\nskills/demo/SKILL.md\nbody\n" - .to_string(), - }], - phase: None, - }; - let subagent_message = ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { - text: "{\"agent_id\":\"a\",\"status\":\"completed\"}" - .to_string(), - }], - phase: None, - }; - - let serialized = serialize_filtered_rollout_response_items(&[ - RolloutItem::ResponseItem(mixed_contextual_message), - RolloutItem::ResponseItem(skill_message), - RolloutItem::ResponseItem(subagent_message.clone()), - ]) - .expect("serialize"); - let parsed: Vec = serde_json::from_str(&serialized).expect("parse"); - - assert_eq!( - parsed, - vec![ - ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { - text: "\n/tmp\n" - .to_string(), - }], - phase: None, - }, - subagent_message, - ] - ); -} - -#[test] -fn serializes_memory_rollout_redacts_secrets_before_prompt_upload() { - let serialized = serialize_filtered_rollout_response_items(&[RolloutItem::ResponseItem( - ResponseItem::FunctionCallOutput { - call_id: "call_123".to_string(), - output: FunctionCallOutputPayload { - body: FunctionCallOutputBody::Text( - r#"{"token":"sk-abcdefghijklmnopqrstuvwxyz123456"}"#.to_string(), - ), - success: Some(true), - }, - }, - )]) - .expect("serialize"); - - assert!(!serialized.contains("sk-abcdefghijklmnopqrstuvwxyz123456")); - assert!(serialized.contains("[REDACTED_SECRET]")); -} - -#[test] -fn count_outcomes_sums_token_usage_across_all_jobs() { - let counts = aggregate_stats(vec![ - JobResult { - outcome: JobOutcome::SucceededWithOutput, - token_usage: Some(TokenUsage { - input_tokens: 10, - cached_input_tokens: 2, - output_tokens: 3, - reasoning_output_tokens: 1, - total_tokens: 13, - }), - }, - JobResult { - outcome: JobOutcome::SucceededNoOutput, - token_usage: Some(TokenUsage { - input_tokens: 7, - cached_input_tokens: 1, - output_tokens: 2, - reasoning_output_tokens: 0, - total_tokens: 9, - }), - }, - JobResult { - outcome: JobOutcome::Failed, - token_usage: None, - }, - ]); - - assert_eq!(counts.claimed, 3); - assert_eq!(counts.succeeded_with_output, 1); - assert_eq!(counts.succeeded_no_output, 1); - assert_eq!(counts.failed, 1); - assert_eq!( - counts.total_token_usage, - Some(TokenUsage { - input_tokens: 17, - cached_input_tokens: 3, - output_tokens: 5, - reasoning_output_tokens: 1, - total_tokens: 22, - }) - ); -} - -#[test] -fn count_outcomes_keeps_usage_empty_when_no_job_reports_it() { - let counts = aggregate_stats(vec![ - JobResult { - outcome: JobOutcome::SucceededWithOutput, - token_usage: None, - }, - JobResult { - outcome: JobOutcome::Failed, - token_usage: None, - }, - ]); - - assert_eq!(counts.claimed, 2); - assert_eq!(counts.total_token_usage, None); -} diff --git a/codex-rs/core/src/memories/start.rs b/codex-rs/core/src/memories/start.rs deleted file mode 100644 index 2cfaf9a87..000000000 --- a/codex-rs/core/src/memories/start.rs +++ /dev/null @@ -1,44 +0,0 @@ -use crate::config::Config; -use crate::memories::phase1; -use crate::memories::phase2; -use crate::session::session::Session; -use codex_features::Feature; -use codex_protocol::protocol::SessionSource; -use std::sync::Arc; -use tracing::warn; - -/// Starts the asynchronous startup memory pipeline for an eligible root session. -/// -/// The pipeline is skipped for ephemeral sessions, disabled feature flags, and -/// subagent sessions. -pub(crate) fn start_memories_startup_task( - session: &Arc, - config: Arc, - source: &SessionSource, -) { - if config.ephemeral - || !config.features.enabled(Feature::MemoryTool) - || matches!(source, SessionSource::SubAgent(_)) - { - return; - } - - if session.services.state_db.is_none() { - warn!("state db unavailable for memories startup pipeline; skipping"); - return; - } - - let weak_session = Arc::downgrade(session); - tokio::spawn(async move { - let Some(session) = weak_session.upgrade() else { - return; - }; - - // Clean memories to make preserve DB size - phase1::prune(&session, &config).await; - // Run phase 1. - phase1::run(&session, &config).await; - // Run phase 2. - phase2::run(&session, config).await; - }); -} diff --git a/codex-rs/core/src/memories/tests.rs b/codex-rs/core/src/memories/tests.rs deleted file mode 100644 index cbf7222b1..000000000 --- a/codex-rs/core/src/memories/tests.rs +++ /dev/null @@ -1,1358 +0,0 @@ -use chrono::TimeZone; -use chrono::Utc; -use codex_config::types::DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION; -use codex_memories_write::clear_memory_roots_contents; -use codex_memories_write::ensure_layout; -use codex_memories_write::memory_root; -use codex_memories_write::raw_memories_file; -use codex_memories_write::rebuild_raw_memories_file_from_memories; -use codex_memories_write::rollout_summaries_dir; -use codex_memories_write::sync_rollout_summaries_from_memories; -use codex_protocol::ThreadId; -use codex_state::Stage1Output; -use codex_utils_absolute_path::AbsolutePathBuf; -use pretty_assertions::assert_eq; -use serde_json::Value; -use std::path::Path; -use std::path::PathBuf; -use tempfile::tempdir; - -#[test] -fn memory_root_uses_shared_global_path() { - let codex_home = AbsolutePathBuf::current_dir().expect("cwd").join("codex"); - assert_eq!(memory_root(&codex_home), codex_home.join("memories")); -} - -#[test] -fn stage_one_output_schema_requires_rollout_slug_and_keeps_it_nullable() { - let schema = crate::memories::phase1::output_schema(); - let properties = schema - .get("properties") - .and_then(Value::as_object) - .expect("properties object"); - let required = schema - .get("required") - .and_then(Value::as_array) - .expect("required array"); - - let mut required_keys = required - .iter() - .map(|key| key.as_str().expect("required key string")) - .collect::>(); - required_keys.sort_unstable(); - - assert!( - properties.contains_key("rollout_slug"), - "schema should declare rollout_slug" - ); - - let rollout_slug_type = properties - .get("rollout_slug") - .and_then(Value::as_object) - .and_then(|schema| schema.get("type")) - .and_then(Value::as_array) - .expect("rollout_slug type array"); - let mut rollout_slug_types = rollout_slug_type - .iter() - .map(|entry| entry.as_str().expect("type entry string")) - .collect::>(); - rollout_slug_types.sort_unstable(); - - assert_eq!( - required_keys, - vec!["raw_memory", "rollout_slug", "rollout_summary"] - ); - assert_eq!(rollout_slug_types, vec!["null", "string"]); -} - -#[tokio::test] -async fn clear_memory_root_contents_preserves_root_directory() { - let dir = tempdir().expect("tempdir"); - let root = dir.path().join("memories"); - let nested_dir = root.join("rollout_summaries"); - tokio::fs::create_dir_all(&nested_dir) - .await - .expect("create rollout summaries dir"); - tokio::fs::write(root.join("MEMORY.md"), "stale memory index\n") - .await - .expect("write memory index"); - tokio::fs::write(nested_dir.join("rollout.md"), "stale rollout\n") - .await - .expect("write rollout summary"); - - clear_memory_roots_contents(dir.path()) - .await - .expect("clear memory root contents"); - - assert!( - tokio::fs::try_exists(&root) - .await - .expect("check memory root existence"), - "memory root should still exist after clearing contents" - ); - let mut entries = tokio::fs::read_dir(&root) - .await - .expect("read memory root after clear"); - assert!( - entries - .next_entry() - .await - .expect("read next entry") - .is_none(), - "memory root should be empty after clearing contents" - ); -} - -#[cfg(unix)] -#[tokio::test] -async fn clear_memory_root_contents_rejects_symlinked_root() { - let dir = tempdir().expect("tempdir"); - let target = dir.path().join("outside"); - tokio::fs::create_dir_all(&target) - .await - .expect("create symlink target dir"); - let target_file = target.join("keep.txt"); - tokio::fs::write(&target_file, "keep\n") - .await - .expect("write target file"); - - let root = dir.path().join("memories"); - std::os::unix::fs::symlink(&target, &root).expect("create memory root symlink"); - - let err = clear_memory_roots_contents(dir.path()) - .await - .expect_err("symlinked memory root should be rejected"); - assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); - assert!( - tokio::fs::try_exists(&target_file) - .await - .expect("check target file existence"), - "rejecting a symlinked memory root should not delete the symlink target" - ); -} - -struct ConsolidatedOutputPaths { - memory_index: PathBuf, - memory_summary: PathBuf, - skill: PathBuf, -} - -async fn write_consolidated_outputs(root: &Path) -> ConsolidatedOutputPaths { - let paths = ConsolidatedOutputPaths { - memory_index: root.join("MEMORY.md"), - memory_summary: root.join("memory_summary.md"), - skill: root.join("skills/demo/SKILL.md"), - }; - - tokio::fs::write(&paths.memory_index, "consolidated memory index\n") - .await - .expect("write memory index"); - tokio::fs::write(&paths.memory_summary, "consolidated memory summary\n") - .await - .expect("write memory summary"); - tokio::fs::create_dir_all(paths.skill.parent().expect("skill parent")) - .await - .expect("create skill dir"); - tokio::fs::write(&paths.skill, "consolidated skill\n") - .await - .expect("write skill"); - - paths -} - -async fn assert_consolidated_outputs_exist(paths: &ConsolidatedOutputPaths, context: &str) { - assert!( - tokio::fs::try_exists(&paths.memory_index) - .await - .expect("check memory index existence"), - "{context} should leave MEMORY.md untouched" - ); - assert!( - tokio::fs::try_exists(&paths.memory_summary) - .await - .expect("check memory summary existence"), - "{context} should leave memory_summary.md untouched" - ); - assert!( - tokio::fs::try_exists(&paths.skill) - .await - .expect("check skill existence"), - "{context} should leave skills untouched" - ); -} - -#[tokio::test] -async fn sync_rollout_summaries_and_raw_memories_file_keeps_latest_memories_only() { - let dir = tempdir().expect("tempdir"); - let root = dir.path().join("memory"); - ensure_layout(&root).await.expect("ensure layout"); - - let keep_id = ThreadId::default().to_string(); - let drop_id = ThreadId::default().to_string(); - let keep_path = rollout_summaries_dir(&root).join(format!("{keep_id}.md")); - let drop_path = rollout_summaries_dir(&root).join(format!("{drop_id}.md")); - tokio::fs::write(&keep_path, "keep") - .await - .expect("write keep"); - tokio::fs::write(&drop_path, "drop") - .await - .expect("write drop"); - - let memories = vec![Stage1Output { - thread_id: ThreadId::try_from(keep_id.clone()).expect("thread id"), - source_updated_at: Utc.timestamp_opt(100, 0).single().expect("timestamp"), - raw_memory: "raw memory".to_string(), - rollout_summary: "short summary".to_string(), - rollout_slug: None, - rollout_path: PathBuf::from("/tmp/rollout-100.jsonl"), - cwd: PathBuf::from("/tmp/workspace"), - git_branch: None, - generated_at: Utc.timestamp_opt(101, 0).single().expect("timestamp"), - }]; - - sync_rollout_summaries_from_memories( - &root, - &memories, - DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION, - ) - .await - .expect("sync rollout summaries"); - rebuild_raw_memories_file_from_memories( - &root, - &memories, - DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION, - ) - .await - .expect("rebuild raw memories"); - - assert!( - !tokio::fs::try_exists(&keep_path) - .await - .expect("check stale keep path"), - "sync should prune stale filename that used thread id only" - ); - assert!( - !tokio::fs::try_exists(&drop_path) - .await - .expect("check stale drop path"), - "sync should prune stale filename for dropped thread" - ); - - let mut dir = tokio::fs::read_dir(rollout_summaries_dir(&root)) - .await - .expect("open rollout summaries dir"); - let mut files = Vec::new(); - while let Some(entry) = dir.next_entry().await.expect("read dir entry") { - files.push(entry.file_name().to_string_lossy().to_string()); - } - files.sort_unstable(); - assert_eq!(files.len(), 1); - let canonical_rollout_summary_file = &files[0]; - - let raw_memories = tokio::fs::read_to_string(raw_memories_file(&root)) - .await - .expect("read raw memories"); - assert!(raw_memories.contains("raw memory")); - assert!(raw_memories.contains(&keep_id)); - assert!(raw_memories.contains("cwd: /tmp/workspace")); - assert!(raw_memories.contains("rollout_path: /tmp/rollout-100.jsonl")); - assert!(raw_memories.contains(&format!( - "rollout_summary_file: {canonical_rollout_summary_file}" - ))); - let thread_header = format!("## Thread `{keep_id}`"); - let thread_pos = raw_memories - .find(&thread_header) - .expect("thread header should exist"); - let updated_pos = raw_memories[thread_pos..] - .find("updated_at: ") - .map(|offset| thread_pos + offset) - .expect("updated_at should exist after thread header"); - let cwd_pos = raw_memories[thread_pos..] - .find("cwd: /tmp/workspace") - .map(|offset| thread_pos + offset) - .expect("cwd should exist after thread header"); - let rollout_path_pos = raw_memories[thread_pos..] - .find("rollout_path: /tmp/rollout-100.jsonl") - .map(|offset| thread_pos + offset) - .expect("rollout_path should exist after thread header"); - let file_pos = raw_memories[thread_pos..] - .find(&format!( - "rollout_summary_file: {canonical_rollout_summary_file}" - )) - .map(|offset| thread_pos + offset) - .expect("rollout_summary_file should exist after thread header"); - assert!(thread_pos < updated_pos); - assert!(updated_pos < cwd_pos); - assert!(cwd_pos < rollout_path_pos); - assert!(rollout_path_pos < file_pos); -} - -#[tokio::test] -async fn sync_empty_inputs_preserves_consolidated_outputs() { - let dir = tempdir().expect("tempdir"); - let root = dir.path().join("memory"); - ensure_layout(&root).await.expect("ensure layout"); - - let stale_rollout_summary_path = rollout_summaries_dir(&root).join("stale.md"); - tokio::fs::write(&stale_rollout_summary_path, "stale summary\n") - .await - .expect("write stale rollout summary"); - let outputs = write_consolidated_outputs(&root).await; - - sync_rollout_summaries_from_memories( - &root, - &[], - DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION, - ) - .await - .expect("sync empty rollout summaries"); - rebuild_raw_memories_file_from_memories( - &root, - &[], - DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION, - ) - .await - .expect("rebuild empty raw memories"); - - assert!( - !tokio::fs::try_exists(&stale_rollout_summary_path) - .await - .expect("check stale rollout summary existence"), - "empty sync should prune stale rollout summaries" - ); - let raw_memories = tokio::fs::read_to_string(raw_memories_file(&root)) - .await - .expect("read raw memories"); - assert_eq!(raw_memories, "# Raw Memories\n\nNo raw memories yet.\n"); - assert_consolidated_outputs_exist(&outputs, "empty sync").await; -} - -#[tokio::test] -async fn sync_rollout_summaries_uses_timestamp_hash_and_sanitized_slug_filename() { - let dir = tempdir().expect("tempdir"); - let root = dir.path().join("memory"); - ensure_layout(&root).await.expect("ensure layout"); - - let thread_id = ThreadId::new(); - let stale_unslugged_path = rollout_summaries_dir(&root).join(format!("{thread_id}.md")); - let stale_old_slug_path = - rollout_summaries_dir(&root).join(format!("{thread_id}--old-slug.md")); - tokio::fs::write(&stale_unslugged_path, "stale") - .await - .expect("write stale unslugged file"); - tokio::fs::write(&stale_old_slug_path, "stale") - .await - .expect("write stale old-slug file"); - - let memories = vec![Stage1Output { - thread_id, - source_updated_at: Utc.timestamp_opt(200, 0).single().expect("timestamp"), - raw_memory: "raw memory".to_string(), - rollout_summary: "short summary".to_string(), - rollout_slug: Some("Unsafe Slug/With Spaces & Symbols + EXTRA_LONG_12345".to_string()), - rollout_path: PathBuf::from("/tmp/rollout-200.jsonl"), - cwd: PathBuf::from("/tmp/workspace"), - git_branch: Some("feature/memory-branch".to_string()), - generated_at: Utc.timestamp_opt(201, 0).single().expect("timestamp"), - }]; - - sync_rollout_summaries_from_memories( - &root, - &memories, - DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION, - ) - .await - .expect("sync rollout summaries"); - - let mut dir = tokio::fs::read_dir(rollout_summaries_dir(&root)) - .await - .expect("open rollout summaries dir"); - let mut files = Vec::new(); - while let Some(entry) = dir.next_entry().await.expect("read dir entry") { - files.push(entry.file_name().to_string_lossy().to_string()); - } - files.sort_unstable(); - - assert_eq!(files.len(), 1); - let file_name = &files[0]; - let stem = file_name - .strip_suffix(".md") - .expect("rollout summary file should end with .md"); - let (prefix, slug) = stem - .rsplit_once('-') - .expect("rollout summary filename should include slug"); - let (timestamp, short_hash) = prefix - .rsplit_once('-') - .expect("rollout summary filename should include short hash"); - - assert_eq!(timestamp.len(), 19, "timestamp should be second precision"); - let parsed_timestamp = chrono::NaiveDateTime::parse_from_str(timestamp, "%Y-%m-%dT%H-%M-%S"); - assert!( - parsed_timestamp.is_ok(), - "timestamp should use YYYY-MM-DDThh-mm-ss" - ); - assert_eq!(short_hash.len(), 4, "short hash should be exactly 4 chars"); - assert!( - short_hash.chars().all(|ch| ch.is_ascii_alphanumeric()), - "short hash should use only alphanumeric chars" - ); - assert!(slug.len() <= 60, "slug should be capped at 60 chars"); - assert!( - slug.chars() - .all(|ch| ch.is_ascii_lowercase() || ch.is_ascii_digit() || ch == '_'), - "slug should be file-safe lowercase ascii with underscores" - ); - - let summary = tokio::fs::read_to_string(rollout_summaries_dir(&root).join(file_name)) - .await - .expect("read rollout summary"); - assert!(summary.contains(&format!("thread_id: {thread_id}"))); - assert!(summary.contains("rollout_path: /tmp/rollout-200.jsonl")); - assert!(summary.contains("git_branch: feature/memory-branch")); - assert!( - !tokio::fs::try_exists(&stale_unslugged_path) - .await - .expect("check stale unslugged path"), - "slugged sync should prune stale unslugged filename for same thread" - ); - assert!( - !tokio::fs::try_exists(&stale_old_slug_path) - .await - .expect("check stale old slug path"), - "slugged sync should prune stale slugged filename for same thread" - ); -} - -#[tokio::test] -async fn rebuild_raw_memories_file_adds_canonical_rollout_summary_file_header() { - let dir = tempdir().expect("tempdir"); - let root = dir.path().join("memory"); - ensure_layout(&root).await.expect("ensure layout"); - - let thread_id = - ThreadId::try_from("0194f5a6-89ab-7cde-8123-456789abcdef").expect("valid thread id"); - let memories = vec![Stage1Output { - thread_id, - source_updated_at: Utc.timestamp_opt(200, 0).single().expect("timestamp"), - raw_memory: "\ ---- -description: Added a migration test -keywords: codex-state, migrations ---- -### Task 1: migration-test -task: add-migration-test -task_group: codex-state -task_outcome: success -- Added regression coverage for migration uniqueness. - -### Task 2: validate-migration -task: validate-migration-ordering -task_group: codex-state -task_outcome: success -- Confirmed no ordering regressions." - .to_string(), - rollout_summary: "short summary".to_string(), - rollout_slug: Some("Unsafe Slug/With Spaces & Symbols + EXTRA_LONG_12345".to_string()), - rollout_path: PathBuf::from("/tmp/rollout-200.jsonl"), - cwd: PathBuf::from("/tmp/workspace"), - git_branch: None, - generated_at: Utc.timestamp_opt(201, 0).single().expect("timestamp"), - }]; - - sync_rollout_summaries_from_memories( - &root, - &memories, - DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION, - ) - .await - .expect("sync rollout summaries"); - rebuild_raw_memories_file_from_memories( - &root, - &memories, - DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION, - ) - .await - .expect("rebuild raw memories"); - - let mut dir = tokio::fs::read_dir(rollout_summaries_dir(&root)) - .await - .expect("open rollout summaries dir"); - let mut files = Vec::new(); - while let Some(entry) = dir.next_entry().await.expect("read dir entry") { - files.push(entry.file_name().to_string_lossy().to_string()); - } - files.sort_unstable(); - assert_eq!(files.len(), 1); - let canonical_rollout_summary_file = &files[0]; - - let raw_memories = tokio::fs::read_to_string(raw_memories_file(&root)) - .await - .expect("read raw memories"); - let summary = tokio::fs::read_to_string( - rollout_summaries_dir(&root).join(canonical_rollout_summary_file), - ) - .await - .expect("read rollout summary"); - assert!(summary.contains("rollout_path: /tmp/rollout-200.jsonl")); - assert!(raw_memories.contains(&format!( - "rollout_summary_file: {canonical_rollout_summary_file}" - ))); - assert!(raw_memories.contains("description: Added a migration test")); - assert!(raw_memories.contains("### Task 1: migration-test")); - assert!(raw_memories.contains("task: add-migration-test")); - assert!(raw_memories.contains("task_group: codex-state")); - assert!(raw_memories.contains("task_outcome: success")); -} - -mod phase2 { - use crate::ThreadManager; - use crate::agent::AgentControl; - use crate::config::Config; - use crate::config::test_config; - use crate::memories::phase2; - use crate::session::session::Session; - use crate::session::tests::make_session_and_context; - use chrono::Duration as ChronoDuration; - use chrono::Utc; - use codex_config::Constrained; - use codex_config::types::McpServerConfig; - use codex_features::Feature; - use codex_login::CodexAuth; - use codex_memories_write::memory_root; - use codex_memories_write::raw_memories_file; - use codex_memories_write::rebuild_raw_memories_file_from_memories; - use codex_memories_write::rollout_summaries_dir; - use codex_memories_write::sync_rollout_summaries_from_memories; - use codex_memories_write::workspace::prepare_memory_workspace; - use codex_protocol::AgentPath; - use codex_protocol::ThreadId; - use codex_protocol::models::PermissionProfile; - use codex_protocol::permissions::FileSystemSandboxPolicy; - use codex_protocol::permissions::NetworkSandboxPolicy; - use codex_protocol::protocol::AskForApproval; - use codex_protocol::protocol::Op; - use codex_protocol::protocol::SandboxPolicy; - use codex_protocol::protocol::SessionSource; - use codex_state::Phase2JobClaimOutcome; - use codex_state::Stage1Output; - use codex_state::ThreadMetadataBuilder; - use std::collections::HashMap; - use std::path::PathBuf; - use std::sync::Arc; - use tempfile::TempDir; - - fn stage1_output_with_source_updated_at(source_updated_at: i64) -> Stage1Output { - Stage1Output { - thread_id: ThreadId::new(), - source_updated_at: chrono::DateTime::::from_timestamp(source_updated_at, 0) - .expect("valid source_updated_at timestamp"), - raw_memory: "raw memory".to_string(), - rollout_summary: "rollout summary".to_string(), - rollout_slug: None, - rollout_path: PathBuf::from("/tmp/rollout-summary.jsonl"), - cwd: PathBuf::from("/tmp/workspace"), - git_branch: None, - generated_at: chrono::DateTime::::from_timestamp(source_updated_at + 1, 0) - .expect("valid generated_at timestamp"), - } - } - - struct DispatchHarness { - _codex_home: TempDir, - config: Arc, - session: Arc, - manager: ThreadManager, - state_db: Arc, - } - - impl DispatchHarness { - async fn new() -> Self { - Self::new_with_config(|_| {}).await - } - - async fn new_with_config(configure: impl FnOnce(&mut Config)) -> Self { - let codex_home = tempfile::tempdir().expect("create temp codex home"); - let mut config = test_config().await; - config.codex_home = - codex_utils_absolute_path::AbsolutePathBuf::from_absolute_path(codex_home.path()) - .expect("codex home is absolute"); - config.cwd = config.codex_home.clone(); - let permission_profile = PermissionProfile::from_runtime_permissions( - &FileSystemSandboxPolicy::unrestricted(), - NetworkSandboxPolicy::Enabled, - ); - config - .permissions - .set_permission_profile(permission_profile) - .expect("permissions are configurable"); - configure(&mut config); - let config = Arc::new(config); - - let state_db = codex_state::StateRuntime::init( - config.codex_home.to_path_buf(), - config.model_provider_id.clone(), - ) - .await - .expect("initialize state db"); - - let manager = ThreadManager::with_models_provider_and_home_for_tests( - CodexAuth::from_api_key("dummy"), - config.model_provider.clone(), - config.codex_home.to_path_buf(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); - let (mut session, _turn_context) = make_session_and_context().await; - session.services.state_db = Some(Arc::clone(&state_db)); - session.services.agent_control = manager.agent_control(); - - Self { - _codex_home: codex_home, - config, - session: Arc::new(session), - manager, - state_db, - } - } - - async fn seed_stage1_output(&self, source_updated_at: i64) -> ThreadId { - let thread_id = ThreadId::new(); - let mut metadata_builder = ThreadMetadataBuilder::new( - thread_id, - self.config - .codex_home - .join(format!("rollout-{thread_id}.jsonl")) - .to_path_buf(), - Utc::now(), - SessionSource::Cli, - ); - metadata_builder.cwd = self.config.cwd.to_path_buf(); - metadata_builder.model_provider = Some(self.config.model_provider_id.clone()); - let metadata = metadata_builder.build(&self.config.model_provider_id); - - self.state_db - .upsert_thread(&metadata) - .await - .expect("upsert thread metadata"); - - let claim = self - .state_db - .try_claim_stage1_job( - thread_id, - self.session.conversation_id, - source_updated_at, - /*lease_seconds*/ 3_600, - /*max_running_jobs*/ 64, - ) - .await - .expect("claim stage-1 job"); - let ownership_token = match claim { - codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, - other => panic!("unexpected stage-1 claim outcome: {other:?}"), - }; - assert!( - self.state_db - .mark_stage1_job_succeeded( - thread_id, - &ownership_token, - source_updated_at, - "raw memory", - "rollout summary", - /*rollout_slug*/ None, - ) - .await - .expect("mark stage-1 success"), - "stage-1 success should enqueue global consolidation" - ); - thread_id - } - - async fn shutdown_threads(&self) { - let report = self - .manager - .shutdown_all_threads_bounded(std::time::Duration::from_secs(10)) - .await; - assert!(report.submit_failed.is_empty()); - assert!(report.timed_out.is_empty()); - } - - fn user_input_ops_count(&self) -> usize { - self.manager - .captured_ops() - .into_iter() - .filter(|(_, op)| matches!(op, Op::UserInput { .. })) - .count() - } - } - - #[test] - fn completion_watermark_never_regresses_below_claimed_input_watermark() { - let stage1_output = stage1_output_with_source_updated_at(/*source_updated_at*/ 123); - - let completion = phase2::get_watermark(/*claimed_watermark*/ 1_000, &[stage1_output]); - pretty_assertions::assert_eq!(completion, 1_000); - } - - #[test] - fn completion_watermark_uses_claimed_watermark_when_there_are_no_memories() { - let completion = phase2::get_watermark(/*claimed_watermark*/ 777, &[]); - pretty_assertions::assert_eq!(completion, 777); - } - - #[test] - fn completion_watermark_uses_latest_memory_timestamp_when_it_is_newer() { - let older = stage1_output_with_source_updated_at(/*source_updated_at*/ 123); - let newer = stage1_output_with_source_updated_at(/*source_updated_at*/ 456); - - let completion = phase2::get_watermark(/*claimed_watermark*/ 200, &[older, newer]); - pretty_assertions::assert_eq!(completion, 456); - } - - #[tokio::test] - async fn dispatch_skips_when_memory_workspace_is_not_dirty() { - let harness = DispatchHarness::new().await; - let root = memory_root(&harness.config.codex_home); - rebuild_raw_memories_file_from_memories( - &root, - &[], - /*max_raw_memories_for_consolidation*/ 0, - ) - .await - .expect("write empty raw memories baseline"); - let outputs = super::write_consolidated_outputs(&root).await; - prepare_memory_workspace(&root) - .await - .expect("commit empty memory workspace as baseline"); - - phase2::run(&harness.session, Arc::clone(&harness.config)).await; - - pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); - super::assert_consolidated_outputs_exist(&outputs, "clean no-input phase2").await; - let thread_ids = harness.manager.list_thread_ids().await; - pretty_assertions::assert_eq!(thread_ids.len(), 0); - } - - #[tokio::test] - async fn dispatch_uses_git_dirty_state_without_db_dirty_watermark() { - let harness = DispatchHarness::new().await; - let root = memory_root(&harness.config.codex_home); - rebuild_raw_memories_file_from_memories( - &root, - &[], - /*max_raw_memories_for_consolidation*/ 0, - ) - .await - .expect("write empty raw memories baseline"); - prepare_memory_workspace(&root) - .await - .expect("commit empty memory workspace as baseline"); - let extension_resource = root - .join("extensions") - .join("chronicle") - .join("resources") - .join("2026-04-22T12-00-00-abcd-10min-memory.md"); - tokio::fs::create_dir_all( - extension_resource - .parent() - .expect("extension resource parent"), - ) - .await - .expect("create extension resource dir"); - tokio::fs::write( - root.join("extensions/chronicle/instructions.md"), - "instructions\n", - ) - .await - .expect("write extension instructions"); - tokio::fs::write(&extension_resource, "extension memory\n") - .await - .expect("write extension resource"); - - phase2::run(&harness.session, Arc::clone(&harness.config)).await; - - pretty_assertions::assert_eq!(harness.user_input_ops_count(), 1); - let workspace_diff = tokio::fs::read_to_string(root.join("phase2_workspace_diff.md")) - .await - .expect("read workspace diff"); - assert!( - workspace_diff.contains("- A extensions/chronicle/instructions.md"), - "git-only extension instructions should dirty phase2: {workspace_diff}" - ); - assert!( - workspace_diff.contains("- A extensions/chronicle/resources/"), - "git-only extension resource should dirty phase2: {workspace_diff}" - ); - let thread_ids = harness.manager.list_thread_ids().await; - pretty_assertions::assert_eq!(thread_ids.len(), 1); - - harness.shutdown_threads().await; - } - - #[tokio::test] - async fn dispatch_skips_when_global_job_is_already_running() { - let harness = DispatchHarness::new().await; - harness - .state_db - .enqueue_global_consolidation(/*input_watermark*/ 123) - .await - .expect("enqueue global consolidation"); - let claimed = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) - .await - .expect("claim running global lock"); - assert!( - matches!(claimed, Phase2JobClaimOutcome::Claimed { .. }), - "precondition should claim the running lock" - ); - - phase2::run(&harness.session, Arc::clone(&harness.config)).await; - - let running_claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) - .await - .expect("claim while lock is still running"); - pretty_assertions::assert_eq!(running_claim, Phase2JobClaimOutcome::SkippedRunning); - pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); - let thread_ids = harness.manager.list_thread_ids().await; - pretty_assertions::assert_eq!(thread_ids.len(), 0); - } - - #[tokio::test] - async fn dispatch_reclaims_stale_global_lock_and_starts_consolidation() { - let harness = DispatchHarness::new_with_config(|config| { - let server: McpServerConfig = - toml::from_str("command = \"docs-server\"").expect("deserialize MCP server"); - config - .mcp_servers - .set(HashMap::from([("docs".to_string(), server)])) - .expect("parent MCP servers are configurable"); - config - .features - .enable(Feature::Apps) - .expect("apps feature is configurable"); - config - .features - .enable(Feature::Plugins) - .expect("plugins feature is configurable"); - config.include_apps_instructions = true; - }) - .await; - harness.seed_stage1_output(Utc::now().timestamp()).await; - - let stale_claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 0) - .await - .expect("claim stale global lock"); - assert!( - matches!(stale_claim, Phase2JobClaimOutcome::Claimed { .. }), - "stale lock precondition should be claimed" - ); - - phase2::run(&harness.session, Arc::clone(&harness.config)).await; - - let post_dispatch_claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) - .await - .expect("claim after stale lock dispatch"); - assert!( - matches!( - post_dispatch_claim, - Phase2JobClaimOutcome::SkippedRunning - | Phase2JobClaimOutcome::SkippedRetryUnavailable - ), - "stale-lock dispatch should either keep the reclaimed job running or finish it before re-claim" - ); - - let user_input_ops = harness.user_input_ops_count(); - pretty_assertions::assert_eq!(user_input_ops, 1); - let thread_ids = harness.manager.list_thread_ids().await; - pretty_assertions::assert_eq!(thread_ids.len(), 1); - let thread_id = thread_ids[0]; - let subagent = harness - .manager - .get_thread(thread_id) - .await - .expect("get consolidation thread"); - let config_snapshot = subagent.config_snapshot().await; - pretty_assertions::assert_eq!(config_snapshot.approval_policy, AskForApproval::Never); - assert!(config_snapshot.ephemeral); - pretty_assertions::assert_eq!( - config_snapshot.cwd.as_path(), - memory_root(&harness.config.codex_home).as_path() - ); - let sandbox_policy = config_snapshot.sandbox_policy(); - match &sandbox_policy { - SandboxPolicy::WorkspaceWrite { network_access, .. } => { - assert!(!*network_access); - let effective_writable_roots: Vec<_> = sandbox_policy - .get_writable_roots_with_cwd(config_snapshot.cwd.as_path()) - .into_iter() - .map(|root| root.root) - .collect(); - pretty_assertions::assert_eq!( - effective_writable_roots.as_slice(), - [memory_root(&harness.config.codex_home)], - "consolidation subagent should only be able to write the memory root" - ); - } - other => panic!("unexpected sandbox policy: {other:?}"), - } - pretty_assertions::assert_eq!( - config_snapshot.session_source.get_agent_path(), - Some(AgentPath::morpheus()) - ); - assert!( - harness - .session - .services - .agent_control - .get_agent_metadata(thread_id) - .is_none(), - "memory consolidation should not be registered in the root collab agent registry" - ); - let turn_context = subagent.codex.session.new_default_turn().await; - let file_system_sandbox_policy = turn_context.file_system_sandbox_policy(); - let legacy_file_system_sandbox_policy = - FileSystemSandboxPolicy::from_legacy_sandbox_policy_for_cwd( - &sandbox_policy, - config_snapshot.cwd.as_path(), - ); - assert!( - file_system_sandbox_policy.is_semantically_equivalent_to( - &legacy_file_system_sandbox_policy, - config_snapshot.cwd.as_path(), - ), - "consolidation subagent split filesystem policy should match the memory-root legacy policy" - ); - assert!( - file_system_sandbox_policy.can_write_path_with_cwd( - memory_root(&harness.config.codex_home).as_path(), - config_snapshot.cwd.as_path(), - ), - "consolidation subagent should be able to write the memory root" - ); - assert!( - !file_system_sandbox_policy.can_write_path_with_cwd( - harness.config.codex_home.join("config.toml").as_path(), - config_snapshot.cwd.as_path(), - ), - "consolidation subagent should not inherit codex_home write access" - ); - pretty_assertions::assert_eq!( - turn_context.network_sandbox_policy(), - NetworkSandboxPolicy::Restricted, - "consolidation subagent split network policy should preserve no-network sandboxing" - ); - assert!( - !turn_context.features.enabled(Feature::MemoryTool), - "consolidation subagent should have the memories feature disabled" - ); - assert!( - turn_context.config.mcp_servers.get().is_empty(), - "consolidation subagent should not inherit configured MCP servers" - ); - assert!( - !subagent - .codex - .session - .services - .mcp_connection_manager - .read() - .await - .has_servers(), - "consolidation subagent should not initialize MCP servers" - ); - assert!( - !turn_context.features.enabled(Feature::Apps), - "consolidation subagent should not expose app-backed MCP" - ); - assert!( - !turn_context.features.enabled(Feature::Plugins), - "consolidation subagent should not expose plugin-backed MCP" - ); - assert!( - !turn_context.config.include_apps_instructions, - "consolidation subagent should not include apps instructions" - ); - assert!( - !turn_context.config.memories.generate_memories, - "consolidation subagent should not generate memories" - ); - assert!( - !turn_context.config.memories.use_memories, - "consolidation subagent should not read memories" - ); - assert!( - subagent.rollout_path().is_none(), - "ephemeral consolidation thread should not materialize a rollout" - ); - let memory_mode = harness - .state_db - .get_thread_memory_mode(thread_id) - .await - .expect("read consolidation thread memory mode"); - pretty_assertions::assert_eq!(memory_mode, None); - - harness.shutdown_threads().await; - } - - #[tokio::test] - async fn dispatch_with_empty_stage1_outputs_spawns_for_workspace_changes() { - let harness = DispatchHarness::new().await; - let root = memory_root(&harness.config.codex_home); - let summaries_dir = rollout_summaries_dir(&root); - tokio::fs::create_dir_all(&summaries_dir) - .await - .expect("create rollout summaries dir"); - - let stale_summary_path = summaries_dir.join(format!("{}.md", ThreadId::new())); - tokio::fs::write(&stale_summary_path, "stale summary\n") - .await - .expect("write stale rollout summary"); - let raw_memories_path = raw_memories_file(&root); - tokio::fs::write(&raw_memories_path, "stale raw memories\n") - .await - .expect("write stale raw memories"); - let outputs = super::write_consolidated_outputs(&root).await; - - harness - .state_db - .enqueue_global_consolidation(/*input_watermark*/ 999) - .await - .expect("enqueue global consolidation"); - - phase2::run(&harness.session, Arc::clone(&harness.config)).await; - - assert!( - !tokio::fs::try_exists(&stale_summary_path) - .await - .expect("check stale summary existence"), - "empty consolidation should prune stale rollout summary files" - ); - let raw_memories = tokio::fs::read_to_string(&raw_memories_path) - .await - .expect("read rebuilt raw memories"); - pretty_assertions::assert_eq!(raw_memories, "# Raw Memories\n\nNo raw memories yet.\n"); - super::assert_consolidated_outputs_exist(&outputs, "empty consolidation").await; - let next_claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) - .await - .expect("claim global job after empty consolidation dispatch"); - pretty_assertions::assert_eq!(next_claim, Phase2JobClaimOutcome::SkippedRunning); - pretty_assertions::assert_eq!(harness.user_input_ops_count(), 1); - let thread_ids = harness.manager.list_thread_ids().await; - pretty_assertions::assert_eq!(thread_ids.len(), 1); - - harness.shutdown_threads().await; - } - - #[tokio::test] - async fn dispatch_with_empty_selected_inputs_preserves_consolidated_outputs() { - let harness = DispatchHarness::new().await; - let source_updated_at = Utc::now().timestamp(); - let thread_id = harness.seed_stage1_output(source_updated_at).await; - let root = memory_root(&harness.config.codex_home); - let selected = harness - .state_db - .get_phase2_input_selection(/*n*/ 1, /*max_unused_days*/ 30) - .await - .expect("load phase2 input selection"); - sync_rollout_summaries_from_memories(&root, &selected, selected.len()) - .await - .expect("sync selected rollout summaries"); - rebuild_raw_memories_file_from_memories(&root, &selected, selected.len()) - .await - .expect("sync selected raw memories"); - let outputs = super::write_consolidated_outputs(&root).await; - prepare_memory_workspace(&root) - .await - .expect("commit current memory workspace as baseline"); - - let claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) - .await - .expect("claim global phase2 job"); - let Phase2JobClaimOutcome::Claimed { - ownership_token, .. - } = claim - else { - panic!("unexpected phase2 claim outcome: {claim:?}"); - }; - assert!( - harness - .state_db - .mark_global_phase2_job_succeeded(&ownership_token, source_updated_at, &selected) - .await - .expect("mark phase2 succeeded"), - "phase2 success should update selected baseline" - ); - assert!( - harness - .state_db - .mark_thread_memory_mode_polluted(thread_id) - .await - .expect("mark thread polluted"), - "polluted selected thread should enqueue phase2 forgetting" - ); - - phase2::run(&harness.session, Arc::clone(&harness.config)).await; - - pretty_assertions::assert_eq!(harness.user_input_ops_count(), 1); - super::assert_consolidated_outputs_exist(&outputs, "empty selected phase2").await; - let workspace_diff = tokio::fs::read_to_string(root.join("phase2_workspace_diff.md")) - .await - .expect("read workspace diff"); - assert!( - workspace_diff.contains("- D rollout_summaries/"), - "empty selected phase2 should surface deleted rollout summaries: {workspace_diff}" - ); - assert!( - !workspace_diff.contains("- D MEMORY.md"), - "empty selected phase2 should not delete MEMORY.md directly: {workspace_diff}" - ); - assert!( - !workspace_diff.contains("- D memory_summary.md"), - "empty selected phase2 should not delete memory_summary.md directly: {workspace_diff}" - ); - assert!( - !workspace_diff.contains("- D skills/demo/SKILL.md"), - "empty selected phase2 should not delete skills directly: {workspace_diff}" - ); - - harness.shutdown_threads().await; - } - - #[tokio::test] - async fn dispatch_with_clean_workspace_rebuilds_selected_phase2_baseline() { - let harness = DispatchHarness::new().await; - let source_updated_at = (Utc::now() - ChronoDuration::days(1)).timestamp(); - let thread_id = harness.seed_stage1_output(source_updated_at).await; - let root = memory_root(&harness.config.codex_home); - let selected = harness - .state_db - .get_phase2_input_selection(/*n*/ 1, /*max_unused_days*/ 30) - .await - .expect("load phase2 input selection"); - - sync_rollout_summaries_from_memories(&root, &selected, selected.len()) - .await - .expect("sync selected rollout summaries"); - rebuild_raw_memories_file_from_memories(&root, &selected, selected.len()) - .await - .expect("sync selected raw memories"); - prepare_memory_workspace(&root) - .await - .expect("commit current memory workspace as baseline"); - - phase2::run(&harness.session, Arc::clone(&harness.config)).await; - - pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); - let pruned = harness - .state_db - .prune_stage1_outputs_for_retention(/*max_unused_days*/ 0, /*limit*/ 10) - .await - .expect("prune stage1 outputs after clean phase2"); - pretty_assertions::assert_eq!(pruned, 0); - - let selected = harness - .state_db - .get_phase2_input_selection(/*n*/ 1, /*max_unused_days*/ 30) - .await - .expect("load phase2 input selection after clean workspace success"); - pretty_assertions::assert_eq!(selected.len(), 1); - pretty_assertions::assert_eq!(selected[0].thread_id, thread_id); - } - - #[tokio::test] - async fn dispatch_marks_job_for_retry_when_sandbox_policy_cannot_be_overridden() { - let harness = DispatchHarness::new().await; - harness - .state_db - .enqueue_global_consolidation(/*input_watermark*/ 99) - .await - .expect("enqueue global consolidation"); - let mut constrained_config = harness.config.as_ref().clone(); - constrained_config.permissions.permission_profile = Constrained::allow_only( - PermissionProfile::from_legacy_sandbox_policy(&SandboxPolicy::DangerFullAccess), - ); - - phase2::run(&harness.session, Arc::new(constrained_config)).await; - - let retry_claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) - .await - .expect("claim global job after sandbox policy failure"); - pretty_assertions::assert_eq!(retry_claim, Phase2JobClaimOutcome::SkippedRetryUnavailable); - pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); - let thread_ids = harness.manager.list_thread_ids().await; - pretty_assertions::assert_eq!(thread_ids.len(), 0); - } - - #[tokio::test] - async fn dispatch_marks_job_for_retry_when_syncing_artifacts_fails() { - let harness = DispatchHarness::new().await; - harness.seed_stage1_output(/*source_updated_at*/ 100).await; - let root = memory_root(&harness.config.codex_home); - tokio::fs::write(&root, "not a directory") - .await - .expect("create file at memory root"); - - phase2::run(&harness.session, Arc::clone(&harness.config)).await; - - let retry_claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) - .await - .expect("claim global job after sync failure"); - pretty_assertions::assert_eq!(retry_claim, Phase2JobClaimOutcome::SkippedRetryUnavailable); - pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); - let thread_ids = harness.manager.list_thread_ids().await; - pretty_assertions::assert_eq!(thread_ids.len(), 0); - } - - #[tokio::test] - async fn dispatch_marks_job_for_retry_when_rebuilding_raw_memories_fails() { - let harness = DispatchHarness::new().await; - harness.seed_stage1_output(/*source_updated_at*/ 100).await; - let root = memory_root(&harness.config.codex_home); - tokio::fs::create_dir_all(raw_memories_file(&root)) - .await - .expect("create raw_memories.md as a directory"); - - phase2::run(&harness.session, Arc::clone(&harness.config)).await; - - let retry_claim = harness - .state_db - .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) - .await - .expect("claim global job after rebuild failure"); - pretty_assertions::assert_eq!(retry_claim, Phase2JobClaimOutcome::SkippedRetryUnavailable); - pretty_assertions::assert_eq!(harness.user_input_ops_count(), 0); - let thread_ids = harness.manager.list_thread_ids().await; - pretty_assertions::assert_eq!(thread_ids.len(), 0); - } - - #[tokio::test] - async fn dispatch_marks_job_for_retry_when_spawn_agent_fails() { - let codex_home = tempfile::tempdir().expect("create temp codex home"); - let mut config = test_config().await; - config.codex_home = - codex_utils_absolute_path::AbsolutePathBuf::from_absolute_path(codex_home.path()) - .expect("codex home is absolute"); - config.cwd = config.codex_home.clone(); - let config = Arc::new(config); - - let state_db = codex_state::StateRuntime::init( - config.codex_home.to_path_buf(), - config.model_provider_id.clone(), - ) - .await - .expect("initialize state db"); - - let (mut session, _turn_context) = make_session_and_context().await; - session.services.state_db = Some(Arc::clone(&state_db)); - session.services.agent_control = AgentControl::default(); - let session = Arc::new(session); - - let thread_id = ThreadId::new(); - let mut metadata_builder = ThreadMetadataBuilder::new( - thread_id, - config - .codex_home - .join(format!("rollout-{thread_id}.jsonl")) - .to_path_buf(), - Utc::now(), - SessionSource::Cli, - ); - metadata_builder.cwd = config.cwd.to_path_buf(); - metadata_builder.model_provider = Some(config.model_provider_id.clone()); - let metadata = metadata_builder.build(&config.model_provider_id); - state_db - .upsert_thread(&metadata) - .await - .expect("upsert thread metadata"); - - let claim = state_db - .try_claim_stage1_job( - thread_id, - session.conversation_id, - /*source_updated_at*/ 100, - /*lease_seconds*/ 3_600, - /*max_running_jobs*/ 64, - ) - .await - .expect("claim stage-1 job"); - let ownership_token = match claim { - codex_state::Stage1JobClaimOutcome::Claimed { ownership_token } => ownership_token, - other => panic!("unexpected stage-1 claim outcome: {other:?}"), - }; - assert!( - state_db - .mark_stage1_job_succeeded( - thread_id, - &ownership_token, - /*source_updated_at*/ 100, - "raw memory", - "rollout summary", - /*rollout_slug*/ None, - ) - .await - .expect("mark stage-1 success"), - "stage-1 success should enqueue global consolidation" - ); - - let chronicle_resources = config - .codex_home - .join("memories/extensions/chronicle/resources"); - tokio::fs::create_dir_all(&chronicle_resources) - .await - .expect("create chronicle resources"); - tokio::fs::write( - config - .codex_home - .join("memories/extensions/chronicle/instructions.md"), - "instructions", - ) - .await - .expect("write chronicle instructions"); - let old_file = chronicle_resources.join(format!( - "{}-abcd-10min-old.md", - (Utc::now() - ChronoDuration::days(8)).format("%Y-%m-%dT%H-%M-%S") - )); - tokio::fs::write(&old_file, "old resource") - .await - .expect("write old extension resource"); - - phase2::run(&session, Arc::clone(&config)).await; - - let retry_claim = state_db - .try_claim_global_phase2_job(ThreadId::new(), /*lease_seconds*/ 3_600) - .await - .expect("claim global job after spawn failure"); - pretty_assertions::assert_eq!( - retry_claim, - Phase2JobClaimOutcome::SkippedRetryUnavailable, - "spawn failures should leave the job in retry backoff instead of running" - ); - assert!( - !tokio::fs::try_exists(&old_file) - .await - .expect("check old extension resource"), - "old extension resources should still be pruned on failed phase2 attempts" - ); - let workspace_diff = - tokio::fs::read_to_string(config.codex_home.join("memories/phase2_workspace_diff.md")) - .await - .expect("read workspace diff"); - assert!( - workspace_diff.contains("- D extensions/chronicle/resources/"), - "spawn failures should keep a retryable workspace diff: {workspace_diff}" - ); - } -} diff --git a/codex-rs/core/src/rollout.rs b/codex-rs/core/src/rollout.rs index 4d282cdb5..d4ac5c699 100644 --- a/codex-rs/core/src/rollout.rs +++ b/codex-rs/core/src/rollout.rs @@ -50,10 +50,6 @@ pub(crate) mod list { pub use codex_rollout::find_thread_path_by_id_str; } -pub(crate) mod policy { - pub use codex_rollout::should_persist_response_item_for_memories; -} - pub(crate) mod recorder { pub use codex_rollout::RolloutRecorder; } diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index d4ca521c9..70900b446 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -670,66 +670,6 @@ pub async fn compact(sess: &Arc, sub_id: String) { .await; } -pub async fn drop_memories(sess: &Arc, config: &Arc, sub_id: String) { - let mut errors = Vec::new(); - - if let Some(state_db) = sess.services.state_db.as_deref() { - if let Err(err) = state_db.clear_memory_data().await { - errors.push(format!("failed clearing memory rows from state db: {err}")); - } - } else { - errors.push("state db unavailable; memory rows were not cleared".to_string()); - } - - if let Err(err) = codex_memories_write::clear_memory_roots_contents(&config.codex_home).await { - errors.push(format!( - "failed clearing memory directories under {}: {err}", - config.codex_home.display() - )); - } - - if errors.is_empty() { - let memory_root = codex_memories_write::memory_root(&config.codex_home); - sess.send_event_raw(Event { - id: sub_id, - msg: EventMsg::Warning(WarningEvent { - message: format!( - "Dropped memories at {} and cleared memory rows from state db.", - memory_root.display() - ), - }), - }) - .await; - return; - } - - sess.send_event_raw(Event { - id: sub_id, - msg: EventMsg::Error(ErrorEvent { - message: format!("Memory drop completed with errors: {}", errors.join("; ")), - codex_error_info: Some(CodexErrorInfo::Other), - }), - }) - .await; -} - -pub async fn update_memories(sess: &Arc, config: &Arc, sub_id: String) { - let session_source = { - let state = sess.state.lock().await; - state.session_configuration.session_source.clone() - }; - - crate::memories::start_memories_startup_task(sess, Arc::clone(config), &session_source); - - sess.send_event_raw(Event { - id: sub_id.clone(), - msg: EventMsg::Warning(WarningEvent { - message: "Memory update triggered.".to_string(), - }), - }) - .await; -} - pub async fn thread_rollback(sess: &Arc, sub_id: String, num_turns: u32) { if num_turns == 0 { sess.send_event_raw(Event { @@ -1181,14 +1121,6 @@ pub(super) async fn submission_loop( compact(&sess, sub.id.clone()).await; false } - Op::DropMemories => { - drop_memories(&sess, &config, sub.id.clone()).await; - false - } - Op::UpdateMemories => { - update_memories(&sess, &config, sub.id.clone()).await; - false - } Op::ThreadRollback { num_turns } => { thread_rollback(&sess, sub.id.clone(), num_turns).await; false diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index cb88f4446..8e76c7fee 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -271,7 +271,6 @@ use crate::context::UserInstructions; use crate::exec_policy::ExecPolicyUpdateError; use crate::guardian::GuardianReviewSessionManager; use crate::mcp::McpManager; -use crate::memories; use crate::network_policy_decision::execpolicy_network_rule_amendment; use crate::plugins::PluginsManager; use crate::rollout::map_session_init_error; @@ -514,9 +513,10 @@ impl Codex { }; let config = Arc::new(config); - let refresh_strategy = match session_source { - SessionSource::SubAgent(_) => codex_models_manager::manager::RefreshStrategy::Offline, - _ => codex_models_manager::manager::RefreshStrategy::OnlineIfUncached, + let refresh_strategy = if session_source.is_non_root_agent() { + codex_models_manager::manager::RefreshStrategy::Offline + } else { + codex_models_manager::manager::RefreshStrategy::OnlineIfUncached }; if config.model.is_none() || !matches!( @@ -1142,10 +1142,10 @@ impl Session { let turn_context = self.new_default_turn().await; let is_subagent = { let state = self.state.lock().await; - matches!( - state.session_configuration.session_source, - SessionSource::SubAgent(_) - ) + state + .session_configuration + .session_source + .is_non_root_agent() }; let has_prior_user_turns = initial_history_has_prior_user_turns(&conversation_history); { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 1d37db1b8..9a13ce335 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -425,10 +425,7 @@ impl Session { session_init.ephemeral = config.ephemeral, )); - let is_subagent = matches!( - session_configuration.session_source, - SessionSource::SubAgent(_) - ); + let is_subagent = session_configuration.session_source.is_non_root_agent(); let history_meta_fut = async { if is_subagent { (0, 0) @@ -989,12 +986,6 @@ impl Session { state.set_pending_session_start_source(Some(session_start_source)); } - memories::start_memories_startup_task( - &sess, - Arc::clone(&config), - &session_configuration.session_source, - ); - Ok(sess) } .await; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 17fb16ad3..5d463f934 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -211,9 +211,10 @@ pub struct ThreadManager { _test_codex_home_guard: Option, } -pub struct StartThreadWithToolsOptions { +pub struct StartThreadOptions { pub config: Config, pub initial_history: InitialHistory, + pub session_source: Option, pub dynamic_tools: Vec, pub persist_extended_history: bool, pub metrics_service_name: Option, @@ -540,34 +541,39 @@ impl ThreadManager { self.state.environment_manager.as_ref(), &config.cwd, ); - Box::pin( - self.start_thread_with_tools_and_service_name(StartThreadWithToolsOptions { - config, - initial_history: InitialHistory::New, - dynamic_tools, - persist_extended_history, - metrics_service_name: None, - parent_trace: None, - environments, - }), - ) + Box::pin(self.start_thread_with_options(StartThreadOptions { + config, + initial_history: InitialHistory::New, + session_source: None, + dynamic_tools, + persist_extended_history, + metrics_service_name: None, + parent_trace: None, + environments, + })) .await } - pub async fn start_thread_with_tools_and_service_name( + pub async fn start_thread_with_options( &self, - options: StartThreadWithToolsOptions, + options: StartThreadOptions, ) -> CodexResult { let thread_store = configured_thread_store(&options.config); - Box::pin(self.state.spawn_thread( + let session_source = options + .session_source + .unwrap_or_else(|| self.state.session_source.clone()); + Box::pin(self.state.spawn_thread_with_source( options.config, thread_store, options.initial_history, Arc::clone(&self.state.auth_manager), self.agent_control(), + session_source, options.dynamic_tools, options.persist_extended_history, options.metrics_service_name, + /*inherited_shell_snapshot*/ None, + /*inherited_exec_policy*/ None, options.parent_trace, options.environments, /*user_shell_override*/ None, @@ -831,16 +837,23 @@ impl ThreadManager { impl ThreadManagerState { pub(crate) async fn list_thread_ids(&self) -> Vec { - self.threads.read().await.keys().copied().collect() + self.threads + .read() + .await + .iter() + .filter_map(|(thread_id, thread)| { + (!thread.session_source.is_internal()).then_some(*thread_id) + }) + .collect() } /// Fetch a thread by ID or return ThreadNotFound. pub(crate) async fn get_thread(&self, thread_id: ThreadId) -> CodexResult> { let threads = self.threads.read().await; - threads - .get(&thread_id) - .cloned() - .ok_or_else(|| CodexErr::ThreadNotFound(thread_id)) + match threads.get(&thread_id) { + Some(thread) if !thread.session_source.is_internal() => Ok(thread.clone()), + Some(_) | None => Err(CodexErr::ThreadNotFound(thread_id)), + } } /// Send an operation to a thread by ID. @@ -1063,6 +1076,7 @@ impl ThreadManagerState { let parent_rollout_thread_trace = self .parent_rollout_thread_trace_for_source(&session_source, &initial_history) .await; + let tracked_session_source = session_source.clone(); let CodexSpawnOk { codex, thread_id, .. } = Codex::spawn(CodexSpawnArgs { @@ -1091,7 +1105,7 @@ impl ThreadManagerState { }) .await?; let new_thread = self - .finalize_thread_spawn(codex, thread_id, watch_registration) + .finalize_thread_spawn(codex, thread_id, tracked_session_source, watch_registration) .await?; if is_resumed_thread && let Err(err) = new_thread.thread.apply_goal_resume_runtime_effects().await @@ -1105,6 +1119,7 @@ impl ThreadManagerState { &self, codex: Codex, thread_id: ThreadId, + session_source: SessionSource, watch_registration: crate::file_watcher::WatchRegistration, ) -> CodexResult { let event = codex.next_event().await?; @@ -1121,6 +1136,7 @@ impl ThreadManagerState { let thread = Arc::new(CodexThread::new( codex, session_configured.rollout_path.clone(), + session_source, watch_registration, )); let mut threads = self.threads.write().await; diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index ef4420d1e..9307927b0 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -13,6 +13,9 @@ use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelsResponse; use codex_protocol::protocol::AgentMessageEvent; +use codex_protocol::protocol::InitialHistory; +use codex_protocol::protocol::InternalSessionSource; +use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::protocol::UserMessageEvent; use core_test_support::PathBufExt; @@ -312,9 +315,10 @@ async fn start_thread_accepts_explicit_environment_when_default_environment_is_d ); let thread = manager - .start_thread_with_tools_and_service_name(StartThreadWithToolsOptions { + .start_thread_with_options(StartThreadOptions { config: config.clone(), initial_history: InitialHistory::New, + session_source: None, dynamic_tools: Vec::new(), persist_extended_history: false, metrics_service_name: None, @@ -330,6 +334,48 @@ async fn start_thread_accepts_explicit_environment_when_default_environment_is_d assert_eq!(manager.list_thread_ids().await, vec![thread.thread_id]); } +#[tokio::test] +async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ); + let thread = manager + .start_thread_with_options(StartThreadOptions { + config, + initial_history: InitialHistory::New, + session_source: Some(SessionSource::Internal( + InternalSessionSource::MemoryConsolidation, + )), + dynamic_tools: Vec::new(), + persist_extended_history: false, + metrics_service_name: None, + parent_trace: None, + environments: Vec::new(), + }) + .await + .expect("internal thread should start"); + + assert_eq!(manager.list_thread_ids().await, Vec::new()); + assert!(manager.get_thread(thread.thread_id).await.is_err()); + + let report = manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await; + assert_eq!(report.completed, vec![thread.thread_id]); + assert!(report.submit_failed.is_empty()); + assert!(report.timed_out.is_empty()); + assert!(manager.list_thread_ids().await.is_empty()); +} + #[tokio::test] async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { let temp_dir = tempdir().expect("tempdir"); @@ -357,9 +403,10 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { let default_cwd = config.cwd.clone(); let source = manager - .start_thread_with_tools_and_service_name(StartThreadWithToolsOptions { + .start_thread_with_options(StartThreadOptions { config: config.clone(), initial_history: InitialHistory::New, + session_source: None, dynamic_tools: Vec::new(), persist_extended_history: false, metrics_service_name: None, diff --git a/codex-rs/core/src/tools/handlers/request_user_input.rs b/codex-rs/core/src/tools/handlers/request_user_input.rs index 69d222aa8..2c94145d3 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input.rs @@ -5,7 +5,6 @@ use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; -use codex_protocol::protocol::SessionSource; use codex_protocol::request_user_input::RequestUserInputArgs; use codex_tools::REQUEST_USER_INPUT_TOOL_NAME; use codex_tools::normalize_request_user_input_args; @@ -40,7 +39,7 @@ impl ToolHandler for RequestUserInputHandler { } }; - if matches!(turn.session_source, SessionSource::SubAgent(_)) { + if turn.session_source.is_non_root_agent() { return Err(FunctionCallError::RespondToModel( "request_user_input can only be used by the root thread".to_string(), )); diff --git a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs index 3c21d6670..811e4a668 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs @@ -4,6 +4,7 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::turn_diff_tracker::TurnDiffTracker; use codex_protocol::ThreadId; +use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use pretty_assertions::assert_eq; use serde_json::json; diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index b4adc6559..bcf4e3d55 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -54,7 +54,6 @@ mod items; mod json_result; mod live_cli; mod live_reload; -mod memories; mod model_overrides; mod model_switching; mod model_visible_layout; diff --git a/codex-rs/memories/write/Cargo.toml b/codex-rs/memories/write/Cargo.toml index 15deafd7e..86ed83129 100644 --- a/codex-rs/memories/write/Cargo.toml +++ b/codex-rs/memories/write/Cargo.toml @@ -14,18 +14,32 @@ workspace = true [dependencies] anyhow = { workspace = true } chrono = { workspace = true } +codex-core = { workspace = true } +codex-config = { workspace = true } +codex-features = { workspace = true } codex-git-utils = { workspace = true } +codex-login = { workspace = true } +codex-otel = { workspace = true } codex-protocol = { workspace = true } +codex-rollout = { workspace = true } +codex-rollout-trace = { workspace = true } +codex-secrets = { workspace = true } codex-state = { workspace = true } +codex-terminal-detection = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-output-truncation = { workspace = true } codex-utils-template = { workspace = true } -tokio = { workspace = true, features = ["fs"] } +futures = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["fs", "rt", "sync", "time"] } tracing = { workspace = true, features = ["log"] } uuid = { workspace = true, features = ["v4", "v5"] } [dev-dependencies] codex-models-manager = { workspace = true } +core_test_support = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["fs", "macros"] } +wiremock = { workspace = true } diff --git a/codex-rs/memories/write/src/control.rs b/codex-rs/memories/write/src/control.rs index 4f09d3e74..bc5d2dc90 100644 --- a/codex-rs/memories/write/src/control.rs +++ b/codex-rs/memories/write/src/control.rs @@ -42,3 +42,75 @@ pub(crate) async fn clear_memory_root_contents(memory_root: &Path) -> std::io::R Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[tokio::test] + async fn clear_memory_root_contents_preserves_root_directory() { + let dir = tempdir().expect("tempdir"); + let root = dir.path().join("memories"); + let nested_dir = root.join("rollout_summaries"); + tokio::fs::create_dir_all(&nested_dir) + .await + .expect("create rollout summaries dir"); + tokio::fs::write(root.join("MEMORY.md"), "stale memory index\n") + .await + .expect("write memory index"); + tokio::fs::write(nested_dir.join("rollout.md"), "stale rollout\n") + .await + .expect("write rollout summary"); + + clear_memory_root_contents(&root) + .await + .expect("clear memory root contents"); + + assert!( + tokio::fs::try_exists(&root) + .await + .expect("check memory root existence"), + "memory root should still exist after clearing contents" + ); + let mut entries = tokio::fs::read_dir(&root) + .await + .expect("read memory root after clear"); + assert!( + entries + .next_entry() + .await + .expect("read next entry") + .is_none(), + "memory root should be empty after clearing contents" + ); + } + + #[cfg(unix)] + #[tokio::test] + async fn clear_memory_root_contents_rejects_symlinked_root() { + let dir = tempdir().expect("tempdir"); + let target = dir.path().join("outside"); + tokio::fs::create_dir_all(&target) + .await + .expect("create symlink target dir"); + let target_file = target.join("keep.txt"); + tokio::fs::write(&target_file, "keep\n") + .await + .expect("write target file"); + + let root = dir.path().join("memories"); + std::os::unix::fs::symlink(&target, &root).expect("create memory root symlink"); + + let err = clear_memory_root_contents(&root) + .await + .expect_err("symlinked memory root should be rejected"); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); + assert!( + tokio::fs::try_exists(&target_file) + .await + .expect("check target file existence"), + "rejecting a symlinked memory root should not delete the symlink target" + ); + } +} diff --git a/codex-rs/memories/write/src/lib.rs b/codex-rs/memories/write/src/lib.rs index f15f58fee..be7966303 100644 --- a/codex-rs/memories/write/src/lib.rs +++ b/codex-rs/memories/write/src/lib.rs @@ -1,12 +1,16 @@ -//! Write-path helpers for Codex memories. +//! Write-path implementation for Codex memories. //! -//! This crate owns the file-backed memory artifact helpers, Phase 1 and Phase -//! 2 prompt rendering, extension pruning, and workspace diffing. Runtime -//! orchestration for Phase 1 and Phase 2 remains in `codex-core`. +//! This crate owns the startup memory pipeline, file-backed memory artifact +//! helpers, Phase 1 and Phase 2 prompt rendering, extension pruning, and +//! workspace diffing. mod control; mod extensions; +mod phase1; +mod phase2; mod prompts; +mod runtime; +mod start; mod storage; pub mod workspace; @@ -18,6 +22,7 @@ pub use control::clear_memory_roots_contents; pub use extensions::prune_old_extension_resources; pub use prompts::build_consolidation_prompt; pub use prompts::build_stage_one_input_message; +pub use start::start_memories_startup_task; pub use storage::rebuild_raw_memories_file_from_memories; pub use storage::rollout_summary_file_stem; pub use storage::sync_rollout_summaries_from_memories; @@ -36,6 +41,9 @@ pub const DEFAULT_STAGE_ONE_ROLLOUT_TOKEN_LIMIT: usize = 150_000; /// and model output. pub const STAGE_ONE_CONTEXT_WINDOW_PERCENT: i64 = 70; +#[cfg(test)] +mod startup_tests; + mod artifacts { pub(super) const EXTENSIONS_SUBDIR: &str = "extensions"; pub(super) const ROLLOUT_SUMMARIES_SUBDIR: &str = "rollout_summaries"; diff --git a/codex-rs/memories/write/src/phase1.rs b/codex-rs/memories/write/src/phase1.rs new file mode 100644 index 000000000..120ff870b --- /dev/null +++ b/codex-rs/memories/write/src/phase1.rs @@ -0,0 +1,803 @@ +use crate::STAGE_ONE_PROMPT; +use crate::build_stage_one_input_message; +use crate::runtime::MemoryStartupContext; +use crate::runtime::StageOneRequestContext; +use codex_config::types::MemoriesConfig; +use codex_core::Prompt; +use codex_core::RolloutRecorder; +use codex_core::config::Config; +use codex_protocol::error::CodexErr; +use codex_protocol::models::BaseInstructions; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::openai_models::ReasoningEffort; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::TokenUsage; +use codex_rollout::INTERACTIVE_SESSION_SOURCES; +use codex_rollout::should_persist_response_item_for_memories; +use codex_secrets::redact_secrets; +use futures::StreamExt; +use serde::Deserialize; +use serde_json::Value; +use serde_json::json; +use std::path::Path; +use std::sync::Arc; +use tracing::info; +use tracing::warn; + +const MODEL: &str = "gpt-5.4-mini"; +const REASONING_EFFORT: ReasoningEffort = ReasoningEffort::Low; +const CONCURRENCY_LIMIT: usize = 8; +const JOB_LEASE_SECONDS: i64 = 3_600; +const JOB_RETRY_DELAY_SECONDS: i64 = 3_600; +const THREAD_SCAN_LIMIT: usize = 5_000; +const PRUNE_BATCH_SIZE: usize = 200; +const MEMORY_PHASE_ONE_JOBS: &str = "codex.memory.phase1"; +const MEMORY_PHASE_ONE_E2E_MS: &str = "codex.memory.phase1.e2e_ms"; +const MEMORY_PHASE_ONE_OUTPUT: &str = "codex.memory.phase1.output"; +const MEMORY_PHASE_ONE_TOKEN_USAGE: &str = "codex.memory.phase1.token_usage"; + +struct JobResult { + outcome: JobOutcome, + token_usage: Option, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum JobOutcome { + SucceededWithOutput, + SucceededNoOutput, + Failed, +} + +struct Stats { + claimed: usize, + succeeded_with_output: usize, + succeeded_no_output: usize, + failed: usize, + total_token_usage: Option, +} + +/// Phase 1 model output payload. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +struct StageOneOutput { + /// Detailed markdown raw memory for a single rollout. + #[serde(rename = "raw_memory")] + pub(crate) raw_memory: String, + /// Compact summary line used for routing and indexing. + #[serde(rename = "rollout_summary")] + pub(crate) rollout_summary: String, + /// Optional slug used to derive rollout summary artifact filenames. + #[serde(default, rename = "rollout_slug")] + pub(crate) rollout_slug: Option, +} + +/// Runs memory phase 1 in strict step order: +/// 1) claim eligible rollout jobs +/// 2) build one stage-1 request context +/// 3) run stage-1 extraction jobs in parallel +/// 4) emit metrics and logs +pub async fn run(context: Arc, config: Arc) { + let stage_one_context = build_request_context(context.as_ref(), config.as_ref()).await; + let _phase_one_e2e_timer = stage_one_context.start_timer(MEMORY_PHASE_ONE_E2E_MS); + + // 1. Claim startup job. + let Some(claimed_candidates) = claim_startup_jobs(context.as_ref(), &config.memories).await + else { + return; + }; + if claimed_candidates.is_empty() { + stage_one_context.counter( + MEMORY_PHASE_ONE_JOBS, + /*inc*/ 1, + &[("status", "skipped_no_candidates")], + ); + return; + } + + // 3. Run the parallel sampling. + let outcomes = run_jobs( + context, + config, + claimed_candidates, + stage_one_context.clone(), + ) + .await; + + // 4. Metrics and logs. + let counts = aggregate_stats(outcomes); + emit_metrics(&stage_one_context, &counts); + info!( + "memory stage-1 extraction complete: {} job(s) claimed, {} succeeded ({} with output, {} no output), {} failed", + counts.claimed, + counts.succeeded_with_output + counts.succeeded_no_output, + counts.succeeded_with_output, + counts.succeeded_no_output, + counts.failed + ); +} + +/// Prune old un-used "dead" raw memories. +pub async fn prune(context: &MemoryStartupContext, config: &Config) { + if let Some(db) = context.state_db() { + let max_unused_days = config.memories.max_unused_days; + match db + .prune_stage1_outputs_for_retention(max_unused_days, PRUNE_BATCH_SIZE) + .await + { + Ok(pruned) => { + if pruned > 0 { + info!( + "memory startup pruned {pruned} stale stage-1 output row(s) older than {max_unused_days} days" + ); + } + } + Err(err) => { + warn!( + "state db prune_stage1_outputs_for_retention failed during memories startup: {err}" + ); + } + } + } +} + +/// JSON schema used to constrain phase-1 model output. +pub fn output_schema() -> Value { + json!({ + "type": "object", + "properties": { + "rollout_summary": { "type": "string" }, + "rollout_slug": { "type": ["string", "null"] }, + "raw_memory": { "type": "string" } + }, + "required": ["rollout_summary", "rollout_slug", "raw_memory"], + "additionalProperties": false + }) +} + +async fn claim_startup_jobs( + context: &MemoryStartupContext, + memories_config: &MemoriesConfig, +) -> Option> { + let Some(state_db) = context.state_db() else { + // This should not happen. + warn!("state db unavailable while claiming phase-1 startup jobs; skipping"); + return None; + }; + + let allowed_sources = INTERACTIVE_SESSION_SOURCES + .iter() + .map(ToString::to_string) + .collect::>(); + + match state_db + .claim_stage1_jobs_for_startup( + context.thread_id(), + codex_state::Stage1StartupClaimParams { + scan_limit: THREAD_SCAN_LIMIT, + max_claimed: memories_config.max_rollouts_per_startup, + max_age_days: memories_config.max_rollout_age_days, + min_rollout_idle_hours: memories_config.min_rollout_idle_hours, + allowed_sources: allowed_sources.as_slice(), + lease_seconds: JOB_LEASE_SECONDS, + }, + ) + .await + { + Ok(claims) => Some(claims), + Err(err) => { + warn!("state db claim_stage1_jobs_for_startup failed during memories startup: {err}"); + None + } + } +} + +async fn build_request_context( + context: &MemoryStartupContext, + config: &Config, +) -> StageOneRequestContext { + let model_name = config + .memories + .extract_model + .clone() + .unwrap_or(MODEL.to_string()); + context + .stage_one_request_context(config, &model_name, REASONING_EFFORT) + .await +} + +async fn run_jobs( + context: Arc, + config: Arc, + claimed_candidates: Vec, + stage_one_context: StageOneRequestContext, +) -> Vec { + futures::stream::iter(claimed_candidates.into_iter()) + .map(|claim| { + let context = Arc::clone(&context); + let config = Arc::clone(&config); + let stage_one_context = stage_one_context.clone(); + async move { + job::run(context.as_ref(), config.as_ref(), claim, &stage_one_context).await + } + }) + .buffer_unordered(CONCURRENCY_LIMIT) + .collect::>() + .await +} + +mod job { + use super::*; + + pub(crate) async fn run( + context: &MemoryStartupContext, + config: &Config, + claim: codex_state::Stage1JobClaim, + stage_one_context: &StageOneRequestContext, + ) -> JobResult { + let claimed_thread = claim.thread; + let (stage_one_output, token_usage) = match sample( + context, + config, + &claimed_thread.rollout_path, + &claimed_thread.cwd, + stage_one_context, + ) + .await + { + Ok(output) => output, + Err(reason) => { + result::failed( + context, + claimed_thread.id, + &claim.ownership_token, + &reason.to_string(), + ) + .await; + return JobResult { + outcome: JobOutcome::Failed, + token_usage: None, + }; + } + }; + + if stage_one_output.raw_memory.is_empty() || stage_one_output.rollout_summary.is_empty() { + return JobResult { + outcome: result::no_output(context, claimed_thread.id, &claim.ownership_token) + .await, + token_usage, + }; + } + + JobResult { + outcome: result::success( + context, + claimed_thread.id, + &claim.ownership_token, + claimed_thread.updated_at.timestamp(), + &stage_one_output.raw_memory, + &stage_one_output.rollout_summary, + stage_one_output.rollout_slug.as_deref(), + ) + .await, + token_usage, + } + } + + /// Extract the rollout and perform the actual sampling. + async fn sample( + context: &MemoryStartupContext, + config: &Config, + rollout_path: &Path, + rollout_cwd: &Path, + stage_one_context: &StageOneRequestContext, + ) -> anyhow::Result<(StageOneOutput, Option)> { + let (rollout_items, _, _) = RolloutRecorder::load_rollout_items(rollout_path).await?; + let rollout_contents = serialize_filtered_rollout_response_items(&rollout_items)?; + + let mut prompt = Prompt::default(); + prompt.input = vec![ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: build_stage_one_input_message( + &stage_one_context.model_info, + rollout_path, + rollout_cwd, + &rollout_contents, + )?, + }], + phase: None, + }]; + prompt.base_instructions = BaseInstructions { + text: STAGE_ONE_PROMPT.to_string(), + }; + prompt.output_schema = Some(output_schema()); + prompt.output_schema_strict = true; + + let (result, token_usage) = context + .stream_stage_one_prompt(config, &prompt, stage_one_context) + .await?; + + let mut output: StageOneOutput = serde_json::from_str(&result)?; + output.raw_memory = redact_secrets(output.raw_memory); + output.rollout_summary = redact_secrets(output.rollout_summary); + output.rollout_slug = output.rollout_slug.map(redact_secrets); + + Ok((output, token_usage)) + } + + mod result { + use super::*; + + pub(crate) async fn failed( + context: &MemoryStartupContext, + thread_id: codex_protocol::ThreadId, + ownership_token: &str, + reason: &str, + ) { + tracing::warn!("Phase 1 job failed for thread {thread_id}: {reason}"); + if let Some(state_db) = context.state_db() { + let _ = state_db + .mark_stage1_job_failed( + thread_id, + ownership_token, + reason, + JOB_RETRY_DELAY_SECONDS, + ) + .await; + } + } + + pub(crate) async fn no_output( + context: &MemoryStartupContext, + thread_id: codex_protocol::ThreadId, + ownership_token: &str, + ) -> JobOutcome { + let Some(state_db) = context.state_db() else { + return JobOutcome::Failed; + }; + + if state_db + .mark_stage1_job_succeeded_no_output(thread_id, ownership_token) + .await + .unwrap_or(false) + { + JobOutcome::SucceededNoOutput + } else { + JobOutcome::Failed + } + } + + pub(crate) async fn success( + context: &MemoryStartupContext, + thread_id: codex_protocol::ThreadId, + ownership_token: &str, + source_updated_at: i64, + raw_memory: &str, + rollout_summary: &str, + rollout_slug: Option<&str>, + ) -> JobOutcome { + let Some(state_db) = context.state_db() else { + return JobOutcome::Failed; + }; + + if state_db + .mark_stage1_job_succeeded( + thread_id, + ownership_token, + source_updated_at, + raw_memory, + rollout_summary, + rollout_slug, + ) + .await + .unwrap_or(false) + { + JobOutcome::SucceededWithOutput + } else { + JobOutcome::Failed + } + } + } + + /// Serializes filtered stage-1 memory items for prompt inclusion. + pub(super) fn serialize_filtered_rollout_response_items( + items: &[RolloutItem], + ) -> codex_protocol::error::Result { + let filtered = items + .iter() + .filter_map(|item| { + if let RolloutItem::ResponseItem(item) = item { + sanitize_response_item_for_memories(item) + } else { + None + } + }) + .collect::>(); + let serialized = serde_json::to_string(&filtered).map_err(|err| { + CodexErr::InvalidRequest(format!("failed to serialize rollout memory: {err}")) + })?; + Ok(redact_secrets(serialized)) + } + + fn sanitize_response_item_for_memories(item: &ResponseItem) -> Option { + let ResponseItem::Message { + id, + role, + content, + phase, + } = item + else { + return should_persist_response_item_for_memories(item).then(|| item.clone()); + }; + + if role == "developer" { + return None; + } + + if role != "user" { + return Some(item.clone()); + } + + let content = content + .iter() + .filter(|content_item| !is_memory_excluded_contextual_user_fragment(content_item)) + .cloned() + .collect::>(); + if content.is_empty() { + return None; + } + + Some(ResponseItem::Message { + id: id.clone(), + role: role.clone(), + content, + phase: phase.clone(), + }) + } + + fn is_memory_excluded_contextual_user_fragment(content_item: &ContentItem) -> bool { + let ContentItem::InputText { text } = content_item else { + return false; + }; + + matches_marked_fragment(text, "# AGENTS.md instructions for ", "") + || matches_marked_fragment(text, "", "") + } + + fn matches_marked_fragment(text: &str, start_marker: &str, end_marker: &str) -> bool { + let trimmed = text.trim_start(); + let starts_with_marker = trimmed + .get(..start_marker.len()) + .is_some_and(|candidate| candidate.eq_ignore_ascii_case(start_marker)); + let trimmed = trimmed.trim_end(); + let ends_with_marker = trimmed + .get(trimmed.len().saturating_sub(end_marker.len())..) + .is_some_and(|candidate| candidate.eq_ignore_ascii_case(end_marker)); + starts_with_marker && ends_with_marker + } + + #[cfg(test)] + mod tests { + use super::*; + + #[test] + fn classifies_memory_excluded_fragments() { + let cases = [ + ( + "# AGENTS.md instructions for /tmp\n\n\nbody\n", + true, + ), + ( + "\ndemo\nskills/demo/SKILL.md\nbody\n", + true, + ), + ( + "\n/tmp\n", + false, + ), + ( + "{\"agent_id\":\"a\",\"status\":\"completed\"}", + false, + ), + ]; + + for (text, expected) in cases { + assert_eq!( + is_memory_excluded_contextual_user_fragment(&ContentItem::InputText { + text: text.to_string(), + }), + expected, + "{text}", + ); + } + } + + #[test] + fn output_schema_requires_rollout_slug_and_keeps_it_nullable() { + let schema = output_schema(); + let properties = schema + .get("properties") + .and_then(Value::as_object) + .expect("properties object"); + let required = schema + .get("required") + .and_then(Value::as_array) + .expect("required array"); + + let mut required_keys = required + .iter() + .map(|key| key.as_str().expect("required key string")) + .collect::>(); + required_keys.sort_unstable(); + + assert!( + properties.contains_key("rollout_slug"), + "schema should declare rollout_slug" + ); + + let rollout_slug_type = properties + .get("rollout_slug") + .and_then(Value::as_object) + .and_then(|entry| entry.get("type")) + .and_then(Value::as_array) + .expect("rollout_slug type array"); + let mut rollout_slug_types = rollout_slug_type + .iter() + .map(|entry| entry.as_str().expect("type entry string")) + .collect::>(); + rollout_slug_types.sort_unstable(); + + assert_eq!( + required_keys, + vec!["raw_memory", "rollout_slug", "rollout_summary"] + ); + assert_eq!(rollout_slug_types, vec!["null", "string"]); + } + } +} + +fn aggregate_stats(outcomes: Vec) -> Stats { + let claimed = outcomes.len(); + let mut succeeded_with_output = 0; + let mut succeeded_no_output = 0; + let mut failed = 0; + let mut total_token_usage = TokenUsage::default(); + let mut has_token_usage = false; + + for outcome in outcomes { + match outcome.outcome { + JobOutcome::SucceededWithOutput => succeeded_with_output += 1, + JobOutcome::SucceededNoOutput => succeeded_no_output += 1, + JobOutcome::Failed => failed += 1, + } + + if let Some(token_usage) = outcome.token_usage { + total_token_usage.add_assign(&token_usage); + has_token_usage = true; + } + } + + Stats { + claimed, + succeeded_with_output, + succeeded_no_output, + failed, + total_token_usage: has_token_usage.then_some(total_token_usage), + } +} + +fn emit_metrics(context: &StageOneRequestContext, counts: &Stats) { + if counts.claimed > 0 { + context.counter( + MEMORY_PHASE_ONE_JOBS, + counts.claimed as i64, + &[("status", "claimed")], + ); + } + if counts.succeeded_with_output > 0 { + context.counter( + MEMORY_PHASE_ONE_JOBS, + counts.succeeded_with_output as i64, + &[("status", "succeeded")], + ); + context.counter( + MEMORY_PHASE_ONE_OUTPUT, + counts.succeeded_with_output as i64, + &[], + ); + } + if counts.succeeded_no_output > 0 { + context.counter( + MEMORY_PHASE_ONE_JOBS, + counts.succeeded_no_output as i64, + &[("status", "succeeded_no_output")], + ); + } + if counts.failed > 0 { + context.counter( + MEMORY_PHASE_ONE_JOBS, + counts.failed as i64, + &[("status", "failed")], + ); + } + if let Some(token_usage) = counts.total_token_usage.as_ref() { + context.histogram( + MEMORY_PHASE_ONE_TOKEN_USAGE, + token_usage.total_tokens.max(0), + &[("token_type", "total")], + ); + context.histogram( + MEMORY_PHASE_ONE_TOKEN_USAGE, + token_usage.input_tokens.max(0), + &[("token_type", "input")], + ); + context.histogram( + MEMORY_PHASE_ONE_TOKEN_USAGE, + token_usage.cached_input(), + &[("token_type", "cached_input")], + ); + context.histogram( + MEMORY_PHASE_ONE_TOKEN_USAGE, + token_usage.output_tokens.max(0), + &[("token_type", "output")], + ); + context.histogram( + MEMORY_PHASE_ONE_TOKEN_USAGE, + token_usage.reasoning_output_tokens.max(0), + &[("token_type", "reasoning_output")], + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + #[test] + fn serializes_memory_rollout_with_agents_removed_but_environment_kept() { + let mixed_contextual_message = ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ + ContentItem::InputText { + text: + "# AGENTS.md instructions for /tmp\n\n\nbody\n" + .to_string(), + }, + ContentItem::InputText { + text: "\n/tmp\n" + .to_string(), + }, + ], + phase: None, + }; + let skill_message = ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: + "\ndemo\nskills/demo/SKILL.md\nbody\n" + .to_string(), + }], + phase: None, + }; + let subagent_message = ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "{\"agent_id\":\"a\",\"status\":\"completed\"}" + .to_string(), + }], + phase: None, + }; + + let serialized = job::serialize_filtered_rollout_response_items(&[ + RolloutItem::ResponseItem(mixed_contextual_message), + RolloutItem::ResponseItem(skill_message), + RolloutItem::ResponseItem(subagent_message.clone()), + ]) + .expect("serialize"); + let parsed: Vec = serde_json::from_str(&serialized).expect("parse"); + + assert_eq!( + parsed, + vec![ + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "\n/tmp\n" + .to_string(), + }], + phase: None, + }, + subagent_message, + ] + ); + } + + #[test] + fn serializes_memory_rollout_redacts_secrets_before_prompt_upload() { + let serialized = + job::serialize_filtered_rollout_response_items(&[RolloutItem::ResponseItem( + ResponseItem::FunctionCallOutput { + call_id: "call_123".to_string(), + output: codex_protocol::models::FunctionCallOutputPayload { + body: codex_protocol::models::FunctionCallOutputBody::Text( + r#"{"token":"sk-abcdefghijklmnopqrstuvwxyz123456"}"#.to_string(), + ), + success: Some(true), + }, + }, + )]) + .expect("serialize"); + + assert!(!serialized.contains("sk-abcdefghijklmnopqrstuvwxyz123456")); + assert!(serialized.contains("[REDACTED_SECRET]")); + } + + #[test] + fn count_outcomes_sums_token_usage_across_all_jobs() { + let counts = aggregate_stats(vec![ + JobResult { + outcome: JobOutcome::SucceededWithOutput, + token_usage: Some(TokenUsage { + input_tokens: 10, + cached_input_tokens: 2, + output_tokens: 3, + reasoning_output_tokens: 1, + total_tokens: 13, + }), + }, + JobResult { + outcome: JobOutcome::SucceededNoOutput, + token_usage: Some(TokenUsage { + input_tokens: 7, + cached_input_tokens: 1, + output_tokens: 2, + reasoning_output_tokens: 0, + total_tokens: 9, + }), + }, + JobResult { + outcome: JobOutcome::Failed, + token_usage: None, + }, + ]); + + assert_eq!(counts.claimed, 3); + assert_eq!(counts.succeeded_with_output, 1); + assert_eq!(counts.succeeded_no_output, 1); + assert_eq!(counts.failed, 1); + assert_eq!( + counts.total_token_usage, + Some(TokenUsage { + input_tokens: 17, + cached_input_tokens: 3, + output_tokens: 5, + reasoning_output_tokens: 1, + total_tokens: 22, + }) + ); + } + + #[test] + fn count_outcomes_keeps_usage_empty_when_no_job_reports_it() { + let counts = aggregate_stats(vec![ + JobResult { + outcome: JobOutcome::SucceededWithOutput, + token_usage: None, + }, + JobResult { + outcome: JobOutcome::Failed, + token_usage: None, + }, + ]); + + assert_eq!(counts.claimed, 2); + assert_eq!(counts.total_token_usage, None); + } +} diff --git a/codex-rs/core/src/memories/phase2.rs b/codex-rs/memories/write/src/phase2.rs similarity index 57% rename from codex-rs/core/src/memories/phase2.rs rename to codex-rs/memories/write/src/phase2.rs index bcad73831..7278207f3 100644 --- a/codex-rs/core/src/memories/phase2.rs +++ b/codex-rs/memories/write/src/phase2.rs @@ -1,26 +1,21 @@ -use crate::agent::AgentStatus; -use crate::agent::status::is_final as is_final_agent_status; -use crate::config::Config; -use crate::memories::metrics; -use crate::memories::phase_two; -use crate::session::emit_subagent_session_started; -use crate::session::session::Session; +use crate::build_consolidation_prompt; +use crate::memory_root; +use crate::prune_old_extension_resources; +use crate::rebuild_raw_memories_file_from_memories; +use crate::runtime::MemoryStartupContext; +use crate::runtime::SpawnedConsolidationAgent; +use crate::sync_rollout_summaries_from_memories; +use crate::workspace::memory_workspace_diff; +use crate::workspace::prepare_memory_workspace; +use crate::workspace::reset_memory_workspace_baseline; +use crate::workspace::write_workspace_diff; use codex_config::Constrained; +use codex_core::config::Config; use codex_features::Feature; -use codex_memories_write::build_consolidation_prompt; -use codex_memories_write::memory_root; -use codex_memories_write::prune_old_extension_resources; -use codex_memories_write::rebuild_raw_memories_file_from_memories; -use codex_memories_write::sync_rollout_summaries_from_memories; -use codex_memories_write::workspace::memory_workspace_diff; -use codex_memories_write::workspace::prepare_memory_workspace; -use codex_memories_write::workspace::reset_memory_workspace_baseline; -use codex_memories_write::workspace::write_workspace_diff; use codex_protocol::ThreadId; +use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::SandboxPolicy; -use codex_protocol::protocol::SessionSource; -use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TokenUsage; use codex_protocol::user_input::UserInput; use codex_state::Stage1Output; @@ -29,8 +24,16 @@ use std::collections::HashMap; use std::path::Path; use std::sync::Arc; use std::time::Duration; -use tokio::sync::watch; -use tracing::warn; +const MODEL: &str = "gpt-5.4"; +const REASONING_EFFORT: codex_protocol::openai_models::ReasoningEffort = + codex_protocol::openai_models::ReasoningEffort::Medium; +const JOB_LEASE_SECONDS: i64 = 3_600; +const JOB_RETRY_DELAY_SECONDS: i64 = 3_600; +const JOB_HEARTBEAT_SECONDS: u64 = 90; +const MEMORY_PHASE_TWO_JOBS: &str = "codex.memory.phase2"; +const MEMORY_PHASE_TWO_E2E_MS: &str = "codex.memory.phase2.e2e_ms"; +const MEMORY_PHASE_TWO_INPUT: &str = "codex.memory.phase2.input"; +const MEMORY_PHASE_TWO_TOKEN_USAGE: &str = "codex.memory.phase2.token_usage"; #[derive(Debug, Clone, Default)] struct Claim { @@ -45,14 +48,10 @@ struct Counters { /// Runs memory phase 2 (aka consolidation) in strict order. The method represents the linear /// flow of the consolidation phase. -pub(super) async fn run(session: &Arc, config: Arc) { - let phase_two_e2e_timer = session - .services - .session_telemetry - .start_timer(metrics::MEMORY_PHASE_TWO_E2E_MS, &[]) - .ok(); +pub async fn run(context: Arc, config: Arc) { + let phase_two_e2e_timer = context.start_timer(MEMORY_PHASE_TWO_E2E_MS); - let Some(db) = session.services.state_db.as_deref() else { + let Some(db) = context.state_db() else { // This should not happen. return; }; @@ -61,14 +60,10 @@ pub(super) async fn run(session: &Arc, config: Arc) { let max_unused_days = config.memories.max_unused_days; // 1. Claim the global Phase 2 lock before touching the memory workspace. - let claim = match job::claim(session, db).await { + let claim = match job::claim(context.as_ref(), db.as_ref()).await { Ok(claim) => claim, Err(e) => { - session.services.session_telemetry.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - /*inc*/ 1, - &[("status", e)], - ); + context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", e)]); return; } }; @@ -76,7 +71,13 @@ pub(super) async fn run(session: &Arc, config: Arc) { // 2. Ensure the memories root has a git baseline repository. if let Err(err) = prepare_memory_workspace(&root).await { tracing::error!("failed preparing memory workspace: {err}"); - job::failed(session, db, &claim, "failed_prepare_workspace").await; + job::failed( + context.as_ref(), + db.as_ref(), + &claim, + "failed_prepare_workspace", + ) + .await; return; } @@ -84,7 +85,13 @@ pub(super) async fn run(session: &Arc, config: Arc) { let Some(agent_config) = agent::get_config(config.as_ref()) else { // If we can't get the config, we can't consolidate. tracing::error!("failed to get agent config"); - job::failed(session, db, &claim, "failed_sandbox_policy").await; + job::failed( + context.as_ref(), + db.as_ref(), + &claim, + "failed_sandbox_policy", + ) + .await; return; }; @@ -96,7 +103,13 @@ pub(super) async fn run(session: &Arc, config: Arc) { Ok(raw_memories) => raw_memories, Err(err) => { tracing::error!("failed to list stage1 outputs from global: {err}"); - job::failed(session, db, &claim, "failed_load_stage1_outputs").await; + job::failed( + context.as_ref(), + db.as_ref(), + &claim, + "failed_load_stage1_outputs", + ) + .await; return; } }; @@ -106,7 +119,13 @@ pub(super) async fn run(session: &Arc, config: Arc) { // 5. Sync the current inputs into the memory workspace. if let Err(err) = sync_phase2_workspace_inputs(&root, &raw_memories).await { tracing::error!("failed syncing phase2 workspace inputs: {err}"); - job::failed(session, db, &claim, "failed_sync_workspace_inputs").await; + job::failed( + context.as_ref(), + db.as_ref(), + &claim, + "failed_sync_workspace_inputs", + ) + .await; return; } @@ -115,7 +134,13 @@ pub(super) async fn run(session: &Arc, config: Arc) { Ok(diff) => diff, Err(err) => { tracing::error!("failed checking memory workspace changes: {err}"); - job::failed(session, db, &claim, "failed_workspace_status").await; + job::failed( + context.as_ref(), + db.as_ref(), + &claim, + "failed_workspace_status", + ) + .await; return; } }; @@ -123,8 +148,8 @@ pub(super) async fn run(session: &Arc, config: Arc) { tracing::error!("Phase 2 no changes"); // We check only after sync of the file system. job::succeed( - session, - db, + context.as_ref(), + db.as_ref(), &claim, new_watermark, &raw_memories, @@ -137,54 +162,38 @@ pub(super) async fn run(session: &Arc, config: Arc) { // 7. Persist the diff for the consolidation agent to inspect. if let Err(err) = write_workspace_diff(&root, &workspace_diff).await { tracing::error!("failed writing memory workspace diff file: {err}"); - job::failed(session, db, &claim, "failed_workspace_diff_file").await; + job::failed( + context.as_ref(), + db.as_ref(), + &claim, + "failed_workspace_diff_file", + ) + .await; return; } // 8. Spawn the consolidation agent. let prompt = agent::get_prompt(&root); - let source = SessionSource::SubAgent(SubAgentSource::MemoryConsolidation); - let agent_control = session.services.agent_control.detached_registry(); - let thread_id = match agent_control - .spawn_agent(agent_config, prompt.into(), Some(source)) + let agent = match context + .spawn_consolidation_agent(agent_config, prompt) .await { - Ok(thread_id) => thread_id, + Ok(agent) => agent, Err(err) => { tracing::error!("failed to spawn global memory consolidation agent: {err}"); - job::failed(session, db, &claim, "failed_spawn_agent").await; + job::failed(context.as_ref(), db.as_ref(), &claim, "failed_spawn_agent").await; return; } }; - if let Some(thread_config) = session - .services - .agent_control - .get_agent_config_snapshot(thread_id) - .await - { - let client_metadata = session.app_server_client_metadata().await; - emit_subagent_session_started( - &session.services.analytics_events_client, - client_metadata, - thread_id, - /*parent_thread_id*/ None, - thread_config, - SubAgentSource::MemoryConsolidation, - ); - } else { - warn!("failed to load memory consolidation thread config for analytics: {thread_id}"); - } - // 9. Hand off completion handling, heartbeats, and baseline reset. agent::handle( - session, + Arc::clone(&context), claim, new_watermark, raw_memories.clone(), root, - thread_id, - agent_control, + agent, phase_two_e2e_timer, ); @@ -192,7 +201,7 @@ pub(super) async fn run(session: &Arc, config: Arc) { let counters = Counters { input: raw_memory_count as i64, }; - emit_metrics(session, counters); + emit_metrics(context.as_ref(), counters); } async fn sync_phase2_workspace_inputs( @@ -210,15 +219,14 @@ mod job { use super::*; pub(super) async fn claim( - session: &Arc, + context: &MemoryStartupContext, db: &StateRuntime, ) -> Result { - let session_telemetry = &session.services.session_telemetry; let claim = db - .try_claim_global_phase2_job(session.conversation_id, phase_two::JOB_LEASE_SECONDS) + .try_claim_global_phase2_job(context.thread_id(), JOB_LEASE_SECONDS) .await .map_err(|e| { - tracing::error!("failed to claim job: {}", e); + tracing::error!("failed to claim job: {e}"); "failed_claim" })?; let (token, watermark) = match claim { @@ -226,8 +234,8 @@ mod job { ownership_token, input_watermark, } => { - session_telemetry.counter( - metrics::MEMORY_PHASE_TWO_JOBS, + context.counter( + MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", "claimed")], ); @@ -243,48 +251,36 @@ mod job { } pub(super) async fn failed( - session: &Arc, + context: &MemoryStartupContext, db: &StateRuntime, claim: &Claim, reason: &'static str, ) { - session.services.session_telemetry.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - /*inc*/ 1, - &[("status", reason)], - ); + context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", reason)]); if matches!( - db.mark_global_phase2_job_failed( - &claim.token, - reason, - phase_two::JOB_RETRY_DELAY_SECONDS, - ) - .await, + db.mark_global_phase2_job_failed(&claim.token, reason, JOB_RETRY_DELAY_SECONDS,) + .await, Ok(false) ) { let _ = db .mark_global_phase2_job_failed_if_unowned( &claim.token, reason, - phase_two::JOB_RETRY_DELAY_SECONDS, + JOB_RETRY_DELAY_SECONDS, ) .await; } } pub(super) async fn succeed( - session: &Arc, + context: &MemoryStartupContext, db: &StateRuntime, claim: &Claim, completion_watermark: i64, selected_outputs: &[codex_state::Stage1Output], reason: &'static str, ) -> bool { - session.services.session_telemetry.counter( - metrics::MEMORY_PHASE_TWO_JOBS, - /*inc*/ 1, - &[("status", reason)], - ); + context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", reason)]); db.mark_global_phase2_job_succeeded(&claim.token, completion_watermark, selected_outputs) .await .unwrap_or(false) @@ -293,6 +289,7 @@ mod job { mod agent { use super::*; + use tracing::warn; pub(super) fn get_config(config: &Config) -> Option { let root = memory_root(&config.codex_home); @@ -307,7 +304,7 @@ mod agent { agent_config.mcp_servers = Constrained::allow_only(HashMap::new()); // Approval policy agent_config.permissions.approval_policy = Constrained::allow_only(AskForApproval::Never); - // Consolidation runs as an internal sub-agent and must not recursively delegate. + // Consolidation runs as an internal worker and must not recursively delegate. let _ = agent_config.features.disable(Feature::SpawnCsv); let _ = agent_config.features.disable(Feature::Collab); let _ = agent_config.features.disable(Feature::MemoryTool); @@ -336,9 +333,9 @@ mod agent { .memories .consolidation_model .clone() - .unwrap_or(phase_two::MODEL.to_string()), + .unwrap_or(MODEL.to_string()), ); - agent_config.model_reasoning_effort = Some(phase_two::REASONING_EFFORT); + agent_config.model_reasoning_effort = Some(REASONING_EFFORT); Some(agent_config) } @@ -354,95 +351,90 @@ mod agent { /// Handle the agent while it is running. #[allow(clippy::too_many_arguments)] pub(super) fn handle( - session: &Arc, + context: Arc, claim: Claim, new_watermark: i64, selected_outputs: Vec, memory_root: codex_utils_absolute_path::AbsolutePathBuf, - thread_id: ThreadId, - agent_control: crate::agent::AgentControl, + agent: SpawnedConsolidationAgent, phase_two_e2e_timer: Option, ) { - let Some(db) = session.services.state_db.clone() else { + let Some(db) = context.state_db() else { return; }; - let session = session.clone(); tokio::spawn(async move { let _phase_two_e2e_timer = phase_two_e2e_timer; - - // TODO(jif) we might have a very small race here. - let rx = match agent_control.subscribe_status(thread_id).await { - Ok(rx) => rx, - Err(err) => { - tracing::error!("agent_control.subscribe_status failed: {err:?}"); - job::failed(&session, &db, &claim, "failed_subscribe_status").await; - return; - } - }; + let SpawnedConsolidationAgent { thread_id, thread } = agent; // Loop the agent until we have the final status. - let final_status = loop_agent(db.clone(), claim.token.clone(), thread_id, rx).await; + let final_status = + loop_agent(db.clone(), claim.token.clone(), thread_id, &thread).await; if matches!(final_status, AgentStatus::Completed(_)) { - if let Some(token_usage) = agent_control.get_total_token_usage(thread_id).await { - emit_token_usage_metrics(&session, &token_usage); + if let Some(token_usage) = thread + .token_usage_info() + .await + .map(|info| info.total_token_usage) + { + emit_token_usage_metrics(context.as_ref(), &token_usage); } // Do not reset the workspace baseline if we lost the lock. - let Ok(still_owns_lock) = db - .heartbeat_global_phase2_job(&claim.token, phase_two::JOB_LEASE_SECONDS) + let still_owns_lock = match db + .heartbeat_global_phase2_job(&claim.token, JOB_LEASE_SECONDS) .await .inspect_err(|err| { tracing::error!( "failed confirming global memory consolidation ownership before resetting workspace baseline: {err}" ); - }) - else { - job::failed(&session, &db, &claim, "failed_confirm_ownership").await; - return; + }) { + Ok(true) => true, + Ok(false) => { + tracing::error!( + "lost global memory consolidation ownership before resetting workspace baseline" + ); + false + } + Err(_) => { + job::failed(context.as_ref(), &db, &claim, "failed_confirm_ownership") + .await; + false + } }; - if !still_owns_lock { - tracing::error!( - "lost global memory consolidation ownership before resetting workspace baseline" - ); - return; - } - - if let Err(err) = reset_memory_workspace_baseline(&memory_root).await { - tracing::error!("failed resetting memory workspace baseline: {err}"); - job::failed(&session, &db, &claim, "failed_workspace_commit").await; - return; - } - if !job::succeed( - &session, - &db, - &claim, - new_watermark, - &selected_outputs, - "succeeded", - ) - .await - { - tracing::error!( - "failed marking global memory consolidation job succeeded after resetting workspace baseline" - ); - } - } else { - job::failed(&session, &db, &claim, "failed_agent").await; - } - - // Fire and forget close of the agent. - if !matches!(final_status, AgentStatus::Shutdown | AgentStatus::NotFound) { - tokio::spawn(async move { - if let Err(err) = agent_control.shutdown_live_agent(thread_id).await { - warn!( - "failed to auto-close global memory consolidation agent {thread_id}: {err}" + if still_owns_lock { + if let Err(err) = reset_memory_workspace_baseline(&memory_root).await { + tracing::error!("failed resetting memory workspace baseline: {err}"); + job::failed(context.as_ref(), &db, &claim, "failed_workspace_commit").await; + } else if !job::succeed( + context.as_ref(), + &db, + &claim, + new_watermark, + &selected_outputs, + "succeeded", + ) + .await + { + tracing::error!( + "failed marking global memory consolidation job succeeded after resetting workspace baseline" ); } - }); + } } else { - tracing::warn!("The agent was already gone"); + job::failed(context.as_ref(), &db, &claim, "failed_agent").await; } + + let cleanup_context = Arc::clone(&context); + tokio::spawn(async move { + if let Err(err) = cleanup_context + .shutdown_consolidation_agent(SpawnedConsolidationAgent { thread_id, thread }) + .await + { + warn!( + "failed to auto-close global memory consolidation agent {thread_id}: {err}" + ); + } + }); }); } @@ -450,42 +442,58 @@ mod agent { db: Arc, token: String, thread_id: ThreadId, - mut rx: watch::Receiver, + thread: &codex_core::CodexThread, ) -> AgentStatus { let mut heartbeat_interval = - tokio::time::interval(Duration::from_secs(phase_two::JOB_HEARTBEAT_SECONDS)); + tokio::time::interval(Duration::from_secs(JOB_HEARTBEAT_SECONDS)); heartbeat_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut status_poll_interval = tokio::time::interval(Duration::from_secs(1)); + status_poll_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let session_termination = thread.wait_until_terminated(); + tokio::pin!(session_termination); loop { - let status = rx.borrow().clone(); + let status = thread.agent_status().await; if is_final_agent_status(&status) { break status; } tokio::select! { - update = rx.changed() => { - if update.is_err() { - tracing::warn!( - "lost status updates for global memory consolidation agent {thread_id}" - ); + _ = &mut session_termination => { + let status = thread.agent_status().await; + if is_final_agent_status(&status) { break status; } + tracing::warn!( + "memory consolidation agent {thread_id} exited before final status; last status was {status:?}" + ); + break AgentStatus::Errored(format!( + "memory consolidation agent exited before final status: {status:?}" + )); + } + _ = status_poll_interval.tick() => { } _ = heartbeat_interval.tick() => { match db .heartbeat_global_phase2_job( &token, - phase_two::JOB_LEASE_SECONDS, + JOB_LEASE_SECONDS, ) .await { Ok(true) => {} Ok(false) => { + tracing::warn!( + "lost global phase-2 ownership during heartbeat for memory consolidation agent {thread_id}" + ); break AgentStatus::Errored( "lost global phase-2 ownership during heartbeat".to_string(), ); } Err(err) => { + tracing::warn!( + "phase-2 heartbeat update failed for memory consolidation agent {thread_id}: {err}" + ); break AgentStatus::Errored(format!( "phase-2 heartbeat update failed: {err}" )); @@ -509,43 +517,48 @@ pub(super) fn get_watermark( .max(claimed_watermark) } -fn emit_metrics(session: &Arc, counters: Counters) { - let otel = session.services.session_telemetry.clone(); +fn is_final_agent_status(status: &AgentStatus) -> bool { + !matches!( + status, + AgentStatus::PendingInit | AgentStatus::Running | AgentStatus::Interrupted + ) +} + +fn emit_metrics(context: &MemoryStartupContext, counters: Counters) { if counters.input > 0 { - otel.counter(metrics::MEMORY_PHASE_TWO_INPUT, counters.input, &[]); + context.counter(MEMORY_PHASE_TWO_INPUT, counters.input, &[]); } - otel.counter( - metrics::MEMORY_PHASE_TWO_JOBS, + context.counter( + MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", "agent_spawned")], ); } -fn emit_token_usage_metrics(session: &Arc, token_usage: &TokenUsage) { - let otel = session.services.session_telemetry.clone(); - otel.histogram( - metrics::MEMORY_PHASE_TWO_TOKEN_USAGE, +fn emit_token_usage_metrics(context: &MemoryStartupContext, token_usage: &TokenUsage) { + context.histogram( + MEMORY_PHASE_TWO_TOKEN_USAGE, token_usage.total_tokens.max(0), &[("token_type", "total")], ); - otel.histogram( - metrics::MEMORY_PHASE_TWO_TOKEN_USAGE, + context.histogram( + MEMORY_PHASE_TWO_TOKEN_USAGE, token_usage.input_tokens.max(0), &[("token_type", "input")], ); - otel.histogram( - metrics::MEMORY_PHASE_TWO_TOKEN_USAGE, + context.histogram( + MEMORY_PHASE_TWO_TOKEN_USAGE, token_usage.cached_input(), &[("token_type", "cached_input")], ); - otel.histogram( - metrics::MEMORY_PHASE_TWO_TOKEN_USAGE, + context.histogram( + MEMORY_PHASE_TWO_TOKEN_USAGE, token_usage.output_tokens.max(0), &[("token_type", "output")], ); - otel.histogram( - metrics::MEMORY_PHASE_TWO_TOKEN_USAGE, + context.histogram( + MEMORY_PHASE_TWO_TOKEN_USAGE, token_usage.reasoning_output_tokens.max(0), &[("token_type", "reasoning_output")], ); diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs new file mode 100644 index 000000000..737fb6787 --- /dev/null +++ b/codex-rs/memories/write/src/runtime.rs @@ -0,0 +1,293 @@ +use codex_core::CodexThread; +use codex_core::ModelClient; +use codex_core::NewThread; +use codex_core::Prompt; +use codex_core::ResponseEvent; +use codex_core::StartThreadOptions; +use codex_core::ThreadManager; +use codex_core::config::Config; +use codex_core::content_items_to_text; +use codex_core::resolve_installation_id; +use codex_features::Feature; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::auth_env_telemetry::collect_auth_env_telemetry; +use codex_login::default_client::originator; +use codex_otel::SessionTelemetry; +use codex_otel::TelemetryAuthMode; +use codex_protocol::ThreadId; +use codex_protocol::config_types::ReasoningSummary; +use codex_protocol::config_types::ServiceTier; +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::openai_models::ReasoningEffort; +use codex_protocol::protocol::InitialHistory; +use codex_protocol::protocol::InternalSessionSource; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TokenUsage; +use codex_protocol::user_input::UserInput; +use codex_rollout_trace::InferenceTraceContext; +use codex_state::StateRuntime; +use codex_terminal_detection::user_agent; +use futures::StreamExt; +use std::sync::Arc; +use std::time::Duration; + +pub(crate) struct SpawnedConsolidationAgent { + pub(crate) thread_id: ThreadId, + pub(crate) thread: Arc, +} + +#[derive(Clone, Debug)] +pub(crate) struct StageOneRequestContext { + pub(crate) model_info: ModelInfo, + pub(crate) session_telemetry: SessionTelemetry, + pub(crate) reasoning_effort: Option, + pub(crate) reasoning_summary: ReasoningSummary, + pub(crate) service_tier: Option, + pub(crate) turn_metadata_header: Option, +} + +impl StageOneRequestContext { + pub(crate) fn start_timer(&self, name: &str) -> Option { + self.session_telemetry.start_timer(name, &[]).ok() + } + + pub(crate) fn counter(&self, name: &str, inc: i64, tags: &[(&str, &str)]) { + self.session_telemetry.counter(name, inc, tags); + } + + pub(crate) fn histogram(&self, name: &str, value: i64, tags: &[(&str, &str)]) { + self.session_telemetry.histogram(name, value, tags); + } +} + +pub(crate) struct MemoryStartupContext { + thread_id: ThreadId, + thread: Arc, + thread_manager: Arc, + auth_manager: Arc, + session_telemetry: SessionTelemetry, +} + +impl MemoryStartupContext { + pub(crate) fn new( + thread_manager: Arc, + auth_manager: Arc, + thread_id: ThreadId, + thread: Arc, + config: &Config, + source: SessionSource, + ) -> Self { + let auth = auth_manager.auth_cached(); + let auth = auth.as_ref(); + let auth_mode = auth.map(CodexAuth::auth_mode).map(TelemetryAuthMode::from); + let account_id = auth.and_then(CodexAuth::get_account_id); + let account_email = auth.and_then(CodexAuth::get_account_email); + let model = config.model.as_deref().unwrap_or("unknown"); + let auth_env_telemetry = collect_auth_env_telemetry( + &config.model_provider, + auth_manager.codex_api_key_env_enabled(), + ); + let session_telemetry = SessionTelemetry::new( + thread_id, + model, + model, + account_id, + account_email, + auth_mode, + originator().value, + config.otel.log_user_prompt, + user_agent(), + source, + ) + .with_auth_env(auth_env_telemetry.to_otel_metadata()); + + Self { + thread_id, + thread, + thread_manager, + auth_manager, + session_telemetry, + } + } + + pub(crate) fn thread_id(&self) -> ThreadId { + self.thread_id + } + + pub(crate) fn state_db(&self) -> Option> { + self.thread.state_db() + } + + pub(crate) fn counter(&self, name: &str, inc: i64, tags: &[(&str, &str)]) { + self.session_telemetry.counter(name, inc, tags); + } + + pub(crate) fn histogram(&self, name: &str, value: i64, tags: &[(&str, &str)]) { + self.session_telemetry.histogram(name, value, tags); + } + + pub(crate) fn start_timer(&self, name: &str) -> Option { + self.session_telemetry.start_timer(name, &[]).ok() + } + + pub(crate) async fn stage_one_request_context( + &self, + config: &Config, + model_name: &str, + reasoning_effort: ReasoningEffort, + ) -> StageOneRequestContext { + let config_snapshot = self.thread.config_snapshot().await; + let model_info = self + .thread_manager + .get_models_manager() + .get_model_info(model_name, &config.to_models_manager_config()) + .await; + let turn_metadata_header = + codex_core::build_turn_metadata_header(&config.cwd, /*sandbox*/ None).await; + let reasoning_summary = config + .model_reasoning_summary + .unwrap_or(model_info.default_reasoning_summary); + + StageOneRequestContext { + model_info, + turn_metadata_header, + session_telemetry: self + .session_telemetry + .clone() + .with_model(model_name, model_name), + reasoning_effort: Some(reasoning_effort), + reasoning_summary, + service_tier: config_snapshot.service_tier, + } + } + + pub(crate) async fn stream_stage_one_prompt( + &self, + config: &Config, + prompt: &Prompt, + context: &StageOneRequestContext, + ) -> anyhow::Result<(String, Option)> { + let installation_id = resolve_installation_id(&config.codex_home).await?; + let session_source = self.thread.config_snapshot().await.session_source; + let model_client = ModelClient::new( + Some(Arc::clone(&self.auth_manager)), + self.thread_id, + installation_id, + config.model_provider.clone(), + session_source, + config.model_verbosity, + config.features.enabled(Feature::EnableRequestCompression), + config.features.enabled(Feature::RuntimeMetrics), + /*beta_features_header*/ None, + ); + + let mut client_session = model_client.new_session(); + let mut stream = client_session + .stream( + prompt, + &context.model_info, + &context.session_telemetry, + context.reasoning_effort, + context.reasoning_summary, + context.service_tier, + context.turn_metadata_header.as_deref(), + &InferenceTraceContext::disabled(), + ) + .await?; + + let mut result = String::new(); + let mut token_usage = None; + while let Some(message) = stream.next().await.transpose()? { + match message { + ResponseEvent::OutputTextDelta(delta) => result.push_str(&delta), + ResponseEvent::OutputItemDone(item) => { + if result.is_empty() + && let codex_protocol::models::ResponseItem::Message { content, .. } = item + && let Some(text) = content_items_to_text(&content) + { + result.push_str(&text); + } + } + ResponseEvent::Completed { + token_usage: usage, .. + } => { + token_usage = usage; + break; + } + _ => {} + } + } + + Ok((result, token_usage)) + } + + pub(crate) async fn spawn_consolidation_agent( + &self, + config: Config, + prompt: Vec, + ) -> anyhow::Result { + let environments = self + .thread_manager + .default_environment_selections(&config.cwd); + let NewThread { + thread_id, thread, .. + } = self + .thread_manager + .start_thread_with_options(StartThreadOptions { + config, + initial_history: InitialHistory::New, + session_source: Some(SessionSource::Internal( + InternalSessionSource::MemoryConsolidation, + )), + dynamic_tools: Vec::new(), + persist_extended_history: false, + metrics_service_name: None, + parent_trace: None, + environments, + }) + .await?; + + let agent = SpawnedConsolidationAgent { thread_id, thread }; + if let Err(err) = agent + .thread + .submit(Op::UserInput { + items: prompt, + environments: None, + final_output_json_schema: None, + responsesapi_client_metadata: None, + }) + .await + { + if let Err(shutdown_err) = self.shutdown_consolidation_agent(agent).await { + tracing::warn!( + "failed to shut down consolidation agent after submit error: {shutdown_err}" + ); + } + return Err(err.into()); + } + + Ok(agent) + } + + pub(crate) async fn shutdown_consolidation_agent( + &self, + agent: SpawnedConsolidationAgent, + ) -> anyhow::Result<()> { + let SpawnedConsolidationAgent { thread_id, thread } = agent; + let thread = self + .thread_manager + .remove_thread(&thread_id) + .await + .unwrap_or(thread); + + tokio::time::timeout(Duration::from_secs(10), thread.shutdown_and_wait()) + .await + .map_err(|_| { + anyhow::anyhow!("memory consolidation agent {thread_id} shutdown timed out") + })??; + + Ok(()) + } +} diff --git a/codex-rs/memories/write/src/start.rs b/codex-rs/memories/write/src/start.rs new file mode 100644 index 000000000..362ce9fa4 --- /dev/null +++ b/codex-rs/memories/write/src/start.rs @@ -0,0 +1,55 @@ +use crate::phase1; +use crate::phase2; +use crate::runtime::MemoryStartupContext; +use codex_core::CodexThread; +use codex_core::ThreadManager; +use codex_core::config::Config; +use codex_features::Feature; +use codex_login::AuthManager; +use codex_protocol::ThreadId; +use codex_protocol::protocol::SessionSource; +use std::sync::Arc; +use tracing::warn; + +/// Starts the asynchronous startup memory pipeline for an eligible root session. +/// +/// The pipeline is skipped for ephemeral sessions, disabled feature flags, and +/// subagent sessions. +pub fn start_memories_startup_task( + thread_manager: Arc, + auth_manager: Arc, + thread_id: ThreadId, + thread: Arc, + config: Arc, + source: &SessionSource, +) { + if config.ephemeral + || !config.features.enabled(Feature::MemoryTool) + || source.is_non_root_agent() + { + return; + } + + let context = Arc::new(MemoryStartupContext::new( + thread_manager, + auth_manager, + thread_id, + thread, + config.as_ref(), + source.clone(), + )); + + if context.state_db().is_none() { + warn!("state db unavailable for memories startup pipeline; skipping"); + return; + } + + tokio::spawn(async move { + // Clean memories to make preserve DB size + phase1::prune(context.as_ref(), &config).await; + // Run phase 1. + phase1::run(Arc::clone(&context), Arc::clone(&config)).await; + // Run phase 2. + phase2::run(context, config).await; + }); +} diff --git a/codex-rs/core/tests/suite/memories.rs b/codex-rs/memories/write/src/startup_tests.rs similarity index 68% rename from codex-rs/core/tests/suite/memories.rs rename to codex-rs/memories/write/src/startup_tests.rs index a585ed24c..d89b68825 100644 --- a/codex-rs/core/tests/suite/memories.rs +++ b/codex-rs/memories/write/src/startup_tests.rs @@ -1,10 +1,10 @@ -use anyhow::Result; -use chrono::Duration as ChronoDuration; -use chrono::Utc; +use crate::start_memories_startup_task; use codex_features::Feature; use codex_git_utils::diff_since_latest_init; use codex_git_utils::reset_git_repository; use codex_protocol::ThreadId; +use codex_protocol::config_types::ServiceTier; +use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::SessionSource; @@ -26,18 +26,18 @@ use tempfile::TempDir; use tokio::time::Duration; use tokio::time::Instant; -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn memories_startup_phase2_tracks_workspace_diff_across_runs() -> Result<()> { +#[tokio::test] +async fn memories_startup_phase2_tracks_workspace_diff_across_runs() -> anyhow::Result<()> { let server = start_mock_server().await; let home = Arc::new(TempDir::new()?); let db = init_state_db(&home).await?; let memory_root = home.path().join("memories"); - let now = Utc::now(); + let now = chrono::Utc::now(); let _thread_a = seed_stage1_output( db.as_ref(), home.path(), - now - ChronoDuration::hours(2), + now - chrono::Duration::hours(2), "raw memory A", "rollout summary A", "rollout-a", @@ -61,7 +61,7 @@ async fn memories_startup_phase2_tracks_workspace_diff_across_runs() -> Result<( let _thread_b = seed_stage1_output( db.as_ref(), home.path(), - now - ChronoDuration::hours(1), + now - chrono::Duration::hours(1), "raw memory B", "rollout summary B", "rollout-b", @@ -78,7 +78,9 @@ async fn memories_startup_phase2_tracks_workspace_diff_across_runs() -> Result<( ) .await; - let codex = build_test_codex(&server, home.clone()).await?; + let test = build_test_codex(&server, home.clone()).await?; + trigger_memories_startup(&test).await; + let request = wait_for_single_request(&phase2).await; let prompt = phase2_prompt_text(&request); assert!( @@ -108,20 +110,20 @@ async fn memories_startup_phase2_tracks_workspace_diff_across_runs() -> Result<( .all(|summary| !summary.contains("rollout summary A")) ); - shutdown_test_codex(&codex).await?; + shutdown_test_codex(&test).await?; Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn memories_startup_phase2_prunes_old_extension_resources() -> Result<()> { +#[tokio::test] +async fn memories_startup_phase2_prunes_old_extension_resources() -> anyhow::Result<()> { let server = start_mock_server().await; let home = Arc::new(TempDir::new()?); let db = init_state_db(&home).await?; - let now = Utc::now(); + let now = chrono::Utc::now(); let _thread_id = seed_stage1_output( db.as_ref(), home.path(), - now - ChronoDuration::hours(1), + now - chrono::Duration::hours(1), "raw memory", "rollout summary", "rollout", @@ -136,15 +138,14 @@ async fn memories_startup_phase2_prunes_old_extension_resources() -> Result<()> "instructions", ) .await?; - let old_file_name = format!( + let old_file = chronicle_resources.join(format!( "{}-abcd-10min-old.md", - (now - ChronoDuration::days(8)).format("%Y-%m-%dT%H-%M-%S") - ); - let old_file = chronicle_resources.join(&old_file_name); + (now - chrono::Duration::days(8)).format("%Y-%m-%dT%H-%M-%S") + )); tokio::fs::write(&old_file, "old resource").await?; let recent_file = chronicle_resources.join(format!( "{}-abcd-10min-recent.md", - (now - ChronoDuration::days(6)).format("%Y-%m-%dT%H-%M-%S") + (now - chrono::Duration::days(6)).format("%Y-%m-%dT%H-%M-%S") )); tokio::fs::write(&recent_file, "recent resource").await?; @@ -158,10 +159,11 @@ async fn memories_startup_phase2_prunes_old_extension_resources() -> Result<()> ) .await; - let codex = build_test_codex(&server, home.clone()).await?; + let test = build_test_codex(&server, home.clone()).await?; + trigger_memories_startup(&test).await; + let request = wait_for_single_request(&phase2).await; let prompt = phase2_prompt_text(&request); - assert!( prompt.contains("phase2_workspace_diff.md"), "expected workspace diff file in prompt: {prompt}" @@ -178,20 +180,20 @@ async fn memories_startup_phase2_prunes_old_extension_resources() -> Result<()> "recent extension resource should be retained" ); - shutdown_test_codex(&codex).await?; + shutdown_test_codex(&test).await?; Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn memories_startup_phase2_prunes_old_extension_resources_without_stage1_input() -> Result<()> -{ +#[tokio::test] +async fn memories_startup_phase2_prunes_old_extension_resources_without_stage1_input() +-> anyhow::Result<()> { let server = start_mock_server().await; let home = Arc::new(TempDir::new()?); let db = init_state_db(&home).await?; db.enqueue_global_consolidation(/*input_watermark*/ 1) .await?; - let now = Utc::now(); + let now = chrono::Utc::now(); let chronicle_resources = home.path().join("memories/extensions/chronicle/resources"); tokio::fs::create_dir_all(&chronicle_resources).await?; tokio::fs::write( @@ -200,11 +202,10 @@ async fn memories_startup_phase2_prunes_old_extension_resources_without_stage1_i "instructions", ) .await?; - let old_file_name = format!( + let old_file = chronicle_resources.join(format!( "{}-abcd-10min-old.md", - (now - ChronoDuration::days(8)).format("%Y-%m-%dT%H-%M-%S") - ); - let old_file = chronicle_resources.join(&old_file_name); + (now - chrono::Duration::days(8)).format("%Y-%m-%dT%H-%M-%S") + )); tokio::fs::write(&old_file, "old resource").await?; let phase2 = mount_sse_once( @@ -217,52 +218,120 @@ async fn memories_startup_phase2_prunes_old_extension_resources_without_stage1_i ) .await; - let codex = build_test_codex(&server, home.clone()).await?; + let test = build_test_codex(&server, home.clone()).await?; + trigger_memories_startup(&test).await; + let request = wait_for_single_request(&phase2).await; let prompt = phase2_prompt_text(&request); - assert!( prompt.contains("phase2_workspace_diff.md"), "expected workspace diff file in prompt: {prompt}" ); + wait_for_file_removed(&old_file).await?; wait_for_phase2_workspace_reset(&home.path().join("memories")).await?; - shutdown_test_codex(&codex).await?; + shutdown_test_codex(&test).await?; Ok(()) } -async fn build_test_codex(server: &wiremock::MockServer, home: Arc) -> Result { - #[allow(clippy::expect_used)] - let mut builder = test_codex().with_home(home).with_config(|config| { - config - .features - .enable(Feature::Sqlite) - .expect("test config should allow feature update"); - config - .features - .enable(Feature::MemoryTool) - .expect("test config should allow feature update"); - config.memories.max_raw_memories_for_consolidation = 1; - }); - builder.build(server).await +#[tokio::test] +async fn memories_startup_phase1_uses_live_thread_service_tier() -> anyhow::Result<()> { + let server = start_mock_server().await; + let home = Arc::new(TempDir::new()?); + let test = build_test_codex(&server, home).await?; + assert_eq!(test.config.service_tier, None); + + test.codex + .submit(Op::OverrideTurnContext { + cwd: None, + approval_policy: None, + approvals_reviewer: None, + sandbox_policy: None, + permission_profile: None, + windows_sandbox_level: None, + model: None, + effort: None, + summary: None, + service_tier: Some(Some(ServiceTier::Fast)), + collaboration_mode: None, + personality: None, + }) + .await?; + + let config_snapshot = wait_for_service_tier(&test, Some(ServiceTier::Fast)).await?; + assert_eq!(config_snapshot.service_tier, Some(ServiceTier::Fast)); + + let context = crate::runtime::MemoryStartupContext::new( + Arc::clone(&test.thread_manager), + test.thread_manager.auth_manager(), + test.session_configured.session_id, + Arc::clone(&test.codex), + &test.config, + config_snapshot.session_source.clone(), + ); + let request_context = context + .stage_one_request_context( + &test.config, + test.config.model.as_deref().unwrap_or("gpt-5.4-mini"), + ReasoningEffort::Low, + ) + .await; + assert_eq!(request_context.service_tier, Some(ServiceTier::Fast)); + + shutdown_test_codex(&test).await?; + Ok(()) } -async fn init_state_db(home: &Arc) -> Result> { +async fn build_test_codex( + server: &wiremock::MockServer, + home: Arc, +) -> anyhow::Result { + test_codex() + .with_home(home) + .with_config(|config| { + config + .features + .enable(Feature::Sqlite) + .expect("test config should allow feature update"); + config.memories.max_raw_memories_for_consolidation = 1; + }) + .build(server) + .await +} + +async fn init_state_db(home: &Arc) -> anyhow::Result> { let db = codex_state::StateRuntime::init(home.path().to_path_buf(), "test-provider".into()).await?; db.mark_backfill_complete(/*last_watermark*/ None).await?; Ok(db) } +async fn trigger_memories_startup(test: &TestCodex) { + let config_snapshot = test.codex.config_snapshot().await; + let mut config = test.config.clone(); + config + .features + .enable(Feature::MemoryTool) + .expect("test config should allow feature update"); + start_memories_startup_task( + Arc::clone(&test.thread_manager), + test.thread_manager.auth_manager(), + test.session_configured.session_id, + Arc::clone(&test.codex), + Arc::new(config), + &config_snapshot.session_source, + ); +} + async fn seed_stage1_output( db: &codex_state::StateRuntime, codex_home: &Path, - updated_at: chrono::DateTime, + updated_at: chrono::DateTime, raw_memory: &str, rollout_summary: &str, rollout_slug: &str, -) -> Result { +) -> anyhow::Result { let thread_id = ThreadId::new(); let mut metadata_builder = codex_state::ThreadMetadataBuilder::new( thread_id, @@ -293,7 +362,7 @@ async fn wait_for_single_request(mock: &ResponseMock) -> ResponsesRequest { wait_for_request(mock, /*expected_count*/ 1).await.remove(0) } -async fn wait_for_file_removed(path: &Path) -> Result<()> { +async fn wait_for_file_removed(path: &Path) -> anyhow::Result<()> { let deadline = Instant::now() + Duration::from_secs(10); loop { if !tokio::fs::try_exists(path).await? { @@ -323,7 +392,25 @@ async fn wait_for_request(mock: &ResponseMock, expected_count: usize) -> Vec, +) -> anyhow::Result { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let config_snapshot = test.codex.config_snapshot().await; + if config_snapshot.service_tier == expected_service_tier { + return Ok(config_snapshot); + } + anyhow::ensure!( + Instant::now() < deadline, + "timed out waiting for service_tier to become {expected_service_tier:?}, current={:?}", + config_snapshot.service_tier + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } +} + fn phase2_prompt_text(request: &ResponsesRequest) -> String { request .message_input_texts("user") @@ -332,7 +419,7 @@ fn phase2_prompt_text(request: &ResponsesRequest) -> String { .expect("phase2 prompt text") } -async fn wait_for_phase2_workspace_reset(memory_root: &Path) -> Result<()> { +async fn wait_for_phase2_workspace_reset(memory_root: &Path) -> anyhow::Result<()> { wait_for_file_removed(&memory_root.join("phase2_workspace_diff.md")).await?; let deadline = Instant::now() + Duration::from_secs(10); loop { @@ -356,7 +443,7 @@ async fn seed_stage1_output_for_existing_thread( raw_memory: &str, rollout_summary: &str, rollout_slug: Option<&str>, -) -> Result<()> { +) -> anyhow::Result<()> { let owner = ThreadId::new(); let claim = db .try_claim_stage1_job( @@ -385,7 +472,7 @@ async fn seed_stage1_output_for_existing_thread( Ok(()) } -async fn read_rollout_summary_bodies(memory_root: &Path) -> Result> { +async fn read_rollout_summary_bodies(memory_root: &Path) -> anyhow::Result> { let mut dir = tokio::fs::read_dir(memory_root.join("rollout_summaries")).await?; let mut summaries = Vec::new(); while let Some(entry) = dir.next_entry().await? { @@ -395,7 +482,7 @@ async fn read_rollout_summary_bodies(memory_root: &Path) -> Result> Ok(summaries) } -async fn shutdown_test_codex(test: &TestCodex) -> Result<()> { +async fn shutdown_test_codex(test: &TestCodex) -> anyhow::Result<()> { test.codex.submit(Op::Shutdown {}).await?; wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::ShutdownComplete)).await; Ok(()) diff --git a/codex-rs/memories/write/src/storage_tests.rs b/codex-rs/memories/write/src/storage_tests.rs index 28344c581..95e57aa3f 100644 --- a/codex-rs/memories/write/src/storage_tests.rs +++ b/codex-rs/memories/write/src/storage_tests.rs @@ -1,10 +1,18 @@ use super::rollout_summary_file_stem; +use crate::ensure_layout; +use crate::raw_memories_file; +use crate::rebuild_raw_memories_file_from_memories; +use crate::rollout_summaries_dir; +use crate::sync_rollout_summaries_from_memories; use chrono::TimeZone; use chrono::Utc; +use codex_config::types::DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION; use codex_protocol::ThreadId; use codex_state::Stage1Output; use pretty_assertions::assert_eq; use std::path::PathBuf; +use tempfile::tempdir; + const FIXED_PREFIX: &str = "2025-02-11T15-35-19-jqmb"; fn stage1_output_with_slug(thread_id: ThreadId, rollout_slug: Option<&str>) -> Stage1Output { @@ -59,3 +67,83 @@ fn rollout_summary_file_stem_uses_uuid_timestamp_and_hash_when_slug_is_empty() { assert_eq!(rollout_summary_file_stem(&memory), FIXED_PREFIX); } + +#[tokio::test] +async fn sync_rollout_summaries_and_raw_memories_file_keeps_latest_memories_only() { + let dir = tempdir().expect("tempdir"); + let root = dir.path().join("memory"); + ensure_layout(&root).await.expect("ensure layout"); + + let keep_id = ThreadId::default().to_string(); + let drop_id = ThreadId::default().to_string(); + let keep_path = rollout_summaries_dir(&root).join(format!("{keep_id}.md")); + let drop_path = rollout_summaries_dir(&root).join(format!("{drop_id}.md")); + tokio::fs::write(&keep_path, "keep") + .await + .expect("write keep"); + tokio::fs::write(&drop_path, "drop") + .await + .expect("write drop"); + + let memories = vec![Stage1Output { + thread_id: ThreadId::try_from(keep_id.clone()).expect("thread id"), + source_updated_at: Utc.timestamp_opt(100, 0).single().expect("timestamp"), + raw_memory: "raw memory".to_string(), + rollout_summary: "short summary".to_string(), + rollout_slug: None, + rollout_path: PathBuf::from("/tmp/rollout-100.jsonl"), + cwd: PathBuf::from("/tmp/workspace"), + git_branch: None, + generated_at: Utc.timestamp_opt(101, 0).single().expect("timestamp"), + }]; + + sync_rollout_summaries_from_memories( + &root, + &memories, + DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION, + ) + .await + .expect("sync rollout summaries"); + rebuild_raw_memories_file_from_memories( + &root, + &memories, + DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_CONSOLIDATION, + ) + .await + .expect("rebuild raw memories"); + + assert!( + !tokio::fs::try_exists(&keep_path) + .await + .expect("check stale keep path"), + "sync should prune stale filename that used thread id only" + ); + assert!( + !tokio::fs::try_exists(&drop_path) + .await + .expect("check stale drop path"), + "sync should prune stale filename for dropped thread" + ); + + let mut dir = tokio::fs::read_dir(rollout_summaries_dir(&root)) + .await + .expect("open rollout summaries dir"); + let mut files = Vec::new(); + while let Some(entry) = dir.next_entry().await.expect("read dir entry") { + files.push(entry.file_name().to_string_lossy().to_string()); + } + files.sort_unstable(); + assert_eq!(files.len(), 1); + let canonical_rollout_summary_file = &files[0]; + + let raw_memories = tokio::fs::read_to_string(raw_memories_file(&root)) + .await + .expect("read raw memories"); + assert!(raw_memories.contains("raw memory")); + assert!(raw_memories.contains(&keep_id)); + assert!(raw_memories.contains("cwd: /tmp/workspace")); + assert!(raw_memories.contains("rollout_path: /tmp/rollout-100.jsonl")); + assert!(raw_memories.contains(&format!( + "rollout_summary_file: {canonical_rollout_summary_file}" + ))); +} diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index c731dfae3..c39c017fe 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -758,12 +758,6 @@ pub enum Op { /// to generate a summary which will be returned as an AgentMessage event. Compact, - /// Drop all persisted memory artifacts and memory-tracking DB rows. - DropMemories, - - /// Trigger a single pass of the startup memory pipeline. - UpdateMemories, - /// Set a user-facing thread name in the persisted rollout metadata. /// This is a local-only operation handled by codex-core; it does not /// involve the model. @@ -906,8 +900,6 @@ impl Op { Self::ReloadUserConfig => "reload_user_config", Self::ListSkills { .. } => "list_skills", Self::Compact => "compact", - Self::DropMemories => "drop_memories", - Self::UpdateMemories => "update_memories", Self::SetThreadName { .. } => "set_thread_name", Self::SetThreadMemoryMode { .. } => "set_thread_memory_mode", Self::Undo => "undo", @@ -2640,11 +2632,19 @@ pub enum SessionSource { Exec, Mcp, Custom(String), + Internal(InternalSessionSource), SubAgent(SubAgentSource), #[serde(other)] Unknown, } +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "snake_case")] +#[ts(rename_all = "snake_case")] +pub enum InternalSessionSource { + MemoryConsolidation, +} + #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS)] #[serde(rename_all = "snake_case")] #[ts(rename_all = "snake_case")] @@ -2673,6 +2673,7 @@ impl fmt::Display for SessionSource { SessionSource::Exec => f.write_str("exec"), SessionSource::Mcp => f.write_str("mcp"), SessionSource::Custom(source) => f.write_str(source), + SessionSource::Internal(source) => write!(f, "internal_{source}"), SessionSource::SubAgent(sub_source) => write!(f, "subagent_{sub_source}"), SessionSource::Unknown => f.write_str("unknown"), } @@ -2701,19 +2702,28 @@ impl SessionSource { pub fn thread_source_name(&self) -> Option<&'static str> { match self { SessionSource::Cli | SessionSource::VSCode | SessionSource::Exec => Some("user"), + SessionSource::Internal(_) => Some("internal"), SessionSource::SubAgent(_) => Some("subagent"), SessionSource::Mcp | SessionSource::Custom(_) | SessionSource::Unknown => None, } } + pub fn is_internal(&self) -> bool { + matches!(self, SessionSource::Internal(_)) + } + + pub fn is_non_root_agent(&self) -> bool { + matches!( + self, + SessionSource::Internal(_) | SessionSource::SubAgent(_) + ) + } + pub fn get_nickname(&self) -> Option { match self { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { agent_nickname, .. }) => { agent_nickname.clone() } - SessionSource::SubAgent(SubAgentSource::MemoryConsolidation) => { - Some("Morpheus".to_string()) - } _ => None, } } @@ -2723,9 +2733,6 @@ impl SessionSource { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { agent_role, .. }) => { agent_role.clone() } - SessionSource::SubAgent(SubAgentSource::MemoryConsolidation) => { - Some("memory builder".to_string()) - } _ => None, } } @@ -2735,9 +2742,6 @@ impl SessionSource { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { agent_path, .. }) => { agent_path.clone() } - SessionSource::SubAgent(SubAgentSource::MemoryConsolidation) => { - Some(AgentPath::morpheus()) - } _ => None, } } @@ -2750,7 +2754,7 @@ impl SessionSource { | SessionSource::Exec | SessionSource::Mcp | SessionSource::Unknown => Some(Product::Codex), - SessionSource::SubAgent(_) => None, + SessionSource::Internal(_) | SessionSource::SubAgent(_) => None, } } @@ -2780,6 +2784,14 @@ impl fmt::Display for SubAgentSource { } } +impl fmt::Display for InternalSessionSource { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + InternalSessionSource::MemoryConsolidation => f.write_str("memory_consolidation"), + } + } +} + /// SessionMeta contains session-level data that doesn't correspond to a specific turn. /// /// NOTE: There used to be an `instructions` field here, which stored user_instructions, but we @@ -4101,6 +4113,10 @@ mod tests { (SessionSource::Cli, Some("user")), (SessionSource::VSCode, Some("user")), (SessionSource::Exec, Some("user")), + ( + SessionSource::Internal(InternalSessionSource::MemoryConsolidation), + Some("internal"), + ), ( SessionSource::SubAgent(SubAgentSource::Review), Some("subagent"), @@ -4143,6 +4159,11 @@ mod tests { SessionSource::SubAgent(SubAgentSource::Review).restriction_product(), None ); + assert_eq!( + SessionSource::Internal(InternalSessionSource::MemoryConsolidation) + .restriction_product(), + None + ); } #[test] diff --git a/codex-rs/thread-store/src/remote/helpers.rs b/codex-rs/thread-store/src/remote/helpers.rs index 5f5bf503a..8b0a73922 100644 --- a/codex-rs/thread-store/src/remote/helpers.rs +++ b/codex-rs/thread-store/src/remote/helpers.rs @@ -124,6 +124,7 @@ pub(super) fn proto_session_source(source: &SessionSource) -> proto::SessionSour sub_agent_other: Some(other.clone()), ..Default::default() }, + SessionSource::Internal(_) => proto_source(proto::SessionSourceKind::Unknown), SessionSource::Unknown => proto_source(proto::SessionSourceKind::Unknown), } }