| 1 | //! Caller snapshots in AttachmentState. No second registry, store or owner token. |
| 2 | use super::protocol::EntryRef; |
| 3 | use serde::{Deserialize, Serialize}; |
| 4 | use std::collections::BTreeMap; |
| 5 | |
| 6 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
| 7 | #[serde(deny_unknown_fields)] |
| 8 | pub struct SelectionRevision { |
| 9 | pub attachment_id: u64, |
| 10 | pub revision: u64, |
| 11 | } |
| 12 | |
| 13 | /// A core-selected entry from the owner's existing reviewed Native inventory. |
| 14 | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] |
| 15 | #[serde(deny_unknown_fields)] |
| 16 | pub struct NativePresetRef { |
| 17 | pub plugin_id: String, |
| 18 | pub content_hash: String, |
| 19 | pub entry: EntryRef, |
| 20 | } |
| 21 | |
| 22 | #[derive(Debug, Clone, Default)] |
| 23 | pub struct CompositionSelection { |
| 24 | pub revision: Option<SelectionRevision>, |
| 25 | pub desired: BTreeMap<String, String>, |
| 26 | pub entries: Vec<NativePresetRef>, |
| 27 | } |
| 28 | impl CompositionSelection { |
| 29 | pub fn includes(&self, plugin_id: &str, content_hash: &str, scope: Option<&EntryRef>) -> bool { |
| 30 | self.desired |
| 31 | .get(plugin_id) |
| 32 | .is_some_and(|hash| hash == content_hash) |
| 33 | && scope.is_none_or(|scope| { |
| 34 | self.entries.iter().any(|entry| { |
| 35 | entry.plugin_id == plugin_id |
| 36 | && entry.content_hash == content_hash |
| 37 | && entry.entry == *scope |
| 38 | }) |
| 39 | }) |
| 40 | } |
| 41 | } |
| 42 | |
| 43 | /// Separate scopes may spell a name alike; core names remain fenced unconditionally. |
| 44 | pub(super) fn name_key(owner: &str, name: &str, scope: Option<&EntryRef>) -> String { |
| 45 | match scope { |
| 46 | None => name.to_string(), |
| 47 | Some(scope) => format!("{owner}\0{}\0{}\0{name}", scope.path, scope.sha256), |
| 48 | } |
| 49 | } |
| 50 | |
| 51 | #[cfg(test)] |
| 52 | mod tests { |
| 53 | use super::*; |
| 54 | use crate::extension_host::protocol::{ |
| 55 | OwnerRef, RegisterKind, RegisterParams, RegisterSpecWire, |
| 56 | }; |
| 57 | use crate::extension_host::registry::{MAX_TOOLS_PER_OWNER, OwnerRegistry}; |
| 58 | use crate::extension_host::tests::fake_authority; |
| 59 | use crate::extension_host::tier::HostTier; |
| 60 | use crate::extension_host::{ExtensionHostManager, ExtensionHostOptions}; |
| 61 | use crate::plugins::PluginRegistry; |
| 62 | use std::sync::Arc; |
| 63 | fn scope(name: &str) -> EntryRef { |
| 64 | EntryRef { |
| 65 | path: format!("/reviewed/{name}.mjs"), |
| 66 | sha256: name.repeat(64), |
| 67 | } |
| 68 | } |
| 69 | fn tool(owner: &OwnerRef, scope: Option<EntryRef>, name: &str) -> RegisterParams { |
| 70 | RegisterParams { |
| 71 | owner: owner.clone(), |
| 72 | scope, |
| 73 | kind: RegisterKind::Tool, |
| 74 | spec: RegisterSpecWire { |
| 75 | name: name.into(), |
| 76 | description: "scoped fixture".into(), |
| 77 | input_schema: Some( |
| 78 | serde_json::json!({"type":"object","properties":{}}) |
| 79 | .as_object() |
| 80 | .unwrap() |
| 81 | .clone(), |
| 82 | ), |
| 83 | argument_hint: None, |
| 84 | }, |
| 85 | } |
| 86 | } |
| 87 | fn owner(registry: &mut OwnerRegistry) -> OwnerRef { |
| 88 | registry |
| 89 | .begin_owner( |
| 90 | HostTier::Plugin, |
| 91 | "local/scoped", |
| 92 | "scoped", |
| 93 | Some(fake_authority("local/scoped")), |
| 94 | "build", |
| 95 | ) |
| 96 | .unwrap() |
| 97 | } |
| 98 | |
| 99 | #[test] |
| 100 | fn scoped_name_identity_includes_owner_even_for_identical_staged_entry_paths() { |
| 101 | let entry = scope("a"); |
| 102 | assert_ne!( |
| 103 | name_key("workspace-a", "echo", Some(&entry)), |
| 104 | name_key("workspace-b", "echo", Some(&entry)) |
| 105 | ); |
| 106 | assert_eq!(name_key("a", "core", None), name_key("b", "core", None)); |
| 107 | } |
| 108 | #[test] |
| 109 | fn composition_selection_requires_exact_owner_build_and_entry() { |
| 110 | let a = scope("a"); |
| 111 | let b = scope("b"); |
| 112 | let selected = CompositionSelection { |
| 113 | revision: Some(SelectionRevision { |
| 114 | attachment_id: 1, |
| 115 | revision: 2, |
| 116 | }), |
| 117 | desired: BTreeMap::from([("owner".into(), "build".into())]), |
| 118 | entries: vec![NativePresetRef { |
| 119 | plugin_id: "owner".into(), |
| 120 | content_hash: "build".into(), |
| 121 | entry: a.clone(), |
| 122 | }], |
| 123 | }; |
| 124 | assert!(selected.includes("owner", "build", Some(&a))); |
| 125 | assert!(!selected.includes("owner", "build", Some(&b))); |
| 126 | assert!(!selected.includes("foreign", "build", Some(&a))); |
| 127 | assert!(!selected.includes("owner", "changed", Some(&a))); |
| 128 | } |
| 129 | #[test] |
| 130 | fn scoped_equal_names_keep_distinct_handles_and_withdraw_exact_entry() { |
| 131 | let mut registry = OwnerRegistry::default(); |
| 132 | let owner = owner(&mut registry); |
| 133 | let a = scope("a"); |
| 134 | let b = scope("b"); |
| 135 | registry.begin_scope(&owner, a.clone()).unwrap(); |
| 136 | registry.begin_scope(&owner, b.clone()).unwrap(); |
| 137 | let first = registry |
| 138 | .register(&tool(&owner, Some(a.clone()), "scoped_echo")) |
| 139 | .unwrap(); |
| 140 | let second = registry |
| 141 | .register(&tool(&owner, Some(b.clone()), "scoped_echo")) |
| 142 | .unwrap(); |
| 143 | assert_ne!(first, second); |
| 144 | assert!(registry.mark_scope_active(&owner, &a)); |
| 145 | assert!(registry.mark_scope_active(&owner, &b)); |
| 146 | assert!(registry.mark_active(&owner)); |
| 147 | assert_eq!(registry.live_tools().len(), 2); |
| 148 | assert!( |
| 149 | registry |
| 150 | .register(&tool(&owner, None, "unscoped_escape")) |
| 151 | .is_err() |
| 152 | ); |
| 153 | let retired = registry.revoke_scope(&owner, &a); |
| 154 | assert_eq!(retired, [first]); |
| 155 | assert!(!registry.is_live(first, &owner)); |
| 156 | assert!(registry.is_live(second, &owner)); |
| 157 | registry.host_exited(HostTier::Builtin, "fixture"); |
| 158 | assert!(registry.is_live(second, &owner)); |
| 159 | registry.host_exited(HostTier::Plugin, "fixture"); |
| 160 | assert!(registry.live_tools().is_empty()); |
| 161 | assert!(registry.owner(&owner.plugin_id).is_none()); |
| 162 | } |
| 163 | #[test] |
| 164 | fn scoped_aggregate_owner_limit_cannot_be_multiplied_by_entries() { |
| 165 | let mut registry = OwnerRegistry::default(); |
| 166 | let owner = owner(&mut registry); |
| 167 | let a = scope("a"); |
| 168 | let b = scope("b"); |
| 169 | registry.begin_scope(&owner, a.clone()).unwrap(); |
| 170 | registry.begin_scope(&owner, b.clone()).unwrap(); |
| 171 | for index in 0..MAX_TOOLS_PER_OWNER { |
| 172 | registry |
| 173 | .register(&tool( |
| 174 | &owner, |
| 175 | Some(if index % 2 == 0 { a.clone() } else { b.clone() }), |
| 176 | &format!("scoped_{index}"), |
| 177 | )) |
| 178 | .unwrap(); |
| 179 | } |
| 180 | assert!( |
| 181 | registry |
| 182 | .register(&tool(&owner, Some(b), "one_too_many")) |
| 183 | .unwrap_err() |
| 184 | .contains("at most") |
| 185 | ); |
| 186 | } |
| 187 | #[test] |
| 188 | fn caller_revision_withdrawal_does_not_revoke_another_attachment() { |
| 189 | let manager = Arc::new(ExtensionHostManager::new(ExtensionHostOptions::default())); |
| 190 | let one = manager.attach(Arc::new(PluginRegistry::empty(std::path::Path::new( |
| 191 | "/workspace", |
| 192 | )))); |
| 193 | let two = manager.attach(Arc::new(PluginRegistry::empty(std::path::Path::new( |
| 194 | "/workspace", |
| 195 | )))); |
| 196 | let entry = scope("a"); |
| 197 | let selected = NativePresetRef { |
| 198 | plugin_id: "owner".into(), |
| 199 | content_hash: "build".into(), |
| 200 | entry: entry.clone(), |
| 201 | }; |
| 202 | { |
| 203 | let mut attachments = manager.shared.attachments.lock().unwrap(); |
| 204 | for state in attachments.values_mut() { |
| 205 | state.selection = CompositionSelection { |
| 206 | revision: state.plugins.caller_selection(), |
| 207 | desired: BTreeMap::from([("owner".into(), "build".into())]), |
| 208 | entries: vec![selected.clone()], |
| 209 | }; |
| 210 | } |
| 211 | } |
| 212 | let old = one.plugin_view(); |
| 213 | let other = two.plugin_view(); |
| 214 | assert!( |
| 215 | manager |
| 216 | .shared |
| 217 | .check_selection( |
| 218 | old.caller_selection(), |
| 219 | Some(&old), |
| 220 | "owner", |
| 221 | "build", |
| 222 | Some(&entry) |
| 223 | ) |
| 224 | .is_ok() |
| 225 | ); |
| 226 | assert!( |
| 227 | manager |
| 228 | .shared |
| 229 | .check_selection( |
| 230 | old.caller_selection(), |
| 231 | Some(&other), |
| 232 | "owner", |
| 233 | "build", |
| 234 | Some(&entry) |
| 235 | ) |
| 236 | .is_err() |
| 237 | ); |
| 238 | one.set_plugins(Arc::new(PluginRegistry::empty(std::path::Path::new( |
| 239 | "/workspace", |
| 240 | )))); |
| 241 | assert!( |
| 242 | manager |
| 243 | .shared |
| 244 | .check_selection( |
| 245 | old.caller_selection(), |
| 246 | Some(&old), |
| 247 | "owner", |
| 248 | "build", |
| 249 | Some(&entry) |
| 250 | ) |
| 251 | .is_err() |
| 252 | ); |
| 253 | assert!( |
| 254 | manager |
| 255 | .shared |
| 256 | .check_selection( |
| 257 | other.caller_selection(), |
| 258 | Some(&other), |
| 259 | "owner", |
| 260 | "build", |
| 261 | Some(&entry) |
| 262 | ) |
| 263 | .is_ok() |
| 264 | ); |
| 265 | } |
| 266 | #[test] |
| 267 | fn failed_scope_requires_reload_and_explicit_null_cannot_erase_scope() { |
| 268 | let mut registry = OwnerRegistry::default(); |
| 269 | let owner = owner(&mut registry); |
| 270 | let a = scope("a"); |
| 271 | registry.begin_scope(&owner, a.clone()).unwrap(); |
| 272 | registry.fail_scope(&owner, &a); |
| 273 | assert!(registry.begin_scope(&owner, a.clone()).is_err()); |
| 274 | registry.forget_inactive(); |
| 275 | assert!(registry.begin_scope(&owner, a.clone()).is_ok()); |
| 276 | let mut wire = serde_json::to_value(tool(&owner, Some(a), "scoped_fixture")).unwrap(); |
| 277 | wire["scope"] = serde_json::Value::Null; |
| 278 | assert!(serde_json::from_value::<RegisterParams>(wire).is_err()); |
| 279 | } |
| 280 | } |
| 281 |