systemprompt_api/services/middleware/context/sources/
payload.rs

1//! Request-context extraction from A2A JSON-RPC bodies.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use axum::body::Body;
7use axum::extract::Request;
8use serde_json::Value;
9use systemprompt_models::execution::{ContextExtractionError, ContextIdSource};
10
11#[derive(Debug, Clone, Copy)]
12pub struct PayloadSource;
13
14impl PayloadSource {
15    pub fn extract_context_source(
16        body_bytes: &[u8],
17    ) -> Result<ContextIdSource, ContextExtractionError> {
18        // JSON: A2A JSON-RPC envelope is an external protocol boundary; the
19        // method name drives which typed field is read, so the shape is dynamic.
20        let payload: Value = serde_json::from_slice(body_bytes).map_err(|e| {
21            ContextExtractionError::InvalidHeaderValue {
22                header: "payload".to_owned(),
23                reason: format!("Invalid JSON: {e}"),
24            }
25        })?;
26
27        let method = payload.get("method").and_then(|m| m.as_str()).unwrap_or("");
28
29        if method.starts_with("tasks/") {
30            let task_id = payload
31                .get("params")
32                .and_then(|p| p.get("id"))
33                .and_then(|id| id.as_str())
34                .map(str::to_owned)
35                .ok_or_else(|| ContextExtractionError::InvalidHeaderValue {
36                    header: "params.id".to_owned(),
37                    reason: "Task ID required for task methods".to_owned(),
38                })?;
39
40            return Ok(ContextIdSource::FromTask {
41                task_id: systemprompt_identifiers::TaskId::new(task_id),
42            });
43        }
44
45        payload
46            .get("params")
47            .and_then(|p| p.get("message"))
48            .and_then(|m| m.get("contextId"))
49            .and_then(|c| c.as_str())
50            .map(|s| ContextIdSource::Direct(s.to_owned()))
51            .ok_or(ContextExtractionError::MissingContextId)
52    }
53
54    pub async fn read_and_reconstruct(
55        request: Request<Body>,
56    ) -> Result<(Vec<u8>, Request<Body>), ContextExtractionError> {
57        let (parts, body) = request.into_parts();
58
59        let body_bytes = axum::body::to_bytes(body, usize::MAX)
60            .await
61            .map_err(|e| ContextExtractionError::InvalidHeaderValue {
62                header: "body".to_owned(),
63                reason: format!("Failed to read body: {e}"),
64            })?
65            .to_vec();
66
67        let new_body = Body::from(body_bytes.clone());
68        let new_request = Request::from_parts(parts, new_body);
69
70        Ok((body_bytes, new_request))
71    }
72}
systemprompt_api/services/middleware/context/sources/payload.rs

systemprompt_api/services/middleware/context/sources/
payload.rs
