feat(supervisor): stage gateway configuration snapshot delivery

Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
This commit is contained in:
Piotr Mlocek
2026-09-10 18:58:44 -07:00
parent 0569c3a20a
commit ea6c6891dc
24 changed files with 4354 additions and 1105 deletions
@@ -989,6 +989,18 @@ impl ComputeRuntime {
}
})?;
if let Err(status) = Box::pin(crate::grpc::policy::initialize_policy_history(
self.store.as_ref(),
&sandbox,
crate::grpc::policy::InitialPolicyHistoryStatus::Pending,
))
.await
{
let _ = self.store.delete(Sandbox::object_type(), &sandbox_id).await;
self.sandbox_index.remove_sandbox(&sandbox_id);
return Err(status);
}
if let Some(token) = sandbox_token
&& let Some(spec) = driver_sandbox.spec.as_mut()
{
@@ -1026,6 +1038,10 @@ impl ComputeRuntime {
Ok(sandbox)
}
Err(status) if status.code() == Code::AlreadyExists => {
let _ = self
.store
.delete_by_scope(POLICY_OBJECT_TYPE, sandbox.object_id())
.await;
let _ = self
.store
.delete(Sandbox::object_type(), sandbox.object_id())
@@ -1034,6 +1050,10 @@ impl ComputeRuntime {
Err(Status::already_exists("sandbox already exists"))
}
Err(status) if status.code() == Code::FailedPrecondition => {
let _ = self
.store
.delete_by_scope(POLICY_OBJECT_TYPE, sandbox.object_id())
.await;
let _ = self
.store
.delete(Sandbox::object_type(), sandbox.object_id())
@@ -1042,6 +1062,10 @@ impl ComputeRuntime {
Err(Status::failed_precondition(status.message().to_string()))
}
Err(err) => {
let _ = self
.store
.delete_by_scope(POLICY_OBJECT_TYPE, sandbox.object_id())
.await;
let _ = self
.store
.delete(Sandbox::object_type(), sandbox.object_id())
@@ -4938,6 +4962,7 @@ pub async fn new_test_runtime_with_driver(
#[cfg(test)]
mod tests {
use super::*;
use crate::policy_store::PolicyStoreExt;
use futures::stream;
use openshell_core::proto::compute::v1::{
CreateSandboxResponse, DeleteSandboxResponse, GetCapabilitiesResponse, GetSandboxRequest,
@@ -11240,6 +11265,36 @@ mod tests {
);
}
#[tokio::test]
async fn create_sandbox_persists_initial_policy_revision() {
let runtime = test_runtime(Arc::new(TestDriver::default())).await;
let mut sandbox = sandbox_record(
"sb-initial-policy",
"initial-policy",
SandboxPhase::Provisioning,
);
let policy = openshell_core::proto::SandboxPolicy::default();
sandbox.spec = Some(SandboxSpec {
policy: Some(policy.clone()),
..Default::default()
});
runtime.create_sandbox(sandbox, None, false).await.unwrap();
let revision = runtime
.store
.get_latest_policy("sb-initial-policy")
.await
.unwrap()
.expect("initial policy revision");
assert_eq!(revision.version, 1);
assert_eq!(
revision.policy_hash,
crate::grpc::policy::deterministic_policy_hash(&policy)
);
assert_eq!(revision.status, "pending");
}
#[tokio::test]
async fn created_sandbox_is_immediately_visible_to_label_selectors() {
let runtime = test_runtime(Arc::new(TestDriver::default())).await;
@@ -0,0 +1,661 @@
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Build and route complete supervisor configuration snapshots.
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::fmt;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use metrics::counter;
use openshell_core::proto::{
ConfigBootstrap, ProviderEnvironmentSnapshot, Sandbox, SandboxConfigSnapshot,
};
use tonic::{Code, Status};
use tracing::warn;
use crate::ServerState;
use crate::grpc::policy::{build_provider_environment_snapshot, build_sandbox_config_snapshot};
use crate::persistence::ObjectWorkspace;
use crate::supervisor_session::SupervisorSessionRegistry;
/// Leaves headroom below tonic's default 4 MiB decode limit for framing and
/// future envelope fields.
pub const MAX_SUPERVISOR_CONFIG_MESSAGE_BYTES: usize = 3 * 1024 * 1024;
const CONFIG_SNAPSHOT_BUILD_TIMEOUT: Duration = Duration::from_secs(45);
const MAX_ACTIVE_FANOUT_WORKERS: usize = 64;
/// One complete configuration component awaiting delivery to a supervisor.
#[derive(Clone)]
pub enum SupervisorConfigMessage {
SandboxConfig(Box<SandboxConfigSnapshot>),
ProviderEnvironment(ProviderEnvironmentSnapshot),
}
impl SupervisorConfigMessage {
pub(crate) fn component_name(&self) -> &'static str {
match self {
Self::SandboxConfig(_) => "sandbox_config",
Self::ProviderEnvironment(_) => "provider_environment",
}
}
}
impl fmt::Debug for SupervisorConfigMessage {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::SandboxConfig(_) => "SandboxConfig(<redacted>)",
Self::ProviderEnvironment(_) => "ProviderEnvironment(<redacted>)",
})
}
}
/// Result of routing one configuration snapshot toward a supervisor session.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DeliveryDisposition {
Enqueued,
NoActiveSession,
QueueFull,
SessionClosed,
PayloadTooLarge,
}
/// Transport boundary for configuration delivery.
#[tonic::async_trait]
pub trait SupervisorConfigRouter: fmt::Debug + Send + Sync {
async fn deliver(
&self,
sandbox_id: &str,
message: SupervisorConfigMessage,
) -> DeliveryDisposition;
async fn routable_sandbox_ids(&self) -> Vec<String>;
}
#[derive(Debug)]
pub struct LocalSupervisorConfigRouter {
sessions: Arc<SupervisorSessionRegistry>,
}
impl LocalSupervisorConfigRouter {
#[must_use]
pub fn new(sessions: Arc<SupervisorSessionRegistry>) -> Self {
Self { sessions }
}
}
#[tonic::async_trait]
impl SupervisorConfigRouter for LocalSupervisorConfigRouter {
async fn deliver(
&self,
sandbox_id: &str,
message: SupervisorConfigMessage,
) -> DeliveryDisposition {
self.sessions.deliver_config(sandbox_id, message)
}
async fn routable_sandbox_ids(&self) -> Vec<String> {
self.sessions.connected_sandbox_ids()
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ConfigComponents {
pub sandbox_config: bool,
pub provider_environment: bool,
}
impl ConfigComponents {
pub const ALL: Self = Self {
sandbox_config: true,
provider_environment: true,
};
pub const SANDBOX_AND_PROVIDER: Self = Self {
sandbox_config: true,
provider_environment: true,
};
pub const SANDBOX_CONFIG: Self = Self {
sandbox_config: true,
provider_environment: false,
};
fn selected(self) -> impl Iterator<Item = ConfigComponentKind> {
[
(self.sandbox_config, ConfigComponentKind::SandboxConfig),
(
self.provider_environment,
ConfigComponentKind::ProviderEnvironment,
),
]
.into_iter()
.filter_map(|(selected, component)| selected.then_some(component))
}
fn only(component: ConfigComponentKind) -> Self {
Self {
sandbox_config: component == ConfigComponentKind::SandboxConfig,
provider_environment: component == ConfigComponentKind::ProviderEnvironment,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum ConfigComponentKind {
SandboxConfig,
ProviderEnvironment,
}
impl ConfigComponentKind {
fn name(self) -> &'static str {
match self {
Self::SandboxConfig => "sandbox_config",
Self::ProviderEnvironment => "provider_environment",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct DeliveryKey {
sandbox_id: String,
component: ConfigComponentKind,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum FanoutScope {
Workspace(String),
AllConnected,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct FanoutKey {
scope: FanoutScope,
component: ConfigComponentKind,
}
/// Coalesces publications and runs one worker per sandbox and component.
///
/// The map entry is also the worker lease. Its boolean is set when another
/// mutation arrives during a build or route operation. The worker then rebuilds
/// the current full snapshot once, regardless of how many mutations arrived.
#[derive(Debug, Default)]
pub struct ConfigDeliveryQueue {
pending: Mutex<HashMap<DeliveryKey, bool>>,
fanout_pending: Mutex<HashMap<FanoutKey, bool>>,
}
impl ConfigDeliveryQueue {
fn enqueue(&self, key: DeliveryKey) -> bool {
let mut pending = self.pending.lock().unwrap();
match pending.entry(key) {
Entry::Occupied(mut entry) => {
*entry.get_mut() = true;
false
}
Entry::Vacant(entry) => {
entry.insert(true);
true
}
}
}
fn take(&self, key: &DeliveryKey) {
let mut pending = self.pending.lock().unwrap();
if let Some(changed) = pending.get_mut(key) {
*changed = false;
}
}
fn finish_pass(&self, key: &DeliveryKey) -> bool {
let mut pending = self.pending.lock().unwrap();
if pending.get(key).is_some_and(|changed| !changed) {
pending.remove(key);
false
} else {
pending.contains_key(key)
}
}
fn enqueue_fanout(&self, key: FanoutKey) -> FanoutEnqueue {
let mut pending = self.fanout_pending.lock().unwrap();
if let Some(changed) = pending.get_mut(&key) {
*changed = true;
return FanoutEnqueue::Coalesced;
}
if pending.len() >= MAX_ACTIVE_FANOUT_WORKERS {
FanoutEnqueue::Full
} else {
pending.insert(key, true);
FanoutEnqueue::StartWorker
}
}
fn take_fanout(&self, key: &FanoutKey) {
let mut pending = self.fanout_pending.lock().unwrap();
if let Some(changed) = pending.get_mut(key) {
*changed = false;
}
}
fn finish_fanout_pass(&self, key: &FanoutKey) -> bool {
let mut pending = self.fanout_pending.lock().unwrap();
if pending.get(key).is_some_and(|changed| !changed) {
pending.remove(key);
false
} else {
pending.contains_key(key)
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum FanoutEnqueue {
StartWorker,
Coalesced,
Full,
}
pub async fn build_config_bootstrap(
state: &Arc<ServerState>,
sandbox: &Sandbox,
) -> Result<ConfigBootstrap, Status> {
tokio::time::timeout(
CONFIG_SNAPSHOT_BUILD_TIMEOUT,
build_consistent_config_bootstrap(state, sandbox),
)
.await
.map_err(|_| Status::deadline_exceeded("supervisor configuration bootstrap timed out"))?
}
async fn build_consistent_config_bootstrap(
state: &Arc<ServerState>,
sandbox: &Sandbox,
) -> Result<ConfigBootstrap, Status> {
const MAX_BUILD_ATTEMPTS: usize = 3;
for _ in 0..MAX_BUILD_ATTEMPTS {
// Components are independent projections. The provider revision is a
// fence for the only overlapping input between sandbox configuration
// and provider environment state.
let (sandbox_config, provider_environment) = tokio::join!(
build_sandbox_config_snapshot(state, sandbox),
build_provider_environment_snapshot(state, sandbox, true),
);
let bootstrap = ConfigBootstrap {
sandbox_config: Some(sandbox_config?),
provider_environment: Some(provider_environment?),
};
if bootstrap_revisions_match(&bootstrap) {
return Ok(bootstrap);
}
counter!("openshell_supervisor_config_bootstrap_revision_mismatches_total").increment(1);
}
Err(Status::aborted(
"configuration changed while building supervisor bootstrap",
))
}
fn bootstrap_revisions_match(bootstrap: &ConfigBootstrap) -> bool {
bootstrap
.sandbox_config
.as_ref()
.zip(bootstrap.provider_environment.as_ref())
.is_some_and(|(sandbox, provider)| {
sandbox.provider_env_revision == provider.provider_env_revision
})
}
pub fn publish_sandbox_components(
state: &Arc<ServerState>,
sandbox_id: &str,
components: ConfigComponents,
) {
enqueue_sandbox(state, sandbox_id, components);
}
fn enqueue_sandbox(state: &Arc<ServerState>, sandbox_id: &str, components: ConfigComponents) {
for component in components.selected() {
let key = DeliveryKey {
sandbox_id: sandbox_id.to_string(),
component,
};
if state.config_delivery_queue.enqueue(key.clone()) {
let state = Arc::clone(state);
tokio::spawn(async move {
loop {
state.config_delivery_queue.take(&key);
publish_sandbox_component_now(&state, &key).await;
if !state.config_delivery_queue.finish_pass(&key) {
break;
}
}
});
}
}
}
async fn publish_sandbox_component_now(state: &Arc<ServerState>, key: &DeliveryKey) {
let sandbox = match state.store.get_message::<Sandbox>(&key.sandbox_id).await {
Ok(Some(sandbox)) => sandbox,
Ok(None) => return,
Err(_) => {
record_build_failure(&key.sandbox_id, "sandbox", Code::Internal);
return;
}
};
let component = key.component.name();
let build = async {
match key.component {
ConfigComponentKind::SandboxConfig => build_sandbox_config_snapshot(state, &sandbox)
.await
.map(|snapshot| SupervisorConfigMessage::SandboxConfig(Box::new(snapshot))),
ConfigComponentKind::ProviderEnvironment => {
build_provider_environment_snapshot(state, &sandbox, true)
.await
.map(SupervisorConfigMessage::ProviderEnvironment)
}
}
};
match tokio::time::timeout(CONFIG_SNAPSHOT_BUILD_TIMEOUT, build).await {
Ok(Ok(message)) => {
let disposition = state
.supervisor_config_router()
.deliver(&key.sandbox_id, message)
.await;
record_delivery(component, disposition);
}
Ok(Err(error)) => {
record_build_failure(&key.sandbox_id, component, error.code());
}
Err(_) => {
record_build_failure(&key.sandbox_id, component, Code::DeadlineExceeded);
}
}
}
pub fn publish_workspace_components(
state: &Arc<ServerState>,
workspace: &str,
components: ConfigComponents,
) {
enqueue_fanout(
state,
FanoutScope::Workspace(workspace.to_string()),
components,
);
}
pub fn publish_all_connected(state: &Arc<ServerState>, components: ConfigComponents) {
enqueue_fanout(state, FanoutScope::AllConnected, components);
}
fn enqueue_fanout(state: &Arc<ServerState>, scope: FanoutScope, components: ConfigComponents) {
for component in components.selected() {
let key = FanoutKey {
scope: scope.clone(),
component,
};
match state.config_delivery_queue.enqueue_fanout(key.clone()) {
FanoutEnqueue::StartWorker => {
let state = Arc::clone(state);
tokio::spawn(async move {
loop {
state.config_delivery_queue.take_fanout(&key);
publish_fanout_now(&state, &key).await;
if !state.config_delivery_queue.finish_fanout_pass(&key) {
break;
}
}
});
}
FanoutEnqueue::Coalesced => {}
FanoutEnqueue::Full => {
counter!("openshell_supervisor_config_fanout_total", "outcome" => "queue_full")
.increment(1);
warn!(
component = component.name(),
"supervisor configuration fanout queue is full"
);
}
}
}
}
async fn publish_fanout_now(state: &Arc<ServerState>, key: &FanoutKey) {
let sandbox_ids = state
.supervisor_config_router()
.routable_sandbox_ids()
.await;
for sandbox_id in sandbox_ids {
if let FanoutScope::Workspace(workspace) = &key.scope {
let sandbox = match state.store.get_message::<Sandbox>(&sandbox_id).await {
Ok(Some(sandbox)) => sandbox,
Ok(None) => continue,
Err(_) => {
record_build_failure(&sandbox_id, "sandbox", Code::Internal);
continue;
}
};
if sandbox.object_workspace() != workspace {
continue;
}
}
enqueue_sandbox(state, &sandbox_id, ConfigComponents::only(key.component));
}
}
fn record_delivery(component: &'static str, disposition: DeliveryDisposition) {
let outcome = match disposition {
DeliveryDisposition::Enqueued => "enqueued",
DeliveryDisposition::NoActiveSession => "no_active_session",
DeliveryDisposition::QueueFull => "queue_full",
DeliveryDisposition::SessionClosed => "session_closed",
DeliveryDisposition::PayloadTooLarge => "payload_too_large",
};
counter!(
"openshell_supervisor_config_deliveries_total",
"component" => component,
"outcome" => outcome,
)
.increment(1);
}
fn record_build_failure(sandbox_id: &str, component: &'static str, error_code: Code) {
counter!(
"openshell_supervisor_config_snapshot_failures_total",
"component" => component,
)
.increment(1);
warn!(
sandbox_id = %sandbox_id,
component,
?error_code,
"failed to build supervisor configuration snapshot"
);
}
#[cfg(test)]
mod tests {
use super::*;
use crate::grpc::{OpenShellService, test_support::test_server_state};
use openshell_core::proto::{
GatewayMessage, ObjectMeta, SandboxSpec, SupervisorHello, SupervisorMessage,
gateway_message, open_shell_client::OpenShellClient, open_shell_server::OpenShellServer,
supervisor_message,
};
use tokio::sync::mpsc;
use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream};
fn key(sandbox_id: &str, component: ConfigComponentKind) -> DeliveryKey {
DeliveryKey {
sandbox_id: sandbox_id.to_string(),
component,
}
}
#[test]
fn queue_coalesces_repeated_component_changes_while_worker_is_active() {
let queue = ConfigDeliveryQueue::default();
let key = key("sb-1", ConfigComponentKind::SandboxConfig);
assert!(queue.enqueue(key.clone()));
queue.take(&key);
assert!(!queue.enqueue(key.clone()));
assert!(!queue.enqueue(key.clone()));
assert!(queue.finish_pass(&key));
queue.take(&key);
assert!(!queue.finish_pass(&key));
}
#[test]
fn queue_runs_components_and_sandboxes_independently() {
let queue = ConfigDeliveryQueue::default();
assert!(queue.enqueue(key("sb-1", ConfigComponentKind::SandboxConfig)));
assert!(queue.enqueue(key("sb-1", ConfigComponentKind::ProviderEnvironment)));
assert!(queue.enqueue(key("sb-2", ConfigComponentKind::SandboxConfig)));
}
#[test]
fn fanout_queue_coalesces_and_bounds_distinct_scopes() {
let queue = ConfigDeliveryQueue::default();
let first = FanoutKey {
scope: FanoutScope::Workspace("workspace-0".into()),
component: ConfigComponentKind::SandboxConfig,
};
assert_eq!(
queue.enqueue_fanout(first.clone()),
FanoutEnqueue::StartWorker
);
queue.take_fanout(&first);
assert_eq!(
queue.enqueue_fanout(first.clone()),
FanoutEnqueue::Coalesced
);
assert!(queue.finish_fanout_pass(&first));
for index in 1..MAX_ACTIVE_FANOUT_WORKERS {
assert_eq!(
queue.enqueue_fanout(FanoutKey {
scope: FanoutScope::Workspace(format!("workspace-{index}")),
component: ConfigComponentKind::SandboxConfig,
}),
FanoutEnqueue::StartWorker
);
}
assert_eq!(
queue.enqueue_fanout(FanoutKey {
scope: FanoutScope::Workspace("overflow".into()),
component: ConfigComponentKind::SandboxConfig,
}),
FanoutEnqueue::Full
);
}
#[test]
fn configuration_message_debug_output_redacts_payloads() {
let message = SupervisorConfigMessage::ProviderEnvironment(ProviderEnvironmentSnapshot {
values: vec![openshell_core::proto::ProviderEnvironmentValue {
name: "TOKEN".into(),
value: "secret-marker".into(),
..Default::default()
}],
..Default::default()
});
assert!(!format!("{message:?}").contains("secret-marker"));
}
#[tokio::test]
async fn session_acceptance_precedes_live_configuration_updates() {
let state = test_server_state().await;
state
.store
.put_message(&Sandbox {
metadata: Some(ObjectMeta {
id: "sandbox".into(),
name: "sandbox".into(),
workspace: "default".into(),
..Default::default()
}),
spec: Some(SandboxSpec::default()),
..Default::default()
})
.await
.unwrap();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(
tonic::transport::Server::builder()
.add_service(OpenShellServer::new(OpenShellService::new(Arc::clone(
&state,
))))
.serve_with_incoming(TcpListenerStream::new(listener)),
);
let mut client = OpenShellClient::connect(format!("http://{address}"))
.await
.unwrap();
let (tx, rx) = mpsc::channel(4);
tx.send(SupervisorMessage {
payload: Some(supervisor_message::Payload::Hello(SupervisorHello {
sandbox_id: "sandbox".into(),
instance_id: "instance".into(),
protocol_revision: openshell_core::proto::SUPERVISOR_PROTOCOL_REVISION,
})),
})
.await
.unwrap();
let mut stream = client
.connect_supervisor(ReceiverStream::new(rx))
.await
.unwrap()
.into_inner();
let first = tokio::time::timeout(Duration::from_secs(5), stream.message())
.await
.unwrap()
.unwrap()
.unwrap();
assert!(matches!(
first.payload,
Some(gateway_message::Payload::SessionAccepted(_))
));
publish_sandbox_components(&state, "sandbox", ConfigComponents::SANDBOX_CONFIG);
let update = tokio::time::timeout(Duration::from_secs(5), stream.message())
.await
.unwrap()
.unwrap()
.unwrap();
assert!(matches!(
update,
GatewayMessage {
payload: Some(gateway_message::Payload::ConfigUpdate(_))
}
));
drop(tx);
server.abort();
}
#[test]
fn bootstrap_requires_matching_provider_revision_fence() {
let mut bootstrap = ConfigBootstrap {
sandbox_config: Some(SandboxConfigSnapshot {
provider_env_revision: 7,
..Default::default()
}),
provider_environment: Some(ProviderEnvironmentSnapshot {
provider_env_revision: 8,
..Default::default()
}),
};
assert!(!bootstrap_revisions_match(&bootstrap));
bootstrap
.provider_environment
.as_mut()
.unwrap()
.provider_env_revision = 7;
assert!(bootstrap_revisions_match(&bootstrap));
}
}
+538 -126
View File
@@ -29,6 +29,8 @@ use crate::storage_proto::StoredProviderCredentialRefreshState;
#[cfg(test)]
use crate::storage_proto::StoredProviderProfile;
use openshell_core::net::{is_always_blocked_ip, is_internal_ip};
#[cfg(test)]
use openshell_core::proto::StaticCredentialBinding;
use openshell_core::proto::policy_merge_operation;
use openshell_core::proto::setting_value;
use openshell_core::proto::{
@@ -42,11 +44,12 @@ use openshell_core::proto::{
GetSandboxPolicyStatusRequest, GetSandboxPolicyStatusResponse,
GetSandboxProviderEnvironmentRequest, GetSandboxProviderEnvironmentResponse,
ListSandboxPoliciesRequest, ListSandboxPoliciesResponse, PolicyChunk, PolicyMergeOperation,
PolicySource, PolicyStatus, PushSandboxLogsRequest, PushSandboxLogsResponse,
PolicySource, PolicyStatus, ProviderEnvironmentSnapshot, ProviderEnvironmentValue,
ProviderEnvironmentValueClassification, PushSandboxLogsRequest, PushSandboxLogsResponse,
RejectDraftChunkRequest, RejectDraftChunkResponse, ReportPolicyStatusRequest,
ReportPolicyStatusResponse, SandboxLogLine, SandboxPolicyRevision, SettingScope, SettingValue,
SubmitPolicyAnalysisRequest, SubmitPolicyAnalysisResponse, UndoDraftChunkRequest,
UndoDraftChunkResponse, UpdateConfigRequest, UpdateConfigResponse,
ReportPolicyStatusResponse, SandboxConfigSnapshot, SandboxLogLine, SandboxPolicyRevision,
SettingScope, SettingValue, SubmitPolicyAnalysisRequest, SubmitPolicyAnalysisResponse,
UndoDraftChunkRequest, UndoDraftChunkResponse, UpdateConfigRequest, UpdateConfigResponse,
};
use openshell_core::proto::{
L7DenyRule, L7Rule, NetworkBinary, NetworkEndpoint, NetworkPolicyRule, Provider, Sandbox,
@@ -1538,6 +1541,11 @@ async fn auto_approve_chunk(
return Err(status);
}
};
crate::config_delivery::publish_sandbox_components(
state,
sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
let chunk_summary = summarize_draft_chunk_rule(&chunk)?;
let now_ms = current_time_ms();
@@ -2374,6 +2382,18 @@ pub(super) async fn handle_get_sandbox_config(
let sandbox =
super::sandbox::fetch_and_authorize_sandbox(state, &principal, &sandbox_id).await?;
let snapshot = build_sandbox_config_snapshot(state, &sandbox).await?;
Ok(Response::new(sandbox_config_response(snapshot)))
}
/// Build the complete effective configuration for one persisted sandbox.
///
/// This is a read-only projection shared by polling and stream delivery.
pub async fn build_sandbox_config_snapshot(
state: &Arc<ServerState>,
sandbox: &Sandbox,
) -> Result<SandboxConfigSnapshot, Status> {
let sandbox_id = sandbox.object_id().to_string();
let workspace = sandbox.object_workspace().to_string();
let sandbox_provider_names = sandbox
.spec
@@ -2398,87 +2418,49 @@ pub(super) async fn handle_get_sandbox_config(
.await
.map_err(|e| Status::internal(format!("fetch policy history failed: {e}")))?;
let (mut policy, version, mut policy_hash, policy_source) = if let Some(global_policy) =
global_policy
{
let version = latest
.as_ref()
.map(|record| u32::try_from(record.version).unwrap_or(0))
.filter(|version| *version > 0)
.unwrap_or(1);
let hash = deterministic_policy_hash(&global_policy);
(Some(global_policy), version, hash, PolicySource::Global)
} else if let Some(record) = latest {
let (policy, hash) = canonical_policy_record_identity(&record)?;
debug!(
sandbox_id = %sandbox_id,
version = record.version,
"GetSandboxConfig served from policy history"
);
(
Some(policy),
u32::try_from(record.version).unwrap_or(0),
hash,
PolicySource::Sandbox,
)
} else {
// Lazy backfill: no policy history exists yet.
let spec = sandbox
.spec
.as_ref()
.ok_or_else(|| Status::internal("sandbox has no spec"))?;
let (mut policy, version, mut policy_hash, policy_source) =
if let Some(global_policy) = global_policy {
let version = latest
.as_ref()
.map(|record| u32::try_from(record.version).unwrap_or(0))
.filter(|version| *version > 0)
.unwrap_or(1);
let hash = deterministic_policy_hash(&global_policy);
(Some(global_policy), version, hash, PolicySource::Global)
} else if let Some(record) = latest {
let (policy, hash) = canonical_policy_record_identity(&record)?;
debug!(
sandbox_id = %sandbox_id,
version = record.version,
"GetSandboxConfig served from policy history"
);
(
Some(policy),
u32::try_from(record.version).unwrap_or(0),
hash,
PolicySource::Sandbox,
)
} else {
// Older sandboxes may have policy only in the sandbox spec. Reading a
// snapshot must not create policy history, so project that baseline as
// version 1 until the startup repair persists it.
let spec = sandbox
.spec
.as_ref()
.ok_or_else(|| Status::internal("sandbox has no spec"))?;
match spec.policy.clone() {
None => {
debug!(
sandbox_id = %sandbox_id,
"GetSandboxConfig: no policy configured, returning empty response"
);
(None, 0, String::new(), PolicySource::Sandbox)
}
Some(spec_policy) => {
// Stored specs may predate the current schema. Validate before
// creating policy history so malformed state is never copied or
// marked loaded, and hash the canonical representation.
let spec_policy = validate_and_canonicalize_stored_policy(
spec_policy,
STORED_POLICY_SOURCE_SPEC,
)?;
let hash = deterministic_policy_hash(&spec_policy);
let payload = spec_policy.encode_to_vec();
let policy_id = uuid::Uuid::new_v4().to_string();
if let Err(e) = state
.store
.put_policy_revision(&policy_id, &sandbox_id, &workspace, 1, &payload, &hash)
.await
{
warn!(
sandbox_id = %sandbox_id,
error = %e,
"Failed to backfill policy version 1"
);
} else if let Err(e) = state
.store
.update_policy_status(&sandbox_id, 1, "loaded", None, None)
.await
{
warn!(
sandbox_id = %sandbox_id,
error = %e,
"Failed to mark backfilled policy as loaded"
);
match spec.policy.clone() {
None => (None, 0, String::new(), PolicySource::Sandbox),
Some(spec_policy) => {
let spec_policy = validate_and_canonicalize_stored_policy(
spec_policy,
STORED_POLICY_SOURCE_SPEC,
)?;
let hash = deterministic_policy_hash(&spec_policy);
(Some(spec_policy), 1, hash, PolicySource::Sandbox)
}
info!(
sandbox_id = %sandbox_id,
"GetSandboxConfig served from spec (backfilled version 1)"
);
(Some(spec_policy), 1, hash, PolicySource::Sandbox)
}
}
};
};
let global_settings = load_global_settings(state.store.as_ref()).await?;
let sandbox_settings =
@@ -2585,7 +2567,7 @@ pub(super) async fn handle_get_sandbox_config(
)
.await?;
Ok(Response::new(GetSandboxConfigResponse {
Ok(SandboxConfigSnapshot {
policy,
version,
policy_hash,
@@ -2602,7 +2584,112 @@ pub(super) async fn handle_get_sandbox_config(
.as_str()
.to_string(),
extension_authentication_enabled: state.sandbox_jwt_issuer.is_some(),
}))
})
}
fn sandbox_config_response(snapshot: SandboxConfigSnapshot) -> GetSandboxConfigResponse {
GetSandboxConfigResponse {
policy: snapshot.policy,
version: snapshot.version,
policy_hash: snapshot.policy_hash,
settings: snapshot.settings,
config_revision: snapshot.config_revision,
policy_source: snapshot.policy_source,
global_policy_version: snapshot.global_policy_version,
provider_env_revision: snapshot.provider_env_revision,
supervisor_middleware_services: snapshot.supervisor_middleware_services,
workspace: snapshot.workspace,
policy_validation_failure_mode: snapshot.policy_validation_failure_mode,
extension_authentication_enabled: snapshot.extension_authentication_enabled,
}
}
#[derive(Clone, Copy)]
pub enum InitialPolicyHistoryStatus {
Pending,
Loaded,
}
impl InitialPolicyHistoryStatus {
fn as_str(self) -> &'static str {
match self {
Self::Pending => "pending",
Self::Loaded => "loaded",
}
}
}
/// Insert the version-one policy baseline if this sandbox still has no policy
/// history. This never modifies an existing revision or apply result.
pub async fn initialize_policy_history(
store: &Store,
sandbox: &Sandbox,
status: InitialPolicyHistoryStatus,
) -> Result<(), Status> {
let Some(policy) = sandbox.spec.as_ref().and_then(|spec| spec.policy.as_ref()) else {
return Ok(());
};
if store
.get_latest_policy(sandbox.object_id())
.await
.map_err(|error| Status::internal(format!("read policy history failed: {error}")))?
.is_some()
{
return Ok(());
}
let policy =
validate_and_canonicalize_stored_policy(policy.clone(), STORED_POLICY_SOURCE_SPEC)?;
store
.put_initial_policy_revision(
&PolicyRecord {
id: uuid::Uuid::new_v4().to_string(),
sandbox_id: sandbox.object_id().to_string(),
version: 1,
policy_payload: policy.encode_to_vec(),
policy_hash: deterministic_policy_hash(&policy),
status: status.as_str().to_string(),
load_error: None,
created_at_ms: current_time_ms(),
loaded_at_ms: None,
provenance: HashMap::new(),
},
sandbox.object_workspace(),
)
.await
.map_err(|error| Status::internal(format!("initialize policy history failed: {error}")))
}
/// Create policy-history baselines for sandboxes written by older gateways.
///
/// Snapshot reads stay pure once this startup repair has completed.
pub async fn backfill_legacy_policy_history(state: &Arc<ServerState>) -> Result<(), Status> {
const PAGE_SIZE: u32 = 1000;
let mut offset = 0;
loop {
let sandboxes = state
.store
.list_all_messages::<Sandbox>(PAGE_SIZE, offset)
.await
.map_err(|error| {
Status::internal(format!("list sandboxes for policy repair failed: {error}"))
})?;
if sandboxes.is_empty() {
return Ok(());
}
let count = u32::try_from(sandboxes.len()).unwrap_or(PAGE_SIZE);
for sandbox in sandboxes {
initialize_policy_history(
state.store.as_ref(),
&sandbox,
InitialPolicyHistoryStatus::Loaded,
)
.await?;
}
if count < PAGE_SIZE {
return Ok(());
}
offset = offset.saturating_add(count);
}
}
#[cfg(test)]
@@ -3105,6 +3192,19 @@ pub(super) async fn handle_get_sandbox_provider_environment(
.await
.map_err(|e| Status::internal(format!("fetch sandbox failed: {e}")))?
.ok_or_else(|| Status::not_found("sandbox not found"))?;
let snapshot =
build_provider_environment_snapshot(state, &sandbox, supports_static_credential_bindings)
.await?;
Ok(Response::new(provider_environment_response(snapshot)))
}
/// Build the complete provider environment for one persisted sandbox.
pub async fn build_provider_environment_snapshot(
state: &Arc<ServerState>,
sandbox: &Sandbox,
supports_static_credential_bindings: bool,
) -> Result<ProviderEnvironmentSnapshot, Status> {
let sandbox_id = sandbox.object_id().to_string();
let workspace = sandbox.object_workspace().to_string();
let spec = sandbox
@@ -3127,7 +3227,7 @@ pub(super) async fn handle_get_sandbox_provider_environment(
state.as_ref(),
&provider_profile_catalog,
&workspace,
&sandbox,
sandbox,
&sandbox_id,
)
.await?;
@@ -3192,21 +3292,81 @@ pub(super) async fn handle_get_sandbox_provider_environment(
"GetSandboxProviderEnvironment request completed successfully"
);
let non_secret_environment_keys = provider_environment
let mut keys = provider_environment
.environment
.keys()
.filter(|key| !provider_environment.static_credential_keys.contains(*key))
.cloned()
.collect();
.collect::<Vec<_>>();
keys.sort();
let mut values = Vec::with_capacity(keys.len());
for name in keys {
let value = provider_environment
.environment
.remove(&name)
.expect("provider environment key came from the same map");
let is_static_credential = provider_environment.static_credential_keys.contains(&name);
let static_credential_binding = provider_environment
.static_credential_bindings
.remove(&name);
if is_static_credential && static_credential_binding.is_none() {
return Err(Status::failed_precondition(format!(
"static provider credential '{name}' has no endpoint binding"
)));
}
values.push(ProviderEnvironmentValue {
name: name.clone(),
value,
expires_at_ms: provider_environment.credential_expires_at_ms.remove(&name),
classification: if is_static_credential {
ProviderEnvironmentValueClassification::StaticCredential.into()
} else {
ProviderEnvironmentValueClassification::NonSecret.into()
},
static_credential_binding,
});
}
Ok(Response::new(GetSandboxProviderEnvironmentResponse {
environment: provider_environment.environment,
Ok(ProviderEnvironmentSnapshot {
provider_env_revision,
credential_expires_at_ms: provider_environment.credential_expires_at_ms,
values,
dynamic_credentials: provider_environment.dynamic_credentials,
static_credential_bindings: provider_environment.static_credential_bindings,
})
}
fn provider_environment_response(
snapshot: ProviderEnvironmentSnapshot,
) -> GetSandboxProviderEnvironmentResponse {
let mut environment = HashMap::with_capacity(snapshot.values.len());
let mut credential_expires_at_ms = HashMap::new();
let mut static_credential_bindings = HashMap::new();
let mut non_secret_environment_keys = Vec::new();
for value in snapshot.values {
environment.insert(value.name.clone(), value.value);
if let Some(expires_at_ms) = value.expires_at_ms {
credential_expires_at_ms.insert(value.name.clone(), expires_at_ms);
}
match ProviderEnvironmentValueClassification::try_from(value.classification)
.unwrap_or_default()
{
ProviderEnvironmentValueClassification::NonSecret => {
non_secret_environment_keys.push(value.name);
}
ProviderEnvironmentValueClassification::StaticCredential => {
if let Some(binding) = value.static_credential_binding {
static_credential_bindings.insert(value.name, binding);
}
}
ProviderEnvironmentValueClassification::Unspecified => {}
}
}
GetSandboxProviderEnvironmentResponse {
environment,
provider_env_revision: snapshot.provider_env_revision,
credential_expires_at_ms,
dynamic_credentials: snapshot.dynamic_credentials,
static_credential_bindings,
non_secret_environment_keys,
}))
}
}
// ---------------------------------------------------------------------------
@@ -3358,6 +3518,10 @@ async fn handle_update_config_inner(
if changed {
global_settings.revision = global_settings.revision.wrapping_add(1);
save_global_settings(state.store.as_ref(), &global_settings).await?;
crate::config_delivery::publish_all_connected(
state,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
}
return Ok(update_config_response(
u32::try_from(current.version).unwrap_or(0),
@@ -3414,6 +3578,10 @@ async fn handle_update_config_inner(
if changed {
global_settings.revision = global_settings.revision.wrapping_add(1);
save_global_settings(state.store.as_ref(), &global_settings).await?;
crate::config_delivery::publish_all_connected(
state,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
}
return Ok(update_config_response(
@@ -3458,6 +3626,10 @@ async fn handle_update_config_inner(
global_settings.revision = global_settings.revision.wrapping_add(1);
save_global_settings(state.store.as_ref(), &global_settings).await?;
crate::config_delivery::publish_all_connected(
state,
crate::config_delivery::ConfigComponents::SANDBOX_CONFIG,
);
if req.delete_setting
&& key == POLICY_SETTING_KEY
@@ -3530,6 +3702,11 @@ async fn handle_update_config_inner(
&sandbox_settings,
)
.await?;
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_CONFIG,
);
}
response_annotations = persist_update_config_annotations(
@@ -3574,6 +3751,11 @@ async fn handle_update_config_inner(
&sandbox_settings,
)
.await?;
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_CONFIG,
);
}
response_annotations = persist_update_config_annotations(
@@ -3632,6 +3814,11 @@ async fn handle_update_config_inner(
Some(&atomic_context),
)
.await?;
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
response_annotations = if let Some(updated_sandbox) = updated_sandbox {
sandbox_metadata_annotations(&updated_sandbox)
} else {
@@ -3851,6 +4038,11 @@ async fn handle_update_config_inner(
})?
};
response_annotations = committed_annotations;
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
state.sandbox_watch_bus.notify(&sandbox_id);
if backfill_policy.is_some() {
@@ -3897,6 +4089,12 @@ async fn handle_update_config_inner(
.await
.map_err(|e| Status::internal(format!("persist policy revision failed: {e}")))?;
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
let _ = state
.store
.supersede_older_policies(&sandbox_id, next_version)
@@ -4891,6 +5089,11 @@ async fn handle_approve_draft_chunk_inner(
return Err(status);
}
};
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
let chunk_summary = summarize_draft_chunk_rule(&chunk)?;
let now_ms = current_time_ms();
@@ -5015,6 +5218,11 @@ async fn handle_reject_draft_chunk_inner(
require_no_global_policy(state).await?;
let (version, hash) =
remove_chunk_from_policy(state, &sandbox_id, &workspace, &chunk).await?;
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
emit_gateway_policy_audit_log(
&sandbox_id,
sandbox.object_name(),
@@ -5306,6 +5514,14 @@ async fn handle_approve_all_draft_chunks_inner(
}
};
if !accepted.is_empty() {
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
}
for (chunk, _, chunk_summary) in &accepted {
let now_ms = current_time_ms();
clear_pending_application_error(state, &chunk.id).await;
@@ -5507,6 +5723,11 @@ async fn handle_undo_draft_chunk_inner(
);
let (version, hash) = remove_chunk_from_policy(state, &sandbox_id, &workspace, &chunk).await?;
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
// Clear any prior rejection_reason on the way back to "pending" so an
// agent reading the chunk via policy.local cannot see a stale guidance
@@ -5862,7 +6083,7 @@ fn canonical_policy_bytes(policy: &ProtoSandboxPolicy) -> Vec<u8> {
/// Compute a deterministic SHA-256 hash of a `SandboxPolicy`, recursively
/// sorting every protobuf map while preserving repeated-field order.
fn deterministic_policy_hash(policy: &ProtoSandboxPolicy) -> String {
pub fn deterministic_policy_hash(policy: &ProtoSandboxPolicy) -> String {
hex::encode(Sha256::digest(canonical_policy_bytes(policy)))
}
@@ -7438,7 +7659,7 @@ mod tests {
}
#[tokio::test]
async fn get_sandbox_config_backfills_canonical_spec_policy_bytes_and_hash() {
async fn startup_repair_backfills_canonical_spec_policy_bytes_and_hash() {
let state = test_server_state().await;
let sandbox_id = "stored-canonical-backfill";
let raw = mcp_policy_with_versions(&["2025-11-25", "2025-03-26", "2025-06-18"]);
@@ -7473,6 +7694,15 @@ mod tests {
assert_eq!(response.policy_hash, canonical_hash);
assert_eq!(response.version, 1);
assert!(
state
.store
.get_latest_policy(sandbox_id)
.await
.unwrap()
.is_none()
);
backfill_legacy_policy_history(&state).await.unwrap();
let persisted = state
.store
.get_latest_policy(sandbox_id)
@@ -7485,7 +7715,7 @@ mod tests {
}
#[tokio::test]
async fn get_sandbox_config_backfills_defaulted_mcp_policy_as_canonical_bytes_and_hash() {
async fn startup_repair_backfills_defaulted_mcp_policy_as_canonical_bytes_and_hash() {
let state = test_server_state().await;
let canonical = validate_and_canonicalize_policy(mcp_policy_with_versions(&["2025-11-25"]))
.expect("explicit default MCP policy must canonicalize");
@@ -7531,6 +7761,15 @@ mod tests {
"{case}"
);
assert!(
state
.store
.get_latest_policy(&sandbox_id)
.await
.unwrap()
.is_none()
);
backfill_legacy_policy_history(&state).await.unwrap();
let persisted = state
.store
.get_latest_policy(&sandbox_id)
@@ -10594,6 +10833,114 @@ mod tests {
);
}
#[tokio::test]
async fn rejected_policy_update_does_not_publish_configuration() {
let state = test_server_state().await;
let mut sandbox = test_sandbox(
"sb-rejected-policy",
"rejected-policy",
ProtoSandboxPolicy::default(),
Vec::new(),
);
sandbox.spec.as_mut().unwrap().policy = None;
state.store.put_message(&sandbox).await.unwrap();
let (tx, mut rx) = tokio::sync::mpsc::channel(1);
let (shutdown_tx, _shutdown_rx) = tokio::sync::oneshot::channel();
state.supervisor_sessions.register(
"sb-rejected-policy".to_string(),
"session-1".to_string(),
tx,
shutdown_tx,
);
let error = handle_update_config(
&state,
with_user(Request::new(UpdateConfigRequest {
name: "rejected-policy".to_string(),
workspace: "default".to_string(),
policy: Some(test_policy_with_rule("rejected", "api.example.com")),
expected_resource_version: u64::MAX,
..Default::default()
})),
)
.await
.expect_err("stale resource version must reject the update");
assert_eq!(error.code(), Code::Aborted);
assert!(rx.try_recv().is_err());
assert!(
state
.store
.get_latest_policy("sb-rejected-policy")
.await
.unwrap()
.is_none()
);
}
#[tokio::test]
async fn committed_policy_update_publishes_complete_snapshot() {
let state = test_server_state().await;
let mut sandbox = test_sandbox(
"sb-published-policy",
"published-policy",
ProtoSandboxPolicy::default(),
Vec::new(),
);
sandbox.spec.as_mut().unwrap().policy = None;
state.store.put_message(&sandbox).await.unwrap();
let (tx, mut rx) = tokio::sync::mpsc::channel(4);
let (shutdown_tx, _shutdown_rx) = tokio::sync::oneshot::channel();
state.supervisor_sessions.register(
"sb-published-policy".to_string(),
"session-1".to_string(),
tx,
shutdown_tx,
);
handle_update_config(
&state,
with_user(Request::new(UpdateConfigRequest {
name: "published-policy".to_string(),
workspace: "default".to_string(),
policy: Some(test_policy_with_rule("published", "api.example.com")),
..Default::default()
})),
)
.await
.unwrap();
let persisted = state
.store
.get_latest_policy("sb-published-policy")
.await
.unwrap()
.expect("committed policy");
let snapshot = tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
let message = rx.recv().await.expect("configuration channel closed");
let Some(openshell_core::proto::gateway_message::Payload::ConfigUpdate(update)) =
message.payload
else {
panic!("expected ConfigUpdate");
};
if let Some(openshell_core::proto::config_update::Component::SandboxConfig(
snapshot,
)) = update.component
{
break snapshot;
}
}
})
.await
.expect("configuration publication timed out");
assert_eq!(snapshot.version, u32::try_from(persisted.version).unwrap());
assert_eq!(snapshot.policy_hash, persisted.policy_hash);
assert!(snapshot.policy.is_some());
}
#[tokio::test]
async fn update_config_accepts_sigv4_covered_by_endpointful_aws_profile() {
let state = test_server_state().await;
@@ -10921,23 +11268,68 @@ mod tests {
.contains_key("_provider_work_github")
);
let persisted = state
assert!(
state
.store
.get_latest_policy("sb-jit")
.await
.unwrap()
.is_none(),
"snapshot reads must not create policy history"
);
}
#[tokio::test]
async fn legacy_policy_history_repair_is_idempotent() {
let state = test_server_state().await;
let policy = test_policy_with_rule("legacy", "legacy.example.com");
state
.store
.get_latest_policy("sb-jit")
.put_message(&test_sandbox(
"sb-legacy-policy",
"legacy-policy",
policy.clone(),
Vec::new(),
))
.await
.unwrap();
backfill_legacy_policy_history(&state).await.unwrap();
let initial = state
.store
.get_latest_policy("sb-legacy-policy")
.await
.unwrap()
.expect("sandbox policy should be lazily backfilled");
let persisted_policy = ProtoSandboxPolicy::decode(persisted.policy_payload.as_slice())
.expect("persisted sandbox policy should decode");
assert!(
persisted_policy
.network_policies
.contains_key("sandbox_only")
.expect("legacy policy baseline");
assert_eq!(initial.status, "loaded");
state
.store
.update_policy_status("sb-legacy-policy", 1, "failed", Some("apply failed"), None)
.await
.unwrap();
backfill_legacy_policy_history(&state).await.unwrap();
let repaired_again = state
.store
.get_latest_policy("sb-legacy-policy")
.await
.unwrap()
.expect("legacy policy baseline");
assert_eq!(repaired_again.version, 1);
assert_eq!(
repaired_again.policy_hash,
deterministic_policy_hash(&policy)
);
assert!(
!persisted_policy
.network_policies
.contains_key("_provider_work_github")
assert_eq!(repaired_again.status, "failed");
assert_eq!(repaired_again.load_error.as_deref(), Some("apply failed"));
assert_eq!(
state
.store
.list_policies("sb-legacy-policy", 10, 0)
.await
.unwrap()
.len(),
1
);
}
@@ -11062,24 +11454,14 @@ mod tests {
assert_eq!(persisted_provider.r#type, provider.r#type);
assert_eq!(persisted_provider.credentials, provider.credentials);
let persisted_policy = state
.store
.get_latest_policy("sb-custom-policy-update")
.await
.unwrap()
.expect("sandbox policy should be lazily backfilled");
let persisted_policy =
ProtoSandboxPolicy::decode(persisted_policy.policy_payload.as_slice())
.expect("persisted sandbox policy should decode");
assert!(
persisted_policy
.network_policies
.contains_key("sandbox_only")
);
assert!(
!persisted_policy
.network_policies
.contains_key("_provider_work_custom")
state
.store
.get_latest_policy("sb-custom-policy-update")
.await
.unwrap()
.is_none(),
"config and profile reads must not create policy history"
);
}
@@ -20869,4 +21251,34 @@ mod tests {
response.unwrap_err()
);
}
#[test]
fn provider_stream_values_expand_to_legacy_polling_response() {
let response = provider_environment_response(ProviderEnvironmentSnapshot {
provider_env_revision: 9,
values: vec![
ProviderEnvironmentValue {
name: "REGION".into(),
value: "west".into(),
classification: ProviderEnvironmentValueClassification::NonSecret.into(),
..Default::default()
},
ProviderEnvironmentValue {
name: "TOKEN".into(),
value: "secret".into(),
expires_at_ms: Some(123),
classification: ProviderEnvironmentValueClassification::StaticCredential.into(),
static_credential_binding: Some(StaticCredentialBinding::default()),
},
],
dynamic_credentials: HashMap::new(),
});
assert_eq!(response.provider_env_revision, 9);
assert_eq!(response.environment["REGION"], "west");
assert_eq!(response.environment["TOKEN"], "secret");
assert_eq!(response.credential_expires_at_ms["TOKEN"], 123);
assert_eq!(response.non_secret_environment_keys, ["REGION"]);
assert!(response.static_credential_bindings.contains_key("TOKEN"));
}
}
+46 -5
View File
@@ -2434,6 +2434,21 @@ async fn authorize_and_resolve_profile_workspace(
}
}
fn publish_provider_change(state: &Arc<ServerState>, workspace: &str) {
if workspace.is_empty() {
crate::config_delivery::publish_all_connected(
state,
crate::config_delivery::ConfigComponents::ALL,
);
} else {
crate::config_delivery::publish_workspace_components(
state,
workspace,
crate::config_delivery::ConfigComponents::ALL,
);
}
}
pub(super) async fn handle_create_provider(
state: &Arc<ServerState>,
request: Request<CreateProviderRequest>,
@@ -2493,6 +2508,7 @@ pub(super) async fn handle_create_provider(
LifecycleOperation::Create,
TelemetryOutcome::Success,
);
publish_provider_change(state, &workspace);
Ok(Response::new(ProviderResponse {
provider: Some(provider),
}))
@@ -2755,6 +2771,7 @@ pub(super) async fn handle_import_provider_profiles(
stored.profile.unwrap_or_default(),
resource_version,
));
publish_provider_change(state, &workspace);
}
Ok(Response::new(ImportProviderProfilesResponse {
@@ -2885,6 +2902,7 @@ pub(super) async fn handle_update_provider_profiles(
}
let resource_version = stored_profile_resource_version(&stored);
let profile = profile_response_payload(stored.profile.unwrap_or_default(), resource_version);
publish_provider_change(state, &workspace);
Ok(Response::new(UpdateProviderProfilesResponse {
diagnostics: Vec::new(),
@@ -2974,6 +2992,9 @@ pub(super) async fn handle_delete_provider_profile(
.delete_by_name(StoredProviderProfile::object_type(), &workspace, &id)
.await
.map_err(|e| Status::internal(format!("delete provider profile failed: {e}")))?;
if deleted {
publish_provider_change(state, &workspace);
}
Ok(Response::new(DeleteProviderProfileResponse { deleted }))
}
@@ -3754,6 +3775,7 @@ pub(super) async fn handle_update_provider(
LifecycleOperation::Update,
TelemetryOutcome::Success,
);
publish_provider_change(state, &workspace);
Ok(Response::new(ProviderResponse {
provider: Some(provider),
}))
@@ -4673,8 +4695,17 @@ pub(super) async fn handle_configure_provider_refresh(
profile_workspace: String::new(),
credential_handles: HashMap::new(),
};
update_provider_record_with_catalog(state.store.as_ref(), &catalog, &workspace, updated)
.await?;
let result = update_provider_record_with_catalog(
state.store.as_ref(),
&catalog,
&workspace,
updated,
)
.await;
publish_provider_change(state, &workspace);
result?;
} else {
publish_provider_change(state, &workspace);
}
Ok(Response::new(ConfigureProviderRefreshResponse {
@@ -4718,6 +4749,7 @@ pub(super) async fn handle_rotate_provider_credential(
credential_key,
)
.await?;
publish_provider_change(state, &workspace);
Ok(Response::new(RotateProviderCredentialResponse {
status: Some(crate::provider_refresh::refresh_status_from_state(
@@ -4800,14 +4832,13 @@ pub(super) async fn handle_delete_provider_refresh(
credential_key,
)
.await?;
// A refresh co-manages the expiry of its primary credential and every pinned
// additional output. Clear each expiry this refresh still owns, leaving
// independently updated ones in place. The equality check and removal run
// inside the CAS closure so they see the current stored provider — deciding
// from the snapshot read above would let a concurrent rotation or provider
// update land between the read and the write and then be clobbered (CWE-362).
if let Some(refresh_state) = existing_refresh_state
let expiry_cleanup = if let Some(refresh_state) = existing_refresh_state
&& refresh_state.expires_at_ms > 0
{
let refresh_expires_at_ms = refresh_state.expires_at_ms;
@@ -4824,8 +4855,15 @@ pub(super) async fn handle_delete_provider_refresh(
Status::internal(format!(
"clear refresh-owned credential expiries failed: {e}"
))
})?;
})
.map(|_| ())
} else {
Ok(())
};
if deleted_refresh_state {
publish_provider_change(state, &workspace);
}
expiry_cleanup?;
Ok(Response::new(DeleteProviderRefreshResponse {
deleted: deleted_refresh_state,
@@ -4866,6 +4904,9 @@ pub(super) async fn handle_delete_provider(
LifecycleOperation::Delete,
outcome,
);
if deleted {
publish_provider_change(state, &workspace);
}
Ok(Response::new(DeleteProviderResponse { deleted }))
}
Err(err) => {
@@ -1175,6 +1175,14 @@ pub(super) async fn handle_attach_sandbox_provider(
.map_err(|e| super::persistence_error_to_status(e, "attach sandbox provider"))?;
let attached = attached.load(Ordering::Relaxed);
if attached {
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
state.sandbox_watch_bus.notify(&sandbox_id);
}
info!(
sandbox_name = %request.sandbox_name,
@@ -1274,6 +1282,14 @@ pub(super) async fn handle_detach_sandbox_provider(
.map_err(|e| super::persistence_error_to_status(e, "detach sandbox provider"))?;
let detached = detached.load(Ordering::Relaxed);
if detached {
crate::config_delivery::publish_sandbox_components(
state,
&sandbox_id,
crate::config_delivery::ConfigComponents::SANDBOX_AND_PROVIDER,
);
state.sandbox_watch_bus.notify(&sandbox_id);
}
info!(
sandbox_name = %request.sandbox_name,
+21
View File
@@ -17,6 +17,7 @@ mod auth;
pub mod certgen;
pub mod cli;
mod compute;
mod config_delivery;
pub mod config_file;
mod credentials;
mod defaults;
@@ -301,6 +302,11 @@ pub struct ServerState {
/// Set once graceful gateway shutdown begins so stream handlers can
/// distinguish expected transport closes from runtime failures.
pub(crate) gateway_shutting_down: AtomicBool,
/// Per-sandbox scheduler for coalesced supervisor configuration delivery.
pub(crate) config_delivery_queue: config_delivery::ConfigDeliveryQueue,
/// Routing boundary for local or remote supervisor configuration delivery.
pub(crate) supervisor_config_router: Arc<dyn config_delivery::SupervisorConfigRouter>,
/// Validated built-in and operator-registered supervisor middleware.
pub middleware_registry: Arc<MiddlewareRegistry>,
@@ -356,6 +362,12 @@ fn is_benign_connection_close(error: &(dyn std::error::Error + 'static)) -> bool
}
impl ServerState {
/// Return the configuration delivery boundary for supervisor sessions.
#[must_use]
pub fn supervisor_config_router(&self) -> Arc<dyn config_delivery::SupervisorConfigRouter> {
Arc::clone(&self.supervisor_config_router)
}
/// Create new server state.
#[must_use]
#[allow(clippy::too_many_arguments)]
@@ -404,6 +416,9 @@ impl ServerState {
.oidc
.as_ref()
.map_or_else(String::new, |oidc| oidc.admin_role.clone());
let supervisor_config_router: Arc<dyn config_delivery::SupervisorConfigRouter> = Arc::new(
config_delivery::LocalSupervisorConfigRouter::new(Arc::clone(&supervisor_sessions)),
);
Self {
config,
store,
@@ -418,6 +433,8 @@ impl ServerState {
settings_mutex: tokio::sync::Mutex::new(()),
supervisor_sessions,
gateway_shutting_down: AtomicBool::new(false),
config_delivery_queue: config_delivery::ConfigDeliveryQueue::default(),
supervisor_config_router,
extension_mint_limiter: auth::extension_mint_limit::ExtensionMintLimiter::default(),
middleware_registry: Arc::new(MiddlewareRegistry::default()),
oidc_cache,
@@ -680,6 +697,10 @@ pub(crate) async fn run_server(
let state = Arc::new(state);
grpc::policy::backfill_legacy_policy_history(&state)
.await
.map_err(|error| Error::execution(error.to_string()))?;
// Reconcile local-driver running intent before watchers spawn so their
// first snapshots observe the post-start backend state. Explicitly stopped
// sandboxes remain stopped.
@@ -945,6 +945,50 @@ VALUES ($1, $2, $3, $4, $5, $6, $7, $7, $8)
Ok(())
}
pub async fn put_initial_policy_revision(
&self,
record: &PolicyRecord,
workspace: &str,
) -> PersistenceResult<()> {
let wrapped_payload = policy_payload_from_record(record)?;
let mut tx = self.pool.begin().await.map_err(|e| map_db_error(&e))?;
let sandbox_exists = sqlx::query(
"SELECT id FROM objects WHERE object_type = 'sandbox' AND id = $1 FOR UPDATE",
)
.bind(&record.sandbox_id)
.fetch_optional(&mut *tx)
.await
.map_err(|e| map_db_error(&e))?
.is_some();
if sandbox_exists {
sqlx::query(
r"
INSERT INTO objects (
object_type, id, scope, version, status, payload, created_at_ms, updated_at_ms, workspace
)
SELECT $1, $2, $3, 1, $4, $5, $6, $6, $7
WHERE NOT EXISTS (SELECT 1 FROM objects WHERE object_type = $1 AND scope = $3)
ON CONFLICT DO NOTHING
",
)
.bind(POLICY_OBJECT_TYPE)
.bind(&record.id)
.bind(&record.sandbox_id)
.bind(&record.status)
.bind(wrapped_payload)
.bind(record.created_at_ms)
.bind(workspace)
.execute(&mut *tx)
.await
.map_err(|e| map_db_error(&e))?;
}
tx.commit().await.map_err(|e| map_db_error(&e))?;
Ok(())
}
pub async fn put_policy_revision_atomic(
&self,
write: &AtomicPolicyRevisionWrite,
@@ -1084,6 +1084,36 @@ VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?7, ?8)
Ok(())
}
pub async fn put_initial_policy_revision(
&self,
record: &PolicyRecord,
workspace: &str,
) -> PersistenceResult<()> {
let wrapped_payload = policy_payload_from_record(record)?;
sqlx::query(
r#"
INSERT INTO "objects" (
"object_type", "id", "scope", "version", "status", "payload", "created_at_ms", "updated_at_ms", "workspace"
)
SELECT ?1, ?2, ?3, 1, ?4, ?5, ?6, ?6, ?7
WHERE EXISTS (SELECT 1 FROM "objects" WHERE "object_type" = 'sandbox' AND "id" = ?3)
AND NOT EXISTS (SELECT 1 FROM "objects" WHERE "object_type" = ?1 AND "scope" = ?3)
ON CONFLICT DO NOTHING
"#,
)
.bind(POLICY_OBJECT_TYPE)
.bind(&record.id)
.bind(&record.sandbox_id)
.bind(&record.status)
.bind(wrapped_payload)
.bind(record.created_at_ms)
.bind(workspace)
.execute(&self.pool)
.await
.map_err(|e| map_db_error(&e))?;
Ok(())
}
pub async fn put_policy_revision_atomic(
&self,
write: &AtomicPolicyRevisionWrite,
@@ -1,7 +1,9 @@
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
use super::{ObjectListQuery, ObjectType, PersistenceError, Store, generate_name, test_store};
use super::{
ObjectListQuery, ObjectType, PersistenceError, PolicyRecord, Store, generate_name, test_store,
};
use crate::policy_store::{AtomicPolicyRevisionWrite, PolicyStoreExt};
use openshell_core::proto::datamodel::v1::ObjectMeta as ProtoObjectMeta;
use openshell_core::proto::{ObjectForTest, Sandbox, SandboxPolicy, SandboxSpec};
@@ -1091,6 +1093,85 @@ fn policy_test_sandbox(id: &str, name: &str) -> Sandbox {
}
}
#[tokio::test]
async fn initial_policy_history_is_insert_only() {
assert_initial_policy_history_is_insert_only(&test_store().await).await;
}
#[tokio::test]
#[ignore = "requires OPENSHELL_TEST_POSTGRES_URL pointing to a test database"]
async fn postgres_initial_policy_history_is_insert_only() {
let url = std::env::var("OPENSHELL_TEST_POSTGRES_URL").expect("test database URL");
let store = Store::connect(&url).await.unwrap();
assert_initial_policy_history_is_insert_only(&store).await;
}
async fn assert_initial_policy_history_is_insert_only(store: &Store) {
let id = uuid::Uuid::new_v4().to_string();
let sandbox = policy_test_sandbox(&id, &id);
let record = PolicyRecord {
id: uuid::Uuid::new_v4().to_string(),
sandbox_id: id.clone(),
version: 1,
policy_payload: SandboxPolicy::default().encode_to_vec(),
policy_hash: "initial-hash".into(),
status: "loaded".into(),
load_error: None,
created_at_ms: 1,
loaded_at_ms: None,
provenance: StdHashMap::new(),
};
store
.put_initial_policy_revision(&record, "default")
.await
.unwrap();
assert!(store.get_latest_policy(&id).await.unwrap().is_none());
store.put_message(&sandbox).await.unwrap();
let (first, second) = tokio::join!(
store.put_initial_policy_revision(&record, "default"),
store.put_initial_policy_revision(&record, "default"),
);
first.unwrap();
second.unwrap();
let initial = store.get_latest_policy(&id).await.unwrap().unwrap();
assert_eq!(initial.status, "loaded");
assert_eq!(initial.policy_hash, record.policy_hash);
assert_eq!(store.list_policies(&id, 10, 0).await.unwrap().len(), 1);
store
.update_policy_status(&id, 1, "failed", Some("apply failed"), None)
.await
.unwrap();
store
.put_initial_policy_revision(&record, "default")
.await
.unwrap();
let failed = store.get_latest_policy(&id).await.unwrap().unwrap();
assert_eq!(failed.status, "failed");
assert_eq!(failed.load_error.as_deref(), Some("apply failed"));
store
.put_policy_revision(
&uuid::Uuid::new_v4().to_string(),
&id,
"default",
2,
&record.policy_payload,
"new-hash",
)
.await
.unwrap();
store
.put_initial_policy_revision(&record, "default")
.await
.unwrap();
let latest = store.get_latest_policy(&id).await.unwrap().unwrap();
assert_eq!(latest.version, 2);
assert_eq!(store.list_policies(&id, 10, 0).await.unwrap().len(), 2);
}
#[tokio::test]
async fn policy_atomic_write_commits_revision_provenance_and_sandbox_projection() {
let store = test_store().await;
@@ -94,6 +94,14 @@ pub fn project_policy_revision_onto_sandbox(
}
pub trait PolicyStoreExt {
/// Insert version-one policy history when the sandbox still has no policy
/// revisions. Existing history and apply status are left untouched.
async fn put_initial_policy_revision(
&self,
record: &PolicyRecord,
workspace: &str,
) -> PersistenceResult<()>;
async fn put_policy_revision(
&self,
id: &str,
@@ -217,6 +225,17 @@ pub trait PolicyStoreExt {
}
impl PolicyStoreExt for Store {
async fn put_initial_policy_revision(
&self,
record: &PolicyRecord,
workspace: &str,
) -> PersistenceResult<()> {
match self {
Self::Postgres(store) => store.put_initial_policy_revision(record, workspace).await,
Self::Sqlite(store) => store.put_initial_policy_revision(record, workspace).await,
}
}
async fn put_policy_revision(
&self,
id: &str,
@@ -1809,14 +1809,25 @@ pub fn spawn_refresh_worker(state: std::sync::Arc<crate::ServerState>, interval:
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
loop {
ticker.tick().await;
if let Err(err) = run_refresh_worker_tick(
match run_refresh_worker_tick(
state.store.as_ref(),
Some(&state.credentials),
Some(&state.compute),
)
.await
{
warn!(error = %err, "provider credential refresh worker tick failed");
Ok(workspaces) => {
for workspace in workspaces {
crate::config_delivery::publish_workspace_components(
&state,
&workspace,
crate::config_delivery::ConfigComponents::ALL,
);
}
}
Err(err) => {
warn!(error = %err, "provider credential refresh worker tick failed");
}
}
}
});
@@ -1835,7 +1846,8 @@ async fn run_refresh_worker_tick(
store: &Store,
credentials: Option<&crate::credentials::CredentialRuntime>,
compute: Option<&crate::compute::ComputeRuntime>,
) -> Result<(), Status> {
) -> Result<std::collections::HashSet<String>, Status> {
let mut changed_workspaces = std::collections::HashSet::new();
let now_ms = current_time_ms();
let states = list_all_refresh_states(store).await.inspect_err(|_| {
crate::otel_tracing::mark_error(&tracing::Span::current());
@@ -1885,6 +1897,8 @@ async fn run_refresh_worker_tick(
error = %err,
"failed to finalize tombstoned provider refresh; retrying on the next sweep"
);
} else {
changed_workspaces.insert(state.object_workspace().to_string());
}
continue;
}
@@ -1968,9 +1982,11 @@ async fn run_refresh_worker_tick(
error = %err,
"provider credential refresh failed"
);
} else {
changed_workspaces.insert(state.object_workspace().to_string());
}
}
Ok(())
Ok(changed_workspaces)
}
#[cfg(test)]
+408 -26
View File
@@ -7,21 +7,29 @@ use std::sync::atomic::Ordering;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use metrics::counter;
use prost::Message;
use tokio::sync::{mpsc, oneshot};
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Request, Response, Status};
use tracing::{debug, info, warn};
use uuid::Uuid;
use openshell_core::proto::SUPERVISOR_PROTOCOL_REVISION;
use openshell_core::proto::{
GatewayMessage, RelayFrame, RelayInit, RelayOpen, ReportMainProcessExitRequest,
ConfigUpdate, GatewayMessage, RelayFrame, RelayInit, RelayOpen, ReportMainProcessExitRequest,
ReportMainProcessExitResponse, Sandbox, SandboxPhase, SessionAccepted, SshRelayTarget,
SupervisorMessage, gateway_message, relay_open, supervisor_message,
SupervisorMessage, config_update, gateway_message, relay_open, supervisor_message,
};
use openshell_core::transport_errors::is_expected_transport_close_status;
use crate::ServerState;
use crate::auth::principal::Principal;
use crate::config_delivery::{
DeliveryDisposition, MAX_SUPERVISOR_CONFIG_MESSAGE_BYTES, SupervisorConfigMessage,
};
#[cfg(test)]
use crate::config_delivery::{LocalSupervisorConfigRouter, SupervisorConfigRouter};
const HEARTBEAT_INTERVAL_SECS: u32 = 15;
const RELAY_PENDING_TIMEOUT: Duration = Duration::from_secs(10);
@@ -52,6 +60,7 @@ struct LiveSession {
/// removing a session that has since been superseded by a reconnect.
session_id: String,
tx: mpsc::Sender<GatewayMessage>,
config_sequences: ConfigSequences,
/// Fires when this session is superseded by a reconnect so the old session
/// task can exit promptly — dropping its own `tx` clone and closing the
/// outbound stream. Without this, a concurrent `open_relay` that grabbed
@@ -65,6 +74,12 @@ struct LiveSession {
connected_at: Instant,
}
#[derive(Debug, Default)]
struct ConfigSequences {
sandbox_config: u64,
provider_environment: u64,
}
/// Holds a oneshot sender that will deliver the upgraded relay stream or a
/// target-open failure reported by the supervisor.
type RelayStreamSender = oneshot::Sender<Result<tokio::io::DuplexStream, Status>>;
@@ -127,6 +142,7 @@ impl SupervisorSessionRegistry {
sandbox_id,
session_id,
tx,
config_sequences: ConfigSequences::default(),
shutdown,
terminal_delivery_finalized: false,
connected_at: Instant::now(),
@@ -233,6 +249,65 @@ impl SupervisorSessionRegistry {
true
}
pub(crate) fn connected_sandbox_ids(&self) -> Vec<String> {
self.sessions.lock().unwrap().keys().cloned().collect()
}
pub(crate) fn deliver_config(
&self,
sandbox_id: &str,
message: SupervisorConfigMessage,
) -> DeliveryDisposition {
let component_name = message.component_name();
let mut sessions = self.sessions.lock().unwrap();
let Some(session) = sessions.get_mut(sandbox_id) else {
return DeliveryDisposition::NoActiveSession;
};
let sequence = match &message {
SupervisorConfigMessage::SandboxConfig(_) => {
&mut session.config_sequences.sandbox_config
}
SupervisorConfigMessage::ProviderEnvironment(_) => {
&mut session.config_sequences.provider_environment
}
};
*sequence = sequence.saturating_add(1);
let component_sequence = *sequence;
let component = match message {
SupervisorConfigMessage::SandboxConfig(snapshot) => {
config_update::Component::SandboxConfig(*snapshot)
}
SupervisorConfigMessage::ProviderEnvironment(snapshot) => {
config_update::Component::ProviderEnvironment(snapshot)
}
};
let gateway_message = GatewayMessage {
payload: Some(gateway_message::Payload::ConfigUpdate(ConfigUpdate {
update_id: Uuid::new_v4().to_string(),
component_sequence,
component: Some(component),
})),
};
if gateway_message.encoded_len() > MAX_SUPERVISOR_CONFIG_MESSAGE_BYTES {
return DeliveryDisposition::PayloadTooLarge;
}
match session.tx.try_send(gateway_message) {
Ok(()) => DeliveryDisposition::Enqueued,
Err(mpsc::error::TrySendError::Full(_)) => {
warn!(
sandbox_id = %sandbox_id,
component = component_name,
"supervisor configuration queue is full"
);
DeliveryDisposition::QueueFull
}
Err(mpsc::error::TrySendError::Closed(_)) => DeliveryDisposition::SessionClosed,
}
}
pub fn is_current_session(&self, sandbox_id: &str, session_id: &str) -> bool {
self.sessions
.lock()
@@ -480,17 +555,13 @@ pub fn spawn_relay_reaper(state: Arc<ServerState>, interval: Duration) {
async fn require_persisted_sandbox(
store: &Arc<crate::persistence::Store>,
sandbox_id: &str,
) -> Result<(), Status> {
) -> Result<Sandbox, Status> {
let sandbox = store
.get_message::<Sandbox>(sandbox_id)
.await
.map_err(|err| Status::internal(format!("failed to load sandbox: {err}")))?;
if sandbox.is_none() {
return Err(Status::not_found("sandbox not found"));
}
Ok(())
sandbox.ok_or_else(|| Status::not_found("sandbox not found"))
}
// ---------------------------------------------------------------------------
@@ -739,10 +810,35 @@ pub async fn handle_connect_supervisor(
if sandbox_id.is_empty() {
return Err(Status::invalid_argument("sandbox_id is required"));
}
validate_protocol_revision(hello.protocol_revision)?;
if let Some(principal) = principal.as_ref() {
crate::auth::guard::ensure_sandbox_principal_scope(principal, &sandbox_id)?;
}
require_persisted_sandbox(&state.store, &sandbox_id).await?;
let sandbox = require_persisted_sandbox(&state.store, &sandbox_id).await?;
let bootstrap = match crate::config_delivery::build_config_bootstrap(state, &sandbox).await {
Ok(bootstrap) => {
counter!(
"openshell_supervisor_config_bootstrap_total",
"outcome" => "built"
)
.increment(1);
Some(bootstrap)
}
Err(error) => {
counter!(
"openshell_supervisor_config_bootstrap_total",
"outcome" => "build_failed"
)
.increment(1);
warn!(
sandbox_id = %sandbox_id,
error_code = ?error.code(),
"failed to build supervisor configuration bootstrap"
);
None
}
};
let session_id = Uuid::new_v4().to_string();
info!(
@@ -752,9 +848,35 @@ pub async fn handle_connect_supervisor(
"supervisor session: accepted"
);
// Step 2: Create and register the outbound channel.
// Step 2: Queue SessionAccepted before the session becomes routable. This
// keeps a concurrent ConfigUpdate from becoming the first stream message.
let (tx, rx) = mpsc::channel::<GatewayMessage>(64);
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
let mut accepted = GatewayMessage {
payload: Some(gateway_message::Payload::SessionAccepted(SessionAccepted {
session_id: session_id.clone(),
heartbeat_interval_secs: HEARTBEAT_INTERVAL_SECS,
bootstrap,
protocol_revision: SUPERVISOR_PROTOCOL_REVISION,
})),
};
if accepted.encoded_len() > MAX_SUPERVISOR_CONFIG_MESSAGE_BYTES {
counter!(
"openshell_supervisor_config_bootstrap_total",
"outcome" => "payload_too_large"
)
.increment(1);
let Some(gateway_message::Payload::SessionAccepted(accepted_payload)) =
accepted.payload.as_mut()
else {
unreachable!("constructed SessionAccepted payload")
};
accepted_payload.bootstrap = None;
}
if tx.send(accepted).await.is_err() {
return Err(Status::internal("failed to send session accepted"));
}
let superseded = state.supervisor_sessions.register(
sandbox_id.clone(),
session_id.clone(),
@@ -769,22 +891,6 @@ pub async fn handle_connect_supervisor(
);
}
// Step 3: Send SessionAccepted.
let accepted = GatewayMessage {
payload: Some(gateway_message::Payload::SessionAccepted(SessionAccepted {
session_id: session_id.clone(),
heartbeat_interval_secs: HEARTBEAT_INTERVAL_SECS,
})),
};
if tx.send(accepted).await.is_err() {
// Only evict ourselves — a faster reconnect may already have
// superseded this registration.
state
.supervisor_sessions
.remove_if_current(&sandbox_id, &session_id);
return Err(Status::internal("failed to send session accepted"));
}
if superseded {
state
.supervisor_sessions
@@ -854,6 +960,16 @@ pub async fn handle_connect_supervisor(
Ok(Response::new(stream))
}
fn validate_protocol_revision(supervisor_revision: u32) -> Result<(), Status> {
if supervisor_revision == SUPERVISOR_PROTOCOL_REVISION {
Ok(())
} else {
Err(Status::failed_precondition(format!(
"supervisor protocol revision mismatch: gateway requires {SUPERVISOR_PROTOCOL_REVISION}, supervisor offered {supervisor_revision}"
)))
}
}
pub async fn handle_report_main_process_exit(
state: &Arc<ServerState>,
request: Request<ReportMainProcessExitRequest>,
@@ -1016,6 +1132,22 @@ fn handle_supervisor_message(
"supervisor session: relay closed by supervisor"
);
}
Some(supervisor_message::Payload::ConfigUpdateResult(result)) => {
debug!(
sandbox_id = %sandbox_id,
session_id = %session_id,
component_sequence = result.component_sequence,
"supervisor session: ignoring configuration result while polling remains authoritative"
);
}
Some(supervisor_message::Payload::ConfigBootstrapResult(result)) => {
debug!(
sandbox_id = %sandbox_id,
session_id = %session_id,
result_count = result.results.len(),
"supervisor session: ignoring bootstrap result while polling remains authoritative"
);
}
_ => {
warn!(
sandbox_id = %sandbox_id,
@@ -1036,6 +1168,256 @@ mod tests {
use crate::auth::identity::{Identity, IdentityProvider};
use crate::auth::principal::{SandboxIdentitySource, SandboxPrincipal, UserPrincipal};
use crate::persistence::Store;
use openshell_core::proto::{
ProviderEnvironmentSnapshot, ProviderEnvironmentValue, SandboxConfigRevision,
SandboxConfigSnapshot,
};
use prost::Message;
#[test]
fn configuration_stream_messages_round_trip() {
let bootstrap = GatewayMessage {
payload: Some(gateway_message::Payload::SessionAccepted(SessionAccepted {
session_id: "session-1".into(),
heartbeat_interval_secs: 15,
bootstrap: Some(openshell_core::proto::ConfigBootstrap {
sandbox_config: Some(SandboxConfigSnapshot::default()),
provider_environment: Some(ProviderEnvironmentSnapshot::default()),
}),
protocol_revision: SUPERVISOR_PROTOCOL_REVISION,
})),
};
let updates = [
config_update::Component::SandboxConfig(SandboxConfigSnapshot::default()),
config_update::Component::ProviderEnvironment(ProviderEnvironmentSnapshot::default()),
]
.into_iter()
.enumerate()
.map(|(index, component)| GatewayMessage {
payload: Some(gateway_message::Payload::ConfigUpdate(ConfigUpdate {
update_id: format!("update-{index}"),
component_sequence: u64::try_from(index + 1).unwrap(),
component: Some(component),
})),
});
for original in std::iter::once(bootstrap).chain(updates) {
let decoded = GatewayMessage::decode(original.encode_to_vec().as_slice()).unwrap();
assert_eq!(decoded, original);
}
let result = SupervisorMessage {
payload: Some(supervisor_message::Payload::ConfigUpdateResult(
openshell_core::proto::ConfigUpdateResult {
update_id: "update-1".into(),
component_sequence: 4,
result: Some(openshell_core::proto::ConfigComponentApplyResult {
component: openshell_core::proto::ConfigComponent::SandboxConfig.into(),
requested_revision: Some(openshell_core::proto::ConfigSnapshotRevision {
component: Some(
openshell_core::proto::config_snapshot_revision::Component::SandboxConfig(
SandboxConfigRevision {
config_revision: 7,
policy_version: 3,
..Default::default()
},
),
),
}),
applied_revision: Some(openshell_core::proto::ConfigSnapshotRevision {
component: Some(
openshell_core::proto::config_snapshot_revision::Component::SandboxConfig(
SandboxConfigRevision {
config_revision: 7,
policy_version: 3,
..Default::default()
},
),
),
}),
outcome: openshell_core::proto::ConfigApplyOutcome::Applied.into(),
failure: None,
}),
},
)),
};
let decoded = SupervisorMessage::decode(result.encode_to_vec().as_slice()).unwrap();
assert_eq!(decoded, result);
}
#[test]
fn supervisor_protocol_revision_must_match_exactly() {
assert!(validate_protocol_revision(SUPERVISOR_PROTOCOL_REVISION).is_ok());
let error = validate_protocol_revision(SUPERVISOR_PROTOCOL_REVISION + 1).unwrap_err();
assert_eq!(error.code(), tonic::Code::FailedPrecondition);
assert!(error.message().contains("revision mismatch"));
}
#[tokio::test]
async fn config_router_reports_missing_session() {
let router = LocalSupervisorConfigRouter::new(Arc::new(SupervisorSessionRegistry::new()));
assert_eq!(
router
.deliver(
"missing",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::NoActiveSession
);
}
#[tokio::test]
async fn config_router_assigns_sequences_per_component() {
let registry = Arc::new(SupervisorSessionRegistry::new());
let router = LocalSupervisorConfigRouter::new(Arc::clone(&registry));
let (tx, mut rx) = mpsc::channel(4);
let (shutdown_tx, _shutdown_rx) = oneshot::channel();
registry.register("sb-1".into(), "session-1".into(), tx, shutdown_tx);
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::Enqueued
);
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::Enqueued
);
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::ProviderEnvironment(
ProviderEnvironmentSnapshot::default(),
),
)
.await,
DeliveryDisposition::Enqueued
);
let first = rx.recv().await.expect("first config update");
let second = rx.recv().await.expect("second config update");
let third = rx.recv().await.expect("provider config update");
let sequence = |message: GatewayMessage| match message.payload {
Some(gateway_message::Payload::ConfigUpdate(update)) => update.component_sequence,
other => panic!("expected config update, got {other:?}"),
};
assert_eq!(sequence(first), 1);
assert_eq!(sequence(second), 2);
assert_eq!(sequence(third), 1);
}
#[tokio::test]
async fn config_router_uses_replacement_session() {
let registry = Arc::new(SupervisorSessionRegistry::new());
let router = LocalSupervisorConfigRouter::new(Arc::clone(&registry));
let (old_tx, mut old_rx) = mpsc::channel(2);
let (old_shutdown_tx, _old_shutdown_rx) = oneshot::channel();
registry.register("sb-1".into(), "old-session".into(), old_tx, old_shutdown_tx);
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::Enqueued
);
let old_update = old_rx.recv().await.expect("old-session config update");
let Some(gateway_message::Payload::ConfigUpdate(old_update)) = old_update.payload else {
panic!("expected config update");
};
assert_eq!(old_update.component_sequence, 1);
let (new_tx, mut new_rx) = mpsc::channel(1);
let (new_shutdown_tx, _new_shutdown_rx) = oneshot::channel();
assert!(registry.register("sb-1".into(), "new-session".into(), new_tx, new_shutdown_tx,));
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::Enqueued
);
assert!(old_rx.try_recv().is_err());
let new_update = new_rx.try_recv().expect("new-session config update");
let Some(gateway_message::Payload::ConfigUpdate(new_update)) = new_update.payload else {
panic!("expected config update");
};
assert_eq!(new_update.component_sequence, 1);
}
#[tokio::test]
async fn config_router_reports_full_and_closed_queues() {
let registry = Arc::new(SupervisorSessionRegistry::new());
let router = LocalSupervisorConfigRouter::new(Arc::clone(&registry));
let (tx, rx) = mpsc::channel(1);
tx.try_send(GatewayMessage::default()).unwrap();
let (shutdown_tx, _shutdown_rx) = oneshot::channel();
registry.register("sb-1".into(), "session-1".into(), tx, shutdown_tx);
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::QueueFull
);
drop(rx);
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::SessionClosed
);
}
#[tokio::test]
async fn config_router_rejects_oversized_messages() {
let registry = Arc::new(SupervisorSessionRegistry::new());
let router = LocalSupervisorConfigRouter::new(Arc::clone(&registry));
let (tx, mut rx) = mpsc::channel(1);
let (shutdown_tx, _shutdown_rx) = oneshot::channel();
registry.register("sb-1".into(), "session-1".into(), tx, shutdown_tx);
let snapshot = ProviderEnvironmentSnapshot {
values: vec![ProviderEnvironmentValue {
name: "TOKEN".into(),
value: "x".repeat(MAX_SUPERVISOR_CONFIG_MESSAGE_BYTES),
..Default::default()
}],
..Default::default()
};
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::ProviderEnvironment(snapshot),
)
.await,
DeliveryDisposition::PayloadTooLarge
);
assert!(rx.try_recv().is_err());
}
use tokio::io::{AsyncReadExt, AsyncWriteExt};
async fn test_store() -> Arc<Store> {