Add Guardian internal session support (#39994)

## What changed

- Add an extension API for spawning host-owned internal sessions and a
  `ThreadReadyInput` lifecycle callback that runs after thread registration.
- Add Guardian reviewer session scaffolding that records the parent thread and
  effective model, plus the under-development `guardian_ext` feature flag.
- Start internal sessions with fresh history while preserving parent lineage,
  shared session controls, and internal-thread visibility rules.
- Scope internal-session prompt cache keys to their source and parent thread,
  and expose `guardian` as an internal session source.

## Testing

- Cover internal-session spawning, parent metadata, history isolation, prompt
  cache keys, and extension spawner argument forwarding.

GitOrigin-RevId: 682dae80397d62bb36247796b447042e760ca364
This commit is contained in:
Won Park
2026-08-21 19:11:05 +00:00
committed by copyberry
parent 79b7606803
commit 56012fafb8
22 changed files with 402 additions and 10 deletions
@@ -2,4 +2,4 @@
// 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";
export type InternalSessionSource = "memory_consolidation" | "guardian";
+21
View File
@@ -19,6 +19,8 @@ use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionRegistry;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::ExtensionWarning;
use codex_extension_api::InternalSessionSpawnFuture;
use codex_extension_api::InternalSessionSpawner;
use codex_goal_extension::GoalExtensionConfig;
use codex_goal_extension::GoalService;
use codex_http_client::HttpClientFactory;
@@ -98,6 +100,7 @@ where
codex_guardian_v2::install(
&mut builder,
guardian_agent_spawner,
internal_session_spawner(thread_manager.clone()),
auth_manager.clone(),
thread_manager,
);
@@ -321,6 +324,24 @@ pub(crate) fn guardian_agent_spawner(
}
}
fn internal_session_spawner(
thread_manager: Weak<ThreadManager>,
) -> impl InternalSessionSpawner<StartThreadOptions, Spawned = NewThread, Error = CodexErr> {
move |parent_thread_id: ThreadId,
options: StartThreadOptions|
-> InternalSessionSpawnFuture<'static, NewThread, CodexErr> {
let thread_manager = thread_manager.clone();
Box::pin(async move {
let thread_manager = thread_manager.upgrade().ok_or_else(|| {
CodexErr::UnsupportedOperation("thread manager dropped".to_string())
})?;
thread_manager
.spawn_internal_session(parent_thread_id, options)
.await
})
}
}
#[cfg(test)]
mod tests {
use codex_protocol::protocol::ThreadGoal as CoreThreadGoal;
@@ -674,6 +674,7 @@ fn source_category(source: &str) -> &'static str {
SessionSource::Internal(InternalSessionSource::MemoryConsolidation) => {
"internal:memory_consolidation"
}
SessionSource::Internal(InternalSessionSource::Guardian) => "internal:guardian",
SessionSource::SubAgent(SubAgentSource::Review) => "subagent:review",
SessionSource::SubAgent(SubAgentSource::Compact) => "subagent:compact",
SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) => "subagent:thread_spawn",
+6
View File
@@ -601,6 +601,9 @@
"guardian_enhanced_node_repl_transcripts": {
"type": "boolean"
},
"guardian_ext": {
"type": "boolean"
},
"guardian_node_repl_transcript_images": {
"type": "boolean"
},
@@ -5562,6 +5565,9 @@
"guardian_enhanced_node_repl_transcripts": {
"type": "boolean"
},
"guardian_ext": {
"type": "boolean"
},
"guardian_node_repl_transcript_images": {
"type": "boolean"
},
+11 -3
View File
@@ -483,9 +483,17 @@ impl ModelClient {
}
fn prompt_cache_key(&self, responses_metadata: &CodexResponsesMetadata) -> String {
self.prompt_cache_key_override
.clone()
.unwrap_or_else(|| responses_metadata.session_id.clone())
if let Some(prompt_cache_key) = &self.prompt_cache_key_override {
return prompt_cache_key.clone();
}
if let SessionSource::Internal(source) = &self.state.session_source
&& let Some(parent_thread_id) = responses_metadata.parent_thread_id
{
return format!("{source}:{parent_thread_id}");
}
responses_metadata.session_id.clone()
}
/// Creates a fresh turn-scoped streaming session.
+18
View File
@@ -481,6 +481,24 @@ fn build_subagent_headers_sets_other_subagent_label() {
assert_eq!(value, Some("memory_consolidation"));
}
#[test]
fn internal_session_prompt_cache_key_is_scoped_to_parent_thread() {
let parent_thread_id = ThreadId::new();
let client = test_model_client(SessionSource::Internal(InternalSessionSource::Guardian));
let metadata = test_responses_metadata_for_client(
&client,
Some("turn-123"),
"window-1".to_string(),
Some(parent_thread_id),
TestCodexResponsesRequestKind::Turn,
);
assert_eq!(
client.prompt_cache_key(&metadata),
format!("guardian:{parent_thread_id}")
);
}
#[test]
fn build_subagent_headers_sets_internal_memory_consolidation_label() {
let client = test_model_client(SessionSource::Internal(
+19
View File
@@ -232,6 +232,25 @@ impl CodexThread {
self.io.session_loop_termination.clone().await;
}
pub(crate) async fn emit_thread_ready_lifecycle(&self) {
let config = self.config().await;
for contributor in self
.session
.services
.extensions
.thread_lifecycle_contributors()
{
contributor
.on_thread_ready(codex_extension_api::ThreadReadyInput {
config: config.as_ref(),
session_source: &self.session_source,
session_store: &self.session.services.session_extension_data,
thread_store: &self.session.services.thread_extension_data,
})
.await;
}
}
pub(crate) async fn emit_thread_resume_lifecycle(&self) {
for contributor in self
.session
+1 -4
View File
@@ -7,7 +7,6 @@ use codex_analytics::CompactionReason;
use codex_analytics::CompactionStrategy;
use codex_analytics::CompactionTrigger;
use codex_protocol::ThreadId;
use codex_protocol::protocol::InternalSessionSource;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::ThreadSource;
@@ -406,9 +405,7 @@ pub(crate) fn subagent_header_value(session_source: &SessionSource) -> Option<St
SubAgentSource::ThreadSpawn { .. } => Some("collab_spawn".to_string()),
SubAgentSource::Other(label) => Some(label.clone()),
},
SessionSource::Internal(InternalSessionSource::MemoryConsolidation) => {
Some("memory_consolidation".to_string())
}
SessionSource::Internal(source) => Some(source.to_string()),
SessionSource::Cli
| SessionSource::VSCode
| SessionSource::Exec
+23
View File
@@ -906,6 +906,28 @@ impl ThreadManager {
Box::pin(self.start_thread_inner(options, /*forked_from_thread_id*/ None)).await
}
/// Starts a fresh internal session associated with an existing parent thread.
pub async fn spawn_internal_session(
&self,
parent_thread_id: ThreadId,
mut options: StartThreadOptions,
) -> CodexResult<NewThread> {
if !matches!(options.session_source, Some(SessionSource::Internal(_))) {
return Err(CodexErr::InvalidRequest(
"internal sessions require an internal session source".to_string(),
));
}
let parent = self.get_thread(parent_thread_id).await?;
options.initial_history = InitialHistory::New;
let mut request = ThreadSpawnRequest::new(
options,
Arc::clone(&parent.session.services.auth_manager),
parent.session.services.agent_control.clone(),
);
request.parent_thread_id = Some(parent_thread_id);
Box::pin(self.state.spawn_thread(request)).await
}
/// Allocates a thread ID before startup so a caller can associate host-owned state with it.
pub fn reserve_thread_id(&self) -> ThreadId {
self.state.thread_id_generator.as_ref()()
@@ -1912,6 +1934,7 @@ impl ThreadManagerState {
let new_thread = self
.finalize_thread_spawn(session, io, tracked_session_source)
.await?;
new_thread.thread.emit_thread_ready_lifecycle().await;
if source_changed_during_startup.load(Ordering::Acquire) {
new_thread.thread.session.request_mcp_runtime_refresh();
}
+81
View File
@@ -874,6 +874,87 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() {
assert!(manager.list_thread_ids().await.is_empty());
}
#[tokio::test]
async fn spawn_internal_session_preserves_parent_lineage_without_forking_history() {
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 parent = manager
.start_thread(StartThreadOptions {
metrics_service_name: Some("codex_work_desktop".to_string()),
..StartThreadOptions::new(config.clone())
})
.await
.expect("start parent thread");
let reviewer = manager
.spawn_internal_session(
parent.thread_id,
StartThreadOptions {
session_source: Some(SessionSource::Internal(InternalSessionSource::Guardian)),
initial_history: InitialHistory::Forked(vec![RolloutItem::ResponseItem(
user_msg("parent history must not be inherited").into(),
)]),
environments: Some(Vec::new()),
..StartThreadOptions::new(config)
},
)
.await
.expect("start internal reviewer");
let reviewer_config = reviewer.thread.config_snapshot().await;
assert_eq!(
reviewer.session_configured.session_id,
parent.session_configured.session_id
);
assert!(std::ptr::eq(
reviewer
.thread
.session
.services
.agent_control
.rollout_budget(),
parent
.thread
.session
.services
.agent_control
.rollout_budget(),
));
assert_eq!(reviewer_config.parent_thread_id, Some(parent.thread_id));
assert_eq!(reviewer_config.forked_from_thread_id, None);
assert_eq!(reviewer_config.originator, "codex_work_desktop");
assert_eq!(
reviewer.session_configured.parent_thread_id,
Some(parent.thread_id)
);
assert_eq!(reviewer.session_configured.forked_from_id, None);
assert_eq!(manager.list_thread_ids().await, vec![parent.thread_id]);
assert!(manager.get_thread(reviewer.thread_id).await.is_err());
assert!(
reviewer
.thread
.session
.clone_history()
.await
.raw_items()
.next()
.is_none()
);
manager
.shutdown_all_threads_bounded(Duration::from_secs(10))
.await;
}
#[tokio::test]
async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() {
struct InitialDataRecorder {
@@ -0,0 +1,38 @@
use codex_protocol::ThreadId;
use super::agent::AgentSpawnFuture;
/// Future returned by one host-owned internal-session spawning request.
pub type InternalSessionSpawnFuture<'a, T, E> = AgentSpawnFuture<'a, T, E>;
/// Constructor-injected host helper for extensions that need private internal sessions.
///
/// The extension owns the request shape and resulting handle types. The host
/// provides the implementation when it constructs the extension.
pub trait InternalSessionSpawner<R>: Send + Sync {
type Spawned;
type Error;
/// Starts a fresh host-owned internal session associated with its parent.
fn spawn_internal_session<'a>(
&'a self,
parent_thread_id: ThreadId,
request: R,
) -> InternalSessionSpawnFuture<'a, Self::Spawned, Self::Error>;
}
impl<R, S, E, F> InternalSessionSpawner<R> for F
where
F: Fn(ThreadId, R) -> InternalSessionSpawnFuture<'static, S, E> + Send + Sync,
{
type Spawned = S;
type Error = E;
fn spawn_internal_session<'a>(
&'a self,
parent_thread_id: ThreadId,
request: R,
) -> InternalSessionSpawnFuture<'a, Self::Spawned, Self::Error> {
self(parent_thread_id, request)
}
}
@@ -1,6 +1,7 @@
mod agent;
mod conversation_history;
mod events;
mod internal_session;
mod metrics;
mod response_items;
@@ -10,6 +11,8 @@ pub use conversation_history::ConversationHistorySnapshot;
pub use events::ExtensionEventSink;
pub use events::ExtensionWarning;
pub use events::NoopExtensionEventSink;
pub use internal_session::InternalSessionSpawnFuture;
pub use internal_session::InternalSessionSpawner;
pub use metrics::ExtensionMetrics;
pub use response_items::NoopResponseItemInjector;
pub use response_items::ResponseItemInjectionFuture;
@@ -34,6 +34,7 @@ pub use skill_invocation::SkillInvocationKind;
pub use thread_lifecycle::ThreadIdleCause;
pub use thread_lifecycle::ThreadIdleInput;
pub use thread_lifecycle::ThreadOriginator;
pub use thread_lifecycle::ThreadReadyInput;
pub use thread_lifecycle::ThreadResumeInput;
pub use thread_lifecycle::ThreadStartInput;
pub use thread_lifecycle::ThreadStopInput;
@@ -133,6 +134,14 @@ pub trait ThreadLifecycleContributor<C: Sync>: Send + Sync {
})
}
/// Called after the initialized thread is registered with its host.
fn on_thread_ready<'a>(&'a self, input: ThreadReadyInput<'a, C>) -> ExtensionFuture<'a, ()> {
Box::pin(async move {
let _self = self;
let _input = input;
})
}
/// Called after the host constructs a runtime from persisted history.
fn on_thread_resume<'a>(&'a self, input: ThreadResumeInput<'a>) -> ExtensionFuture<'a, ()> {
Box::pin(async move {
@@ -34,6 +34,18 @@ pub struct ThreadStartInput<'a, C> {
pub thread_store: &'a ExtensionData,
}
/// Input supplied after the host has registered a fully initialized thread.
pub struct ThreadReadyInput<'a, C> {
/// Host configuration visible after thread registration.
pub config: &'a C,
/// Source that created the session for this thread.
pub session_source: &'a SessionSource,
/// Store scoped to the host session runtime.
pub session_store: &'a ExtensionData,
/// Store scoped to this thread runtime.
pub thread_store: &'a ExtensionData,
}
/// Input supplied when the host resumes an existing thread.
pub struct ThreadResumeInput<'a> {
/// Store scoped to the host session runtime.
+3
View File
@@ -10,6 +10,8 @@ pub use capabilities::ConversationHistorySnapshot;
pub use capabilities::ExtensionEventSink;
pub use capabilities::ExtensionMetrics;
pub use capabilities::ExtensionWarning;
pub use capabilities::InternalSessionSpawnFuture;
pub use capabilities::InternalSessionSpawner;
pub use capabilities::NoopExtensionEventSink;
pub use capabilities::NoopResponseItemInjector;
pub use capabilities::ResponseItemInjectionFuture;
@@ -55,6 +57,7 @@ pub use contributors::ThreadIdleCause;
pub use contributors::ThreadIdleInput;
pub use contributors::ThreadLifecycleContributor;
pub use contributors::ThreadOriginator;
pub use contributors::ThreadReadyInput;
pub use contributors::ThreadResumeInput;
pub use contributors::ThreadStartInput;
pub use contributors::ThreadStopInput;
@@ -3,6 +3,8 @@ use std::sync::Mutex;
use codex_extension_api::AgentSpawnFuture;
use codex_extension_api::AgentSpawner;
use codex_extension_api::InternalSessionSpawnFuture;
use codex_extension_api::InternalSessionSpawner;
use codex_extension_api::NoopResponseItemInjector;
use codex_extension_api::ResponseItemInjector;
use codex_protocol::ThreadId;
@@ -54,3 +56,30 @@ async fn closure_agent_spawner_forwards_arguments_and_result() {
[(thread_id, "delegate this".to_string())]
);
}
#[tokio::test]
async fn closure_internal_session_spawner_forwards_arguments_and_result() {
let calls = Arc::new(Mutex::new(Vec::new()));
let recorded_calls = Arc::clone(&calls);
let spawner = move |thread_id: ThreadId,
request: String|
-> InternalSessionSpawnFuture<'static, usize, &'static str> {
recorded_calls
.lock()
.expect("agent spawn calls lock")
.push((thread_id, request.clone()));
Box::pin(async move { Ok(request.len()) })
};
let thread_id =
ThreadId::from_string("11111111-1111-4111-8111-111111111111").expect("valid thread id");
let spawned = spawner
.spawn_internal_session(thread_id, "delegate this".to_string())
.await;
assert_eq!(spawned, Ok(13));
assert_eq!(
calls.lock().expect("agent spawn calls lock").as_slice(),
[(thread_id, "delegate this".to_string())]
);
}
+8 -2
View File
@@ -13,8 +13,11 @@ use codex_login::AuthManager;
use codex_protocol::ThreadId;
mod async_scorer;
mod sync_reviewer;
pub use async_scorer::StrictReviewReason;
pub use sync_reviewer::GuardianExtension as GuardianReviewerExtension;
pub use sync_reviewer::GuardianThreadContext as GuardianReviewerThreadContext;
/// Guardian extension dependencies supplied by the host at construction time.
#[derive(Clone, Debug)]
@@ -76,14 +79,17 @@ where
}
/// Installs the guardian contributors into the extension registry.
pub fn install<S>(
pub fn install<S, I>(
registry: &mut ExtensionRegistryBuilder<Config>,
agent_spawner: S,
internal_session_spawner: I,
auth_manager: Arc<AuthManager>,
thread_manager: Weak<ThreadManager>,
) where
S: Send + Sync + 'static,
I: Send + Sync + 'static,
{
registry.thread_lifecycle_contributor(Arc::new(GuardianExtension::new(agent_spawner)));
async_scorer::install(registry, auth_manager, thread_manager);
async_scorer::install(registry, auth_manager, thread_manager.clone());
sync_reviewer::install(registry, thread_manager, internal_session_spawner);
}
@@ -0,0 +1,108 @@
use std::sync::Arc;
use std::sync::Weak;
use codex_core::ThreadManager;
use codex_core::config::Config;
use codex_extension_api::ExtensionFuture;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::InternalSessionSpawnFuture;
use codex_extension_api::InternalSessionSpawner;
use codex_extension_api::ThreadLifecycleContributor;
use codex_extension_api::ThreadReadyInput;
use codex_protocol::ThreadId;
/// Guardian extension dependencies supplied by the host at construction time.
#[derive(Clone, Debug)]
pub struct GuardianExtension<S> {
thread_manager: Weak<ThreadManager>,
internal_session_spawner: S,
}
impl<S> GuardianExtension<S> {
/// Creates a guardian extension with its host-owned internal-session spawner.
pub fn new(thread_manager: Weak<ThreadManager>, internal_session_spawner: S) -> Self {
Self {
thread_manager,
internal_session_spawner,
}
}
/// Delegates a fresh internal-session request to the host helper.
pub fn spawn_internal_session<'a, R>(
&'a self,
parent_thread_id: ThreadId,
request: R,
) -> InternalSessionSpawnFuture<
'a,
<S as InternalSessionSpawner<R>>::Spawned,
<S as InternalSessionSpawner<R>>::Error,
>
where
S: InternalSessionSpawner<R>,
{
self.internal_session_spawner
.spawn_internal_session(parent_thread_id, request)
}
}
/// Thread-local guardian state captured after the host registers a thread.
#[derive(Clone, Debug)]
pub struct GuardianThreadContext {
parent_thread_id: ThreadId,
parent_model: String,
}
impl GuardianThreadContext {
/// Returns the parent thread associated with future Guardian reviewer sessions.
pub fn parent_thread_id(&self) -> ThreadId {
self.parent_thread_id
}
/// Returns the parent's effective model after provider fallback.
pub fn parent_model(&self) -> &str {
&self.parent_model
}
}
impl<S> ThreadLifecycleContributor<Config> for GuardianExtension<S>
where
S: Send + Sync,
{
fn on_thread_ready<'a>(
&'a self,
input: ThreadReadyInput<'a, Config>,
) -> ExtensionFuture<'a, ()> {
Box::pin(async move {
if input.session_source.is_internal() {
return;
}
let Ok(parent_thread_id) = ThreadId::from_string(input.thread_store.level_id()) else {
return;
};
let Some(thread_manager) = self.thread_manager.upgrade() else {
return;
};
let Ok(parent) = thread_manager.get_thread(parent_thread_id).await else {
return;
};
input.thread_store.insert(GuardianThreadContext {
parent_thread_id,
parent_model: parent.config_snapshot().await.model,
});
})
}
}
/// Installs the guardian contributors into the extension registry.
pub fn install<S>(
registry: &mut ExtensionRegistryBuilder<Config>,
thread_manager: Weak<ThreadManager>,
internal_session_spawner: S,
) where
S: Send + Sync + 'static,
{
registry.thread_lifecycle_contributor(Arc::new(GuardianExtension::new(
thread_manager,
internal_session_spawner,
)));
}
+8
View File
@@ -279,6 +279,8 @@ pub enum Feature {
GuardianNodeReplTranscriptImages,
/// Enable Guardian V2 automatic approval reviews.
GuardianV2,
/// Enable the extension-owned synchronous Guardian reviewer.
GuardianExt,
/// Enable persisted thread goals and automatic goal continuation.
Goals,
/// Add current context-window metadata to model-visible context.
@@ -1411,6 +1413,12 @@ pub const FEATURES: &[FeatureSpec] = &[
stage: Stage::UnderDevelopment,
default_enabled: false,
},
FeatureSpec {
id: Feature::GuardianExt,
key: "guardian_ext",
stage: Stage::UnderDevelopment,
default_enabled: false,
},
FeatureSpec {
id: Feature::Goals,
key: "goals",
+2
View File
@@ -2650,6 +2650,7 @@ impl FromStr for ThreadSource {
#[ts(rename_all = "snake_case")]
pub enum InternalSessionSource {
MemoryConsolidation,
Guardian,
}
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS)]
@@ -2822,6 +2823,7 @@ impl fmt::Display for InternalSessionSource {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
InternalSessionSource::MemoryConsolidation => f.write_str("memory_consolidation"),
InternalSessionSource::Guardian => f.write_str("guardian"),
}
}
}