1use cpex_core::cmf::MessagePayload;
26use cpex_core::error::PluginViolation;
27use cpex_core::executor::ErasedResultFields;
28use cpex_core::hooks::payload::PluginPayload;
29use serde::{Deserialize, Serialize};
30use serde_json::Value;
31
32use crate::error::HostError;
33use crate::legacy::payloads::{
34 IdentityResolvePayload, PromptPostFetchPayload, PromptPreFetchPayload,
35 ResourcePostFetchPayload, ResourcePreFetchPayload, TokenDelegatePayload, ToolPostInvokePayload,
36 ToolPreInvokePayload,
37};
38
39pub const CMF_PREFIX: &str = "cmf.";
41
42#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct GenericPayload {
51 pub value: Value,
52}
53
54cpex_core::impl_plugin_payload!(GenericPayload);
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum PayloadKind {
59 CmfMessage,
61 ToolPreInvoke,
62 ToolPostInvoke,
63 PromptPreFetch,
64 PromptPostFetch,
65 ResourcePreFetch,
66 ResourcePostFetch,
67 IdentityResolve,
68 TokenDelegate,
69 Generic,
71}
72
73pub fn payload_kind_for_hook(hook_name: &str) -> PayloadKind {
78 use cpex_core::hooks::types::hook_names;
79
80 if hook_name.starts_with(CMF_PREFIX) {
81 return PayloadKind::CmfMessage;
82 }
83
84 match hook_name {
85 hook_names::TOOL_PRE_INVOKE => PayloadKind::ToolPreInvoke,
86 hook_names::TOOL_POST_INVOKE => PayloadKind::ToolPostInvoke,
87 hook_names::PROMPT_PRE_FETCH => PayloadKind::PromptPreFetch,
88 hook_names::PROMPT_POST_FETCH => PayloadKind::PromptPostFetch,
89 hook_names::RESOURCE_PRE_FETCH => PayloadKind::ResourcePreFetch,
90 hook_names::RESOURCE_POST_FETCH => PayloadKind::ResourcePostFetch,
91 hook_names::IDENTITY_RESOLVE => PayloadKind::IdentityResolve,
92 hook_names::TOKEN_DELEGATE => PayloadKind::TokenDelegate,
93 _ => PayloadKind::Generic,
94 }
95}
96
97pub fn is_credential_hook(hook_name: &str) -> bool {
105 matches!(
106 payload_kind_for_hook(hook_name),
107 PayloadKind::IdentityResolve | PayloadKind::TokenDelegate
108 )
109}
110
111pub fn serialize_payload(payload: &dyn PluginPayload) -> Result<Value, HostError> {
122 let any = payload.as_any();
123
124 macro_rules! try_downcast {
127 ($($ty:ty),+ $(,)?) => {
128 $(
129 if let Some(concrete) = any.downcast_ref::<$ty>() {
130 return serde_json::to_value(concrete).map_err(|e| HostError::Protocol {
131 message: format!(
132 "could not serialize a {} payload: {e}",
133 stringify!($ty)
134 ),
135 });
136 }
137 )+
138 };
139 }
140
141 try_downcast!(
142 MessagePayload,
143 ToolPreInvokePayload,
144 ToolPostInvokePayload,
145 PromptPreFetchPayload,
146 PromptPostFetchPayload,
147 ResourcePreFetchPayload,
148 ResourcePostFetchPayload,
149 IdentityResolvePayload,
150 TokenDelegatePayload,
151 );
152
153 if let Some(generic) = any.downcast_ref::<GenericPayload>() {
156 return Ok(generic.value.clone());
157 }
158
159 Err(HostError::Protocol {
160 message: "payload type is not one this host can forward to a Python worker \
161 (expected a CMF message payload, a legacy typed payload, or the generic payload)"
162 .into(),
163 })
164}
165
166pub fn deserialize_payload(
172 hook_name: &str,
173 value: Value,
174) -> Result<Box<dyn PluginPayload>, HostError> {
175 fn parse<T>(value: Value, hook_name: &str) -> Result<Box<dyn PluginPayload>, HostError>
176 where
177 T: PluginPayload + serde::de::DeserializeOwned,
178 {
179 serde_json::from_value::<T>(value)
180 .map(|p| Box::new(p) as Box<dyn PluginPayload>)
181 .map_err(|e| HostError::Protocol {
182 message: format!("worker returned a modified payload that does not match hook '{hook_name}': {e}"),
183 })
184 }
185
186 match payload_kind_for_hook(hook_name) {
187 PayloadKind::CmfMessage => parse::<MessagePayload>(value, hook_name),
188 PayloadKind::ToolPreInvoke => parse::<ToolPreInvokePayload>(value, hook_name),
189 PayloadKind::ToolPostInvoke => parse::<ToolPostInvokePayload>(value, hook_name),
190 PayloadKind::PromptPreFetch => parse::<PromptPreFetchPayload>(value, hook_name),
191 PayloadKind::PromptPostFetch => parse::<PromptPostFetchPayload>(value, hook_name),
192 PayloadKind::ResourcePreFetch => parse::<ResourcePreFetchPayload>(value, hook_name),
193 PayloadKind::ResourcePostFetch => parse::<ResourcePostFetchPayload>(value, hook_name),
194 PayloadKind::IdentityResolve => parse::<IdentityResolvePayload>(value, hook_name),
195 PayloadKind::TokenDelegate => parse::<TokenDelegatePayload>(value, hook_name),
196 PayloadKind::Generic => Ok(Box::new(GenericPayload { value })),
197 }
198}
199
200pub fn response_to_result(
213 hook_name: &str,
214 response: Value,
215 inbound: &cpex_core::extensions::Extensions,
216) -> Result<ErasedResultFields, HostError> {
217 let continue_processing = response
218 .get("continue_processing")
219 .and_then(Value::as_bool)
220 .unwrap_or(true);
221
222 let violation = match response.get("violation") {
223 Some(Value::Null) | None => None,
224 Some(raw) => Some(parse_violation(raw.clone())?),
225 };
226
227 let modified_payload = match response.get("modified_payload") {
228 Some(Value::Null) | None => None,
229 Some(raw) => Some(deserialize_payload(hook_name, raw.clone())?),
230 };
231
232 let modified_extensions = match response.get(crate::extensions::MODIFIED_EXTENSIONS_FIELD) {
233 Some(Value::Null) | None => None,
234 Some(raw) => crate::extensions::owned_from_returned_slot(raw, inbound)?,
235 };
236
237 Ok(ErasedResultFields {
238 continue_processing,
239 modified_payload,
240 modified_extensions,
241 violation,
242 })
243}
244
245fn parse_violation(raw: Value) -> Result<PluginViolation, HostError> {
253 let obj = raw.as_object().ok_or_else(|| HostError::Protocol {
254 message: "worker returned a violation that is not a JSON object".into(),
255 })?;
256
257 let code = obj.get("code").and_then(Value::as_str).unwrap_or_default();
258 let reason = obj
259 .get("reason")
260 .and_then(Value::as_str)
261 .unwrap_or_default();
262
263 let mut violation = PluginViolation::new(code, reason);
264 violation.description = obj
265 .get("description")
266 .and_then(Value::as_str)
267 .map(str::to_string);
268
269 if let Some(details) = obj.get("details").and_then(Value::as_object) {
270 violation.details = details.clone().into_iter().collect();
271 }
272
273 violation.proto_error_code = obj
276 .get("mcp_error_code")
277 .or_else(|| obj.get("proto_error_code"))
278 .and_then(Value::as_i64);
279
280 Ok(violation)
281}
282
283#[cfg(test)]
284mod tests {
285 use super::*;
286 use std::collections::HashMap;
287
288 fn no_inbound() -> cpex_core::extensions::Extensions {
296 cpex_core::extensions::Extensions::default()
297 }
298
299 #[test]
302 fn cmf_hooks_route_to_the_message_payload() {
303 assert_eq!(
307 payload_kind_for_hook("cmf.tool_pre_invoke"),
308 PayloadKind::CmfMessage
309 );
310 assert_eq!(
311 payload_kind_for_hook("cmf.llm_input"),
312 PayloadKind::CmfMessage
313 );
314 assert_eq!(
315 payload_kind_for_hook("cmf.some_future_hook"),
316 PayloadKind::CmfMessage,
317 "any cmf.* name is a message payload, even one this host predates"
318 );
319 }
320
321 #[test]
322 fn every_legacy_hook_resolves_to_its_own_typed_payload() {
323 let expected = [
324 ("tool_pre_invoke", PayloadKind::ToolPreInvoke),
325 ("tool_post_invoke", PayloadKind::ToolPostInvoke),
326 ("prompt_pre_fetch", PayloadKind::PromptPreFetch),
327 ("prompt_post_fetch", PayloadKind::PromptPostFetch),
328 ("resource_pre_fetch", PayloadKind::ResourcePreFetch),
329 ("resource_post_fetch", PayloadKind::ResourcePostFetch),
330 ("identity_resolve", PayloadKind::IdentityResolve),
331 ("token_delegate", PayloadKind::TokenDelegate),
332 ];
333
334 for (hook, kind) in expected {
335 assert_eq!(
336 payload_kind_for_hook(hook),
337 kind,
338 "hook '{hook}' resolved wrongly"
339 );
340 }
341
342 let kinds: std::collections::HashSet<_> =
345 expected.iter().map(|(_, k)| format!("{k:?}")).collect();
346 assert_eq!(
347 kinds.len(),
348 8,
349 "each legacy hook needs its own payload kind"
350 );
351 }
352
353 #[test]
354 fn an_unknown_hook_name_falls_back_to_the_generic_payload() {
355 assert_eq!(
356 payload_kind_for_hook("some_custom_hook"),
357 PayloadKind::Generic
358 );
359 assert_eq!(payload_kind_for_hook(""), PayloadKind::Generic);
360 }
361
362 #[test]
363 fn only_identity_and_delegation_are_credential_hooks() {
364 assert!(is_credential_hook("identity_resolve"));
365 assert!(is_credential_hook("token_delegate"));
366
367 for hook in [
368 "tool_pre_invoke",
369 "tool_post_invoke",
370 "prompt_pre_fetch",
371 "resource_pre_fetch",
372 "cmf.tool_pre_invoke",
373 "unknown",
374 ] {
375 assert!(
376 !is_credential_hook(hook),
377 "'{hook}' must not receive credentials"
378 );
379 }
380 }
381
382 #[test]
385 fn a_tool_pre_invoke_payload_serializes_to_the_worker_field_shape() {
386 let payload = ToolPreInvokePayload {
387 name: "search".into(),
388 args: Some(HashMap::from([("q".into(), serde_json::json!("rust"))])),
389 headers: Some(HashMap::from([("X-Tenant".into(), "acme".into())])),
390 };
391
392 let json = serialize_payload(&payload).expect("serializes");
393 assert_eq!(json["name"], "search");
394 assert_eq!(json["args"]["q"], "rust");
395 assert_eq!(json["headers"]["X-Tenant"], "acme");
396 }
397
398 #[test]
399 fn the_generic_payload_serializes_to_its_inner_value() {
400 let payload = GenericPayload {
402 value: serde_json::json!({"anything": [1, 2, 3]}),
403 };
404 let json = serialize_payload(&payload).expect("serializes");
405 assert_eq!(json["anything"][2], 3);
406 assert!(
407 json.get("value").is_none(),
408 "the wrapper must not appear on the wire"
409 );
410 }
411
412 #[test]
413 fn a_cmf_message_payload_serializes_as_itself() {
414 use cpex_core::cmf::{Message, Role};
415
416 let payload = MessagePayload {
417 message: Message::text(Role::User, "what is the weather?"),
418 };
419 let json = serialize_payload(&payload).expect("serializes");
420 assert_eq!(json["message"]["role"], "user");
421 }
422
423 #[test]
424 fn an_unforwardable_payload_type_errors_rather_than_sending_null() {
425 #[derive(Debug, Clone)]
426 struct ForeignPayload;
427 cpex_core::impl_plugin_payload!(ForeignPayload);
428
429 let err =
430 serialize_payload(&ForeignPayload).expect_err("an unknown type cannot be forwarded");
431 assert!(matches!(err, HostError::Protocol { .. }), "got {err:?}");
432 }
433
434 #[test]
437 fn a_modified_payload_deserializes_into_the_hooks_typed_payload() {
438 let boxed =
439 deserialize_payload("tool_pre_invoke", serde_json::json!({ "name": "redacted" }))
440 .expect("parses");
441 let typed = boxed
442 .as_any()
443 .downcast_ref::<ToolPreInvokePayload>()
444 .expect("must come back as the tool payload, or the executor's downcast fails");
445 assert_eq!(typed.name, "redacted");
446 }
447
448 #[test]
449 fn a_cmf_modified_payload_deserializes_as_a_message_payload() {
450 use cpex_core::cmf::{Message, Role};
451
452 let original = MessagePayload {
455 message: Message::text(Role::Assistant, "redacted"),
456 };
457 let json = serde_json::to_value(&original).unwrap();
458
459 let boxed = deserialize_payload("cmf.tool_pre_invoke", json).expect("parses");
460 let typed = boxed
461 .as_any()
462 .downcast_ref::<MessagePayload>()
463 .expect("a cmf hook's modified payload must come back as a MessagePayload");
464 assert_eq!(typed.message.role, Role::Assistant);
465 }
466
467 #[test]
468 fn an_unknown_hooks_modified_payload_becomes_a_generic_payload() {
469 let boxed =
470 deserialize_payload("mystery_hook", serde_json::json!({"x": 1})).expect("parses");
471 let generic = boxed
472 .as_any()
473 .downcast_ref::<GenericPayload>()
474 .expect("generic");
475 assert_eq!(generic.value["x"], 1);
476 }
477
478 #[test]
479 fn a_mismatched_modified_payload_errors_and_names_the_hook() {
480 let err = deserialize_payload("tool_pre_invoke", serde_json::json!({ "unexpected": true }))
483 .expect_err("a payload missing required fields must be rejected");
484 assert!(err.to_string().contains("tool_pre_invoke"), "{err}");
485 }
486
487 #[test]
490 fn an_allow_response_becomes_an_allow_result() {
491 let fields = response_to_result(
492 "tool_pre_invoke",
493 serde_json::json!({ "continue_processing": true }),
494 &no_inbound(),
495 )
496 .expect("converts");
497
498 assert!(fields.continue_processing);
499 assert!(fields.violation.is_none());
500 assert!(fields.modified_payload.is_none());
501 assert!(fields.modified_extensions.is_none());
502 }
503
504 #[test]
505 fn a_response_without_continue_processing_defaults_to_allow() {
506 let fields = response_to_result("tool_pre_invoke", serde_json::json!({}), &no_inbound())
509 .expect("converts");
510 assert!(fields.continue_processing);
511 }
512
513 #[test]
514 fn a_deny_response_carries_its_violation() {
515 let fields = response_to_result(
516 "tool_pre_invoke",
517 serde_json::json!({
518 "continue_processing": false,
519 "violation": {
520 "code": "PII_DETECTED",
521 "reason": "email address in args",
522 "description": "matched the email pattern",
523 "details": {"field": "q"},
524 "mcp_error_code": -32603
525 }
526 }),
527 &no_inbound(),
528 )
529 .expect("converts");
530
531 assert!(!fields.continue_processing);
532 let violation = fields.violation.expect("a deny must carry its violation");
533 assert_eq!(violation.code, "PII_DETECTED");
534 assert_eq!(violation.reason, "email address in args");
535 assert_eq!(
536 violation.description.as_deref(),
537 Some("matched the email pattern")
538 );
539 assert_eq!(violation.details["field"], "q");
540 assert_eq!(
541 violation.proto_error_code,
542 Some(-32603),
543 "the Python mcp_error_code maps onto the Rust proto_error_code"
544 );
545 }
546
547 #[test]
548 fn a_violation_with_only_a_reason_still_denies() {
549 let fields = response_to_result(
553 "tool_pre_invoke",
554 serde_json::json!({ "continue_processing": false, "violation": { "reason": "nope" } }),
555 &no_inbound(),
556 )
557 .expect("a sparse violation is still a violation");
558
559 assert!(!fields.continue_processing);
560 let violation = fields.violation.unwrap();
561 assert_eq!(violation.reason, "nope");
562 assert_eq!(violation.code, "");
563 }
564
565 #[test]
566 fn a_modified_payload_survives_the_round_trip() {
567 let fields = response_to_result(
568 "tool_pre_invoke",
569 serde_json::json!({
570 "continue_processing": true,
571 "modified_payload": { "name": "search", "args": {"q": "[REDACTED]"} }
572 }),
573 &no_inbound(),
574 )
575 .expect("converts");
576
577 let payload = fields
578 .modified_payload
579 .expect("the modification must survive");
580 let typed = payload
581 .as_any()
582 .downcast_ref::<ToolPreInvokePayload>()
583 .unwrap();
584 assert_eq!(typed.args.as_ref().unwrap()["q"], "[REDACTED]");
585 }
586
587 #[test]
588 fn modified_extensions_survive_the_round_trip() {
589 let fields = response_to_result(
590 "tool_pre_invoke",
591 serde_json::json!({
592 "continue_processing": true,
593 "modified_extensions": { "custom": { "seen_by": "py-plugin" } }
594 }),
595 &no_inbound(),
596 )
597 .expect("converts");
598
599 let extensions = fields.modified_extensions.expect("extensions must survive");
600 assert_eq!(
601 extensions.custom.expect("custom slot")["seen_by"],
602 "py-plugin"
603 );
604 }
605
606 #[test]
607 fn a_gated_slot_write_without_the_capability_is_dropped() {
608 for slot in ["http", "security", "delegation"] {
619 let fields = response_to_result(
620 "tool_pre_invoke",
621 serde_json::json!({
622 "continue_processing": true,
623 "modified_extensions": { slot: {"labels": ["PII"]} }
624 }),
625 &no_inbound(),
626 )
627 .expect("an unauthorized gated write is dropped, not a protocol error");
628
629 assert!(
630 fields.modified_extensions.is_none(),
631 "the '{slot}' write had no token behind it, so nothing should merge"
632 );
633 }
634 }
635
636 #[test]
637 fn an_extensions_object_with_only_gated_nulls_is_no_modification() {
638 let fields = response_to_result(
642 "tool_pre_invoke",
643 serde_json::json!({
644 "continue_processing": true,
645 "modified_extensions": { "http": null, "security": null, "custom": null }
646 }),
647 &no_inbound(),
648 )
649 .expect("all-null slots are not a write");
650 assert!(fields.modified_extensions.is_none());
651 }
652
653 #[test]
654 fn explicit_nulls_are_treated_as_absent() {
655 let fields = response_to_result(
659 "tool_pre_invoke",
660 serde_json::json!({
661 "continue_processing": true,
662 "modified_payload": null,
663 "modified_extensions": null,
664 "violation": null
665 }),
666 &no_inbound(),
667 )
668 .expect("nulls are absent, not malformed");
669
670 assert!(fields.modified_payload.is_none());
671 assert!(fields.modified_extensions.is_none());
672 assert!(fields.violation.is_none());
673 }
674
675 #[test]
676 fn a_malformed_modified_payload_surfaces_as_a_protocol_error() {
677 let Err(err) = response_to_result(
678 "tool_pre_invoke",
679 serde_json::json!({ "continue_processing": true, "modified_payload": "not an object" }),
680 &no_inbound(),
681 ) else {
682 panic!("a payload that cannot be rebuilt must not be silently dropped");
683 };
684 assert!(matches!(err, HostError::Protocol { .. }), "got {err:?}");
685 }
686}