use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::path::PathBuf;
use std::sync::{Arc, Mutex, MutexGuard, OnceLock};
use std::time::{Duration, Instant};
use serde::{Deserialize, Serialize};
use crate::{
find_live_runtime, forget_live_runtime, list_live_runtimes, resolve_live_runtime,
DiscoveryQuery, FrontendActions, FrontendConnectionState, FrontendRuntimeDescriptor,
FrontendTurnState, HarnessCatalog, HttpFrontendRuntime, LiveRuntimeEndpoint,
RuntimeAuthorization, RuntimeClientId, RuntimeControllerLease, RuntimeObserverLease,
RuntimePermission, SdkError, SdkOperation, Session,
};
use supercode_interchange::catalog::StorageLocator;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default)]
pub struct RuntimeRegistryQuery {
pub persisted: DiscoveryQuery,
pub include_live: bool,
pub include_persisted: bool,
}
impl Default for RuntimeRegistryQuery {
fn default() -> Self {
Self {
persisted: DiscoveryQuery::default(),
include_live: true,
include_persisted: true,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RuntimeRegistryState {
Persisted,
Idle,
Busy,
ShuttingDown,
}
impl RuntimeRegistryState {
pub fn as_str(self) -> &'static str {
match self {
Self::Persisted => "persisted",
Self::Idle => "idle",
Self::Busy => "busy",
Self::ShuttingDown => "shutting_down",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RuntimeRegistryOwner {
pub pid: u32,
pub controller: Option<RuntimeControllerLease>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RuntimeRegistryEntry {
pub id: String,
pub runtime_id: Option<String>,
pub source_session_id: String,
pub source_workspace: Option<PathBuf>,
pub source_harness: String,
pub profile: Option<String>,
pub state: RuntimeRegistryState,
pub model: Option<String>,
pub owner: Option<RuntimeRegistryOwner>,
pub observers: Vec<RuntimeObserverLease>,
pub started_at_ms: Option<u128>,
pub updated_at_ms: Option<u64>,
pub endpoint: Option<LiveRuntimeEndpoint>,
pub endpoint_capabilities: Vec<String>,
pub actions: Option<FrontendActions>,
pub persistence_location: Option<PathBuf>,
pub supervisor: Option<crate::LiveRuntimeSupervisor>,
pub title: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum RuntimeRegistryEvent {
Added {
entry: RuntimeRegistryEntry,
},
Updated {
entry: RuntimeRegistryEntry,
},
Removed {
id: String,
},
Error {
message: String,
},
}
pub struct RuntimeRegistryWatch {
receiver: tokio::sync::mpsc::Receiver<RuntimeRegistryEvent>,
task: tokio::task::JoinHandle<()>,
}
impl RuntimeRegistryWatch {
pub async fn next(&mut self) -> Option<RuntimeRegistryEvent> {
self.receiver.recv().await
}
}
impl Drop for RuntimeRegistryWatch {
fn drop(&mut self) {
self.task.abort();
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct LocalRuntimeRegistry;
impl LocalRuntimeRegistry {
pub fn new() -> Self {
Self
}
pub async fn list(
&self,
query: &RuntimeRegistryQuery,
authorization: &RuntimeAuthorization,
) -> Result<Vec<RuntimeRegistryEntry>, SdkError> {
require_permission(authorization, RuntimePermission::Observe)?;
let mut entries = BTreeMap::<String, RuntimeRegistryEntry>::new();
if query.include_persisted {
let persisted = HarnessCatalog::new()
.discover(&query.persisted)
.map_err(|error| SdkError::Execution {
operation: SdkOperation::Discover,
message: error.to_string(),
})?;
for descriptor in persisted {
let id = format!(
"{}:{}",
descriptor.locator.harness.as_str(),
descriptor.locator.session_id
);
let persistence_location = Some(match &descriptor.locator.storage {
StorageLocator::File { path } | StorageLocator::Sqlite { path, .. } => {
path.clone()
}
});
entries.insert(
id.clone(),
RuntimeRegistryEntry {
id,
runtime_id: None,
source_session_id: descriptor.locator.session_id,
source_workspace: None,
source_harness: descriptor.locator.harness.0,
profile: None,
state: RuntimeRegistryState::Persisted,
model: descriptor.model,
owner: None,
observers: Vec::new(),
started_at_ms: None,
updated_at_ms: descriptor.updated_at_ms,
endpoint: None,
endpoint_capabilities: Vec::new(),
actions: None,
persistence_location,
supervisor: None,
title: descriptor.title,
},
);
}
}
if query.include_live {
for record in list_live_runtimes().map_err(registry_receipt_error)? {
let Some((probe, descriptor)) = probe_receipt(&record).await? else {
continue;
};
let leases = probe.lease_snapshot().await?;
let entry = live_entry(record, descriptor, leases);
if entries.insert(entry.id.clone(), entry).is_some() {
return Err(SdkError::Execution {
operation: SdkOperation::Discover,
message: "duplicate stable runtime id in live registry".into(),
});
}
}
}
Ok(entries.into_values().collect())
}
pub async fn describe(
&self,
id: &str,
query: &RuntimeRegistryQuery,
authorization: &RuntimeAuthorization,
) -> Result<RuntimeRegistryEntry, SdkError> {
self.list(query, authorization)
.await?
.into_iter()
.find(|entry| entry.id == id)
.ok_or_else(|| SdkError::NotFound {
operation: SdkOperation::Discover,
message: format!("runtime or persisted session `{id}`"),
})
}
pub async fn source_state(
&self,
harness: &str,
session_id: &str,
authorization: &RuntimeAuthorization,
) -> Result<Option<RuntimeRegistryState>, SdkError> {
require_permission(authorization, RuntimePermission::Observe)?;
for record in list_live_runtimes().map_err(registry_receipt_error)? {
if record.source.harness != harness || record.source.session_id != session_id {
continue;
}
let Some((_probe, descriptor)) = probe_receipt(&record).await? else {
continue;
};
return Ok(Some(reconciled_state(&descriptor)));
}
Ok(None)
}
pub(crate) async fn source_states(
&self,
sources: &BTreeSet<(String, String)>,
authorization: &RuntimeAuthorization,
) -> Result<BTreeMap<(String, String), RuntimeRegistryState>, SdkError> {
require_permission(authorization, RuntimePermission::Observe)?;
if sources.is_empty() {
return Ok(BTreeMap::new());
}
let mut runtimes = BTreeMap::new();
for record in list_live_runtimes().map_err(registry_receipt_error)? {
let key = (
record.source.harness.clone(),
record.source.session_id.clone(),
);
if !sources.contains(&key) {
continue;
}
let Some((_probe, descriptor)) = probe_receipt(&record).await? else {
continue;
};
if runtimes
.insert(
record.runtime_session_id,
(key, reconciled_state(&descriptor)),
)
.is_some()
{
return Err(SdkError::Execution {
operation: SdkOperation::Discover,
message: "duplicate stable runtime id in live registry".into(),
});
}
}
Ok(runtimes.into_values().collect())
}
pub async fn attach(
&self,
runtime_id: &str,
client_id: RuntimeClientId,
authorization: RuntimeAuthorization,
) -> Result<Arc<HttpFrontendRuntime>, SdkError> {
require_permission(&authorization, RuntimePermission::Observe)?;
let record = find_live_runtime(runtime_id)
.map_err(registry_receipt_error)?
.ok_or_else(|| SdkError::NotFound {
operation: SdkOperation::Resume,
message: format!("live runtime `{runtime_id}`"),
})?;
let resolved = resolve_live_runtime(&record.endpoint, &record.source)
.map_err(registry_receipt_error)?;
let attached = HttpFrontendRuntime::connect_with_authorization(
resolved.base_url,
resolved.token,
client_id,
authorization,
)
.await?;
note_reachable(&record.endpoint);
Ok(attached)
}
pub fn load_persisted(
&self,
id: &str,
query: &RuntimeRegistryQuery,
authorization: &RuntimeAuthorization,
) -> Result<Session, SdkError> {
require_permission(authorization, RuntimePermission::Observe)?;
let descriptor = HarnessCatalog::new()
.discover(&query.persisted)
.map_err(|error| SdkError::Execution {
operation: SdkOperation::Discover,
message: error.to_string(),
})?
.into_iter()
.find(|descriptor| {
format!(
"{}:{}",
descriptor.locator.harness.as_str(),
descriptor.locator.session_id
) == id
})
.ok_or_else(|| SdkError::NotFound {
operation: SdkOperation::Load,
message: format!("persisted session `{id}`"),
})?;
HarnessCatalog::new()
.load(&descriptor.locator)
.map_err(|error| SdkError::Execution {
operation: SdkOperation::Load,
message: error.to_string(),
})
}
pub fn watch(
&self,
query: RuntimeRegistryQuery,
authorization: RuntimeAuthorization,
poll_interval: Duration,
) -> Result<RuntimeRegistryWatch, SdkError> {
require_permission(&authorization, RuntimePermission::Observe)?;
let (sender, receiver) = tokio::sync::mpsc::channel(128);
let registry = *self;
let interval = poll_interval.max(Duration::from_millis(25));
let task = tokio::spawn(async move {
let mut previous = BTreeMap::<String, RuntimeRegistryEntry>::new();
let mut ticker = tokio::time::interval(interval);
loop {
ticker.tick().await;
let current = match registry.list(&query, &authorization).await {
Ok(entries) => entries
.into_iter()
.map(|entry| (entry.id.clone(), entry))
.collect::<BTreeMap<_, _>>(),
Err(error) => {
if sender
.send(RuntimeRegistryEvent::Error {
message: error.to_string(),
})
.await
.is_err()
{
return;
}
continue;
}
};
for (id, entry) in ¤t {
let event = match previous.get(id) {
None => Some(RuntimeRegistryEvent::Added {
entry: entry.clone(),
}),
Some(prior) if prior != entry => Some(RuntimeRegistryEvent::Updated {
entry: entry.clone(),
}),
Some(_) => None,
};
if let Some(event) = event {
if sender.send(event).await.is_err() {
return;
}
}
}
for id in previous.keys().filter(|id| !current.contains_key(*id)) {
if sender
.send(RuntimeRegistryEvent::Removed { id: id.clone() })
.await
.is_err()
{
return;
}
}
previous = current;
}
});
Ok(RuntimeRegistryWatch { receiver, task })
}
}
const FORGET_AFTER_FAILED_PROBES: u32 = 3;
const FORGET_AFTER_UNREACHABLE_FOR: Duration = Duration::from_secs(2);
async fn probe_receipt(
record: &crate::LiveRuntimeRecord,
) -> Result<Option<(Arc<HttpFrontendRuntime>, FrontendRuntimeDescriptor)>, SdkError> {
let Ok(resolved) = resolve_live_runtime(&record.endpoint, &record.source) else {
note_unreachable(&record.endpoint);
return Ok(None);
};
let probe_id = registry_probe_id(&record.endpoint)?;
match HttpFrontendRuntime::probe_described(resolved.base_url, resolved.token, probe_id).await {
Ok(probed) => {
note_reachable(&record.endpoint);
Ok(Some(probed))
}
Err(_) => {
note_unreachable(&record.endpoint);
Ok(None)
}
}
}
struct Outage {
started: Instant,
latest: Instant,
failures: u32,
}
fn outages() -> &'static Mutex<HashMap<String, Outage>> {
static OUTAGES: OnceLock<Mutex<HashMap<String, Outage>>> = OnceLock::new();
OUTAGES.get_or_init(|| Mutex::new(HashMap::new()))
}
fn lock_outages() -> MutexGuard<'static, HashMap<String, Outage>> {
outages()
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
fn note_reachable(endpoint: &LiveRuntimeEndpoint) {
lock_outages().remove(endpoint.as_str());
}
fn note_unreachable(endpoint: &LiveRuntimeEndpoint) {
let now = Instant::now();
let corroborated = {
let mut outages = lock_outages();
outages
.retain(|_, outage| now.duration_since(outage.latest) <= FORGET_AFTER_UNREACHABLE_FOR);
let outage = outages
.entry(endpoint.as_str().to_string())
.or_insert(Outage {
started: now,
latest: now,
failures: 0,
});
outage.failures += 1;
outage.latest = now;
let corroborated = outage.failures >= FORGET_AFTER_FAILED_PROBES
&& now.duration_since(outage.started) >= FORGET_AFTER_UNREACHABLE_FOR;
if corroborated {
outages.remove(endpoint.as_str());
}
corroborated
};
if corroborated {
let _ = forget_live_runtime(endpoint);
}
}
fn reconciled_state(descriptor: &FrontendRuntimeDescriptor) -> RuntimeRegistryState {
if descriptor.connection_state == FrontendConnectionState::ShuttingDown {
RuntimeRegistryState::ShuttingDown
} else if descriptor.turn_state == FrontendTurnState::Busy {
RuntimeRegistryState::Busy
} else {
RuntimeRegistryState::Idle
}
}
fn registry_probe_id(endpoint: &LiveRuntimeEndpoint) -> Result<RuntimeClientId, SdkError> {
RuntimeClientId::parse(format!(
"registry-{}",
endpoint.as_str().rsplit('/').next().unwrap_or("probe")
))
.map_err(|error| SdkError::InvalidArgument {
operation: SdkOperation::Discover,
message: error.to_string(),
})
}
fn live_entry(
record: crate::LiveRuntimeRecord,
descriptor: FrontendRuntimeDescriptor,
leases: crate::RuntimeLeaseSnapshot,
) -> RuntimeRegistryEntry {
let state = reconciled_state(&descriptor);
RuntimeRegistryEntry {
id: record.runtime_session_id.clone(),
runtime_id: Some(record.runtime_session_id),
source_session_id: record.source.session_id,
source_workspace: Some(record.source.workspace),
source_harness: record.source.harness,
profile: descriptor
.emulation_profile
.or(record.metadata.profile.clone()),
state,
model: Some(descriptor.model),
owner: Some(RuntimeRegistryOwner {
pid: record.pid,
controller: leases.controller,
}),
observers: leases.observers,
started_at_ms: Some(record.created_at_ms),
updated_at_ms: None,
endpoint: Some(record.endpoint),
endpoint_capabilities: record.metadata.endpoint_capabilities,
actions: Some(descriptor.actions),
persistence_location: record.metadata.persistence_location,
supervisor: record.metadata.supervisor,
title: None,
}
}
fn require_permission(
authorization: &RuntimeAuthorization,
permission: RuntimePermission,
) -> Result<(), SdkError> {
if authorization.allows(permission) {
Ok(())
} else {
Err(SdkError::Unauthorized {
permission: permission.as_str().into(),
})
}
}
fn registry_receipt_error(error: crate::LiveRuntimeReceiptError) -> SdkError {
SdkError::Execution {
operation: SdkOperation::Discover,
message: error.to_string(),
}
}