cratesの整理
This commit is contained in:
@@ -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"] }
|
||||
@@ -0,0 +1,69 @@
|
||||
//! Minimal example: Pod running a single prompt with persistence.
|
||||
//!
|
||||
//! Demonstrates the core insomnia abstraction — a TOML manifest drives
|
||||
//! provider selection, model config, and system prompt, while FsStore
|
||||
//! persists the session to disk automatically.
|
||||
//!
|
||||
//! ## Usage
|
||||
//!
|
||||
//! ```bash
|
||||
//! echo "ANTHROPIC_API_KEY=your-key" > .env
|
||||
//! cargo run -p pod --example pod_cli
|
||||
//! ```
|
||||
|
||||
use pod::{Pod, PodManifest, PodRunResult};
|
||||
use llm_worker_persistence::FsStore;
|
||||
|
||||
const MANIFEST_TOML: &str = r#"
|
||||
[pod]
|
||||
name = "hello-pod"
|
||||
|
||||
[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();
|
||||
|
||||
// 1. Parse the manifest
|
||||
let manifest = PodManifest::from_toml(MANIFEST_TOML)?;
|
||||
println!("Pod: {}", manifest.pod.name);
|
||||
|
||||
// 2. Create a persistent store (temp dir for demo)
|
||||
let tmp = tempfile::tempdir()?;
|
||||
let store = FsStore::new(tmp.path()).await?;
|
||||
|
||||
// 3. Build the Pod from manifest
|
||||
let mut pod = Pod::from_manifest(manifest, store, None).await?;
|
||||
println!("Session: {}", pod.session_id());
|
||||
|
||||
// 4. Run a prompt
|
||||
let result = pod.run("What is the capital of France?").await?;
|
||||
match result {
|
||||
PodRunResult::Finished => println!("(finished)"),
|
||||
PodRunResult::Paused => println!("(paused)"),
|
||||
}
|
||||
|
||||
// 5. Extract the assistant's reply from history
|
||||
let history = pod.session_mut().worker.history();
|
||||
if let Some(text) = history
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|item| item.is_assistant_message())
|
||||
.and_then(|item| item.as_text())
|
||||
{
|
||||
println!("\nAssistant: {text}");
|
||||
}
|
||||
|
||||
// 6. Session ID for potential restore
|
||||
println!("\nSession ID: {}", pod.session_id());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -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;
|
||||
@@ -0,0 +1,179 @@
|
||||
use llm_worker::llm_client::client::LlmClient;
|
||||
use llm_worker::llm_client::RequestConfig;
|
||||
use llm_worker::Worker;
|
||||
use llm_worker_persistence::{
|
||||
Session, SessionConfig, SessionError, SessionId, Store, StoreError,
|
||||
};
|
||||
|
||||
use manifest::{PodManifest, Scope, WorkerManifest};
|
||||
|
||||
/// Pod identifier. UUID v7 (time-ordered).
|
||||
pub type PodId = uuid::Uuid;
|
||||
|
||||
/// Generate a new Pod ID.
|
||||
pub fn new_pod_id() -> PodId {
|
||||
uuid::Uuid::now_v7()
|
||||
}
|
||||
|
||||
/// An independent agent execution unit.
|
||||
///
|
||||
/// Wraps a persistent [`Session`] with manifest metadata and an optional
|
||||
/// directory scope. This is the primary abstraction in insomnia.
|
||||
pub struct Pod<C: LlmClient, St: Store> {
|
||||
id: PodId,
|
||||
manifest: PodManifest,
|
||||
session: Session<C, St>,
|
||||
scope: Option<Scope>,
|
||||
}
|
||||
|
||||
impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
/// Create a new Pod from a pre-built Worker and store.
|
||||
///
|
||||
/// The caller is responsible for constructing the `LlmClient` from the
|
||||
/// manifest's provider config. This keeps Pod free of provider-specific
|
||||
/// dependencies.
|
||||
pub async fn new(
|
||||
manifest: PodManifest,
|
||||
worker: Worker<C>,
|
||||
store: St,
|
||||
scope: Option<Scope>,
|
||||
) -> Result<Self, PodError> {
|
||||
let session = Session::new(worker, store, SessionConfig::default()).await?;
|
||||
Ok(Self {
|
||||
id: new_pod_id(),
|
||||
manifest,
|
||||
session,
|
||||
scope,
|
||||
})
|
||||
}
|
||||
|
||||
/// Restore a Pod from a persisted session.
|
||||
pub async fn restore(
|
||||
id: PodId,
|
||||
session_id: SessionId,
|
||||
manifest: PodManifest,
|
||||
client: C,
|
||||
store: St,
|
||||
scope: Option<Scope>,
|
||||
) -> Result<Self, PodError> {
|
||||
let session = Session::restore(client, store, session_id, SessionConfig::default()).await?;
|
||||
Ok(Self {
|
||||
id,
|
||||
manifest,
|
||||
session,
|
||||
scope,
|
||||
})
|
||||
}
|
||||
|
||||
/// The Pod's unique identifier.
|
||||
pub fn id(&self) -> PodId {
|
||||
self.id
|
||||
}
|
||||
|
||||
/// The session ID used for persistence.
|
||||
pub fn session_id(&self) -> SessionId {
|
||||
self.session.session_id()
|
||||
}
|
||||
|
||||
/// The Pod's manifest.
|
||||
pub fn manifest(&self) -> &PodManifest {
|
||||
&self.manifest
|
||||
}
|
||||
|
||||
/// The Pod's directory scope, if any.
|
||||
pub fn scope(&self) -> Option<&Scope> {
|
||||
self.scope.as_ref()
|
||||
}
|
||||
|
||||
/// Direct access to the underlying session.
|
||||
///
|
||||
/// Use this to register tools, hooks, or subscribers on the worker
|
||||
/// before calling [`run`](Self::run).
|
||||
pub fn session_mut(&mut self) -> &mut Session<C, St> {
|
||||
&mut self.session
|
||||
}
|
||||
|
||||
/// Send user input and run until the LLM turn completes.
|
||||
pub async fn run(&mut self, input: impl Into<String>) -> Result<PodRunResult, PodError> {
|
||||
let result = self.session.run(input).await?;
|
||||
Ok(result.into())
|
||||
}
|
||||
|
||||
/// Resume from a paused state.
|
||||
pub async fn resume(&mut self) -> Result<PodRunResult, PodError> {
|
||||
let result = self.session.resume().await?;
|
||||
Ok(result.into())
|
||||
}
|
||||
}
|
||||
|
||||
impl<St: Store> Pod<Box<dyn LlmClient>, St> {
|
||||
/// Create a Pod entirely from a manifest.
|
||||
///
|
||||
/// Builds the LLM client from the provider config, applies worker
|
||||
/// settings, and creates a new persistent session.
|
||||
pub async fn from_manifest(
|
||||
manifest: PodManifest,
|
||||
store: St,
|
||||
scope: Option<Scope>,
|
||||
) -> Result<Self, PodError> {
|
||||
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?;
|
||||
Ok(Self {
|
||||
id: new_pod_id(),
|
||||
manifest,
|
||||
session,
|
||||
scope,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply worker-level manifest settings to a Worker.
|
||||
pub fn apply_worker_manifest<C: LlmClient>(worker: &mut Worker<C>, wm: &WorkerManifest) {
|
||||
if let Some(ref prompt) = wm.system_prompt {
|
||||
worker.set_system_prompt(prompt);
|
||||
}
|
||||
let mut config = RequestConfig::new();
|
||||
if let Some(max_tokens) = wm.max_tokens {
|
||||
config.max_tokens = Some(max_tokens);
|
||||
}
|
||||
if let Some(temperature) = wm.temperature {
|
||||
config.temperature = Some(temperature);
|
||||
}
|
||||
worker.set_request_config(config);
|
||||
}
|
||||
|
||||
/// Result of a Pod run.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum PodRunResult {
|
||||
/// The LLM finished its turn normally.
|
||||
Finished,
|
||||
/// The LLM paused (e.g. awaiting user confirmation via a hook).
|
||||
Paused,
|
||||
}
|
||||
|
||||
impl From<llm_worker::WorkerResult> for PodRunResult {
|
||||
fn from(r: llm_worker::WorkerResult) -> Self {
|
||||
match r {
|
||||
llm_worker::WorkerResult::Finished => PodRunResult::Finished,
|
||||
llm_worker::WorkerResult::Paused => PodRunResult::Paused,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Pod errors.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum PodError {
|
||||
#[error(transparent)]
|
||||
Session(#[from] SessionError),
|
||||
|
||||
#[error(transparent)]
|
||||
Store(#[from] StoreError),
|
||||
|
||||
#[error("scope violation: {path} is outside the allowed directory")]
|
||||
ScopeViolation { path: 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");
|
||||
}
|
||||
Reference in New Issue
Block a user