a3s_code_core/agent/
model_middleware_obs.rs1use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::Arc;
5
6#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq, Default)]
11#[serde(rename_all = "camelCase", deny_unknown_fields)]
12pub struct ModelMiddlewareHealthSnapshot {
13 pub trust_admitted: u64,
15 pub trust_rejected: u64,
17 pub provider_calls: u64,
19 pub usage_recorded: u64,
21 pub completion_calls: u64,
23 pub streaming_calls: u64,
25 pub trusted_tool_results: u64,
27 pub workspace_data_tool_results: u64,
29 pub external_tool_results: u64,
31 pub external_redaction_reviewed: u64,
33}
34
35#[derive(Debug, Default)]
36pub(crate) struct ModelMiddlewareObs {
37 trust_admitted: AtomicU64,
38 trust_rejected: AtomicU64,
39 provider_calls: AtomicU64,
40 usage_recorded: AtomicU64,
41 completion_calls: AtomicU64,
42 streaming_calls: AtomicU64,
43 trusted_tool_results: AtomicU64,
44 workspace_data_tool_results: AtomicU64,
45 external_tool_results: AtomicU64,
46 external_redaction_reviewed: AtomicU64,
47}
48
49impl ModelMiddlewareObs {
50 pub(crate) fn shared() -> Arc<Self> {
51 Arc::new(Self::default())
52 }
53
54 pub(crate) fn snapshot(&self) -> ModelMiddlewareHealthSnapshot {
55 ModelMiddlewareHealthSnapshot {
56 trust_admitted: self.trust_admitted.load(Ordering::Relaxed),
57 trust_rejected: self.trust_rejected.load(Ordering::Relaxed),
58 provider_calls: self.provider_calls.load(Ordering::Relaxed),
59 usage_recorded: self.usage_recorded.load(Ordering::Relaxed),
60 completion_calls: self.completion_calls.load(Ordering::Relaxed),
61 streaming_calls: self.streaming_calls.load(Ordering::Relaxed),
62 trusted_tool_results: self.trusted_tool_results.load(Ordering::Relaxed),
63 workspace_data_tool_results: self.workspace_data_tool_results.load(Ordering::Relaxed),
64 external_tool_results: self.external_tool_results.load(Ordering::Relaxed),
65 external_redaction_reviewed: self.external_redaction_reviewed.load(Ordering::Relaxed),
66 }
67 }
68
69 fn bump(counter: &AtomicU64, by: u64) {
70 if by == 0 {
71 return;
72 }
73 counter.fetch_add(by, Ordering::Relaxed);
74 }
75
76 pub(crate) fn record_trust_admitted(
77 &self,
78 streaming: bool,
79 trusted: u32,
80 workspace_data: u32,
81 external: u32,
82 external_reviewed: u32,
83 ) {
84 self.trust_admitted.fetch_add(1, Ordering::Relaxed);
85 if streaming {
86 self.streaming_calls.fetch_add(1, Ordering::Relaxed);
87 } else {
88 self.completion_calls.fetch_add(1, Ordering::Relaxed);
89 }
90 Self::bump(&self.trusted_tool_results, u64::from(trusted));
91 Self::bump(&self.workspace_data_tool_results, u64::from(workspace_data));
92 Self::bump(&self.external_tool_results, u64::from(external));
93 Self::bump(
94 &self.external_redaction_reviewed,
95 u64::from(external_reviewed),
96 );
97 }
98
99 pub(crate) fn record_trust_rejected(&self) {
100 self.trust_rejected.fetch_add(1, Ordering::Relaxed);
101 }
102
103 pub(crate) fn record_provider_call(&self) {
104 self.provider_calls.fetch_add(1, Ordering::Relaxed);
105 }
106
107 pub(crate) fn record_usage(&self) {
108 self.usage_recorded.fetch_add(1, Ordering::Relaxed);
109 }
110}