Skip to content

Commit f288cdc

Browse files
ualtinokalfonso-aft
ualtinok
and
alfonso-aft
committed
mason: add per-root memory attribution
Co-authored-by: Alfonso <289616620+alfonso-aft@users.noreply.github.com>
1 parent fc0ca95 commit f288cdc

20 files changed

Lines changed: 1551 additions & 85 deletions

File tree

‎crates/aft/src/bash_background/registry.rs‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,40 @@ struct RecoveryContext {
150150
include_stderr_path: bool,
151151
}
152152

153+
fn optional_string_bytes(value: Option<&String>) -> u64 {
154+
value
155+
.map(|value| crate::memory::usize_to_u64(value.len()))
156+
.unwrap_or(0)
157+
}
158+
159+
fn terminal_output_cache_estimated_bytes(cache: &TerminalOutputCache) -> u64 {
160+
let recovery_bytes = cache
161+
.recovery
162+
.as_ref()
163+
.map(|recovery| {
164+
crate::memory::usize_to_u64(recovery.dropped_by_class.len())
165+
.saturating_mul(
166+
(std::mem::size_of::<DropClass>() + std::mem::size_of::<usize>()) as u64,
167+
)
168+
.saturating_add(optional_string_bytes(recovery.output_path.as_ref()))
169+
.saturating_add(optional_string_bytes(recovery.stderr_path.as_ref()))
170+
})
171+
.unwrap_or(0);
172+
(std::mem::size_of::<TerminalOutputCache>() as u64)
173+
.saturating_add(crate::memory::usize_to_u64(cache.output_preview.len()))
174+
.saturating_add(optional_string_bytes(cache.output_path.as_ref()))
175+
.saturating_add(optional_string_bytes(cache.stderr_path.as_ref()))
176+
.saturating_add(recovery_bytes)
177+
}
178+
179+
fn completion_estimated_bytes(completion: &BgCompletion) -> u64 {
180+
(std::mem::size_of::<BgCompletion>() as u64)
181+
.saturating_add(crate::memory::usize_to_u64(completion.task_id.len()))
182+
.saturating_add(crate::memory::usize_to_u64(completion.session_id.len()))
183+
.saturating_add(crate::memory::usize_to_u64(completion.command.len()))
184+
.saturating_add(crate::memory::usize_to_u64(completion.output_preview.len()))
185+
}
186+
153187
impl RecoveryContext {
154188
fn has_visible_drop(&self) -> bool {
155189
self.byte_truncated
@@ -2740,6 +2774,47 @@ impl BgTaskRegistry {
27402774
})
27412775
}
27422776

2777+
/// Estimate resident bash output caches without reading disk-backed task
2778+
/// streams. Spill files are deliberately excluded because they do not
2779+
/// occupy the daemon heap.
2780+
pub fn estimated_memory(&self) -> crate::memory::MemoryEstimate {
2781+
let tasks = match self.inner.tasks.try_lock() {
2782+
Ok(tasks) => tasks.values().cloned().collect::<Vec<_>>(),
2783+
Err(_) => return crate::memory::MemoryEstimate::busy(),
2784+
};
2785+
let mut bytes = 0u64;
2786+
let mut terminal_output_caches = 0usize;
2787+
let mut sessions = HashSet::new();
2788+
for task in &tasks {
2789+
sessions.insert(task.session_id.clone());
2790+
let state = match task.state.try_lock() {
2791+
Ok(state) => state,
2792+
Err(_) => return crate::memory::MemoryEstimate::busy(),
2793+
};
2794+
if let Some(cache) = state.terminal_output_cache.as_ref() {
2795+
terminal_output_caches = terminal_output_caches.saturating_add(1);
2796+
bytes = bytes.saturating_add(terminal_output_cache_estimated_bytes(cache));
2797+
}
2798+
}
2799+
let completion_count = match self.inner.completions.try_lock() {
2800+
Ok(completions) => {
2801+
for completion in completions.iter() {
2802+
sessions.insert(completion.session_id.clone());
2803+
bytes = bytes.saturating_add(completion_estimated_bytes(completion));
2804+
}
2805+
completions.len()
2806+
}
2807+
Err(_) => return crate::memory::MemoryEstimate::busy(),
2808+
};
2809+
2810+
crate::memory::MemoryEstimate::estimated(bytes)
2811+
.count("tasks", tasks.len())
2812+
.count("sessions", sessions.len())
2813+
.count("terminal_output_caches", terminal_output_caches)
2814+
.count("completion_caches", completion_count)
2815+
.count_u64("output_ring_bytes", 0)
2816+
}
2817+
27432818
fn running_count(&self) -> usize {
27442819
self.inner
27452820
.tasks
@@ -4059,6 +4134,33 @@ mod tests {
40594134
#[cfg(windows)]
40604135
const LONG_RUNNING_COMMAND: &str = "cmd /c timeout /t 5 /nobreak > nul";
40614136

4137+
#[test]
4138+
fn bash_memory_estimate_is_zero_when_empty_and_nonzero_for_completion_cache() {
4139+
let registry = BgTaskRegistry::default();
4140+
assert_eq!(registry.estimated_memory().estimated_bytes, Some(0));
4141+
registry
4142+
.inner
4143+
.completions
4144+
.lock()
4145+
.unwrap()
4146+
.push_back(BgCompletion {
4147+
task_id: "bash-memory".to_string(),
4148+
session_id: "session-memory".to_string(),
4149+
status: BgTaskStatus::Completed,
4150+
exit_code: Some(0),
4151+
command: "printf memory".to_string(),
4152+
output_preview: "resident completion output".to_string(),
4153+
output_truncated: false,
4154+
original_tokens: None,
4155+
compressed_tokens: None,
4156+
tokens_skipped: false,
4157+
});
4158+
let estimate = registry.estimated_memory();
4159+
assert!(estimate.estimated_bytes.unwrap() > 0);
4160+
assert_eq!(estimate.counts["completion_caches"], 1);
4161+
assert_eq!(estimate.counts["sessions"], 1);
4162+
}
4163+
40624164
#[test]
40634165
fn gh_structured_detection_rejects_piped_commands() {
40644166
assert!(is_gh_structured_command(

‎crates/aft/src/callgraph_store/mod.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2102,6 +2102,15 @@ impl ReadonlyCallGraphStore {
21022102
self.inner.sqlite_path()
21032103
}
21042104

2105+
/// Report the open generation handle without guessing at SQLite's internal
2106+
/// page or prepared-statement caches.
2107+
pub fn estimated_memory(&self) -> crate::memory::MemoryEstimate {
2108+
crate::memory::MemoryEstimate::not_estimated()
2109+
.count("open_generation_handles", 1)
2110+
.gap("sqlite_internal_bytes")
2111+
.gap("prepared_statement_cache_entries")
2112+
}
2113+
21052114
/// Whether this reader is temporarily serving a legacy harness partition.
21062115
pub fn is_legacy_fallback(&self) -> bool {
21072116
self.inner.is_legacy_fallback()

‎crates/aft/src/commands/status.rs‎

Lines changed: 93 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -45,13 +45,10 @@ impl AppContext {
4545
pub fn build_status_snapshot_for_session(&self, session_id: &str) -> StatusPayload {
4646
let config = self.config();
4747

48-
// Search index status
49-
let search_index_info = {
50-
let index = self
51-
.search_index()
52-
.read()
53-
.unwrap_or_else(std::sync::PoisonError::into_inner);
54-
match index.as_ref() {
48+
// Search index status. Status is a control-path snapshot, so lock
49+
// pressure is represented directly instead of delaying the caller.
50+
let search_index_info = match self.search_index().try_read() {
51+
Ok(index) => match index.as_ref() {
5552
Some(idx) if idx.ready => {
5653
let file_count = idx.file_count();
5754
let trigram_count = idx.trigram_count();
@@ -70,78 +67,80 @@ impl AppContext {
7067
};
7168
serde_json::json!({ "status": status })
7269
}
73-
}
70+
},
71+
Err(_) => serde_json::json!({ "status": "busy" }),
7472
};
7573

76-
// Semantic index status
77-
let semantic_index_info = {
78-
let status = self
79-
.semantic_index_status()
80-
.read()
81-
.unwrap_or_else(std::sync::PoisonError::into_inner)
82-
.clone();
83-
let refreshing_count = status.refreshing_count();
84-
let index = self
85-
.semantic_index()
86-
.read()
87-
.unwrap_or_else(std::sync::PoisonError::into_inner);
88-
match index.as_ref() {
89-
Some(idx) => {
90-
let status_label = match status {
91-
SemanticIndexStatus::Ready { .. } => "ready",
92-
_ => idx.status_label(),
93-
};
94-
serde_json::json!({
95-
"status": status_label,
96-
"state": status_label,
97-
"refreshing_count": refreshing_count,
98-
"entries": idx.entry_count(),
99-
"dimension": idx.dimension(),
100-
"backend": idx.backend_label().unwrap_or(config.semantic_backend_label()),
101-
"model": idx.model_label().unwrap_or(config.semantic.model.as_str()),
102-
})
74+
let semantic_status = self
75+
.semantic_index_status()
76+
.try_read()
77+
.ok()
78+
.map(|status| status.clone());
79+
let semantic_index_info = match semantic_status {
80+
None => serde_json::json!({ "status": "busy", "state": "busy" }),
81+
Some(status) => match self.semantic_index().try_read() {
82+
Err(_) => serde_json::json!({ "status": "busy", "state": "busy" }),
83+
Ok(index) => {
84+
let refreshing_count = status.refreshing_count();
85+
match index.as_ref() {
86+
Some(idx) => {
87+
let status_label = match status {
88+
SemanticIndexStatus::Ready { .. } => "ready",
89+
_ => idx.status_label(),
90+
};
91+
serde_json::json!({
92+
"status": status_label,
93+
"state": status_label,
94+
"refreshing_count": refreshing_count,
95+
"entries": idx.entry_count(),
96+
"dimension": idx.dimension(),
97+
"backend": idx.backend_label().unwrap_or(config.semantic_backend_label()),
98+
"model": idx.model_label().unwrap_or(config.semantic.model.as_str()),
99+
})
100+
}
101+
None => match status {
102+
SemanticIndexStatus::Disabled => serde_json::json!({
103+
"status": "disabled",
104+
"state": "disabled",
105+
"refreshing_count": 0,
106+
"backend": config.semantic_backend_label(),
107+
"model": config.semantic.model.as_str(),
108+
}),
109+
SemanticIndexStatus::Building {
110+
stage,
111+
files,
112+
entries_done,
113+
entries_total,
114+
} => serde_json::json!({
115+
"status": "loading",
116+
"state": "loading",
117+
"refreshing_count": 0,
118+
"stage": stage,
119+
"files": files,
120+
"entries_done": entries_done,
121+
"entries_total": entries_total,
122+
"backend": config.semantic_backend_label(),
123+
"model": config.semantic.model.as_str(),
124+
}),
125+
SemanticIndexStatus::Ready { refreshing, .. } => serde_json::json!({
126+
"status": "ready",
127+
"state": "ready",
128+
"refreshing_count": refreshing.len(),
129+
"backend": config.semantic_backend_label(),
130+
"model": config.semantic.model.as_str(),
131+
}),
132+
SemanticIndexStatus::Failed(error) => serde_json::json!({
133+
"status": "failed",
134+
"state": "failed",
135+
"refreshing_count": 0,
136+
"error": error,
137+
"backend": config.semantic_backend_label(),
138+
"model": config.semantic.model.as_str(),
139+
}),
140+
},
141+
}
103142
}
104-
None => match status {
105-
SemanticIndexStatus::Disabled => serde_json::json!({
106-
"status": "disabled",
107-
"state": "disabled",
108-
"refreshing_count": 0,
109-
"backend": config.semantic_backend_label(),
110-
"model": config.semantic.model.as_str(),
111-
}),
112-
SemanticIndexStatus::Building {
113-
stage,
114-
files,
115-
entries_done,
116-
entries_total,
117-
} => serde_json::json!({
118-
"status": "loading",
119-
"state": "loading",
120-
"refreshing_count": 0,
121-
"stage": stage,
122-
"files": files,
123-
"entries_done": entries_done,
124-
"entries_total": entries_total,
125-
"backend": config.semantic_backend_label(),
126-
"model": config.semantic.model.as_str(),
127-
}),
128-
SemanticIndexStatus::Ready { refreshing, .. } => serde_json::json!({
129-
"status": "ready",
130-
"state": "ready",
131-
"refreshing_count": refreshing.len(),
132-
"backend": config.semantic_backend_label(),
133-
"model": config.semantic.model.as_str(),
134-
}),
135-
SemanticIndexStatus::Failed(error) => serde_json::json!({
136-
"status": "failed",
137-
"state": "failed",
138-
"refreshing_count": 0,
139-
"error": error,
140-
"backend": config.semantic_backend_label(),
141-
"model": config.semantic.model.as_str(),
142-
}),
143-
},
144-
}
143+
},
145144
};
146145

147146
// Disk cache sizes — scoped to the **current project** only.
@@ -258,6 +257,11 @@ impl AppContext {
258257
}),
259258
None => serde_json::Value::Null,
260259
};
260+
let memory_root = self
261+
.canonical_cache_root_opt()
262+
.or_else(|| config.project_root.clone());
263+
let memory = serde_json::to_value(self.memory_snapshot(memory_root.as_deref()))
264+
.unwrap_or(serde_json::Value::Null);
261265

262266
serde_json::json!({
263267
"version": env!("CARGO_PKG_VERSION"),
@@ -282,6 +286,7 @@ impl AppContext {
282286
"disk": disk_info,
283287
"lsp_servers": lsp_count,
284288
"symbol_cache": symbol_cache_stats,
289+
"memory": memory,
285290
"compression": compression,
286291
"storage_dir": storage_dir,
287292
// Project-wide (all sessions): total in-memory checkpoint count.
@@ -396,6 +401,18 @@ mod tests {
396401
assert_eq!(response.data["cache_role"], "worktree");
397402
}
398403

404+
#[test]
405+
fn memory_snapshot_reports_contended_subsystem_as_busy() {
406+
let ctx = AppContext::new(Box::new(TreeSitterProvider::new()), Config::default());
407+
let _semantic_writer = ctx.semantic_index().write().unwrap();
408+
let status = ctx.build_status_snapshot();
409+
assert_eq!(status["semantic_index"]["status"], "busy");
410+
assert_eq!(
411+
status["memory"]["roots"]["<unconfigured>"]["semantic"]["status"],
412+
"busy"
413+
);
414+
}
415+
399416
#[test]
400417
fn status_status_bar_is_null_until_tier2_populated() {
401418
let ctx = AppContext::new(Box::new(TreeSitterProvider::new()), Config::default());

0 commit comments

Comments
 (0)