Activate prompt hooks

This commit is contained in:
Abhinav Vedmala
2026-06-03 17:51:34 -07:00
parent d546866b45
commit c621ca2c2f
13 changed files with 556 additions and 51 deletions
+22
View File
@@ -1057,6 +1057,27 @@
},
{
"properties": {
"continueOnBlock": {
"default": false,
"type": "boolean"
},
"model": {
"default": null,
"type": "string"
},
"prompt": {
"type": "string"
},
"statusMessage": {
"default": null,
"type": "string"
},
"timeout": {
"default": null,
"format": "uint64",
"minimum": 0.0,
"type": "integer"
},
"type": {
"enum": [
"prompt"
@@ -1065,6 +1086,7 @@
}
},
"required": [
"prompt",
"type"
],
"type": "object"
+18 -3
View File
@@ -238,6 +238,7 @@ pub struct ModelClient {
pub struct ModelClientSession {
client: ModelClient,
websocket_session: WebsocketSession,
cache_websocket_session_on_drop: bool,
/// Turn state for sticky routing.
///
/// This is an `OnceLock` that stores the turn state value received from the server
@@ -382,6 +383,18 @@ impl ModelClient {
ModelClientSession {
client: self.clone(),
websocket_session: self.take_cached_websocket_session(),
cache_websocket_session_on_drop: true,
turn_state: Arc::new(OnceLock::new()),
}
}
/// Creates a fresh streaming session that does not read or write the session's cached
/// WebSocket state.
pub(crate) fn new_isolated_session(&self) -> ModelClientSession {
ModelClientSession {
client: self.clone(),
websocket_session: WebsocketSession::default(),
cache_websocket_session_on_drop: false,
turn_state: Arc::new(OnceLock::new()),
}
}
@@ -951,9 +964,11 @@ impl ModelClient {
impl Drop for ModelClientSession {
fn drop(&mut self) {
let websocket_session = std::mem::take(&mut self.websocket_session);
self.client
.store_cached_websocket_session(websocket_session);
if self.cache_websocket_session_on_drop {
let websocket_session = std::mem::take(&mut self.websocket_session);
self.client
.store_cached_websocket_session(websocket_session);
}
}
}
+19
View File
@@ -2,6 +2,7 @@ use super::AuthRequestTelemetryContext;
use super::ModelClient;
use super::PendingUnauthorizedRetry;
use super::UnauthorizedRecoveryExecution;
use super::WebsocketSession;
use super::X_CODEX_INSTALLATION_ID_HEADER;
use super::X_CODEX_PARENT_THREAD_ID_HEADER;
use super::X_CODEX_TURN_METADATA_HEADER;
@@ -116,6 +117,24 @@ fn test_model_info() -> ModelInfo {
.expect("deserialize test model info")
}
#[test]
fn isolated_session_does_not_touch_cached_websocket_state() {
let model_client = test_model_client(SessionSource::Cli);
let cached_websocket_session = WebsocketSession::default();
cached_websocket_session.set_connection_reused(/*connection_reused*/ true);
model_client.store_cached_websocket_session(cached_websocket_session);
let isolated_session = model_client.new_isolated_session();
assert!(!isolated_session.websocket_session.connection_reused());
drop(isolated_session);
assert!(
model_client
.take_cached_websocket_session()
.connection_reused()
);
}
fn test_session_telemetry() -> SessionTelemetry {
SessionTelemetry::new(
ThreadId::new(),
+160
View File
@@ -0,0 +1,160 @@
use std::sync::Arc;
use codex_hooks::PromptHookRequest;
use codex_hooks::PromptHookRunner;
use codex_models_manager::ModelsManagerConfig;
use codex_models_manager::manager::SharedModelsManager;
use codex_otel::SessionTelemetry;
use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_rollout_trace::InferenceTraceContext;
use futures::StreamExt;
use serde_json::json;
use crate::client::ModelClient;
use crate::client_common::Prompt;
use crate::client_common::ResponseEvent;
use crate::config::Config;
use crate::stream_events_utils::raw_assistant_output_text_from_item;
const PROMPT_HOOK_BASE_INSTRUCTIONS: &str = r#"You evaluate a Codex prompt hook.
The user message contains:
1. The hook author's instructions.
2. The hook input JSON.
Decide whether the hook input satisfies the hook author's instructions.
Return only JSON:
{"ok": true}
or
{"ok": false, "reason": "concise actionable reason"}
Use ok:false only when the hook criteria fail. Do not answer the user's task. Do not include Markdown or extra text."#;
pub(crate) fn build_prompt_hook_runner(
model_client: ModelClient,
models_manager: SharedModelsManager,
config: &Config,
session_telemetry: SessionTelemetry,
service_tier: Option<String>,
) -> PromptHookRunner {
let models_manager_config = config.to_models_manager_config();
PromptHookRunner::new(move |request| {
let model_client = model_client.clone();
let models_manager = Arc::clone(&models_manager);
let models_manager_config = models_manager_config.clone();
let session_telemetry = session_telemetry.clone();
let service_tier = service_tier.clone();
async move {
run_prompt_hook(
model_client,
models_manager,
models_manager_config,
session_telemetry,
service_tier,
request,
)
.await
}
})
}
/// Run the hook as an isolated model request: resolve the configured model,
/// send the rendered hook prompt as the only user message, disable all
/// tools/personality/reasoning summaries, and constrain the final answer to the
/// small `{ ok, reason }` contract that `codex-hooks` maps back into existing
/// hook output semantics.
async fn run_prompt_hook(
model_client: ModelClient,
models_manager: SharedModelsManager,
models_manager_config: ModelsManagerConfig,
session_telemetry: SessionTelemetry,
service_tier: Option<String>,
request: PromptHookRequest,
) -> anyhow::Result<String> {
let model_info = models_manager
.get_model_info(request.model.as_str(), &models_manager_config)
.await;
let prompt = Prompt {
input: vec![ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: request.prompt,
}],
phase: None,
}],
tools: Vec::new(),
parallel_tool_calls: false,
base_instructions: BaseInstructions {
text: PROMPT_HOOK_BASE_INSTRUCTIONS.to_string(),
},
personality: None,
output_schema: Some(prompt_hook_output_schema()),
output_schema_strict: false,
};
let disabled_trace = InferenceTraceContext::disabled();
let mut client_session = model_client.new_isolated_session();
let mut stream = client_session
.stream(
&prompt,
&model_info,
&session_telemetry,
/*effort*/ None,
ReasoningSummaryConfig::None,
service_tier,
/*turn_metadata_header*/ None,
&disabled_trace,
)
.await?;
let mut delta_text = String::new();
let mut item_texts = Vec::new();
while let Some(event) = stream.next().await {
match event? {
ResponseEvent::OutputItemDone(item) => {
if let Some(text) = raw_assistant_output_text_from_item(&item) {
item_texts.push(text);
}
}
ResponseEvent::OutputTextDelta(delta) => delta_text.push_str(delta.as_str()),
ResponseEvent::Completed { .. } => break,
ResponseEvent::Created
| ResponseEvent::OutputItemAdded(_)
| ResponseEvent::ServerModel(_)
| ResponseEvent::ModelVerifications(_)
| ResponseEvent::ServerReasoningIncluded(_)
| ResponseEvent::ToolCallInputDelta { .. }
| ResponseEvent::ReasoningSummaryDelta { .. }
| ResponseEvent::ReasoningContentDelta { .. }
| ResponseEvent::ReasoningSummaryPartAdded { .. }
| ResponseEvent::RateLimits(_)
| ResponseEvent::ModelsEtag(_) => {}
}
}
if item_texts.is_empty() {
Ok(delta_text)
} else {
Ok(item_texts.join(""))
}
}
fn prompt_hook_output_schema() -> serde_json::Value {
json!({
"type": "object",
"properties": {
"ok": {
"type": "boolean"
},
"reason": {
"type": "string"
}
},
"required": ["ok"],
"additionalProperties": false
})
}
+1
View File
@@ -41,6 +41,7 @@ mod goals;
pub use goals::ExternalGoalPreviousStatus;
pub use goals::ExternalGoalSet;
mod guardian;
mod hook_prompt;
mod hook_runtime;
mod installation_id;
pub(crate) mod landlock;
+14 -3
View File
@@ -58,6 +58,7 @@ use codex_features::Feature;
use codex_features::unstable_features_warning_event;
use codex_hooks::Hooks;
use codex_hooks::HooksConfig;
use codex_hooks::PromptHookRunner;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::auth_env_telemetry::collect_auth_env_telemetry;
@@ -1482,10 +1483,11 @@ impl Session {
// layers such as request/session overrides that were present when this session
// was created.
let notify_config_contributors = !self.services.extensions.config_contributors().is_empty();
let (previous_config, new_config, config) = {
let (previous_config, new_config, config, service_tier) = {
let mut state = self.state.lock().await;
let previous_config = notify_config_contributors
.then(|| Self::build_effective_session_config(&state.session_configuration));
let service_tier = state.session_configuration.service_tier.clone();
let mut config = (*state.session_configuration.original_config_do_not_use).clone();
config.config_layer_stack = config
.config_layer_stack
@@ -1496,15 +1498,23 @@ impl Session {
state.session_configuration.original_config_do_not_use = Arc::clone(&config);
let new_config = notify_config_contributors
.then(|| Self::build_effective_session_config(&state.session_configuration));
(previous_config, new_config, config)
(previous_config, new_config, config, service_tier)
};
self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref());
self.services.skills_manager.clear_cache();
self.services.plugins_manager.clear_cache();
let prompt_hook_runner = crate::hook_prompt::build_prompt_hook_runner(
self.services.model_client.clone(),
Arc::clone(&self.services.models_manager),
config.as_ref(),
self.services.session_telemetry.clone(),
service_tier,
);
let hooks = build_hooks_for_config(
config.as_ref(),
self.services.plugins_manager.as_ref(),
self.services.user_shell.as_ref(),
Some(prompt_hook_runner),
)
.await;
@@ -3382,6 +3392,7 @@ async fn build_hooks_for_config(
config: &Config,
plugins_manager: &PluginsManager,
user_shell: &crate::shell::Shell,
prompt_hook_runner: Option<PromptHookRunner>,
) -> Hooks {
let mut hook_shell_argv = user_shell.derive_exec_args("", /*use_login_shell*/ false);
let hook_shell_program = hook_shell_argv.remove(0);
@@ -3399,7 +3410,7 @@ async fn build_hooks_for_config(
plugin_hook_load_warnings,
shell_program: Some(hook_shell_program),
shell_args: hook_shell_argv,
prompt_hook_runner: None,
prompt_hook_runner,
})
}
+44 -34
View File
@@ -942,17 +942,6 @@ impl Session {
(None, None)
};
let hooks =
build_hooks_for_config(&config, plugins_manager.as_ref(), &default_shell).await;
for warning in hooks.startup_warnings() {
post_session_configured_events.push(Event {
id: INITIAL_SUBMIT_ID.to_owned(),
msg: EventMsg::Warning(WarningEvent {
message: warning.clone(),
}),
});
}
let analytics_events_client = analytics_events_client.unwrap_or_else(|| {
AnalyticsEventsClient::new(
Arc::clone(&auth_manager),
@@ -966,6 +955,49 @@ impl Session {
SessionId::from(thread_id)
};
let agent_control = agent_control.with_session_id(session_id);
let model_client = ModelClient::new(
Some(Arc::clone(&auth_manager)),
session_id,
thread_id,
installation_id.clone(),
session_configuration.provider.clone(),
session_configuration.session_source.clone(),
session_configuration.parent_thread_id,
config.model_verbosity,
config.features.enabled(Feature::EnableRequestCompression),
config.features.enabled(Feature::RuntimeMetrics),
Self::build_model_client_beta_features_header(config.as_ref()),
attestation_provider.clone(),
)
.with_prompt_cache_key_override(
crate::guardian::prompt_cache_key_override_for_review_session(
&session_configuration.session_source,
session_configuration.parent_thread_id,
),
);
model_client.set_window_generation(window_generation);
let prompt_hook_runner = crate::hook_prompt::build_prompt_hook_runner(
model_client.clone(),
Arc::clone(&models_manager),
config.as_ref(),
session_telemetry.clone(),
session_configuration.service_tier.clone(),
);
let hooks = build_hooks_for_config(
&config,
plugins_manager.as_ref(),
&default_shell,
Some(prompt_hook_runner),
)
.await;
for warning in hooks.startup_warnings() {
post_session_configured_events.push(Event {
id: INITIAL_SUBMIT_ID.to_owned(),
msg: EventMsg::Warning(WarningEvent {
message: warning.clone(),
}),
});
}
let session_extension_data =
codex_extension_api::ExtensionData::new(session_id.to_string());
let thread_extension_data =
@@ -1031,32 +1063,10 @@ impl Session {
live_thread: live_thread_init.as_ref().cloned(),
thread_store: Arc::clone(&thread_store),
attestation_provider: attestation_provider.clone(),
model_client: ModelClient::new(
Some(Arc::clone(&auth_manager)),
session_id,
thread_id,
installation_id.clone(),
session_configuration.provider.clone(),
session_configuration.session_source.clone(),
session_configuration.parent_thread_id,
config.model_verbosity,
config.features.enabled(Feature::EnableRequestCompression),
config.features.enabled(Feature::RuntimeMetrics),
Self::build_model_client_beta_features_header(config.as_ref()),
attestation_provider,
)
.with_prompt_cache_key_override(
crate::guardian::prompt_cache_key_override_for_review_session(
&session_configuration.session_source,
session_configuration.parent_thread_id,
),
),
model_client,
code_mode_service: crate::tools::code_mode::CodeModeService::new(),
environment_manager,
};
services
.model_client
.set_window_generation(window_generation);
let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) =
watch::channel(false);
+80
View File
@@ -47,6 +47,7 @@ use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use std::sync::Arc;
use std::time::Duration;
use tempfile::TempDir;
@@ -242,6 +243,23 @@ if payload.get("prompt") == {blocked_prompt_json}:
Ok(())
}
fn write_user_prompt_submit_prompt_hook(home: &Path) -> Result<()> {
let hooks = serde_json::json!({
"hooks": {
"UserPromptSubmit": [{
"hooks": [{
"type": "prompt",
"prompt": "Reject prompts that mention secrets: $ARGUMENTS",
"statusMessage": "checking prompt",
}]
}]
}
});
fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?;
Ok(())
}
fn write_session_start_and_user_prompt_submit_order_hooks(home: &Path) -> Result<()> {
let session_start_script_path = home.join("session_start_order_hook.py");
let user_prompt_submit_script_path = home.join("user_prompt_submit_order_hook.py");
@@ -1713,6 +1731,68 @@ async fn multiple_blocking_stop_hooks_persist_multiple_hook_prompt_fragments() -
Ok(())
}
#[tokio::test]
async fn user_prompt_submit_prompt_hook_runs_isolated_request() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let response = mount_sse_once(
&server,
sse(vec![
ev_response_created("prompt-hook-resp"),
ev_assistant_message(
"prompt-hook-msg",
r#"{"ok":false,"reason":"do not mention secrets"}"#,
),
ev_completed("prompt-hook-resp"),
]),
)
.await;
let mut builder = test_codex()
.with_model("gpt-5.4")
.with_pre_build_hook(|home| {
if let Err(error) = write_user_prompt_submit_prompt_hook(home) {
panic!("failed to write user prompt submit prompt hook fixture: {error}");
}
})
.with_config(trust_discovered_hooks);
let test = builder.build(&server).await?;
let prompt = "secret launch codes";
test.submit_turn(prompt).await?;
let request = response.single_request();
let body = request.body_json();
assert_eq!(body["model"], json!("gpt-5.4"));
assert_eq!(body["tools"], json!([]));
assert!(
request
.instructions_text()
.starts_with("You evaluate a Codex prompt hook."),
"prompt hook request should use isolated hook instructions",
);
let user_inputs = request.message_input_texts("user");
assert_eq!(
user_inputs.len(),
1,
"prompt hook request should contain exactly one user message",
);
assert_eq!(
request.message_input_texts("developer"),
Vec::<String>::new(),
"prompt hook request should not inherit developer context",
);
let hook_input = user_inputs.first().context("prompt hook user input")?;
assert!(hook_input.contains("Reject prompts that mention secrets:"));
assert!(
hook_input.contains(prompt),
"prompt hook request should include the submitted prompt",
);
Ok(())
}
#[tokio::test]
async fn blocked_user_prompt_submit_persists_additional_context_for_next_turn() -> Result<()> {
skip_if_no_network!(Ok(()));