cratesの整理
This commit is contained in:
@@ -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"] }
|
||||
@@ -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" }
|
||||
@@ -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;
|
||||
@@ -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"] }
|
||||
@@ -1,3 +0,0 @@
|
||||
fn main() {
|
||||
println!("insomnia");
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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#"
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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(),
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
@@ -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"
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
@@ -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"
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user