use std::collections::HashMap;
use std::convert::Infallible;
use std::fs;
use std::path::{Path as FsPath, PathBuf};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use axum::extract::{DefaultBodyLimit, Path, State};
use axum::http::{header, HeaderMap, HeaderValue, StatusCode};
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::response::{IntoResponse, Response};
use axum::routing::{delete, get, post, put};
use axum::{Json, Router};
use axum::middleware::{self, Next};
use candle_core::{DType, Device, Tensor};
use futures::stream::{self, Stream};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use tokio::task::{spawn_blocking, JoinSet};
use tower_http::cors::CorsLayer;
use tower_http::decompression::RequestDecompressionLayer;
use uuid::Uuid;
use crate::anthropic_forwarder::{
build_compressed_payload, partition_messages, whitespace_token_count, AnthropicForwarder,
ClientAuth, ForwarderError,
};
use crate::backend_router::{Router as BackendRouter, TaskKind};
use crate::claude_backend::{ChatTurn, ClaudeBackend};
use crate::cluster::StateDeltaUpdate;
use crate::config::AxiomConfig;
use crate::cvm_store::CvmStore;
use crate::context_compressor::{
adapt_session_blocking, extract_memory_vector_blocking, feedback_adaptation_text,
should_retry_uncompressed, CompressionControls, CompressorConfig, MemoryFingerprint,
SessionStates, TttSessionStore,
};
use crate::digest::Digestor;
use crate::dwe::{
extract_delta_fragment, record_applied_fragment, record_rejected_fragment, DweBus, DweTelemetry,
};
use crate::hamiltonian::QuantumRuntimeStatus;
use crate::inference::InferencePipeline;
use crate::metrics;
use crate::openai_forwarder::{OpenAiClientAuth, OpenAiForwarder, OpenAiForwarderError};
use crate::poly_jit::{PolyJitEngine, PolyJitReport, PolyJitRunRequest, PolyJitStatus};
use crate::quantization::{NF4QuantizedDescriptor, NF4Quantizer};
use crate::responses_compressor::{apply_plan, plan_compression};
use crate::sandbox::{SandboxController, SandboxDiagnostic};
use crate::session_awareness::AwarenessStore;
use crate::surprisal::{ExactAttentionResidualCache, ExactResidualTelemetry};
use crate::swarm_route::{LocalSwarmRouteMatrix, SwarmMatrixState};
use crate::swarm_router::{SwarmChatResult, SwarmRouter};
use crate::vfs::{NeuralVfs, VfsMountReport, VfsReadReport, VfsStats};
use crate::vibe_memory::MasterVibe;
use crate::weight_merge::{
fleet_dare_ties, merge_checkpoint_files_with, MergeMethod, MergeSummary,
};
const MAX_ACTIVE_VRAM_SESSIONS: usize = 32;
const MAX_RESPONSES_RUN_CONCURRENCY: usize = 4;
fn unix_now() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
fn context_hash(source: &str) -> String {
let digest = Sha256::digest(source.as_bytes());
format!("{digest:x}")
}
fn compression_cache_path() -> PathBuf {
std::env::var("AXIOM_COMPRESSION_CACHE")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from("checkpoints/memory/axiom_compression_cache.bin"))
}
enum SessionResidency {
Active(Vec<Tensor>),
Quantized(Vec<NF4QuantizedDescriptor>),
}
#[derive(Clone, Copy)]
struct SequenceState {
version: u64,
timestamp: i64,
}
struct SessionData {
residency: SessionResidency,
created_at: u64,
last_used: u64,
}
impl SessionData {
fn new_active(states: Vec<Tensor>, created_at: u64) -> Self {
Self {
residency: SessionResidency::Active(states),
created_at,
last_used: created_at,
}
}
fn replace_states(&mut self, states: Vec<Tensor>) {
self.residency = SessionResidency::Active(states);
}
fn is_quantized(&self) -> bool {
matches!(self.residency, SessionResidency::Quantized(_))
}
fn states_clone(&self) -> candle_core::Result<Vec<Tensor>> {
match &self.residency {
SessionResidency::Active(states) => Ok(states.clone()),
SessionResidency::Quantized(_) => {
candle_core::bail!("session is parked in compressed form")
}
}
}
fn ensure_active(&mut self, device: &Device) -> candle_core::Result<Vec<Tensor>> {
if let SessionResidency::Quantized(descriptors) = &self.residency {
let mut states = Vec::with_capacity(descriptors.len());
for (descriptor_idx, descriptor) in descriptors.iter().enumerate() {
if descriptor.packed_width == 0 {
candle_core::bail!("descriptor {descriptor_idx}: packed width must be non-zero")
}
if descriptor.scales.is_empty() {
candle_core::bail!("descriptor {descriptor_idx}: scale list cannot be empty")
}
let num_blocks = descriptor.scales.len();
let packed = Tensor::from_vec(
descriptor.packed_indices.clone(),
(num_blocks, descriptor.packed_width),
device,
)?;
let scales = Tensor::from_vec(descriptor.scales.clone(), (num_blocks,), device)?;
let mut data = NF4Quantizer::dequantize_state(&packed, &scales)?
.to_dtype(DType::F32)?
.contiguous()?
.flatten_all()?
.to_vec1::<f32>()?;
let total: usize = descriptor.shape.iter().product();
if data.len() < total {
candle_core::bail!(
"dequantized data too small: expected at least {}, got {}",
total,
data.len()
);
}
data.truncate(total);
let tensor = Tensor::from_vec(data, (total,), device)?
.reshape(descriptor.shape.as_slice())?;
states.push(tensor);
}
self.residency = SessionResidency::Active(states);
}
self.states_clone()
}
fn merge_delta(
&mut self,
layer_index: usize,
delta: &Tensor,
device: &Device,
) -> candle_core::Result<()> {
let mut states = self.ensure_active(device)?;
if layer_index >= states.len() {
candle_core::bail!(
"layer index {layer_index} out of range for {} session layers",
states.len()
);
}
let previous = states[layer_index].clone();
let merged = states[layer_index].add(&delta.to_dtype(DType::F32)?)?;
if !tensor_is_finite(&merged)? {
eprintln!(
"[emergency] non-finite tensor detected in merge_delta; session update discarded"
);
states[layer_index] = previous;
candle_core::bail!("delta merge produced non-finite values");
}
states[layer_index] = merged;
self.replace_states(states);
Ok(())
}
}
#[derive(Clone)]
pub struct AppState {
pub pipeline: Arc<RwLock<InferencePipeline>>,
pub device: Device,
sessions: Arc<RwLock<HashMap<String, SessionData>>>,
sequence_versions: Arc<RwLock<HashMap<String, SequenceState>>>,
pub model_id: String,
pub claude_backend: Arc<Option<ClaudeBackend>>,
pub router: Arc<Option<BackendRouter>>,
pub mcp: Arc<Option<crate::mcp_stdio::McpContext>>,
pub mcp_token: Arc<Option<String>>,
pub api_key: Arc<Option<String>>,
pub jit_exec_enabled: bool,
pub ttt_sessions: Arc<TttSessionStore>,
pub anthropic_forwarder: Arc<Option<AnthropicForwarder>>,
pub openai_forwarder: Arc<Option<OpenAiForwarder>>,
pub swarm_router: Arc<Option<SwarmRouter>>,
pub compressor_config: Arc<CompressorConfig>,
pub master_vibe: Arc<Mutex<Option<MasterVibe>>>,
pub source_store: Arc<RwLock<HashMap<String, String>>>,
adapted_context_hashes: Arc<RwLock<HashMap<String, String>>>,
cache_safety_memo: Arc<RwLock<HashMap<String, String>>>,
prefix_diet_last: Arc<RwLock<HashMap<String, crate::prefix_diet::DietReport>>>,
digest_turn: Arc<RwLock<HashMap<String, u64>>>,
digest_page_turn: Arc<RwLock<HashMap<(String, String), u64>>>,
pss_prefix_hash: Arc<RwLock<HashMap<String, (usize, String)>>>,
pss_gap: Arc<RwLock<HashMap<String, (u64, u32)>>>,
pss_local_last: Arc<RwLock<HashMap<String, (String, u64)>>>,
pss_route_cooldown: Arc<RwLock<HashMap<String, u32>>>,
pub controls: Arc<CompressionControls>,
pub neural_vfs: Arc<NeuralVfs>,
pub savings: Arc<Mutex<HashMap<String, (u64, u64)>>>,
pub poly_jit: Arc<PolyJitEngine>,
pub exact_residual_cache: Arc<ExactAttentionResidualCache>,
pub dwe_bus: Arc<DweBus>,
pub swarm_matrix: Arc<LocalSwarmRouteMatrix>,
pub heal_memory_path: Arc<Option<std::path::PathBuf>>,
pub awareness: AwarenessStore,
pub cvm_store: Arc<CvmStore>,
pub keepalive: crate::keepalive::KeepaliveManager,
}
impl AppState {
pub fn new(pipeline: InferencePipeline, model_id: String) -> Self {
let device = pipeline.device().clone();
Self {
pipeline: Arc::new(RwLock::new(pipeline)),
device,
sessions: Arc::new(RwLock::new(HashMap::new())),
sequence_versions: Arc::new(RwLock::new(HashMap::new())),
model_id,
claude_backend: Arc::new(None),
router: Arc::new(None),
mcp: Arc::new(None),
mcp_token: Arc::new(None),
api_key: Arc::new(None),
jit_exec_enabled: false,
ttt_sessions: Arc::new(TttSessionStore::new()),
anthropic_forwarder: Arc::new(None),
openai_forwarder: Arc::new(None),
swarm_router: Arc::new(None),
compressor_config: Arc::new(CompressorConfig::default()),
master_vibe: Arc::new(Mutex::new(None)),
source_store: Arc::new(RwLock::new(HashMap::new())),
adapted_context_hashes: Arc::new(RwLock::new(HashMap::new())),
cache_safety_memo: Arc::new(RwLock::new(HashMap::new())),
prefix_diet_last: Arc::new(RwLock::new(HashMap::new())),
digest_turn: Arc::new(RwLock::new(HashMap::new())),
digest_page_turn: Arc::new(RwLock::new(HashMap::new())),
pss_prefix_hash: Arc::new(RwLock::new(HashMap::new())),
pss_gap: Arc::new(RwLock::new(HashMap::new())),
pss_local_last: Arc::new(RwLock::new(HashMap::new())),
pss_route_cooldown: Arc::new(RwLock::new(HashMap::new())),
controls: Arc::new(CompressionControls::from_config(
&CompressorConfig::default(),
)),
neural_vfs: Arc::new(NeuralVfs::new()),
savings: Arc::new(Mutex::new(HashMap::new())),
poly_jit: Arc::new(PolyJitEngine::default()),
exact_residual_cache: Arc::new(ExactAttentionResidualCache::default()),
dwe_bus: Arc::new(DweBus::from_env()),
swarm_matrix: Arc::new(LocalSwarmRouteMatrix::new()),
heal_memory_path: Arc::new(None),
awareness: AwarenessStore::new(),
cvm_store: Arc::new(Self::open_cvm_store()),
keepalive: crate::keepalive::KeepaliveManager::disabled(),
}
}
fn open_cvm_store() -> CvmStore {
let root =
std::env::var("AXIOM_CVM_DIR").unwrap_or_else(|_| "checkpoints/cvm".to_string());
CvmStore::open(&root).unwrap_or_else(|e| {
eprintln!(
"[axiom-cvm] failed to open CVM store at {root}: {e} -- falling back to a scratch temp dir"
);
let fallback = std::env::temp_dir().join("axiom-cvm-fallback");
CvmStore::open(fallback).expect("temp dir CVM store open must succeed")
})
}
pub fn with_keepalive_manager(mut self, manager: crate::keepalive::KeepaliveManager) -> Self {
self.keepalive = manager;
self
}
pub fn with_heal_memory_path(mut self, path: Option<std::path::PathBuf>) -> Self {
self.heal_memory_path = Arc::new(path);
self
}
pub fn with_cvm_store(mut self, store: CvmStore) -> Self {
self.cvm_store = Arc::new(store);
self
}
pub fn with_dwe_bus(mut self, bus: DweBus) -> Self {
self.dwe_bus = Arc::new(bus);
self
}
pub fn store_source(&self, session_id: &str, source: String) {
if let Ok(mut map) = self.source_store.write() {
if map.len() >= 256 {
map.clear();
}
map.insert(session_id.to_string(), source);
}
}
fn should_adapt_heavy_context(&self, session_id: &str, source: &str) -> bool {
let hash = context_hash(source);
self.adapted_context_hashes
.read()
.map(|map| map.get(session_id) != Some(&hash))
.unwrap_or(true)
}
fn mark_heavy_context_adapted(&self, session_id: &str, source: &str) {
if let Ok(mut map) = self.adapted_context_hashes.write() {
if map.len() >= 256 {
map.clear();
}
map.insert(session_id.to_string(), context_hash(source));
}
}
pub(crate) fn cache_safety_memo_get(&self, key: &str) -> Option<String> {
self.cache_safety_memo
.read()
.ok()
.and_then(|map| map.get(key).cloned())
}
pub(crate) fn cache_safety_memo_set(&self, key: String, messages_json: String) {
if let Ok(mut map) = self.cache_safety_memo.write() {
if map.len() >= 512 {
map.clear();
}
map.insert(key, messages_json);
}
}
pub(crate) fn prefix_diet_last_get(
&self,
session_id: &str,
) -> Option<crate::prefix_diet::DietReport> {
self.prefix_diet_last
.read()
.ok()
.and_then(|map| map.get(session_id).copied())
}
pub(crate) fn prefix_diet_last_set(
&self,
session_id: String,
report: crate::prefix_diet::DietReport,
) {
if let Ok(mut map) = self.prefix_diet_last.write() {
if map.len() >= 512 {
map.clear();
}
map.insert(session_id, report);
}
}
pub(crate) fn digest_turn_next(&self, session_id: &str) -> u64 {
let Ok(mut map) = self.digest_turn.write() else {
return 0;
};
let entry = map.entry(session_id.to_string()).or_insert(0);
*entry += 1;
*entry
}
pub(crate) fn digest_page_turn_set(&self, session_id: &str, page_id: &str, turn: u64) {
if let Ok(mut map) = self.digest_page_turn.write() {
if map.len() >= 4096 {
map.clear();
}
map.insert((session_id.to_string(), page_id.to_string()), turn);
}
}
pub(crate) fn digest_turns_since(&self, session_id: &str, page_id: &str) -> u64 {
let created_turn = self
.digest_page_turn
.read()
.ok()
.and_then(|m| m.get(&(session_id.to_string(), page_id.to_string())).copied())
.unwrap_or(0);
let current_turn = self
.digest_turn
.read()
.ok()
.and_then(|m| m.get(session_id).copied())
.unwrap_or(0);
current_turn.saturating_sub(created_turn)
}
pub(crate) fn pss_detect_break(&self, session_id: &str, frozen: &[Value]) -> bool {
let Ok(mut map) = self.pss_prefix_hash.write() else {
return false;
};
if map.len() >= 512 {
map.clear();
}
let fingerprint = crate::rebase::frozen_fingerprint(frozen);
let prev = map.insert(session_id.to_string(), fingerprint);
match prev {
Some((prev_len, prev_hash)) => {
crate::rebase::is_genuine_break(prev_len, &prev_hash, frozen)
}
None => false,
}
}
pub(crate) fn pss_note_prefix(&self, session_id: &str, frozen: &[Value]) {
if let Ok(mut map) = self.pss_prefix_hash.write() {
map.insert(
session_id.to_string(),
crate::rebase::frozen_fingerprint(frozen),
);
}
}
pub(crate) fn pss_gap_tick(
&self,
session_id: &str,
now_unix: u64,
gap_threshold_secs: u64,
) -> u32 {
let Ok(mut map) = self.pss_gap.write() else {
return 0;
};
if map.len() >= 512 {
map.clear();
}
let entry = map.entry(session_id.to_string()).or_insert((now_unix, 0));
let (last_ts, count) = *entry;
let new_count =
crate::rebase::next_long_gap_count(last_ts, now_unix, count, gap_threshold_secs);
*entry = (now_unix, new_count);
new_count
}
pub(crate) fn pss_local_admit(
&self,
session_id: &str,
inbound_hash: &str,
now: u64,
cooldown_secs: u64,
) -> bool {
let Ok(mut map) = self.pss_local_last.write() else {
return false;
};
let prior = map.get(session_id).cloned();
if let Some((h, ts)) = prior {
if h == inbound_hash && now.saturating_sub(ts) < cooldown_secs {
map.insert(session_id.to_string(), (inbound_hash.to_string(), now));
return false;
}
}
if map.len() >= 512 {
map.clear();
}
map.insert(session_id.to_string(), (inbound_hash.to_string(), now));
true
}
pub(crate) fn pss_route_cooldown_tick(&self, session_id: &str, had_error: bool) -> u32 {
let Ok(mut map) = self.pss_route_cooldown.write() else {
return 0;
};
if map.len() >= 512 {
map.clear();
}
let entry = map.entry(session_id.to_string()).or_insert(0);
if had_error {
*entry = 3;
} else if *entry > 0 {
*entry -= 1;
}
*entry
}
async fn adapt_feedback_to_cache(
&self,
req: TttFeedbackRequest,
cache_path: &FsPath,
) -> Result<TttFeedbackResponse, ApiError> {
let session_id = req.session_id.trim().to_string();
let message = req.message.trim().to_string();
if session_id.is_empty() || message.is_empty() {
return Err(ApiError::BadRequest(
"session_id and message are required".into(),
));
}
let kind = req.feedback_type.as_deref().unwrap_or("execution_feedback");
let feedback_text = feedback_adaptation_text(kind, &message, req.trace.as_deref());
let pipeline_arc = self.pipeline.clone();
let store = self.ttt_sessions.clone();
let session_id_for_task = session_id.clone();
let feedback_for_task = feedback_text.clone();
let dwe_sequence = unix_now();
let adapt_result = spawn_blocking(move || {
let pipeline = pipeline_arc
.read()
.map_err(|_| ApiError::Internal("pipeline lock poisoned".into()))?;
let handle = store
.get_or_create(&session_id_for_task, &pipeline)
.map_err(|e| ApiError::Internal(format!("session allocation failed: {e}")))?;
let mut states = handle.blocking_lock();
let baseline = states.clone();
let tokens = pipeline.encode_text(&feedback_for_task);
adapt_session_blocking(&pipeline, &mut states, &tokens)
.map_err(|e| ApiError::Internal(format!("feedback TTT adapt failed: {e}")))?;
let fragment =
extract_delta_fragment(&session_id_for_task, dwe_sequence, &states, &baseline).ok();
Ok::<_, ApiError>((tokens.len(), fragment))
})
.await
.map_err(|e| ApiError::Internal(format!("blocking task join failed: {e}")))??;
let (token_count, fragment) = adapt_result;
if let Some(fragment) = fragment {
self.dwe_bus.broadcast(fragment);
}
self.mark_heavy_context_adapted(&session_id, &feedback_text);
self.persist_compression_cache_to(cache_path)
.await
.map_err(ApiError::Internal)?;
Ok(TttFeedbackResponse {
session_id,
feedback_tokens: token_count,
persisted: true,
cache_path: cache_path.display().to_string(),
})
}
async fn sandbox_local_synthesis(&self, session_id: &str, content: &str) {
if SandboxController::rust_code_blocks(content).is_empty() {
return;
}
let Some(sandbox) = SandboxController::from_env() else {
return;
};
let state = self.clone();
let cache_path = compression_cache_path();
let result = sandbox
.verify_rust_code_blocks_with_feedback(session_id, content, move |diag| {
let state = state.clone();
let cache_path = cache_path.clone();
async move { state.adapt_sandbox_diagnostic(diag, &cache_path).await }
})
.await;
match result {
Ok(report) if report.passed => eprintln!(
"[sandbox] session={} rust_blocks={} passed attempts={}",
report.session_id, report.blocks_checked, report.attempts
),
Ok(report) => eprintln!(
"[sandbox] session={} rust_blocks={} failed diagnostics={} attempts={}",
report.session_id,
report.blocks_checked,
report.diagnostics.len(),
report.attempts
),
Err(e) => eprintln!("[sandbox] verification skipped: {e}"),
}
}
async fn adapt_sandbox_diagnostic(
&self,
diag: SandboxDiagnostic,
cache_path: &FsPath,
) -> Result<(), String> {
let message = format!(
"sandbox compiler check failed at step {} with status {:?}",
diag.step, diag.status_code
);
let trace = diag.feedback_trace();
self.adapt_feedback_to_cache(
TttFeedbackRequest {
session_id: diag.session_id,
message,
feedback_type: Some("sandbox_compilation_error".to_string()),
trace: Some(trace),
},
cache_path,
)
.await
.map(|_| ())
.map_err(|e| format!("{e:?}"))
}
async fn persist_compression_cache(&self) -> Result<(), String> {
self.persist_compression_cache_to(&compression_cache_path())
.await
}
async fn hydrate_compression_cache(&self) -> Result<usize, String> {
self.hydrate_compression_cache_from(&compression_cache_path())
.await
}
async fn persist_compression_cache_to(&self, path: &FsPath) -> Result<(), String> {
let hashes = self
.adapted_context_hashes
.read()
.map_err(|_| "adapted context hash lock poisoned".to_string())?
.clone();
if hashes.is_empty() {
return Ok(());
}
let mut entries = Vec::new();
for (session_id, handle) in self.ttt_sessions.snapshot_handles() {
let Some(context_hash) = hashes.get(&session_id).cloned() else {
continue;
};
let states = handle.lock().await;
let layers = states
.iter()
.map(tensor_to_layer_checkpoint)
.collect::<candle_core::Result<Vec<_>>>()
.map_err(|e| format!("checkpoint encode failed: {e}"))?;
entries.push(PersistedCompressionEntry {
session_id: session_id.clone(),
context_hash,
checkpoint: SessionCheckpoint {
session_id,
version: 1,
created_at: unix_now(),
layers,
},
});
}
let payload = PersistedCompressionCache {
version: 1,
entries,
};
let bytes = crate::bincode_compat::serialize(&payload)
.map_err(|e| format!("compression cache serialize failed: {e}"))?;
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)
.map_err(|e| format!("compression cache mkdir failed: {e}"))?;
}
let tmp = path.with_extension("tmp");
fs::write(&tmp, bytes).map_err(|e| format!("compression cache write failed: {e}"))?;
fs::rename(&tmp, path).map_err(|e| format!("compression cache rename failed: {e}"))?;
Ok(())
}
async fn hydrate_compression_cache_from(&self, path: &FsPath) -> Result<usize, String> {
if !path.exists() {
return Ok(0);
}
let bytes = fs::read(path).map_err(|e| format!("compression cache read failed: {e}"))?;
let payload: PersistedCompressionCache = crate::bincode_compat::deserialize(&bytes)
.map_err(|e| format!("compression cache decode failed: {e}"))?;
if payload.version != 1 {
return Err(format!(
"unsupported compression cache version {}",
payload.version
));
}
let mut restored = 0usize;
for entry in payload.entries {
if entry.checkpoint.version != 1 {
continue;
}
let states = entry
.checkpoint
.layers
.iter()
.map(|lc| layer_checkpoint_to_tensor(lc, &self.device))
.collect::<candle_core::Result<Vec<_>>>()
.map_err(|e| format!("compression cache tensor restore failed: {e}"))?;
self.ttt_sessions
.insert_states(entry.session_id.clone(), states);
if let Ok(mut hashes) = self.adapted_context_hashes.write() {
hashes.insert(entry.session_id, entry.context_hash);
}
restored += 1;
}
Ok(restored)
}
pub fn with_master_vibe(mut self, vibe: Option<MasterVibe>) -> Self {
self.master_vibe = Arc::new(Mutex::new(vibe));
self
}
async fn flush_session_to_vibe(&self, handle: &SessionStates) {
let states = { handle.lock().await.clone() };
if let Ok(mut guard) = self.master_vibe.lock() {
if let Some(vibe) = guard.as_mut() {
if let Err(e) = vibe.commit_and_save(&states) {
eprintln!("[vibe] auto-commit skipped: {e}");
}
}
}
}
pub async fn flush_all_sessions_to_vibe(&self) {
let enabled = self
.master_vibe
.lock()
.map(|g| g.is_some())
.unwrap_or(false);
if !enabled {
return;
}
let handles = self.ttt_sessions.snapshot_handles();
if handles.is_empty() {
return;
}
eprintln!(
"[vibe] flushing {} live session(s) into master vibe",
handles.len()
);
for (_, handle) in handles {
self.flush_session_to_vibe(&handle).await;
}
}
pub fn with_claude_backend(mut self, backend: Option<ClaudeBackend>) -> Self {
self.claude_backend = Arc::new(backend);
self
}
pub fn with_router(mut self, router: Option<BackendRouter>) -> Self {
self.router = Arc::new(router);
self
}
pub fn with_mcp(mut self, mcp: Option<crate::mcp_stdio::McpContext>) -> Self {
self.mcp = Arc::new(mcp);
self
}
pub fn with_mcp_token(mut self, token: Option<String>) -> Self {
self.mcp_token = Arc::new(token);
self
}
pub fn with_api_key(mut self, key: Option<String>) -> Self {
self.api_key = Arc::new(key);
self
}
pub fn with_jit_exec_enabled(mut self, enabled: bool) -> Self {
self.jit_exec_enabled = enabled;
self
}
pub fn with_anthropic_forwarder(mut self, forwarder: Option<AnthropicForwarder>) -> Self {
self.anthropic_forwarder = Arc::new(forwarder);
self
}
pub fn with_openai_forwarder(mut self, forwarder: Option<OpenAiForwarder>) -> Self {
self.openai_forwarder = Arc::new(forwarder);
self
}
pub fn with_swarm_router(mut self, router: Option<SwarmRouter>) -> Self {
self.swarm_router = Arc::new(router);
self
}
pub fn with_compressor_config(mut self, config: CompressorConfig) -> Self {
self.controls = Arc::new(CompressionControls::from_config(&config));
self.compressor_config = Arc::new(config);
self
}
pub fn compression_active(&self) -> bool {
self.controls.enabled() && self.anthropic_forwarder.is_some()
}
pub fn openai_compression_active(&self) -> bool {
self.controls.enabled() && self.openai_forwarder.is_some()
}
fn refresh_session_metrics(&self) -> Result<(), ApiError> {
let sessions = self
.sessions
.read()
.map_err(|_| ApiError::Internal("session lock poisoned".into()))?;
let active = sessions
.values()
.filter(|session| !session.is_quantized())
.count();
let quantized = sessions
.values()
.filter(|session| session.is_quantized())
.count();
metrics::set_active_sessions(active);
metrics::set_quantized_sessions(quantized);
Ok(())
}
fn trigger_lru_vram_budget(&self) {
let sessions = self.sessions.clone();
tokio::spawn(async move {
if let Err(err) = enforce_lru_vram_budget_async(sessions).await {
eprintln!("[emergency] LRU budget enforcement failed: {err}");
}
});
}
}
async fn enforce_lru_vram_budget_async(
sessions: Arc<RwLock<HashMap<String, SessionData>>>,
) -> std::result::Result<(), String> {
loop {
let candidate = {
let sessions_guard = sessions
.read()
.map_err(|_| "session lock poisoned".to_string())?;
let active = sessions_guard
.iter()
.filter_map(|(id, session)| {
if session.is_quantized() {
None
} else {
Some((id.clone(), session.last_used))
}
})
.collect::<Vec<_>>();
if active.len() <= MAX_ACTIVE_VRAM_SESSIONS {
None
} else {
active.into_iter().min_by_key(|(_, last_used)| *last_used)
}
};
let Some((evict_session_id, baseline_last_used)) = candidate else {
refresh_session_metrics_from_sessions(&sessions)
.map_err(|e| format!("session metrics refresh failed: {e}"))?;
return Ok(());
};
let raw_layers = {
let sessions_guard = sessions
.read()
.map_err(|_| "session lock poisoned".to_string())?;
let Some(session) = sessions_guard.get(&evict_session_id) else {
continue;
};
match &session.residency {
SessionResidency::Quantized(_) => continue,
SessionResidency::Active(states) => states
.iter()
.map(|state| {
let cpu_f32 = state.to_device(&Device::Cpu)?.to_dtype(DType::F32)?;
let shape = cpu_f32.dims().to_vec();
let data = cpu_f32.contiguous()?.flatten_all()?.to_vec1::<f32>()?;
Ok((shape, data))
})
.collect::<candle_core::Result<Vec<_>>>()
.map_err(|e| format!("state staging failed: {e}"))?,
}
};
let quantized_layers = spawn_blocking(move || {
raw_layers
.into_iter()
.map(|(shape, data)| {
let total: usize = shape.iter().product();
let staged = Tensor::from_vec(data, (total,), &Device::Cpu)
.and_then(|t| t.reshape(shape.as_slice()))
.map_err(|e| e.to_string())?;
let (packed_indices, scale) =
NF4Quantizer::quantize_state(&staged).map_err(|e| e.to_string())?;
let (num_blocks, packed_width) =
packed_indices.dims2().map_err(|e| e.to_string())?;
let packed_indices = packed_indices
.to_dtype(DType::U8)
.and_then(|t| t.contiguous())
.and_then(|t| t.flatten_all())
.and_then(|t| t.to_vec1::<u8>())
.map_err(|e| e.to_string())?;
let scales = scale
.to_dtype(DType::F32)
.and_then(|t| t.contiguous())
.and_then(|t| t.flatten_all())
.and_then(|t| t.to_vec1::<f32>())
.map_err(|e| e.to_string())?;
if scales.len() != num_blocks {
return Err(format!(
"invalid scale length for shape {:?}: expected {num_blocks}, got {}",
shape,
scales.len(),
));
}
Ok(NF4QuantizedDescriptor {
shape,
packed_indices,
scales,
packed_width,
})
})
.collect::<std::result::Result<Vec<NF4QuantizedDescriptor>, String>>()
})
.await
.map_err(|e| format!("state offload task join failed: {e}"))?
.map_err(|e| format!("state offload quantization failed: {e}"))?;
let mut sessions_guard = sessions
.write()
.map_err(|_| "session lock poisoned".to_string())?;
if let Some(session) = sessions_guard.get_mut(&evict_session_id) {
let still_active = !session.is_quantized();
let unchanged_clock = session.last_used == baseline_last_used;
if still_active && unchanged_clock {
session.residency = SessionResidency::Quantized(quantized_layers);
metrics::mark_session_quantized(&evict_session_id, true);
}
}
}
}
fn refresh_session_metrics_from_sessions(
sessions: &Arc<RwLock<HashMap<String, SessionData>>>,
) -> std::result::Result<(), String> {
let sessions_guard = sessions
.read()
.map_err(|_| "session lock poisoned".to_string())?;
let active = sessions_guard
.values()
.filter(|session| !session.is_quantized())
.count();
let quantized = sessions_guard
.values()
.filter(|session| session.is_quantized())
.count();
metrics::set_active_sessions(active);
metrics::set_quantized_sessions(quantized);
Ok(())
}
#[derive(Debug)]
enum ApiError {
Internal(String),
NotFound(String),
BadRequest(String),
Conflict(String),
Forbidden(String),
Upstream {
status: u16,
message: String,
},
}
impl IntoResponse for ApiError {
fn into_response(self) -> Response {
match self {
ApiError::Internal(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg).into_response(),
ApiError::NotFound(msg) => (StatusCode::NOT_FOUND, msg).into_response(),
ApiError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg).into_response(),
ApiError::Conflict(msg) => (StatusCode::CONFLICT, msg).into_response(),
ApiError::Forbidden(msg) => (StatusCode::FORBIDDEN, msg).into_response(),
ApiError::Upstream { status, message } => {
let code = StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY);
(code, message).into_response()
}
}
}
}
#[derive(Debug, Serialize)]
struct ModelInfo {
id: String,
object: String,
created: u64,
owned_by: String,
}
#[derive(Debug, Serialize)]
struct ListModelsResponse {
object: String,
data: Vec<ModelInfo>,
}
#[derive(Debug, Deserialize)]
pub struct CompletionRequest {
#[serde(default)]
pub model: Option<String>,
pub prompt: String,
pub max_tokens: Option<usize>,
pub session_id: Option<String>,
}
#[derive(Debug, Serialize)]
struct CompletionChoice {
text: String,
index: usize,
finish_reason: String,
}
#[derive(Debug, Serialize)]
struct CompletionResponse {
id: String,
object: String,
created: u64,
model: String,
choices: Vec<CompletionChoice>,
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct ChatMessage {
pub role: String,
pub content: String,
}
#[derive(Debug, Deserialize)]
pub struct ChatCompletionRequest {
#[serde(default)]
pub model: Option<String>,
pub messages: Vec<ChatMessage>,
pub max_tokens: Option<usize>,
pub session_id: Option<String>,
#[serde(default)]
pub stream: Option<bool>,
}
#[derive(Debug, Serialize)]
struct ChatChoice {
index: usize,
message: ChatMessage,
finish_reason: String,
}
#[derive(Debug, Serialize)]
struct ChatCompletionResponse {
id: String,
object: String,
created: u64,
model: String,
choices: Vec<ChatChoice>,
}
#[derive(Debug, Deserialize)]
pub struct CreateSessionRequest {
#[serde(default)]
pub model: Option<String>,
}
#[derive(Debug, Serialize)]
struct CreateSessionResponse {
session_id: String,
object: String,
created: u64,
model: String,
}
#[derive(Debug, Serialize)]
struct DeleteSessionResponse {
session_id: String,
deleted: bool,
}
#[derive(Debug, Deserialize)]
pub struct AdaptRequest {
pub corpus: Vec<String>,
pub steps: Option<usize>,
pub session_id: Option<String>,
}
#[derive(Debug, Serialize)]
struct AdaptResponse {
session_id: String,
object: String,
steps_per_token: usize,
corpus_documents: usize,
}
#[derive(Debug, Deserialize)]
#[serde(untagged)]
pub enum AnthropicContent {
Text(String),
Blocks(Vec<AnthropicInputBlock>),
}
#[derive(Debug, Deserialize)]
pub struct AnthropicInputBlock {
#[serde(rename = "type", default)]
block_type: String,
#[serde(default)]
text: String,
}
impl AnthropicContent {
fn to_text(&self) -> String {
match self {
AnthropicContent::Text(s) => s.clone(),
AnthropicContent::Blocks(blocks) => blocks
.iter()
.filter(|b| b.block_type == "text" || b.block_type.is_empty())
.map(|b| b.text.clone())
.collect::<Vec<_>>()
.concat(),
}
}
}
#[derive(Debug, Deserialize)]
pub struct AnthropicMessagesRequest {
#[serde(default)]
pub model: Option<String>,
#[serde(default = "default_anthropic_max_tokens")]
pub max_tokens: usize,
pub messages: Vec<AnthropicInputMessage>,
#[serde(default)]
pub system: Option<AnthropicContent>,
#[serde(default)]
pub session_id: Option<String>,
}
fn default_anthropic_max_tokens() -> usize {
1024
}
#[derive(Debug, Deserialize)]
pub struct AnthropicInputMessage {
pub role: String,
pub content: AnthropicContent,
}
#[derive(Debug, Serialize)]
struct AnthropicOutputBlock {
#[serde(rename = "type")]
block_type: String,
text: String,
}
#[derive(Debug, Serialize)]
struct AnthropicUsage {
input_tokens: usize,
output_tokens: usize,
}
#[derive(Debug, Serialize)]
struct AnthropicMessagesResponse {
id: String,
#[serde(rename = "type")]
response_type: String,
role: String,
content: Vec<AnthropicOutputBlock>,
model: String,
stop_reason: String,
stop_sequence: Option<String>,
usage: AnthropicUsage,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct LayerCheckpoint {
pub shape: Vec<usize>,
pub data: Vec<f32>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SessionCheckpoint {
pub session_id: String,
pub version: u32,
pub created_at: u64,
pub layers: Vec<LayerCheckpoint>,
}
#[derive(Debug, Serialize, Deserialize)]
struct PersistedCompressionCache {
version: u32,
entries: Vec<PersistedCompressionEntry>,
}
#[derive(Debug, Serialize, Deserialize)]
struct PersistedCompressionEntry {
session_id: String,
context_hash: String,
checkpoint: SessionCheckpoint,
}
#[derive(Debug, Deserialize)]
struct TttFeedbackRequest {
session_id: String,
message: String,
#[serde(default)]
feedback_type: Option<String>,
#[serde(default)]
trace: Option<String>,
}
#[derive(Debug, Serialize)]
struct TttFeedbackResponse {
session_id: String,
feedback_tokens: usize,
persisted: bool,
cache_path: String,
}
#[derive(Debug, Deserialize)]
struct ClusterMergeRequest {
inputs: Vec<String>,
output: String,
#[serde(default)]
alpha: Option<f32>,
#[serde(default)]
method: Option<String>,
}
#[derive(Debug, Deserialize)]
struct HypervisorMountRequest {
root: String,
#[serde(default)]
session_id: Option<String>,
#[serde(default)]
warm_paths: Vec<String>,
}
#[derive(Debug, Serialize)]
struct HypervisorMountResponse {
mount: VfsMountReport,
warmed: Vec<VfsReadReport>,
vfs: VfsStats,
}
#[derive(Debug, Deserialize)]
struct HypervisorReadRequest {
path: String,
#[serde(default)]
session_id: Option<String>,
}
#[derive(Debug, Serialize)]
struct HypervisorReadResponse {
read: VfsReadReport,
vfs: VfsStats,
}
#[derive(Debug, Deserialize)]
struct HypervisorJitRunRequest {
#[serde(default)]
session_id: Option<String>,
command: String,
#[serde(default)]
args: Vec<String>,
#[serde(default)]
working_dir: Option<String>,
#[serde(default)]
source_path: Option<String>,
}
#[derive(Debug, Deserialize)]
struct BudgetRequest {
remaining_tokens: usize,
model: Option<String>,
session_id: Option<String>,
}
#[derive(Debug, Serialize)]
struct HypervisorJitRunResponse {
report: PolyJitReport,
source_restored: bool,
jit: PolyJitStatus,
}
#[derive(Debug, Serialize)]
struct HypervisorJitStatusResponse {
jit: PolyJitStatus,
vfs: VfsStats,
}
#[derive(Debug, Serialize)]
struct HypervisorQuantumStateResponse {
quantum: QuantumRuntimeStatus,
}
#[derive(Debug, Serialize)]
struct SwarmMatrixStateResponse {
matrix: SwarmMatrixState,
dwe: DweTelemetry,
exact_residual: ExactResidualTelemetry,
}
fn tensor_to_layer_checkpoint(t: &Tensor) -> candle_core::Result<LayerCheckpoint> {
let shape = t.dims().to_vec();
let data = t
.to_dtype(candle_core::DType::F32)?
.contiguous()?
.flatten_all()?
.to_vec1::<f32>()?;
Ok(LayerCheckpoint { shape, data })
}
fn layer_checkpoint_to_tensor(
lc: &LayerCheckpoint,
device: &Device,
) -> candle_core::Result<Tensor> {
let total: usize = lc.shape.iter().product();
if total != lc.data.len() {
candle_core::bail!(
"checkpoint shape {:?} implies {} elements but data has {}",
lc.shape,
total,
lc.data.len()
);
}
Tensor::from_vec(lc.data.clone(), (total,), device)?.reshape(lc.shape.as_slice())
}
fn tensor_is_finite(tensor: &Tensor) -> candle_core::Result<bool> {
let values = tensor
.to_dtype(DType::F32)?
.contiguous()?
.flatten_all()?
.to_vec1::<f32>()?;
Ok(values.into_iter().all(f32::is_finite))
}