1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
//! Duplex approval / denial round-trip tests for the client host surface
//! (Phase 106, HOST-04).
//!
//! Proves the two host-side denial paths keep the connection alive:
//! * a **preflight-denied** sampling request returns a sanitized `-32603`
//! policy denial to the server, the LLM handler is NEVER invoked, and a
//! subsequent client request still completes; and
//! * a **known-but-unhandled** host method (`elicitation/create` with no
//! registered handler) returns `-32601`, and a subsequent client request
//! still completes.
//!
//! Both drive the *server* side by hand over the duplex transport (see the
//! sibling `client_host_roundtrip.rs` for why a raw pump is used instead of
//! `Server::run` + `PeerHandle`): the fake server answers `initialize`, waits
//! for the client's in-flight `tools/list`, injects a raw inbound request,
//! captures the client's `Response`, answers the first call, then answers a
//! SECOND `tools/list` to prove the connection survived the error.
#![cfg(not(target_arch = "wasm32"))]
#[path = "common/duplex.rs"]
mod duplex;
#[path = "common/host_pump.rs"]
mod host_pump;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use async_trait::async_trait;
use serde_json::json;
use pmcp::client::host::{ApprovalDecision, HostSamplingHandler};
use pmcp::types::jsonrpc::{JSONRPCResponse, ResponsePayload};
use pmcp::types::protocol::{ClientRequest, Request, ServerRequest};
use pmcp::types::sampling::{CreateMessageParams, CreateMessageResult};
use pmcp::types::{ClientCapabilities, Content, RequestId};
use pmcp::ClientBuilder;
// ---------------------------------------------------------------------------
// Mock host sampling handler that records whether it was invoked.
// ---------------------------------------------------------------------------
struct TrackingSampling {
invoked: Arc<AtomicBool>,
}
#[async_trait]
impl HostSamplingHandler for TrackingSampling {
async fn handle_create_message(
&self,
_params: CreateMessageParams,
) -> pmcp::Result<CreateMessageResult> {
self.invoked.store(true, Ordering::SeqCst);
Ok(CreateMessageResult::new(
Content::text("host completion"),
"host-model",
))
}
}
// ---------------------------------------------------------------------------
// Raw duplex pump driver (survival variant).
// ---------------------------------------------------------------------------
/// Drive the shared [`host_pump::inject_and_capture`] prefix (initialize, inject
/// `inbound`, capture the client's `Response`, answer the first call), then add
/// the connection-survival probe: answer a SECOND `tools/list` to prove the
/// connection survived the inbound error. Returns the captured inbound
/// `Response`.
async fn pump_inbound_then_survive(
mut server_t: duplex::DuplexTransport,
inbound_id: RequestId,
inbound: Request,
) -> JSONRPCResponse {
let inbound_response = host_pump::inject_and_capture(&mut server_t, inbound_id, inbound).await;
// Connection-survival probe: answer a SECOND `tools/list`.
let call_id_2 = host_pump::recv_request_id(&mut server_t).await;
host_pump::send_success(&mut server_t, call_id_2, json!({ "tools": [] })).await;
inbound_response
}
fn assert_error_code(response: &JSONRPCResponse, expected: i32) -> String {
match &response.payload {
ResponsePayload::Error(e) => {
assert_eq!(e.code, expected, "unexpected JSON-RPC error code");
e.message.clone()
},
ResponsePayload::Result(r) => panic!("expected error {expected}, got result: {r:?}"),
}
}
// ---------------------------------------------------------------------------
// HOST-04: preflight denial -> sanitized -32603 -> connection survives.
// ---------------------------------------------------------------------------
#[tokio::test]
async fn sampling_preflight_deny_survives_connection() {
let invoked = Arc::new(AtomicBool::new(false));
let (client_t, server_t) = duplex::DuplexTransport::pair();
let inbound_id = RequestId::from("sample-deny-1".to_string());
// Inbound sampling parses as the CLIENT-alias variant (parse ambiguity).
let inbound = Request::Client(Box::new(ClientRequest::CreateMessage(Box::new(
CreateMessageParams::new(Vec::new()),
))));
let server_task = tokio::spawn(pump_inbound_then_survive(
server_t,
inbound_id.clone(),
inbound,
));
let mut client = ClientBuilder::new(client_t)
.on_sampling(TrackingSampling {
invoked: invoked.clone(),
})
.on_sampling_approval(|_params| async {
ApprovalDecision::Deny("local-policy-secret".to_string())
})
.build();
client
.initialize(ClientCapabilities::default())
.await
.expect("initialize");
// First call: the server injects a sampling request mid-flight; the client
// must deny it (preflight) and still complete this call.
let _ = client.list_tools(None).await.expect("first list_tools");
// Second call: proves the connection survived the denial.
let _ = client.list_tools(None).await.expect("second list_tools");
let inbound_response = server_task.await.expect("server task");
assert_eq!(inbound_response.id, inbound_id);
// The denial is a sanitized -32603 whose message is the generic policy
// string — the raw deny reason must never cross the wire.
let message = assert_error_code(&inbound_response, -32603);
assert_eq!(message, "request denied by host policy");
assert!(
!message.contains("local-policy-secret"),
"deny reason must be logged locally, not forwarded: {message}"
);
// The LLM handler was NEVER invoked — the real denial-of-wallet fix.
assert!(
!invoked.load(Ordering::SeqCst),
"preflight Deny must prevent the LLM call entirely"
);
}
// ---------------------------------------------------------------------------
// Orchestrator addendum: a KNOWN host method with no registered handler
// returns -32601, and a subsequent request still completes.
// ---------------------------------------------------------------------------
#[tokio::test]
async fn unhandled_known_method_returns_method_not_found_and_survives() {
let (client_t, server_t) = duplex::DuplexTransport::pair();
let inbound_id = RequestId::from("elicit-unhandled-1".to_string());
// `elicitation/create` is a KNOWN host method, but no handler is registered.
let inbound = Request::Server(Box::new(ServerRequest::ElicitationCreate(Box::new(
pmcp::types::elicitation::ElicitRequestParams::Form {
message: "confirm?".to_string(),
requested_schema: json!({"type": "object"}),
},
))));
let server_task = tokio::spawn(pump_inbound_then_survive(
server_t,
inbound_id.clone(),
inbound,
));
// No elicitation handler registered.
let mut client = ClientBuilder::new(client_t).build();
client
.initialize(ClientCapabilities::default())
.await
.expect("initialize");
let _ = client.list_tools(None).await.expect("first list_tools");
let _ = client.list_tools(None).await.expect("second list_tools");
let inbound_response = server_task.await.expect("server task");
assert_eq!(inbound_response.id, inbound_id);
// Known-but-unhandled => -32601 (method not found), connection preserved.
let _ = assert_error_code(&inbound_response, -32601);
}