From b8602087bd19f9a67478632037050aa1c1a5fa16 Mon Sep 17 00:00:00 2001 From: Lawrence Chen Date: Sun, 3 May 2026 23:44:24 -0700 Subject: [PATCH] Add SubscriptionExhausted hook for auth rotation --- codex-rs/analytics/src/events.rs | 1 + .../schema/json/ServerNotification.json | 3 +- .../codex_app_server_protocol.schemas.json | 10 +- .../codex_app_server_protocol.v2.schemas.json | 10 +- .../v2/ConfigRequirementsReadResponse.json | 7 + .../json/v2/HookCompletedNotification.json | 3 +- .../json/v2/HookStartedNotification.json | 3 +- .../schema/json/v2/HooksListResponse.json | 3 +- .../schema/typescript/v2/HookEventName.ts | 2 +- .../typescript/v2/ManagedHooksRequirements.ts | 2 +- .../app-server-protocol/src/protocol/v2.rs | 5 +- codex-rs/app-server/src/config_api.rs | 3 + codex-rs/config/src/hook_config.rs | 12 +- codex-rs/core/config.schema.json | 7 + codex-rs/core/src/client.rs | 9 + codex-rs/core/src/hook_runtime.rs | 23 +- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/session/turn.rs | 29 ++ codex-rs/core/src/subscription_exhaustion.rs | 114 ++++++++ codex-rs/core/src/thread_manager.rs | 2 +- codex-rs/core/tests/common/test_codex.rs | 15 +- codex-rs/core/tests/suite/quota_exceeded.rs | 275 ++++++++++++++++++ ...mission-request.command.output.schema.json | 3 +- .../post-tool-use.command.output.schema.json | 3 +- .../pre-tool-use.command.output.schema.json | 3 +- .../session-start.command.output.schema.json | 3 +- ...iption-exhausted.command.input.schema.json | 79 +++++ ...ption-exhausted.command.output.schema.json | 24 ++ ...r-prompt-submit.command.output.schema.json | 3 +- codex-rs/hooks/src/engine/discovery.rs | 1 + codex-rs/hooks/src/engine/dispatcher.rs | 7 +- codex-rs/hooks/src/engine/mod.rs | 19 ++ codex-rs/hooks/src/engine/output_parser.rs | 13 + codex-rs/hooks/src/engine/schema_loader.rs | 20 ++ codex-rs/hooks/src/events/common.rs | 8 +- codex-rs/hooks/src/events/mod.rs | 1 + .../src/events/subscription_exhausted.rs | 195 +++++++++++++ codex-rs/hooks/src/lib.rs | 5 +- codex-rs/hooks/src/registry.rs | 16 + codex-rs/hooks/src/schema.rs | 72 +++++ codex-rs/login/src/auth/manager.rs | 4 + codex-rs/protocol/src/protocol.rs | 1 + .../tui/src/bottom_pane/hooks_browser_view.rs | 2 + codex-rs/tui/src/chatwidget/tests/helpers.rs | 1 + codex-rs/tui/src/history_cell/hook_cell.rs | 1 + 45 files changed, 999 insertions(+), 24 deletions(-) create mode 100644 codex-rs/core/src/subscription_exhaustion.rs create mode 100644 codex-rs/hooks/schema/generated/subscription-exhausted.command.input.schema.json create mode 100644 codex-rs/hooks/schema/generated/subscription-exhausted.command.output.schema.json create mode 100644 codex-rs/hooks/src/events/subscription_exhausted.rs diff --git a/codex-rs/analytics/src/events.rs b/codex-rs/analytics/src/events.rs index 8bd94402997..832a09e408e 100644 --- a/codex-rs/analytics/src/events.rs +++ b/codex-rs/analytics/src/events.rs @@ -679,6 +679,7 @@ fn analytics_hook_event_name(event_name: HookEventName) -> &'static str { HookEventName::SessionStart => "SessionStart", HookEventName::UserPromptSubmit => "UserPromptSubmit", HookEventName::Stop => "Stop", + HookEventName::SubscriptionExhausted => "SubscriptionExhausted", } } diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index 82914f3a6f2..3e00f9d539c 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -1738,7 +1738,8 @@ "postToolUse", "sessionStart", "userPromptSubmit", - "stop" + "stop", + "subscriptionExhausted" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index f856b43d660..6770606e1f1 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -9741,7 +9741,8 @@ "postToolUse", "sessionStart", "userPromptSubmit", - "stop" + "stop", + "subscriptionExhausted" ], "type": "string" }, @@ -10574,6 +10575,12 @@ }, "type": "array" }, + "SubscriptionExhausted": { + "items": { + "$ref": "#/definitions/v2/ConfiguredHookMatcherGroup" + }, + "type": "array" + }, "UserPromptSubmit": { "items": { "$ref": "#/definitions/v2/ConfiguredHookMatcherGroup" @@ -10599,6 +10606,7 @@ "PreToolUse", "SessionStart", "Stop", + "SubscriptionExhausted", "UserPromptSubmit" ], "type": "object" diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index c17efe7a453..a88cddc7fe6 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -6350,7 +6350,8 @@ "postToolUse", "sessionStart", "userPromptSubmit", - "stop" + "stop", + "subscriptionExhausted" ], "type": "string" }, @@ -7227,6 +7228,12 @@ }, "type": "array" }, + "SubscriptionExhausted": { + "items": { + "$ref": "#/definitions/ConfiguredHookMatcherGroup" + }, + "type": "array" + }, "UserPromptSubmit": { "items": { "$ref": "#/definitions/ConfiguredHookMatcherGroup" @@ -7252,6 +7259,7 @@ "PreToolUse", "SessionStart", "Stop", + "SubscriptionExhausted", "UserPromptSubmit" ], "type": "object" diff --git a/codex-rs/app-server-protocol/schema/json/v2/ConfigRequirementsReadResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ConfigRequirementsReadResponse.json index 545d8dc9b40..9678d67a424 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ConfigRequirementsReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ConfigRequirementsReadResponse.json @@ -237,6 +237,12 @@ }, "type": "array" }, + "SubscriptionExhausted": { + "items": { + "$ref": "#/definitions/ConfiguredHookMatcherGroup" + }, + "type": "array" + }, "UserPromptSubmit": { "items": { "$ref": "#/definitions/ConfiguredHookMatcherGroup" @@ -262,6 +268,7 @@ "PreToolUse", "SessionStart", "Stop", + "SubscriptionExhausted", "UserPromptSubmit" ], "type": "object" diff --git a/codex-rs/app-server-protocol/schema/json/v2/HookCompletedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/HookCompletedNotification.json index d55c059a735..173f101837f 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/HookCompletedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/HookCompletedNotification.json @@ -12,7 +12,8 @@ "postToolUse", "sessionStart", "userPromptSubmit", - "stop" + "stop", + "subscriptionExhausted" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/HookStartedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/HookStartedNotification.json index 03d2998ca5f..d25c560eb8f 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/HookStartedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/HookStartedNotification.json @@ -12,7 +12,8 @@ "postToolUse", "sessionStart", "userPromptSubmit", - "stop" + "stop", + "subscriptionExhausted" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/HooksListResponse.json b/codex-rs/app-server-protocol/schema/json/v2/HooksListResponse.json index 5190b227118..5b9227b8fc3 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/HooksListResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/HooksListResponse.json @@ -27,7 +27,8 @@ "postToolUse", "sessionStart", "userPromptSubmit", - "stop" + "stop", + "subscriptionExhausted" ], "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/HookEventName.ts b/codex-rs/app-server-protocol/schema/typescript/v2/HookEventName.ts index 28657d22821..f6ab47c3d59 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/HookEventName.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/HookEventName.ts @@ -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 HookEventName = "preToolUse" | "permissionRequest" | "postToolUse" | "sessionStart" | "userPromptSubmit" | "stop"; +export type HookEventName = "preToolUse" | "permissionRequest" | "postToolUse" | "sessionStart" | "userPromptSubmit" | "stop" | "subscriptionExhausted"; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ManagedHooksRequirements.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ManagedHooksRequirements.ts index 3386d16ec32..aa6b6ac52b3 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/ManagedHooksRequirements.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ManagedHooksRequirements.ts @@ -3,4 +3,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { ConfiguredHookMatcherGroup } from "./ConfiguredHookMatcherGroup"; -export type ManagedHooksRequirements = { managedDir: string | null, windowsManagedDir: string | null, PreToolUse: Array, PermissionRequest: Array, PostToolUse: Array, SessionStart: Array, UserPromptSubmit: Array, Stop: Array, }; +export type ManagedHooksRequirements = { managedDir: string | null, windowsManagedDir: string | null, PreToolUse: Array, PermissionRequest: Array, PostToolUse: Array, SessionStart: Array, UserPromptSubmit: Array, Stop: Array, SubscriptionExhausted: Array, }; diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index 4eb33fa8500..cb4ee263858 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -445,7 +445,7 @@ v2_enum_from_core!( v2_enum_from_core!( pub enum HookEventName from CoreHookEventName { - PreToolUse, PermissionRequest, PostToolUse, SessionStart, UserPromptSubmit, Stop + PreToolUse, PermissionRequest, PostToolUse, SessionStart, UserPromptSubmit, Stop, SubscriptionExhausted } ); @@ -988,6 +988,9 @@ pub struct ManagedHooksRequirements { #[serde(rename = "Stop")] #[ts(rename = "Stop")] pub stop: Vec, + #[serde(rename = "SubscriptionExhausted")] + #[ts(rename = "SubscriptionExhausted")] + pub subscription_exhausted: Vec, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] diff --git a/codex-rs/app-server/src/config_api.rs b/codex-rs/app-server/src/config_api.rs index 4b6cbdd1934..89d5b319f58 100644 --- a/codex-rs/app-server/src/config_api.rs +++ b/codex-rs/app-server/src/config_api.rs @@ -308,6 +308,7 @@ fn map_hooks_requirements_to_api(hooks: ManagedHooksRequirementsToml) -> Managed session_start, user_prompt_submit, stop, + subscription_exhausted, } = hooks; ManagedHooksRequirements { @@ -319,6 +320,7 @@ fn map_hooks_requirements_to_api(hooks: ManagedHooksRequirementsToml) -> Managed session_start: map_hook_matcher_groups_to_api(session_start), user_prompt_submit: map_hook_matcher_groups_to_api(user_prompt_submit), stop: map_hook_matcher_groups_to_api(stop), + subscription_exhausted: map_hook_matcher_groups_to_api(subscription_exhausted), } } @@ -627,6 +629,7 @@ mod tests { session_start: Vec::new(), user_prompt_submit: Vec::new(), stop: Vec::new(), + subscription_exhausted: Vec::new(), }), ); assert_eq!( diff --git a/codex-rs/config/src/hook_config.rs b/codex-rs/config/src/hook_config.rs index d947ebb8678..64998fbc548 100644 --- a/codex-rs/config/src/hook_config.rs +++ b/codex-rs/config/src/hook_config.rs @@ -41,6 +41,8 @@ pub struct HookEventsToml { pub user_prompt_submit: Vec, #[serde(rename = "Stop", default)] pub stop: Vec, + #[serde(rename = "SubscriptionExhausted", default)] + pub subscription_exhausted: Vec, } impl HookEventsToml { @@ -52,6 +54,7 @@ impl HookEventsToml { session_start, user_prompt_submit, stop, + subscription_exhausted, } = self; pre_tool_use.is_empty() && permission_request.is_empty() @@ -59,6 +62,7 @@ impl HookEventsToml { && session_start.is_empty() && user_prompt_submit.is_empty() && stop.is_empty() + && subscription_exhausted.is_empty() } pub fn handler_count(&self) -> usize { @@ -69,6 +73,7 @@ impl HookEventsToml { session_start, user_prompt_submit, stop, + subscription_exhausted, } = self; [ pre_tool_use, @@ -77,6 +82,7 @@ impl HookEventsToml { session_start, user_prompt_submit, stop, + subscription_exhausted, ] .into_iter() .flatten() @@ -84,7 +90,7 @@ impl HookEventsToml { .sum() } - pub fn into_matcher_groups(self) -> [(HookEventName, Vec); 6] { + pub fn into_matcher_groups(self) -> [(HookEventName, Vec); 7] { [ (HookEventName::PreToolUse, self.pre_tool_use), (HookEventName::PermissionRequest, self.permission_request), @@ -92,6 +98,10 @@ impl HookEventsToml { (HookEventName::SessionStart, self.session_start), (HookEventName::UserPromptSubmit, self.user_prompt_submit), (HookEventName::Stop, self.stop), + ( + HookEventName::SubscriptionExhausted, + self.subscription_exhausted, + ), ] } } diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index a30a3ed9254..143ffa6ca18 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1037,6 +1037,13 @@ }, "type": "array" }, + "SubscriptionExhausted": { + "default": [], + "items": { + "$ref": "#/definitions/MatcherGroup" + }, + "type": "array" + }, "UserPromptSubmit": { "default": [], "items": { diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index ba81b451a74..f39bbd8930d 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -828,6 +828,15 @@ impl ModelClientSession { .set_connection_reused(/*connection_reused*/ false); } + /// Clears transport and sticky-routing state after an out-of-band auth or account switch. + /// + /// Reusing the previous `x-codex-turn-state` can keep a retry pinned to the same exhausted + /// route even after the hook has selected a different account. + pub(crate) fn reset_after_subscription_exhaustion_recovery(&mut self) { + self.reset_websocket_session(); + self.turn_state = Arc::new(OnceLock::new()); + } + fn build_responses_request( &self, provider: &codex_api::Provider, diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 9a928545152..6c34f27787b 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -12,6 +12,8 @@ use codex_hooks::PostToolUseRequest; use codex_hooks::PreToolUseOutcome; use codex_hooks::PreToolUseRequest; use codex_hooks::SessionStartOutcome; +use codex_hooks::SubscriptionExhaustedOutcome; +use codex_hooks::SubscriptionExhaustedRequest; use codex_hooks::UserPromptSubmitOutcome; use codex_hooks::UserPromptSubmitRequest; use codex_otel::HOOK_RUN_DURATION_METRIC; @@ -278,6 +280,24 @@ pub(crate) async fn run_user_prompt_submit_hooks( .await } +pub(crate) async fn run_subscription_exhausted_hooks( + sess: &Arc, + turn_context: &Arc, + request: SubscriptionExhaustedRequest, +) -> bool { + let hooks = sess.hooks(); + let preview_runs = hooks.preview_subscription_exhausted(&request); + emit_hook_started_events(sess, turn_context, preview_runs).await; + + let SubscriptionExhaustedOutcome { hook_events } = + hooks.run_subscription_exhausted(request).await; + let succeeded = hook_events + .iter() + .any(|event| event.run.status == HookRunStatus::Completed); + emit_hook_completed_events(sess, turn_context, hook_events).await; + succeeded +} + pub(crate) async fn inspect_pending_input( sess: &Arc, turn_context: &Arc, @@ -470,6 +490,7 @@ fn hook_run_metric_tags(run: &HookRunSummary) -> [(&'static str, &'static str); HookEventName::SessionStart => "SessionStart", HookEventName::UserPromptSubmit => "UserPromptSubmit", HookEventName::Stop => "Stop", + HookEventName::SubscriptionExhausted => "SubscriptionExhausted", }; let hook_source = match run.source { HookSource::System => "system", @@ -498,7 +519,7 @@ fn hook_run_metric_tags(run: &HookRunSummary) -> [(&'static str, &'static str); ] } -fn hook_permission_mode(turn_context: &TurnContext) -> String { +pub(crate) fn hook_permission_mode(turn_context: &TurnContext) -> String { match turn_context.approval_policy.value() { AskForApproval::Never => "bypassPermissions", AskForApproval::UnlessTrusted diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index a396851f98d..092cf1e5fdf 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -105,6 +105,7 @@ pub(crate) use skills::resolve_skill_dependencies_for_turn; pub(crate) use skills::skills_load_input_from_config; mod skills_watcher; mod stream_events_utils; +mod subscription_exhaustion; pub mod test_support; mod unified_exec; pub mod windows_sandbox; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 2b37372a33f..3492d5486ee 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -49,6 +49,7 @@ use crate::stream_events_utils::last_assistant_message_from_item; use crate::stream_events_utils::mark_thread_memory_mode_polluted_if_external_context; use crate::stream_events_utils::raw_assistant_output_text_from_item; use crate::stream_events_utils::record_completed_response_item; +use crate::subscription_exhaustion; use crate::tools::ToolRouter; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::parallel::ToolCallRuntime; @@ -372,6 +373,7 @@ pub(crate) async fn run_turn( // 1. At the start of a turn, so the fresh user prompt in `input` gets sampled first. // 2. After auto-compact, when model/tool continuation needs to resume before any steer. let mut can_drain_pending_input = input.is_empty(); + let mut subscription_exhaustion_hook_attempted = false; loop { if run_pending_session_start_hooks(&sess, &turn_context).await { @@ -652,6 +654,33 @@ pub(crate) async fn run_turn( break; } Err(e) => { + if subscription_exhaustion::should_run_hook_for_error(&e) + && !subscription_exhaustion_hook_attempted + { + subscription_exhaustion_hook_attempted = true; + if let Some(recovery) = subscription_exhaustion::recover_with_hook_if_available( + &sess, + &turn_context, + &e, + ) + .await + { + client_session.reset_after_subscription_exhaustion_recovery(); + let message = if recovery.auth_changed { + "A SubscriptionExhausted hook refreshed Codex auth. Reconnecting and retrying request." + } else { + "A SubscriptionExhausted hook completed. Reconnecting and retrying request." + }; + sess.send_event( + &turn_context, + EventMsg::Warning(WarningEvent { + message: message.to_string(), + }), + ) + .await; + continue; + } + } info!("Turn error: {e:#}"); let event = EventMsg::Error(e.to_error_event(/*message_prefix*/ None)); sess.send_event(&turn_context, event).await; diff --git a/codex-rs/core/src/subscription_exhaustion.rs b/codex-rs/core/src/subscription_exhaustion.rs new file mode 100644 index 00000000000..2755d736442 --- /dev/null +++ b/codex-rs/core/src/subscription_exhaustion.rs @@ -0,0 +1,114 @@ +use std::sync::Arc; + +use codex_hooks::SubscriptionExhaustedRequest; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_protocol::auth::PlanType; +use codex_protocol::error::CodexErr; + +use crate::hook_runtime::run_subscription_exhausted_hooks; +use crate::session::session::Session; +use crate::session::turn_context::TurnContext; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct SubscriptionExhaustionRecovery { + pub auth_changed: bool, +} + +#[derive(PartialEq, Eq)] +struct AuthSnapshot { + mode: codex_app_server_protocol::AuthMode, + account_id: Option, + token: Option, +} + +pub(crate) fn should_run_hook_for_error(err: &CodexErr) -> bool { + matches!( + err, + CodexErr::UsageLimitReached(_) | CodexErr::QuotaExceeded + ) +} + +pub(crate) async fn recover_with_hook_if_available( + sess: &Arc, + turn_context: &Arc, + err: &CodexErr, +) -> Option { + let auth_manager = sess.services.model_client.auth_manager()?; + let auth_before = auth_manager.auth_cached()?; + + if !can_refresh_managed_chatgpt_auth(&auth_before) { + return None; + } + + let before = auth_snapshot(&auth_before); + let request = + build_hook_request(sess, turn_context, auth_manager.as_ref(), &auth_before, err).await; + + if !run_subscription_exhausted_hooks(sess, turn_context, request).await { + return None; + } + + auth_manager.reload().await; + + let auth_changed = auth_manager + .auth_cached() + .is_some_and(|auth_after| auth_snapshot(&auth_after) != before); + + Some(SubscriptionExhaustionRecovery { auth_changed }) +} + +fn can_refresh_managed_chatgpt_auth(auth: &CodexAuth) -> bool { + auth.is_chatgpt_auth() && !auth.is_external_chatgpt_tokens() +} + +async fn build_hook_request( + sess: &Arc, + turn_context: &Arc, + auth_manager: &AuthManager, + auth: &CodexAuth, + err: &CodexErr, +) -> SubscriptionExhaustedRequest { + let (error_kind, plan_type, resets_at) = subscription_error_fields(err); + + SubscriptionExhaustedRequest { + session_id: sess.conversation_id, + turn_id: turn_context.sub_id.clone(), + cwd: turn_context.cwd.clone(), + codex_home: auth_manager.codex_home().to_path_buf(), + transcript_path: sess.hook_transcript_path().await, + model: turn_context.model_info.slug.clone(), + permission_mode: crate::hook_runtime::hook_permission_mode(turn_context), + error_kind, + plan_type, + resets_at, + account_id: auth.get_account_id(), + } +} + +fn subscription_error_fields(err: &CodexErr) -> (String, Option, Option) { + match err { + CodexErr::UsageLimitReached(err) => ( + "usage_limit_reached".to_string(), + err.plan_type.clone().map(plan_type_wire_value), + err.resets_at.map(|resets_at| resets_at.timestamp()), + ), + CodexErr::QuotaExceeded => ("quota_exceeded".to_string(), None, None), + _ => unreachable!("subscription hook request only builds for subscription errors"), + } +} + +fn plan_type_wire_value(plan_type: PlanType) -> String { + match plan_type { + PlanType::Known(plan) => plan.raw_value().to_string(), + PlanType::Unknown(raw) => raw, + } +} + +fn auth_snapshot(auth: &CodexAuth) -> AuthSnapshot { + AuthSnapshot { + mode: auth.api_auth_mode(), + account_id: auth.get_account_id(), + token: auth.get_token().ok(), + } +} diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index eb7419076d0..7809700fe3a 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -349,7 +349,7 @@ impl ThreadManager { environment_manager: Arc, ) -> Self { set_thread_manager_test_mode_for_tests(/*enabled*/ true); - let auth_manager = AuthManager::from_auth_for_testing(auth); + let auth_manager = AuthManager::from_auth_for_testing_with_home(auth, codex_home.clone()); let skills_codex_home = match AbsolutePathBuf::from_absolute_path_checked(&codex_home) { Ok(codex_home) => codex_home, Err(err) => panic!("test codex_home should be absolute: {err}"), diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 291a0795ce8..f948c4b1345 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -426,7 +426,10 @@ impl TestCodexBuilder { let thread_manager = if config.model_catalog.is_some() { ThreadManager::new( &config, - codex_core::test_support::auth_manager_from_auth(auth.clone()), + codex_core::test_support::auth_manager_from_auth_with_home( + auth.clone(), + config.codex_home.to_path_buf(), + ), SessionSource::Exec, Arc::clone(&environment_manager), /*analytics_events_client*/ None, @@ -445,7 +448,10 @@ impl TestCodexBuilder { let new_conversation = match (resume_from, user_shell_override) { (Some(path), Some(user_shell_override)) => { - let auth_manager = codex_core::test_support::auth_manager_from_auth(auth); + let auth_manager = codex_core::test_support::auth_manager_from_auth_with_home( + auth, + config.codex_home.to_path_buf(), + ); Box::pin( codex_core::test_support::resume_thread_from_rollout_with_user_shell_override( thread_manager.as_ref(), @@ -458,7 +464,10 @@ impl TestCodexBuilder { .await? } (Some(path), None) => { - let auth_manager = codex_core::test_support::auth_manager_from_auth(auth); + let auth_manager = codex_core::test_support::auth_manager_from_auth_with_home( + auth, + config.codex_home.to_path_buf(), + ); Box::pin(thread_manager.resume_thread_from_rollout( config.clone(), path, diff --git a/codex-rs/core/tests/suite/quota_exceeded.rs b/codex-rs/core/tests/suite/quota_exceeded.rs index 4c0677e69a5..61ac0afa678 100644 --- a/codex-rs/core/tests/suite/quota_exceeded.rs +++ b/codex-rs/core/tests/suite/quota_exceeded.rs @@ -1,16 +1,32 @@ use anyhow::Result; +#[cfg(unix)] +use codex_features::Feature; +#[cfg(unix)] +use codex_login::CodexAuth; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::user_input::UserInput; +#[cfg(unix)] +use core_test_support::responses::ev_assistant_message; +#[cfg(unix)] +use core_test_support::responses::ev_completed; use core_test_support::responses::ev_response_created; +#[cfg(unix)] +use core_test_support::responses::mount_response_sequence; use core_test_support::responses::mount_sse_once; use core_test_support::responses::sse; +#[cfg(unix)] +use core_test_support::responses::sse_failed; +#[cfg(unix)] +use core_test_support::responses::sse_response; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::json; +#[cfg(unix)] +use wiremock::ResponseTemplate; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn quota_exceeded_emits_single_error_event() -> Result<()> { @@ -74,3 +90,262 @@ async fn quota_exceeded_emits_single_error_event() -> Result<()> { Ok(()) } + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn quota_exceeded_runs_successful_subscription_exhausted_hook_and_retries_with_fresh_route() +-> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let tempdir = tempfile::tempdir()?; + let same_auth_path = tempdir.path().join("same-auth.json"); + std::fs::write( + &same_auth_path, + auth_json_for_chatgpt("Access Token", "account_id"), + )?; + + let command_path = tempdir.path().join("successful_hook.py"); + std::fs::write( + &command_path, + format!( + r#"import json +import shutil +import sys +from pathlib import Path + +payload = json.load(sys.stdin) +codex_home = Path(payload["codex_home"]) +codex_home.mkdir(parents=True, exist_ok=True) +with (codex_home / "subscription_exhausted_hook_log.jsonl").open("a", encoding="utf-8") as handle: + handle.write(json.dumps(payload) + "\n") +shutil.copyfile(r"{same_auth_path}", codex_home / "auth.json") +"#, + same_auth_path = same_auth_path.display() + ), + )?; + + let quota_exceeded = sse_response(sse_failed("resp-1", "insufficient_quota", "quota exceeded")) + .insert_header("x-codex-turn-state", "stale-route"); + let success = ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "ok"), + ev_completed("resp-2"), + ])); + let request_log = mount_response_sequence(&server, vec![quota_exceeded, success]).await; + + let command = format!("python3 {}", command_path.display()); + let mut builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_pre_build_hook(move |home| { + let hooks = json!({ + "hooks": { + "SubscriptionExhausted": [{ + "hooks": [{ + "type": "command", + "command": command, + "statusMessage": "selecting a fresh subscription route", + "timeout": 5 + }] + }] + } + }); + std::fs::write(home.join("hooks.json"), hooks.to_string()).expect("write hooks.json"); + }) + .with_config(|config| { + config + .features + .enable(Feature::CodexHooks) + .expect("test config should allow hooks feature"); + }); + let test = builder.build(&server).await?; + + test.submit_turn("hello").await?; + + let requests = request_log.requests(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].header("x-codex-turn-state"), None); + assert_eq!(requests[1].header("x-codex-turn-state"), None); + assert_eq!( + requests[0].header("authorization").as_deref(), + Some("Bearer Access Token") + ); + assert_eq!( + requests[1].header("authorization").as_deref(), + Some("Bearer Access Token") + ); + + let hook_log_path = test + .codex_home_path() + .join("subscription_exhausted_hook_log.jsonl"); + let hook_log = std::fs::read_to_string(hook_log_path)?; + let hook_payload: serde_json::Value = serde_json::from_str( + hook_log + .lines() + .next() + .expect("subscription hook should log one invocation"), + )?; + assert_eq!(hook_payload["hook_event_name"], "SubscriptionExhausted"); + assert_eq!(hook_payload["error_kind"], "quota_exceeded"); + + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn usage_limit_reached_runs_subscription_exhausted_hook_and_retries_with_rotated_auth() +-> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let tempdir = tempfile::tempdir()?; + let rotated_auth_path = tempdir.path().join("rotated-auth.json"); + std::fs::write( + &rotated_auth_path, + auth_json_for_chatgpt("rotated-access-token", "rotated-account"), + )?; + + let command_path = tempdir.path().join("rotate_auth.py"); + std::fs::write( + &command_path, + format!( + r#"import json +import shutil +import sys +from pathlib import Path + +payload = json.load(sys.stdin) +codex_home = Path(payload["codex_home"]) +codex_home.mkdir(parents=True, exist_ok=True) +with (codex_home / "subscription_exhausted_hook_log.jsonl").open("a", encoding="utf-8") as handle: + handle.write(json.dumps(payload) + "\n") +shutil.copyfile(r"{rotated_auth_path}", codex_home / "auth.json") +"#, + rotated_auth_path = rotated_auth_path.display() + ), + )?; + + let usage_limit = ResponseTemplate::new(429).set_body_json(json!({ + "error": { + "type": "usage_limit_reached", + "message": "limit reached", + "resets_at": 1704067242, + "plan_type": "pro" + } + })); + let success = ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "ok"), + ev_completed("resp-2"), + ])); + let request_log = mount_response_sequence(&server, vec![usage_limit, success]).await; + + let command = format!("python3 {}", command_path.display()); + let mut builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_pre_build_hook(move |home| { + let hooks = json!({ + "hooks": { + "SubscriptionExhausted": [{ + "hooks": [{ + "type": "command", + "command": command, + "statusMessage": "rotating exhausted subscription auth", + "timeout": 5 + }] + }] + } + }); + std::fs::write(home.join("hooks.json"), hooks.to_string()).expect("write hooks.json"); + }) + .with_config(|config| { + config + .features + .enable(Feature::CodexHooks) + .expect("test config should allow hooks feature"); + }); + let test = builder.build(&server).await?; + let listed_hooks = codex_hooks::list_hooks(codex_hooks::HooksConfig { + feature_enabled: true, + config_layer_stack: Some(test.config.config_layer_stack.clone()), + ..codex_hooks::HooksConfig::default() + }); + assert_eq!(listed_hooks.hooks.len(), 1); + + test.submit_turn("hello").await?; + + let hook_log_path = test + .codex_home_path() + .join("subscription_exhausted_hook_log.jsonl"); + assert!( + hook_log_path.exists(), + "subscription exhausted hook should have run" + ); + let requests = request_log.requests(); + assert_eq!(requests.len(), 2); + assert_eq!( + requests[0].header("authorization").as_deref(), + Some("Bearer Access Token") + ); + assert_eq!( + requests[1].header("authorization").as_deref(), + Some("Bearer rotated-access-token") + ); + assert_eq!( + requests[1].header("chatgpt-account-id").as_deref(), + Some("rotated-account") + ); + let hook_log = std::fs::read_to_string(hook_log_path)?; + let hook_payload: serde_json::Value = serde_json::from_str( + hook_log + .lines() + .next() + .expect("subscription hook should log one invocation"), + )?; + assert_eq!(hook_payload["hook_event_name"], "SubscriptionExhausted"); + assert_eq!(hook_payload["error_kind"], "usage_limit_reached"); + assert_eq!(hook_payload["plan_type"], "pro"); + assert_eq!(hook_payload["resets_at"], 1704067242); + assert_eq!(hook_payload["account_id"], "account_id"); + assert_eq!( + hook_payload["codex_home"], + test.codex_home_path().display().to_string() + ); + + Ok(()) +} + +#[cfg(unix)] +fn auth_json_for_chatgpt(access_token: &str, account_id: &str) -> String { + use base64::Engine as _; + + let header = json!({ "alg": "none", "typ": "JWT" }); + let payload = json!({ + "email": "rotated@example.com", + "https://api.openai.com/auth": { + "chatgpt_plan_type": "plus", + "chatgpt_account_id": account_id + } + }); + let encode = |value: &serde_json::Value| { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(value.to_string()) + }; + let id_token = format!("{}.{}.sig", encode(&header), encode(&payload)); + + json!({ + "auth_mode": "chatgpt", + "OPENAI_API_KEY": null, + "tokens": { + "id_token": id_token, + "access_token": access_token, + "refresh_token": "rotated-refresh-token", + "account_id": account_id + }, + "last_refresh": chrono::Utc::now() + }) + .to_string() +} diff --git a/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json b/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json index 21d45382b0e..bdd98fa99c8 100644 --- a/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json +++ b/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json @@ -9,7 +9,8 @@ "PostToolUse", "SessionStart", "UserPromptSubmit", - "Stop" + "Stop", + "SubscriptionExhausted" ], "type": "string" }, diff --git a/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json b/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json index 43a2a4828e3..cc47e842d79 100644 --- a/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json +++ b/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json @@ -15,7 +15,8 @@ "PostToolUse", "SessionStart", "UserPromptSubmit", - "Stop" + "Stop", + "SubscriptionExhausted" ], "type": "string" }, diff --git a/codex-rs/hooks/schema/generated/pre-tool-use.command.output.schema.json b/codex-rs/hooks/schema/generated/pre-tool-use.command.output.schema.json index ba6bb3401b6..2baadb0df33 100644 --- a/codex-rs/hooks/schema/generated/pre-tool-use.command.output.schema.json +++ b/codex-rs/hooks/schema/generated/pre-tool-use.command.output.schema.json @@ -9,7 +9,8 @@ "PostToolUse", "SessionStart", "UserPromptSubmit", - "Stop" + "Stop", + "SubscriptionExhausted" ], "type": "string" }, diff --git a/codex-rs/hooks/schema/generated/session-start.command.output.schema.json b/codex-rs/hooks/schema/generated/session-start.command.output.schema.json index 30d2b0f9779..aedfa37d17b 100644 --- a/codex-rs/hooks/schema/generated/session-start.command.output.schema.json +++ b/codex-rs/hooks/schema/generated/session-start.command.output.schema.json @@ -9,7 +9,8 @@ "PostToolUse", "SessionStart", "UserPromptSubmit", - "Stop" + "Stop", + "SubscriptionExhausted" ], "type": "string" }, diff --git a/codex-rs/hooks/schema/generated/subscription-exhausted.command.input.schema.json b/codex-rs/hooks/schema/generated/subscription-exhausted.command.input.schema.json new file mode 100644 index 00000000000..3f062e94470 --- /dev/null +++ b/codex-rs/hooks/schema/generated/subscription-exhausted.command.input.schema.json @@ -0,0 +1,79 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "account_id": { + "$ref": "#/definitions/NullableString" + }, + "codex_home": { + "type": "string" + }, + "cwd": { + "type": "string" + }, + "error_kind": { + "enum": [ + "usage_limit_reached", + "quota_exceeded" + ], + "type": "string" + }, + "hook_event_name": { + "const": "SubscriptionExhausted", + "type": "string" + }, + "model": { + "type": "string" + }, + "permission_mode": { + "enum": [ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions" + ], + "type": "string" + }, + "plan_type": { + "$ref": "#/definitions/NullableString" + }, + "resets_at": { + "format": "int64", + "type": "integer" + }, + "session_id": { + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "account_id", + "codex_home", + "cwd", + "error_kind", + "hook_event_name", + "model", + "permission_mode", + "plan_type", + "session_id", + "transcript_path", + "turn_id" + ], + "title": "subscription-exhausted.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/subscription-exhausted.command.output.schema.json b/codex-rs/hooks/schema/generated/subscription-exhausted.command.output.schema.json new file mode 100644 index 00000000000..8747f63bf14 --- /dev/null +++ b/codex-rs/hooks/schema/generated/subscription-exhausted.command.output.schema.json @@ -0,0 +1,24 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "subscription-exhausted.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/user-prompt-submit.command.output.schema.json b/codex-rs/hooks/schema/generated/user-prompt-submit.command.output.schema.json index 26353d97db1..37d2b1ea9c1 100644 --- a/codex-rs/hooks/schema/generated/user-prompt-submit.command.output.schema.json +++ b/codex-rs/hooks/schema/generated/user-prompt-submit.command.output.schema.json @@ -15,7 +15,8 @@ "PostToolUse", "SessionStart", "UserPromptSubmit", - "Stop" + "Stop", + "SubscriptionExhausted" ], "type": "string" }, diff --git a/codex-rs/hooks/src/engine/discovery.rs b/codex-rs/hooks/src/engine/discovery.rs index 8c520b74970..1d5246b518f 100644 --- a/codex-rs/hooks/src/engine/discovery.rs +++ b/codex-rs/hooks/src/engine/discovery.rs @@ -461,6 +461,7 @@ fn hook_event_key_label(event_name: codex_protocol::protocol::HookEventName) -> codex_protocol::protocol::HookEventName::SessionStart => "session_start", codex_protocol::protocol::HookEventName::UserPromptSubmit => "user_prompt_submit", codex_protocol::protocol::HookEventName::Stop => "stop", + codex_protocol::protocol::HookEventName::SubscriptionExhausted => "subscription_exhausted", } } diff --git a/codex-rs/hooks/src/engine/dispatcher.rs b/codex-rs/hooks/src/engine/dispatcher.rs index c44b1fe69d7..2a7f745a34c 100644 --- a/codex-rs/hooks/src/engine/dispatcher.rs +++ b/codex-rs/hooks/src/engine/dispatcher.rs @@ -55,7 +55,9 @@ pub(crate) fn select_handlers_for_matcher_inputs( .any(|input| matches_matcher(handler.matcher.as_deref(), Some(input))) } } - HookEventName::UserPromptSubmit | HookEventName::Stop => true, + HookEventName::UserPromptSubmit + | HookEventName::Stop + | HookEventName::SubscriptionExhausted => true, }) .cloned() .collect() @@ -133,7 +135,8 @@ fn scope_for_event(event_name: HookEventName) -> HookScope { | HookEventName::PermissionRequest | HookEventName::PostToolUse | HookEventName::UserPromptSubmit - | HookEventName::Stop => HookScope::Turn, + | HookEventName::Stop + | HookEventName::SubscriptionExhausted => HookScope::Turn, } } diff --git a/codex-rs/hooks/src/engine/mod.rs b/codex-rs/hooks/src/engine/mod.rs index c06c9fabd7b..ffd7d3de323 100644 --- a/codex-rs/hooks/src/engine/mod.rs +++ b/codex-rs/hooks/src/engine/mod.rs @@ -24,6 +24,8 @@ use crate::events::session_start::SessionStartOutcome; use crate::events::session_start::SessionStartRequest; use crate::events::stop::StopOutcome; use crate::events::stop::StopRequest; +use crate::events::subscription_exhausted::SubscriptionExhaustedOutcome; +use crate::events::subscription_exhausted::SubscriptionExhaustedRequest; use crate::events::user_prompt_submit::UserPromptSubmitOutcome; use crate::events::user_prompt_submit::UserPromptSubmitRequest; @@ -64,6 +66,9 @@ impl ConfiguredHandler { codex_protocol::protocol::HookEventName::SessionStart => "session-start", codex_protocol::protocol::HookEventName::UserPromptSubmit => "user-prompt-submit", codex_protocol::protocol::HookEventName::Stop => "stop", + codex_protocol::protocol::HookEventName::SubscriptionExhausted => { + "subscription-exhausted" + } } } } @@ -197,6 +202,20 @@ impl ClaudeHooksEngine { pub(crate) async fn run_stop(&self, request: StopRequest) -> StopOutcome { crate::events::stop::run(&self.handlers, &self.shell, request).await } + + pub(crate) fn preview_subscription_exhausted( + &self, + request: &SubscriptionExhaustedRequest, + ) -> Vec { + crate::events::subscription_exhausted::preview(&self.handlers, request) + } + + pub(crate) async fn run_subscription_exhausted( + &self, + request: SubscriptionExhaustedRequest, + ) -> SubscriptionExhaustedOutcome { + crate::events::subscription_exhausted::run(&self.handlers, &self.shell, request).await + } } #[cfg(test)] diff --git a/codex-rs/hooks/src/engine/output_parser.rs b/codex-rs/hooks/src/engine/output_parser.rs index 0a3a994e19d..fa1bd072c49 100644 --- a/codex-rs/hooks/src/engine/output_parser.rs +++ b/codex-rs/hooks/src/engine/output_parser.rs @@ -59,6 +59,11 @@ pub(crate) struct StopOutput { pub invalid_block_reason: Option, } +#[derive(Debug, Clone)] +pub(crate) struct SubscriptionExhaustedOutput { + pub universal: UniversalOutput, +} + use crate::schema::BlockDecisionWire; use crate::schema::HookUniversalOutputWire; use crate::schema::PermissionRequestBehaviorWire; @@ -70,6 +75,7 @@ use crate::schema::PreToolUseDecisionWire; use crate::schema::PreToolUsePermissionDecisionWire; use crate::schema::SessionStartCommandOutputWire; use crate::schema::StopCommandOutputWire; +use crate::schema::SubscriptionExhaustedCommandOutputWire; use crate::schema::UserPromptSubmitCommandOutputWire; pub(crate) fn parse_session_start(stdout: &str) -> Option { @@ -232,6 +238,13 @@ pub(crate) fn parse_stop(stdout: &str) -> Option { }) } +pub(crate) fn parse_subscription_exhausted(stdout: &str) -> Option { + let wire: SubscriptionExhaustedCommandOutputWire = parse_json(stdout)?; + Some(SubscriptionExhaustedOutput { + universal: UniversalOutput::from(wire.universal), + }) +} + impl From for UniversalOutput { fn from(value: HookUniversalOutputWire) -> Self { Self { diff --git a/codex-rs/hooks/src/engine/schema_loader.rs b/codex-rs/hooks/src/engine/schema_loader.rs index d8bdb141c28..f232c3a43f8 100644 --- a/codex-rs/hooks/src/engine/schema_loader.rs +++ b/codex-rs/hooks/src/engine/schema_loader.rs @@ -16,6 +16,8 @@ pub(crate) struct GeneratedHookSchemas { pub user_prompt_submit_command_output: Value, pub stop_command_input: Value, pub stop_command_output: Value, + pub subscription_exhausted_command_input: Value, + pub subscription_exhausted_command_output: Value, } pub(crate) fn generated_hook_schemas() -> &'static GeneratedHookSchemas { @@ -69,6 +71,16 @@ pub(crate) fn generated_hook_schemas() -> &'static GeneratedHookSchemas { "stop.command.output", include_str!("../../schema/generated/stop.command.output.schema.json"), ), + subscription_exhausted_command_input: parse_json_schema( + "subscription-exhausted.command.input", + include_str!("../../schema/generated/subscription-exhausted.command.input.schema.json"), + ), + subscription_exhausted_command_output: parse_json_schema( + "subscription-exhausted.command.output", + include_str!( + "../../schema/generated/subscription-exhausted.command.output.schema.json" + ), + ), }) } @@ -98,5 +110,13 @@ mod tests { assert_eq!(schemas.user_prompt_submit_command_output["type"], "object"); assert_eq!(schemas.stop_command_input["type"], "object"); assert_eq!(schemas.stop_command_output["type"], "object"); + assert_eq!( + schemas.subscription_exhausted_command_input["type"], + "object" + ); + assert_eq!( + schemas.subscription_exhausted_command_output["type"], + "object" + ); } } diff --git a/codex-rs/hooks/src/events/common.rs b/codex-rs/hooks/src/events/common.rs index de3f3292acd..104dfe02179 100644 --- a/codex-rs/hooks/src/events/common.rs +++ b/codex-rs/hooks/src/events/common.rs @@ -104,7 +104,9 @@ pub(crate) fn matcher_pattern_for_event( | HookEventName::PermissionRequest | HookEventName::PostToolUse | HookEventName::SessionStart => matcher, - HookEventName::UserPromptSubmit | HookEventName::Stop => None, + HookEventName::UserPromptSubmit + | HookEventName::Stop + | HookEventName::SubscriptionExhausted => None, } } @@ -251,6 +253,10 @@ mod tests { matcher_pattern_for_event(HookEventName::Stop, Some("^done$")), None ); + assert_eq!( + matcher_pattern_for_event(HookEventName::SubscriptionExhausted, Some("^pro$")), + None + ); } #[test] diff --git a/codex-rs/hooks/src/events/mod.rs b/codex-rs/hooks/src/events/mod.rs index 52148324e2f..d2a6d304a7f 100644 --- a/codex-rs/hooks/src/events/mod.rs +++ b/codex-rs/hooks/src/events/mod.rs @@ -4,4 +4,5 @@ pub mod post_tool_use; pub mod pre_tool_use; pub mod session_start; pub mod stop; +pub mod subscription_exhausted; pub mod user_prompt_submit; diff --git a/codex-rs/hooks/src/events/subscription_exhausted.rs b/codex-rs/hooks/src/events/subscription_exhausted.rs new file mode 100644 index 00000000000..0ab87e61ea7 --- /dev/null +++ b/codex-rs/hooks/src/events/subscription_exhausted.rs @@ -0,0 +1,195 @@ +use std::path::PathBuf; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use codex_utils_absolute_path::AbsolutePathBuf; + +use super::common; +use crate::engine::CommandShell; +use crate::engine::ConfiguredHandler; +use crate::engine::command_runner::CommandRunResult; +use crate::engine::dispatcher; +use crate::engine::output_parser; +use crate::schema::NullableString; +use crate::schema::SubscriptionExhaustedCommandInput; + +#[derive(Debug, Clone)] +pub struct SubscriptionExhaustedRequest { + pub session_id: ThreadId, + pub turn_id: String, + pub cwd: AbsolutePathBuf, + pub codex_home: PathBuf, + pub transcript_path: Option, + pub model: String, + pub permission_mode: String, + pub error_kind: String, + pub plan_type: Option, + pub resets_at: Option, + pub account_id: Option, +} + +#[derive(Debug)] +pub struct SubscriptionExhaustedOutcome { + pub hook_events: Vec, +} + +#[derive(Debug, Default, PartialEq, Eq)] +struct SubscriptionExhaustedHandlerData; + +pub(crate) fn preview( + handlers: &[ConfiguredHandler], + _request: &SubscriptionExhaustedRequest, +) -> Vec { + dispatcher::select_handlers( + handlers, + HookEventName::SubscriptionExhausted, + /*matcher_input*/ None, + ) + .into_iter() + .map(|handler| dispatcher::running_summary(&handler)) + .collect() +} + +pub(crate) async fn run( + handlers: &[ConfiguredHandler], + shell: &CommandShell, + request: SubscriptionExhaustedRequest, +) -> SubscriptionExhaustedOutcome { + let matched = dispatcher::select_handlers( + handlers, + HookEventName::SubscriptionExhausted, + /*matcher_input*/ None, + ); + if matched.is_empty() { + return SubscriptionExhaustedOutcome { + hook_events: Vec::new(), + }; + } + + let input_json = match serde_json::to_string(&SubscriptionExhaustedCommandInput { + session_id: request.session_id.to_string(), + turn_id: request.turn_id.clone(), + transcript_path: NullableString::from_path(request.transcript_path.clone()), + cwd: request.cwd.display().to_string(), + codex_home: request.codex_home.display().to_string(), + hook_event_name: "SubscriptionExhausted".to_string(), + model: request.model.clone(), + permission_mode: request.permission_mode.clone(), + error_kind: request.error_kind.clone(), + plan_type: NullableString::from_string(request.plan_type.clone()), + resets_at: request.resets_at, + account_id: NullableString::from_string(request.account_id.clone()), + }) { + Ok(input_json) => input_json, + Err(error) => { + return SubscriptionExhaustedOutcome { + hook_events: common::serialization_failure_hook_events( + matched, + Some(request.turn_id), + format!("failed to serialize subscription exhausted hook input: {error}"), + ), + }; + } + }; + + let results = dispatcher::execute_handlers( + shell, + matched, + input_json, + request.cwd.as_path(), + Some(request.turn_id), + parse_completed, + ) + .await; + + SubscriptionExhaustedOutcome { + hook_events: results.into_iter().map(|result| result.completed).collect(), + } +} + +fn parse_completed( + handler: &ConfiguredHandler, + run_result: CommandRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler { + let mut entries = Vec::new(); + let mut status = HookRunStatus::Completed; + + match run_result.error.as_deref() { + Some(error) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error.to_string(), + }); + } + None => match run_result.exit_code { + Some(0) => { + let trimmed_stdout = run_result.stdout.trim(); + if trimmed_stdout.is_empty() { + } else if let Some(parsed) = + output_parser::parse_subscription_exhausted(&run_result.stdout) + { + if !parsed.universal.suppress_output + && let Some(system_message) = parsed.universal.system_message + { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: system_message, + }); + } + if !parsed.universal.continue_processing { + status = HookRunStatus::Stopped; + if let Some(stop_reason) = parsed.universal.stop_reason { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: stop_reason, + }); + } + } + } else if trimmed_stdout.starts_with('{') || trimmed_stdout.starts_with('[') { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook returned invalid subscription exhausted JSON output" + .to_string(), + }); + } else { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: trimmed_stdout.to_string(), + }); + } + } + Some(exit_code) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: format!("hook exited with code {exit_code}"), + }); + } + None => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook exited without a status code".to_string(), + }); + } + }, + } + + let completed = HookCompletedEvent { + turn_id, + run: dispatcher::completed_summary(handler, &run_result, status, entries), + }; + + dispatcher::ParsedHandler { + completed, + data: SubscriptionExhaustedHandlerData, + } +} diff --git a/codex-rs/hooks/src/lib.rs b/codex-rs/hooks/src/lib.rs index 4e16969a587..33656ff46e6 100644 --- a/codex-rs/hooks/src/lib.rs +++ b/codex-rs/hooks/src/lib.rs @@ -8,13 +8,14 @@ mod types; pub use engine::HookListEntry; /// Hook event names as they appear in hooks JSON and config files. -pub const HOOK_EVENT_NAMES: [&str; 6] = [ +pub const HOOK_EVENT_NAMES: [&str; 7] = [ "PreToolUse", "PermissionRequest", "PostToolUse", "SessionStart", "UserPromptSubmit", "Stop", + "SubscriptionExhausted", ]; /// Hook event names whose matcher fields are meaningful during dispatch. @@ -40,6 +41,8 @@ pub use events::session_start::SessionStartRequest; pub use events::session_start::SessionStartSource; pub use events::stop::StopOutcome; pub use events::stop::StopRequest; +pub use events::subscription_exhausted::SubscriptionExhaustedOutcome; +pub use events::subscription_exhausted::SubscriptionExhaustedRequest; pub use events::user_prompt_submit::UserPromptSubmitOutcome; pub use events::user_prompt_submit::UserPromptSubmitRequest; pub use legacy_notify::legacy_notify_json; diff --git a/codex-rs/hooks/src/registry.rs b/codex-rs/hooks/src/registry.rs index ae80015729b..0f74ac8791d 100644 --- a/codex-rs/hooks/src/registry.rs +++ b/codex-rs/hooks/src/registry.rs @@ -15,6 +15,8 @@ use crate::events::session_start::SessionStartOutcome; use crate::events::session_start::SessionStartRequest; use crate::events::stop::StopOutcome; use crate::events::stop::StopRequest; +use crate::events::subscription_exhausted::SubscriptionExhaustedOutcome; +use crate::events::subscription_exhausted::SubscriptionExhaustedRequest; use crate::events::user_prompt_submit::UserPromptSubmitOutcome; use crate::events::user_prompt_submit::UserPromptSubmitRequest; use crate::types::Hook; @@ -178,6 +180,20 @@ impl Hooks { pub async fn run_stop(&self, request: StopRequest) -> StopOutcome { self.engine.run_stop(request).await } + + pub fn preview_subscription_exhausted( + &self, + request: &SubscriptionExhaustedRequest, + ) -> Vec { + self.engine.preview_subscription_exhausted(request) + } + + pub async fn run_subscription_exhausted( + &self, + request: SubscriptionExhaustedRequest, + ) -> SubscriptionExhaustedOutcome { + self.engine.run_subscription_exhausted(request).await + } } pub fn list_hooks(config: HooksConfig) -> HookListOutcome { diff --git a/codex-rs/hooks/src/schema.rs b/codex-rs/hooks/src/schema.rs index d08cce6ee29..084f4bb1886 100644 --- a/codex-rs/hooks/src/schema.rs +++ b/codex-rs/hooks/src/schema.rs @@ -25,6 +25,10 @@ const USER_PROMPT_SUBMIT_INPUT_FIXTURE: &str = "user-prompt-submit.command.input const USER_PROMPT_SUBMIT_OUTPUT_FIXTURE: &str = "user-prompt-submit.command.output.schema.json"; const STOP_INPUT_FIXTURE: &str = "stop.command.input.schema.json"; const STOP_OUTPUT_FIXTURE: &str = "stop.command.output.schema.json"; +const SUBSCRIPTION_EXHAUSTED_INPUT_FIXTURE: &str = + "subscription-exhausted.command.input.schema.json"; +const SUBSCRIPTION_EXHAUSTED_OUTPUT_FIXTURE: &str = + "subscription-exhausted.command.output.schema.json"; #[derive(Debug, Clone, Serialize)] #[serde(transparent)] @@ -81,6 +85,8 @@ pub(crate) enum HookEventNameWire { UserPromptSubmit, #[serde(rename = "Stop")] Stop, + #[serde(rename = "SubscriptionExhausted")] + SubscriptionExhausted, } #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] @@ -326,6 +332,15 @@ pub(crate) struct StopCommandOutputWire { pub reason: Option, } +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "subscription-exhausted.command.output")] +pub(crate) struct SubscriptionExhaustedCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, +} + #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)] pub(crate) enum BlockDecisionWire { #[serde(rename = "block")] @@ -404,6 +419,28 @@ pub(crate) struct StopCommandInput { pub last_assistant_message: NullableString, } +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "subscription-exhausted.command.input")] +pub(crate) struct SubscriptionExhaustedCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + pub transcript_path: NullableString, + pub cwd: String, + pub codex_home: String, + #[schemars(schema_with = "subscription_exhausted_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "permission_mode_schema")] + pub permission_mode: String, + #[schemars(schema_with = "subscription_exhausted_error_kind_schema")] + pub error_kind: String, + pub plan_type: NullableString, + pub resets_at: Option, + pub account_id: NullableString, +} + pub fn write_schema_fixtures(schema_root: &Path) -> anyhow::Result<()> { let generated_dir = schema_root.join(GENERATED_DIR); ensure_empty_dir(&generated_dir)?; @@ -456,6 +493,14 @@ pub fn write_schema_fixtures(schema_root: &Path) -> anyhow::Result<()> { &generated_dir.join(STOP_OUTPUT_FIXTURE), schema_json::()?, )?; + write_schema( + &generated_dir.join(SUBSCRIPTION_EXHAUSTED_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(SUBSCRIPTION_EXHAUSTED_OUTPUT_FIXTURE), + schema_json::()?, + )?; Ok(()) } @@ -535,6 +580,10 @@ fn stop_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { string_const_schema("Stop") } +fn subscription_exhausted_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("SubscriptionExhausted") +} + fn permission_mode_schema(_gen: &mut SchemaGenerator) -> Schema { string_enum_schema(&[ "default", @@ -549,6 +598,10 @@ fn session_start_source_schema(_gen: &mut SchemaGenerator) -> Schema { string_enum_schema(&["startup", "resume", "clear"]) } +fn subscription_exhausted_error_kind_schema(_gen: &mut SchemaGenerator) -> Schema { + string_enum_schema(&["usage_limit_reached", "quota_exceeded"]) +} + fn string_const_schema(value: &str) -> Schema { let mut schema = SchemaObject { instance_type: Some(InstanceType::String.into()), @@ -591,7 +644,10 @@ mod tests { use super::SESSION_START_OUTPUT_FIXTURE; use super::STOP_INPUT_FIXTURE; use super::STOP_OUTPUT_FIXTURE; + use super::SUBSCRIPTION_EXHAUSTED_INPUT_FIXTURE; + use super::SUBSCRIPTION_EXHAUSTED_OUTPUT_FIXTURE; use super::StopCommandInput; + use super::SubscriptionExhaustedCommandInput; use super::USER_PROMPT_SUBMIT_INPUT_FIXTURE; use super::USER_PROMPT_SUBMIT_OUTPUT_FIXTURE; use super::UserPromptSubmitCommandInput; @@ -639,6 +695,14 @@ mod tests { STOP_OUTPUT_FIXTURE => { include_str!("../schema/generated/stop.command.output.schema.json") } + SUBSCRIPTION_EXHAUSTED_INPUT_FIXTURE => { + include_str!("../schema/generated/subscription-exhausted.command.input.schema.json") + } + SUBSCRIPTION_EXHAUSTED_OUTPUT_FIXTURE => { + include_str!( + "../schema/generated/subscription-exhausted.command.output.schema.json" + ) + } _ => panic!("unexpected fixture name: {name}"), } } @@ -666,6 +730,8 @@ mod tests { USER_PROMPT_SUBMIT_OUTPUT_FIXTURE, STOP_INPUT_FIXTURE, STOP_OUTPUT_FIXTURE, + SUBSCRIPTION_EXHAUSTED_INPUT_FIXTURE, + SUBSCRIPTION_EXHAUSTED_OUTPUT_FIXTURE, ] { let expected = normalize_newlines(expected_fixture(fixture)); let actual = std::fs::read_to_string(schema_root.join("generated").join(fixture)) @@ -702,6 +768,11 @@ mod tests { &schema_json::().expect("serialize stop input schema"), ) .expect("parse stop input schema"); + let subscription_exhausted: Value = serde_json::from_slice( + &schema_json::() + .expect("serialize subscription exhausted input schema"), + ) + .expect("parse subscription exhausted input schema"); for schema in [ &pre_tool_use, @@ -709,6 +780,7 @@ mod tests { &post_tool_use, &user_prompt_submit, &stop, + &subscription_exhausted, ] { assert_eq!(schema["properties"]["turn_id"]["type"], "string"); assert!( diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 29897db7bea..2291db5dc9c 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1387,6 +1387,10 @@ impl AuthManager { self.inner.read().ok().and_then(|c| c.auth.clone()) } + pub fn codex_home(&self) -> &Path { + &self.codex_home + } + pub fn refresh_failure_for_auth(&self, auth: &CodexAuth) -> Option { self.inner.read().ok().and_then(|cached| { cached diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 95f61eba880..dbb6cb4881c 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1519,6 +1519,7 @@ pub enum HookEventName { SessionStart, UserPromptSubmit, Stop, + SubscriptionExhausted, } #[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] diff --git a/codex-rs/tui/src/bottom_pane/hooks_browser_view.rs b/codex-rs/tui/src/bottom_pane/hooks_browser_view.rs index 2f4c6a8a0d4..f756a0d0e41 100644 --- a/codex-rs/tui/src/bottom_pane/hooks_browser_view.rs +++ b/codex-rs/tui/src/bottom_pane/hooks_browser_view.rs @@ -539,6 +539,7 @@ fn event_label(event_name: HookEventName) -> &'static str { HookEventName::SessionStart => "SessionStart", HookEventName::UserPromptSubmit => "UserPromptSubmit", HookEventName::Stop => "Stop", + HookEventName::SubscriptionExhausted => "SubscriptionExhausted", } } @@ -550,6 +551,7 @@ fn event_description(event_name: HookEventName) -> &'static str { HookEventName::SessionStart => "When a new session starts", HookEventName::UserPromptSubmit => "When the user submits a prompt", HookEventName::Stop => "Right before Codex ends its turn", + HookEventName::SubscriptionExhausted => "When a subscription or quota is exhausted", } } diff --git a/codex-rs/tui/src/chatwidget/tests/helpers.rs b/codex-rs/tui/src/chatwidget/tests/helpers.rs index 04f7e3d9071..59993dc13d6 100644 --- a/codex-rs/tui/src/chatwidget/tests/helpers.rs +++ b/codex-rs/tui/src/chatwidget/tests/helpers.rs @@ -1656,5 +1656,6 @@ fn hook_event_label(event_name: codex_app_server_protocol::HookEventName) -> &'s codex_app_server_protocol::HookEventName::SessionStart => "SessionStart", codex_app_server_protocol::HookEventName::UserPromptSubmit => "UserPromptSubmit", codex_app_server_protocol::HookEventName::Stop => "Stop", + codex_app_server_protocol::HookEventName::SubscriptionExhausted => "SubscriptionExhausted", } } diff --git a/codex-rs/tui/src/history_cell/hook_cell.rs b/codex-rs/tui/src/history_cell/hook_cell.rs index ec5f1ca18e2..b6cec1f68f1 100644 --- a/codex-rs/tui/src/history_cell/hook_cell.rs +++ b/codex-rs/tui/src/history_cell/hook_cell.rs @@ -714,6 +714,7 @@ fn hook_event_label(event_name: HookEventName) -> &'static str { HookEventName::SessionStart => "SessionStart", HookEventName::UserPromptSubmit => "UserPromptSubmit", HookEventName::Stop => "Stop", + HookEventName::SubscriptionExhausted => "SubscriptionExhausted", } }