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

1use axum::body::Body;
2use axum::extract::Request;
3use serde_json::Value;
4use systemprompt_models::execution::{ContextExtractionError, ContextIdSource};
5
6#[derive(Debug, Clone, Copy)]
7pub struct PayloadSource;
8
9impl PayloadSource {
10    pub fn extract_context_source(
11        body_bytes: &[u8],
12    ) -> Result<ContextIdSource, ContextExtractionError> {
13        let payload: Value = serde_json::from_slice(body_bytes).map_err(|e| {
14            ContextExtractionError::InvalidHeaderValue {
15                header: "payload".to_string(),
16                reason: format!("Invalid JSON: {e}"),
17            }
18        })?;
19
20        let method = payload.get("method").and_then(|m| m.as_str()).unwrap_or("");
21
22        if method.starts_with("tasks/") {
23            let task_id = payload
24                .get("params")
25                .and_then(|p| p.get("id"))
26                .and_then(|id| id.as_str())
27                .map(ToString::to_string)
28                .ok_or_else(|| ContextExtractionError::InvalidHeaderValue {
29                    header: "params.id".to_string(),
30                    reason: "Task ID required for task methods".to_string(),
31                })?;
32
33            return Ok(ContextIdSource::FromTask { task_id });
34        }
35
36        payload
37            .get("params")
38            .and_then(|p| p.get("message"))
39            .and_then(|m| m.get("contextId"))
40            .and_then(|c| c.as_str())
41            .map(|s| ContextIdSource::Direct(s.to_string()))
42            .ok_or(ContextExtractionError::MissingContextId)
43    }
44
45    pub async fn read_and_reconstruct(
46        request: Request<Body>,
47    ) -> Result<(Vec<u8>, Request<Body>), ContextExtractionError> {
48        let (parts, body) = request.into_parts();
49
50        let body_bytes = axum::body::to_bytes(body, usize::MAX)
51            .await
52            .map_err(|e| ContextExtractionError::InvalidHeaderValue {
53                header: "body".to_string(),
54                reason: format!("Failed to read body: {e}"),
55            })?
56            .to_vec();
57
58        let new_body = Body::from(body_bytes.clone());
59        let new_request = Request::from_parts(parts, new_body);
60
61        Ok((body_bytes, new_request))
62    }
63}
systemprompt_api/services/middleware/context/sources/payload.rs

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