mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
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:
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -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";
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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"
|
||||
},
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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())]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
)));
|
||||
}
|
||||
@@ -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",
|
||||
|
||||
@@ -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"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user