agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: split memories part 2 (#19860)

Browse Source 
Keep extracting memories out of core and moving the write trigger in the
app-server
This is temporary and it should move at the client level as a follow-up
This makes core fully independant from `codex-memories-write`

---------

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
jif-oai
2026-04-28 13:03:28 +02:00
committed by GitHub
Unverified
parent fd36838cf3
commit 431ebeaef7
47 changed files with 1932 additions and 2735 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+16 -2
View File
@@ -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]]
codex-rs/agent-identity/src/lib.rs
+1
View File
@@ -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(),
},
codex-rs/analytics/src/reducer.rs
+1
View File
@@ -106,6 +106,7 @@ impl ThreadMetadataState {
| SessionSource::Exec
| SessionSource::Mcp
| SessionSource::Custom(_)
| SessionSource::Internal(_)
| SessionSource::Unknown => (None, None),
};
Self {
codex-rs/app-server-protocol/schema/typescript/InternalSessionSource.ts
Generated
+5
View File
@@ -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";
codex-rs/app-server-protocol/schema/typescript/SessionSource.ts
Generated
+2 -1
View File
@@ -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";
codex-rs/app-server-protocol/schema/typescript/index.ts
Generated
+1
View File
@@ -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";
codex-rs/app-server-protocol/src/protocol/v2.rs
+2
View File
@@ -1996,6 +1996,8 @@ impl From<CoreSessionSource> 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,
}
codex-rs/app-server/Cargo.toml
+1
View File
@@ -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 }
codex-rs/app-server/src/codex_message_processor.rs
+38 -6
View File
@@ -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<AuthManager>,
thread_manager: Arc<ThreadManager>,
thread_state_manager: ThreadStateManager,
outgoing: Arc<OutgoingMessageSender>,
@@ -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<ThreadLoadedListResponse, JSONRPCErrorError> {
let ThreadLoadedListParams { cursor, limit } = params;
let mut data = self
let mut data: Vec<String> = self
.thread_manager
.list_thread_ids()
.await
.into_iter()
.map(|thread_id| thread_id.to_string())
.collect::<Vec<_>>();
.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<EnsureConversationListenerResult, JSONRPCErrorError> {
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,
codex-rs/cli/Cargo.toml
+1
View File
@@ -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 }
codex-rs/cli/src/main.rs
+1 -1
View File
@@ -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;
codex-rs/core/Cargo.toml
-2
View File
@@ -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 }
codex-rs/core/src/agent/control.rs
+1 -20
View File
@@ -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<TokenUsage> {
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,
codex-rs/core/src/client.rs
+20 -10
View File
@@ -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<String> {
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<Strin
| SessionSource::Exec
| SessionSource::Mcp
| SessionSource::Custom(_)
| SessionSource::Internal(_)
| SessionSource::SubAgent(_)
| SessionSource::Unknown => None,
}
codex-rs/core/src/client_tests.rs
+13
View File
@@ -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();
codex-rs/core/src/codex_thread.rs
+8 -5
View File
@@ -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<PathBuf>,
out_of_band_elicitation_count: Mutex<u64>,
_watch_registration: WatchRegistration,
@@ -97,10 +97,12 @@ impl CodexThread {
pub(crate) fn new(
codex: Codex,
rollout_path: Option<PathBuf>,
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<TokenUsage> {
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
codex-rs/core/src/config/mod.rs
+1 -1
View File
@@ -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;
codex-rs/core/src/context/contextual_user_message.rs
-22
View File
@@ -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;
codex-rs/core/src/context/contextual_user_message_tests.rs
-32
View File
@@ -33,38 +33,6 @@ fn ignores_regular_user_text() {
}));
}
#[test]
fn classifies_memory_excluded_fragments() {
let cases = [
(
"# AGENTS.md instructions for /tmp\n\n<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>",
true,
),
(
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>",
true,
),
(
"<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>",
false,
),
(
"<subagent_notification>{\"agent_id\":\"a\",\"status\":\"completed\"}</subagent_notification>",
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(
codex-rs/core/src/context/mod.rs
-1
View File
@@ -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;
codex-rs/core/src/installation_id.rs
+1 -1
View File
@@ -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<String> {
pub async fn resolve_installation_id(codex_home: &AbsolutePathBuf) -> Result<String> {
let path = codex_home.join(INSTALLATION_ID_FILENAME);
fs::create_dir_all(codex_home).await?;
tokio::task::spawn_blocking(move || {
codex-rs/core/src/lib.rs
+2 -3
View File
@@ -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;
codex-rs/core/src/memories/mod.rs
-73
View File
@@ -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";
}
codex-rs/core/src/memories/phase1.rs
-620
View File
@@ -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<ReasoningEffortConfig>,
pub(in crate::memories) reasoning_summary: ReasoningSummaryConfig,
pub(in crate::memories) service_tier: Option<ServiceTier>,
pub(in crate::memories) turn_metadata_header: Option<String>,
}
struct JobResult {
outcome: JobOutcome,
token_usage: Option<TokenUsage>,
}
#[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<TokenUsage>,
}
/// 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<String>,
}
/// 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<Session>, 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<Session>, 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<String>,
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<Session>,
memories_config: &MemoriesConfig,
) -> Option<Vec<codex_state::Stage1JobClaim>> {
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::<Vec<_>>();
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<Session>, 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<Session>,
claimed_candidates: Vec<codex_state::Stage1JobClaim>,
stage_one_context: RequestContext,
) -> Vec<JobResult> {
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::<Vec<_>>()
.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<TokenUsage>)> {
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<String> {
let filtered = items
.iter()
.filter_map(|item| {
if let RolloutItem::ResponseItem(item) = item {
sanitize_response_item_for_memories(item)
} else {
None
}
})
.collect::<Vec<_>>();
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<ResponseItem> {
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::<Vec<_>>();
if content.is_empty() {
return None;
}
Some(ResponseItem::Message {
id: id.clone(),
role: role.clone(),
content,
phase: phase.clone(),
})
}
}
fn aggregate_stats(outcomes: Vec<JobResult>) -> 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;
codex-rs/core/src/memories/phase1_tests.rs
-152
View File
@@ -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<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>"
.to_string(),
},
ContentItem::InputText {
text: "<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>".to_string(),
},
],
phase: None,
};
let skill_message = ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>"
.to_string(),
}],
phase: None,
};
let subagent_message = ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<subagent_notification>{\"agent_id\":\"a\",\"status\":\"completed\"}</subagent_notification>"
.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<ResponseItem> = serde_json::from_str(&serialized).expect("parse");
assert_eq!(
parsed,
vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>"
.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);
}
codex-rs/core/src/memories/start.rs
-44
View File
@@ -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<Session>,
config: Arc<Config>,
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;
});
}
codex-rs/core/src/memories/tests.rs
-1358
View File
File diff suppressed because it is too large Load Diff
codex-rs/core/src/rollout.rs
-4
View File
@@ -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;
}
codex-rs/core/src/session/handlers.rs
-68
View File
@@ -670,66 +670,6 @@ pub async fn compact(sess: &Arc<Session>, sub_id: String) {
.await;
}
pub async fn drop_memories(sess: &Arc<Session>, config: &Arc<Config>, 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<Session>, config: &Arc<Config>, 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<Session>, 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
codex-rs/core/src/session/mod.rs
+8 -8
View File
@@ -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);
{
codex-rs/core/src/session/session.rs
+1 -10
View File
@@ -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;
codex-rs/core/src/thread_manager.rs
+37 -21
View File
@@ -211,9 +211,10 @@ pub struct ThreadManager {
_test_codex_home_guard: Option<TempCodexHomeGuard>,
}
pub struct StartThreadWithToolsOptions {
pub struct StartThreadOptions {
pub config: Config,
pub initial_history: InitialHistory,
pub session_source: Option<SessionSource>,
pub dynamic_tools: Vec<codex_protocol::dynamic_tools::DynamicToolSpec>,
pub persist_extended_history: bool,
pub metrics_service_name: Option<String>,
@@ -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<NewThread> {
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<ThreadId> {
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<Arc<CodexThread>> {
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<NewThread> {
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;
codex-rs/core/src/thread_manager_tests.rs
+49 -2
View File
@@ -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,
codex-rs/core/src/tools/handlers/request_user_input.rs
+1 -2
View File
@@ -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(),
));
codex-rs/core/src/tools/handlers/request_user_input_tests.rs
+1
View File
@@ -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;
codex-rs/core/tests/suite/mod.rs
-1
View File
@@ -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;
codex-rs/memories/write/Cargo.toml
+15 -1
View File
@@ -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 }
codex-rs/memories/write/src/control.rs
+72
View File
@@ -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"
);
}
}
codex-rs/memories/write/src/lib.rs
+12 -4
View File
@@ -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";
codex-rs/memories/write/src/phase1.rs
+803
View File
@@ -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<TokenUsage>,
}
#[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<TokenUsage>,
}
/// 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<String>,
}
/// 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<MemoryStartupContext>, config: Arc<Config>) {
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<Vec<codex_state::Stage1JobClaim>> {
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::<Vec<_>>();
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<MemoryStartupContext>,
config: Arc<Config>,
claimed_candidates: Vec<codex_state::Stage1JobClaim>,
stage_one_context: StageOneRequestContext,
) -> Vec<JobResult> {
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::<Vec<_>>()
.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<TokenUsage>)> {
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<String> {
let filtered = items
.iter()
.filter_map(|item| {
if let RolloutItem::ResponseItem(item) = item {
sanitize_response_item_for_memories(item)
} else {
None
}
})
.collect::<Vec<_>>();
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<ResponseItem> {
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::<Vec<_>>();
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 ", "</INSTRUCTIONS>")
|| matches_marked_fragment(text, "<skill>", "</skill>")
}
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<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>",
true,
),
(
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>",
true,
),
(
"<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>",
false,
),
(
"<subagent_notification>{\"agent_id\":\"a\",\"status\":\"completed\"}</subagent_notification>",
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::<Vec<_>>();
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::<Vec<_>>();
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<JobResult>) -> 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<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>"
.to_string(),
},
ContentItem::InputText {
text: "<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>"
.to_string(),
},
],
phase: None,
};
let skill_message = ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text:
"<skill>\n<name>demo</name>\n<path>skills/demo/SKILL.md</path>\nbody\n</skill>"
.to_string(),
}],
phase: None,
};
let subagent_message = ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<subagent_notification>{\"agent_id\":\"a\",\"status\":\"completed\"}</subagent_notification>"
.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<ResponseItem> = serde_json::from_str(&serialized).expect("parse");
assert_eq!(
parsed,
vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<environment_context>\n<cwd>/tmp</cwd>\n</environment_context>"
.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);
}
}
codex-rs/core/src/memories/phase2.rs → codex-rs/memories/write/src/phase2.rs
+199 -186
View File
@@ -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<Session>, config: Arc<Config>) {
let phase_two_e2e_timer = session
.services
.session_telemetry
.start_timer(metrics::MEMORY_PHASE_TWO_E2E_MS, &[])
.ok();
pub async fn run(context: Arc<MemoryStartupContext>, config: Arc<Config>) {
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<Session>, config: Arc<Config>) {
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<Session>, config: Arc<Config>) {
// 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<Session>, config: Arc<Config>) {
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<Session>, config: Arc<Config>) {
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<Session>, config: Arc<Config>) {
// 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<Session>, config: Arc<Config>) {
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<Session>, config: Arc<Config>) {
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<Session>, config: Arc<Config>) {
// 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<Session>, config: Arc<Config>) {
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<Session>,
context: &MemoryStartupContext,
db: &StateRuntime,
) -> Result<Claim, &'static str> {
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<Session>,
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<Session>,
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<Config> {
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<Session>,
context: Arc<MemoryStartupContext>,
claim: Claim,
new_watermark: i64,
selected_outputs: Vec<codex_state::Stage1Output>,
memory_root: codex_utils_absolute_path::AbsolutePathBuf,
thread_id: ThreadId,
agent_control: crate::agent::AgentControl,
agent: SpawnedConsolidationAgent,
phase_two_e2e_timer: Option<codex_otel::Timer>,
) {
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<StateRuntime>,
token: String,
thread_id: ThreadId,
mut rx: watch::Receiver<AgentStatus>,
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<Session>, 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<Session>, 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")],
);
codex-rs/memories/write/src/runtime.rs
+293
View File
@@ -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<CodexThread>,
}
#[derive(Clone, Debug)]
pub(crate) struct StageOneRequestContext {
pub(crate) model_info: ModelInfo,
pub(crate) session_telemetry: SessionTelemetry,
pub(crate) reasoning_effort: Option<ReasoningEffort>,
pub(crate) reasoning_summary: ReasoningSummary,
pub(crate) service_tier: Option<ServiceTier>,
pub(crate) turn_metadata_header: Option<String>,
}
impl StageOneRequestContext {
pub(crate) fn start_timer(&self, name: &str) -> Option<codex_otel::Timer> {
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<CodexThread>,
thread_manager: Arc<ThreadManager>,
auth_manager: Arc<AuthManager>,
session_telemetry: SessionTelemetry,
}
impl MemoryStartupContext {
pub(crate) fn new(
thread_manager: Arc<ThreadManager>,
auth_manager: Arc<AuthManager>,
thread_id: ThreadId,
thread: Arc<CodexThread>,
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<Arc<StateRuntime>> {
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<codex_otel::Timer> {
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<TokenUsage>)> {
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<UserInput>,
) -> anyhow::Result<SpawnedConsolidationAgent> {
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(())
}
}
codex-rs/memories/write/src/start.rs
+55
View File
@@ -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<ThreadManager>,
auth_manager: Arc<AuthManager>,
thread_id: ThreadId,
thread: Arc<CodexThread>,
config: Arc<Config>,
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;
});
}
codex-rs/core/tests/suite/memories.rs → codex-rs/memories/write/src/startup_tests.rs
+143 -56
View File
@@ -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<TempDir>) -> Result<TestCodex> {
#[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<TempDir>) -> Result<Arc<codex_state::StateRuntime>> {
async fn build_test_codex(
server: &wiremock::MockServer,
home: Arc<TempDir>,
) -> anyhow::Result<TestCodex> {
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<TempDir>) -> anyhow::Result<Arc<codex_state::StateRuntime>> {
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<Utc>,
updated_at: chrono::DateTime<chrono::Utc>,
raw_memory: &str,
rollout_summary: &str,
rollout_slug: &str,
) -> Result<ThreadId> {
) -> anyhow::Result<ThreadId> {
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<Res
}
}
#[allow(clippy::expect_used)]
async fn wait_for_service_tier(
test: &TestCodex,
expected_service_tier: Option<ServiceTier>,
) -> anyhow::Result<codex_core::ThreadConfigSnapshot> {
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<Vec<String>> {
async fn read_rollout_summary_bodies(memory_root: &Path) -> anyhow::Result<Vec<String>> {
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<Vec<String>>
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(())
codex-rs/memories/write/src/storage_tests.rs
+88
View File
@@ -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}"
)));
}
codex-rs/protocol/src/protocol.rs
+39 -18
View File
@@ -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<String> {
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]
codex-rs/thread-store/src/remote/helpers.rs
+1
View File
@@ -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),
}
}