cratesの整理

This commit is contained in:
2026-04-11 02:48:50 +09:00
parent cff082ff3a
commit 4c3f81b4fa
34 changed files with 2524 additions and 669 deletions
@@ -1,10 +1,10 @@
[package]
name = "insomnia"
name = "daemon"
version = "0.1.0"
edition.workspace = true
license.workspace = true
[dependencies]
insomnia-core = { path = "../insomnia-core" }
insomnia-daemon = { path = "../insomnia-daemon" }
manifest = { path = "../manifest" }
protocol = { path = "../protocol" }
tokio = { version = "1.49", features = ["full"] }
-20
View File
@@ -1,20 +0,0 @@
[package]
name = "insomnia-core"
version = "0.1.0"
edition.workspace = true
license.workspace = true
[dependencies]
llm-worker = { path = "../llm-worker" }
llm-worker-persistence = { path = "../llm-worker-persistence" }
serde = { version = "1.0", features = ["derive"] }
toml = "0.8"
uuid = { version = "1", features = ["v7", "serde"] }
thiserror = "2.0"
tokio = { version = "1.49", features = ["fs"] }
[dev-dependencies]
tokio = { version = "1.49", features = ["macros", "rt-multi-thread"] }
tempfile = "3.24"
dotenv = "0.15"
llm-worker-persistence = { path = "../llm-worker-persistence" }
-9
View File
@@ -1,9 +0,0 @@
pub mod manifest;
pub mod pod;
pub mod provider;
pub mod scope;
pub use manifest::{PodManifest, ProviderConfig, ProviderKind};
pub use pod::{Pod, PodError, PodId, PodRunResult, apply_worker_manifest, new_pod_id};
pub use provider::build_client;
pub use scope::Scope;
-10
View File
@@ -1,10 +0,0 @@
[package]
name = "insomnia-daemon"
version = "0.1.0"
edition.workspace = true
license.workspace = true
[dependencies]
insomnia-core = { path = "../insomnia-core" }
llm-worker-persistence = { path = "../llm-worker-persistence" }
tokio = { version = "1.49", features = ["full"] }
-3
View File
@@ -1,3 +0,0 @@
fn main() {
println!("insomnia");
}
+12
View File
@@ -0,0 +1,12 @@
[package]
name = "manifest"
version = "0.1.0"
edition.workspace = true
license.workspace = true
[dependencies]
serde = { version = "1.0.228", features = ["derive"] }
toml = "1.1.2"
[dev-dependencies]
tempfile = "3.27.0"
@@ -1,12 +1,16 @@
mod scope;
pub use scope::Scope;
use std::path::PathBuf;
use serde::Deserialize;
use serde::{Deserialize, Serialize};
/// Declarative configuration for a Pod.
///
/// Parsed from a TOML manifest file. Describes the provider, model,
/// system prompt, and optional directory scope.
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PodManifest {
pub pod: PodMeta,
pub provider: ProviderConfig,
@@ -16,13 +20,13 @@ pub struct PodManifest {
}
/// Pod metadata.
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PodMeta {
pub name: String,
}
/// LLM provider configuration.
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderConfig {
pub kind: ProviderKind,
pub model: String,
@@ -35,7 +39,7 @@ pub struct ProviderConfig {
}
/// Supported LLM providers.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ProviderKind {
Anthropic,
@@ -45,7 +49,7 @@ pub enum ProviderKind {
}
/// Worker-level configuration embedded in the manifest.
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerManifest {
#[serde(default)]
pub system_prompt: Option<String>,
@@ -56,7 +60,7 @@ pub struct WorkerManifest {
}
/// Directory scope configuration.
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScopeConfig {
pub root: PathBuf,
}
+26
View File
@@ -0,0 +1,26 @@
[package]
name = "pod"
version = "0.1.0"
edition.workspace = true
license.workspace = true
[dependencies]
llm-worker = { version = "0.2.1", path = "../llm-worker" }
llm-worker-persistence = { version = "0.1.0", path = "../llm-worker-persistence" }
manifest = { version = "0.1.0", path = "../manifest" }
protocol = { version = "0.1.0", path = "../protocol" }
provider = { version = "0.1.0", path = "../provider" }
serde = { version = "1.0.228", features = ["derive"] }
serde_json = "1.0.149"
thiserror = "2.0"
tokio = { version = "1.49", features = ["fs", "io-util", "net", "sync"] }
toml = "1.1.2"
uuid = { version = "1.23.0", features = ["v7", "serde"] }
[dev-dependencies]
async-trait = "0.1.89"
dotenv = "0.15.0"
futures = "0.3.32"
llm-worker-persistence = { path = "../llm-worker-persistence" }
tempfile = "3.27.0"
tokio = { version = "1.49", features = ["macros", "rt-multi-thread", "time"] }
@@ -8,10 +8,10 @@
//!
//! ```bash
//! echo "ANTHROPIC_API_KEY=your-key" > .env
//! cargo run -p insomnia-core --example pod_cli
//! cargo run -p pod --example pod_cli
//! ```
use insomnia_core::{Pod, PodManifest, PodRunResult};
use pod::{Pod, PodManifest, PodRunResult};
use llm_worker_persistence::FsStore;
const MANIFEST_TOML: &str = r#"
+99
View File
@@ -0,0 +1,99 @@
//! Pod Protocol example: control a Pod via PodHandle and stream events.
//!
//! ```bash
//! echo "ANTHROPIC_API_KEY=your-key" > .env
//! cargo run -p pod --example pod_protocol
//! ```
use pod::{Event, Method, PodController, PodManifest};
use llm_worker_persistence::FsStore;
const MANIFEST_TOML: &str = r#"
[pod]
name = "protocol-demo"
[provider]
kind = "anthropic"
model = "claude-sonnet-4-20250514"
api_key_env = "ANTHROPIC_API_KEY"
[worker]
system_prompt = "You are a concise assistant. Reply in one or two sentences."
max_tokens = 256
"#;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
dotenv::dotenv().ok();
let manifest = PodManifest::from_toml(MANIFEST_TOML)?;
let tmp = tempfile::tempdir()?;
let store = FsStore::new(tmp.path()).await?;
let pod = pod::Pod::from_manifest(manifest, store, None).await?;
let runtime_tmp = tempfile::tempdir()?;
let handle = PodController::spawn(pod, runtime_tmp.path()).await?;
// Check initial status via shared state
println!("[shared_state] {}", handle.shared_state.status_json());
// Check runtime directory files
println!("[runtime_dir] {:?}", handle.runtime_dir.path());
// Spawn event listener
let mut rx = handle.subscribe();
let shared = handle.shared_state.clone();
let listener = tokio::spawn(async move {
while let Ok(event) = rx.recv().await {
match &event {
Event::TurnStart { turn } => {
println!("[turn {turn}] start");
}
Event::TextDelta { text } => {
print!("{text}");
}
Event::TextDone { .. } => {
println!();
}
Event::TurnEnd { turn, result } => {
println!("[turn {turn}] end ({result:?})");
println!("[shared_state] {}", shared.status_json());
}
Event::ToolCallStart { name, .. } => {
println!("[tool] {name}");
}
Event::Usage {
input_tokens,
output_tokens,
} => {
println!(
"[usage] in={} out={}",
input_tokens.unwrap_or(0),
output_tokens.unwrap_or(0)
);
}
Event::Error { code, message } => {
println!("[error] {code:?}: {message}");
}
_ => {}
}
}
});
// Send a run method
handle
.send(Method::Run {
input: "What is the capital of France?".into(),
})
.await?;
// Wait for completion
tokio::time::sleep(std::time::Duration::from_secs(15)).await;
println!("\n[shared_state] final: {}", handle.shared_state.status_json());
println!("[history] {} bytes", handle.shared_state.history_json().len());
drop(handle);
let _ = listener.await;
Ok(())
}
+329
View File
@@ -0,0 +1,329 @@
use std::path::Path;
use std::sync::Arc;
use llm_worker::hook::ToolCall;
use llm_worker::llm_client::client::LlmClient;
use llm_worker::subscriber::WorkerSubscriber;
use llm_worker::timeline::event::{ErrorEvent, UsageEvent};
use llm_worker::timeline::{TextBlockEvent, ToolUseBlockEvent};
use llm_worker::WorkerError;
use llm_worker_persistence::Store;
use tokio::sync::{broadcast, mpsc};
use crate::pod::{Pod, PodRunResult, PodError};
use protocol::{ErrorCode, Event, Method, TurnResult};
use crate::runtime_dir::RuntimeDir;
use crate::shared_state::{PodSharedState, PodStatus};
use crate::socket_server::SocketServer;
// ---------------------------------------------------------------------------
// PodHandle — client-facing, Clone-able
// ---------------------------------------------------------------------------
#[derive(Clone)]
pub struct PodHandle {
method_tx: mpsc::Sender<Method>,
event_tx: broadcast::Sender<Event>,
pub shared_state: Arc<PodSharedState>,
pub runtime_dir: Arc<RuntimeDir>,
}
impl PodHandle {
pub async fn send(&self, method: Method) -> Result<(), mpsc::error::SendError<Method>> {
self.method_tx.send(method).await
}
pub fn subscribe(&self) -> broadcast::Receiver<Event> {
self.event_tx.subscribe()
}
/// Broadcast an event to all listeners (including socket clients).
pub fn send_event(&self, event: Event) -> Result<usize, broadcast::error::SendError<Event>> {
self.event_tx.send(event)
}
}
// ---------------------------------------------------------------------------
// PodController — actor that owns a Pod
// ---------------------------------------------------------------------------
pub struct PodController;
impl PodController {
pub async fn spawn<C, St>(
mut pod: Pod<C, St>,
runtime_base: &Path,
) -> Result<PodHandle, std::io::Error>
where
C: LlmClient + 'static,
St: Store + 'static,
{
let (method_tx, mut method_rx) = mpsc::channel::<Method>(32);
let (event_tx, _) = broadcast::channel::<Event>(256);
let manifest_toml = toml::to_string_pretty(pod.manifest()).unwrap_or_default();
let shared_state = Arc::new(PodSharedState::new(
pod.manifest().pod.name.clone(),
pod.session_id(),
manifest_toml.clone(),
));
// Create runtime directory and write initial files
let runtime_dir = RuntimeDir::create(runtime_base, &pod.manifest().pod.name).await?;
runtime_dir.write_manifest(&manifest_toml).await?;
runtime_dir.write_status(&shared_state).await?;
runtime_dir.write_history(&shared_state).await?;
let runtime_dir = Arc::new(runtime_dir);
let handle = PodHandle {
method_tx,
event_tx: event_tx.clone(),
shared_state: shared_state.clone(),
runtime_dir: runtime_dir.clone(),
};
// Start socket server (lives as a background task, cleaned up on drop via RuntimeDir)
let _socket_server = SocketServer::start(&handle).await?;
// Keep the server alive by moving it into the controller task
// (it will be dropped when the task ends)
// Register the event bridge subscriber on the worker
let bridge = EventBridgeSubscriber {
event_tx: event_tx.clone(),
};
pod.session_mut().worker.subscribe(bridge);
// Clone cancel sender before moving pod
let cancel_tx = pod.session_mut().worker.cancel_sender();
tokio::spawn(async move {
// Hold socket server alive for the lifetime of the controller task
let _socket_server = _socket_server;
loop {
let method = match method_rx.recv().await {
Some(m) => m,
None => break,
};
match method {
Method::Run { input } => {
if shared_state.get_status() != PodStatus::Idle {
let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Pod is already executing a turn".into(),
});
continue;
}
shared_state.set_status(PodStatus::Running);
let _ = runtime_dir.write_status(&shared_state).await;
let new_status = run_with_cancel_support(
pod.run(&input),
&mut method_rx,
&event_tx,
&cancel_tx,
&shared_state,
)
.await;
let items = pod.session_mut().worker.history().to_vec();
shared_state.update_history(items);
shared_state.set_status(new_status);
let _ = runtime_dir.write_status(&shared_state).await;
let _ = runtime_dir.write_history(&shared_state).await;
}
Method::Resume => {
if shared_state.get_status() != PodStatus::Paused {
let _ = event_tx.send(Event::Error {
code: ErrorCode::NotPaused,
message: "Pod is not paused".into(),
});
continue;
}
shared_state.set_status(PodStatus::Running);
let _ = runtime_dir.write_status(&shared_state).await;
let new_status = run_with_cancel_support(
pod.resume(),
&mut method_rx,
&event_tx,
&cancel_tx,
&shared_state,
)
.await;
let items = pod.session_mut().worker.history().to_vec();
shared_state.update_history(items);
shared_state.set_status(new_status);
let _ = runtime_dir.write_status(&shared_state).await;
let _ = runtime_dir.write_history(&shared_state).await;
}
Method::Cancel => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::NotRunning,
message: "Pod is not running".into(),
});
}
}
}
});
Ok(handle)
}
}
/// Runs a Pod future while concurrently processing incoming methods.
/// Only `Cancel` is handled during execution; `Run` and `Resume` get errors.
async fn run_with_cancel_support<F>(
pod_future: F,
method_rx: &mut mpsc::Receiver<Method>,
event_tx: &broadcast::Sender<Event>,
cancel_tx: &mpsc::Sender<()>,
shared_state: &Arc<PodSharedState>,
) -> PodStatus
where
F: std::future::Future<Output = Result<PodRunResult, PodError>>,
{
tokio::pin!(pod_future);
loop {
tokio::select! {
result = &mut pod_future => {
return match result {
Ok(r) => match r {
PodRunResult::Finished => PodStatus::Idle,
PodRunResult::Paused => PodStatus::Paused,
},
Err(e) => {
let code = worker_error_code(&e);
let _ = event_tx.send(Event::Error {
code,
message: e.to_string(),
});
PodStatus::Idle
}
};
}
method = method_rx.recv() => {
match method {
Some(Method::Cancel) => {
let _ = cancel_tx.try_send(());
}
Some(Method::Run { .. } | Method::Resume) => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Pod is already executing a turn".into(),
});
}
None => {
let _ = cancel_tx.try_send(());
shared_state.set_status(PodStatus::Idle);
return PodStatus::Idle;
}
}
}
}
}
}
fn worker_error_code(e: &PodError) -> ErrorCode {
match e {
PodError::Session(se) => {
use llm_worker_persistence::SessionError;
match se {
SessionError::Worker(we) => match we {
WorkerError::Tool(_) => ErrorCode::ToolError,
WorkerError::Client(_) => ErrorCode::ProviderError,
_ => ErrorCode::Internal,
},
_ => ErrorCode::Internal,
}
}
PodError::Provider(_) => ErrorCode::ProviderError,
_ => ErrorCode::Internal,
}
}
// ---------------------------------------------------------------------------
// EventBridgeSubscriber — bridges Worker events to broadcast channel
// ---------------------------------------------------------------------------
struct EventBridgeSubscriber {
event_tx: broadcast::Sender<Event>,
}
impl WorkerSubscriber for EventBridgeSubscriber {
type TextBlockScope = ();
type ToolUseBlockScope = ();
fn on_turn_start(&mut self, turn: usize) {
let _ = self.event_tx.send(Event::TurnStart { turn });
}
fn on_turn_end(&mut self, turn: usize) {
let _ = self.event_tx.send(Event::TurnEnd {
turn,
result: TurnResult::Finished,
});
}
fn on_text_block(&mut self, _scope: &mut (), event: &TextBlockEvent) {
match event {
TextBlockEvent::Delta(text) => {
let _ = self.event_tx.send(Event::TextDelta {
text: text.clone(),
});
}
TextBlockEvent::Start(_) | TextBlockEvent::Stop(_) => {}
}
}
fn on_text_complete(&mut self, text: &str) {
let _ = self.event_tx.send(Event::TextDone {
text: text.to_owned(),
});
}
fn on_tool_use_block(&mut self, _scope: &mut (), event: &ToolUseBlockEvent) {
match event {
ToolUseBlockEvent::Start(start) => {
let _ = self.event_tx.send(Event::ToolCallStart {
id: start.id.clone(),
name: start.name.clone(),
});
}
ToolUseBlockEvent::InputJsonDelta(json) => {
let _ = self.event_tx.send(Event::ToolCallArgsDelta {
id: String::new(),
json: json.clone(),
});
}
ToolUseBlockEvent::Stop(_) => {}
}
}
fn on_tool_call_complete(&mut self, call: &ToolCall) {
let _ = self.event_tx.send(Event::ToolCallDone {
id: call.id.clone(),
name: call.name.clone(),
arguments: call.input.to_string(),
});
}
fn on_usage(&mut self, event: &UsageEvent) {
let _ = self.event_tx.send(Event::Usage {
input_tokens: event.input_tokens,
output_tokens: event.output_tokens,
});
}
fn on_error(&mut self, event: &ErrorEvent) {
let _ = self.event_tx.send(Event::Error {
code: ErrorCode::ProviderError,
message: event.message.clone(),
});
}
}
+15
View File
@@ -0,0 +1,15 @@
pub mod controller;
pub mod runtime_dir;
pub mod shared_state;
pub mod socket_server;
mod pod;
pub use controller::{PodController, PodHandle};
pub use manifest::{PodManifest, ProviderConfig, ProviderKind, Scope};
pub use pod::{Pod, PodError, PodId, PodRunResult, apply_worker_manifest, new_pod_id};
pub use protocol::{ErrorCode, Event, Method, TurnResult};
pub use provider::{ProviderError, build_client};
pub use runtime_dir::RuntimeDir;
pub use shared_state::{PodSharedState, PodStatus};
pub use socket_server::SocketServer;
@@ -5,8 +5,7 @@ use llm_worker_persistence::{
Session, SessionConfig, SessionError, SessionId, Store, StoreError,
};
use crate::manifest::{PodManifest, WorkerManifest};
use crate::scope::Scope;
use manifest::{PodManifest, Scope, WorkerManifest};
/// Pod identifier. UUID v7 (time-ordered).
pub type PodId = uuid::Uuid;
@@ -117,7 +116,7 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
store: St,
scope: Option<Scope>,
) -> Result<Self, PodError> {
let client = crate::provider::build_client(&manifest.provider)?;
let client = provider::build_client(&manifest.provider)?;
let mut worker = Worker::new(client);
apply_worker_manifest(&mut worker, &manifest.worker);
let session = Session::new(worker, store, SessionConfig::default()).await?;
@@ -175,6 +174,6 @@ pub enum PodError {
#[error("scope violation: {path} is outside the allowed directory")]
ScopeViolation { path: String },
#[error("provider configuration error: {0}")]
ProviderConfig(String),
#[error(transparent)]
Provider(#[from] provider::ProviderError),
}
+194
View File
@@ -0,0 +1,194 @@
use std::io;
use std::path::{Path, PathBuf};
use tokio::fs;
use crate::shared_state::PodSharedState;
/// Manages the Pod's runtime directory on tmpfs.
///
/// ```text
/// $XDG_RUNTIME_DIR/insomnia/{pod_name}/
/// ├── pid
/// ├── status.json
/// ├── manifest.toml
/// ├── history.json
/// └── sock (created by socket listener, not by RuntimeDir)
/// ```
///
/// Files are written atomically (write tmp → rename).
/// The directory is removed on drop.
pub struct RuntimeDir {
path: PathBuf,
}
impl RuntimeDir {
/// Create the runtime directory and write the PID file.
pub async fn create(base: &Path, pod_name: &str) -> Result<Self, io::Error> {
let path = base.join(pod_name);
fs::create_dir_all(&path).await?;
let pid = std::process::id().to_string();
fs::write(path.join("pid"), pid.as_bytes()).await?;
Ok(Self { path })
}
/// Create in the default base directory.
///
/// Uses `$XDG_RUNTIME_DIR/insomnia/` if available,
/// otherwise falls back to `~/.insomnia/run/`.
pub async fn create_default(pod_name: &str) -> Result<Self, io::Error> {
let base = default_base()?;
Self::create(&base, pod_name).await
}
/// Write status.json atomically.
pub async fn write_status(&self, state: &PodSharedState) -> Result<(), io::Error> {
let content = state.status_json();
atomic_write(&self.path.join("status.json"), content.as_bytes()).await
}
/// Write manifest.toml (typically once at startup).
pub async fn write_manifest(&self, toml: &str) -> Result<(), io::Error> {
atomic_write(&self.path.join("manifest.toml"), toml.as_bytes()).await
}
/// Write history.json atomically.
pub async fn write_history(&self, state: &PodSharedState) -> Result<(), io::Error> {
let content = state.history_json();
atomic_write(&self.path.join("history.json"), content.as_bytes()).await
}
/// Path to this Pod's runtime directory.
pub fn path(&self) -> &Path {
&self.path
}
/// Path where the Unix socket should be created.
pub fn socket_path(&self) -> PathBuf {
self.path.join("sock")
}
}
impl Drop for RuntimeDir {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.path);
}
}
/// Atomic write: write to a temp file, then rename.
async fn atomic_write(target: &Path, content: &[u8]) -> Result<(), io::Error> {
let tmp = target.with_extension("tmp");
fs::write(&tmp, content).await?;
fs::rename(&tmp, target).await?;
Ok(())
}
/// Resolve the default base directory for runtime data.
fn default_base() -> Result<PathBuf, io::Error> {
if let Ok(runtime_dir) = std::env::var("XDG_RUNTIME_DIR") {
Ok(PathBuf::from(runtime_dir).join("insomnia"))
} else if let Ok(home) = std::env::var("HOME") {
Ok(PathBuf::from(home).join(".insomnia").join("run"))
} else {
Err(io::Error::new(
io::ErrorKind::NotFound,
"neither XDG_RUNTIME_DIR nor HOME is set",
))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::shared_state::{PodSharedState, PodStatus};
fn test_state() -> PodSharedState {
PodSharedState::new(
"test-pod".into(),
llm_worker_persistence::new_session_id(),
"[pod]\nname = \"test-pod\"".into(),
)
}
#[tokio::test]
async fn creates_directory_and_pid() {
let tmp = tempfile::tempdir().unwrap();
let rt = RuntimeDir::create(tmp.path(), "my-pod").await.unwrap();
assert!(rt.path().join("pid").exists());
let pid = std::fs::read_to_string(rt.path().join("pid")).unwrap();
assert_eq!(pid, std::process::id().to_string());
}
#[tokio::test]
async fn write_status_creates_file() {
let tmp = tempfile::tempdir().unwrap();
let rt = RuntimeDir::create(tmp.path(), "my-pod").await.unwrap();
let state = test_state();
rt.write_status(&state).await.unwrap();
let content = std::fs::read_to_string(rt.path().join("status.json")).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&content).unwrap();
assert_eq!(parsed["state"], "idle");
assert_eq!(parsed["pod_name"], "test-pod");
}
#[tokio::test]
async fn write_status_reflects_changes() {
let tmp = tempfile::tempdir().unwrap();
let rt = RuntimeDir::create(tmp.path(), "my-pod").await.unwrap();
let state = test_state();
state.set_status(PodStatus::Running);
rt.write_status(&state).await.unwrap();
let content = std::fs::read_to_string(rt.path().join("status.json")).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&content).unwrap();
assert_eq!(parsed["state"], "running");
}
#[tokio::test]
async fn write_manifest_creates_file() {
let tmp = tempfile::tempdir().unwrap();
let rt = RuntimeDir::create(tmp.path(), "my-pod").await.unwrap();
rt.write_manifest("[pod]\nname = \"test\"").await.unwrap();
let content = std::fs::read_to_string(rt.path().join("manifest.toml")).unwrap();
assert_eq!(content, "[pod]\nname = \"test\"");
}
#[tokio::test]
async fn write_history_creates_file() {
let tmp = tempfile::tempdir().unwrap();
let rt = RuntimeDir::create(tmp.path(), "my-pod").await.unwrap();
let state = test_state();
rt.write_history(&state).await.unwrap();
let content = std::fs::read_to_string(rt.path().join("history.json")).unwrap();
assert_eq!(content, "[]");
}
#[tokio::test]
async fn socket_path() {
let tmp = tempfile::tempdir().unwrap();
let rt = RuntimeDir::create(tmp.path(), "my-pod").await.unwrap();
assert_eq!(rt.socket_path(), rt.path().join("sock"));
}
#[tokio::test]
async fn drop_removes_directory() {
let tmp = tempfile::tempdir().unwrap();
let dir_path;
{
let rt = RuntimeDir::create(tmp.path(), "my-pod").await.unwrap();
dir_path = rt.path().to_owned();
assert!(dir_path.exists());
}
assert!(!dir_path.exists());
}
}
+149
View File
@@ -0,0 +1,149 @@
use std::sync::RwLock;
use llm_worker::llm_client::types::Item;
use llm_worker_persistence::SessionId;
use serde::{Deserialize, Serialize};
/// Shared state between PodController and runtime directory.
///
/// Controller updates this in-memory; RuntimeDir writes it to disk.
/// Wrapped in `Arc` for sharing.
pub struct PodSharedState {
pub pod_name: String,
pub session_id: SessionId,
pub manifest_toml: String,
pub status: RwLock<PodStatus>,
pub history: RwLock<Vec<Item>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PodStatus {
Idle,
Running,
Paused,
}
impl PodSharedState {
pub fn new(
pod_name: String,
session_id: SessionId,
manifest_toml: String,
) -> Self {
Self {
pod_name,
session_id,
manifest_toml,
status: RwLock::new(PodStatus::Idle),
history: RwLock::new(Vec::new()),
}
}
pub fn set_status(&self, status: PodStatus) {
if let Ok(mut s) = self.status.write() {
*s = status;
}
}
pub fn get_status(&self) -> PodStatus {
self.status.read().map(|s| *s).unwrap_or(PodStatus::Idle)
}
pub fn update_history(&self, items: Vec<Item>) {
if let Ok(mut h) = self.history.write() {
*h = items;
}
}
/// Serialize status as JSON.
pub fn status_json(&self) -> String {
let status = self.get_status();
serde_json::json!({
"state": status,
"session_id": self.session_id.to_string(),
"pod_name": self.pod_name,
})
.to_string()
}
/// Serialize history as JSON.
pub fn history_json(&self) -> String {
if let Ok(h) = self.history.read() {
serde_json::to_string(&*h).unwrap_or_else(|_| "[]".into())
} else {
"[]".into()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use llm_worker::llm_client::types::{ContentPart, Item, Role};
fn test_state() -> PodSharedState {
PodSharedState::new(
"test-pod".into(),
llm_worker_persistence::new_session_id(),
"[pod]\nname = \"test-pod\"".into(),
)
}
#[test]
fn initial_status_is_idle() {
let state = test_state();
assert_eq!(state.get_status(), PodStatus::Idle);
}
#[test]
fn set_and_get_status() {
let state = test_state();
state.set_status(PodStatus::Running);
assert_eq!(state.get_status(), PodStatus::Running);
state.set_status(PodStatus::Paused);
assert_eq!(state.get_status(), PodStatus::Paused);
}
#[test]
fn status_json_contains_fields() {
let state = test_state();
let json = state.status_json();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["state"], "idle");
assert_eq!(parsed["pod_name"], "test-pod");
assert!(parsed["session_id"].is_string());
}
#[test]
fn status_json_reflects_changes() {
let state = test_state();
state.set_status(PodStatus::Running);
let json = state.status_json();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["state"], "running");
}
#[test]
fn history_json_empty_initially() {
let state = test_state();
assert_eq!(state.history_json(), "[]");
}
#[test]
fn history_json_after_update() {
let state = test_state();
let items = vec![Item::Message {
id: None,
role: Role::Assistant,
content: vec![ContentPart::Text {
text: "Hello".into(),
}],
status: None,
}];
state.update_history(items);
let json = state.history_json();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert!(parsed.is_array());
assert_eq!(parsed[0]["role"], "assistant");
}
}
+102
View File
@@ -0,0 +1,102 @@
use std::io;
use std::path::PathBuf;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::UnixListener;
use tokio::task::JoinHandle;
use crate::controller::PodHandle;
use protocol::{Event, Method};
/// Unix socket server for Pod Protocol.
///
/// Listens on the Pod's runtime directory socket path.
/// Each client connection gets bidirectional JSONL:
/// - Client writes Method lines → forwarded to PodController
/// - Pod events → written as Event lines to all connected clients
pub struct SocketServer {
_accept_task: JoinHandle<()>,
path: PathBuf,
}
impl SocketServer {
/// Start listening on the PodHandle's socket path.
pub async fn start(handle: &PodHandle) -> Result<Self, io::Error> {
let path = handle.runtime_dir.socket_path();
// Remove stale socket file if it exists
let _ = tokio::fs::remove_file(&path).await;
let listener = UnixListener::bind(&path)?;
let handle = handle.clone();
let _accept_task = tokio::spawn(async move {
loop {
match listener.accept().await {
Ok((stream, _)) => {
let handle = handle.clone();
tokio::spawn(handle_connection(stream, handle));
}
Err(_) => break,
}
}
});
Ok(Self {
_accept_task,
path,
})
}
/// The socket file path.
pub fn path(&self) -> &std::path::Path {
&self.path
}
}
impl Drop for SocketServer {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.path);
}
}
async fn handle_connection(stream: tokio::net::UnixStream, handle: PodHandle) {
let (reader, mut writer) = stream.into_split();
let mut lines = BufReader::new(reader).lines();
let mut rx = handle.subscribe();
// Event writer: broadcast events → socket
let write_task = tokio::spawn(async move {
while let Ok(event) = rx.recv().await {
if let Ok(line) = event.to_json_line() {
let mut buf = line.into_bytes();
buf.push(b'\n');
if writer.write_all(&buf).await.is_err() {
break;
}
}
}
});
// Method reader: socket → controller
while let Ok(Some(line)) = lines.next_line().await {
if line.is_empty() {
continue;
}
match Method::from_json_line(&line) {
Ok(method) => {
let _ = handle.send(method).await;
}
Err(e) => {
// Send parse error back as an event
let _ = handle.send_event(Event::Error {
code: protocol::ErrorCode::Internal,
message: format!("invalid method: {e}"),
});
}
}
}
// Client disconnected — stop the write task
write_task.abort();
}
+429
View File
@@ -0,0 +1,429 @@
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use async_trait::async_trait;
use futures::Stream;
use llm_worker::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
use llm_worker::llm_client::{ClientError, LlmClient, Request};
use llm_worker::Worker;
use llm_worker_persistence::FsStore;
use pod::{
Event, Method, Pod, PodController, PodManifest, PodStatus,
};
// ---------------------------------------------------------------------------
// Mock LLM Client
// ---------------------------------------------------------------------------
#[derive(Clone)]
struct MockClient {
responses: Arc<Vec<Vec<LlmEvent>>>,
call_count: Arc<AtomicUsize>,
}
impl MockClient {
fn new(events: Vec<LlmEvent>) -> Self {
Self {
responses: Arc::new(vec![events]),
call_count: Arc::new(AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl LlmClient for MockClient {
async fn stream(
&self,
_request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
if count >= self.responses.len() {
return Err(ClientError::Api {
status: Some(500),
code: Some("mock".into()),
message: "No more responses".into(),
});
}
let events = self.responses[count].clone();
let stream = futures::stream::iter(events.into_iter().map(Ok));
Ok(Box::pin(stream))
}
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
fn simple_text_events() -> Vec<LlmEvent> {
vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, "Hello"),
LlmEvent::text_delta(0, " World"),
LlmEvent::text_block_stop(0, None),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
const MANIFEST_TOML: &str = r#"
[pod]
name = "test-pod"
[provider]
kind = "anthropic"
model = "test-model"
[worker]
max_tokens = 100
"#;
async fn make_pod(client: MockClient) -> Pod<MockClient, FsStore> {
let manifest = PodManifest::from_toml(MANIFEST_TOML).unwrap();
let tmp = tempfile::tempdir().unwrap();
let store = FsStore::new(tmp.path()).await.unwrap();
// Leak tempdir to keep it alive
std::mem::forget(tmp);
let worker = Worker::new(client);
Pod::new(manifest, worker, store, None).await.unwrap()
}
use pod::PodHandle;
async fn spawn_controller(pod: Pod<MockClient, FsStore>) -> PodHandle {
let tmp = tempfile::tempdir().unwrap();
let runtime_base = tmp.path().to_owned();
// Leak tempdir so it survives the test
std::mem::forget(tmp);
PodController::spawn(pod, &runtime_base).await.unwrap()
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[tokio::test]
async fn shared_state_starts_idle() {
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
assert_eq!(handle.shared_state.get_status(), PodStatus::Idle);
}
#[tokio::test]
async fn run_updates_shared_state_to_idle_after_completion() {
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
handle
.send(Method::Run {
input: "Hello".into(),
})
.await
.unwrap();
// Wait for the run to complete
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(handle.shared_state.get_status(), PodStatus::Idle);
}
#[tokio::test]
async fn run_populates_history() {
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
handle
.send(Method::Run {
input: "Hello".into(),
})
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let history = handle.shared_state.history_json();
assert_ne!(history, "[]");
let parsed: serde_json::Value = serde_json::from_str(&history).unwrap();
assert!(parsed.is_array());
assert!(parsed.as_array().unwrap().len() >= 2); // user + assistant
}
#[tokio::test]
async fn events_are_broadcast() {
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
let mut rx = handle.subscribe();
handle
.send(Method::Run {
input: "Hello".into(),
})
.await
.unwrap();
let mut saw_turn_start = false;
let mut saw_text_delta = false;
let mut saw_text_done = false;
let mut saw_turn_end = false;
// Collect events with a timeout
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
tokio::select! {
event = rx.recv() => {
match event {
Ok(Event::TurnStart { .. }) => saw_turn_start = true,
Ok(Event::TextDelta { .. }) => saw_text_delta = true,
Ok(Event::TextDone { .. }) => saw_text_done = true,
Ok(Event::TurnEnd { .. }) => {
saw_turn_end = true;
break;
}
Err(_) => break,
_ => {}
}
}
_ = tokio::time::sleep_until(deadline) => break,
}
}
assert!(saw_turn_start, "should see turn_start");
assert!(saw_text_delta, "should see text_delta");
assert!(saw_text_done, "should see text_done");
assert!(saw_turn_end, "should see turn_end");
}
#[tokio::test]
async fn double_run_returns_error() {
// Create a client that streams slowly
let events = vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, "slow..."),
// No stop/completed — the stream will end but without proper completion
];
let client = MockClient::new(events);
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
let mut rx = handle.subscribe();
// Send first run
handle
.send(Method::Run {
input: "first".into(),
})
.await
.unwrap();
// Immediately send second run (should get error)
handle
.send(Method::Run {
input: "second".into(),
})
.await
.unwrap();
// Look for the error event
let mut saw_already_running = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
tokio::select! {
event = rx.recv() => {
match event {
Ok(Event::Error { code, .. }) => {
if code == pod::ErrorCode::AlreadyRunning {
saw_already_running = true;
break;
}
}
Err(_) => break,
_ => {}
}
}
_ = tokio::time::sleep_until(deadline) => break,
}
}
assert!(saw_already_running, "should see already_running error");
}
#[tokio::test]
async fn resume_without_pause_returns_error() {
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
let mut rx = handle.subscribe();
handle.send(Method::Resume).await.unwrap();
let mut saw_not_paused = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1);
loop {
tokio::select! {
event = rx.recv() => {
match event {
Ok(Event::Error { code, .. }) if code == pod::ErrorCode::NotPaused => {
saw_not_paused = true;
break;
}
Err(_) => break,
_ => {}
}
}
_ = tokio::time::sleep_until(deadline) => break,
}
}
assert!(saw_not_paused, "should see not_paused error");
}
#[tokio::test]
async fn cancel_without_run_returns_error() {
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
let mut rx = handle.subscribe();
handle.send(Method::Cancel).await.unwrap();
let mut saw_not_running = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1);
loop {
tokio::select! {
event = rx.recv() => {
match event {
Ok(Event::Error { code, .. }) if code == pod::ErrorCode::NotRunning => {
saw_not_running = true;
break;
}
Err(_) => break,
_ => {}
}
}
_ = tokio::time::sleep_until(deadline) => break,
}
}
assert!(saw_not_running, "should see not_running error");
}
#[tokio::test]
async fn status_json_reflects_pod_name() {
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
let json = handle.shared_state.status_json();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["pod_name"], "test-pod");
}
// ---------------------------------------------------------------------------
// Socket transport tests
// ---------------------------------------------------------------------------
#[tokio::test]
async fn socket_run_receives_events() {
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::UnixStream;
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
// Give the socket server a moment to bind
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let sock_path = handle.runtime_dir.socket_path();
let stream = UnixStream::connect(&sock_path).await.unwrap();
let (reader, mut writer) = stream.into_split();
let mut lines = BufReader::new(reader).lines();
// Send run method via socket
writer
.write_all(b"{\"method\":\"run\",\"params\":{\"input\":\"Hello\"}}\n")
.await
.unwrap();
// Collect events
let mut saw_turn_start = false;
let mut saw_text_delta = false;
let mut saw_turn_end = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
tokio::select! {
line = lines.next_line() => {
match line {
Ok(Some(line)) => {
let parsed: serde_json::Value = serde_json::from_str(&line).unwrap();
match parsed["event"].as_str() {
Some("turn_start") => saw_turn_start = true,
Some("text_delta") => saw_text_delta = true,
Some("turn_end") => {
saw_turn_end = true;
break;
}
_ => {}
}
}
_ => break,
}
}
_ = tokio::time::sleep_until(deadline) => break,
}
}
assert!(saw_turn_start, "should see turn_start via socket");
assert!(saw_text_delta, "should see text_delta via socket");
assert!(saw_turn_end, "should see turn_end via socket");
}
#[tokio::test]
async fn socket_invalid_method_returns_error() {
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::UnixStream;
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let sock_path = handle.runtime_dir.socket_path();
let stream = UnixStream::connect(&sock_path).await.unwrap();
let (reader, mut writer) = stream.into_split();
let mut lines = BufReader::new(reader).lines();
// Send garbage
writer.write_all(b"{\"bad\":\"json\"}\n").await.unwrap();
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1);
let mut saw_error = false;
loop {
tokio::select! {
line = lines.next_line() => {
match line {
Ok(Some(line)) => {
let parsed: serde_json::Value = serde_json::from_str(&line).unwrap();
if parsed["event"] == "error" {
saw_error = true;
break;
}
}
_ => break,
}
}
_ = tokio::time::sleep_until(deadline) => break,
}
}
assert!(saw_error, "should see error for invalid method");
}
+9
View File
@@ -0,0 +1,9 @@
[package]
name = "protocol"
version = "0.1.0"
edition.workspace = true
license.workspace = true
[dependencies]
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
+140
View File
@@ -0,0 +1,140 @@
use serde::{Deserialize, Serialize};
// ---------------------------------------------------------------------------
// Method (Client → Pod via Unix Socket)
// ---------------------------------------------------------------------------
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "method", content = "params", rename_all = "snake_case")]
pub enum Method {
Run { input: String },
Resume,
Cancel,
}
impl Method {
pub fn from_json_line(line: &str) -> Result<Self, serde_json::Error> {
serde_json::from_str(line)
}
}
// ---------------------------------------------------------------------------
// Event (Pod → Client via Unix Socket broadcast)
// ---------------------------------------------------------------------------
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "event", content = "data", rename_all = "snake_case")]
pub enum Event {
TurnStart {
turn: usize,
},
TurnEnd {
turn: usize,
result: TurnResult,
},
TextDelta {
text: String,
},
TextDone {
text: String,
},
ToolCallStart {
id: String,
name: String,
},
ToolCallArgsDelta {
id: String,
json: String,
},
ToolCallDone {
id: String,
name: String,
arguments: String,
},
ToolResult {
id: String,
output: String,
is_error: bool,
},
Usage {
input_tokens: Option<u64>,
output_tokens: Option<u64>,
},
Error {
code: ErrorCode,
message: String,
},
}
impl Event {
pub fn to_json_line(&self) -> Result<String, serde_json::Error> {
serde_json::to_string(self)
}
}
// ---------------------------------------------------------------------------
// Supporting types
// ---------------------------------------------------------------------------
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TurnResult {
Finished,
Paused,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ErrorCode {
AlreadyRunning,
NotRunning,
NotPaused,
ProviderError,
ToolError,
Internal,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn method_run_json_roundtrip() {
let json = r#"{"method":"run","params":{"input":"Hello"}}"#;
let method = Method::from_json_line(json).unwrap();
assert!(matches!(method, Method::Run { ref input } if input == "Hello"));
let serialized = serde_json::to_string(&method).unwrap();
assert_eq!(serialized, json);
}
#[test]
fn method_without_params() {
let json = r#"{"method":"resume"}"#;
let method = Method::from_json_line(json).unwrap();
assert!(matches!(method, Method::Resume));
}
#[test]
fn event_text_delta_format() {
let event = Event::TextDelta {
text: "Hello".into(),
};
let json = event.to_json_line().unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["event"], "text_delta");
assert_eq!(parsed["data"]["text"], "Hello");
}
#[test]
fn event_error_format() {
let event = Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Pod is already executing a turn".into(),
};
let json = event.to_json_line().unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["event"], "error");
assert_eq!(parsed["data"]["code"], "already_running");
}
}
+10
View File
@@ -0,0 +1,10 @@
[package]
name = "provider"
version = "0.1.0"
edition.workspace = true
license.workspace = true
[dependencies]
llm-worker = { version = "0.2.1", path = "../llm-worker" }
manifest = { version = "0.1.0", path = "../manifest" }
thiserror = "2.0"
@@ -4,24 +4,30 @@ use llm_worker::llm_client::providers::gemini::GeminiClient;
use llm_worker::llm_client::providers::ollama::OllamaClient;
use llm_worker::llm_client::providers::openai::OpenAIClient;
use crate::manifest::{ProviderConfig, ProviderKind};
use crate::pod::PodError;
use manifest::{ProviderConfig, ProviderKind};
/// Errors from provider client construction.
#[derive(Debug, thiserror::Error)]
pub enum ProviderError {
#[error("provider configuration error: {0}")]
Config(String),
}
/// Build an [`LlmClient`] from a [`ProviderConfig`].
///
/// Resolves the API key from the environment variable specified in the config.
pub fn build_client(config: &ProviderConfig) -> Result<Box<dyn LlmClient>, PodError> {
pub fn build_client(config: &ProviderConfig) -> Result<Box<dyn LlmClient>, ProviderError> {
let api_key = config
.api_key_env
.as_deref()
.map(std::env::var)
.transpose()
.map_err(|e| PodError::ProviderConfig(format!("env var: {e}")))?;
.map_err(|e| ProviderError::Config(format!("env var: {e}")))?;
match config.kind {
ProviderKind::Anthropic => {
let key = api_key.ok_or_else(|| {
PodError::ProviderConfig("anthropic requires api_key_env".into())
ProviderError::Config("anthropic requires api_key_env".into())
})?;
let mut client = AnthropicClient::new(key, &config.model);
if let Some(ref url) = config.base_url {
@@ -31,7 +37,7 @@ pub fn build_client(config: &ProviderConfig) -> Result<Box<dyn LlmClient>, PodEr
}
ProviderKind::Openai => {
let key = api_key.ok_or_else(|| {
PodError::ProviderConfig("openai requires api_key_env".into())
ProviderError::Config("openai requires api_key_env".into())
})?;
let mut client = OpenAIClient::new(key, &config.model);
if let Some(ref url) = config.base_url {
@@ -41,7 +47,7 @@ pub fn build_client(config: &ProviderConfig) -> Result<Box<dyn LlmClient>, PodEr
}
ProviderKind::Gemini => {
let key = api_key.ok_or_else(|| {
PodError::ProviderConfig("gemini requires api_key_env".into())
ProviderError::Config("gemini requires api_key_env".into())
})?;
let mut client = GeminiClient::new(key, &config.model);
if let Some(ref url) = config.base_url {
+12
View File
@@ -0,0 +1,12 @@
[package]
name = "tui"
version = "0.1.0"
edition.workspace = true
license.workspace = true
[dependencies]
protocol = { path = "../protocol" }
ratatui = "0.29"
crossterm = "0.28"
tokio = { version = "1.49", features = ["rt-multi-thread", "macros", "net", "io-util", "sync", "time"] }
serde_json = "1.0"
+227
View File
@@ -0,0 +1,227 @@
use protocol::{Event, Method};
pub struct App {
pub pod_name: String,
pub connected: bool,
pub messages: Vec<Message>,
pub current_text: String,
pub input: String,
pub cursor: usize,
pub scroll: u16,
pub quit: bool,
}
pub struct Message {
pub kind: MessageKind,
pub content: String,
}
#[derive(Clone, Copy)]
pub enum MessageKind {
User,
Assistant,
Tool,
Error,
Status,
}
impl App {
pub fn new(pod_name: String) -> Self {
Self {
pod_name,
connected: false,
messages: Vec::new(),
current_text: String::new(),
input: String::new(),
cursor: 0,
scroll: 0,
quit: false,
}
}
pub fn submit_input(&mut self) -> Option<Method> {
let text = self.input.trim().to_owned();
if text.is_empty() {
return None;
}
self.messages.push(Message {
kind: MessageKind::User,
content: text.clone(),
});
self.input.clear();
self.cursor = 0;
self.scroll_to_bottom();
Some(Method::Run { input: text })
}
pub fn handle_pod_event(&mut self, event: Event) {
match event {
Event::TurnStart { turn } => {
self.push_status(format!("[turn {turn}] start"));
}
Event::TextDelta { text } => {
self.current_text.push_str(&text);
}
Event::TextDone { .. } => {
let text = std::mem::take(&mut self.current_text);
if !text.is_empty() {
self.messages.push(Message {
kind: MessageKind::Assistant,
content: text,
});
self.scroll_to_bottom();
}
}
Event::TurnEnd { turn, result } => {
// Flush any remaining text delta
if !self.current_text.is_empty() {
let text = std::mem::take(&mut self.current_text);
self.messages.push(Message {
kind: MessageKind::Assistant,
content: text,
});
}
self.push_status(format!("[turn {turn}] end ({result:?})"));
}
Event::ToolCallStart { name, .. } => {
self.messages.push(Message {
kind: MessageKind::Tool,
content: format!("[tool] {name}"),
});
self.scroll_to_bottom();
}
Event::ToolCallDone {
name, arguments, ..
} => {
self.messages.push(Message {
kind: MessageKind::Tool,
content: format!("[tool] {name} done ({} bytes)", arguments.len()),
});
self.scroll_to_bottom();
}
Event::ToolResult {
output, is_error, ..
} => {
let prefix = if is_error { "[tool error]" } else { "[tool result]" };
let display = if output.len() > 200 {
format!("{}...", &output[..200])
} else {
output
};
self.messages.push(Message {
kind: MessageKind::Tool,
content: format!("{prefix} {display}"),
});
self.scroll_to_bottom();
}
Event::Usage {
input_tokens,
output_tokens,
} => {
self.push_status(format!(
"[usage] in={} out={}",
input_tokens.unwrap_or(0),
output_tokens.unwrap_or(0),
));
}
Event::Error { code, message } => {
self.messages.push(Message {
kind: MessageKind::Error,
content: format!("[{code:?}] {message}"),
});
self.scroll_to_bottom();
}
Event::ToolCallArgsDelta { .. } => {}
}
}
pub fn insert_char(&mut self, c: char) {
self.input.insert(self.cursor, c);
self.cursor += c.len_utf8();
}
pub fn delete_char_before(&mut self) {
if self.cursor > 0 {
let prev = self.input[..self.cursor]
.char_indices()
.next_back()
.map(|(i, _)| i)
.unwrap_or(0);
self.input.drain(prev..self.cursor);
self.cursor = prev;
}
}
pub fn delete_char_after(&mut self) {
if self.cursor < self.input.len() {
let next = self.input[self.cursor..]
.char_indices()
.nth(1)
.map(|(i, _)| self.cursor + i)
.unwrap_or(self.input.len());
self.input.drain(self.cursor..next);
}
}
pub fn move_cursor_left(&mut self) {
if self.cursor > 0 {
self.cursor = self.input[..self.cursor]
.char_indices()
.next_back()
.map(|(i, _)| i)
.unwrap_or(0);
}
}
pub fn move_cursor_right(&mut self) {
if self.cursor < self.input.len() {
self.cursor = self.input[self.cursor..]
.char_indices()
.nth(1)
.map(|(i, _)| self.cursor + i)
.unwrap_or(self.input.len());
}
}
pub fn move_cursor_home(&mut self) {
self.cursor = 0;
}
pub fn move_cursor_end(&mut self) {
self.cursor = self.input.len();
}
pub fn scroll_up(&mut self) {
self.scroll = self.scroll.saturating_sub(3);
}
pub fn scroll_down(&mut self) {
self.scroll = self.scroll.saturating_add(3);
}
/// Total visible lines (for rendering the in-progress text as part of output).
pub fn display_lines(&self) -> Vec<(&MessageKind, &str)> {
let mut lines: Vec<(&MessageKind, &str)> = self
.messages
.iter()
.map(|m| (&m.kind, m.content.as_str()))
.collect();
if !self.current_text.is_empty() {
lines.push((&MessageKind::Assistant, &self.current_text));
}
lines
}
fn push_status(&mut self, content: String) {
self.messages.push(Message {
kind: MessageKind::Status,
content,
});
self.scroll_to_bottom();
}
fn scroll_to_bottom(&mut self) {
// Will be clamped during rendering
self.scroll = u16::MAX;
}
}
+50
View File
@@ -0,0 +1,50 @@
use std::io;
use std::path::Path;
use protocol::{Event, Method};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::UnixStream;
use tokio::sync::mpsc;
pub struct PodClient {
writer: tokio::io::WriteHalf<UnixStream>,
event_rx: mpsc::Receiver<Event>,
}
impl PodClient {
pub async fn connect(path: &Path) -> Result<Self, io::Error> {
let stream = UnixStream::connect(path).await?;
let (reader, writer) = tokio::io::split(stream);
let (event_tx, event_rx) = mpsc::channel::<Event>(256);
tokio::spawn(async move {
let mut lines = BufReader::new(reader).lines();
while let Ok(Some(line)) = lines.next_line().await {
if line.is_empty() {
continue;
}
if let Ok(event) = serde_json::from_str::<Event>(&line) {
if event_tx.send(event).await.is_err() {
break;
}
}
}
});
Ok(Self { writer, event_rx })
}
pub async fn send(&mut self, method: &Method) -> Result<(), io::Error> {
let json = serde_json::to_string(method)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
self.writer.write_all(json.as_bytes()).await?;
self.writer.write_all(b"\n").await?;
self.writer.flush().await?;
Ok(())
}
pub async fn next_event(&mut self) -> Option<Event> {
self.event_rx.recv().await
}
}
+217
View File
@@ -0,0 +1,217 @@
mod app;
mod client;
mod ui;
use std::io;
use std::path::PathBuf;
use crossterm::event::{self, Event as TermEvent, KeyCode, KeyEvent, KeyModifiers};
use crossterm::terminal::{self, EnterAlternateScreen, LeaveAlternateScreen};
use crossterm::{execute};
use protocol::Method;
use ratatui::backend::CrosstermBackend;
use ratatui::Terminal;
use crate::app::App;
use crate::client::PodClient;
fn resolve_socket(pod_name: &str, override_path: Option<PathBuf>) -> PathBuf {
if let Some(p) = override_path {
return p;
}
if let Ok(rd) = std::env::var("XDG_RUNTIME_DIR") {
PathBuf::from(rd).join("insomnia").join(pod_name).join("sock")
} else if let Ok(home) = std::env::var("HOME") {
PathBuf::from(home)
.join(".insomnia")
.join("run")
.join(pod_name)
.join("sock")
} else {
PathBuf::from("/tmp")
.join("insomnia")
.join(pod_name)
.join("sock")
}
}
fn parse_args() -> (String, Option<PathBuf>) {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("usage: tui <pod_name> [--socket <path>]");
std::process::exit(1);
}
let pod_name = args[1].clone();
let socket = args
.windows(2)
.find(|w| w[0] == "--socket")
.map(|w| PathBuf::from(&w[1]));
(pod_name, socket)
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let (pod_name, socket_override) = parse_args();
let socket_path = resolve_socket(&pod_name, socket_override);
// Install panic hook to restore terminal
let original_hook = std::panic::take_hook();
std::panic::set_hook(Box::new(move |info| {
let _ = terminal::disable_raw_mode();
let _ = execute!(io::stdout(), LeaveAlternateScreen);
original_hook(info);
}));
// Setup terminal
terminal::enable_raw_mode()?;
let mut stdout = io::stdout();
execute!(stdout, EnterAlternateScreen)?;
let backend = CrosstermBackend::new(stdout);
let mut terminal = Terminal::new(backend)?;
let mut app = App::new(pod_name);
// Connect to pod
match PodClient::connect(&socket_path).await {
Ok(client) => {
app.connected = true;
run_loop(&mut terminal, &mut app, client).await?;
}
Err(e) => {
app.messages.push(app::Message {
kind: app::MessageKind::Error,
content: format!("Failed to connect to {}: {e}", socket_path.display()),
});
// Show error and wait for quit
run_disconnected(&mut terminal, &mut app)?;
}
}
// Restore terminal
terminal::disable_raw_mode()?;
execute!(terminal.backend_mut(), LeaveAlternateScreen)?;
Ok(())
}
async fn run_loop(
terminal: &mut Terminal<CrosstermBackend<io::Stdout>>,
app: &mut App,
mut client: PodClient,
) -> Result<(), Box<dyn std::error::Error>> {
loop {
terminal.draw(|f| ui::draw(f, app))?;
if app.quit {
break;
}
tokio::select! {
// Terminal input
_ = tokio::task::spawn_blocking(|| event::poll(std::time::Duration::from_millis(50))) => {
while event::poll(std::time::Duration::ZERO)? {
if let TermEvent::Key(key) = event::read()? {
if let Some(method) = handle_key(app, key) {
client.send(&method).await?;
}
if app.quit {
break;
}
}
}
}
// Pod events
event = client.next_event() => {
match event {
Some(ev) => app.handle_pod_event(ev),
None => {
app.connected = false;
app.messages.push(app::Message {
kind: app::MessageKind::Error,
content: "Connection lost".into(),
});
}
}
}
}
}
Ok(())
}
fn run_disconnected(
terminal: &mut Terminal<CrosstermBackend<io::Stdout>>,
app: &mut App,
) -> Result<(), Box<dyn std::error::Error>> {
loop {
terminal.draw(|f| ui::draw(f, app))?;
if event::poll(std::time::Duration::from_millis(100))? {
if let TermEvent::Key(key) = event::read()? {
match key.code {
KeyCode::Esc => break,
KeyCode::Char('c') if key.modifiers.contains(KeyModifiers::CONTROL) => break,
_ => {}
}
}
}
}
Ok(())
}
fn handle_key(app: &mut App, key: KeyEvent) -> Option<Method> {
match key.code {
KeyCode::Esc => {
app.quit = true;
None
}
KeyCode::Char('c') if key.modifiers.contains(KeyModifiers::CONTROL) => {
app.quit = true;
None
}
KeyCode::Char('r') if key.modifiers.contains(KeyModifiers::CONTROL) => {
Some(Method::Resume)
}
KeyCode::Char('x') if key.modifiers.contains(KeyModifiers::CONTROL) => {
Some(Method::Cancel)
}
KeyCode::Enter => app.submit_input(),
KeyCode::Backspace => {
app.delete_char_before();
None
}
KeyCode::Delete => {
app.delete_char_after();
None
}
KeyCode::Left => {
app.move_cursor_left();
None
}
KeyCode::Right => {
app.move_cursor_right();
None
}
KeyCode::Home => {
app.move_cursor_home();
None
}
KeyCode::End => {
app.move_cursor_end();
None
}
KeyCode::PageUp => {
app.scroll_up();
None
}
KeyCode::PageDown => {
app.scroll_down();
None
}
KeyCode::Char(c) => {
app.insert_char(c);
None
}
_ => None,
}
}
+90
View File
@@ -0,0 +1,90 @@
use ratatui::layout::{Constraint, Layout, Position};
use ratatui::style::{Color, Modifier, Style};
use ratatui::text::{Line, Span};
use ratatui::widgets::{Block, Borders, Paragraph, Wrap};
use ratatui::Frame;
use crate::app::{App, MessageKind};
pub fn draw(frame: &mut Frame, app: &mut App) {
let chunks = Layout::vertical([
Constraint::Length(1),
Constraint::Min(3),
Constraint::Length(3),
])
.split(frame.area());
draw_status_bar(frame, app, chunks[0]);
draw_output(frame, app, chunks[1]);
draw_input(frame, app, chunks[2]);
}
fn draw_status_bar(frame: &mut Frame, app: &App, area: ratatui::layout::Rect) {
let conn_style = if app.connected {
Style::default().fg(Color::Green)
} else {
Style::default().fg(Color::Red)
};
let conn_text = if app.connected {
"connected"
} else {
"disconnected"
};
let line = Line::from(vec![
Span::styled(&app.pod_name, Style::default().add_modifier(Modifier::BOLD)),
Span::raw(" | "),
Span::styled(conn_text, conn_style),
]);
frame.render_widget(Paragraph::new(line), area);
}
fn draw_output(frame: &mut Frame, app: &mut App, area: ratatui::layout::Rect) {
let display = app.display_lines();
let lines: Vec<Line> = display
.iter()
.flat_map(|(kind, content)| {
let style = kind_style(kind);
content.lines().map(move |l| Line::from(Span::styled(l.to_owned(), style)))
})
.collect();
let total = lines.len() as u16;
let visible = area.height.saturating_sub(2); // block borders
let max_scroll = total.saturating_sub(visible);
if app.scroll > max_scroll {
app.scroll = max_scroll;
}
let block = Block::default().borders(Borders::ALL);
let paragraph = Paragraph::new(lines)
.block(block)
.wrap(Wrap { trim: false })
.scroll((app.scroll, 0));
frame.render_widget(paragraph, area);
}
fn draw_input(frame: &mut Frame, app: &App, area: ratatui::layout::Rect) {
let display = format!("> {}", app.input);
let block = Block::default().borders(Borders::ALL).title("Input");
let paragraph = Paragraph::new(display).block(block);
frame.render_widget(paragraph, area);
// Cursor position: "> " is 2 chars, plus cursor offset in the input
let cursor_x = area.x + 1 + 2 + app.input[..app.cursor].chars().count() as u16;
let cursor_y = area.y + 1;
frame.set_cursor_position(Position::new(cursor_x, cursor_y));
}
fn kind_style(kind: &MessageKind) -> Style {
match kind {
MessageKind::User => Style::default().fg(Color::Green),
MessageKind::Assistant => Style::default().fg(Color::White),
MessageKind::Tool => Style::default().fg(Color::Cyan),
MessageKind::Error => Style::default().fg(Color::Red),
MessageKind::Status => Style::default().fg(Color::DarkGray),
}
}