mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-27 07:03:17 +08:00
feat(supervisor): stage gateway configuration snapshot delivery
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
This commit is contained in:
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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(®istry));
|
||||
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(®istry));
|
||||
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(®istry));
|
||||
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(®istry));
|
||||
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> {
|
||||
|
||||
Reference in New Issue
Block a user