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