worker: route sub-worker tools through internal sessions
This commit is contained in:
@@ -62,15 +62,23 @@ impl Tool for SubWorkerListTool {
|
|||||||
let _input: SubWorkerListInput = serde_json::from_str(input_json).map_err(|error| {
|
let _input: SubWorkerListInput = serde_json::from_str(input_json).map_err(|error| {
|
||||||
ToolError::InvalidArgument(format!("invalid SubWorkerList input: {error}"))
|
ToolError::InvalidArgument(format!("invalid SubWorkerList input: {error}"))
|
||||||
})?;
|
})?;
|
||||||
let items = self
|
let mut items = self
|
||||||
.registry
|
.registry
|
||||||
.list()
|
.list_internal()
|
||||||
.await
|
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|record| SubWorkerListItem {
|
.map(|record| SubWorkerListItem {
|
||||||
name: record.worker_name,
|
name: record.worker_name,
|
||||||
})
|
})
|
||||||
.collect::<Vec<_>>();
|
.collect::<Vec<_>>();
|
||||||
|
items.extend(
|
||||||
|
self.registry
|
||||||
|
.list()
|
||||||
|
.await
|
||||||
|
.into_iter()
|
||||||
|
.map(|record| SubWorkerListItem {
|
||||||
|
name: record.worker_name,
|
||||||
|
}),
|
||||||
|
);
|
||||||
let count = items.len();
|
let count = items.len();
|
||||||
let content = serde_json::to_string_pretty(&serde_json::json!({ "sub_workers": items }))
|
let content = serde_json::to_string_pretty(&serde_json::json!({ "sub_workers": items }))
|
||||||
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
|
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
|
||||||
@@ -125,6 +133,15 @@ impl Tool for SubWorkerSendTool {
|
|||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let input: SubWorkerSendInput = serde_json::from_str(input_json)
|
let input: SubWorkerSendInput = serde_json::from_str(input_json)
|
||||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid SubWorkerSend input: {e}")))?;
|
.map_err(|e| ToolError::InvalidArgument(format!("invalid SubWorkerSend input: {e}")))?;
|
||||||
|
if let Some(record) = self.registry.get_internal(&input.name) {
|
||||||
|
record.session.send(input.message).await.map_err(|error| {
|
||||||
|
ToolError::ExecutionFailed(format!("send to `{}`: {error}", input.name))
|
||||||
|
})?;
|
||||||
|
return Ok(ToolOutput {
|
||||||
|
summary: format!("sent message to `{}`", input.name),
|
||||||
|
content: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
let record = self
|
let record = self
|
||||||
.registry
|
.registry
|
||||||
.get(&input.name)
|
.get(&input.name)
|
||||||
@@ -191,6 +208,35 @@ impl Tool for SubWorkerReadOutputTool {
|
|||||||
let input: NameInput = serde_json::from_str(input_json).map_err(|e| {
|
let input: NameInput = serde_json::from_str(input_json).map_err(|e| {
|
||||||
ToolError::InvalidArgument(format!("invalid SubWorkerReadOutput input: {e}"))
|
ToolError::InvalidArgument(format!("invalid SubWorkerReadOutput input: {e}"))
|
||||||
})?;
|
})?;
|
||||||
|
if let Some(record) = self.registry.get_internal(&input.name) {
|
||||||
|
let entries = record.session.entries();
|
||||||
|
let cursor = self.registry.cursor(&input.name).await;
|
||||||
|
let new_entries = if cursor >= entries.len() {
|
||||||
|
&[] as &[LogEntry]
|
||||||
|
} else {
|
||||||
|
&entries[cursor..]
|
||||||
|
};
|
||||||
|
let values = new_entries
|
||||||
|
.iter()
|
||||||
|
.filter_map(|entry| serde_json::to_value(entry).ok())
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let new_text = extract_assistant_text(&values);
|
||||||
|
self.registry.set_cursor(&input.name, entries.len()).await;
|
||||||
|
let status = format!("{:?}", record.session.status()).to_lowercase();
|
||||||
|
let summary = if new_text.is_empty() {
|
||||||
|
format!("worker `{}` {status}; no new assistant text", input.name)
|
||||||
|
} else {
|
||||||
|
format!(
|
||||||
|
"worker `{}` {status}: {} new line(s) of assistant text",
|
||||||
|
input.name,
|
||||||
|
new_text.lines().count()
|
||||||
|
)
|
||||||
|
};
|
||||||
|
return Ok(ToolOutput {
|
||||||
|
summary,
|
||||||
|
content: (!new_text.is_empty()).then_some(new_text),
|
||||||
|
});
|
||||||
|
}
|
||||||
let record = self
|
let record = self
|
||||||
.registry
|
.registry
|
||||||
.get(&input.name)
|
.get(&input.name)
|
||||||
@@ -269,6 +315,22 @@ impl Tool for SubWorkerStopTool {
|
|||||||
) -> Result<ToolOutput, ToolError> {
|
) -> Result<ToolOutput, ToolError> {
|
||||||
let input: NameInput = serde_json::from_str(input_json)
|
let input: NameInput = serde_json::from_str(input_json)
|
||||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid SubWorkerStop input: {e}")))?;
|
.map_err(|e| ToolError::InvalidArgument(format!("invalid SubWorkerStop input: {e}")))?;
|
||||||
|
if let Some(record) = self.registry.get_internal(&input.name) {
|
||||||
|
record.session.stop().await.map_err(|error| {
|
||||||
|
ToolError::ExecutionFailed(format!("stop `{}`: {error}", input.name))
|
||||||
|
})?;
|
||||||
|
self.registry
|
||||||
|
.remove_internal(&input.name)
|
||||||
|
.await
|
||||||
|
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
|
||||||
|
return Ok(ToolOutput {
|
||||||
|
summary: format!(
|
||||||
|
"stopped worker `{}` and reclaimed delegated scope",
|
||||||
|
input.name
|
||||||
|
),
|
||||||
|
content: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
let record = self
|
let record = self
|
||||||
.registry
|
.registry
|
||||||
.get(&input.name)
|
.get(&input.name)
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ use tokio::net::UnixStream;
|
|||||||
use tokio::sync::Mutex;
|
use tokio::sync::Mutex;
|
||||||
use tracing::warn;
|
use tracing::warn;
|
||||||
|
|
||||||
|
use crate::internal_worker::InternalWorkerSessionHandle;
|
||||||
use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
|
use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
|
||||||
use crate::runtime::worker_allocation;
|
use crate::runtime::worker_allocation;
|
||||||
|
|
||||||
@@ -38,8 +39,16 @@ type RegistryReclaimWriter = Arc<dyn Fn(&SpawnedWorkerRecord) -> io::Result<()>
|
|||||||
const RESTORE_REACHABILITY_TIMEOUT: Duration = Duration::from_millis(500);
|
const RESTORE_REACHABILITY_TIMEOUT: Duration = Duration::from_millis(500);
|
||||||
const REGISTRY_CLEANUP_TIMEOUT: Duration = Duration::from_secs(15);
|
const REGISTRY_CLEANUP_TIMEOUT: Duration = Duration::from_secs(15);
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub(crate) struct InternalSpawnedWorkerRecord {
|
||||||
|
pub worker_name: String,
|
||||||
|
pub scope_delegated: Vec<ScopeRule>,
|
||||||
|
pub session: InternalWorkerSessionHandle,
|
||||||
|
}
|
||||||
|
|
||||||
pub struct SpawnedWorkerRegistry {
|
pub struct SpawnedWorkerRegistry {
|
||||||
records: Mutex<Vec<SpawnedWorkerRecord>>,
|
records: Mutex<Vec<SpawnedWorkerRecord>>,
|
||||||
|
internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>,
|
||||||
cursors: Mutex<HashMap<String, usize>>,
|
cursors: Mutex<HashMap<String, usize>>,
|
||||||
mutations: Mutex<()>,
|
mutations: Mutex<()>,
|
||||||
runtime_dir: Arc<RuntimeDir>,
|
runtime_dir: Arc<RuntimeDir>,
|
||||||
@@ -58,6 +67,7 @@ impl SpawnedWorkerRegistry {
|
|||||||
pub fn new(runtime_dir: Arc<RuntimeDir>) -> Arc<Self> {
|
pub fn new(runtime_dir: Arc<RuntimeDir>) -> Arc<Self> {
|
||||||
Arc::new(Self {
|
Arc::new(Self {
|
||||||
records: Mutex::new(Vec::new()),
|
records: Mutex::new(Vec::new()),
|
||||||
|
internal_records: std::sync::Mutex::new(Vec::new()),
|
||||||
cursors: Mutex::new(HashMap::new()),
|
cursors: Mutex::new(HashMap::new()),
|
||||||
mutations: Mutex::new(()),
|
mutations: Mutex::new(()),
|
||||||
runtime_dir,
|
runtime_dir,
|
||||||
@@ -170,6 +180,7 @@ impl SpawnedWorkerRegistry {
|
|||||||
Ok(SpawnedWorkerRegistryLoad {
|
Ok(SpawnedWorkerRegistryLoad {
|
||||||
registry: Arc::new(Self {
|
registry: Arc::new(Self {
|
||||||
records: Mutex::new(records),
|
records: Mutex::new(records),
|
||||||
|
internal_records: std::sync::Mutex::new(Vec::new()),
|
||||||
cursors: Mutex::new(HashMap::new()),
|
cursors: Mutex::new(HashMap::new()),
|
||||||
mutations: Mutex::new(()),
|
mutations: Mutex::new(()),
|
||||||
runtime_dir,
|
runtime_dir,
|
||||||
@@ -182,7 +193,73 @@ impl SpawnedWorkerRegistry {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Append a new record and persist the full list. Returns an I/O
|
pub(crate) fn add_internal(&self, record: InternalSpawnedWorkerRecord) -> io::Result<()> {
|
||||||
|
let mut records = self
|
||||||
|
.internal_records
|
||||||
|
.lock()
|
||||||
|
.map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))?;
|
||||||
|
if records
|
||||||
|
.iter()
|
||||||
|
.any(|existing| existing.worker_name == record.worker_name)
|
||||||
|
{
|
||||||
|
return Err(io::Error::new(
|
||||||
|
io::ErrorKind::AlreadyExists,
|
||||||
|
format!(
|
||||||
|
"spawned worker `{}` is already registered",
|
||||||
|
record.worker_name
|
||||||
|
),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
records.push(record);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn get_internal(&self, worker_name: &str) -> Option<InternalSpawnedWorkerRecord> {
|
||||||
|
self.internal_records
|
||||||
|
.lock()
|
||||||
|
.ok()?
|
||||||
|
.iter()
|
||||||
|
.find(|record| record.worker_name == worker_name)
|
||||||
|
.cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn list_internal(&self) -> Vec<InternalSpawnedWorkerRecord> {
|
||||||
|
self.internal_records
|
||||||
|
.lock()
|
||||||
|
.map(|records| records.clone())
|
||||||
|
.unwrap_or_default()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn remove_internal(
|
||||||
|
&self,
|
||||||
|
worker_name: &str,
|
||||||
|
) -> io::Result<Option<InternalSpawnedWorkerRecord>> {
|
||||||
|
let removed = {
|
||||||
|
let mut records = self
|
||||||
|
.internal_records
|
||||||
|
.lock()
|
||||||
|
.map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))?;
|
||||||
|
records
|
||||||
|
.iter()
|
||||||
|
.position(|record| record.worker_name == worker_name)
|
||||||
|
.map(|index| records.remove(index))
|
||||||
|
};
|
||||||
|
self.cursors.lock().await.remove(worker_name);
|
||||||
|
if let (Some(record), Some(parent_scope)) = (&removed, &self.parent_scope) {
|
||||||
|
let write_rules = record
|
||||||
|
.scope_delegated
|
||||||
|
.iter()
|
||||||
|
.filter(|rule| rule.permission == Permission::Write)
|
||||||
|
.cloned()
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
parent_scope
|
||||||
|
.update(|current| current.with_removed_deny_rules(write_rules))
|
||||||
|
.map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?;
|
||||||
|
}
|
||||||
|
Ok(removed)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append a new legacy process record and persist the full list.
|
||||||
/// error if either persisted write fails; the in-memory state is still
|
/// error if either persisted write fails; the in-memory state is still
|
||||||
/// updated in that case — the next successful write will reconcile.
|
/// updated in that case — the next successful write will reconcile.
|
||||||
pub async fn add(&self, record: SpawnedWorkerRecord) -> io::Result<()> {
|
pub async fn add(&self, record: SpawnedWorkerRecord) -> io::Result<()> {
|
||||||
|
|||||||
Reference in New Issue
Block a user