use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::{Context, Result};
use tokio::sync::RwLock;
use tracing::{info, warn};
use crate::control::actor::{
DEFAULT_BROADCAST_CAPACITY, SessionActorConfig, SessionActorHandle, spawn_actor,
};
use crate::control::admission::{AdmissionDecision, AdmissionError, check_admission};
use crate::control::backend::stream_json::StreamJsonBackend;
use crate::control::backend::tmux::TmuxBackend;
use crate::control::classifier::{ActivityClassifier, ClassifierGate, OpenRouterClassifier};
use crate::control::config::ControlPlaneConfig;
use crate::control::event::BackendKind;
use crate::control::id::{ControlSessionId, SessionCounter};
pub struct RunParams {
pub project_id: String,
pub workdir: PathBuf,
pub backend: BackendKind,
pub prompt_file: Option<PathBuf>,
pub claude_cmd: Option<String>,
}
#[derive(Clone)]
pub struct SessionRegistry {
inner: Arc<RwLock<RegistryInner>>,
config: Arc<ControlPlaneConfig>,
classifier: Option<Arc<dyn ActivityClassifier>>,
}
impl std::fmt::Debug for SessionRegistry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SessionRegistry")
.field("config", &self.config)
.field("classifier", &self.classifier.is_some())
.finish()
}
}
struct RegistryInner {
actors: HashMap<ControlSessionId, SessionActorHandle>,
counter: SessionCounter,
}
impl std::fmt::Debug for RegistryInner {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RegistryInner")
.field("actor_count", &self.actors.len())
.finish()
}
}
impl SessionRegistry {
pub fn new() -> Self {
Self::with_config(ControlPlaneConfig::default())
}
pub fn with_config(config: ControlPlaneConfig) -> Self {
Self {
inner: Arc::new(RwLock::new(RegistryInner {
actors: HashMap::new(),
counter: SessionCounter::default(),
})),
config: Arc::new(config),
classifier: None,
}
}
pub fn with_classifier(
mut self,
provider: Arc<dyn crate::core::sm::providers::LlmProvider>,
model: impl Into<String>,
) -> Self {
self.classifier = Some(Arc::new(OpenRouterClassifier::new(provider, model)));
self
}
fn actor_config(&self) -> SessionActorConfig {
let gate = ClassifierGate::from_config(
&self.config.observability,
self.classifier.is_some(),
false,
);
let classifier = if gate.should_classify() {
self.classifier.clone()
} else {
None
};
SessionActorConfig {
auth_timeout: self.config.auth_timeout(),
classifier,
}
}
pub fn config(&self) -> &ControlPlaneConfig {
&self.config
}
pub async fn live_count(&self) -> usize {
self.inner.read().await.actors.len()
}
pub async fn register(&self, id: ControlSessionId, handle: SessionActorHandle) {
let mut inner = self.inner.write().await;
inner.actors.insert(id, handle);
}
pub async fn deregister(&self, id: &ControlSessionId) {
let mut inner = self.inner.write().await;
if inner.actors.remove(id).is_none() {
warn!(%id, "deregister: session not found in registry");
}
}
pub async fn get(&self, id: &ControlSessionId) -> Option<SessionActorHandle> {
let inner = self.inner.read().await;
inner.actors.get(id).cloned()
}
pub async fn list_ids(&self) -> Vec<ControlSessionId> {
let inner = self.inner.read().await;
inner.actors.keys().cloned().collect()
}
pub async fn admission(&self) -> AdmissionDecision {
let live = self.live_count().await as u32;
check_admission(
live,
self.config.max_concurrent_sessions,
self.config.launch_stagger_ms,
)
}
pub async fn run_session(&self, params: RunParams) -> Result<ControlSessionId> {
{
let inner = self.inner.read().await;
let live = inner.actors.len() as u32;
if let AdmissionDecision::Reject { live, cap } =
check_admission(live, self.config.max_concurrent_sessions, 0)
{
warn!(
live,
cap, "session launch rejected: concurrency cap reached (§9.2)"
);
return Err(AdmissionError::CapReached { live, cap }.into());
}
}
let stagger = self.config.stagger_duration();
if !stagger.is_zero() {
tokio::time::sleep(stagger).await;
}
let mut inner = self.inner.write().await;
let live = inner.actors.len() as u32;
if let AdmissionDecision::Reject { live, cap } =
check_admission(live, self.config.max_concurrent_sessions, 0)
{
warn!(
live,
cap,
"session launch rejected at write-lock recheck (concurrent launch raced, §9.2)"
);
return Err(AdmissionError::CapReached { live, cap }.into());
}
let session_id = inner.counter.next(¶ms.project_id);
info!(
session_id = %session_id,
project_id = %params.project_id,
backend = ?params.backend,
workdir = %params.workdir.display(),
"spawning session"
);
let handle = match params.backend {
BackendKind::StreamJson => {
let backend = StreamJsonBackend::spawn(
session_id.clone(),
params.workdir,
params.prompt_file,
)
.with_context(|| format!("stream-json spawn failed for session {session_id}"))?;
spawn_actor(
session_id.clone(),
params.project_id,
BackendKind::StreamJson,
backend,
DEFAULT_BROADCAST_CAPACITY,
self.actor_config(),
)
}
BackendKind::Tmux => {
let claude_cmd = params.claude_cmd.unwrap_or_else(|| "claude".into());
let backend = TmuxBackend::new(
session_id.clone(),
params.workdir,
params.prompt_file,
claude_cmd,
100, )
.with_context(|| format!("tmux backend spawn failed for session {session_id}"))?;
spawn_actor(
session_id.clone(),
params.project_id,
BackendKind::Tmux,
backend,
DEFAULT_BROADCAST_CAPACITY,
self.actor_config(),
)
}
};
inner.actors.insert(session_id.clone(), handle);
Ok(session_id)
}
}
impl Default for SessionRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::control::actor::SessionActorHandle;
use crate::control::event::BackendKind;
use crate::control::id::ControlSessionId;
use crate::control::state::SessionMetadata;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use tokio::sync::{RwLock, broadcast, mpsc};
fn make_handle(id: &ControlSessionId) -> SessionActorHandle {
let (command_tx, _) = mpsc::channel(1);
let (event_tx, _) = broadcast::channel(4);
SessionActorHandle {
command_tx,
event_tx,
write_lock_held: Arc::new(AtomicBool::new(false)),
metadata: Arc::new(RwLock::new(SessionMetadata::new(
id.clone(),
"proj".into(),
BackendKind::StreamJson,
))),
}
}
#[tokio::test]
async fn registry_register_deregister() {
let registry = SessionRegistry::new();
let id = ControlSessionId::new("proj", 0);
let handle = make_handle(&id);
registry.register(id.clone(), handle).await;
assert!(registry.get(&id).await.is_some());
registry.deregister(&id).await;
assert!(registry.get(&id).await.is_none());
}
#[tokio::test]
async fn registry_list() {
let registry = SessionRegistry::new();
let id0 = ControlSessionId::new("proj", 0);
let id1 = ControlSessionId::new("proj", 1);
registry.register(id0.clone(), make_handle(&id0)).await;
registry.register(id1.clone(), make_handle(&id1)).await;
let ids = registry.list_ids().await;
assert_eq!(ids.len(), 2);
assert!(ids.contains(&id0));
assert!(ids.contains(&id1));
}
#[tokio::test]
async fn registry_monotonic_ids() {
let registry = SessionRegistry::new();
let id0 = {
let mut inner = registry.inner.write().await;
inner.counter.next("p")
};
let id1 = {
let mut inner = registry.inner.write().await;
inner.counter.next("p")
};
assert_eq!(id0.as_str(), "p-0");
assert_eq!(id1.as_str(), "p-1");
}
#[tokio::test]
async fn registry_run_session_id_visible_immediately() {
let registry = SessionRegistry::new();
let id = ControlSessionId::new("p", 0);
let handle = make_handle(&id);
{
let mut inner = registry.inner.write().await;
inner.actors.insert(id.clone(), handle);
}
assert!(
registry.get(&id).await.is_some(),
"session must be visible immediately after registration"
);
}
#[tokio::test]
async fn registry_actor_config_default_no_classifier() {
let registry = SessionRegistry::new();
let cfg = registry.actor_config();
assert_eq!(cfg.auth_timeout, std::time::Duration::from_secs(30));
assert!(
cfg.classifier.is_none(),
"default registry must attach NO classifier (AC-9)"
);
}
#[tokio::test]
async fn registry_actor_config_threads_auth_timeout() {
let control = ControlPlaneConfig {
auth_timeout_secs: 7,
..Default::default()
};
let registry = SessionRegistry::with_config(control);
assert_eq!(
registry.actor_config().auth_timeout,
std::time::Duration::from_secs(7)
);
}
#[tokio::test]
async fn registry_actor_config_gate_off_drops_classifier() {
use crate::core::sm::providers::{LlmProvider, LlmRequest, LlmResponse, error::SmLlmError};
struct NoopProvider;
#[async_trait::async_trait]
impl LlmProvider for NoopProvider {
fn name(&self) -> &str {
"noop"
}
async fn complete(&self, _req: LlmRequest) -> Result<LlmResponse, SmLlmError> {
Err(SmLlmError::Degraded("test".into()))
}
}
let registry =
SessionRegistry::new().with_classifier(Arc::new(NoopProvider), "openai/gpt-4o-mini");
assert!(
registry.actor_config().classifier.is_none(),
"gate off (config flag false) must drop the classifier even when a provider is wired"
);
}
#[tokio::test]
async fn registry_actor_config_gate_on_attaches_classifier() {
use crate::control::config::ObservabilityConfig;
use crate::core::sm::providers::{LlmProvider, LlmRequest, LlmResponse, error::SmLlmError};
struct NoopProvider;
#[async_trait::async_trait]
impl LlmProvider for NoopProvider {
fn name(&self) -> &str {
"noop"
}
async fn complete(&self, _req: LlmRequest) -> Result<LlmResponse, SmLlmError> {
Err(SmLlmError::Degraded("test".into()))
}
}
let control = ControlPlaneConfig {
observability: ObservabilityConfig {
llm_classifier: true,
},
..Default::default()
};
let registry = SessionRegistry::with_config(control)
.with_classifier(Arc::new(NoopProvider), "openai/gpt-4o-mini");
assert!(
registry.actor_config().classifier.is_some(),
"gate on (flag true + provider wired + no SM) must attach the classifier"
);
}
#[tokio::test]
async fn registry_live_count_and_config() {
let registry = SessionRegistry::new();
assert_eq!(registry.live_count().await, 0);
assert_eq!(registry.config().max_concurrent_sessions, 5);
assert!(!registry.config().observability.llm_classifier);
let id = ControlSessionId::new("p", 0);
registry.register(id.clone(), make_handle(&id)).await;
assert_eq!(registry.live_count().await, 1);
}
#[tokio::test]
async fn registry_admission_check_under_cap_admits() {
let cfg = ControlPlaneConfig {
max_concurrent_sessions: 2,
launch_stagger_ms: 0,
..Default::default()
};
let registry = SessionRegistry::with_config(cfg);
assert!(matches!(
registry.admission().await,
AdmissionDecision::Admit { stagger_ms: 0 }
));
for n in 0..2 {
let id = ControlSessionId::new("p", n);
registry.register(id.clone(), make_handle(&id)).await;
}
assert!(matches!(
registry.admission().await,
AdmissionDecision::Reject { live: 2, cap: 2 }
));
}
#[tokio::test]
async fn registry_run_session_rejects_when_capped() {
let cfg = ControlPlaneConfig {
max_concurrent_sessions: 1,
launch_stagger_ms: 0,
..Default::default()
};
let registry = SessionRegistry::with_config(cfg);
let id = ControlSessionId::new("p", 0);
registry.register(id.clone(), make_handle(&id)).await;
let params = RunParams {
project_id: "p".into(),
workdir: PathBuf::from("/tmp"),
backend: BackendKind::StreamJson,
prompt_file: None,
claude_cmd: None,
};
let err = registry
.run_session(params)
.await
.expect_err("run_session must reject when at cap");
let admission_err = err
.downcast_ref::<crate::control::admission::AdmissionError>()
.expect("error must be an AdmissionError");
assert_eq!(
*admission_err,
crate::control::admission::AdmissionError::CapReached { live: 1, cap: 1 }
);
}
#[tokio::test]
async fn registry_get_clones_under_read_lock() {
let registry = SessionRegistry::new();
let id = ControlSessionId::new("proj", 0);
let handle = make_handle(&id);
registry.register(id.clone(), handle).await;
let h1 = registry.get(&id).await.expect("should exist");
let h2 = registry.get(&id).await.expect("should exist");
assert!(h1.try_acquire_write_lock());
assert!(!h2.try_acquire_write_lock()); h1.release_write_lock();
assert!(h2.try_acquire_write_lock()); }
#[tokio::test]
async fn registry_cap_reject_is_immediate() {
use std::time::{Duration, Instant};
let cfg = ControlPlaneConfig {
max_concurrent_sessions: 1,
launch_stagger_ms: 5_000,
..Default::default()
};
let registry = SessionRegistry::with_config(cfg);
let id = ControlSessionId::new("p", 0);
registry.register(id.clone(), make_handle(&id)).await;
let t0 = Instant::now();
let params = RunParams {
project_id: "p".into(),
workdir: PathBuf::from("/tmp"),
backend: BackendKind::StreamJson,
prompt_file: None,
claude_cmd: None,
};
let err = registry
.run_session(params)
.await
.expect_err("must reject at cap");
let elapsed = t0.elapsed();
let admission_err = err
.downcast_ref::<crate::control::admission::AdmissionError>()
.expect("error must be AdmissionError");
assert_eq!(
*admission_err,
crate::control::admission::AdmissionError::CapReached { live: 1, cap: 1 }
);
assert!(
elapsed < Duration::from_secs(1),
"rejected launch must return in <1 s, not {elapsed:?} \
(5 s stagger must not be paid on reject path)"
);
}
}