Publish harness and TUI open-source
initial sync from the monorepo
This commit is contained in:
grokkybara[bot]
commit
c68e39f604
2734 changed files with 1437016 additions and 0 deletions
335
crates/common/xai-computer-hub-core/tests/compound_resolver.rs
Normal file
335
crates/common/xai-computer-hub-core/tests/compound_resolver.rs
Normal file
|
|
@ -0,0 +1,335 @@
|
|||
//! `CompoundResolver` and `ResolvedTool` coverage. Exercises local-only,
|
||||
//! local-shadows-remote, remote-fallback, and cross-session scenarios.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use futures::StreamExt;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use xai_computer_hub_core::{
|
||||
CompoundResolver, ConnectionCleanupReport, ErasedTool, ResolvedTool, SessionCleanupReport,
|
||||
ToolHandle, ToolRegistry, ToolSessionBindOutcome, ToolSessionUnbindOutcome,
|
||||
};
|
||||
use xai_tool_protocol::{
|
||||
ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId,
|
||||
ToolRegistration, ToolServerRegistration, TransportKind, UserId,
|
||||
};
|
||||
use xai_tool_runtime::{
|
||||
SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolError, ToolStreamItem,
|
||||
};
|
||||
use xai_tool_types::ToolDescription;
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
struct EmptyArgs {}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct StubTool {
|
||||
id: ToolId,
|
||||
}
|
||||
|
||||
impl Tool for StubTool {
|
||||
type Args = EmptyArgs;
|
||||
type Output = serde_json::Value;
|
||||
|
||||
fn id(&self) -> ToolId {
|
||||
self.id.clone()
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &::xai_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new(self.id.as_str(), format!("stub for {}", self.id))
|
||||
}
|
||||
|
||||
async fn run(
|
||||
&self,
|
||||
_ctx: ToolCallContext,
|
||||
_args: Self::Args,
|
||||
) -> Result<Self::Output, ToolError> {
|
||||
Ok(serde_json::json!({"id": self.id.as_str()}))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct PlaneRegistry {
|
||||
// Set once at construction; `TransportKind` is `Copy` so a direct
|
||||
// field is the obvious choice — no interior mutability required.
|
||||
transport_kind: TransportKind,
|
||||
entries: DashMap<(SessionId, ToolId), ToolRegistration>,
|
||||
handles: DashMap<ToolId, Arc<dyn ToolHandle>>,
|
||||
}
|
||||
|
||||
impl PlaneRegistry {
|
||||
fn new(kind: TransportKind) -> Self {
|
||||
Self {
|
||||
transport_kind: kind,
|
||||
entries: DashMap::new(),
|
||||
handles: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn install(&self, session: &SessionId, id: &ToolId) {
|
||||
let reg = ToolRegistration {
|
||||
tool_id: id.clone(),
|
||||
sessions: Some(vec![session.clone()]),
|
||||
user_id: UserId::new("alice").expect("user id"),
|
||||
server_id: None,
|
||||
description: ToolDescription::new(id.as_str(), format!("stub for {id}")),
|
||||
input_schema: None,
|
||||
capabilities: None,
|
||||
notification_schemas: None,
|
||||
transport_kind: self.transport_kind,
|
||||
if_match_generation: None,
|
||||
metadata: None,
|
||||
};
|
||||
self.entries.insert((session.clone(), id.clone()), reg);
|
||||
self.handles.insert(
|
||||
id.clone(),
|
||||
Arc::new(ErasedTool::new(StubTool { id: id.clone() })),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolRegistry for PlaneRegistry {
|
||||
async fn register_tool(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolRegistration,
|
||||
) -> RegistrationOutcome {
|
||||
unreachable!("resolver tests pre-populate via install()")
|
||||
}
|
||||
async fn register_server(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolServerRegistration,
|
||||
) -> Vec<RegistrationOutcome> {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_tool(&self, _connection_id: &ConnectionId, _tool: &ToolId) -> bool {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_server(&self, _connection_id: &ConnectionId, _server: &ServerId) -> usize {
|
||||
unreachable!()
|
||||
}
|
||||
async fn bind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionBindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unbind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionUnbindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn drop_connection(&self, _connection_id: &ConnectionId) -> ConnectionCleanupReport {
|
||||
ConnectionCleanupReport::default()
|
||||
}
|
||||
|
||||
fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option<ResolvedTool> {
|
||||
let registration = self
|
||||
.entries
|
||||
.get(&(session.clone(), tool.clone()))?
|
||||
.value()
|
||||
.clone();
|
||||
let handle = self.handles.get(tool)?.value().clone();
|
||||
match registration.transport_kind {
|
||||
TransportKind::Local => Some(ResolvedTool::Local {
|
||||
tool: handle,
|
||||
registration,
|
||||
}),
|
||||
TransportKind::Remote => Some(ResolvedTool::Remote {
|
||||
proxy: handle,
|
||||
registration,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn list_tools(&self, _session: &SessionId, _mode: &ToolDefinitionMode) -> Vec<ToolDescription> {
|
||||
vec![]
|
||||
}
|
||||
fn list_servers(&self, _session: &SessionId) -> Vec<ServerSummary> {
|
||||
vec![]
|
||||
}
|
||||
fn search(&self, _session: &SessionId, _query: &str, _limit: usize) -> SearchSnapshot {
|
||||
SearchSnapshot {
|
||||
results: vec![],
|
||||
total_hidden_tools: 0,
|
||||
is_ready: true,
|
||||
}
|
||||
}
|
||||
async fn unregister_session(&self, _session: &SessionId) -> SessionCleanupReport {
|
||||
SessionCleanupReport::default()
|
||||
}
|
||||
fn tool_sessions(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
) -> std::collections::HashSet<SessionId> {
|
||||
std::collections::HashSet::new()
|
||||
}
|
||||
|
||||
fn list_servers_for_user(
|
||||
&self,
|
||||
_user_id: &xai_tool_protocol::UserId,
|
||||
) -> Vec<xai_computer_hub_core::registry::ServerRecord> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn get_server_record(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
) -> Option<xai_computer_hub_core::registry::ServerRecord> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("tool id")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_only_resolves_local_hits() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
local.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
match resolver.resolve(&sid("sess-1"), &tid("foo")) {
|
||||
Some(ResolvedTool::Local { registration, .. }) => {
|
||||
assert_eq!(registration.tool_id, tid("foo"));
|
||||
}
|
||||
other => panic!("expected Local, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_only_returns_none_for_unknown() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
assert!(resolver.resolve(&sid("sess-1"), &tid("missing")).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compound_falls_through_to_remote_when_local_misses() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
let remote = Arc::new(PlaneRegistry::new(TransportKind::Remote));
|
||||
remote.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::compound(
|
||||
local as Arc<dyn ToolRegistry>,
|
||||
remote as Arc<dyn ToolRegistry>,
|
||||
);
|
||||
match resolver.resolve(&sid("sess-1"), &tid("foo")) {
|
||||
Some(ResolvedTool::Remote { registration, .. }) => {
|
||||
assert_eq!(registration.tool_id, tid("foo"));
|
||||
assert_eq!(registration.transport_kind, TransportKind::Remote);
|
||||
}
|
||||
other => panic!("expected Remote, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_shadows_same_id_remote() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
local.install(&sid("sess-1"), &tid("foo"));
|
||||
let remote = Arc::new(PlaneRegistry::new(TransportKind::Remote));
|
||||
remote.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::compound(
|
||||
local as Arc<dyn ToolRegistry>,
|
||||
remote as Arc<dyn ToolRegistry>,
|
||||
);
|
||||
match resolver.resolve(&sid("sess-1"), &tid("foo")) {
|
||||
Some(ResolvedTool::Local { registration, .. }) => {
|
||||
assert_eq!(registration.transport_kind, TransportKind::Local);
|
||||
}
|
||||
other => panic!("expected local resolution to shadow remote, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cross_session_lookup_returns_none() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
local.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
assert!(resolver.resolve(&sid("sess-other"), &tid("foo")).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compound_returns_none_when_neither_plane_holds_id() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
let remote = Arc::new(PlaneRegistry::new(TransportKind::Remote));
|
||||
let resolver = CompoundResolver::compound(
|
||||
local as Arc<dyn ToolRegistry>,
|
||||
remote as Arc<dyn ToolRegistry>,
|
||||
);
|
||||
assert!(resolver.resolve(&sid("sess-1"), &tid("foo")).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolved_tool_helpers_borrow_active_handle_and_registration() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
local.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
let resolved = resolver.resolve(&sid("sess-1"), &tid("foo")).expect("hit");
|
||||
assert_eq!(resolved.registration().tool_id, tid("foo"));
|
||||
assert_eq!(resolved.handle().id(), tid("foo"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_and_dispatch_drives_the_resolved_handle() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
local.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
let mut stream = resolver
|
||||
.resolve_and_dispatch(
|
||||
&sid("sess-1"),
|
||||
tid("foo"),
|
||||
serde_json::json!({}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert_eq!(typed.value, serde_json::json!({"id": "foo"}));
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
assert!(stream.next().await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_and_dispatch_misses_yield_terminal_not_found() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
let mut stream = resolver
|
||||
.resolve_and_dispatch(
|
||||
&sid("sess-1"),
|
||||
tid("missing"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == xai_tool_runtime::ToolErrorKind::NotFound =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("missing"),
|
||||
"detail should mention tool id: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(Err(NotFound)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
317
crates/common/xai-computer-hub-core/tests/inner_dispatch.rs
Normal file
317
crates/common/xai-computer-hub-core/tests/inner_dispatch.rs
Normal file
|
|
@ -0,0 +1,317 @@
|
|||
//! `InnerDispatchForResolver` coverage. Verifies the cycle-safe `Weak`
|
||||
//! resolver semantics and the session-bound resolution path.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::StreamExt;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use xai_computer_hub_core::{
|
||||
CompoundResolver, ConnectionCleanupReport, ErasedTool, InnerDispatchForResolver, ResolvedTool,
|
||||
SessionCleanupReport, ToolHandle, ToolRegistry, ToolSessionBindOutcome,
|
||||
ToolSessionUnbindOutcome,
|
||||
};
|
||||
use xai_tool_protocol::{
|
||||
ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId,
|
||||
ToolRegistration, ToolServerRegistration, TransportKind, UserId,
|
||||
};
|
||||
use xai_tool_runtime::{
|
||||
SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolDispatch, ToolError, ToolStreamItem,
|
||||
};
|
||||
use xai_tool_types::ToolDescription;
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
struct EchoArgs {
|
||||
payload: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct EchoTool;
|
||||
|
||||
impl Tool for EchoTool {
|
||||
type Args = EchoArgs;
|
||||
type Output = serde_json::Value;
|
||||
|
||||
fn id(&self) -> ToolId {
|
||||
ToolId::new("echo").expect("tool id")
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &::xai_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new("echo", "Echoes its input.")
|
||||
}
|
||||
|
||||
async fn run(
|
||||
&self,
|
||||
_ctx: ToolCallContext,
|
||||
args: Self::Args,
|
||||
) -> Result<Self::Output, ToolError> {
|
||||
Ok(serde_json::json!({"echoed": args.payload}))
|
||||
}
|
||||
}
|
||||
|
||||
type RegistryEntry = (ToolRegistration, Arc<dyn ToolHandle>);
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct InMemRegistry {
|
||||
entries: DashMap<(SessionId, ToolId), RegistryEntry>,
|
||||
}
|
||||
|
||||
impl InMemRegistry {
|
||||
fn install(&self, session: SessionId, tool: ToolId, handle: Arc<dyn ToolHandle>) {
|
||||
let registration = ToolRegistration {
|
||||
tool_id: tool.clone(),
|
||||
sessions: Some(vec![session.clone()]),
|
||||
user_id: UserId::new("alice").expect("user id"),
|
||||
server_id: None,
|
||||
description: handle.description(&xai_tool_runtime::ListToolsContext::default()),
|
||||
input_schema: None,
|
||||
capabilities: Some(handle.capabilities()),
|
||||
notification_schemas: None,
|
||||
transport_kind: TransportKind::Local,
|
||||
if_match_generation: None,
|
||||
metadata: None,
|
||||
};
|
||||
self.entries.insert((session, tool), (registration, handle));
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolRegistry for InMemRegistry {
|
||||
async fn register_tool(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolRegistration,
|
||||
) -> RegistrationOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn register_server(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolServerRegistration,
|
||||
) -> Vec<RegistrationOutcome> {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_tool(&self, _connection_id: &ConnectionId, _tool: &ToolId) -> bool {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_server(&self, _connection_id: &ConnectionId, _server: &ServerId) -> usize {
|
||||
unreachable!()
|
||||
}
|
||||
async fn bind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionBindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unbind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionUnbindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn drop_connection(&self, _connection_id: &ConnectionId) -> ConnectionCleanupReport {
|
||||
ConnectionCleanupReport::default()
|
||||
}
|
||||
fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option<ResolvedTool> {
|
||||
let (registration, handle) = self
|
||||
.entries
|
||||
.get(&(session.clone(), tool.clone()))?
|
||||
.value()
|
||||
.clone();
|
||||
Some(ResolvedTool::Local {
|
||||
tool: handle,
|
||||
registration,
|
||||
})
|
||||
}
|
||||
fn list_tools(&self, _session: &SessionId, _mode: &ToolDefinitionMode) -> Vec<ToolDescription> {
|
||||
vec![]
|
||||
}
|
||||
fn list_servers(&self, _session: &SessionId) -> Vec<ServerSummary> {
|
||||
vec![]
|
||||
}
|
||||
fn search(&self, _session: &SessionId, _query: &str, _limit: usize) -> SearchSnapshot {
|
||||
SearchSnapshot {
|
||||
results: vec![],
|
||||
total_hidden_tools: 0,
|
||||
is_ready: true,
|
||||
}
|
||||
}
|
||||
async fn unregister_session(&self, _session: &SessionId) -> SessionCleanupReport {
|
||||
SessionCleanupReport::default()
|
||||
}
|
||||
fn tool_sessions(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
) -> std::collections::HashSet<SessionId> {
|
||||
std::collections::HashSet::new()
|
||||
}
|
||||
|
||||
fn list_servers_for_user(
|
||||
&self,
|
||||
_user_id: &xai_tool_protocol::UserId,
|
||||
) -> Vec<xai_computer_hub_core::registry::ServerRecord> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn get_server_record(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
) -> Option<xai_computer_hub_core::registry::ServerRecord> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("tool id")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inner_dispatch_resolves_through_bound_session() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
registry.install(
|
||||
sid("sess-1"),
|
||||
tid("echo"),
|
||||
Arc::new(ErasedTool::new(EchoTool)),
|
||||
);
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let inner = InnerDispatchForResolver::new(Arc::downgrade(&resolver), sid("sess-1"));
|
||||
assert_eq!(inner.session_id(), &sid("sess-1"));
|
||||
|
||||
let result = inner
|
||||
.call_terminal(
|
||||
tid("echo"),
|
||||
serde_json::json!({"payload": "x"}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await
|
||||
.expect("terminal ok");
|
||||
assert_eq!(result.value, serde_json::json!({"echoed": "x"}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inner_dispatch_returns_not_found_when_tool_absent() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let inner = InnerDispatchForResolver::new(Arc::downgrade(&resolver), sid("sess-1"));
|
||||
let mut stream = inner
|
||||
.call(
|
||||
tid("ghost"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == xai_tool_runtime::ToolErrorKind::NotFound =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("ghost"),
|
||||
"detail should mention tool id: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(NotFound), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inner_dispatch_uses_bound_session_not_context_session() {
|
||||
// Even if the context were to carry a different session, the inner
|
||||
// dispatch handle resolves against its construction-time session.
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
registry.install(
|
||||
sid("sess-A"),
|
||||
tid("echo"),
|
||||
Arc::new(ErasedTool::new(EchoTool)),
|
||||
);
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let inner = InnerDispatchForResolver::new(Arc::downgrade(&resolver), sid("sess-B"));
|
||||
let mut stream = inner
|
||||
.call(
|
||||
tid("echo"),
|
||||
serde_json::json!({"payload": "x"}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == xai_tool_runtime::ToolErrorKind::NotFound => {}
|
||||
other => panic!("session-A registration must not be visible from session-B, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inner_dispatch_after_resolver_drop_returns_computer_hub_dropped() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let weak = Arc::downgrade(&resolver);
|
||||
let inner = InnerDispatchForResolver::new(weak, sid("sess-1"));
|
||||
drop(resolver);
|
||||
let mut stream = inner
|
||||
.call(
|
||||
tid("echo"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == xai_tool_runtime::ToolErrorKind::Custom =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("computer_hub_dropped")
|
||||
|| e.details
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("code"))
|
||||
.and_then(|c| c.as_str())
|
||||
== Some("computer_hub_dropped"),
|
||||
"expected computer_hub_dropped code, got: {:?}",
|
||||
e
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(Custom(computer_hub_dropped)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inner_dispatch_implements_object_safe_tool_dispatch() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let inner: Arc<dyn ToolDispatch> = Arc::new(InnerDispatchForResolver::new(
|
||||
Arc::downgrade(&resolver),
|
||||
sid("sess-1"),
|
||||
));
|
||||
let result = inner
|
||||
.call_terminal(
|
||||
tid("ghost"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(result, Err(ref e) if e.kind == xai_tool_runtime::ToolErrorKind::NotFound));
|
||||
}
|
||||
|
|
@ -0,0 +1,65 @@
|
|||
//! Decode-side coverage for `ToolErrorWire::Internal`'s optional `detail`:
|
||||
//! a populated detail must become the reconstructed `ToolError`'s message,
|
||||
//! and its absence (frames from older peers) must fall back to the historic
|
||||
//! constant.
|
||||
|
||||
use serde_json::json;
|
||||
use xai_computer_hub_core::{error_from_envelope, tool_error_from_wire};
|
||||
use xai_tool_protocol::{JsonRpcError, RequestId, ToolErrorWire};
|
||||
use xai_tool_runtime::ToolErrorKind;
|
||||
|
||||
#[test]
|
||||
fn internal_with_detail_reconstructs_the_wire_detail() {
|
||||
let err = tool_error_from_wire(ToolErrorWire::Internal {
|
||||
request_id: None,
|
||||
detail: Some("cross-instance tool.call timed out".to_owned()),
|
||||
});
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert_eq!(err.detail, "cross-instance tool.call timed out");
|
||||
// The `internal_error` code survives so callers can still classify it.
|
||||
assert_eq!(
|
||||
err.details
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("code"))
|
||||
.and_then(|v| v.as_str()),
|
||||
Some("internal_error"),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internal_without_detail_falls_back_to_the_historic_constant() {
|
||||
let err = tool_error_from_wire(ToolErrorWire::Internal {
|
||||
request_id: None,
|
||||
detail: None,
|
||||
});
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert_eq!(err.detail, "internal router error");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internal_with_request_id_keeps_both_code_and_request_id() {
|
||||
let err = tool_error_from_wire(ToolErrorWire::Internal {
|
||||
request_id: Some(RequestId::new("req-7").unwrap()),
|
||||
detail: Some("relay publish failed".to_owned()),
|
||||
});
|
||||
assert_eq!(err.detail, "relay publish failed");
|
||||
let details = err.details.expect("details present");
|
||||
assert_eq!(details["code"], json!("internal_error"));
|
||||
assert_eq!(details["request_id"], json!("req-7"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn envelope_with_internal_data_prefers_data_detail_over_message() {
|
||||
// The hub's `-32000 "internal error"` envelope keeps its constant message;
|
||||
// the harness must read the cause from `error.data`, not the message.
|
||||
let err = error_from_envelope(JsonRpcError {
|
||||
code: -32000,
|
||||
message: "internal error".to_owned(),
|
||||
data: Some(json!({
|
||||
"code": "internal_error",
|
||||
"detail": "cross-instance call cancelled",
|
||||
})),
|
||||
});
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert_eq!(err.detail, "cross-instance call cancelled");
|
||||
}
|
||||
367
crates/common/xai-computer-hub-core/tests/local_transport.rs
Normal file
367
crates/common/xai-computer-hub-core/tests/local_transport.rs
Normal file
|
|
@ -0,0 +1,367 @@
|
|||
//! `LocalTransport` end-to-end coverage. Verifies that the transport
|
||||
//! resolves through the bound resolver, drives both blocking and
|
||||
//! streaming tools, and surfaces missing tools as `Terminal(NotFound)`.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::StreamExt;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use xai_computer_hub_core::{
|
||||
CompoundResolver, ConnectionCleanupReport, ErasedTool, LocalTransport, ResolvedTool,
|
||||
SessionCleanupReport, ToolHandle, ToolRegistry, ToolSessionBindOutcome,
|
||||
ToolSessionUnbindOutcome, Transport, TransportKind,
|
||||
};
|
||||
use xai_tool_protocol::{
|
||||
ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId,
|
||||
ToolRegistration, ToolServerRegistration, TransportKind as WireTransportKind, UserId,
|
||||
};
|
||||
use xai_tool_runtime::{
|
||||
SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolError, ToolProgress, ToolStream,
|
||||
ToolStreamItem, terminal_only, with_progress,
|
||||
};
|
||||
use xai_tool_types::ToolDescription;
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
struct EchoArgs {
|
||||
payload: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct EchoTool;
|
||||
|
||||
impl Tool for EchoTool {
|
||||
type Args = EchoArgs;
|
||||
type Output = serde_json::Value;
|
||||
|
||||
fn id(&self) -> ToolId {
|
||||
ToolId::new("echo").expect("tool id")
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &::xai_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new("echo", "Echoes its input.")
|
||||
}
|
||||
|
||||
async fn run(
|
||||
&self,
|
||||
_ctx: ToolCallContext,
|
||||
args: Self::Args,
|
||||
) -> Result<Self::Output, ToolError> {
|
||||
Ok(serde_json::json!({ "echoed": args.payload }))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct StreamerTool;
|
||||
|
||||
impl Tool for StreamerTool {
|
||||
type Args = EchoArgs;
|
||||
type Output = serde_json::Value;
|
||||
|
||||
fn id(&self) -> ToolId {
|
||||
ToolId::new("streamer").expect("tool id")
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &::xai_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new("streamer", "Emits three progress chunks.")
|
||||
}
|
||||
|
||||
async fn execute(&self, _ctx: ToolCallContext, args: Self::Args) -> ToolStream<Self::Output> {
|
||||
let chunks = futures::stream::iter(vec![
|
||||
ToolProgress::Text {
|
||||
text: "tick".to_string(),
|
||||
},
|
||||
ToolProgress::Text {
|
||||
text: "tock".to_string(),
|
||||
},
|
||||
ToolProgress::Text {
|
||||
text: "boom".to_string(),
|
||||
},
|
||||
]);
|
||||
with_progress(chunks, async move {
|
||||
Ok(serde_json::json!({ "echoed": args.payload }))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type RegistryEntry = (ToolRegistration, Arc<dyn ToolHandle>);
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct InMemRegistry {
|
||||
entries: DashMap<(SessionId, ToolId), RegistryEntry>,
|
||||
}
|
||||
|
||||
impl InMemRegistry {
|
||||
fn install(&self, session: SessionId, tool_id: ToolId, handle: Arc<dyn ToolHandle>) {
|
||||
let registration = ToolRegistration {
|
||||
tool_id: tool_id.clone(),
|
||||
sessions: Some(vec![session.clone()]),
|
||||
user_id: UserId::new("alice").expect("user id"),
|
||||
server_id: None,
|
||||
description: handle.description(&xai_tool_runtime::ListToolsContext::default()),
|
||||
input_schema: None,
|
||||
capabilities: Some(handle.capabilities()),
|
||||
notification_schemas: None,
|
||||
transport_kind: WireTransportKind::Local,
|
||||
if_match_generation: None,
|
||||
metadata: None,
|
||||
};
|
||||
self.entries
|
||||
.insert((session, tool_id), (registration, handle));
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolRegistry for InMemRegistry {
|
||||
async fn register_tool(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolRegistration,
|
||||
) -> RegistrationOutcome {
|
||||
unreachable!("transport tests pre-populate via install()")
|
||||
}
|
||||
async fn register_server(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolServerRegistration,
|
||||
) -> Vec<RegistrationOutcome> {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_tool(&self, _connection_id: &ConnectionId, _tool: &ToolId) -> bool {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_server(&self, _connection_id: &ConnectionId, _server: &ServerId) -> usize {
|
||||
unreachable!()
|
||||
}
|
||||
async fn bind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionBindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unbind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionUnbindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn drop_connection(&self, _connection_id: &ConnectionId) -> ConnectionCleanupReport {
|
||||
ConnectionCleanupReport::default()
|
||||
}
|
||||
fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option<ResolvedTool> {
|
||||
let (registration, handle) = self
|
||||
.entries
|
||||
.get(&(session.clone(), tool.clone()))?
|
||||
.value()
|
||||
.clone();
|
||||
Some(ResolvedTool::Local {
|
||||
tool: handle,
|
||||
registration,
|
||||
})
|
||||
}
|
||||
fn list_tools(&self, _session: &SessionId, _mode: &ToolDefinitionMode) -> Vec<ToolDescription> {
|
||||
vec![]
|
||||
}
|
||||
fn list_servers(&self, _session: &SessionId) -> Vec<ServerSummary> {
|
||||
vec![]
|
||||
}
|
||||
fn search(&self, _session: &SessionId, _query: &str, _limit: usize) -> SearchSnapshot {
|
||||
SearchSnapshot {
|
||||
results: vec![],
|
||||
total_hidden_tools: 0,
|
||||
is_ready: true,
|
||||
}
|
||||
}
|
||||
async fn unregister_session(&self, _session: &SessionId) -> SessionCleanupReport {
|
||||
SessionCleanupReport::default()
|
||||
}
|
||||
fn tool_sessions(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
) -> std::collections::HashSet<SessionId> {
|
||||
std::collections::HashSet::new()
|
||||
}
|
||||
|
||||
fn list_servers_for_user(
|
||||
&self,
|
||||
_user_id: &xai_tool_protocol::UserId,
|
||||
) -> Vec<xai_computer_hub_core::registry::ServerRecord> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn get_server_record(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
) -> Option<xai_computer_hub_core::registry::ServerRecord> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("tool id")
|
||||
}
|
||||
|
||||
fn uid(s: &str) -> UserId {
|
||||
UserId::new(s).expect("user id")
|
||||
}
|
||||
|
||||
async fn collect(
|
||||
stream: &mut ToolStream<xai_tool_runtime::TypedToolOutput>,
|
||||
) -> Vec<ToolStreamItem<xai_tool_runtime::TypedToolOutput>> {
|
||||
let mut items = Vec::new();
|
||||
while let Some(item) = stream.next().await {
|
||||
items.push(item);
|
||||
}
|
||||
items
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatches_blocking_tool_to_terminal_value() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
registry.install(
|
||||
sid("sess-1"),
|
||||
tid("echo"),
|
||||
Arc::new(ErasedTool::new(EchoTool)),
|
||||
);
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1"));
|
||||
|
||||
let mut stream = transport
|
||||
.call(
|
||||
tid("echo"),
|
||||
serde_json::json!({"payload": "hi"}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let items = collect(&mut stream).await;
|
||||
assert_eq!(items.len(), 1);
|
||||
match &items[0] {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert_eq!(typed.value, serde_json::json!({"echoed": "hi"}));
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatches_streaming_tool_with_three_progress_then_terminal() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
registry.install(
|
||||
sid("sess-1"),
|
||||
tid("streamer"),
|
||||
Arc::new(ErasedTool::new(StreamerTool)),
|
||||
);
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1"));
|
||||
|
||||
let mut stream = transport
|
||||
.call(
|
||||
tid("streamer"),
|
||||
serde_json::json!({"payload": "hi"}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let items = collect(&mut stream).await;
|
||||
assert_eq!(items.len(), 4);
|
||||
for item in &items[..3] {
|
||||
assert!(matches!(item, ToolStreamItem::Progress(_)));
|
||||
}
|
||||
assert!(matches!(items[3], ToolStreamItem::Terminal(Ok(_))));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_tool_resolves_as_terminal_not_found() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1"));
|
||||
|
||||
let mut stream = transport
|
||||
.call(
|
||||
tid("ghost"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let items = collect(&mut stream).await;
|
||||
assert_eq!(items.len(), 1);
|
||||
match &items[0] {
|
||||
ToolStreamItem::Terminal(Err(e)) if e.kind == xai_tool_runtime::ToolErrorKind::NotFound => {
|
||||
assert!(
|
||||
e.detail.contains("ghost"),
|
||||
"detail should mention tool id: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(Err(NotFound)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_arguments_surface_as_terminal_error() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
registry.install(
|
||||
sid("sess-1"),
|
||||
tid("echo"),
|
||||
Arc::new(ErasedTool::new(EchoTool)),
|
||||
);
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1"));
|
||||
|
||||
let mut stream = transport
|
||||
.call(
|
||||
tid("echo"),
|
||||
// Missing required `payload` field.
|
||||
serde_json::json!({}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let items = collect(&mut stream).await;
|
||||
assert_eq!(items.len(), 1);
|
||||
match &items[0] {
|
||||
ToolStreamItem::Terminal(Err(e))
|
||||
if e.kind == xai_tool_runtime::ToolErrorKind::InvalidArguments => {}
|
||||
other => panic!("expected Terminal(Err(InvalidArguments)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn authorize_returns_bound_principal_with_invoke_scope() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1"));
|
||||
let principal = transport.authorize().await.expect("authorize");
|
||||
assert_eq!(principal.user_id, uid("alice"));
|
||||
assert!(principal.authorizes_session(&sid("sess-1")));
|
||||
assert!(principal.has_scope(xai_computer_hub_core::LOCAL_INVOKE_SCOPE));
|
||||
assert_eq!(transport.kind(), TransportKind::Local);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unused_helpers_silenced() {
|
||||
// `terminal_only` is re-exported for adapter authors; touch it here so
|
||||
// a future refactor that drops the import does not silently break the
|
||||
// re-export surface.
|
||||
let _: ToolStream<serde_json::Value> = terminal_only(Ok(serde_json::Value::Null));
|
||||
}
|
||||
727
crates/common/xai-computer-hub-core/tests/remote_proxy.rs
Normal file
727
crates/common/xai-computer-hub-core/tests/remote_proxy.rs
Normal file
|
|
@ -0,0 +1,727 @@
|
|||
//! `RemoteToolProxy` and `RemoteTransport` coverage. A channel-backed
|
||||
//! mock `ConnectionClient` lets the test inspect outgoing frames and
|
||||
//! drive synthetic responses + progress without any tokio I/O.
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::StreamExt;
|
||||
use futures::channel::{mpsc, oneshot};
|
||||
use futures::stream::BoxStream;
|
||||
|
||||
use xai_computer_hub_core::{
|
||||
ConnectionClient, RemoteToolProxy, RemoteTransport, ToolHandle, Transport, TransportKind,
|
||||
};
|
||||
use xai_tool_protocol::{
|
||||
JsonRpcError, JsonRpcId, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, JsonRpcVersion,
|
||||
Method, ResponseOutcome, SessionId, ToolCallId, ToolCallParams, ToolCallProgressFrame,
|
||||
ToolCallResult, ToolCapabilities, ToolErrorWire, ToolId, ToolOutputWire, UserId,
|
||||
};
|
||||
use xai_tool_runtime::{
|
||||
ContentBlock, ToolCallContext, ToolError, ToolOutput, ToolProgress, ToolStreamItem,
|
||||
};
|
||||
use xai_tool_types::ToolDescription;
|
||||
|
||||
/// Programmable `ConnectionClient`. Each request gets a pre-staged
|
||||
/// response; progress frames are pushed through per-call senders.
|
||||
#[derive(Debug, Default)]
|
||||
struct MockConnection {
|
||||
/// Senders keyed by `tool_call_id`. Pulled out of the inner state so
|
||||
/// per-call subscription touches a lock-free DashMap rather than the
|
||||
/// shared Mutex that guards the rest of the queue + capture state.
|
||||
progress_senders: DashMap<ToolCallId, mpsc::UnboundedSender<ToolCallProgressFrame>>,
|
||||
/// Three-Vec state guarded by one Mutex. The lock provides atomic
|
||||
/// pop-from-`responses` + push-to-`captured_requests` semantics that
|
||||
/// some tests rely on.
|
||||
inner: Mutex<MockState>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct MockState {
|
||||
/// FIFO queue of responses to return for each `request` call.
|
||||
responses: Vec<MockResponse>,
|
||||
/// Captured outgoing requests so tests can assert on them.
|
||||
captured_requests: Vec<JsonRpcRequest>,
|
||||
/// Captured one-way notifications.
|
||||
captured_notifications: Vec<JsonRpcNotification>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for MockState {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("MockState")
|
||||
.field("responses_len", &self.responses.len())
|
||||
.field("captured_reqs", &self.captured_requests.len())
|
||||
.field("captured_notifs", &self.captured_notifications.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
enum MockResponse {
|
||||
Ok(serde_json::Value),
|
||||
Err(JsonRpcError),
|
||||
/// Resolves a oneshot when the request arrives so the test can
|
||||
/// release progress before allowing the response.
|
||||
Gated {
|
||||
gate: oneshot::Receiver<()>,
|
||||
body: serde_json::Value,
|
||||
},
|
||||
/// Fail at the transport layer (e.g. socket dropped).
|
||||
Network(String),
|
||||
}
|
||||
|
||||
impl MockConnection {
|
||||
fn enqueue_ok(&self, body: serde_json::Value) {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.responses
|
||||
.push(MockResponse::Ok(body));
|
||||
}
|
||||
|
||||
fn enqueue_err(&self, code: i32, message: impl Into<String>, data: Option<serde_json::Value>) {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.responses
|
||||
.push(MockResponse::Err(JsonRpcError {
|
||||
code,
|
||||
message: message.into(),
|
||||
data,
|
||||
}));
|
||||
}
|
||||
|
||||
fn enqueue_gated(&self, gate: oneshot::Receiver<()>, body: serde_json::Value) {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.responses
|
||||
.push(MockResponse::Gated { gate, body });
|
||||
}
|
||||
|
||||
fn enqueue_network_failure(&self, message: impl Into<String>) {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.responses
|
||||
.push(MockResponse::Network(message.into()));
|
||||
}
|
||||
|
||||
fn last_request(&self) -> Option<JsonRpcRequest> {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.captured_requests
|
||||
.last()
|
||||
.cloned()
|
||||
}
|
||||
|
||||
fn captured_request_count(&self) -> usize {
|
||||
self.inner.lock().expect("mutex").captured_requests.len()
|
||||
}
|
||||
|
||||
fn push_progress(&self, tool_call_id: &ToolCallId, frame: ToolCallProgressFrame) {
|
||||
if let Some(tx) = self.progress_senders.get(tool_call_id) {
|
||||
let _ = tx.value().unbounded_send(frame);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ConnectionClient for MockConnection {
|
||||
async fn request(&self, request: JsonRpcRequest) -> Result<JsonRpcResponse, ToolError> {
|
||||
let response = {
|
||||
let mut guard = self.inner.lock().expect("mutex");
|
||||
guard.captured_requests.push(request.clone());
|
||||
if guard.responses.is_empty() {
|
||||
return Err(ToolError::custom(
|
||||
"mock_response_missing",
|
||||
"no response staged",
|
||||
));
|
||||
}
|
||||
guard.responses.remove(0)
|
||||
};
|
||||
match response {
|
||||
MockResponse::Ok(body) => Ok(JsonRpcResponse {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id: request.id,
|
||||
session_id: request.session_id,
|
||||
outcome: ResponseOutcome::Result(body),
|
||||
}),
|
||||
MockResponse::Err(err) => Ok(JsonRpcResponse {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id: request.id,
|
||||
session_id: request.session_id,
|
||||
outcome: ResponseOutcome::Error(err),
|
||||
}),
|
||||
MockResponse::Gated { gate, body } => {
|
||||
let _ = gate.await;
|
||||
Ok(JsonRpcResponse {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id: request.id,
|
||||
session_id: request.session_id,
|
||||
outcome: ResponseOutcome::Result(body),
|
||||
})
|
||||
}
|
||||
MockResponse::Network(msg) => Err(ToolError::network_error(msg)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn subscribe_progress(
|
||||
&self,
|
||||
tool_call_id: ToolCallId,
|
||||
) -> BoxStream<'static, ToolCallProgressFrame> {
|
||||
let (tx, rx) = mpsc::unbounded();
|
||||
self.progress_senders.insert(tool_call_id, tx);
|
||||
rx.boxed()
|
||||
}
|
||||
|
||||
async fn notify(&self, notification: JsonRpcNotification) -> Result<(), ToolError> {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.captured_notifications
|
||||
.push(notification);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("tool id")
|
||||
}
|
||||
|
||||
fn uid(s: &str) -> UserId {
|
||||
UserId::new(s).expect("user id")
|
||||
}
|
||||
|
||||
fn description_for(name: &str) -> ToolDescription {
|
||||
ToolDescription::new(name, format!("desc for {name}"))
|
||||
}
|
||||
|
||||
fn ok_call_result(call_id: &ToolCallId, output: ToolOutputWire) -> serde_json::Value {
|
||||
serde_json::to_value(ToolCallResult {
|
||||
tool_call_id: call_id.clone(),
|
||||
output,
|
||||
follow_ups: vec![],
|
||||
reminders: vec![],
|
||||
chat_completion_output: None,
|
||||
})
|
||||
.expect("serialise call result")
|
||||
}
|
||||
|
||||
fn ok_call_result_with_cco(
|
||||
call_id: &ToolCallId,
|
||||
output: ToolOutputWire,
|
||||
chat_completion_output: serde_json::Value,
|
||||
) -> serde_json::Value {
|
||||
serde_json::to_value(ToolCallResult {
|
||||
tool_call_id: call_id.clone(),
|
||||
output,
|
||||
follow_ups: vec![],
|
||||
reminders: vec![],
|
||||
chat_completion_output: Some(chat_completion_output),
|
||||
})
|
||||
.expect("serialise call result")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_sends_well_formed_tool_call_request() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
conn.enqueue_ok(ok_call_result(
|
||||
&call_id,
|
||||
ToolOutputWire::Text("hello".to_string()),
|
||||
));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!({"k": "v"})).await;
|
||||
while stream.next().await.is_some() {}
|
||||
let req = conn.last_request().expect("captured request");
|
||||
assert_eq!(req.method, Method::ToolCallRequest.as_wire_str());
|
||||
let params: ToolCallParams = serde_json::from_value(req.params).expect("decode params");
|
||||
assert_eq!(params.tool_id, tid("foo"));
|
||||
assert_eq!(params.tool_call_id, call_id);
|
||||
assert_eq!(params.arguments, serde_json::json!({"k": "v"}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn progress_then_terminal_orders_correctly() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
let (gate_tx, gate_rx) = oneshot::channel();
|
||||
conn.enqueue_gated(
|
||||
gate_rx,
|
||||
ok_call_result(&call_id, ToolOutputWire::Text("done".to_string())),
|
||||
);
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
|
||||
// Push two progress frames before the terminal is unblocked.
|
||||
conn.push_progress(
|
||||
&call_id,
|
||||
ToolCallProgressFrame {
|
||||
tool_call_id: call_id.clone(),
|
||||
kind: "log".to_string(),
|
||||
body: serde_json::json!({"text": "tick"}),
|
||||
dropped_count: None,
|
||||
},
|
||||
);
|
||||
conn.push_progress(
|
||||
&call_id,
|
||||
ToolCallProgressFrame {
|
||||
tool_call_id: call_id.clone(),
|
||||
kind: "log".to_string(),
|
||||
body: serde_json::json!({"text": "tock"}),
|
||||
dropped_count: None,
|
||||
},
|
||||
);
|
||||
|
||||
let first = stream.next().await.expect("first item");
|
||||
let second = stream.next().await.expect("second item");
|
||||
match (&first, &second) {
|
||||
(ToolStreamItem::Progress(p1), ToolStreamItem::Progress(p2)) => {
|
||||
match p1 {
|
||||
ToolProgress::Custom { subkind, payload } => {
|
||||
assert_eq!(subkind, "log");
|
||||
assert_eq!(payload, &serde_json::json!({"text": "tick"}));
|
||||
}
|
||||
other => panic!("expected Custom progress, got {other:?}"),
|
||||
}
|
||||
match p2 {
|
||||
ToolProgress::Custom { subkind, .. } => assert_eq!(subkind, "log"),
|
||||
other => panic!("expected Custom progress, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected two Progress items, got {other:?}"),
|
||||
}
|
||||
|
||||
// Release the response and consume the terminal.
|
||||
let _ = gate_tx.send(());
|
||||
let terminal = stream.next().await.expect("terminal");
|
||||
match terminal {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert_eq!(typed.value, serde_json::json!("done"));
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
assert!(stream.next().await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn json_rpc_error_response_decodes_into_tool_error() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let wire = ToolErrorWire::ToolNotFound {
|
||||
tool_id: tid("foo"),
|
||||
};
|
||||
conn.enqueue_err(
|
||||
-32011,
|
||||
"tool not found",
|
||||
Some(serde_json::to_value(&wire).unwrap()),
|
||||
);
|
||||
let mut stream = proxy
|
||||
.execute(ToolCallContext::default(), serde_json::json!(null))
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == xai_tool_runtime::ToolErrorKind::NotFound =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("foo"),
|
||||
"detail should mention tool id: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(NotFound), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn network_failure_surfaces_as_terminal_network_error() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
conn.enqueue_network_failure("socket closed");
|
||||
let mut stream = proxy
|
||||
.execute(ToolCallContext::default(), serde_json::json!(null))
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == xai_tool_runtime::ToolErrorKind::NetworkError =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("socket closed"),
|
||||
"detail should mention cause: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(NetworkError), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mcp_output_re_serialises_into_blocks_value() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
let blocks = vec![xai_tool_protocol::McpBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
}];
|
||||
conn.enqueue_ok(ok_call_result(&call_id, ToolOutputWire::Mcp { blocks }));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
// The wire blocks round-trip through ContentBlock; assert the
|
||||
// text value survives the transformation.
|
||||
let blocks_value = typed
|
||||
.value
|
||||
.get("blocks")
|
||||
.and_then(|v| v.as_array())
|
||||
.cloned()
|
||||
.expect("blocks array");
|
||||
assert_eq!(blocks_value.len(), 1);
|
||||
let block: ContentBlock =
|
||||
serde_json::from_value(blocks_value[0].clone()).expect("decode runtime block");
|
||||
match block {
|
||||
ContentBlock::Text { text } => assert_eq!(text, "hello"),
|
||||
other => panic!("expected Text block, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_carries_chat_completion_output_from_wire() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("bash"),
|
||||
sid("sess-1"),
|
||||
description_for("bash"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
let cco = serde_json::json!({
|
||||
"result": {
|
||||
"sender": "assistant",
|
||||
"message": "",
|
||||
"code_execution_result": {
|
||||
"stdout": "hi\n",
|
||||
"stderr": "",
|
||||
"exit_code": 0,
|
||||
"command_timed_out": false
|
||||
}
|
||||
}
|
||||
});
|
||||
conn.enqueue_ok(ok_call_result_with_cco(
|
||||
&call_id,
|
||||
ToolOutputWire::Json(serde_json::json!({"stdout": "hi\n"})),
|
||||
cco,
|
||||
));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
let response = typed
|
||||
.chat_completion_output()
|
||||
.expect("chat completion output survives the wire");
|
||||
let completion = response.result.expect("completion result present");
|
||||
let exec = completion
|
||||
.code_execution_result
|
||||
.expect("code execution result present");
|
||||
assert_eq!(exec.stdout, "hi\n");
|
||||
assert_eq!(exec.exit_code, 0);
|
||||
assert!(!exec.command_timed_out);
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_without_chat_completion_output_is_none() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("bash"),
|
||||
sid("sess-1"),
|
||||
description_for("bash"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
conn.enqueue_ok(ok_call_result(
|
||||
&call_id,
|
||||
ToolOutputWire::Json(serde_json::json!({"stdout": "hi\n"})),
|
||||
));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert!(typed.chat_completion_output().is_none());
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn malformed_inner_chat_completion_output_degrades_to_none() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("bash"),
|
||||
sid("sess-1"),
|
||||
description_for("bash"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
conn.enqueue_ok(ok_call_result_with_cco(
|
||||
&call_id,
|
||||
ToolOutputWire::Json(serde_json::json!({"stdout": "hi\n"})),
|
||||
serde_json::json!({"result": "not-a-completion-object"}),
|
||||
));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert_eq!(typed.value, serde_json::json!({"stdout": "hi\n"}));
|
||||
assert!(typed.chat_completion_output().is_none());
|
||||
}
|
||||
other => panic!("expected Terminal(Ok) with degraded cco, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bare_non_enveloped_success_body_passes_through() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("bash"),
|
||||
sid("sess-1"),
|
||||
description_for("bash"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let ctx = ToolCallContext::new(ToolCallId::new_v7());
|
||||
conn.enqueue_ok(serde_json::json!({"stdout": "hi\n"}));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert_eq!(typed.value, serde_json::json!({"stdout": "hi\n"}));
|
||||
assert!(typed.chat_completion_output().is_none());
|
||||
}
|
||||
other => panic!("expected Terminal(Ok) passthrough, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn malformed_envelope_with_tool_call_id_surfaces_decode_error() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("bash"),
|
||||
sid("sess-1"),
|
||||
description_for("bash"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let ctx = ToolCallContext::new(ToolCallId::new_v7());
|
||||
conn.enqueue_ok(serde_json::json!({"tool_call_id": "call_x", "output": 123}));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == xai_tool_runtime::ToolErrorKind::Custom =>
|
||||
{
|
||||
let code = e
|
||||
.details
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("code"))
|
||||
.and_then(|c| c.as_str());
|
||||
assert_eq!(code, Some("response_decoding"), "error: {e:?}");
|
||||
}
|
||||
other => panic!("expected Terminal(Err) decode failure, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_transport_call_dispatches_via_connection() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let transport = RemoteTransport::new(conn.clone(), sid("sess-1"), uid("alice"));
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
conn.enqueue_ok(ok_call_result(&call_id, ToolOutputWire::Text("hi".into())));
|
||||
let mut stream = transport
|
||||
.call(tid("foo"), serde_json::json!({"k": "v"}), ctx)
|
||||
.await;
|
||||
let _ = stream.next().await;
|
||||
assert_eq!(conn.captured_request_count(), 1);
|
||||
assert_eq!(transport.kind(), TransportKind::Remote);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_transport_authorize_returns_bound_principal() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let transport = RemoteTransport::new(conn, sid("sess-1"), uid("alice"));
|
||||
let principal = transport.authorize().await.expect("authorize");
|
||||
assert_eq!(principal.user_id, uid("alice"));
|
||||
assert!(principal.authorizes_session(&sid("sess-1")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_subscribe_happens_before_request_send() {
|
||||
// Locks in BOTH halves of the subscribe-before-send contract:
|
||||
// 1. the subscription IS active by the time `execute` returns;
|
||||
// 2. the request HAS NOT been sent yet at that point.
|
||||
// A future refactor that eagerly sent the request inside
|
||||
// `execute` would still satisfy (1) but would break (2).
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
conn.enqueue_ok(ok_call_result(&call_id, ToolOutputWire::Text("ok".into())));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
{
|
||||
// The DashMap subscription read and the captured-requests check
|
||||
// are individually atomic. Single-threaded `#[tokio::test]`
|
||||
// execution means no other task can mutate either between the
|
||||
// two checks, so the pair is observationally simultaneous.
|
||||
assert!(
|
||||
conn.progress_senders.contains_key(&call_id),
|
||||
"subscription must be active before request send"
|
||||
);
|
||||
let guard = conn.inner.lock().expect("mutex");
|
||||
assert!(
|
||||
guard.captured_requests.is_empty(),
|
||||
"request must not be sent before stream is polled"
|
||||
);
|
||||
}
|
||||
// Polling the stream is what actually drives the request future,
|
||||
// so the captured-requests vec only fills in once we start consuming.
|
||||
while stream.next().await.is_some() {}
|
||||
{
|
||||
let guard = conn.inner.lock().expect("mutex");
|
||||
assert_eq!(
|
||||
guard.captured_requests.len(),
|
||||
1,
|
||||
"request must have been sent during stream polling"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn notify_round_trips_through_connection_client() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let notification = JsonRpcNotification {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
session_id: Some(sid("sess-1")),
|
||||
seq: None,
|
||||
method: Method::Hook.as_wire_str().to_string(),
|
||||
params: serde_json::json!({
|
||||
"session_id": "sess-1",
|
||||
"tool_id": "foo",
|
||||
"call_id": "call-1",
|
||||
"event": { "type": "Cancel" }
|
||||
}),
|
||||
};
|
||||
let trait_handle: &dyn ConnectionClient = conn.as_ref();
|
||||
trait_handle
|
||||
.notify(notification.clone())
|
||||
.await
|
||||
.expect("notify succeeds");
|
||||
let guard = conn.inner.lock().expect("mutex");
|
||||
assert_eq!(guard.captured_notifications.len(), 1);
|
||||
let captured = &guard.captured_notifications[0];
|
||||
assert_eq!(captured.method, Method::Hook.as_wire_str());
|
||||
assert_eq!(captured.session_id, Some(sid("sess-1")));
|
||||
assert_eq!(
|
||||
captured.params.get("event").and_then(|v| v.get("type")),
|
||||
Some(&serde_json::Value::String("Cancel".to_string()))
|
||||
);
|
||||
assert_eq!(captured, ¬ification);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn json_rpc_id_is_unique_per_call() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let transport = RemoteTransport::new(conn.clone(), sid("sess-1"), uid("alice"));
|
||||
let call_a = ToolCallId::new_v7();
|
||||
let call_b = ToolCallId::new_v7();
|
||||
conn.enqueue_ok(ok_call_result(&call_a, ToolOutputWire::Text("a".into())));
|
||||
conn.enqueue_ok(ok_call_result(&call_b, ToolOutputWire::Text("b".into())));
|
||||
|
||||
let mut s1 = transport
|
||||
.call(
|
||||
tid("foo"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::new(call_a.clone()),
|
||||
)
|
||||
.await;
|
||||
while s1.next().await.is_some() {}
|
||||
let mut s2 = transport
|
||||
.call(
|
||||
tid("foo"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::new(call_b.clone()),
|
||||
)
|
||||
.await;
|
||||
while s2.next().await.is_some() {}
|
||||
|
||||
let guard = conn.inner.lock().expect("mutex");
|
||||
assert_eq!(guard.captured_requests.len(), 2);
|
||||
let id_a = match &guard.captured_requests[0].id {
|
||||
JsonRpcId::String(s) => s.clone(),
|
||||
JsonRpcId::Number(n) => n.to_string(),
|
||||
};
|
||||
let id_b = match &guard.captured_requests[1].id {
|
||||
JsonRpcId::String(s) => s.clone(),
|
||||
JsonRpcId::Number(n) => n.to_string(),
|
||||
};
|
||||
assert_ne!(id_a, id_b, "envelope ids must differ across calls");
|
||||
}
|
||||
607
crates/common/xai-computer-hub-core/tests/tool_registry.rs
Normal file
607
crates/common/xai-computer-hub-core/tests/tool_registry.rs
Normal file
|
|
@ -0,0 +1,607 @@
|
|||
//! `ToolRegistry` trait coverage via a per-test mock backed by `DashMap`
|
||||
//! — lock-free per-key concurrent access mirrors the production
|
||||
//! direction even at the test layer. The mock implements the
|
||||
//! connection-scoped `ToolRegistry` trait surface.
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use dashmap::DashMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use xai_computer_hub_core::{
|
||||
ConnectionCleanupReport, ErasedTool, ResolvedTool, SessionCleanupReport, ToolHandle,
|
||||
ToolRegistry, ToolSessionBindOutcome, ToolSessionUnbindOutcome, resolver::CompoundResolver,
|
||||
};
|
||||
use xai_tool_protocol::{
|
||||
ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId,
|
||||
ToolRegistration, ToolServerRegistration, TransportKind, UserId,
|
||||
};
|
||||
use xai_tool_runtime::{SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolError};
|
||||
use xai_tool_types::ToolDescription;
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
struct EmptyArgs {}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct StubTool {
|
||||
id: ToolId,
|
||||
}
|
||||
|
||||
impl Tool for StubTool {
|
||||
type Args = EmptyArgs;
|
||||
type Output = serde_json::Value;
|
||||
|
||||
fn id(&self) -> ToolId {
|
||||
self.id.clone()
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &::xai_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new(self.id.as_str(), format!("stub for {}", self.id))
|
||||
}
|
||||
|
||||
async fn run(
|
||||
&self,
|
||||
_ctx: ToolCallContext,
|
||||
_args: Self::Args,
|
||||
) -> Result<Self::Output, ToolError> {
|
||||
unreachable!("registry tests do not exercise execution")
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct MockEntry {
|
||||
registration: ToolRegistration,
|
||||
sessions: HashSet<SessionId>,
|
||||
}
|
||||
|
||||
/// Mock registry. Last-write-wins on duplicate registrations within a
|
||||
/// `(connection, tool_id)` slot — pinned here so the trait contract has
|
||||
/// a clear test fixture.
|
||||
#[derive(Debug, Default)]
|
||||
struct MockRegistry {
|
||||
entries: DashMap<(ConnectionId, ToolId), MockEntry>,
|
||||
by_session: DashMap<(SessionId, ToolId), ConnectionId>,
|
||||
handles: DashMap<ToolId, Arc<dyn ToolHandle>>,
|
||||
}
|
||||
|
||||
impl MockRegistry {
|
||||
fn install_handle(&self, tool: Arc<dyn ToolHandle>) {
|
||||
self.handles.insert(tool.id(), tool);
|
||||
}
|
||||
}
|
||||
|
||||
fn build_registration(tool: &ToolId, sessions: &[SessionId]) -> ToolRegistration {
|
||||
ToolRegistration {
|
||||
tool_id: tool.clone(),
|
||||
sessions: Some(sessions.to_vec()),
|
||||
user_id: UserId::new("alice").expect("valid user id"),
|
||||
server_id: None,
|
||||
description: ToolDescription::new(tool.as_str(), format!("desc for {tool}")),
|
||||
input_schema: None,
|
||||
capabilities: None,
|
||||
notification_schemas: None,
|
||||
transport_kind: TransportKind::Local,
|
||||
if_match_generation: None,
|
||||
metadata: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolRegistry for MockRegistry {
|
||||
async fn register_tool(
|
||||
&self,
|
||||
connection_id: ConnectionId,
|
||||
reg: ToolRegistration,
|
||||
) -> RegistrationOutcome {
|
||||
let key = (connection_id.clone(), reg.tool_id.clone());
|
||||
let sessions: HashSet<SessionId> = reg
|
||||
.sessions
|
||||
.as_ref()
|
||||
.map(|v| v.iter().cloned().collect())
|
||||
.unwrap_or_default();
|
||||
let updated = self
|
||||
.entries
|
||||
.insert(
|
||||
key,
|
||||
MockEntry {
|
||||
registration: reg.clone(),
|
||||
sessions: sessions.clone(),
|
||||
},
|
||||
)
|
||||
.is_some();
|
||||
for session in &sessions {
|
||||
self.by_session.insert(
|
||||
(session.clone(), reg.tool_id.clone()),
|
||||
connection_id.clone(),
|
||||
);
|
||||
}
|
||||
if updated {
|
||||
RegistrationOutcome::Updated {
|
||||
tool_id: reg.tool_id,
|
||||
generation: 1,
|
||||
}
|
||||
} else {
|
||||
RegistrationOutcome::Registered {
|
||||
tool_id: reg.tool_id,
|
||||
generation: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn register_server(
|
||||
&self,
|
||||
connection_id: ConnectionId,
|
||||
reg: ToolServerRegistration,
|
||||
) -> Vec<RegistrationOutcome> {
|
||||
let mut outcomes = Vec::with_capacity(reg.tools.len());
|
||||
for tool in reg.tools {
|
||||
let tool_id = tool
|
||||
.derive_tool_id()
|
||||
.expect("test descriptions have valid tool ids");
|
||||
let registration = ToolRegistration {
|
||||
tool_id: tool_id.clone(),
|
||||
sessions: reg.sessions.clone(),
|
||||
user_id: reg.user_id.clone(),
|
||||
server_id: Some(reg.server_id.clone()),
|
||||
description: tool.description,
|
||||
input_schema: tool.input_schema,
|
||||
capabilities: tool.capabilities,
|
||||
notification_schemas: tool.notification_schemas,
|
||||
transport_kind: TransportKind::Remote,
|
||||
if_match_generation: None,
|
||||
metadata: None,
|
||||
};
|
||||
outcomes.push(
|
||||
self.register_tool(connection_id.clone(), registration)
|
||||
.await,
|
||||
);
|
||||
}
|
||||
outcomes
|
||||
}
|
||||
|
||||
async fn unregister_tool(&self, connection_id: &ConnectionId, tool: &ToolId) -> bool {
|
||||
let Some((_, removed)) = self.entries.remove(&(connection_id.clone(), tool.clone())) else {
|
||||
return false;
|
||||
};
|
||||
for session in &removed.sessions {
|
||||
self.by_session
|
||||
.remove_if(&(session.clone(), tool.clone()), |_, owner| {
|
||||
owner == connection_id
|
||||
});
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn unregister_server(&self, connection_id: &ConnectionId, server: &ServerId) -> usize {
|
||||
let to_remove: Vec<ToolId> = self
|
||||
.entries
|
||||
.iter()
|
||||
.filter(|r| {
|
||||
r.key().0 == *connection_id
|
||||
&& r.value().registration.server_id.as_ref() == Some(server)
|
||||
})
|
||||
.map(|r| r.key().1.clone())
|
||||
.collect();
|
||||
let mut removed = 0usize;
|
||||
for tool in to_remove {
|
||||
if self.unregister_tool(connection_id, &tool).await {
|
||||
removed += 1;
|
||||
}
|
||||
}
|
||||
removed
|
||||
}
|
||||
|
||||
async fn bind_tool_session(
|
||||
&self,
|
||||
connection_id: &ConnectionId,
|
||||
tool: &ToolId,
|
||||
session_id: &SessionId,
|
||||
) -> ToolSessionBindOutcome {
|
||||
let key = (connection_id.clone(), tool.clone());
|
||||
let Some(mut entry) = self.entries.get_mut(&key) else {
|
||||
return ToolSessionBindOutcome::UnknownTool;
|
||||
};
|
||||
if !entry.value_mut().sessions.insert(session_id.clone()) {
|
||||
return ToolSessionBindOutcome::AlreadyBound;
|
||||
}
|
||||
self.by_session
|
||||
.insert((session_id.clone(), tool.clone()), connection_id.clone());
|
||||
ToolSessionBindOutcome::Bound
|
||||
}
|
||||
|
||||
async fn unbind_tool_session(
|
||||
&self,
|
||||
connection_id: &ConnectionId,
|
||||
tool: &ToolId,
|
||||
session_id: &SessionId,
|
||||
) -> ToolSessionUnbindOutcome {
|
||||
let key = (connection_id.clone(), tool.clone());
|
||||
let Some(mut entry) = self.entries.get_mut(&key) else {
|
||||
return ToolSessionUnbindOutcome::UnknownTool;
|
||||
};
|
||||
if !entry.value_mut().sessions.remove(session_id) {
|
||||
return ToolSessionUnbindOutcome::NotBound;
|
||||
}
|
||||
self.by_session
|
||||
.remove_if(&(session_id.clone(), tool.clone()), |_, owner| {
|
||||
owner == connection_id
|
||||
});
|
||||
ToolSessionUnbindOutcome::Unbound
|
||||
}
|
||||
|
||||
async fn drop_connection(&self, connection_id: &ConnectionId) -> ConnectionCleanupReport {
|
||||
let to_remove: Vec<ToolId> = self
|
||||
.entries
|
||||
.iter()
|
||||
.filter(|r| r.key().0 == *connection_id)
|
||||
.map(|r| r.key().1.clone())
|
||||
.collect();
|
||||
let mut report = ConnectionCleanupReport::default();
|
||||
for tool in to_remove {
|
||||
if let Some((_, removed)) = self.entries.remove(&(connection_id.clone(), tool.clone()))
|
||||
{
|
||||
report.tools_dropped += 1;
|
||||
for session in removed.sessions {
|
||||
if self
|
||||
.by_session
|
||||
.remove_if(&(session, tool.clone()), |_, owner| owner == connection_id)
|
||||
.is_some()
|
||||
{
|
||||
report.session_bindings_cleared += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
report
|
||||
}
|
||||
|
||||
fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option<ResolvedTool> {
|
||||
let owner = self
|
||||
.by_session
|
||||
.get(&(session.clone(), tool.clone()))?
|
||||
.value()
|
||||
.clone();
|
||||
let entry = self.entries.get(&(owner, tool.clone()))?;
|
||||
let registration = entry.value().registration.clone();
|
||||
let handle = self.handles.get(tool)?.value().clone();
|
||||
match registration.transport_kind {
|
||||
TransportKind::Local => Some(ResolvedTool::Local {
|
||||
tool: handle,
|
||||
registration,
|
||||
}),
|
||||
TransportKind::Remote => Some(ResolvedTool::Remote {
|
||||
proxy: handle,
|
||||
registration,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn list_tools(&self, session: &SessionId, _mode: &ToolDefinitionMode) -> Vec<ToolDescription> {
|
||||
self.by_session
|
||||
.iter()
|
||||
.filter(|r| r.key().0 == *session)
|
||||
.filter_map(|r| {
|
||||
let owner = r.value().clone();
|
||||
let tool_id = r.key().1.clone();
|
||||
self.entries
|
||||
.get(&(owner, tool_id))
|
||||
.map(|e| e.value().registration.description.clone())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn list_servers(&self, session: &SessionId) -> Vec<ServerSummary> {
|
||||
let mut by_server: HashMap<ServerId, Vec<String>> = HashMap::new();
|
||||
for r in self.by_session.iter().filter(|r| r.key().0 == *session) {
|
||||
let owner = r.value().clone();
|
||||
let tool_id = r.key().1.clone();
|
||||
if let Some(entry) = self.entries.get(&(owner, tool_id)) {
|
||||
let reg = &entry.value().registration;
|
||||
let server = reg
|
||||
.server_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| ServerId::synthesize_for_tool(r.value(), ®.tool_id));
|
||||
by_server
|
||||
.entry(server)
|
||||
.or_default()
|
||||
.push(reg.tool_id.as_str().to_string());
|
||||
}
|
||||
}
|
||||
by_server
|
||||
.into_iter()
|
||||
.map(|(server, mut names)| {
|
||||
names.sort();
|
||||
ServerSummary {
|
||||
name: server.into_inner(),
|
||||
description: None,
|
||||
tool_names: names,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn search(&self, session: &SessionId, query: &str, limit: usize) -> SearchSnapshot {
|
||||
let matches: Vec<_> = self
|
||||
.by_session
|
||||
.iter()
|
||||
.filter(|r| r.key().0 == *session)
|
||||
.filter_map(|r| {
|
||||
let owner = r.value().clone();
|
||||
let tool_id = r.key().1.clone();
|
||||
let entry = self.entries.get(&(owner, tool_id))?;
|
||||
let reg = &entry.value().registration;
|
||||
if reg.tool_id.as_str().contains(query) {
|
||||
Some(xai_tool_runtime::ToolSearchResult {
|
||||
tool_name: reg.tool_id.as_str().to_string(),
|
||||
server_name: reg
|
||||
.server_id
|
||||
.as_ref()
|
||||
.map(|s| s.as_str().to_string())
|
||||
.unwrap_or_default(),
|
||||
description: reg.description.description.clone(),
|
||||
score: 1.0,
|
||||
parameters: vec![],
|
||||
input_schema: serde_json::Value::Null,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.take(limit)
|
||||
.collect();
|
||||
SearchSnapshot {
|
||||
results: matches,
|
||||
total_hidden_tools: 0,
|
||||
is_ready: true,
|
||||
}
|
||||
}
|
||||
|
||||
async fn unregister_session(&self, session: &SessionId) -> SessionCleanupReport {
|
||||
let pairs: Vec<(ToolId, ConnectionId)> = self
|
||||
.by_session
|
||||
.iter()
|
||||
.filter(|r| r.key().0 == *session)
|
||||
.map(|r| (r.key().1.clone(), r.value().clone()))
|
||||
.collect();
|
||||
let mut report = SessionCleanupReport::default();
|
||||
for (tool_id, owner) in pairs {
|
||||
self.by_session
|
||||
.remove_if(&(session.clone(), tool_id.clone()), |_, value| {
|
||||
value == &owner
|
||||
});
|
||||
if let Some(mut entry) = self.entries.get_mut(&(owner, tool_id)) {
|
||||
entry.value_mut().sessions.remove(session);
|
||||
report.tools_touched += 1;
|
||||
if entry.value().sessions.is_empty() {
|
||||
report.tools_left_orphaned += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
report
|
||||
}
|
||||
|
||||
fn tool_sessions(&self, connection_id: &ConnectionId, tool: &ToolId) -> HashSet<SessionId> {
|
||||
self.entries
|
||||
.get(&(connection_id.clone(), tool.clone()))
|
||||
.map(|r| r.value().sessions.clone())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn list_servers_for_user(
|
||||
&self,
|
||||
_user_id: &xai_tool_protocol::UserId,
|
||||
) -> Vec<xai_computer_hub_core::registry::ServerRecord> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn get_server_record(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
) -> Option<xai_computer_hub_core::registry::ServerRecord> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("valid session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("valid tool id")
|
||||
}
|
||||
|
||||
fn cid(s: &str) -> ConnectionId {
|
||||
ConnectionId::new(s).expect("valid connection id")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn register_then_find_returns_local_resolution() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
let outcome = reg
|
||||
.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
assert!(matches!(outcome, RegistrationOutcome::Registered { .. }));
|
||||
let resolved = reg
|
||||
.find_tool(&sid("sess-1"), &tid("foo"))
|
||||
.expect("registration found");
|
||||
match resolved {
|
||||
ResolvedTool::Local { registration, .. } => {
|
||||
assert_eq!(registration.tool_id, tid("foo"));
|
||||
assert!(
|
||||
registration
|
||||
.sessions
|
||||
.as_ref()
|
||||
.is_some_and(|s| s.contains(&sid("sess-1")))
|
||||
);
|
||||
}
|
||||
other => panic!("expected Local, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn find_in_other_session_returns_none() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
assert!(reg.find_tool(&sid("sess-2"), &tid("foo")).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn duplicate_registration_yields_updated_outcome() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
let first = reg
|
||||
.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
let second = reg
|
||||
.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
assert!(matches!(first, RegistrationOutcome::Registered { .. }));
|
||||
assert!(matches!(second, RegistrationOutcome::Updated { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unregister_tool_removes_only_that_entry() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("bar"), &[sid("sess-1")]))
|
||||
.await;
|
||||
assert!(reg.unregister_tool(&cid("c1"), &tid("foo")).await);
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none());
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("bar")).is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unregister_session_drops_session_binding_and_leaves_orphan_count() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
reg.register_tool(
|
||||
cid("c1"),
|
||||
build_registration(&tid("bar"), &[sid("sess-1"), sid("sess-2")]),
|
||||
)
|
||||
.await;
|
||||
let report = reg.unregister_session(&sid("sess-1")).await;
|
||||
assert_eq!(report.tools_touched, 2);
|
||||
// `foo` had only sess-1 → orphaned. `bar` had sess-2 left → not orphaned.
|
||||
assert_eq!(report.tools_left_orphaned, 1);
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none());
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("bar")).is_none());
|
||||
assert!(reg.find_tool(&sid("sess-2"), &tid("bar")).is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_tools_filters_by_session() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("bar"), &[sid("sess-2")]))
|
||||
.await;
|
||||
let s1 = reg.list_tools(&sid("sess-1"), &ToolDefinitionMode::Full);
|
||||
let s2 = reg.list_tools(&sid("sess-2"), &ToolDefinitionMode::Full);
|
||||
assert_eq!(s1.len(), 1);
|
||||
assert_eq!(s1[0].name, "foo");
|
||||
assert_eq!(s2.len(), 1);
|
||||
assert_eq!(s2[0].name, "bar");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_servers_groups_by_owning_server() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
let summaries = reg.list_servers(&sid("sess-1"));
|
||||
assert_eq!(summaries.len(), 1);
|
||||
assert_eq!(summaries[0].tool_count(), 1);
|
||||
assert_eq!(summaries[0].tool_names[0], "foo");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn search_returns_substring_matches() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foobar") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
reg.register_tool(
|
||||
cid("c1"),
|
||||
build_registration(&tid("foobar"), &[sid("sess-1")]),
|
||||
)
|
||||
.await;
|
||||
let snap = reg.search(&sid("sess-1"), "foo", 10);
|
||||
assert_eq!(snap.results.len(), 2);
|
||||
assert!(snap.is_ready);
|
||||
assert_eq!(snap.total_hidden_tools, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn registry_drives_compound_resolver() {
|
||||
let registry = Arc::new(MockRegistry::default());
|
||||
registry.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
registry
|
||||
.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
let resolver = CompoundResolver::local_only(registry as Arc<dyn ToolRegistry>);
|
||||
assert!(resolver.resolve(&sid("sess-1"), &tid("foo")).is_some());
|
||||
assert!(resolver.resolve(&sid("sess-1"), &tid("missing")).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bind_and_unbind_tool_session_round_trips_visibility() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[]))
|
||||
.await;
|
||||
// Empty sessions: tool is registered but unreachable.
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none());
|
||||
let outcome = reg
|
||||
.bind_tool_session(&cid("c1"), &tid("foo"), &sid("sess-1"))
|
||||
.await;
|
||||
assert_eq!(outcome, ToolSessionBindOutcome::Bound);
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_some());
|
||||
let again = reg
|
||||
.bind_tool_session(&cid("c1"), &tid("foo"), &sid("sess-1"))
|
||||
.await;
|
||||
assert_eq!(again, ToolSessionBindOutcome::AlreadyBound);
|
||||
let unbind = reg
|
||||
.unbind_tool_session(&cid("c1"), &tid("foo"), &sid("sess-1"))
|
||||
.await;
|
||||
assert_eq!(unbind, ToolSessionUnbindOutcome::Unbound);
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none());
|
||||
let unknown = reg
|
||||
.bind_tool_session(&cid("c1"), &tid("missing"), &sid("sess-1"))
|
||||
.await;
|
||||
assert_eq!(unknown, ToolSessionBindOutcome::UnknownTool);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drop_connection_releases_every_owned_tool() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
reg.register_tool(
|
||||
cid("c1"),
|
||||
build_registration(&tid("bar"), &[sid("sess-1"), sid("sess-2")]),
|
||||
)
|
||||
.await;
|
||||
let report = reg.drop_connection(&cid("c1")).await;
|
||||
assert_eq!(report.tools_dropped, 2);
|
||||
assert_eq!(report.session_bindings_cleared, 3);
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none());
|
||||
assert!(reg.find_tool(&sid("sess-2"), &tid("bar")).is_none());
|
||||
assert!(reg.tool_sessions(&cid("c1"), &tid("foo")).is_empty());
|
||||
}
|
||||
137
crates/common/xai-computer-hub-core/tests/transport_trait.rs
Normal file
137
crates/common/xai-computer-hub-core/tests/transport_trait.rs
Normal file
|
|
@ -0,0 +1,137 @@
|
|||
//! Behavioural coverage for the `Transport` trait, `Principal` builder,
|
||||
//! and `TransportKind` re-export.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use xai_computer_hub_core::{Principal, Transport, TransportKind};
|
||||
use xai_tool_protocol::{SessionId, ToolId, UserId};
|
||||
use xai_tool_runtime::{
|
||||
ToolCallContext, ToolError, ToolStream, ToolStreamItem, TypedToolOutput, terminal_only,
|
||||
};
|
||||
|
||||
fn uid(s: &str) -> UserId {
|
||||
UserId::new(s).expect("test user id")
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("test session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("test tool id")
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct EchoTransport {
|
||||
kind: TransportKind,
|
||||
user: UserId,
|
||||
session: SessionId,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Transport for EchoTransport {
|
||||
fn kind(&self) -> TransportKind {
|
||||
self.kind
|
||||
}
|
||||
|
||||
async fn authorize(&self) -> Result<Principal, ToolError> {
|
||||
Ok(Principal::new(self.user.clone())
|
||||
.with_session(self.session.clone())
|
||||
.with_scope("tool.invoke"))
|
||||
}
|
||||
|
||||
async fn call(
|
||||
&self,
|
||||
tool_id: ToolId,
|
||||
args: Value,
|
||||
_ctx: ToolCallContext,
|
||||
) -> ToolStream<TypedToolOutput> {
|
||||
terminal_only(Ok(TypedToolOutput::from_value(tool_id, args)))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn boxed_transport_compiles_and_dispatches() {
|
||||
let boxed: Box<dyn Transport> = Box::new(EchoTransport {
|
||||
kind: TransportKind::Local,
|
||||
user: uid("alice"),
|
||||
session: sid("sess-1"),
|
||||
});
|
||||
let mut stream = boxed
|
||||
.call(tid("echo"), json!({"k": "v"}), ToolCallContext::default())
|
||||
.await;
|
||||
let item = futures::StreamExt::next(&mut stream)
|
||||
.await
|
||||
.expect("at least one item");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => assert_eq!(typed.value, json!({"k": "v"})),
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn kind_distinguishes_local_and_remote() {
|
||||
let local = EchoTransport {
|
||||
kind: TransportKind::Local,
|
||||
user: uid("alice"),
|
||||
session: sid("sess-1"),
|
||||
};
|
||||
let remote = EchoTransport {
|
||||
kind: TransportKind::Remote,
|
||||
user: uid("alice"),
|
||||
session: sid("sess-1"),
|
||||
};
|
||||
assert_eq!(local.kind(), TransportKind::Local);
|
||||
assert_eq!(remote.kind(), TransportKind::Remote);
|
||||
assert_ne!(local.kind(), remote.kind());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn authorize_returns_bound_principal() {
|
||||
let t = EchoTransport {
|
||||
kind: TransportKind::Local,
|
||||
user: uid("alice"),
|
||||
session: sid("sess-1"),
|
||||
};
|
||||
let principal = t.authorize().await.expect("authorize succeeds");
|
||||
assert_eq!(principal.user_id, uid("alice"));
|
||||
assert!(principal.authorizes_session(&sid("sess-1")));
|
||||
assert!(!principal.authorizes_session(&sid("sess-other")));
|
||||
assert!(principal.has_scope("tool.invoke"));
|
||||
assert!(!principal.has_scope("admin"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn principal_builder_chains_in_order() {
|
||||
let principal = Principal::new(uid("alice"))
|
||||
.with_session(sid("sess-a"))
|
||||
.with_session(sid("sess-b"))
|
||||
.with_scope("tool.invoke")
|
||||
.with_scope("tool.search")
|
||||
.with_audience("dispatcher.example");
|
||||
assert_eq!(principal.session_ids, vec![sid("sess-a"), sid("sess-b")]);
|
||||
assert_eq!(principal.scopes, vec!["tool.invoke", "tool.search"]);
|
||||
assert_eq!(principal.audiences, vec!["dispatcher.example"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn principal_supports_multi_session_tokens() {
|
||||
let p = Principal::new(uid("alice"))
|
||||
.with_session(sid("sess-1"))
|
||||
.with_session(sid("sess-2"));
|
||||
assert!(p.authorizes_session(&sid("sess-1")));
|
||||
assert!(p.authorizes_session(&sid("sess-2")));
|
||||
assert!(!p.authorizes_session(&sid("sess-3")));
|
||||
assert_eq!(p.session_ids.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn principal_default_state_is_empty() {
|
||||
let p = Principal::new(uid("alice"));
|
||||
assert!(p.session_ids.is_empty());
|
||||
assert!(p.scopes.is_empty());
|
||||
assert!(p.audiences.is_empty());
|
||||
assert!(!p.has_scope("anything"));
|
||||
assert!(!p.authorizes_session(&sid("sess")));
|
||||
}
|
||||
|
|
@ -0,0 +1,203 @@
|
|||
//! `is_workspace_unavailable` recognizer coverage, pinned against the real
|
||||
//! wire decode path (`error_from_envelope` / `tool_error_from_wire`).
|
||||
|
||||
use serde_json::json;
|
||||
use xai_computer_hub_core::{error_from_envelope, is_workspace_unavailable, tool_error_from_wire};
|
||||
use xai_tool_protocol::{
|
||||
JsonRpcError, ToolErrorWire, WORKSPACE_UNAVAILABLE_SUBCODE, WorkspaceGonePhase,
|
||||
WorkspaceGoneReason, WorkspaceUnavailableDetails, workspace_unavailable_wire,
|
||||
};
|
||||
use xai_tool_runtime::{ToolError, ToolErrorKind};
|
||||
|
||||
const REASONS: [WorkspaceGoneReason; 5] = [
|
||||
WorkspaceGoneReason::IdleTimeout,
|
||||
WorkspaceGoneReason::Disconnect,
|
||||
WorkspaceGoneReason::Shutdown,
|
||||
WorkspaceGoneReason::NotBound,
|
||||
WorkspaceGoneReason::InstanceGone,
|
||||
];
|
||||
const PHASES: [WorkspaceGonePhase; 2] = [
|
||||
WorkspaceGonePhase::InFlightCancelled,
|
||||
WorkspaceGonePhase::RouteMissing,
|
||||
];
|
||||
|
||||
fn envelope_for(wire: &ToolErrorWire) -> JsonRpcError {
|
||||
JsonRpcError {
|
||||
// -32005 is the best-effort numeric companion (`tool_server_gone`);
|
||||
// recognition keys on `data.details.code`, not the numeric.
|
||||
code: -32005,
|
||||
message: "workspace server gone".to_owned(),
|
||||
data: Some(serde_json::to_value(wire).unwrap()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_trip_through_envelope_is_recognized_for_every_reason_and_phase() {
|
||||
for reason in REASONS {
|
||||
for phase in PHASES {
|
||||
let wire = workspace_unavailable_wire(reason, phase);
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
|
||||
assert!(
|
||||
is_workspace_unavailable(&err),
|
||||
"should recognize {reason:?}/{phase:?}",
|
||||
);
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
|
||||
// The full structured payload survives into `ToolError::details`,
|
||||
// so a caller can branch on code/reason/phase/retryable.
|
||||
let details: WorkspaceUnavailableDetails =
|
||||
serde_json::from_value(err.details.expect("details survive")).unwrap();
|
||||
assert_eq!(
|
||||
details,
|
||||
WorkspaceUnavailableDetails {
|
||||
code: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(),
|
||||
reason,
|
||||
phase,
|
||||
retryable: true,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_error_from_wire_directly_is_recognized() {
|
||||
let wire = workspace_unavailable_wire(
|
||||
WorkspaceGoneReason::Disconnect,
|
||||
WorkspaceGonePhase::RouteMissing,
|
||||
);
|
||||
let err = tool_error_from_wire(wire);
|
||||
assert!(is_workspace_unavailable(&err));
|
||||
let details = err.details.expect("details survive");
|
||||
assert_eq!(details["code"], json!(WORKSPACE_UNAVAILABLE_SUBCODE));
|
||||
assert_eq!(details["reason"], json!("disconnect"));
|
||||
assert_eq!(details["phase"], json!("route_missing"));
|
||||
assert_eq!(details["retryable"], json!(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wire_to_tool_error_to_wire_preserves_outer_subcode() {
|
||||
// Keying the identity on details.code lets From<ToolError> for ToolErrorWire
|
||||
// rebuild the outer subcode on re-serialization.
|
||||
let original = workspace_unavailable_wire(
|
||||
WorkspaceGoneReason::IdleTimeout,
|
||||
WorkspaceGonePhase::InFlightCancelled,
|
||||
);
|
||||
let tool_error = tool_error_from_wire(original);
|
||||
let back: ToolErrorWire = tool_error.into();
|
||||
let ToolErrorWire::Custom { subcode, .. } = back else {
|
||||
panic!("expected Custom variant");
|
||||
};
|
||||
assert_eq!(subcode, WORKSPACE_UNAVAILABLE_SUBCODE);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recognized_with_unknown_reason_and_phase() {
|
||||
// Recognition is decoupled from the typed reason/phase enums: a newer hub
|
||||
// emitting unknown values is still recognized (it keys only on `code`).
|
||||
let wire = ToolErrorWire::Custom {
|
||||
subcode: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(),
|
||||
message: "from a newer hub".to_owned(),
|
||||
details: Some(json!({
|
||||
"code": WORKSPACE_UNAVAILABLE_SUBCODE,
|
||||
"reason": "brand_new_reason",
|
||||
"phase": "brand_new_phase",
|
||||
"retryable": true,
|
||||
})),
|
||||
};
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
assert!(is_workspace_unavailable(&err));
|
||||
// End-to-end decode → typed-parse → `Unknown`, the path consumers read by.
|
||||
let details: WorkspaceUnavailableDetails =
|
||||
serde_json::from_value(err.details.expect("details survive")).unwrap();
|
||||
assert_eq!(details.reason, WorkspaceGoneReason::Unknown);
|
||||
assert_eq!(details.phase, WorkspaceGonePhase::Unknown);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoded_custom_with_none_details_is_recognized_via_canonical_code() {
|
||||
// Wire `details: None` decodes through `ToolError::custom`, which repopulates
|
||||
// `details = {"code": subcode}`, so it IS recognized — contrast the hand-built
|
||||
// no-details case in `custom_error_without_any_details_is_not_recognized`.
|
||||
let wire = ToolErrorWire::Custom {
|
||||
subcode: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(),
|
||||
message: "no structured details".to_owned(),
|
||||
details: None,
|
||||
};
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert!(is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoded_custom_without_code_key_is_not_recognized() {
|
||||
// The central correctness property: recognition keys on the surviving
|
||||
// `details.code`, NOT the outer `Custom.subcode`. Here the outer subcode
|
||||
// matches, but `with_details` overwrote the auto-populated `code`, so the
|
||||
// decoded error must NOT be recognized.
|
||||
let wire = ToolErrorWire::Custom {
|
||||
subcode: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(),
|
||||
message: "details lack code".to_owned(),
|
||||
details: Some(json!({ "reason": "disconnect" })),
|
||||
};
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_custom_code_is_not_recognized() {
|
||||
let wire = ToolErrorWire::Custom {
|
||||
subcode: "some_other_error".to_owned(),
|
||||
message: "nope".to_owned(),
|
||||
details: Some(json!({ "code": "some_other_error" })),
|
||||
};
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn numeric_only_tool_server_gone_without_data_is_not_recognized() {
|
||||
// Recognition is by the data payload, never the numeric code: a bare -32005
|
||||
// with no `data` decodes to a `jsonrpc_-32005` custom error, not recognized.
|
||||
let err = error_from_envelope(JsonRpcError {
|
||||
code: -32005,
|
||||
message: "tool server gone".to_owned(),
|
||||
data: None,
|
||||
});
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_error_without_any_details_is_not_recognized() {
|
||||
// Hand-built Custom with no `details` (no `code`) — unlike a wire `details:
|
||||
// None`, nothing repopulates `code` here, so it is not recognized.
|
||||
let err = ToolError::new(ToolErrorKind::Custom, "no details at all");
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_custom_error_with_matching_code_is_not_recognized() {
|
||||
// The kind guard matters: a non-Custom error carrying a matching
|
||||
// `details.code` must still be rejected.
|
||||
let err = ToolError::new(ToolErrorKind::NetworkError, "socket closed")
|
||||
.with_details(json!({ "code": WORKSPACE_UNAVAILABLE_SUBCODE }));
|
||||
assert_ne!(err.kind, ToolErrorKind::Custom);
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_custom_decoded_error_is_not_recognized() {
|
||||
let wire = ToolErrorWire::ToolNotFound {
|
||||
tool_id: xai_tool_protocol::ToolId::new("ns:tool").unwrap(),
|
||||
};
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
assert_ne!(err.kind, ToolErrorKind::Custom);
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
|
||||
assert!(!is_workspace_unavailable(&ToolError::network_error(
|
||||
"socket closed"
|
||||
)));
|
||||
}
|
||||
Loading…
Reference in a new issue