diff --git a/architecture/sandbox.md b/architecture/sandbox.md index 580d8f96d8..30d1bf4783 100644 --- a/architecture/sandbox.md +++ b/architecture/sandbox.md @@ -171,6 +171,38 @@ the structured 403 and authors the narrowest rule. Mechanistically mapping L7 would either over-broaden rules or require path-templating logic that rots quickly. +## Policy Revision Acknowledgement + +When the supervisor loads a sandbox-scoped policy from the gateway, it retains +the version, hash, source, and configuration revision returned with that exact +policy snapshot. After the OPA engine is built successfully, the supervisor +reports that revision as `LOADED`, which advances +`SandboxStatus.current_policy_version` and moves the revision out of `Pending`. +If policy construction fails, it reports the captured revision as `FAILED` with +the original construction error. It never infers revision identity by comparing +policy structure. + +This holds even when the initial policy is enriched with baseline paths during +startup: the enriched revision the supervisor synced back to the gateway is the +revision it acknowledges, so a successfully constructed initial policy never +remains `Pending`. If the first poll returns a different revision, the supervisor +processes it through the normal reload path instead of treating it as already +loaded. + +Policy status delivery uses a FIFO background worker. Retryable delivery +failures retain the ordered update and retry with capped exponential backoff; +terminal errors are logged and discarded. The outbox is nonblocking and does +not discard updates because of a fixed queue capacity, so status endpoint +outages cannot block policy polling, enforcement, settings, or provider +refreshes and cannot permanently lose the initial acknowledgement. + +Only sandbox-scoped revisions (`PolicySource::Sandbox`, version greater than +zero) are acknowledged. Global policies and local-file development policies do +not use the sandbox revision API and produce no acknowledgement. When explicit +local Rego and data files are configured, the supervisor continues polling the +gateway for settings and provider refreshes but never replaces the local OPA +engine with a gateway policy revision. + ## Failure Behavior - If gateway config polling fails, the sandbox keeps its last-known-good policy. diff --git a/crates/openshell-cli/tests/mtls_integration.rs b/crates/openshell-cli/tests/mtls_integration.rs index 28a4e6c9c3..3f51dd6044 100644 --- a/crates/openshell-cli/tests/mtls_integration.rs +++ b/crates/openshell-cli/tests/mtls_integration.rs @@ -668,6 +668,7 @@ async fn gateway_add_mtls_loopback_explicit_name_does_not_fallback_to_openshell_ #[tokio::test] async fn cli_connects_with_client_cert() { + let _env = EnvVarGuard::set(&[]); install_rustls_provider(); let (ca, ca_key) = build_ca(); @@ -741,6 +742,7 @@ async fn run_server_no_client_auth( #[tokio::test] async fn cli_connects_with_gateway_insecure() { + let _env = EnvVarGuard::set(&[]); install_rustls_provider(); let (ca, ca_key) = build_ca(); diff --git a/crates/openshell-core/src/grpc_client.rs b/crates/openshell-core/src/grpc_client.rs index 96158a1d10..57f6bdd816 100644 --- a/crates/openshell-core/src/grpc_client.rs +++ b/crates/openshell-core/src/grpc_client.rs @@ -573,11 +573,23 @@ pub async fn fetch_policy(endpoint: &str, sandbox_id: &str) -> Result Result { + debug!(endpoint = %endpoint, sandbox_id = %sandbox_id, "Connecting to fetch OpenShell settings snapshot"); + let mut client = connect(endpoint).await?; + fetch_settings_snapshot_with_client(&mut client, sandbox_id).await +} + +async fn fetch_settings_snapshot_with_client( client: &mut OpenShellClient, sandbox_id: &str, -) -> Result> { +) -> Result { let response = client .get_sandbox_config(GetSandboxConfigRequest { sandbox_id: sandbox_id.to_string(), @@ -585,14 +597,22 @@ async fn fetch_policy_with_client( .await .into_diagnostic()?; - let inner = response.into_inner(); + Ok(settings_poll_result(response.into_inner())) +} + +/// Fetch sandbox policy using an existing client connection. +async fn fetch_policy_with_client( + client: &mut OpenShellClient, + sandbox_id: &str, +) -> Result> { + let snapshot = fetch_settings_snapshot_with_client(client, sandbox_id).await?; // version 0 with no policy means the sandbox was created without one. - if inner.version == 0 && inner.policy.is_none() { + if snapshot.version == 0 && snapshot.policy.is_none() { return Ok(None); } - Ok(Some(inner.policy.ok_or_else(|| { + Ok(Some(snapshot.policy.ok_or_else(|| { miette::miette!("Server returned non-zero version but empty policy") })?)) } @@ -661,6 +681,18 @@ pub async fn sync_policy(endpoint: &str, sandbox: &str, policy: &ProtoSandboxPol sync_policy_with_client(&mut client, sandbox, policy).await } +/// Sync an enriched policy and return the authoritative revision snapshot. +pub async fn sync_policy_and_fetch_snapshot( + endpoint: &str, + sandbox_id: &str, + sandbox: &str, + policy: &ProtoSandboxPolicy, +) -> Result { + let mut client = connect(endpoint).await?; + sync_policy_with_client(&mut client, sandbox, policy).await?; + fetch_settings_snapshot_with_client(&mut client, sandbox_id).await +} + /// Fetch provider environment variables for a sandbox from `OpenShell` server via gRPC. /// /// Returns a map of environment variable names to values derived from provider @@ -700,6 +732,7 @@ pub struct CachedOpenShellClient { } /// Settings poll result returned by [`CachedOpenShellClient::poll_settings`]. +#[derive(Clone, Debug)] pub struct SettingsPollResult { pub policy: Option, pub version: u32, @@ -713,6 +746,20 @@ pub struct SettingsPollResult { pub provider_env_revision: u64, } +fn settings_poll_result(inner: crate::proto::GetSandboxConfigResponse) -> SettingsPollResult { + SettingsPollResult { + policy: inner.policy, + version: inner.version, + policy_hash: inner.policy_hash, + config_revision: inner.config_revision, + policy_source: PolicySource::try_from(inner.policy_source) + .unwrap_or(PolicySource::Unspecified), + settings: inner.settings, + global_policy_version: inner.global_policy_version, + provider_env_revision: inner.provider_env_revision, + } +} + pub struct ProviderEnvironmentResult { pub environment: HashMap, pub provider_env_revision: u64, @@ -743,19 +790,7 @@ impl CachedOpenShellClient { .await .into_diagnostic()?; - let inner = response.into_inner(); - - Ok(SettingsPollResult { - policy: inner.policy, - version: inner.version, - policy_hash: inner.policy_hash, - config_revision: inner.config_revision, - policy_source: PolicySource::try_from(inner.policy_source) - .unwrap_or(PolicySource::Unspecified), - settings: inner.settings, - global_policy_version: inner.global_policy_version, - provider_env_revision: inner.provider_env_revision, - }) + Ok(settings_poll_result(response.into_inner())) } /// Submit denial summaries and/or agent-authored proposals for policy analysis. diff --git a/crates/openshell-sandbox/src/lib.rs b/crates/openshell-sandbox/src/lib.rs index 53b1eba582..cf193c63ce 100644 --- a/crates/openshell-sandbox/src/lib.rs +++ b/crates/openshell-sandbox/src/lib.rs @@ -128,7 +128,7 @@ pub async fn run_sandbox( // Load policy and initialize OPA engine let openshell_endpoint_for_proxy = openshell_endpoint.clone(); let sandbox_name_for_agg = sandbox.clone(); - let (mut policy, opa_engine, retained_proto) = load_policy( + let (mut policy, opa_engine, retained_proto, loaded_policy_origin) = load_policy( sandbox_id.clone(), sandbox, openshell_endpoint.clone(), @@ -418,6 +418,7 @@ pub async fn run_sandbox( endpoint: poll_endpoint, sandbox_id: poll_id, opa_engine: poll_engine, + loaded_policy_origin, entrypoint_pid: poll_pid, interval_secs: poll_interval_secs, ocsf_enabled: poll_ocsf_enabled, @@ -1370,6 +1371,7 @@ async fn load_policy( SandboxPolicy, Option>, Option, + LoadedPolicyOrigin, )> { // File mode: load OPA engine from rego rules + YAML data (dev override) if let (Some(policy_file), Some(data_file)) = (&policy_rules, &policy_data) { @@ -1399,7 +1401,12 @@ async fn load_policy( process: config.process, }; enrich_sandbox_baseline_paths(&mut policy); - return Ok((policy, Some(Arc::new(engine)), None)); + return Ok(( + policy, + Some(Arc::new(engine)), + None, + LoadedPolicyOrigin::LocalOverride, + )); } // gRPC mode: fetch typed proto policy, construct OPA engine from baked rules + proto data @@ -1409,12 +1416,12 @@ async fn load_policy( endpoint = %endpoint, "Fetching sandbox policy via gRPC" ); - let proto_policy = grpc_retry("Policy fetch", || { - openshell_core::grpc_client::fetch_policy(endpoint, id) + let mut snapshot = grpc_retry("Policy fetch", || { + openshell_core::grpc_client::fetch_settings_snapshot(endpoint, id) }) .await?; - let mut proto_policy = if let Some(p) = proto_policy { + let mut proto_policy = if let Some(p) = snapshot.policy.clone() { p } else { // No policy configured on the server. Discover from disk or @@ -1442,33 +1449,68 @@ async fn load_policy( // Sync and re-fetch over a single connection to avoid extra // TLS handshakes. - grpc_retry("Policy discovery sync", || { - openshell_core::grpc_client::discover_and_sync_policy( + snapshot = grpc_retry("Policy discovery sync", || { + openshell_core::grpc_client::sync_policy_and_fetch_snapshot( endpoint, id, sandbox, &discovered, ) }) - .await? + .await?; + snapshot.policy.clone().ok_or_else(|| { + miette::miette!("Server still returned no policy after sync — this is a bug") + })? }; + // True only while `snapshot` describes the exact policy that will be + // constructed below. If enrichment cannot be synced and re-fetched, + // the policy remains enforceable but cannot be acknowledged by + // inferred structural equality. + let mut policy_bound_to_snapshot = true; + // Ensure baseline filesystem paths are present for proxy-mode // sandboxes. If the policy was enriched, sync the updated version // back to the gateway so users can see the effective policy. let enriched = enrich_proto_baseline_paths(&mut proto_policy); let sync_policy = proto_sync_payload_for_enriched_policy(&proto_policy, enriched); - if let Some(sync_policy) = sync_policy - && let Some(sandbox_name) = sandbox.as_deref() - && let Err(e) = - openshell_core::grpc_client::sync_policy(endpoint, sandbox_name, &sync_policy).await - { - warn!( - error = %e, - "Failed to sync enriched policy back to gateway (non-fatal)" - ); + if let Some(sync_policy) = sync_policy { + if let Some(sandbox_name) = sandbox.as_deref() { + match openshell_core::grpc_client::sync_policy_and_fetch_snapshot( + endpoint, + id, + sandbox_name, + &sync_policy, + ) + .await + { + Ok(canonical) => { + if let Some(policy) = canonical.policy.clone() { + proto_policy = policy; + snapshot = canonical; + } else { + policy_bound_to_snapshot = false; + warn!( + "Gateway returned no policy after enrichment sync; initial revision will be reconciled" + ); + } + } + Err(e) => { + policy_bound_to_snapshot = false; + warn!( + error = %e, + "Failed to sync enriched policy back to gateway; initial revision will be reconciled" + ); + } + } + } else { + policy_bound_to_snapshot = false; + } } + let loaded_policy_revision = + policy_bound_to_snapshot.then(|| LoadedPolicyRevision::from_snapshot(&snapshot)); + // Build OPA engine from baked-in rules + typed proto data. // In cluster mode, proxy networking is always enabled so OPA is // always required for allow/deny decisions. @@ -1476,10 +1518,31 @@ async fn load_policy( // container hasn't started yet. After the entrypoint spawns, the // engine is rebuilt with the real PID for symlink resolution. info!("Creating OPA engine from proto policy data"); - let opa_engine = Some(Arc::new(OpaEngine::from_proto(&proto_policy)?)); + let opa_engine = match OpaEngine::from_proto(&proto_policy) { + Ok(engine) => Some(Arc::new(engine)), + Err(e) => { + report_initial_policy_failure(endpoint, id, loaded_policy_revision.as_ref(), &e) + .await; + return Err(e); + } + }; - let policy = SandboxPolicy::try_from(proto_policy.clone())?; - return Ok((policy, opa_engine, Some(proto_policy))); + let policy = match SandboxPolicy::try_from(proto_policy.clone()) { + Ok(policy) => policy, + Err(e) => { + report_initial_policy_failure(endpoint, id, loaded_policy_revision.as_ref(), &e) + .await; + return Err(e); + } + }; + return Ok(( + policy, + opa_engine, + Some(proto_policy), + LoadedPolicyOrigin::Gateway { + revision: loaded_policy_revision, + }, + )); } // No policy source available @@ -1598,6 +1661,271 @@ fn discover_policy_from_path(path: &std::path::Path) -> openshell_core::proto::S } } +/// Identity returned with the exact policy snapshot used to construct OPA. +#[derive(Clone, Debug, PartialEq, Eq)] +struct LoadedPolicyRevision { + version: u32, + policy_hash: String, + config_revision: u64, + policy_source: openshell_core::proto::PolicySource, +} + +/// Identifies where the policy currently loaded into OPA came from. +/// +/// A missing gateway revision means the policy was loaded from the gateway but +/// could not be bound to an authoritative snapshot (for example, enrichment +/// sync failed). That state must reconcile on the first successful poll. A +/// local-file override is different: gateway policy revisions are observed for +/// settings/provider refreshes but must never replace the explicit local OPA +/// policy. +#[derive(Clone, Debug, PartialEq, Eq)] +enum LoadedPolicyOrigin { + LocalOverride, + Gateway { + revision: Option, + }, +} + +impl LoadedPolicyOrigin { + fn allows_gateway_policy_reload(&self) -> bool { + matches!(self, Self::Gateway { .. }) + } +} + +impl LoadedPolicyRevision { + fn from_snapshot(snapshot: &openshell_core::grpc_client::SettingsPollResult) -> Self { + Self { + version: snapshot.version, + policy_hash: snapshot.policy_hash.clone(), + config_revision: snapshot.config_revision, + policy_source: snapshot.policy_source, + } + } +} + +/// A sandbox-scoped policy revision that was constructed successfully at +/// startup and must be acknowledged to the gateway exactly once. +#[derive(Clone, Debug, PartialEq, Eq)] +struct InitialPolicyAck { + version: u32, + policy_hash: String, + config_revision: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct PolicyStatusUpdate { + version: u32, + loaded: bool, + error: String, + initial_policy_hash: Option, +} + +impl PolicyStatusUpdate { + fn initial_loaded(ack: &InitialPolicyAck) -> Self { + Self { + version: ack.version, + loaded: true, + error: String::new(), + initial_policy_hash: Some(ack.policy_hash.clone()), + } + } + + fn loaded(version: u32) -> Self { + Self { + version, + loaded: true, + error: String::new(), + initial_policy_hash: None, + } + } + + fn failed(version: u32, error: String) -> Self { + Self { + version, + loaded: false, + error, + initial_policy_hash: None, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +enum InitialPollDisposition { + Acknowledge(InitialPolicyAck), + Reconcile, + TrackOnly, +} + +/// Determine whether the initially loaded policy corresponds to an +/// authoritative sandbox-scoped revision that must be acknowledged. +/// +/// Returns `Some` only for sandbox-sourced revisions (version > 0) whose +/// captured gateway identity matches the current version and hash. Global +/// policies, local-file development policies, version zero, and changed +/// identities yield `None`, so those paths never emit a sandbox-revision +/// acknowledgement. +fn initial_policy_ack_candidate( + loaded: Option<&LoadedPolicyRevision>, + canonical: &openshell_core::grpc_client::SettingsPollResult, +) -> Option { + let loaded = loaded?; + if loaded.policy_source != openshell_core::proto::PolicySource::Sandbox + || canonical.policy_source != openshell_core::proto::PolicySource::Sandbox + { + return None; + } + if loaded.version == 0 || canonical.version == 0 { + return None; + } + if loaded.version != canonical.version + || loaded.policy_hash != canonical.policy_hash + || canonical.config_revision < loaded.config_revision + { + return None; + } + Some(InitialPolicyAck { + version: loaded.version, + policy_hash: loaded.policy_hash.clone(), + config_revision: canonical.config_revision, + }) +} + +fn initial_poll_disposition( + origin: &LoadedPolicyOrigin, + canonical: &openshell_core::grpc_client::SettingsPollResult, +) -> InitialPollDisposition { + match origin { + LoadedPolicyOrigin::LocalOverride => InitialPollDisposition::TrackOnly, + LoadedPolicyOrigin::Gateway { revision } => { + initial_policy_ack_candidate(revision.as_ref(), canonical).map_or( + InitialPollDisposition::Reconcile, + InitialPollDisposition::Acknowledge, + ) + } + } +} + +/// Deliver policy status updates independently from policy reconciliation. +/// +/// The channel is FIFO, so a delayed older status can never arrive after a +/// newer status and move the gateway's active version backward. Delivery uses +/// the existing bounded retry, but failures never delay policy enforcement. +async fn run_policy_status_reporter( + client: openshell_core::grpc_client::CachedOpenShellClient, + sandbox_id: String, + mut updates: tokio::sync::mpsc::UnboundedReceiver, +) { + 'updates: while let Some(update) = updates.recv().await { + let operation = if update.initial_policy_hash.is_some() { + "Initial policy acknowledgement" + } else { + "Policy status report" + }; + let mut attempt = 1_u32; + loop { + let sandbox_id = sandbox_id.clone(); + let error = update.error.clone(); + let client = client.clone(); + match client + .report_policy_status(&sandbox_id, update.version, update.loaded, &error) + .await + { + Ok(()) => break, + Err(error) if is_retryable_error(&error) => { + let backoff = Duration::from_secs(1_u64 << attempt.saturating_sub(1).min(5)); + warn!( + %error, + attempt, + version = update.version, + loaded = update.loaded, + retry_in_secs = backoff.as_secs(), + "{operation} failed transiently; retaining ordered update" + ); + tokio::time::sleep(backoff).await; + attempt = attempt.saturating_add(1); + } + Err(error) => { + warn!( + %error, + version = update.version, + loaded = update.loaded, + "Discarding terminal policy status update" + ); + continue 'updates; + } + } + } + + if let Some(policy_hash) = update.initial_policy_hash { + ocsf_emit!( + ConfigStateChangeBuilder::new(ocsf_ctx()) + .severity(SeverityId::Informational) + .status(StatusId::Success) + .state(StateId::Enabled, "loaded") + .unmapped("version", serde_json::json!(update.version)) + .unmapped("policy_hash", serde_json::json!(policy_hash)) + .message(format!( + "Acknowledged initial policy revision as loaded [version:{}]", + update.version + )) + .build() + ); + } + } +} + +fn enqueue_policy_status(sender: &UnboundedSender, update: PolicyStatusUpdate) { + let version = update.version; + if let Err(error) = sender.send(update) { + warn!( + %error, + version, + "Policy status reporter unavailable during shutdown" + ); + } +} + +/// Best-effort `FAILED` acknowledgement when initial policy construction or +/// conversion fails. +/// +/// Uses the revision identity captured with the policy that failed to build, +/// and preserves the original construction error as the reported message. A +/// delivery failure here is swallowed so it can never mask that error. +async fn report_initial_policy_failure( + endpoint: &str, + sandbox_id: &str, + revision: Option<&LoadedPolicyRevision>, + error: &miette::Report, +) { + let Some(revision) = revision.filter(|revision| { + revision.version > 0 + && revision.policy_source == openshell_core::proto::PolicySource::Sandbox + }) else { + return; + }; + let client = match openshell_core::grpc_client::CachedOpenShellClient::connect(endpoint).await { + Ok(client) => client, + Err(e) => { + warn!(error = %e, "Failed to connect to report initial policy failure"); + return; + } + }; + let message = error.to_string(); + if let Err(e) = grpc_retry("Initial policy failure report", || { + let client = client.clone(); + let message = message.clone(); + async move { + client + .report_policy_status(sandbox_id, revision.version, false, &message) + .await + } + }) + .await + { + warn!(error = %e, version = revision.version, "Failed to report initial policy failure"); + } +} + /// Background loop that polls the server for policy updates. /// /// When a new version is detected, attempts to reload the OPA engine via @@ -1610,6 +1938,10 @@ struct PolicyPollLoopContext { endpoint: String, sandbox_id: String, opa_engine: Arc, + /// Source of the policy currently loaded into OPA. This distinguishes an + /// explicit local-file override from an unbound gateway revision so the + /// former is never replaced by policy polling. + loaded_policy_origin: LoadedPolicyOrigin, entrypoint_pid: Arc, interval_secs: u64, ocsf_enabled: Arc, @@ -1623,6 +1955,13 @@ async fn run_policy_poll_loop(ctx: PolicyPollLoopContext) -> Result<()> { use std::sync::atomic::Ordering; let client = CachedOpenShellClient::connect(&ctx.endpoint).await?; + let (status_sender, status_receiver) = tokio::sync::mpsc::unbounded_channel(); + tokio::spawn(run_policy_status_reporter( + client.clone(), + ctx.sandbox_id.clone(), + status_receiver, + )); + let mut current_config_revision: u64 = 0; let mut current_provider_env_revision: u64 = ctx.provider_credentials.snapshot().revision; let mut current_policy_hash = String::new(); @@ -1631,18 +1970,42 @@ async fn run_policy_poll_loop(ctx: PolicyPollLoopContext) -> Result<()> { openshell_core::proto::EffectiveSetting, > = std::collections::HashMap::new(); - // Initialize revision from the first poll. + // A first poll that does not match the policy already loaded into OPA must + // pass through the normal reconciliation path immediately. It must never + // seed the applied-state trackers before OPA actually loads it. + let mut pending_result = None; + + // Initialize revision from the first poll and acknowledge the initial + // policy revision the supervisor actually loaded. A mismatched result is + // reconciled below instead of being recorded as already applied. match client.poll_settings(&ctx.sandbox_id).await { - Ok(result) => { - apply_ocsf_json_setting(&ctx.ocsf_enabled, &result.settings); - current_config_revision = result.config_revision; - current_policy_hash = result.policy_hash.clone(); - current_settings = result.settings; - debug!( - config_revision = current_config_revision, - "Settings poll: initial config revision" - ); - } + Ok(result) => match initial_poll_disposition(&ctx.loaded_policy_origin, &result) { + InitialPollDisposition::Acknowledge(candidate) => { + apply_ocsf_json_setting(&ctx.ocsf_enabled, &result.settings); + current_config_revision = candidate.config_revision; + current_policy_hash.clone_from(&candidate.policy_hash); + current_settings = result.settings; + enqueue_policy_status( + &status_sender, + PolicyStatusUpdate::initial_loaded(&candidate), + ); + debug!( + config_revision = current_config_revision, + "Settings poll: initial policy matches loaded revision" + ); + } + InitialPollDisposition::Reconcile => pending_result = Some(result), + InitialPollDisposition::TrackOnly => { + apply_ocsf_json_setting(&ctx.ocsf_enabled, &result.settings); + current_config_revision = result.config_revision; + current_policy_hash = result.policy_hash.clone(); + current_settings = result.settings; + debug!( + config_revision = current_config_revision, + "Settings poll: tracking gateway config while preserving local policy override" + ); + } + }, Err(e) => { warn!(error = %e, "Settings poll: failed to fetch initial version, will retry"); } @@ -1650,13 +2013,16 @@ async fn run_policy_poll_loop(ctx: PolicyPollLoopContext) -> Result<()> { let interval = Duration::from_secs(ctx.interval_secs); loop { - tokio::time::sleep(interval).await; - - let result = match client.poll_settings(&ctx.sandbox_id).await { - Ok(r) => r, - Err(e) => { - debug!(error = %e, "Settings poll: server unreachable, will retry"); - continue; + let result = if let Some(result) = pending_result.take() { + result + } else { + tokio::time::sleep(interval).await; + match client.poll_settings(&ctx.sandbox_id).await { + Ok(result) => result, + Err(e) => { + debug!(error = %e, "Settings poll: server unreachable, will retry"); + continue; + } } }; @@ -1725,7 +2091,7 @@ async fn run_policy_poll_loop(ctx: PolicyPollLoopContext) -> Result<()> { } // Only reload OPA when the policy payload actually changed. - if policy_changed { + if policy_changed && ctx.loaded_policy_origin.allows_gateway_policy_reload() { let Some(policy) = result.policy.as_ref() else { ocsf_emit!(ConfigStateChangeBuilder::new(ocsf_ctx()) .severity(SeverityId::Medium) @@ -1772,13 +2138,11 @@ async fn run_policy_poll_loop(ctx: PolicyPollLoopContext) -> Result<()> { .build() ); } - if result.version > 0 - && result.policy_source == PolicySource::Sandbox - && let Err(e) = client - .report_policy_status(&ctx.sandbox_id, result.version, true, "") - .await - { - warn!(error = %e, "Failed to report policy load success"); + if result.version > 0 && result.policy_source == PolicySource::Sandbox { + enqueue_policy_status( + &status_sender, + PolicyStatusUpdate::loaded(result.version), + ); } } Err(e) => { @@ -1793,18 +2157,11 @@ async fn run_policy_poll_loop(ctx: PolicyPollLoopContext) -> Result<()> { result.version )) .build()); - if result.version > 0 - && result.policy_source == PolicySource::Sandbox - && let Err(report_err) = client - .report_policy_status( - &ctx.sandbox_id, - result.version, - false, - &e.to_string(), - ) - .await - { - warn!(error = %report_err, "Failed to report policy load failure"); + if result.version > 0 && result.policy_source == PolicySource::Sandbox { + enqueue_policy_status( + &status_sender, + PolicyStatusUpdate::failed(result.version, e.to_string()), + ); } } } @@ -2087,4 +2444,176 @@ filesystem_policy: let local_policy = SandboxPolicy::try_from(proto).expect("conversion should succeed"); assert!(matches!(local_policy.network.mode, NetworkMode::Proxy)); } + + // ---- Initial policy acknowledgement tests ---- + + fn proto_policy_fixture() -> openshell_core::proto::SandboxPolicy { + openshell_policy::restrictive_default_policy() + } + + fn settings_poll_result( + policy: Option, + version: u32, + source: openshell_core::proto::PolicySource, + ) -> openshell_core::grpc_client::SettingsPollResult { + openshell_core::grpc_client::SettingsPollResult { + policy, + version, + policy_hash: format!("hash-v{version}"), + config_revision: u64::from(version) * 100, + policy_source: source, + settings: std::collections::HashMap::new(), + global_policy_version: 0, + provider_env_revision: 0, + } + } + + #[test] + fn initial_ack_candidate_matches_sandbox_revision() { + let canonical = settings_poll_result( + Some(proto_policy_fixture()), + 2, + openshell_core::proto::PolicySource::Sandbox, + ); + let loaded = LoadedPolicyRevision::from_snapshot(&canonical); + + let ack = initial_policy_ack_candidate(Some(&loaded), &canonical) + .expect("sandbox-sourced matching revision should be acknowledged"); + + assert_eq!(ack.version, 2); + assert_eq!(ack.policy_hash, "hash-v2"); + assert_eq!(ack.config_revision, 200); + } + + #[test] + fn initial_ack_candidate_ignores_global_policy() { + let canonical = settings_poll_result( + Some(proto_policy_fixture()), + 1, + openshell_core::proto::PolicySource::Global, + ); + let loaded = LoadedPolicyRevision::from_snapshot(&canonical); + + assert!(initial_policy_ack_candidate(Some(&loaded), &canonical).is_none()); + } + + #[test] + fn initial_ack_candidate_ignores_version_zero() { + let canonical = settings_poll_result( + Some(proto_policy_fixture()), + 0, + openshell_core::proto::PolicySource::Sandbox, + ); + let loaded = LoadedPolicyRevision::from_snapshot(&canonical); + + assert!(initial_policy_ack_candidate(Some(&loaded), &canonical).is_none()); + } + + #[test] + fn initial_ack_candidate_ignores_local_file_mode() { + // Local-file mode retains no proto policy, so there is nothing to + // acknowledge to the gateway. + let canonical = settings_poll_result( + Some(proto_policy_fixture()), + 1, + openshell_core::proto::PolicySource::Sandbox, + ); + + assert!(initial_policy_ack_candidate(None, &canonical).is_none()); + } + + #[test] + fn initial_ack_candidate_rejects_mismatched_identity() { + let loaded_snapshot = settings_poll_result( + Some(proto_policy_fixture()), + 1, + openshell_core::proto::PolicySource::Sandbox, + ); + let loaded = LoadedPolicyRevision::from_snapshot(&loaded_snapshot); + let canonical = settings_poll_result( + Some(proto_policy_fixture()), + 2, + openshell_core::proto::PolicySource::Sandbox, + ); + + assert!(initial_policy_ack_candidate(Some(&loaded), &canonical).is_none()); + } + + #[test] + fn initial_poll_reconciles_provider_composition_that_was_not_loaded() { + let loaded_snapshot = settings_poll_result( + Some(proto_policy_fixture()), + 1, + openshell_core::proto::PolicySource::Sandbox, + ); + let loaded = LoadedPolicyRevision::from_snapshot(&loaded_snapshot); + let mut newer = proto_policy_fixture(); + newer.network_policies.insert( + "_provider_work_github".to_string(), + openshell_core::proto::NetworkPolicyRule::default(), + ); + let canonical = + settings_poll_result(Some(newer), 1, openshell_core::proto::PolicySource::Sandbox); + let canonical = openshell_core::grpc_client::SettingsPollResult { + policy_hash: "hash-provider-change".to_string(), + config_revision: loaded.config_revision + 1, + ..canonical + }; + + assert_eq!( + initial_poll_disposition( + &LoadedPolicyOrigin::Gateway { + revision: Some(loaded), + }, + &canonical, + ), + InitialPollDisposition::Reconcile + ); + } + + #[test] + fn initial_poll_tracks_local_override_without_reconciliation() { + let canonical = settings_poll_result( + Some(proto_policy_fixture()), + 2, + openshell_core::proto::PolicySource::Sandbox, + ); + + assert_eq!( + initial_poll_disposition(&LoadedPolicyOrigin::LocalOverride, &canonical), + InitialPollDisposition::TrackOnly + ); + assert!(!LoadedPolicyOrigin::LocalOverride.allows_gateway_policy_reload()); + } + + #[test] + fn initial_poll_reconciles_unbound_gateway_policy() { + let canonical = settings_poll_result( + Some(proto_policy_fixture()), + 2, + openshell_core::proto::PolicySource::Sandbox, + ); + let origin = LoadedPolicyOrigin::Gateway { revision: None }; + + assert_eq!( + initial_poll_disposition(&origin, &canonical), + InitialPollDisposition::Reconcile + ); + assert!(origin.allows_gateway_policy_reload()); + } + + #[test] + fn policy_status_outbox_preserves_all_revision_order() { + let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel(); + for version in 1..=128 { + enqueue_policy_status(&sender, PolicyStatusUpdate::loaded(version)); + } + + for version in 1..=128 { + assert_eq!( + receiver.try_recv().unwrap(), + PolicyStatusUpdate::loaded(version) + ); + } + } } diff --git a/docs/sandboxes/manage-sandboxes.mdx b/docs/sandboxes/manage-sandboxes.mdx index fdbf497f5d..8d47754c5d 100644 --- a/docs/sandboxes/manage-sandboxes.mdx +++ b/docs/sandboxes/manage-sandboxes.mdx @@ -221,6 +221,21 @@ openshell sandbox list --selector env=dev openshell sandbox list --selector env=dev,team=platform ``` +The Python SDK accepts the same gateway labels and selectors. Labels passed to +`create` are stored on the gateway sandbox object and are returned on +`SandboxRef.labels`, so selector-based listing finds Python-created sandboxes: + +```python +from openshell import SandboxClient + +with SandboxClient.from_active_cluster() as client: + sandbox = client.create(name="deep-research-1", labels={"env": "dev", "team": "platform"}) + assert sandbox.labels["team"] == "platform" + + matches = client.list(label_selector="env=dev,team=platform") + assert sandbox.id in {s.id for s in matches} +``` + ## Expose Long Running Services Service forwarding makes a long-running process inside a sandbox reachable through a gateway-managed URL. Use it for development servers, notebooks, dashboards, or other services that keep listening after the sandbox starts. Run the service on loopback inside the sandbox, expose its port, then open the URL printed by OpenShell. diff --git a/e2e/python/test_sandbox_api.py b/e2e/python/test_sandbox_api.py index 80a7d5d6b7..929728d7c3 100644 --- a/e2e/python/test_sandbox_api.py +++ b/e2e/python/test_sandbox_api.py @@ -58,3 +58,58 @@ def read(self, path: str) -> str: ) assert verify_file.exit_code == 0 assert verify_file.stdout.strip() == "ok" + + +def test_sandbox_labels_and_selectors(sandbox_client: SandboxClient) -> None: + import contextlib + import uuid + + suffix = uuid.uuid4().hex[:8] + job_a = f"aiq-labels-a-{suffix}" + job_b = f"aiq-labels-b-{suffix}" + group_selector = f"aiq-test={suffix}" + primary_selector = f"aiq-test={suffix},role=primary" + + created: list[str] = [] + try: + ref_a = sandbox_client.create( + name=job_a, labels={"aiq-test": suffix, "role": "primary"} + ) + created.append(ref_a.name) + ref_b = sandbox_client.create( + name=job_b, labels={"aiq-test": suffix, "role": "secondary"} + ) + created.append(ref_b.name) + + # Labels round-trip through create and get. + assert ref_a.labels["role"] == "primary" + assert dict(sandbox_client.get(job_a).labels)["role"] == "primary" + assert dict(sandbox_client.get(job_b).labels)["role"] == "secondary" + + # A specific selector filters to exactly the primary sandbox. + assert { + s.name for s in sandbox_client.list(label_selector=primary_selector) + } == {job_a} + # The shared group label returns both. + assert {s.name for s in sandbox_client.list(label_selector=group_selector)} == { + job_a, + job_b, + } + + # Deleting one removes only it from selector results. + assert sandbox_client.delete(job_a) + sandbox_client.wait_deleted(job_a) + created.remove(job_a) + assert {s.name for s in sandbox_client.list(label_selector=group_selector)} == { + job_b + } + + # Final deletion leaves no matching sandboxes. + assert sandbox_client.delete(job_b) + sandbox_client.wait_deleted(job_b) + created.remove(job_b) + assert not sandbox_client.list(label_selector=group_selector) + finally: + for name in created: + with contextlib.suppress(Exception): + sandbox_client.delete(name) diff --git a/e2e/rust/src/harness/sandbox.rs b/e2e/rust/src/harness/sandbox.rs index af71a38ee5..43fa4ee396 100644 --- a/e2e/rust/src/harness/sandbox.rs +++ b/e2e/rust/src/harness/sandbox.rs @@ -126,13 +126,30 @@ impl SandboxGuard { pub async fn create_keep( command: &[&str], ready_marker: &str, + ) -> Result { + Self::create_keep_with_args(&[], command, ready_marker).await + } + + /// Like [`SandboxGuard::create_keep`], but forwards extra flags to + /// `sandbox create` (e.g. `--policy `, `--name `) before the + /// `-- ` separator. + /// + /// # Errors + /// + /// Returns an error if the process exits prematurely, the ready marker is + /// not seen within [`SANDBOX_READY_TIMEOUT`], or the sandbox name cannot + /// be parsed. + pub async fn create_keep_with_args( + create_args: &[&str], + command: &[&str], + ready_marker: &str, ) -> Result { let mut cmd = openshell_cmd(); - cmd.arg("sandbox") - .arg("create") - .arg("--keep") - .arg("--") - .args(command); + cmd.arg("sandbox").arg("create").arg("--keep"); + for arg in create_args { + cmd.arg(arg); + } + cmd.arg("--").args(command); cmd.stdout(Stdio::piped()).stderr(Stdio::piped()); let mut child = cmd diff --git a/e2e/rust/tests/live_policy_update.rs b/e2e/rust/tests/live_policy_update.rs index c60b295487..c3bad7e802 100644 --- a/e2e/rust/tests/live_policy_update.rs +++ b/e2e/rust/tests/live_policy_update.rs @@ -32,6 +32,30 @@ use openshell_e2e::harness::output::{extract_field, strip_ansi}; use openshell_e2e::harness::sandbox::SandboxGuard; use tempfile::NamedTempFile; +#[cfg(feature = "e2e-docker")] +const LOCAL_OVERRIDE_REGO: &str = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../crates/openshell-supervisor-network/data/sandbox-policy.rego" +)); + +#[cfg(feature = "e2e-docker")] +const LOCAL_OVERRIDE_DOCKERFILE: &str = r#"FROM public.ecr.aws/docker/library/python:3.13-slim + +RUN apt-get update && apt-get install -y --no-install-recommends iproute2 \ + && rm -rf /var/lib/apt/lists/* +RUN groupadd -g 1000660000 sandbox && \ + useradd -m -u 1000660000 -g sandbox sandbox + +COPY local-policy.rego /etc/openshell/local-policy.rego +COPY local-policy.yaml /etc/openshell/local-policy.yaml + +ENV OPENSHELL_POLICY_RULES=/etc/openshell/local-policy.rego +ENV OPENSHELL_POLICY_DATA=/etc/openshell/local-policy.yaml +ENV OPENSHELL_POLICY_POLL_INTERVAL_SECS=1 + +CMD ["sleep", "infinity"] +"#; + // --------------------------------------------------------------------------- // Policy YAML builders // --------------------------------------------------------------------------- @@ -130,6 +154,44 @@ process: Ok(file) } +#[cfg(feature = "e2e-docker")] +fn write_local_override_image() -> Result { + let dir = tempfile::tempdir().map_err(|e| format!("create image context: {e}"))?; + std::fs::write(dir.path().join("Dockerfile"), LOCAL_OVERRIDE_DOCKERFILE) + .map_err(|e| format!("write local override Dockerfile: {e}"))?; + std::fs::write(dir.path().join("local-policy.rego"), LOCAL_OVERRIDE_REGO) + .map_err(|e| format!("write local override Rego policy: {e}"))?; + std::fs::write( + dir.path().join("local-policy.yaml"), + r"version: 1 + +filesystem_policy: + include_workdir: true + read_only: + - /usr + - /lib + - /proc + - /dev/urandom + - /etc + read_write: + - /sandbox + - /tmp + - /dev/null + +landlock: + compatibility: best_effort + +process: + run_as_user: sandbox + run_as_group: sandbox + +network_policies: {} +", + ) + .map_err(|e| format!("write local override policy data: {e}"))?; + Ok(dir) +} + // --------------------------------------------------------------------------- // CLI helpers // --------------------------------------------------------------------------- @@ -421,3 +483,177 @@ async fn live_policy_update_from_empty_network_policies() { guard.cleanup().await; } + +/// Regression for #2159: a sparse initial policy that the supervisor enriches +/// with baseline filesystem paths during startup must have its resulting +/// revision acknowledged as loaded, not left `Pending`. +/// +/// Reproduction (matches the maintainer's triage): create a sandbox with the +/// network-only `examples/policy-advisor/sandbox-policy.yaml`. The supervisor +/// adds baseline filesystem paths, syncs the enriched policy back to the +/// gateway (creating revision 2, superseding revision 1), builds the OPA engine +/// and then acknowledges revision 2 as loaded. Before the fix, revision 2 +/// stayed `Pending` even though the sandbox was `Ready` and the policy was +/// effective. +/// +/// NOTE: This exercises the Docker-backed supervisor built from this branch. +/// The exact `policy list` status wording ("Loaded"/"Superseded") may differ by +/// CLI version; the assertions below key on the effective version reaching 2 and +/// no revision remaining `Pending` once the acknowledgement lands. +#[tokio::test] +async fn initial_sparse_policy_is_acknowledged_as_loaded() { + // Repo-relative path to the sparse network-only policy fixture. + let sparse_policy = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../examples/policy-advisor/sandbox-policy.yaml" + ); + + let mut guard = SandboxGuard::create_keep_with_args( + &["--name", "e2e-2159-sparse-enrich", "--policy", sparse_policy, "--no-tty"], + &["sh", "-c", "echo Ready && sleep infinity"], + "Ready", + ) + .await + .expect("create keep sandbox with sparse policy"); + + // The enriched revision (2) is synced during startup; the acknowledgement + // (LOADED) is delivered by the supervisor's poll loop shortly after Ready. + // Poll until the effective policy is version 2 and no revision is Pending. + let mut acknowledged = false; + let mut last_list = String::new(); + for _ in 0..30 { + let get = run_cli(&["policy", "get", &guard.name]).await; + let version = get.success.then(|| extract_version(&get.output)).flatten(); + + let list = run_cli(&["policy", "list", &guard.name]).await; + last_list = list.output.clone(); + let pending = list.output.to_lowercase().contains("pending"); + + if version == Some(2) && list.success && !pending { + acknowledged = true; + break; + } + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + } + + assert!( + acknowledged, + "enriched initial policy should reach revision 2 with no Pending revision.\n\ + last `policy list` output:\n{last_list}" + ); + + // Both the superseded original (1) and the loaded enriched revision (2) + // must appear in the revision history. + assert!( + list_output_contains_version(&last_list, 2), + "policy list should contain revision 2:\n{last_list}" + ); + assert!( + list_output_contains_version(&last_list, 1), + "policy list should contain revision 1:\n{last_list}" + ); + + guard.cleanup().await; +} + +/// An explicit local Rego/data override remains authoritative even when the +/// sandbox has a gateway policy and that policy changes while it is running. +/// Gateway polling must continue for settings and providers without replacing +/// the locally loaded OPA engine. +#[cfg(feature = "e2e-docker")] +#[tokio::test] +async fn local_policy_override_survives_gateway_policy_polls() { + let image_context = write_local_override_image().expect("write local override image"); + let dockerfile = image_context.path().join("Dockerfile"); + let dockerfile = dockerfile.to_str().expect("Dockerfile path should be utf-8"); + + let gateway_policy_a_file = write_policy(&["example.com"]).expect("write gateway policy A"); + let gateway_policy_a_path = gateway_policy_a_file + .path() + .to_str() + .expect("gateway policy A path should be utf-8") + .to_string(); + let gateway_policy_b_file = + write_policy(&["example.com", "api.anthropic.com"]).expect("write gateway policy B"); + let gateway_policy_b_path = gateway_policy_b_file + .path() + .to_str() + .expect("gateway policy B path should be utf-8") + .to_string(); + + let mut guard = SandboxGuard::create_keep_with_args( + &[ + "--name", + "e2e-local-policy-override", + "--from", + dockerfile, + "--policy", + &gateway_policy_a_path, + "--no-tty", + ], + &["sh", "-c", "echo Ready && sleep infinity"], + "Ready", + ) + .await + .expect("create sandbox with local policy override"); + + // Allow several one-second poll intervals. Before the fix, the first poll + // immediately reloaded gateway policy A over the local override. + tokio::time::sleep(std::time::Duration::from_secs(4)).await; + let initial_logs = run_cli(&[ + "logs", + &guard.name, + "-n", + "500", + "--since", + "1m", + "--source", + "sandbox", + ]) + .await; + assert!(initial_logs.success, "fetch initial sandbox logs:\n{}", initial_logs.output); + assert!( + initial_logs + .output + .contains("Loading OPA policy engine from local files"), + "sandbox should load the explicit local policy:\n{}", + initial_logs.output + ); + assert!( + !initial_logs.output.contains("Policy reloaded successfully"), + "the first gateway poll must not replace the local policy:\n{}", + initial_logs.output + ); + + let update = run_cli(&[ + "policy", + "set", + &guard.name, + "--policy", + &gateway_policy_b_path, + ]) + .await; + assert!(update.success, "publish gateway policy B:\n{}", update.output); + + // A later gateway revision must also remain observational in local mode. + tokio::time::sleep(std::time::Duration::from_secs(4)).await; + let updated_logs = run_cli(&[ + "logs", + &guard.name, + "-n", + "500", + "--since", + "1m", + "--source", + "sandbox", + ]) + .await; + assert!(updated_logs.success, "fetch updated sandbox logs:\n{}", updated_logs.output); + assert!( + !updated_logs.output.contains("Policy reloaded successfully"), + "gateway policy updates must not replace the local override:\n{}", + updated_logs.output + ); + + guard.cleanup().await; +} diff --git a/python/openshell/sandbox.py b/python/openshell/sandbox.py index 3b64089deb..9b00ebfb70 100644 --- a/python/openshell/sandbox.py +++ b/python/openshell/sandbox.py @@ -13,8 +13,8 @@ import threading import time from collections import namedtuple -from dataclasses import dataclass -from typing import TYPE_CHECKING, cast +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Never, SupportsIndex, cast from urllib.parse import urlparse import grpc @@ -130,11 +130,42 @@ class SandboxStatusRef: current_policy_version: int +class _ImmutableLabels(dict[str, str]): + """A read-only, copy- and pickle-safe label mapping.""" + + def _deny_mutation(self, *args: object, **kwargs: object) -> Never: + del args, kwargs + raise TypeError("sandbox labels are immutable") + + __setitem__ = _deny_mutation + __delitem__ = _deny_mutation + clear = _deny_mutation + pop = _deny_mutation + popitem = _deny_mutation + setdefault = _deny_mutation + update = _deny_mutation + __ior__ = _deny_mutation + + def __deepcopy__(self, memo: dict[int, object]) -> _ImmutableLabels: + del memo + return type(self)(self) + + def __reduce_ex__(self, protocol: SupportsIndex, /) -> str | tuple[Any, ...]: + del protocol + return type(self), (dict(self),) + + @dataclass(frozen=True) class SandboxRef: id: str name: str status: SandboxStatusRef + # Excluded from equality/hash to preserve the original identity while the + # immutable mapping remains safe for deepcopy, pickle, and asdict. + labels: Mapping[str, str] = field(default_factory=_ImmutableLabels, compare=False) + + def __post_init__(self) -> None: + object.__setattr__(self, "labels", _ImmutableLabels(self.labels)) @property def phase(self) -> int: @@ -397,10 +428,16 @@ def create( self, *, spec: openshell_pb2.SandboxSpec | None = None, + name: str | None = None, + labels: Mapping[str, str] | None = None, ) -> SandboxRef: request_spec = spec if spec is not None else _default_spec() response = self._stub.CreateSandbox( - openshell_pb2.CreateSandboxRequest(spec=request_spec), + openshell_pb2.CreateSandboxRequest( + spec=request_spec, + name=name or "", + labels=dict(labels) if labels else {}, + ), timeout=self._timeout, ) sandbox_ref = _sandbox_ref(response.sandbox) @@ -412,8 +449,10 @@ def create_session( self, *, spec: openshell_pb2.SandboxSpec | None = None, + name: str | None = None, + labels: Mapping[str, str] | None = None, ) -> SandboxSession: - return SandboxSession(self, self.create(spec=spec)) + return SandboxSession(self, self.create(spec=spec, name=name, labels=labels)) def get(self, sandbox_name: str) -> SandboxRef: response = self._stub.GetSandbox( @@ -425,15 +464,36 @@ def get(self, sandbox_name: str) -> SandboxRef: def get_session(self, sandbox_name: str) -> SandboxSession: return SandboxSession(self, self.get(sandbox_name)) - def list(self, *, limit: int = 100, offset: int = 0) -> builtins.list[SandboxRef]: + def list( + self, + *, + limit: int = 100, + offset: int = 0, + label_selector: str | None = None, + ) -> builtins.list[SandboxRef]: response = self._stub.ListSandboxes( - openshell_pb2.ListSandboxesRequest(limit=limit, offset=offset), + openshell_pb2.ListSandboxesRequest( + limit=limit, + offset=offset, + label_selector=label_selector or "", + ), timeout=self._timeout, ) return [_sandbox_ref(item) for item in response.sandboxes] - def list_ids(self, *, limit: int = 100, offset: int = 0) -> builtins.list[str]: - return [item.id for item in self.list(limit=limit, offset=offset)] + def list_ids( + self, + *, + limit: int = 100, + offset: int = 0, + label_selector: str | None = None, + ) -> builtins.list[str]: + return [ + item.id + for item in self.list( + limit=limit, offset=offset, label_selector=label_selector + ) + ] def delete(self, sandbox_name: str) -> bool: response = self._stub.DeleteSandbox( @@ -646,6 +706,8 @@ def __init__( sandbox: str | SandboxRef | None = None, delete_on_exit: bool = True, spec: openshell_pb2.SandboxSpec | None = None, + name: str | None = None, + labels: Mapping[str, str] | None = None, timeout: float = 30.0, ready_timeout_seconds: float = 120.0, auto_refresh: bool = True, @@ -665,6 +727,9 @@ def __init__( self._sandbox_input = sandbox self._delete_on_exit = delete_on_exit self._spec = spec + self._name = name + # Copy so later caller mutation cannot change what gets sent on enter. + self._labels = dict(labels) if labels is not None else None self._timeout = timeout self._ready_timeout_seconds = ready_timeout_seconds self._auto_refresh = auto_refresh @@ -686,6 +751,15 @@ def sandbox(self) -> SandboxRef: return self._session.sandbox def __enter__(self) -> Sandbox: + # Creation metadata cannot be applied when attaching to an existing + # sandbox; reject it before opening a connection. + if self._sandbox_input is not None and ( + self._name is not None or self._labels is not None + ): + raise SandboxError( + "name and labels cannot be set when attaching to an existing sandbox" + ) + client = SandboxClient.from_active_cluster( cluster=self._cluster, timeout=self._timeout, @@ -696,7 +770,9 @@ def __enter__(self) -> Sandbox: self._client = client if self._sandbox_input is None: - self._session = client.create_session(spec=self._spec) + self._session = client.create_session( + spec=self._spec, name=self._name, labels=self._labels + ) elif isinstance(self._sandbox_input, SandboxRef): self._session = SandboxSession(client, self._sandbox_input) else: @@ -813,6 +889,7 @@ def _sandbox_ref(sandbox: openshell_pb2.Sandbox) -> SandboxRef: phase=status.phase if status else 0, current_policy_version=status.current_policy_version if status else 0, ), + labels=sandbox.metadata.labels if sandbox.metadata else {}, ) diff --git a/python/openshell/sandbox_test.py b/python/openshell/sandbox_test.py index 82e86eefc1..700a9d1ed6 100644 --- a/python/openshell/sandbox_test.py +++ b/python/openshell/sandbox_test.py @@ -5,11 +5,17 @@ import json import os +import pickle import threading import time +from copy import deepcopy +from dataclasses import asdict from pathlib import Path +from types import SimpleNamespace from typing import Any, cast +import pytest + from openshell._proto import openshell_pb2 from openshell.sandbox import ( _PYTHON_CLOUDPICKLE_BOOTSTRAP, @@ -18,6 +24,8 @@ Sandbox, SandboxClient, SandboxError, + SandboxRef, + SandboxStatusRef, TlsConfig, _BearerAuthInterceptor, _load_cluster_bearer_token, @@ -25,6 +33,7 @@ _normalize_bearer, _OidcRefresher, _read_oidc_token_bundle, + _sandbox_ref, ) @@ -60,7 +69,7 @@ class _Response: return _Response() -def _client_with_fake_stub(stub: _FakeStub) -> SandboxClient: +def _client_with_fake_stub(stub: object) -> SandboxClient: client = cast("SandboxClient", object.__new__(SandboxClient)) client._timeout = 30.0 client._stub = cast("Any", stub) @@ -1413,3 +1422,272 @@ def test_from_active_cluster_reads_utf8_bytes_from_active_gateway_and_metadata( assert client._endpoint == "tést.example:8080" finally: client.close() + + +# ---- Sandbox label / selector API tests ---- + + +def _make_sandbox_proto( + id_: str, + name: str, + labels: dict[str, str] | None = None, + phase: openshell_pb2.SandboxPhase = openshell_pb2.SANDBOX_PHASE_READY, + version: int = 0, +) -> openshell_pb2.Sandbox: + sandbox = openshell_pb2.Sandbox() + sandbox.metadata.id = id_ + sandbox.metadata.name = name + for key, value in (labels or {}).items(): + sandbox.metadata.labels[key] = value + sandbox.status.phase = phase + sandbox.status.current_policy_version = version + return sandbox + + +class _FakeSandboxStub: + def __init__(self, listed: list[openshell_pb2.Sandbox] | None = None) -> None: + self.create_request: openshell_pb2.CreateSandboxRequest | None = None + self.list_request: openshell_pb2.ListSandboxesRequest | None = None + self._listed = listed or [] + + def CreateSandbox( + self, + request: openshell_pb2.CreateSandboxRequest, + timeout: float | None = None, + ) -> Any: + self.create_request = request + _ = timeout + return SimpleNamespace( + sandbox=_make_sandbox_proto( + "sandbox-1", request.name or "generated", dict(request.labels) + ) + ) + + def ListSandboxes( + self, + request: openshell_pb2.ListSandboxesRequest, + timeout: float | None = None, + ) -> Any: + self.list_request = request + _ = timeout + return SimpleNamespace(sandboxes=list(self._listed)) + + +class _RecordingHighLevelClient: + """A stand-in for SandboxClient used to observe high-level forwarding.""" + + def __init__(self) -> None: + self.create_kwargs: dict[str, Any] | None = None + + def create_session( + self, + *, + spec: Any = None, + name: str | None = None, + labels: Any = None, + ) -> Any: + self.create_kwargs = {"spec": spec, "name": name, "labels": labels} + return SimpleNamespace(sandbox=SimpleNamespace(name=name or "generated")) + + def wait_ready(self, name: str, *, timeout_seconds: float = 300.0) -> SandboxRef: + _ = timeout_seconds + return SandboxRef( + id="sandbox-1", + name=name, + status=SandboxStatusRef(phase=2, current_policy_version=0), + ) + + +def test_create_forwards_name_and_labels() -> None: + stub = _FakeSandboxStub() + client = _client_with_fake_stub(stub) + + ref = client.create(name="job-1", labels={"aiq": "deep-research"}) + + assert stub.create_request is not None + assert stub.create_request.name == "job-1" + assert dict(stub.create_request.labels) == {"aiq": "deep-research"} + assert dict(ref.labels) == {"aiq": "deep-research"} + + +def test_create_without_args_sends_empty_metadata() -> None: + stub = _FakeSandboxStub() + client = _client_with_fake_stub(stub) + + client.create() + + assert stub.create_request is not None + assert stub.create_request.name == "" + assert dict(stub.create_request.labels) == {} + + +def test_create_copies_caller_labels() -> None: + stub = _FakeSandboxStub() + client = _client_with_fake_stub(stub) + + caller_labels = {"aiq": "deep-research"} + client.create(labels=caller_labels) + caller_labels["aiq"] = "mutated" + + assert stub.create_request is not None + assert dict(stub.create_request.labels) == {"aiq": "deep-research"} + + +def test_create_session_forwards_name_and_labels() -> None: + stub = _FakeSandboxStub() + client = _client_with_fake_stub(stub) + + session = client.create_session(name="job-2", labels={"team": "aiq"}) + + assert stub.create_request is not None + assert stub.create_request.name == "job-2" + assert dict(stub.create_request.labels) == {"team": "aiq"} + assert session.sandbox.name == "job-2" + + +def test_list_forwards_label_selector() -> None: + stub = _FakeSandboxStub() + client = _client_with_fake_stub(stub) + + client.list(label_selector="aiq=deep-research") + + assert stub.list_request is not None + assert stub.list_request.label_selector == "aiq=deep-research" + + +def test_list_without_selector_sends_empty_string() -> None: + stub = _FakeSandboxStub() + client = _client_with_fake_stub(stub) + + client.list() + + assert stub.list_request is not None + assert stub.list_request.label_selector == "" + + +def test_list_ids_forwards_label_selector() -> None: + stub = _FakeSandboxStub(listed=[_make_sandbox_proto("sandbox-1", "job-1")]) + client = _client_with_fake_stub(stub) + + ids = client.list_ids(label_selector="aiq=deep-research") + + assert stub.list_request is not None + assert stub.list_request.label_selector == "aiq=deep-research" + assert ids == ["sandbox-1"] + + +def test_sandbox_ref_retains_gateway_labels() -> None: + proto = _make_sandbox_proto( + "sandbox-1", "job-1", {"aiq": "deep-research", "env": "dev"} + ) + + ref = _sandbox_ref(proto) + + assert dict(ref.labels) == {"aiq": "deep-research", "env": "dev"} + + +def test_returned_labels_are_immutable() -> None: + proto = _make_sandbox_proto("sandbox-1", "job-1", {"aiq": "deep-research"}) + ref = _sandbox_ref(proto) + + with pytest.raises(TypeError): + ref.labels["mutated"] = "nope" # type: ignore[index] + + +def test_direct_sandbox_ref_construction_defaults_labels() -> None: + ref = SandboxRef( + id="sandbox-1", + name="job-1", + status=SandboxStatusRef(phase=2, current_policy_version=0), + ) + + assert dict(ref.labels) == {} + + +def test_sandbox_ref_stays_hashable_with_labels_excluded_from_identity() -> None: + ref_a = _sandbox_ref(_make_sandbox_proto("sandbox-1", "job-1", {"aiq": "a"})) + ref_b = _sandbox_ref(_make_sandbox_proto("sandbox-1", "job-1", {"aiq": "b"})) + + # Frozen dataclass must remain hashable despite the immutable labels field. + assert hash(ref_a) == hash(ref_b) + # Labels are excluded from identity: same (id, name, status) compares equal. + assert ref_a == ref_b + assert {ref_a, ref_b} == {ref_a} + + +def test_sandbox_ref_labels_support_standard_serialization() -> None: + ref = _sandbox_ref( + _make_sandbox_proto("sandbox-1", "job-1", {"aiq": "deep-research"}) + ) + + assert asdict(ref)["labels"] == {"aiq": "deep-research"} + assert dict(deepcopy(ref).labels) == {"aiq": "deep-research"} + assert dict(pickle.loads(pickle.dumps(ref)).labels) == {"aiq": "deep-research"} + + +def test_default_sandbox_ref_labels_support_standard_serialization() -> None: + ref = SandboxRef( + id="sandbox-1", + name="job-1", + status=SandboxStatusRef(phase=2, current_policy_version=0), + ) + + assert asdict(ref)["labels"] == {} + assert dict(deepcopy(ref).labels) == {} + assert dict(pickle.loads(pickle.dumps(ref)).labels) == {} + + +def test_direct_sandbox_ref_copies_and_freezes_labels() -> None: + labels = {"aiq": "deep-research"} + ref = SandboxRef( + id="sandbox-1", + name="job-1", + status=SandboxStatusRef(phase=2, current_policy_version=0), + labels=labels, + ) + labels["aiq"] = "mutated" + + assert dict(ref.labels) == {"aiq": "deep-research"} + with pytest.raises(TypeError): + ref.labels["mutated"] = "nope" # type: ignore[index] + + +def test_high_level_creation_forwards_name_and_labels( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recording = _RecordingHighLevelClient() + monkeypatch.setattr( + SandboxClient, + "from_active_cluster", + classmethod(lambda _cls, **_kwargs: recording), + ) + + sandbox = Sandbox( + name="job-1", labels={"aiq": "deep-research"}, delete_on_exit=False + ) + sandbox.__enter__() + + assert recording.create_kwargs == { + "spec": None, + "name": "job-1", + "labels": {"aiq": "deep-research"}, + } + + +def test_high_level_attach_rejects_name() -> None: + sandbox = Sandbox(sandbox="existing-sandbox", name="job-1") + + with pytest.raises(SandboxError): + sandbox.__enter__() + + +def test_high_level_attach_rejects_labels() -> None: + ref = SandboxRef( + id="sandbox-1", + name="existing", + status=SandboxStatusRef(phase=2, current_policy_version=0), + ) + sandbox = Sandbox(sandbox=ref, labels={"aiq": "deep-research"}) + + with pytest.raises(SandboxError): + sandbox.__enter__()