use anyhow::{bail, ensure, Context, Result};
use clap::{Parser, Subcommand};
use futures::{SinkExt, StreamExt};
use klbr_core::{config::Config, models::LlmClient};
use klbr_daemon::{
daemon,
ingress::{admit_input, SourceRef},
session_registry::{ConfiguredPlacementSource, SessionRegistry},
};
use klbr_ipc::v2::{V2ClientMsg, V2ServerMsg};
use klbr_runtime::{
behavior::Release, protocol::SessionId, remote::placement::PlacementConfig, store::Store,
worker::WorkerConfig,
};
use std::{path::PathBuf, sync::Arc};
use tokio::sync::{broadcast, mpsc, watch};
use tokio_tungstenite::{connect_async, tungstenite::Message};
#[derive(Debug, Parser)]
#[command(name = "klbr", about = "one dogfood runtime, one v2 control path")]
struct Cli {
#[command(subcommand)]
command: Option<CliCommand>,
}
#[derive(Debug, Subcommand)]
enum CliCommand {
/// Start the runtime and bundled web UI. This is the default command.
Serve,
/// Create the state layout, schema, and first content-addressed behavior release.
Init,
/// Validate packaged paths, Python worker imports, behavior, storage, and model config.
Doctor,
/// Read live session state through the same v2 socket as the UI.
Status { session_id: Option<String> },
/// Hard-stop one session independently of model or policy callbacks.
Stop { session_id: String },
/// Finish accepted durable work, then stop one session.
Drain { session_id: String },
/// Replace one Python kernel; in-memory heap state is deliberately lost.
Reset { session_id: String },
/// Install a due context frontier for one session.
Compact { session_id: String },
/// Atomically activate a validated content-addressed behavior release.
Activate {
session_id: String,
expected_epoch: i64,
release_path: PathBuf,
},
}
#[derive(Debug)]
struct RuntimePaths {
data: PathBuf,
workspace: PathBuf,
behavior_seed: PathBuf,
skills_seed: PathBuf,
releases: PathBuf,
runner: PathBuf,
web_root: Option<PathBuf>,
python: PathBuf,
bind_addr: String,
}
impl RuntimePaths {
fn resolve() -> Result<Self> {
let workspace = env_path("KLBR_WORKSPACE")
.or_else(|| env_path("KLBR_SOURCE_ROOT"))
.unwrap_or(std::env::current_dir().context("cannot resolve current directory")?);
ensure!(
workspace.is_dir(),
"KLBR_WORKSPACE is not a directory: {}",
workspace.display()
);
let workspace = workspace
.canonicalize()
.context("cannot resolve KLBR_WORKSPACE")?;
let data = env_path("KLBR_DATA_DIR")
.or_else(|| dirs::data_dir().map(|path| path.join("klbr")))
.context("cannot resolve data directory; set KLBR_DATA_DIR")?;
let behavior_seed =
env_path("KLBR_BEHAVIOR_DIR").unwrap_or_else(|| workspace.join("behavior"));
ensure!(
behavior_seed.is_dir(),
"behavior seed is not a directory: {}",
behavior_seed.display()
);
let behavior_seed = behavior_seed
.canonicalize()
.context("cannot resolve behavior seed")?;
// Builtin skills are read where the daemon was started, and are allowed to
// be absent: a checkout that ships no skills is a runtime with no builtin
// namespace, not a broken one. The managed root is prepared instead, so an
// agent can write a skill without inventing the directory first.
let skills_seed = env_path("KLBR_SKILLS_DIR").unwrap_or_else(|| workspace.join("skills"));
let skills_seed = skills_seed.canonicalize().unwrap_or(skills_seed);
let runner = env_path("KLBR_WORKER_RUNNER")
.unwrap_or_else(|| workspace.join("python/klbr-runtime/src/klbr_runtime/__main__.py"));
ensure!(
runner.is_file(),
"Python worker runner is not a file: {}",
runner.display()
);
let runner = runner
.canonicalize()
.context("cannot resolve Python worker runner")?;
let web_root = env_path("KLBR_WEB_DIR")
.or_else(|| Some(workspace.join("klbr-web/dist")))
.filter(|path| path.join("index.html").is_file())
.map(|path| path.canonicalize())
.transpose()
.context("cannot resolve bundled web UI")?;
let bind_addr = std::env::var("KLBR_BIND_ADDR").unwrap_or_else(|_| "127.0.0.1:8765".into());
let socket: std::net::SocketAddr = bind_addr
.parse()
.context("KLBR_BIND_ADDR must be a numeric IP address and port")?;
ensure!(
socket.ip().is_loopback(),
"KLBR_BIND_ADDR must remain loopback-only; put an authenticated proxy in front"
);
Ok(Self {
releases: data.join("releases"),
data,
workspace,
behavior_seed,
skills_seed,
runner,
web_root,
python: env_path("KLBR_PYTHON").unwrap_or_else(|| "python3".into()),
bind_addr,
})
}
fn prepare(&self) -> Result<()> {
std::fs::create_dir_all(&self.data)
.with_context(|| format!("cannot create data directory {}", self.data.display()))?;
std::fs::create_dir_all(self.data.join("drafts"))?;
std::fs::create_dir_all(self.data.join("artifacts"))?;
std::fs::create_dir_all(self.data.join("skills"))?;
Ok(())
}
fn skills(&self) -> klbr_runtime::skills::SkillRoots {
klbr_runtime::skills::SkillRoots {
builtin: self.skills_seed.clone(),
managed: Some(self.data.join("skills")),
}
}
fn store(&self) -> PathBuf {
self.data.join("runtime.sqlite")
}
}
fn env_path(name: &str) -> Option<PathBuf> {
std::env::var_os(name)
.filter(|value| !value.is_empty())
.map(PathBuf::from)
}
#[tokio::main]
async fn main() -> Result<()> {
tracing_subscriber::fmt().try_init().ok();
let command = Cli::parse().command.unwrap_or(CliCommand::Serve);
let paths = RuntimePaths::resolve()?;
match command {
CliCommand::Serve => serve(paths).await,
CliCommand::Init => init(&paths),
CliCommand::Doctor => doctor(&paths).await,
CliCommand::Status { session_id } => status(&paths, session_id).await,
CliCommand::Stop { session_id } => control(&paths, V2ClientMsg::Stop { session_id }).await,
CliCommand::Drain { session_id } => {
control(&paths, V2ClientMsg::Drain { session_id }).await
}
CliCommand::Reset { session_id } => {
control(&paths, V2ClientMsg::ResetKernel { session_id }).await
}
CliCommand::Compact { session_id } => {
control(&paths, V2ClientMsg::Compact { session_id }).await
}
CliCommand::Activate {
session_id,
expected_epoch,
release_path,
} => {
control(
&paths,
V2ClientMsg::Activate {
session_id,
expected_epoch,
release_path: release_path.display().to_string(),
},
)
.await
}
}
}
fn init(paths: &RuntimePaths) -> Result<()> {
paths.prepare()?;
let release = Release::publish(&paths.behavior_seed, &paths.releases)
.context("failed to publish initial behavior release")?;
drop(Store::open(paths.store())?);
println!(
"{}",
serde_json::to_string_pretty(&serde_json::json!({
"status": "initialized",
"data_dir": paths.data,
"store": paths.store(),
"behavior_revision": release.id,
"behavior_release": release.path,
}))?
);
Ok(())
}
async fn doctor(paths: &RuntimePaths) -> Result<()> {
let scratch = tempfile::tempdir().context("cannot create doctor scratch directory")?;
let release = Release::publish(&paths.behavior_seed, &scratch.path().join("releases"))
.context("behavior source failed content-addressed publication")?;
drop(Store::open(scratch.path().join("runtime.sqlite"))?);
let python = tokio::process::Command::new(&paths.python)
.arg("-I")
.arg("-B")
.arg(&paths.runner)
.arg("--help")
.current_dir(&paths.workspace)
.output()
.await
.context("failed to start configured Python")?;
ensure!(
python.status.success(),
"Python worker import failed: {}",
String::from_utf8_lossy(&python.stderr)
);
let web_root = paths
.web_root
.as_ref()
.context("bundled web UI not found; set KLBR_WEB_DIR or build klbr-web/dist")?;
let config = Config::load().context("model configuration is invalid")?;
let skills = paths.skills();
let skill_catalog = skills.catalog();
let placement = placement_config_path()
.map(|path| PlacementConfig::from_path(&path).map(|_| path))
.transpose()?;
println!(
"{}",
serde_json::to_string_pretty(&serde_json::json!({
"status": "healthy",
"workspace": paths.workspace,
"data_dir": paths.data,
"python": paths.python,
"worker_runner": paths.runner,
"behavior_revision": release.id,
"schema": klbr_runtime::store::CURRENT_SCHEMA_VERSION,
"web_root": web_root,
"bind_addr": paths.bind_addr,
"models": {
"llm": config.models.llm.url,
"embedder": config.models.embedder.url,
"reranker": config.models.reranker.url,
},
"skills": {
"builtin_root": skills.builtin_root(),
"managed_root": skills.managed_root(),
"ids": skill_catalog.ids(),
"refused": skill_catalog.refused.iter().map(|refusal| serde_json::json!({
"id": refusal.id,
"reason": refusal.reason,
})).collect::<Vec<_>>(),
},
"placement": placement,
}))?
);
Ok(())
}
async fn serve(paths: RuntimePaths) -> Result<()> {
paths.prepare()?;
let config = Config::load()?;
tracing::info!(?paths, config = ?config, "loaded dogfood runtime");
let initial_release = Release::publish(&paths.behavior_seed, &paths.releases)
.context("failed to publish initial behavior release")?;
let default_soul = klbr_daemon::maintenance::install_standing_responsibility_with_release(
klbr_core::config::DEFAULT_SOUL,
Some(&initial_release),
);
let runtime_store = Store::open(paths.store())?;
let skills = paths.skills();
let skill_catalog = skills.catalog();
if !skill_catalog.skills.is_empty() || !skill_catalog.refused.is_empty() {
tracing::info!(
ids = ?skill_catalog.ids(),
refused = skill_catalog.refused.len(),
builtin = %skills.builtin_root().display(),
managed = ?skills.managed_root().map(|path| path.display().to_string()),
"loaded skill index"
);
}
let worker_config = WorkerConfig::installed(
paths.python.clone(),
paths.runner.clone(),
paths.workspace.clone(),
)
.with_skills(skills);
let placement = match placement_config_path() {
Some(path) => {
let config = PlacementConfig::from_path(&path)?;
tracing::info!(
path = %path.display(),
primary = ?config.primary.as_ref().map(|entry| entry.target.id.to_string()),
auxiliary = config.auxiliary.len(),
"loaded placement config"
);
Some(ConfiguredPlacementSource::build(config)?)
}
None => None,
};
let session_registry = Arc::new({
let registry = SessionRegistry::new(runtime_store, worker_config)
.with_guidance_budget_from_watermark(config.watermark_tokens)
.with_history_window(config.history_window);
match placement {
Some(source) => registry.with_placement(Arc::new(source)),
None => registry,
}
});
let llm = LlmClient::new(config.models);
let (v2_event_tx, _) = broadcast::channel::<V2ServerMsg>(4096);
let (shutdown_tx, shutdown_rx) = watch::channel(false);
session_registry.enable_default_peer_host(
Some(llm.clone()),
Some(initial_release.clone()),
Some(v2_event_tx.clone()),
);
let discord_tasks = if let Some(discord) = klbr_discord::DiscordRuntime::from_env().await? {
let (ingress_tx, ingress_rx) = mpsc::channel(1024);
let gateway = discord.spawn_gateway(ingress_tx);
let admission = tokio::spawn(forward_discord_ingress(
session_registry.clone(),
ingress_rx,
));
Some((gateway, admission))
} else {
None
};
if paths.web_root.is_none() {
tracing::warn!("bundled UI is absent; the v2 WebSocket remains available");
}
let mut service = tokio::spawn(daemon::serve(
paths.bind_addr,
paths.web_root,
v2_event_tx,
default_soul,
session_registry,
Some(llm),
Some(initial_release),
shutdown_rx,
));
tokio::select! {
result = &mut service => result??,
_ = shutdown_signal() => {
tracing::info!("shutdown signal received");
let _ = shutdown_tx.send(true);
service.await??;
}
}
if let Some((gateway, admission)) = discord_tasks {
gateway.abort();
admission.abort();
}
Ok(())
}
async fn forward_discord_ingress(
registry: Arc<SessionRegistry>,
mut ingress_rx: mpsc::Receiver<klbr_discord::DiscordIngress>,
) {
let session_id: SessionId = "default".parse().expect("static session id");
while let Some(batch) = ingress_rx.recv().await {
let source = SourceRef::from(batch);
if let Err(error) = admit_input(®istry, &session_id, source).await {
tracing::error!(%error, "failed to admit Discord input into runtime store");
}
}
}
async fn status(paths: &RuntimePaths, session_id: Option<String>) -> Result<()> {
let request = match session_id {
Some(session_id) => V2ClientMsg::SessionStatus { session_id },
None => V2ClientMsg::SessionList,
};
let reply = request_reply(paths, request, |message| {
matches!(
message,
V2ServerMsg::SessionStatus { .. } | V2ServerMsg::SessionList { .. }
)
})
.await?;
println!("{}", serde_json::to_string_pretty(&reply)?);
Ok(())
}
async fn control(paths: &RuntimePaths, request: V2ClientMsg) -> Result<()> {
let reply = request_reply(paths, request, |message| {
matches!(message, V2ServerMsg::OperatorAck { .. })
})
.await?;
if let V2ServerMsg::OperatorAck {
success: false,
message,
..
} = &reply
{
bail!(message
.clone()
.unwrap_or_else(|| "operator command failed".into()));
}
println!("{}", serde_json::to_string_pretty(&reply)?);
Ok(())
}
async fn request_reply(
paths: &RuntimePaths,
request: V2ClientMsg,
accept: impl Fn(&V2ServerMsg) -> bool,
) -> Result<V2ServerMsg> {
let url =
std::env::var("KLBR_WS_URL").unwrap_or_else(|_| format!("ws://{}/ws", paths.bind_addr));
let (mut socket, _) = connect_async(&url)
.await
.with_context(|| format!("cannot connect to {url}"))?;
socket
.send(Message::Text(serde_json::to_string(&request)?))
.await?;
tokio::time::timeout(std::time::Duration::from_secs(5), async {
while let Some(frame) = socket.next().await {
let frame = frame?;
let Message::Text(text) = frame else { continue };
let reply: V2ServerMsg = serde_json::from_str(&text)?;
if accept(&reply) {
return Ok(reply);
}
if let V2ServerMsg::Error { code, message } = reply {
bail!("{code}: {message}");
}
}
bail!("runtime closed the control connection")
})
.await
.context("runtime did not answer within five seconds")?
}
/// Path to the packaged placement declaration, or `None` for a local-only service.
fn placement_config_path() -> Option<PathBuf> {
env_path("KLBR_REMOTE_PLACEMENT")
}
async fn shutdown_signal() {
#[cfg(unix)]
{
let mut sigterm =
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()).ok();
tokio::select! {
_ = tokio::signal::ctrl_c() => {}
_ = async {
if let Some(sigterm) = sigterm.as_mut() {
sigterm.recv().await;
} else {
std::future::pending::<()>().await;
}
} => {}
}
}
#[cfg(not(unix))]
{
let _ = tokio::signal::ctrl_c().await;
}
}