mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
feat: split memories part 2 (#19860)
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
committed by
GitHub
Unverified
GitHub
parent
fd36838cf3
commit
431ebeaef7
@@ -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,
|
||||
|
||||
+20
-10
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 || {
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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";
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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);
|
||||
}
|
||||
@@ -1,552 +0,0 @@
|
||||
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 codex_config::Constrained;
|
||||
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::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;
|
||||
use codex_state::StateRuntime;
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::watch;
|
||||
use tracing::warn;
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct Claim {
|
||||
token: String,
|
||||
watermark: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct Counters {
|
||||
input: i64,
|
||||
}
|
||||
|
||||
/// 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();
|
||||
|
||||
let Some(db) = session.services.state_db.as_deref() else {
|
||||
// This should not happen.
|
||||
return;
|
||||
};
|
||||
let root = memory_root(&config.codex_home);
|
||||
let max_raw_memories = config.memories.max_raw_memories_for_consolidation;
|
||||
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 {
|
||||
Ok(claim) => claim,
|
||||
Err(e) => {
|
||||
session.services.session_telemetry.counter(
|
||||
metrics::MEMORY_PHASE_TWO_JOBS,
|
||||
/*inc*/ 1,
|
||||
&[("status", e)],
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// 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;
|
||||
return;
|
||||
}
|
||||
|
||||
// 3. Build the locked-down config used by the consolidation agent.
|
||||
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;
|
||||
return;
|
||||
};
|
||||
|
||||
// 4. Load current DB-backed Phase 2 inputs.
|
||||
let raw_memories = match db
|
||||
.get_phase2_input_selection(max_raw_memories, max_unused_days)
|
||||
.await
|
||||
{
|
||||
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;
|
||||
return;
|
||||
}
|
||||
};
|
||||
let raw_memory_count = raw_memories.len();
|
||||
let new_watermark = get_watermark(claim.watermark, &raw_memories);
|
||||
|
||||
// 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;
|
||||
return;
|
||||
}
|
||||
|
||||
// 6. Use git to decide whether the synced workspace actually changed.
|
||||
let workspace_diff = match memory_workspace_diff(&root).await {
|
||||
Ok(diff) => diff,
|
||||
Err(err) => {
|
||||
tracing::error!("failed checking memory workspace changes: {err}");
|
||||
job::failed(session, db, &claim, "failed_workspace_status").await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
if !workspace_diff.has_changes() {
|
||||
tracing::error!("Phase 2 no changes");
|
||||
// We check only after sync of the file system.
|
||||
job::succeed(
|
||||
session,
|
||||
db,
|
||||
&claim,
|
||||
new_watermark,
|
||||
&raw_memories,
|
||||
"succeeded_no_workspace_changes",
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
// 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;
|
||||
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))
|
||||
.await
|
||||
{
|
||||
Ok(thread_id) => thread_id,
|
||||
Err(err) => {
|
||||
tracing::error!("failed to spawn global memory consolidation agent: {err}");
|
||||
job::failed(session, db, &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,
|
||||
claim,
|
||||
new_watermark,
|
||||
raw_memories.clone(),
|
||||
root,
|
||||
thread_id,
|
||||
agent_control,
|
||||
phase_two_e2e_timer,
|
||||
);
|
||||
|
||||
// 10. Emit dispatch metrics.
|
||||
let counters = Counters {
|
||||
input: raw_memory_count as i64,
|
||||
};
|
||||
emit_metrics(session, counters);
|
||||
}
|
||||
|
||||
async fn sync_phase2_workspace_inputs(
|
||||
root: &Path,
|
||||
raw_memories: &[Stage1Output],
|
||||
) -> std::io::Result<()> {
|
||||
let raw_memory_count = raw_memories.len();
|
||||
sync_rollout_summaries_from_memories(root, raw_memories, raw_memory_count).await?;
|
||||
rebuild_raw_memories_file_from_memories(root, raw_memories, raw_memory_count).await?;
|
||||
prune_old_extension_resources(root).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
mod job {
|
||||
use super::*;
|
||||
|
||||
pub(super) async fn claim(
|
||||
session: &Arc<Session>,
|
||||
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)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!("failed to claim job: {}", e);
|
||||
"failed_claim"
|
||||
})?;
|
||||
let (token, watermark) = match claim {
|
||||
codex_state::Phase2JobClaimOutcome::Claimed {
|
||||
ownership_token,
|
||||
input_watermark,
|
||||
} => {
|
||||
session_telemetry.counter(
|
||||
metrics::MEMORY_PHASE_TWO_JOBS,
|
||||
/*inc*/ 1,
|
||||
&[("status", "claimed")],
|
||||
);
|
||||
(ownership_token, input_watermark)
|
||||
}
|
||||
codex_state::Phase2JobClaimOutcome::SkippedRetryUnavailable => {
|
||||
return Err("skipped_retry_unavailable");
|
||||
}
|
||||
codex_state::Phase2JobClaimOutcome::SkippedRunning => return Err("skipped_running"),
|
||||
};
|
||||
|
||||
Ok(Claim { token, watermark })
|
||||
}
|
||||
|
||||
pub(super) async fn failed(
|
||||
session: &Arc<Session>,
|
||||
db: &StateRuntime,
|
||||
claim: &Claim,
|
||||
reason: &'static str,
|
||||
) {
|
||||
session.services.session_telemetry.counter(
|
||||
metrics::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,
|
||||
Ok(false)
|
||||
) {
|
||||
let _ = db
|
||||
.mark_global_phase2_job_failed_if_unowned(
|
||||
&claim.token,
|
||||
reason,
|
||||
phase_two::JOB_RETRY_DELAY_SECONDS,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn succeed(
|
||||
session: &Arc<Session>,
|
||||
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)],
|
||||
);
|
||||
db.mark_global_phase2_job_succeeded(&claim.token, completion_watermark, selected_outputs)
|
||||
.await
|
||||
.unwrap_or(false)
|
||||
}
|
||||
}
|
||||
|
||||
mod agent {
|
||||
use super::*;
|
||||
|
||||
pub(super) fn get_config(config: &Config) -> Option<Config> {
|
||||
let root = memory_root(&config.codex_home);
|
||||
let mut agent_config = config.clone();
|
||||
|
||||
agent_config.cwd = root.clone();
|
||||
// Consolidation threads must never feed back into phase-1 memory generation.
|
||||
agent_config.ephemeral = true;
|
||||
agent_config.memories.generate_memories = false;
|
||||
agent_config.memories.use_memories = false;
|
||||
agent_config.include_apps_instructions = false;
|
||||
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.
|
||||
let _ = agent_config.features.disable(Feature::SpawnCsv);
|
||||
let _ = agent_config.features.disable(Feature::Collab);
|
||||
let _ = agent_config.features.disable(Feature::MemoryTool);
|
||||
let _ = agent_config.features.disable(Feature::Apps);
|
||||
let _ = agent_config.features.disable(Feature::Plugins);
|
||||
let _ = agent_config
|
||||
.features
|
||||
.disable(Feature::SkillMcpDependencyInstall);
|
||||
|
||||
// Sandbox policy
|
||||
let writable_roots = vec![root];
|
||||
// The consolidation agent only needs local memory-root write access and no network.
|
||||
let consolidation_sandbox_policy = SandboxPolicy::WorkspaceWrite {
|
||||
writable_roots,
|
||||
network_access: false,
|
||||
exclude_tmpdir_env_var: true,
|
||||
exclude_slash_tmp: true,
|
||||
};
|
||||
agent_config
|
||||
.permissions
|
||||
.set_legacy_sandbox_policy(consolidation_sandbox_policy, agent_config.cwd.as_path())
|
||||
.ok()?;
|
||||
|
||||
agent_config.model = Some(
|
||||
config
|
||||
.memories
|
||||
.consolidation_model
|
||||
.clone()
|
||||
.unwrap_or(phase_two::MODEL.to_string()),
|
||||
);
|
||||
agent_config.model_reasoning_effort = Some(phase_two::REASONING_EFFORT);
|
||||
|
||||
Some(agent_config)
|
||||
}
|
||||
|
||||
pub(super) fn get_prompt(root: &Path) -> Vec<UserInput> {
|
||||
let prompt = build_consolidation_prompt(root);
|
||||
vec![UserInput::Text {
|
||||
text: prompt,
|
||||
text_elements: vec![],
|
||||
}]
|
||||
}
|
||||
|
||||
/// Handle the agent while it is running.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) fn handle(
|
||||
session: &Arc<Session>,
|
||||
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,
|
||||
phase_two_e2e_timer: Option<codex_otel::Timer>,
|
||||
) {
|
||||
let Some(db) = session.services.state_db.clone() 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;
|
||||
}
|
||||
};
|
||||
|
||||
// Loop the agent until we have the final status.
|
||||
let final_status = loop_agent(db.clone(), claim.token.clone(), thread_id, rx).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);
|
||||
}
|
||||
// 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)
|
||||
.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;
|
||||
};
|
||||
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}"
|
||||
);
|
||||
}
|
||||
});
|
||||
} else {
|
||||
tracing::warn!("The agent was already gone");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async fn loop_agent(
|
||||
db: Arc<StateRuntime>,
|
||||
token: String,
|
||||
thread_id: ThreadId,
|
||||
mut rx: watch::Receiver<AgentStatus>,
|
||||
) -> AgentStatus {
|
||||
let mut heartbeat_interval =
|
||||
tokio::time::interval(Duration::from_secs(phase_two::JOB_HEARTBEAT_SECONDS));
|
||||
heartbeat_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
|
||||
loop {
|
||||
let status = rx.borrow().clone();
|
||||
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}"
|
||||
);
|
||||
break status;
|
||||
}
|
||||
}
|
||||
_ = heartbeat_interval.tick() => {
|
||||
match db
|
||||
.heartbeat_global_phase2_job(
|
||||
&token,
|
||||
phase_two::JOB_LEASE_SECONDS,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(true) => {}
|
||||
Ok(false) => {
|
||||
break AgentStatus::Errored(
|
||||
"lost global phase-2 ownership during heartbeat".to_string(),
|
||||
);
|
||||
}
|
||||
Err(err) => {
|
||||
break AgentStatus::Errored(format!(
|
||||
"phase-2 heartbeat update failed: {err}"
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn get_watermark(
|
||||
claimed_watermark: i64,
|
||||
latest_memories: &[codex_state::Stage1Output],
|
||||
) -> i64 {
|
||||
latest_memories
|
||||
.iter()
|
||||
.map(|memory| memory.source_updated_at.timestamp())
|
||||
.max()
|
||||
.unwrap_or(claimed_watermark)
|
||||
.max(claimed_watermark)
|
||||
}
|
||||
|
||||
fn emit_metrics(session: &Arc<Session>, counters: Counters) {
|
||||
let otel = session.services.session_telemetry.clone();
|
||||
if counters.input > 0 {
|
||||
otel.counter(metrics::MEMORY_PHASE_TWO_INPUT, counters.input, &[]);
|
||||
}
|
||||
|
||||
otel.counter(
|
||||
metrics::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,
|
||||
token_usage.total_tokens.max(0),
|
||||
&[("token_type", "total")],
|
||||
);
|
||||
otel.histogram(
|
||||
metrics::MEMORY_PHASE_TWO_TOKEN_USAGE,
|
||||
token_usage.input_tokens.max(0),
|
||||
&[("token_type", "input")],
|
||||
);
|
||||
otel.histogram(
|
||||
metrics::MEMORY_PHASE_TWO_TOKEN_USAGE,
|
||||
token_usage.cached_input(),
|
||||
&[("token_type", "cached_input")],
|
||||
);
|
||||
otel.histogram(
|
||||
metrics::MEMORY_PHASE_TWO_TOKEN_USAGE,
|
||||
token_usage.output_tokens.max(0),
|
||||
&[("token_type", "output")],
|
||||
);
|
||||
otel.histogram(
|
||||
metrics::MEMORY_PHASE_TWO_TOKEN_USAGE,
|
||||
token_usage.reasoning_output_tokens.max(0),
|
||||
&[("token_type", "reasoning_output")],
|
||||
);
|
||||
}
|
||||
@@ -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;
|
||||
});
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
{
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
));
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user