feat(TASK-061): add rag.sync.* and mcp.server.* hook events plus COYOTE_ROLE
- Register five new events (26 -> 31): rag.sync.started/completed/failed and mcp.server.connected/failed, whitelist-gated for agent global hooks like every non-escalation event - Bracket every real RAG document sync with rag.sync.* via pre-resolved RagSyncHooks snapshots carried into the build funnels; cache hits, pure loads, and attaches fire nothing, and a user abort mid-sync counts as a failure - Fire mcp.server.connected only on real spawns in McpFactory::acquire (reusing a live server is silent), marking respawns of a previously live key with COYOTE_MCP_RECONNECT; mcp.server.failed carries COYOTE_ERROR and COYOTE_MCP_AUTH_REQUIRED on auth failures. Boot-time registry connects deliberately fire nothing - Carry COYOTE_ROLE on session.* and turn.* events, resolved at fire time so mid-session role switches and temp roles are reflected
This commit is contained in:
9 files changed
+986
-37
No files matched your search
+17
-1
@@ -1,6 +1,6 @@
|
||||
use super::*;
|
||||
|
||||
use crate::hooks::HooksMap;
|
||||
use crate::hooks::{HooksMap, RagSyncHooks};
|
||||
use crate::{
|
||||
client::Model,
|
||||
config::memory,
|
||||
@@ -356,6 +356,16 @@ impl Agent {
|
||||
}
|
||||
};
|
||||
|
||||
// Resolved from the config under construction: no context carries
|
||||
// this agent yet, but its whitelist gate and hooks must still apply
|
||||
// to the RAG builds below.
|
||||
let rag_sync_hooks = RagSyncHooks::resolve_for_agent(
|
||||
&app.hooks,
|
||||
&agent_config.global_hooks,
|
||||
&agent_config.hooks,
|
||||
name,
|
||||
);
|
||||
|
||||
let rag = if rag_path.exists() {
|
||||
let key = RagKey::Agent(name.to_string());
|
||||
let app_clone = app.clone();
|
||||
@@ -388,6 +398,7 @@ impl Agent {
|
||||
let app_clone = app.clone();
|
||||
let rag_path_clone = rag_path.clone();
|
||||
let abort = abort_signal.clone();
|
||||
let sync_hooks = rag_sync_hooks.clone();
|
||||
let rag = app_state
|
||||
.rag_cache
|
||||
.load_with(key, || async move {
|
||||
@@ -398,6 +409,7 @@ impl Agent {
|
||||
&document_paths,
|
||||
abort,
|
||||
true,
|
||||
sync_hooks,
|
||||
)
|
||||
.await
|
||||
})
|
||||
@@ -421,6 +433,7 @@ impl Agent {
|
||||
&loaders,
|
||||
info_flag,
|
||||
abort_signal.clone(),
|
||||
&rag_sync_hooks,
|
||||
)
|
||||
.await?
|
||||
}
|
||||
@@ -1303,6 +1316,7 @@ async fn init_graph_rags(
|
||||
loaders: &HashMap<String, String>,
|
||||
info_flag: bool,
|
||||
abort_signal: AbortSignal,
|
||||
sync_hooks: &RagSyncHooks,
|
||||
) -> Result<HashMap<String, Arc<Rag>>> {
|
||||
let mut rags = HashMap::new();
|
||||
if info_flag {
|
||||
@@ -1379,6 +1393,7 @@ async fn init_graph_rags(
|
||||
let path_clone = rag_path.clone();
|
||||
let name_clone = node_id.clone();
|
||||
let abort = abort_signal.clone();
|
||||
let sync_hooks = sync_hooks.clone();
|
||||
app_state
|
||||
.rag_cache
|
||||
.load_with(key, || async move {
|
||||
@@ -1389,6 +1404,7 @@ async fn init_graph_rags(
|
||||
&document_paths,
|
||||
&config,
|
||||
abort,
|
||||
sync_hooks,
|
||||
)
|
||||
.await
|
||||
})
|
||||
|
||||
+254
-11
@@ -1,5 +1,6 @@
|
||||
use crate::hooks::McpServerHooks;
|
||||
use crate::mcp::{
|
||||
ConnectedServer, JsonField, McpAuthRequired, McpServer, McpTransportType,
|
||||
ConnectedServer, HttpAuth, JsonField, McpAuthRequired, McpServer, McpTransportType,
|
||||
is_auth_required_error, resolve_http_auth, spawn_mcp_server,
|
||||
};
|
||||
|
||||
@@ -95,34 +96,91 @@ impl McpFactory {
|
||||
name: &str,
|
||||
spec: &McpServer,
|
||||
log_path: Option<&Path>,
|
||||
hooks: &McpServerHooks,
|
||||
) -> Result<Arc<ConnectedServer>> {
|
||||
let key = McpServerKey::from_spec(name, spec);
|
||||
|
||||
// Reuse of a live server fires no hooks; only a real spawn below
|
||||
// reports anything.
|
||||
if let Some(existing) = self.try_get_active(&key) {
|
||||
return Ok(existing);
|
||||
}
|
||||
|
||||
let (auth, auth_reason) = resolve_http_auth(name, spec).await;
|
||||
let handle = spawn_mcp_server(spec, log_path, auth).await.map_err(|e| {
|
||||
if is_auth_required_error(&e) {
|
||||
e.context(McpAuthRequired {
|
||||
server: name.to_string(),
|
||||
reason: auth_reason,
|
||||
})
|
||||
} else {
|
||||
e
|
||||
// The live probe failed, so an entry still present for this key is a
|
||||
// dead weak: the key was connected earlier in this process and the
|
||||
// spawn below is a reconnect.
|
||||
let reconnect = self.active.lock().contains_key(&key);
|
||||
|
||||
let transport = transport_label(&spec.transport_type);
|
||||
let handle = match spawn_server(name, spec, log_path).await {
|
||||
Ok(handle) => handle,
|
||||
Err(e) => {
|
||||
hooks.fire_failed(name, transport, &e, is_auth_required_error(&e));
|
||||
return Err(e);
|
||||
}
|
||||
})?;
|
||||
};
|
||||
self.insert_active(key, &handle);
|
||||
hooks.fire_connected(name, transport, reconnect);
|
||||
Ok(handle)
|
||||
}
|
||||
}
|
||||
|
||||
fn transport_label(transport: &McpTransportType) -> &'static str {
|
||||
match transport {
|
||||
McpTransportType::Stdio => "stdio",
|
||||
McpTransportType::Http => "http",
|
||||
McpTransportType::Sse => "sse",
|
||||
}
|
||||
}
|
||||
|
||||
async fn spawn_server(
|
||||
name: &str,
|
||||
spec: &McpServer,
|
||||
log_path: Option<&Path>,
|
||||
) -> Result<Arc<ConnectedServer>> {
|
||||
let (auth, auth_reason) = resolve_http_auth(name, spec).await;
|
||||
spawn_transport(spec, log_path, auth).await.map_err(|e| {
|
||||
if is_auth_required_error(&e) {
|
||||
e.context(McpAuthRequired {
|
||||
server: name.to_string(),
|
||||
reason: auth_reason,
|
||||
})
|
||||
} else {
|
||||
e
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn spawn_transport(
|
||||
spec: &McpServer,
|
||||
log_path: Option<&Path>,
|
||||
auth: HttpAuth,
|
||||
) -> Result<Arc<ConnectedServer>> {
|
||||
#[cfg(test)]
|
||||
if let Some(result) = STUB_SPAWNS.lock().pop() {
|
||||
return result;
|
||||
}
|
||||
spawn_mcp_server(spec, log_path, auth).await
|
||||
}
|
||||
|
||||
/// Injected spawn outcomes, popped instead of spawning a real transport:
|
||||
/// nothing in a test environment speaks MCP over stdio or HTTP, so the
|
||||
/// connect path is exercised with fixture handles. Empty means spawn for
|
||||
/// real.
|
||||
#[cfg(test)]
|
||||
static STUB_SPAWNS: Mutex<Vec<Result<Arc<ConnectedServer>>>> = Mutex::new(Vec::new());
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::test_fixtures::FixtureServer;
|
||||
use crate::config::{AppConfig, AppState, RequestContext, WorkingMode};
|
||||
use crate::hooks::{HookDef, HooksMap, test_sink};
|
||||
use crate::mcp::{JsonField, McpServer, McpTransportType};
|
||||
use anyhow::anyhow;
|
||||
use indexmap::IndexMap;
|
||||
use rmcp::ServiceExt;
|
||||
use serial_test::serial;
|
||||
use std::collections::HashMap;
|
||||
|
||||
fn stdio_spec(
|
||||
@@ -321,4 +379,189 @@ mod tests {
|
||||
let map = factory.active.lock();
|
||||
assert!(map.is_empty());
|
||||
}
|
||||
|
||||
/// Clears injected spawn outcomes on drop so a panicking test cannot
|
||||
/// leak stubs into later tests.
|
||||
struct StubSpawnsGuard;
|
||||
|
||||
impl StubSpawnsGuard {
|
||||
fn push(result: Result<Arc<ConnectedServer>>) -> Self {
|
||||
STUB_SPAWNS.lock().push(result);
|
||||
StubSpawnsGuard
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for StubSpawnsGuard {
|
||||
fn drop(&mut self) {
|
||||
STUB_SPAWNS.lock().clear();
|
||||
}
|
||||
}
|
||||
|
||||
/// An in-process client/server pair; the returned server half must stay
|
||||
/// alive for the client handle to keep working.
|
||||
async fn fixture_handle() -> (Arc<ConnectedServer>, impl Sized) {
|
||||
let (client_io, server_io) = tokio::io::duplex(4096);
|
||||
let (server, client) = tokio::join!(
|
||||
FixtureServer::default().serve(server_io),
|
||||
().serve(client_io)
|
||||
);
|
||||
(Arc::new(client.unwrap()), server.unwrap())
|
||||
}
|
||||
|
||||
fn ctx_with_mcp_hooks(marker: &str) -> RequestContext {
|
||||
let mut hooks_map = HooksMap::default();
|
||||
for event in ["mcp.server.connected", "mcp.server.failed"] {
|
||||
hooks_map.insert(
|
||||
event.to_string(),
|
||||
vec![HookDef {
|
||||
name: format!("{marker}-{event}"),
|
||||
command: "true".to_string(),
|
||||
}],
|
||||
);
|
||||
}
|
||||
let mut app = AppState::test_default();
|
||||
app.config = Arc::new(AppConfig {
|
||||
hooks: hooks_map,
|
||||
..Default::default()
|
||||
});
|
||||
RequestContext::new(Arc::new(app), WorkingMode::Cmd)
|
||||
}
|
||||
|
||||
fn mcp_captures(marker: &str) -> Vec<test_sink::Capture> {
|
||||
test_sink::snapshot()
|
||||
.into_iter()
|
||||
.filter(|capture| capture.hook_name.starts_with(marker))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
async fn acquire_fires_connected_only_on_real_spawns_and_marks_reconnects() {
|
||||
let _sink = test_sink::install();
|
||||
let marker = "mcp-spawn-k7d";
|
||||
let ctx = ctx_with_mcp_hooks(marker);
|
||||
let hooks = McpServerHooks::resolve(&ctx);
|
||||
let factory = McpFactory::default();
|
||||
let spec = stdio_spec("fixture-server-cmd", None, None);
|
||||
|
||||
let (first, _first_server) = fixture_handle().await;
|
||||
let stub = StubSpawnsGuard::push(Ok(first.clone()));
|
||||
let connected = factory.acquire("srv", &spec, None, &hooks).await.unwrap();
|
||||
drop(stub);
|
||||
assert!(Arc::ptr_eq(&connected, &first));
|
||||
let captures = mcp_captures(marker);
|
||||
assert_eq!(captures.len(), 1, "{captures:?}");
|
||||
let capture = &captures[0];
|
||||
assert_eq!(capture.hook_name, format!("{marker}-mcp.server.connected"));
|
||||
assert_eq!(
|
||||
capture.envs.get("COYOTE_MCP_SERVER").map(String::as_str),
|
||||
Some("srv")
|
||||
);
|
||||
assert_eq!(
|
||||
capture.envs.get("COYOTE_MCP_TRANSPORT").map(String::as_str),
|
||||
Some("stdio")
|
||||
);
|
||||
assert!(
|
||||
!capture.envs.contains_key("COYOTE_MCP_RECONNECT"),
|
||||
"a first connect must not be marked as a reconnect"
|
||||
);
|
||||
|
||||
let reused = factory.acquire("srv", &spec, None, &hooks).await.unwrap();
|
||||
assert!(Arc::ptr_eq(&reused, &first));
|
||||
assert_eq!(
|
||||
mcp_captures(marker).len(),
|
||||
1,
|
||||
"reusing a live server must fire nothing"
|
||||
);
|
||||
|
||||
// Every holder dropped: the weak in the factory dies, and the next
|
||||
// acquire is a real spawn again — now marked as a reconnect.
|
||||
drop(connected);
|
||||
drop(reused);
|
||||
drop(first);
|
||||
let (second, _second_server) = fixture_handle().await;
|
||||
let _stub = StubSpawnsGuard::push(Ok(second.clone()));
|
||||
let respawned = factory.acquire("srv", &spec, None, &hooks).await.unwrap();
|
||||
assert!(Arc::ptr_eq(&respawned, &second));
|
||||
let captures = mcp_captures(marker);
|
||||
assert_eq!(captures.len(), 2, "{captures:?}");
|
||||
assert_eq!(
|
||||
captures[1]
|
||||
.envs
|
||||
.get("COYOTE_MCP_RECONNECT")
|
||||
.map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
async fn acquire_spawn_failure_fires_failed_with_the_error() {
|
||||
let _sink = test_sink::install();
|
||||
let marker = "mcp-fail-v2q";
|
||||
let ctx = ctx_with_mcp_hooks(marker);
|
||||
let hooks = McpServerHooks::resolve(&ctx);
|
||||
let factory = McpFactory::default();
|
||||
// No stub: the nonexistent binary drives the real spawn-error path.
|
||||
let spec = stdio_spec("coyote-nonexistent-mcp-binary-v2q", None, None);
|
||||
|
||||
let result = factory.acquire("bad-srv", &spec, None, &hooks).await;
|
||||
|
||||
assert!(result.is_err());
|
||||
let captures = mcp_captures(marker);
|
||||
assert_eq!(captures.len(), 1, "{captures:?}");
|
||||
let capture = &captures[0];
|
||||
assert_eq!(capture.hook_name, format!("{marker}-mcp.server.failed"));
|
||||
assert_eq!(
|
||||
capture.envs.get("COYOTE_MCP_SERVER").map(String::as_str),
|
||||
Some("bad-srv")
|
||||
);
|
||||
assert_eq!(
|
||||
capture.envs.get("COYOTE_MCP_TRANSPORT").map(String::as_str),
|
||||
Some("stdio")
|
||||
);
|
||||
assert!(
|
||||
capture
|
||||
.envs
|
||||
.get("COYOTE_ERROR")
|
||||
.is_some_and(|error| !error.is_empty())
|
||||
);
|
||||
assert!(
|
||||
!capture.envs.contains_key("COYOTE_MCP_AUTH_REQUIRED"),
|
||||
"a plain spawn failure must not be flagged as auth-required"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
async fn acquire_auth_failure_sets_the_auth_required_flag() {
|
||||
let _sink = test_sink::install();
|
||||
let marker = "mcp-auth-h4n";
|
||||
let ctx = ctx_with_mcp_hooks(marker);
|
||||
let hooks = McpServerHooks::resolve(&ctx);
|
||||
let factory = McpFactory::default();
|
||||
let spec = stdio_spec("oauth-cmd", None, None);
|
||||
let _stub = StubSpawnsGuard::push(Err(anyhow!("Auth required: no stored token")));
|
||||
|
||||
let err = factory
|
||||
.acquire("oauth-srv", &spec, None, &hooks)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(
|
||||
err.downcast_ref::<McpAuthRequired>().is_some(),
|
||||
"the auth error must be wrapped in McpAuthRequired context"
|
||||
);
|
||||
let captures = mcp_captures(marker);
|
||||
assert_eq!(captures.len(), 1, "{captures:?}");
|
||||
let capture = &captures[0];
|
||||
assert_eq!(capture.hook_name, format!("{marker}-mcp.server.failed"));
|
||||
assert_eq!(
|
||||
capture
|
||||
.envs
|
||||
.get("COYOTE_MCP_AUTH_REQUIRED")
|
||||
.map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -47,3 +47,61 @@ impl RagCache {
|
||||
Ok(arc)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::AppConfig;
|
||||
use crate::hooks::test_sink;
|
||||
use crate::rag::RagData;
|
||||
use anyhow::anyhow;
|
||||
use serial_test::serial;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
/// A cache hit must never reach the loader: the loader closure is where
|
||||
/// a real RAG build (and its `rag.sync.*` hook dispatch) happens, so
|
||||
/// returning the cached entry is what keeps warm-cache lookups silent.
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
async fn cache_hit_returns_the_cached_rag_without_running_the_loader() {
|
||||
let _sink = test_sink::install();
|
||||
test_sink::drain();
|
||||
let unique = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
let dir = std::env::temp_dir().join(format!("coyote-rag-cache-hit-{unique}"));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let path = dir.join("kb.yaml");
|
||||
let data = RagData {
|
||||
driver: "yaml".to_string(),
|
||||
// The only embedding model resolvable in tests; see
|
||||
// `seed_model_registries` in the request_context tests.
|
||||
embedding_model: "test-seeded:test-embedder".to_string(),
|
||||
chunk_size: 1000,
|
||||
chunk_overlap: 100,
|
||||
top_k: 5,
|
||||
..Default::default()
|
||||
};
|
||||
std::fs::write(&path, serde_yaml::to_string(&data).unwrap()).unwrap();
|
||||
let rag = Arc::new(Rag::load(&AppConfig::default(), "kb", &path).await.unwrap());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
|
||||
let cache = RagCache::default();
|
||||
let key = RagKey::Named("kb".to_string());
|
||||
cache.insert(key.clone(), &rag);
|
||||
|
||||
let loaded = cache
|
||||
.load_with(key, || async {
|
||||
Err(anyhow!("a cache hit must not rebuild"))
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(Arc::ptr_eq(&loaded, &rag));
|
||||
assert!(
|
||||
test_sink::drain().is_empty(),
|
||||
"a cache hit must dispatch no hooks"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1077,7 +1077,12 @@ impl RequestContext {
|
||||
// COYOTE_SESSION_ID and session-held role hooks still resolve; no
|
||||
// session means no event.
|
||||
if self.session.is_some() {
|
||||
hooks::fire(HookEvent::SessionEnded, self, &[], None);
|
||||
hooks::fire(
|
||||
HookEvent::SessionEnded,
|
||||
self,
|
||||
&hooks::role_extras(self),
|
||||
None,
|
||||
);
|
||||
}
|
||||
if let Some(mut session) = self.session.take() {
|
||||
let sessions_dir = self.sessions_dir();
|
||||
@@ -4406,6 +4411,7 @@ impl RequestContext {
|
||||
|
||||
if !server_ids.is_empty() {
|
||||
let app_ref = &self.app;
|
||||
let mcp_hooks = hooks::McpServerHooks::resolve(self);
|
||||
let acquire_all = async {
|
||||
let mut handles = Vec::new();
|
||||
let mut auth_required = Vec::new();
|
||||
@@ -4413,7 +4419,7 @@ impl RequestContext {
|
||||
if let Some(spec) = mcp_config.mcp_servers.get(id) {
|
||||
match app_ref
|
||||
.mcp_factory
|
||||
.acquire(id, spec, app_ref.mcp_log_path.as_deref())
|
||||
.acquire(id, spec, app_ref.mcp_log_path.as_deref(), &mcp_hooks)
|
||||
.await
|
||||
{
|
||||
Ok(handle) => handles.push((id.clone(), handle)),
|
||||
@@ -4738,9 +4744,19 @@ impl RequestContext {
|
||||
self.refresh_mcp_tool_filters();
|
||||
self.init_agent_session_variables(new_session)?;
|
||||
if created_new_session {
|
||||
hooks::fire(HookEvent::SessionStarted, self, &[], None);
|
||||
hooks::fire(
|
||||
HookEvent::SessionStarted,
|
||||
self,
|
||||
&hooks::role_extras(self),
|
||||
None,
|
||||
);
|
||||
} else {
|
||||
hooks::fire(HookEvent::SessionResumed, self, &[], None);
|
||||
hooks::fire(
|
||||
HookEvent::SessionResumed,
|
||||
self,
|
||||
&hooks::role_extras(self),
|
||||
None,
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -5334,6 +5350,7 @@ impl RequestContext {
|
||||
let vault = self.app.vault.clone();
|
||||
let rag_cache = self.rag_cache();
|
||||
let working_mode = self.working_mode;
|
||||
let sync_hooks = hooks::RagSyncHooks::resolve(self);
|
||||
|
||||
let (rag, rag_key): (Arc<Rag>, Option<RagKey>) = match rag {
|
||||
None => {
|
||||
@@ -5352,6 +5369,7 @@ impl RequestContext {
|
||||
&[],
|
||||
abort_signal.clone(),
|
||||
false,
|
||||
sync_hooks,
|
||||
)
|
||||
.await?,
|
||||
),
|
||||
@@ -5373,8 +5391,16 @@ impl RequestContext {
|
||||
if working_mode.is_cmd() {
|
||||
bail!("Unknown RAG '{name}'");
|
||||
}
|
||||
Rag::init(&app, name, &rag_path, &[], abort_signal.clone(), true)
|
||||
.await
|
||||
Rag::init(
|
||||
&app,
|
||||
name,
|
||||
&rag_path,
|
||||
&[],
|
||||
abort_signal.clone(),
|
||||
true,
|
||||
sync_hooks,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
Rag::load_async(&app, &vault, name, &rag_path).await
|
||||
}
|
||||
@@ -5458,6 +5484,7 @@ impl RequestContext {
|
||||
false,
|
||||
&self.app.config,
|
||||
abort_signal,
|
||||
hooks::RagSyncHooks::resolve(self),
|
||||
)
|
||||
.await?;
|
||||
self.rag = Some(Arc::new(rag));
|
||||
@@ -5488,8 +5515,15 @@ impl RequestContext {
|
||||
This will call the embedding API and may take a while.",
|
||||
rag.file_count()
|
||||
);
|
||||
rag.refresh_document_paths(&document_paths, true, true, &self.app.config, abort_signal)
|
||||
.await?;
|
||||
rag.refresh_document_paths(
|
||||
&document_paths,
|
||||
true,
|
||||
true,
|
||||
&self.app.config,
|
||||
abort_signal,
|
||||
hooks::RagSyncHooks::resolve(self),
|
||||
)
|
||||
.await?;
|
||||
self.rag = Some(Arc::new(rag));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user