1use serde_json::Value;
11
12use super::FindingSink;
13use super::support::{has_rfc3986_scheme, is_base64, server_capability};
14use crate::context::TraceContext;
15
16pub(super) fn capability_declared(context: &TraceContext<'_>, sink: &mut FindingSink) {
20 if server_capability(context, &["prompts"]) != Some(false) {
21 return;
22 }
23 for exchange in context.exchanges() {
24 if matches!(exchange.method, "prompts/list" | "prompts/get") && exchange.result.is_some() {
25 sink.push(
26 Some(exchange.response.seq),
27 format!(
28 "server answered {:?} without declaring the prompts capability",
29 exchange.method
30 ),
31 );
32 }
33 }
34}
35
36fn prompt_content_items<'a>(context: &TraceContext<'a>) -> Vec<(u64, &'a Value)> {
38 let mut items = Vec::new();
39 for exchange in context.exchanges_for("prompts/get") {
40 let messages = exchange
41 .result
42 .and_then(|result| result.get("messages"))
43 .and_then(Value::as_array);
44 for message in messages.into_iter().flatten() {
45 if let Some(content) = message.get("content") {
46 items.push((exchange.response.seq, content));
47 }
48 }
49 }
50 items
51}
52
53fn declared_required_arguments<'a>(
58 context: &TraceContext<'a>,
59) -> std::collections::BTreeMap<&'a str, Vec<&'a str>> {
60 let mut declared = std::collections::BTreeMap::new();
61 for exchange in context.exchanges_for("prompts/list") {
62 let prompts = exchange
63 .result
64 .and_then(|result| result.get("prompts"))
65 .and_then(Value::as_array);
66 for prompt in prompts.into_iter().flatten() {
67 let Some(name) = prompt.get("name").and_then(Value::as_str) else {
68 continue;
69 };
70 let required: Vec<&str> = prompt
71 .get("arguments")
72 .and_then(Value::as_array)
73 .into_iter()
74 .flatten()
75 .filter(|argument| argument.get("required").and_then(Value::as_bool) == Some(true))
76 .filter_map(|argument| argument.get("name").and_then(Value::as_str))
77 .collect();
78 if !required.is_empty() {
79 declared.insert(name, required);
81 }
82 }
83 }
84 declared
85}
86
87pub(super) fn arguments_validated(context: &TraceContext<'_>, sink: &mut FindingSink) {
106 let declared = declared_required_arguments(context);
107 if declared.is_empty() {
108 return;
109 }
110 for exchange in context.exchanges_for("prompts/get") {
111 if exchange.result.is_none() {
112 continue;
113 }
114 let Some(params) = exchange.params else {
115 continue;
116 };
117 let Some(name) = params.get("name").and_then(Value::as_str) else {
118 continue;
119 };
120 let Some(required) = declared.get(name) else {
121 continue;
122 };
123 let supplied = params.get("arguments");
124 let missing: Vec<&str> = required
125 .iter()
126 .filter(|argument| {
127 supplied
128 .and_then(|arguments| arguments.get(*argument))
129 .is_none()
130 })
131 .copied()
132 .collect();
133 if missing.is_empty() {
134 continue;
135 }
136 sink.push(
137 Some(exchange.response.seq),
138 format!(
139 "server returned a result for prompts/get {name:?} although the request \
140 omitted the required argument(s) {missing:?} it declared in prompts/list"
141 ),
142 );
143 }
144}
145
146pub(super) fn image_content_encoding(context: &TraceContext<'_>, sink: &mut FindingSink) {
148 binary_content_encoding(context, sink, "image");
149}
150
151pub(super) fn audio_content_encoding(context: &TraceContext<'_>, sink: &mut FindingSink) {
153 binary_content_encoding(context, sink, "audio");
154}
155
156fn binary_content_encoding(context: &TraceContext<'_>, sink: &mut FindingSink, kind: &str) {
157 for (seq, content) in prompt_content_items(context) {
158 if content.get("type").and_then(Value::as_str) != Some(kind) {
159 continue;
160 }
161 let data_valid = content
162 .get("data")
163 .and_then(Value::as_str)
164 .is_some_and(is_base64);
165 if !data_valid {
166 sink.push(
167 Some(seq),
168 format!("{kind} content data is not valid base64"),
169 );
170 }
171 let mime_present = content
172 .get("mimeType")
173 .and_then(Value::as_str)
174 .is_some_and(|mime| {
175 mime.split_once('/')
176 .is_some_and(|(t, s)| !t.is_empty() && !s.is_empty())
177 });
178 if !mime_present {
179 sink.push(
180 Some(seq),
181 format!("{kind} content lacks a valid mimeType (expected type/subtype)"),
182 );
183 }
184 }
185}
186
187pub(super) fn embedded_resource_shape(context: &TraceContext<'_>, sink: &mut FindingSink) {
190 for (seq, content) in prompt_content_items(context) {
191 if content.get("type").and_then(Value::as_str) != Some("resource") {
192 continue;
193 }
194 let Some(resource) = content.get("resource") else {
195 sink.push(
196 Some(seq),
197 "embedded resource content lacks the resource member".to_owned(),
198 );
199 continue;
200 };
201 let uri_ok = resource
202 .get("uri")
203 .and_then(Value::as_str)
204 .is_some_and(has_rfc3986_scheme);
205 if !uri_ok {
206 sink.push(
207 Some(seq),
208 "embedded resource lacks a valid resource URI".to_owned(),
209 );
210 }
211 if resource.get("mimeType").and_then(Value::as_str).is_none() {
212 sink.push(Some(seq), "embedded resource lacks a mimeType".to_owned());
213 }
214 let text = resource.get("text").and_then(Value::as_str);
215 let blob = resource.get("blob").and_then(Value::as_str);
216 match (text, blob) {
217 (Some(_), None) => {}
218 (None, Some(blob)) if is_base64(blob) => {}
219 (None, Some(_)) => sink.push(
220 Some(seq),
221 "embedded resource blob is not valid base64".to_owned(),
222 ),
223 (Some(_), Some(_)) => sink.push(
224 Some(seq),
225 "embedded resource carries both text and blob; expected exactly one".to_owned(),
226 ),
227 (None, None) => sink.push(
228 Some(seq),
229 "embedded resource carries neither text nor blob data".to_owned(),
230 ),
231 }
232 }
233}
234
235#[cfg(test)]
236#[allow(clippy::unwrap_used)]
237mod tests {
238 use crate::checks;
239 use crate::context::TraceContext;
240 use crate::reader::{Limits, parse_trace};
241
242 fn findings_for(check: &str, trace: &str) -> Vec<String> {
243 let events = parse_trace(trace, &Limits::default()).unwrap();
244 let context = TraceContext::new(&events);
245 checks::find(check)
246 .unwrap()
247 .run(&context)
248 .into_iter()
249 .map(|finding| finding.detail)
250 .collect()
251 }
252
253 const HANDSHAKE: &str = r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"t","version":"0"}}}}
254{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2025-11-25","capabilities":{"prompts":{}},"serverInfo":{"name":"s","version":"0"}}}}
255{"seq":2,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","method":"notifications/initialized"}}"#;
256
257 fn get_prompt_with_content(content: &str) -> String {
258 let request = r#"{"seq":3,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":2,"method":"prompts/get","params":{"name":"p"}}}"#;
259 let result = format!(
260 r#"{{"seq":4,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{{"jsonrpc":"2.0","id":2,"result":{{"messages":[{{"role":"user","content":{content}}}]}}}}}}"#
261 );
262 format!("{HANDSHAKE}\n{request}\n{result}")
263 }
264
265 fn listed_then_called(arguments: &str, response: &str) -> String {
268 let list = r#"{"seq":3,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":2,"method":"prompts/list"}}
269{"seq":4,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":2,"result":{"prompts":[{"name":"review","arguments":[{"name":"diff","required":true},{"name":"tone","required":false}]}]}}}"#;
270 let get = format!(
271 r#"{{"seq":5,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{{"jsonrpc":"2.0","id":3,"method":"prompts/get","params":{{"name":"review","arguments":{arguments}}}}}}}"#
272 );
273 format!("{HANDSHAKE}\n{list}\n{get}\n{response}")
274 }
275
276 const RESULT: &str = r#"{"seq":6,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":3,"result":{"messages":[]}}}"#;
277 const ERROR: &str = r#"{"seq":6,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":3,"error":{"code":-32602,"message":"missing required argument: diff"}}}"#;
278
279 #[test]
280 fn serving_a_result_despite_a_missing_required_argument_is_a_finding() {
281 let trace = listed_then_called(r#"{"tone":"terse"}"#, RESULT);
282 let findings = findings_for("prompts.arguments-validated", &trace);
283 assert_eq!(findings.len(), 1, "{findings:?}");
284 assert!(findings[0].contains("\"diff\""), "{findings:?}");
285 assert!(findings[0].contains("review"), "{findings:?}");
286 }
287
288 #[test]
289 fn supplying_the_required_argument_is_clean() {
290 let trace = listed_then_called(r#"{"diff":"--- a\n+++ b"}"#, RESULT);
291 assert!(findings_for("prompts.arguments-validated", &trace).is_empty());
292 }
293
294 #[test]
295 fn an_omitted_optional_argument_is_not_a_finding() {
296 let trace = listed_then_called(r#"{"diff":"d"}"#, RESULT);
298 assert!(findings_for("prompts.arguments-validated", &trace).is_empty());
299 }
300
301 #[test]
302 fn rejecting_the_call_satisfies_the_clause() {
303 let trace = listed_then_called(r#"{"tone":"terse"}"#, ERROR);
306 assert!(findings_for("prompts.arguments-validated", &trace).is_empty());
307 }
308
309 #[test]
310 fn a_prompt_never_listed_is_not_judged() {
311 let get = r#"{"seq":3,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":2,"method":"prompts/get","params":{"name":"unlisted"}}}"#;
314 let result = r#"{"seq":4,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":2,"result":{"messages":[]}}}"#;
315 let trace = format!("{HANDSHAKE}\n{get}\n{result}");
316 assert!(findings_for("prompts.arguments-validated", &trace).is_empty());
317 }
318
319 #[test]
320 fn a_missing_arguments_object_entirely_is_still_a_finding() {
321 let get = r#"{"seq":5,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":3,"method":"prompts/get","params":{"name":"review"}}}"#;
324 let list = r#"{"seq":3,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":2,"method":"prompts/list"}}
325{"seq":4,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":2,"result":{"prompts":[{"name":"review","arguments":[{"name":"diff","required":true}]}]}}}"#;
326 let trace = format!("{HANDSHAKE}\n{list}\n{get}\n{RESULT}");
327 assert_eq!(findings_for("prompts.arguments-validated", &trace).len(), 1);
328 }
329
330 #[test]
331 fn image_and_audio_checks_are_type_scoped() {
332 let trace = get_prompt_with_content(
334 r#"{"type":"audio","data":"not base64!","mimeType":"audio/wav"}"#,
335 );
336 assert!(findings_for("prompts.image-content-encoding", &trace).is_empty());
337 let findings = findings_for("prompts.audio-content-encoding", &trace);
338 assert_eq!(findings.len(), 1, "{findings:?}");
339 assert!(findings[0].contains("audio content data"), "{findings:?}");
340 }
341
342 #[test]
343 fn mime_type_must_be_type_slash_subtype() {
344 for (mime, expect_finding) in [
345 (r#""image/png""#, false),
346 (r#""image/""#, true),
347 (r#""png""#, true),
348 ("42", true),
349 ] {
350 let trace = get_prompt_with_content(&format!(
351 r#"{{"type":"image","data":"QUJDRA==","mimeType":{mime}}}"#
352 ));
353 let findings = findings_for("prompts.image-content-encoding", &trace);
354 assert_eq!(!findings.is_empty(), expect_finding, "{mime}: {findings:?}");
355 }
356 }
357
358 #[test]
359 fn embedded_resource_shape_flags_each_defect_once() {
360 let trace = get_prompt_with_content(
361 r#"{"type":"resource","resource":{"uri":"no scheme","text":"x","blob":"QUJDRA=="}}"#,
362 );
363 let findings = findings_for("prompts.embedded-resource-shape", &trace);
364 assert_eq!(findings.len(), 3, "{findings:?}");
366 }
367
368 #[test]
369 fn well_formed_embedded_resource_passes() {
370 let trace = get_prompt_with_content(
371 r#"{"type":"resource","resource":{"uri":"file:///a.txt","mimeType":"text/plain","text":"hello"}}"#,
372 );
373 assert!(findings_for("prompts.embedded-resource-shape", &trace).is_empty());
374 }
375
376 #[test]
377 fn blob_only_embedded_resources_hinge_on_base64_validity() {
378 let valid = get_prompt_with_content(
380 r#"{"type":"resource","resource":{"uri":"file:///a.png","mimeType":"image/png","blob":"QUJDRA=="}}"#,
381 );
382 assert!(findings_for("prompts.embedded-resource-shape", &valid).is_empty());
383
384 let invalid = get_prompt_with_content(
386 r#"{"type":"resource","resource":{"uri":"file:///a.png","mimeType":"image/png","blob":"not base64!"}}"#,
387 );
388 let findings = findings_for("prompts.embedded-resource-shape", &invalid);
389 assert_eq!(findings.len(), 1, "{findings:?}");
390 assert!(findings[0].contains("not valid base64"), "{findings:?}");
391 }
392}