Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

feat: add peer pod handshake command

Browse Source
This commit is contained in:
Hare
2026-06-02 19:43:20 +09:00
parent b51c32401c
commit 093a84fc83
12 changed files with 663 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/pod/src/controller.rs
+25 -4
View File
@@ -8,7 +8,7 @@ use pod_store::PodMetadataStore;
use session_store::Store;
use tokio::sync::{broadcast, mpsc, oneshot};
use crate::discovery::{PodDiscovery, list_pods_tool, restore_pod_tool};
use crate::discovery::{PodDiscovery, list_pods_tool, restore_pod_tool, send_to_peer_pod_tool};
use crate::ipc::alerter::Alerter;
use crate::ipc::notify_buffer::NotifyBuffer;
use crate::ipc::server::SocketServer;
@@ -561,7 +561,8 @@ where
worker.register_tool(stop_pod_tool(spawned_registry.clone()));
let discovery = PodDiscovery::new(pod_store, spawner_name, runtime_base, pwd, spawned_registry);
worker.register_tool(list_pods_tool(discovery.clone()));
worker.register_tool(restore_pod_tool(discovery));
worker.register_tool(restore_pod_tool(discovery.clone()));
worker.register_tool(send_to_peer_pod_tool(discovery));
pod.attach_tracker(tracker);
fs_for_view
}
@@ -862,6 +863,26 @@ async fn controller_loop<C, St>(
}
},
Method::RegisterPeer { name } => match discovery.register_peer(&name) {
Ok(result) => match serde_json::to_value(result) {
Ok(result) => {
let _ = event_tx.send(Event::PeerRegistered { result });
}
Err(error) => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: format!("serialize peer registration result: {error}"),
});
}
},
Err(error) => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
}
},
// ListCompletions is handled at the socket layer (direct
// response). If it reaches the controller, ignore it.
Method::ListCompletions { .. } => {}
@@ -1063,10 +1084,10 @@ where
notify_buffer.push_notify(message);
}
Some(Method::ListCompletions { .. }) => {}
Some(Method::ListPods | Method::RestorePod { .. }) => {
Some(Method::ListPods | Method::RestorePod { .. } | Method::RegisterPeer { .. }) => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Pod discovery requests are only handled while the Pod is idle or paused"
message: "Pod discovery/control requests are only handled while the Pod is idle or paused"
.into(),
});
}
crates/pod/src/discovery.rs
+333 -8
View File
@@ -17,9 +17,9 @@ use async_trait::async_trait;
use client::PodRuntimeCommand;
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use manifest::{Permission, ScopeRule};
use pod_store::{PodActiveSegmentRef, PodMetadata, PodMetadataStore};
use protocol::stream::JsonLineReader;
use protocol::{Event, PodStatus};
use pod_store::{PodActiveSegmentRef, PodMetadata, PodMetadataStore, validate_pod_name};
use protocol::stream::{JsonLineReader, JsonLineWriter};
use protocol::{Event, Method, PodStatus};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use session_store::{SegmentId, SessionId};
@@ -172,6 +172,41 @@ where
})
}
pub fn register_peer(
&self,
peer_name: &str,
) -> Result<PeerRegistrationResult, PodDiscoveryError> {
validate_pod_name(peer_name)?;
if peer_name == self.self_pod_name {
return Err(PodDiscoveryError::SelfPeer {
pod_name: peer_name.to_string(),
});
}
let self_exists = self.store.read_by_name(&self.self_pod_name)?.is_some();
if !self_exists {
return Err(PodDiscoveryError::StateMissing {
pod_name: self.self_pod_name.clone(),
});
}
let peer_exists = self.store.read_by_name(peer_name)?.is_some();
if !peer_exists {
return Err(PodDiscoveryError::MissingPod {
pod_name: peer_name.to_string(),
});
}
self.store.add_peer(&self.self_pod_name, peer_name)?;
if let Err(error) = self.store.add_peer(peer_name, &self.self_pod_name) {
let _ = self.store.remove_peer(&self.self_pod_name, peer_name);
return Err(PodDiscoveryError::PodStore(error));
}
Ok(PeerRegistrationResult {
source: self.self_pod_name.clone(),
peer: peer_name.to_string(),
})
}
async fn visibility(&self) -> Result<VisibilitySet, PodDiscoveryError> {
let mut visible = BTreeMap::new();
let mut child_sockets = BTreeMap::new();
@@ -187,6 +222,11 @@ where
child_sockets.insert(child.pod_name.clone(), child.socket_path.clone());
comm_registry.insert(child.pod_name.clone(), comm_info_from_spawned_child(&child));
}
for peer in metadata.peers {
visible
.entry(peer.pod_name)
.or_insert(VisibilityReason::Peer);
}
}
// The live in-memory registry covers just-spawned children even if a
@@ -379,6 +419,7 @@ where
pub enum VisibilityReason {
SelfPod,
SpawnedChild,
Peer,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
@@ -525,6 +566,12 @@ pub enum RestoreResult {
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct PeerRegistrationResult {
pub source: String,
pub peer: String,
}
#[derive(Debug, thiserror::Error)]
pub enum PodDiscoveryError {
#[error("pod state missing for `{pod_name}`")]
@@ -533,6 +580,10 @@ pub enum PodDiscoveryError {
NotVisible { pod_name: String },
#[error("pod `{pod_name}` is not restorable: {reason}")]
NotRestorable { pod_name: String, reason: String },
#[error("pod `{pod_name}` cannot be registered as a peer of itself")]
SelfPeer { pod_name: String },
#[error("pod `{pod_name}` does not exist")]
MissingPod { pod_name: String },
#[error(
"pod `{pod_name}` segment {segment_id} is locked by `{owner_pod}` pid {pid} at {socket_path}"
)]
@@ -683,6 +734,14 @@ struct PodNameInput {
name: String,
}
#[derive(Debug, Deserialize, JsonSchema)]
struct SendToPeerPodInput {
/// Target peer Pod name.
name: String,
/// Text delivered to the peer as a peer notification.
message: String,
}
struct ListPodsTool<St> {
discovery: PodDiscovery<St>,
}
@@ -776,6 +835,83 @@ where
})
}
const SEND_TO_PEER_POD_DESCRIPTION: &str = "Send a text message to a peer Pod made visible by an explicit peer handshake. The message is delivered as a peer notification through the target Pod's durable notification/history path. This does not grant delegated scope, create a spawned-child output cursor, imply parent ownership, or produce child completion notifications. Fails if the target is not a visible live peer.";
struct SendToPeerPodTool<St> {
discovery: PodDiscovery<St>,
}
#[async_trait]
impl<St> Tool for SendToPeerPodTool<St>
where
St: PodMetadataStore + Clone + Send + Sync + 'static,
{
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
let input: SendToPeerPodInput = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid SendToPeerPod input: {e}")))?;
let detail = self
.discovery
.inspect(&input.name)
.await
.map_err(discovery_error_to_tool_error)?;
if detail.visibility != VisibilityReason::Peer {
return Err(ToolError::InvalidArgument(format!(
"pod `{}` is visible as {:?}, not as a peer",
input.name, detail.visibility
)));
}
if !detail.live.reachable {
return Err(ToolError::ExecutionFailed(format!(
"peer pod `{}` is not live/reachable; restore it before sending",
input.name
)));
}
let message = format!(
"[Peer message from `{}`]\n{}",
self.discovery.self_pod_name, input.message
);
send_peer_notify(&detail.live.socket_path, message)
.await
.map_err(|error| {
ToolError::ExecutionFailed(format!("send to peer `{}`: {error}", input.name))
})?;
Ok(ToolOutput {
summary: format!("sent peer message to `{}`", input.name),
content: None,
})
}
}
pub fn send_to_peer_pod_tool<St>(discovery: PodDiscovery<St>) -> ToolDefinition
where
St: PodMetadataStore + Clone + Send + Sync + 'static,
{
Arc::new(move || {
let meta = ToolMeta::new("SendToPeerPod")
.description(SEND_TO_PEER_POD_DESCRIPTION)
.input_schema(serde_json::to_value(schemars::schema_for!(SendToPeerPodInput)).unwrap());
let tool: Arc<dyn Tool> = Arc::new(SendToPeerPodTool {
discovery: discovery.clone(),
});
(meta, tool)
})
}
async fn send_peer_notify(socket_path: &Path, message: String) -> io::Result<()> {
let stream = tokio::time::timeout(Duration::from_secs(5), UnixStream::connect(socket_path))
.await
.map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "connect timed out"))??;
let mut writer = JsonLineWriter::new(stream);
tokio::time::timeout(
Duration::from_secs(5),
writer.write(&Method::Notify { message }),
)
.await
.map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "write timed out"))?
}
fn json_content<T: Serialize>(value: &T) -> Result<String, ToolError> {
serde_json::to_string_pretty(value)
.map_err(|e| ToolError::Internal(format!("serialize pod discovery output: {e}")))
@@ -785,7 +921,9 @@ fn discovery_error_to_tool_error(error: PodDiscoveryError) -> ToolError {
match error {
PodDiscoveryError::StateMissing { .. }
| PodDiscoveryError::NotVisible { .. }
| PodDiscoveryError::NotRestorable { .. } => ToolError::InvalidArgument(error.to_string()),
| PodDiscoveryError::NotRestorable { .. }
| PodDiscoveryError::SelfPeer { .. }
| PodDiscoveryError::MissingPod { .. } => ToolError::InvalidArgument(error.to_string()),
PodDiscoveryError::LockConflict { .. }
| PodDiscoveryError::Store(_)
| PodDiscoveryError::PodStore(_)
@@ -805,7 +943,7 @@ mod tests {
use manifest::{Permission, ScopeRule};
use pod_store::{FsPodStore, PodSpawnedChild, PodSpawnedScopeRule};
use protocol::stream::JsonLineWriter;
use protocol::{Alert, AlertLevel, AlertSource, Greeting};
use protocol::{Alert, AlertLevel, AlertSource};
use session_store::{new_segment_id, new_session_id};
use tempfile::TempDir;
use tokio::net::UnixListener;
@@ -844,6 +982,9 @@ mod tests {
child("child-pending", &pending_socket),
],
reclaimed_children: Vec::new(),
peers: vec![pod_store::PodPeer {
pod_name: "peer".into(),
}],
resolved_manifest_snapshot: None,
};
store.write(&parent).unwrap();
@@ -856,6 +997,7 @@ mod tests {
)),
spawned_children: Vec::new(),
reclaimed_children: Vec::new(),
peers: Vec::new(),
resolved_manifest_snapshot: None,
})
.unwrap();
@@ -868,6 +1010,7 @@ mod tests {
)),
spawned_children: Vec::new(),
reclaimed_children: Vec::new(),
peers: Vec::new(),
resolved_manifest_snapshot: None,
})
.unwrap();
@@ -877,6 +1020,7 @@ mod tests {
active: Some(PodActiveSegmentRef::pending_segment(pending_session_id)),
spawned_children: Vec::new(),
reclaimed_children: Vec::new(),
peers: Vec::new(),
resolved_manifest_snapshot: None,
})
.unwrap();
@@ -889,6 +1033,19 @@ mod tests {
)),
spawned_children: Vec::new(),
reclaimed_children: Vec::new(),
peers: Vec::new(),
resolved_manifest_snapshot: None,
})
.unwrap();
store
.write(&PodMetadata {
pod_name: "peer".into(),
active: None,
spawned_children: Vec::new(),
reclaimed_children: Vec::new(),
peers: vec![pod_store::PodPeer {
pod_name: "parent".into(),
}],
resolved_manifest_snapshot: None,
})
.unwrap();
@@ -913,14 +1070,37 @@ mod tests {
let restore_tool_def = restore_pod_tool(discovery.clone());
let (restore_meta, _) = restore_tool_def();
assert_eq!(restore_meta.name, "RestorePod");
let send_peer_tool_def = send_to_peer_pod_tool(discovery.clone());
let (send_peer_meta, _) = send_peer_tool_def();
assert_eq!(send_peer_meta.name, "SendToPeerPod");
let list = discovery.list_visible().await.unwrap();
let names: Vec<_> = list.iter().map(|p| p.pod_name.as_str()).collect();
assert_eq!(
names,
vec!["child-live", "child-pending", "child-stale", "parent"]
vec![
"child-live",
"child-pending",
"child-stale",
"parent",
"peer"
]
);
assert!(!names.contains(&"hidden"));
assert_eq!(
list.iter()
.find(|p| p.pod_name == "peer")
.unwrap()
.visibility,
VisibilityReason::Peer
);
assert_eq!(
list.iter()
.find(|p| p.pod_name == "child-live")
.unwrap()
.visibility,
VisibilityReason::SpawnedChild
);
assert!(
list.iter()
.find(|p| p.pod_name == "child-live")
@@ -983,6 +1163,151 @@ mod tests {
live_listener.abort();
}
#[tokio::test(flavor = "current_thread")]
async fn register_peer_persists_reciprocal_metadata() {
let root = TempDir::new().unwrap();
let store_dir = root.path().join("store");
let runtime_base = root.path().join("runtime");
std::fs::create_dir_all(&runtime_base).unwrap();
let store = FsPodStore::new(&store_dir).unwrap();
store.write(&PodMetadata::new("source", None)).unwrap();
store.write(&PodMetadata::new("target", None)).unwrap();
let runtime_dir = Arc::new(RuntimeDir::create(&runtime_base, "source").await.unwrap());
let discovery = PodDiscovery::new(
store.clone(),
"source".into(),
runtime_base.clone(),
root.path().to_path_buf(),
SpawnedPodRegistry::new(runtime_dir),
);
let result = discovery.register_peer("target").unwrap();
assert_eq!(result.source, "source");
assert_eq!(result.peer, "target");
let source = store.read_by_name("source").unwrap().unwrap();
let target = store.read_by_name("target").unwrap().unwrap();
assert_eq!(source.peers[0].pod_name, "target");
assert_eq!(target.peers[0].pod_name, "source");
let list = discovery.list_visible().await.unwrap();
assert_eq!(
list.iter()
.find(|item| item.pod_name == "target")
.unwrap()
.visibility,
VisibilityReason::Peer
);
}
#[tokio::test(flavor = "current_thread")]
async fn register_peer_rejects_self_and_missing_target() {
let root = TempDir::new().unwrap();
let store_dir = root.path().join("store");
let runtime_base = root.path().join("runtime");
std::fs::create_dir_all(&runtime_base).unwrap();
let store = FsPodStore::new(&store_dir).unwrap();
store.write(&PodMetadata::new("source", None)).unwrap();
let runtime_dir = Arc::new(RuntimeDir::create(&runtime_base, "source").await.unwrap());
let discovery = PodDiscovery::new(
store,
"source".into(),
runtime_base,
root.path().to_path_buf(),
SpawnedPodRegistry::new(runtime_dir),
);
let self_err = discovery.register_peer("source").unwrap_err();
assert!(matches!(self_err, PodDiscoveryError::SelfPeer { .. }));
let missing_err = discovery.register_peer("missing").unwrap_err();
assert!(matches!(missing_err, PodDiscoveryError::MissingPod { .. }));
}
#[tokio::test(flavor = "current_thread")]
async fn send_to_peer_pod_delivers_notify_without_child_registry() {
let root = TempDir::new().unwrap();
let store_dir = root.path().join("store");
let runtime_base = root.path().join("runtime");
std::fs::create_dir_all(runtime_base.join("target")).unwrap();
let store = FsPodStore::new(&store_dir).unwrap();
store
.write(&PodMetadata {
pod_name: "source".into(),
active: None,
spawned_children: Vec::new(),
reclaimed_children: Vec::new(),
peers: vec![pod_store::PodPeer {
pod_name: "target".into(),
}],
resolved_manifest_snapshot: None,
})
.unwrap();
store
.write(&PodMetadata {
pod_name: "target".into(),
active: None,
spawned_children: Vec::new(),
reclaimed_children: Vec::new(),
peers: vec![pod_store::PodPeer {
pod_name: "source".into(),
}],
resolved_manifest_snapshot: None,
})
.unwrap();
let runtime_dir = Arc::new(RuntimeDir::create(&runtime_base, "source").await.unwrap());
let discovery = PodDiscovery::new(
store,
"source".into(),
runtime_base.clone(),
root.path().to_path_buf(),
SpawnedPodRegistry::new(runtime_dir),
);
let socket = runtime_base.join("target").join("sock");
let listener = UnixListener::bind(&socket).unwrap();
let (tx, mut rx) = tokio::sync::mpsc::channel(1);
let target = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut writer = JsonLineWriter::new(stream);
writer
.write(&Event::Snapshot {
entries: Vec::new(),
greeting: protocol::Greeting {
pod_name: "target".into(),
cwd: "/tmp".into(),
provider: "test".into(),
model: "test".into(),
scope_summary: String::new(),
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
},
status: PodStatus::Idle,
})
.await
.unwrap();
let (stream, _) = listener.accept().await.unwrap();
let mut reader = JsonLineReader::new(stream);
let method = reader.next::<Method>().await.unwrap().unwrap();
if let Method::Notify { message } = method {
tx.send(message).await.unwrap();
} else {
panic!("expected Notify, got {method:?}");
}
});
let (_, tool) = send_to_peer_pod_tool(discovery)();
let output = tool
.execute(r#"{"name":"target","message":"hello"}"#)
.await
.unwrap();
assert_eq!(output.summary, "sent peer message to `target`");
let message = rx.recv().await.unwrap();
assert_eq!(message, "[Peer message from `source`]\nhello");
target.await.unwrap();
}
#[tokio::test(flavor = "current_thread")]
async fn probe_socket_reads_status_after_replayed_alert() {
let root = TempDir::new().unwrap();
@@ -1003,7 +1328,7 @@ mod tests {
writer
.write(&Event::Snapshot {
entries: Vec::new(),
greeting: Greeting {
greeting: protocol::Greeting {
pod_name: "alerted".into(),
cwd: "/tmp".into(),
provider: "test".into(),
@@ -1051,7 +1376,7 @@ mod tests {
let _ = writer
.write(&Event::Snapshot {
entries: Vec::new(),
greeting: Greeting {
greeting: protocol::Greeting {
pod_name: "child-live".into(),
cwd: "/tmp".into(),
provider: "test".into(),