Skip to main content

stasis/application/runtime/
runtime_diagnostics_helpers.rs

1use serde_json::{Value as JsonValue, json};
2
3use crate::ports::outbound::memory::memory_models::{
4    MemoryNode, MemoryRecallResponse, MemoryStoreResponse,
5};
6
7pub fn memory_nodes_json(nodes: &[MemoryNode]) -> JsonValue {
8    JsonValue::Array(
9        nodes
10            .iter()
11            .map(|node| {
12                json!({
13                    "sync_key": node.sync_key,
14                    "session_id": node.session_id,
15                    "tier": node.tier,
16                    "raw": node.raw,
17                    "context_summary": node.context_summary,
18                    "psi": node.psi,
19                    "rho": node.rho,
20                    "kappa": node.kappa,
21                })
22            })
23            .collect(),
24    )
25}
26
27pub struct RuntimeMemoryDiagnosticsBundle {
28    pub retrieved_count: usize,
29    pub retrieval_path: Option<String>,
30    pub fallback_triggered: bool,
31    pub fallback_reason: Option<String>,
32    pub store_valid: bool,
33    pub store_node_id: Option<String>,
34    pub memory_recall: JsonValue,
35    pub memory_store: JsonValue,
36    pub identity_context: JsonValue,
37}
38
39pub struct RuntimeMemoryRecallDiagnosticsInput {
40    pub attempted: bool,
41    pub response: Option<MemoryRecallResponse>,
42    pub query_id: Option<String>,
43    pub query_fingerprint: Option<String>,
44    pub error: Option<String>,
45}
46
47pub struct RuntimeMemoryStoreDiagnosticsInput {
48    pub attempted: bool,
49    pub response: Option<MemoryStoreResponse>,
50    pub error: Option<String>,
51}
52
53pub struct RuntimeIdentityDiagnosticsInput {
54    pub attempted: bool,
55    pub summary: Option<String>,
56    pub error: Option<String>,
57}
58
59pub fn build_runtime_memory_diagnostics_bundle(
60    memory_recall: RuntimeMemoryRecallDiagnosticsInput,
61    memory_store: RuntimeMemoryStoreDiagnosticsInput,
62    identity: RuntimeIdentityDiagnosticsInput,
63) -> RuntimeMemoryDiagnosticsBundle {
64    let retrieved_count = memory_recall
65        .response
66        .as_ref()
67        .map(|value| value.retrieved)
68        .unwrap_or_default();
69    let retrieval_path = memory_recall
70        .response
71        .as_ref()
72        .and_then(|value| value.retrieval_path.clone());
73    let fallback_triggered = memory_recall
74        .response
75        .as_ref()
76        .map(|value| value.fallback_triggered)
77        .unwrap_or(false);
78    let fallback_reason = memory_recall
79        .response
80        .as_ref()
81        .and_then(|value| value.fallback_reason.clone());
82
83    let store_valid = memory_store
84        .response
85        .as_ref()
86        .map(|value| value.valid)
87        .unwrap_or(false);
88    let store_node_id = memory_store.response.as_ref().map(|value| value.node_id.clone());
89
90    let memory_recall_section = json!({
91        "attempted": memory_recall.attempted,
92        "query_id": memory_recall.query_id,
93        "query_fingerprint": memory_recall.query_fingerprint,
94        "retrieved": retrieved_count,
95        "retrieval_path": retrieval_path,
96        "fallback_triggered": fallback_triggered,
97        "fallback_reason": fallback_reason,
98        "error": memory_recall.error,
99    });
100
101    let memory_store_section = json!({
102        "attempted": memory_store.attempted,
103        "node_id": store_node_id,
104        "valid": store_valid,
105        "error": memory_store.error,
106    });
107
108    let identity_context_section = json!({
109        "attempted": identity.attempted,
110        "summary": identity.summary,
111        "error": identity.error,
112    });
113
114    RuntimeMemoryDiagnosticsBundle {
115        retrieved_count,
116        retrieval_path,
117        fallback_triggered,
118        fallback_reason,
119        store_valid,
120        store_node_id,
121        memory_recall: memory_recall_section,
122        memory_store: memory_store_section,
123        identity_context: identity_context_section,
124    }
125}
126
127pub fn build_runtime_failure_memory_recall_section(
128    attempted: bool,
129    error: Option<String>,
130) -> JsonValue {
131    json!({
132        "attempted": attempted,
133        "error": error,
134    })
135}
136
137pub fn build_runtime_failure_identity_context_section(
138    attempted: bool,
139    summary: Option<String>,
140    error: Option<String>,
141) -> JsonValue {
142    json!({
143        "attempted": attempted,
144        "summary": summary,
145        "error": error,
146    })
147}
148
149#[derive(Clone, Debug, Default)]
150pub struct RuntimeDiagnosticsEnvelope {
151    pub guardrail_code: Option<String>,
152    pub policy_reason: Option<String>,
153    pub duration_ms: Option<u64>,
154    pub input_memory_query_id: Option<String>,
155    pub input_memory_query_fingerprint: Option<String>,
156    pub output_memory_node_id: Option<String>,
157    pub retrieval_path: Option<String>,
158    pub thread_id: Option<String>,
159}
160
161pub fn extract_runtime_diagnostics_envelope(diagnostics: Option<&str>) -> RuntimeDiagnosticsEnvelope {
162    let Some(raw) = diagnostics else {
163        return RuntimeDiagnosticsEnvelope::default();
164    };
165
166    let Ok(json) = serde_json::from_str::<JsonValue>(raw) else {
167        return RuntimeDiagnosticsEnvelope::default();
168    };
169
170    RuntimeDiagnosticsEnvelope {
171        guardrail_code: json
172            .get("guardrail_code")
173            .and_then(|v| v.as_str())
174            .map(str::to_owned),
175        policy_reason: json
176            .get("policy_reason")
177            .and_then(|v| v.as_str())
178            .map(str::to_owned),
179        duration_ms: json.get("duration_ms").and_then(|v| v.as_u64()),
180        input_memory_query_id: json
181            .get("input_memory_query_id")
182            .and_then(|v| v.as_str())
183            .map(str::to_owned),
184        input_memory_query_fingerprint: json
185            .get("input_memory_query_fingerprint")
186            .and_then(|v| v.as_str())
187            .map(str::to_owned)
188            .or_else(|| {
189                json.get("memory_recall")
190                    .and_then(|v| v.get("query_fingerprint"))
191                    .and_then(|v| v.as_str())
192                    .map(str::to_owned)
193            }),
194        output_memory_node_id: json
195            .get("output_memory_node_id")
196            .and_then(|v| v.as_str())
197            .map(str::to_owned)
198            .or_else(|| {
199                json.get("memory_store_node_id")
200                    .and_then(|v| v.as_str())
201                    .map(str::to_owned)
202            }),
203        retrieval_path: json
204            .get("memory_retrieval_path")
205            .and_then(|v| v.as_str())
206            .map(str::to_owned),
207        thread_id: json
208            .get("thread_id")
209            .and_then(|v| v.as_str())
210            .map(str::to_owned),
211    }
212}
213
214pub fn extract_diagnostics_fields(
215    diagnostics: Option<&str>,
216) -> (Option<String>, Option<String>, Option<u64>) {
217    let envelope = extract_runtime_diagnostics_envelope(diagnostics);
218    (
219        envelope.guardrail_code,
220        envelope.policy_reason,
221        envelope.duration_ms,
222    )
223}
224
225pub fn extract_memory_lineage_fields(
226    diagnostics: Option<&str>,
227) -> (
228    Option<String>,
229    Option<String>,
230    Option<String>,
231    Option<String>,
232) {
233    let envelope = extract_runtime_diagnostics_envelope(diagnostics);
234
235    (
236        envelope.input_memory_query_id,
237        envelope.input_memory_query_fingerprint,
238        envelope.output_memory_node_id,
239        envelope.retrieval_path,
240    )
241}
242
243pub fn extract_thread_id(diagnostics: Option<&str>) -> Option<String> {
244    extract_runtime_diagnostics_envelope(diagnostics).thread_id
245}