| 1 | //! Programmable admission listeners reuse the native strict hook fold. |
| 2 | //! The host returns proposals, never tool or approval authority. |
| 3 | |
| 4 | use std::time::{Duration, Instant}; |
| 5 | |
| 6 | use serde_json::{Value, json}; |
| 7 | |
| 8 | use super::protocol::{CoreRequest, HookCallPayload, HookEvaluateParams, HookVerdictWire}; |
| 9 | use super::{HostAttachment, ManagerShared}; |
| 10 | use crate::hooks::HookResult; |
| 11 | |
| 12 | const HOOK_DEADLINE: Duration = Duration::from_secs(5); |
| 13 | |
| 14 | impl HostAttachment { |
| 15 | /// Only the reviewed owners this engine's own snapshot desires see its |
| 16 | /// call. A shared host never broadcasts another workspace's arguments. |
| 17 | pub(crate) async fn tool_before_hooks(&self, mut payload: HookCallPayload) -> Vec<HookResult> { |
| 18 | let shared = &self.manager.shared; |
| 19 | let desired = shared |
| 20 | .attachments |
| 21 | .lock() |
| 22 | .expect("attachments lock") |
| 23 | .get(&self.id) |
| 24 | .map(|state| state.selection.clone()) |
| 25 | .unwrap_or_default(); |
| 26 | let hooks: Vec<_> = shared |
| 27 | .registry |
| 28 | .lock() |
| 29 | .expect("registry lock") |
| 30 | .live_hooks() |
| 31 | .into_iter() |
| 32 | .filter(|hook| { |
| 33 | desired.includes( |
| 34 | &hook.owner.plugin_id, |
| 35 | &hook.content_hash, |
| 36 | hook.scope.as_ref(), |
| 37 | ) |
| 38 | }) |
| 39 | .collect(); |
| 40 | let mut results = Vec::with_capacity(hooks.len()); |
| 41 | let batch_started = Instant::now(); |
| 42 | for hook in hooks { |
| 43 | let started = Instant::now(); |
| 44 | let answer = async { |
| 45 | let host = shared.live_host_for_hook(&hook).await?; |
| 46 | let still_desired = desired.revision.is_some_and(|revision| { |
| 47 | shared.selection_current( |
| 48 | revision, |
| 49 | &hook.owner.plugin_id, |
| 50 | &hook.content_hash, |
| 51 | hook.scope.as_ref(), |
| 52 | ) |
| 53 | }); |
| 54 | if !still_desired { |
| 55 | return Err("extension hook was withdrawn".to_string()); |
| 56 | } |
| 57 | let remaining_ms = HOOK_DEADLINE |
| 58 | .saturating_sub(batch_started.elapsed()) |
| 59 | .as_millis() as u64; |
| 60 | if remaining_ms == 0 { |
| 61 | return Err("extension hook batch deadline elapsed".to_string()); |
| 62 | } |
| 63 | let value = host |
| 64 | .call( |
| 65 | CoreRequest::HookEvaluate(HookEvaluateParams { |
| 66 | handle: hook.handle, |
| 67 | event: hook.event.clone(), |
| 68 | payload: payload.clone(), |
| 69 | deadline_ms: remaining_ms, |
| 70 | }), |
| 71 | Some(hook.owner.plugin_id.clone()), |
| 72 | ) |
| 73 | .await |
| 74 | .map_err(|_| "extension hook did not answer".to_string())?; |
| 75 | // A late response cannot restore an owner revoked while the |
| 76 | // callback waited. Revalidate its receipt and generation too. |
| 77 | shared.live_host_for_hook(&hook).await?; |
| 78 | let still_desired = desired.revision.is_some_and(|revision| { |
| 79 | shared.selection_current( |
| 80 | revision, |
| 81 | &hook.owner.plugin_id, |
| 82 | &hook.content_hash, |
| 83 | hook.scope.as_ref(), |
| 84 | ) |
| 85 | }); |
| 86 | if !still_desired { |
| 87 | return Err("extension hook was withdrawn".to_string()); |
| 88 | } |
| 89 | let verdict: HookVerdictWire = serde_json::from_value(value) |
| 90 | .map_err(|_| "extension hook returned a malformed verdict".to_string())?; |
| 91 | let stdout = verdict_stdout(verdict.clone())?; |
| 92 | if let HookVerdictWire::Revise { input } = verdict { |
| 93 | payload.input = Value::Object(input); |
| 94 | } |
| 95 | Ok(stdout) |
| 96 | } |
| 97 | .await; |
| 98 | let answered = answer.is_ok(); |
| 99 | let denies = answer.as_ref().is_ok_and(|stdout| { |
| 100 | crate::hooks::parse_tool_call_before_stdout(stdout).decision |
| 101 | == Some(crate::hooks::ToolCallDecision::Deny) |
| 102 | }); |
| 103 | results.push(HookResult { |
| 104 | name: Some(format!("extension:{}:pre-execute", hook.plugin_name)), |
| 105 | success: answered, |
| 106 | background: false, |
| 107 | strict: true, |
| 108 | exit_code: answered.then_some(0), |
| 109 | stdout: answer.unwrap_or_default(), |
| 110 | stderr: String::new(), |
| 111 | duration: started.elapsed(), |
| 112 | error: (!answered).then(|| "extension hook returned no valid verdict".to_string()), |
| 113 | }); |
| 114 | if !answered || denies { |
| 115 | break; |
| 116 | } |
| 117 | } |
| 118 | results |
| 119 | } |
| 120 | } |
| 121 | |
| 122 | impl ManagerShared { |
| 123 | async fn live_host_for_hook( |
| 124 | &self, |
| 125 | hook: &super::registry::HookRegistration, |
| 126 | ) -> Result<std::sync::Arc<super::supervisor::HostProcess>, String> { |
| 127 | self.live_host(hook.tier, |registry| { |
| 128 | registry |
| 129 | .is_live_hook(hook.handle, &hook.owner) |
| 130 | .then(|| hook.owner.clone()) |
| 131 | .ok_or_else(|| "extension hook is no longer registered".to_string()) |
| 132 | }) |
| 133 | .await |
| 134 | } |
| 135 | } |
| 136 | |
| 137 | fn verdict_stdout(verdict: HookVerdictWire) -> Result<String, String> { |
| 138 | let revises = matches!(verdict, HookVerdictWire::Revise { .. }); |
| 139 | let value = match verdict { |
| 140 | HookVerdictWire::Abstain => json!({}), |
| 141 | HookVerdictWire::Deny { reason } => json!({"decision":"deny", "reason":reason}), |
| 142 | HookVerdictWire::Ask { reason } => json!({"decision":"ask", "reason":reason}), |
| 143 | HookVerdictWire::Annotate { text } => json!({"additionalContext":text}), |
| 144 | HookVerdictWire::Revise { input } => json!({"updatedInput":Value::Object(input)}), |
| 145 | }; |
| 146 | let stdout = serde_json::to_string(&value) |
| 147 | .map_err(|_| "extension hook verdict is not JSON".to_string())?; |
| 148 | // Reuse the existing validation and size limit. A rejected revision is |
| 149 | // a strict no-verdict, never a silently ignored plugin proposal. |
| 150 | if revises |
| 151 | && crate::hooks::parse_tool_call_before_stdout(&stdout) |
| 152 | .updated_input |
| 153 | .is_none() |
| 154 | { |
| 155 | return Err("extension hook input revision exceeded the limit".to_string()); |
| 156 | } |
| 157 | Ok(stdout) |
| 158 | } |
| 159 | |
| 160 | #[cfg(test)] |
| 161 | mod tests { |
| 162 | use super::*; |
| 163 | use crate::core::engine::turn_loop::run_tool_call_before_hooks; |
| 164 | use crate::extension_host::tests::{FixturePlugins, node_for_tests}; |
| 165 | use codewhale_config::AppMode; |
| 166 | use std::sync::Arc; |
| 167 | |
| 168 | #[tokio::test] |
| 169 | async fn real_host_hook_timeout_and_owner_revocation_fail_closed_and_do_not_cross_workspaces() { |
| 170 | let Some(node) = node_for_tests("real_host_hook_timeout_and_revocation") else { |
| 171 | return; |
| 172 | }; |
| 173 | let _policy = crate::plugins::activation::TestPolicyGuard::extension_host(true); |
| 174 | let fixture = FixturePlugins::new(&["hook-policy"]).await; |
| 175 | let manager = Arc::new(super::super::ExtensionHostManager::new( |
| 176 | super::super::ExtensionHostOptions { |
| 177 | runtime: crate::config::ExtensionHostRuntime::Node, |
| 178 | node_override: Some(node), |
| 179 | root: Some(fixture.root.clone()), |
| 180 | supervision: super::super::SupervisionOptions { |
| 181 | heartbeat_interval: Duration::from_secs(60), |
| 182 | ..Default::default() |
| 183 | }, |
| 184 | ..Default::default() |
| 185 | }, |
| 186 | )); |
| 187 | let source = Arc::new(manager.attach(fixture.registry())); |
| 188 | source.sync().await.unwrap(); |
| 189 | let timeout = run_tool_call_before_hooks( |
| 190 | None, |
| 191 | Some(&source), |
| 192 | "read", |
| 193 | "timeout", |
| 194 | &json!({"path":"timeout.txt"}), |
| 195 | AppMode::Agent, |
| 196 | fixture.workspace(), |
| 197 | "fixture", |
| 198 | ) |
| 199 | .await; |
| 200 | assert!( |
| 201 | timeout |
| 202 | .unwrap_err() |
| 203 | .to_string() |
| 204 | .contains("returned no verdict") |
| 205 | ); |
| 206 | |
| 207 | let other_workspace = tempfile::tempdir().unwrap(); |
| 208 | let other = manager.attach(Arc::new(crate::plugins::PluginRegistry::empty( |
| 209 | other_workspace.path(), |
| 210 | ))); |
| 211 | other.sync().await.unwrap(); |
| 212 | let unrelated = run_tool_call_before_hooks( |
| 213 | None, |
| 214 | Some(&other), |
| 215 | "read", |
| 216 | "other", |
| 217 | &json!({"path":"blocked.txt"}), |
| 218 | AppMode::Agent, |
| 219 | other_workspace.path(), |
| 220 | "fixture", |
| 221 | ) |
| 222 | .await |
| 223 | .unwrap(); |
| 224 | assert_eq!( |
| 225 | unrelated, |
| 226 | Default::default(), |
| 227 | "a workspace does not receive another workspace's listener" |
| 228 | ); |
| 229 | |
| 230 | let sent = manager.host_requests_started().unwrap(); |
| 231 | let pending_source = Arc::clone(&source); |
| 232 | let workspace = fixture.workspace().to_path_buf(); |
| 233 | let pending = tokio::spawn(async move { |
| 234 | run_tool_call_before_hooks( |
| 235 | None, |
| 236 | Some(&pending_source), |
| 237 | "read", |
| 238 | "held", |
| 239 | &json!({"path":"held.txt"}), |
| 240 | AppMode::Agent, |
| 241 | &workspace, |
| 242 | "fixture", |
| 243 | ) |
| 244 | .await |
| 245 | }); |
| 246 | tokio::time::timeout(Duration::from_secs(2), async { |
| 247 | while manager.host_requests_started() == Some(sent) { |
| 248 | tokio::task::yield_now().await; |
| 249 | } |
| 250 | }) |
| 251 | .await |
| 252 | .unwrap(); |
| 253 | source.set_plugins(fixture.disable("hook-policy")); |
| 254 | source.sync().await.unwrap(); |
| 255 | let error = tokio::time::timeout(Duration::from_secs(2), pending) |
| 256 | .await |
| 257 | .expect("revocation cancels a held hook") |
| 258 | .unwrap() |
| 259 | .unwrap_err(); |
| 260 | assert!(error.to_string().contains("returned no verdict")); |
| 261 | let after = run_tool_call_before_hooks( |
| 262 | None, |
| 263 | Some(&source), |
| 264 | "read", |
| 265 | "after", |
| 266 | &json!({"path":"before.txt"}), |
| 267 | AppMode::Agent, |
| 268 | fixture.workspace(), |
| 269 | "fixture", |
| 270 | ) |
| 271 | .await |
| 272 | .unwrap(); |
| 273 | assert_eq!(after, Default::default()); |
| 274 | manager.shutdown().await; |
| 275 | } |
| 276 | } |
| 277 |