use std::sync::Arc;
use meerkat_mobkit::identity_first::{
ContinuitySessionStoreAdapter, ContinuityStore, LocalContinuityStore, SessionRuntimeState,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SelectionMode {
ContentDedup,
KeepNewest(usize),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct ToolResultBounds {
truncate_max_bytes: Option<usize>,
drop_older_than: Option<usize>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RepairMode {
System(SelectionMode),
ToolResults(ToolResultBounds),
}
const TOOL_RESULT_MARKER_PREFIX: &str = "[mobkit-repair:tool-result-";
#[derive(Debug, serde::Serialize)]
struct SessionReport {
session_id: String,
identity: Option<String>,
outcome: String,
#[serde(skip_serializing_if = "Option::is_none")]
detail: Option<String>,
rows_before: Option<usize>,
rows_after: Option<usize>,
bytes_before: Option<usize>,
bytes_after: Option<usize>,
system_rows_before: Option<usize>,
system_rows_after: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_result_messages: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_results_bounded: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_result_bytes_before: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_result_bytes_after: Option<usize>,
}
#[derive(Debug, Clone, Copy)]
struct ToolResultAccounting {
messages: usize,
bounded: usize,
bytes_before: usize,
bytes_after: usize,
}
#[derive(Debug, serde::Serialize)]
struct RepairReport {
mode: String,
applied: bool,
sessions: Vec<SessionReport>,
refused: usize,
bytes_semantics: &'static str,
}
const BYTES_SEMANTICS: &str = "bytes_before/bytes_after measure the live HEAD transcript (what \
a resume serves and the provider reads), not the on-disk file: audited rewrites retain \
prior strands by design, so the store file does not shrink after apply";
#[tokio::main(flavor = "current_thread")]
async fn main() {
match run().await {
Ok(exit) => std::process::exit(exit),
Err(error) => {
eprintln!("mobkit-repair: {error}");
std::process::exit(2);
}
}
}
async fn run() -> Result<i32, Box<dyn std::error::Error>> {
let mut db: Option<String> = None;
let mut sessions: Vec<String> = Vec::new();
let mut all_sessions = false;
let mut keep_newest: Option<usize> = None;
let mut truncate_tool_results: Option<usize> = None;
let mut drop_tool_results_older_than: Option<usize> = None;
let mut apply = false;
let mut report_path: Option<String> = None;
let mut args = std::env::args().skip(1);
while let Some(arg) = args.next() {
match arg.as_str() {
"--db" => db = args.next(),
"--session" => {
sessions.push(args.next().ok_or("--session requires a value")?);
}
"--all-sessions" => all_sessions = true,
"--keep-newest" => {
let value = args.next().ok_or("--keep-newest requires a value")?;
let n: usize = value
.parse()
.map_err(|_| format!("--keep-newest must be a number: {value}"))?;
if n == 0 {
return Err("--keep-newest must be >= 1 (a session keeps at least \
its newest System row)"
.into());
}
keep_newest = Some(n);
}
"--truncate-tool-results" => {
let value = args
.next()
.ok_or("--truncate-tool-results requires a value")?;
let n: usize = value.parse().map_err(|_| {
format!("--truncate-tool-results must be a byte count: {value}")
})?;
if n == 0 {
return Err("--truncate-tool-results must be >= 1 byte".into());
}
truncate_tool_results = Some(n);
}
"--drop-tool-results-older-than" => {
let value = args
.next()
.ok_or("--drop-tool-results-older-than requires a value")?;
let n: usize = value.parse().map_err(|_| {
format!("--drop-tool-results-older-than must be a message count: {value}")
})?;
if n == 0 {
return Err("--drop-tool-results-older-than must be >= 1 (the \
tail row's own results always stay)"
.into());
}
drop_tool_results_older_than = Some(n);
}
"--apply" => apply = true,
"--report" => report_path = args.next(),
other => return Err(format!("unknown argument: {other}").into()),
}
}
let db = db.ok_or("--db <path to the continuity store> is required")?;
if all_sessions && !sessions.is_empty() {
return Err("--all-sessions and --session are mutually exclusive".into());
}
if !all_sessions && sessions.is_empty() {
return Err("select sessions with --session <uuid> (repeatable) or --all-sessions".into());
}
let tool_result_bounds = match (truncate_tool_results, drop_tool_results_older_than) {
(None, None) => None,
(truncate_max_bytes, drop_older_than) => Some(ToolResultBounds {
truncate_max_bytes,
drop_older_than,
}),
};
if keep_newest.is_some() && tool_result_bounds.is_some() {
return Err("--keep-newest selects System rows; it cannot combine with \
the tool-result bounding flags"
.into());
}
let mode = match tool_result_bounds {
Some(bounds) => RepairMode::ToolResults(bounds),
None => RepairMode::System(
keep_newest.map_or(SelectionMode::ContentDedup, SelectionMode::KeepNewest),
),
};
let local = Arc::new(LocalContinuityStore::open(&db)?);
let targets: Vec<(Option<String>, meerkat_core::types::SessionId)> = if all_sessions {
local
.list_session_bindings()
.await?
.into_iter()
.map(|(identity, session_id)| (Some(identity.to_string()), session_id))
.collect()
} else {
let mut targets = Vec::new();
for raw in &sessions {
let session_id = meerkat_core::types::SessionId::parse(raw)
.map_err(|e| format!("--session must be a session UUID: {e}"))?;
targets.push((None, session_id));
}
targets
};
if targets.is_empty() {
return Err("the continuity store holds no session bindings".into());
}
let continuity: Arc<dyn ContinuityStore> = local;
let mut reports = Vec::with_capacity(targets.len());
let mut refused = 0usize;
for (identity_hint, session_id) in targets {
eprintln!("mobkit-repair: session {session_id} ...");
match repair_session(&continuity, &session_id, mode, apply).await {
Ok(report) => reports.push(report),
Err(error) => {
refused += 1;
eprintln!("mobkit-repair: session {session_id} REFUSED: {error}");
reports.push(SessionReport {
session_id: session_id.to_string(),
identity: identity_hint,
outcome: "refused".to_string(),
detail: Some(error.to_string()),
rows_before: None,
rows_after: None,
bytes_before: None,
bytes_after: None,
system_rows_before: None,
system_rows_after: None,
tool_result_messages: None,
tool_results_bounded: None,
tool_result_bytes_before: None,
tool_result_bytes_after: None,
});
}
}
}
let report = RepairReport {
mode: match mode {
RepairMode::System(SelectionMode::ContentDedup) => "content_dedup".to_string(),
RepairMode::System(SelectionMode::KeepNewest(n)) => format!("keep_newest_{n}"),
RepairMode::ToolResults(bounds) => {
let mut parts = Vec::new();
if let Some(n) = bounds.truncate_max_bytes {
parts.push(format!("truncate_tool_results_{n}"));
}
if let Some(n) = bounds.drop_older_than {
parts.push(format!("drop_tool_results_older_than_{n}"));
}
parts.join("+")
}
},
applied: apply,
sessions: reports,
refused,
bytes_semantics: BYTES_SEMANTICS,
};
let rendered = serde_json::to_string_pretty(&report)?;
println!("{rendered}");
if let Some(path) = report_path {
std::fs::write(&path, &rendered)?;
eprintln!("mobkit-repair: report written to {path}");
}
if apply
&& report
.sessions
.iter()
.any(|session| session.outcome == "applied")
{
eprintln!();
eprintln!(
"mobkit-repair: APPLY COMPLETE - ONE STEP REMAINS BEFORE RESTARTING THE GATEWAY:"
);
eprintln!(
"mobkit-repair: remove the runtime scratch store (the file set next to this \
continuity store; `rkat storage doctor` names it)."
);
eprintln!(
"mobkit-repair: It still projects the PRE-repair transcript; the next boot must \
mint runtime authority from the healed durable rows (the reset-reseed lane). \
Skipping this step makes the repair appear to have not happened."
);
eprintln!(
"mobkit-repair: Note: the store FILE does not shrink - audited rewrites retain \
prior strands; the healed HEAD is what got smaller (see bytes_semantics in the \
report)."
);
}
Ok(i32::from(refused != 0))
}
async fn repair_session(
continuity: &Arc<dyn ContinuityStore>,
session_id: &meerkat_core::types::SessionId,
mode: RepairMode,
apply: bool,
) -> Result<SessionReport, Box<dyn std::error::Error>> {
let adapter = Arc::new(ContinuitySessionStoreAdapter::new(Arc::clone(continuity)));
let (record, fencing_token, fence_current) = continuity
.resolve_record_by_session(session_id)
.await?
.ok_or("no continuity record binds this session")?;
adapter
.register_session(
session_id,
SessionRuntimeState {
identity: record.identity.clone(),
generation: record.generation,
fencing_token,
checkpoint_version: fence_current,
},
)
.await?;
let mut session = meerkat::SessionStore::load(adapter.as_ref(), session_id)
.await?
.ok_or("no durable row for this session")?;
let channel = continuity
.as_incremental_sessions()
.ok_or("the continuity store provides no incremental channel")?;
let rewrite_records = channel.load_rewrites(session_id).await?;
if let Some(validated) =
meerkat_core::ValidatedTranscriptHistory::from_rewrite_records_with_proved(
rewrite_records,
None,
)?
{
session.install_validated_audited_transcript_history_preserving_live(validated)?;
}
let messages = session.messages();
let serialized: Vec<Vec<u8>> = messages
.iter()
.map(serde_json::to_vec)
.collect::<Result<_, _>>()?;
let values: Vec<serde_json::Value> = serialized
.iter()
.map(|bytes| serde_json::from_slice(bytes))
.collect::<Result<_, _>>()?;
let is_system = |value: &serde_json::Value| -> bool {
value
.get("role")
.and_then(|r| r.as_str())
.map(|r| r == "system")
.unwrap_or(false)
};
let system_positions: Vec<usize> = values
.iter()
.enumerate()
.filter(|(_, value)| is_system(value))
.map(|(index, _)| index)
.collect();
let rows_before = messages.len();
let bytes_before: usize = serialized.iter().map(Vec::len).sum();
let system_before = system_positions.len();
let (cleaned, system_after, tool_accounting) = match mode {
RepairMode::System(selection) => {
let duplicate_indices: Vec<usize> = match selection {
SelectionMode::ContentDedup => {
let mut seen_groups: std::collections::HashMap<String, usize> =
std::collections::HashMap::new();
let mut duplicates = Vec::new();
for &index in &system_positions {
let value = &values[index];
let key = serde_json::to_string(&serde_json::json!({
"content": value.get("content"),
"identity": value.get("identity"),
}))?;
match seen_groups.entry(key) {
std::collections::hash_map::Entry::Vacant(slot) => {
slot.insert(index);
}
std::collections::hash_map::Entry::Occupied(_) => {
duplicates.push(index);
}
}
}
duplicates
}
SelectionMode::KeepNewest(n) => {
let cut = system_positions.len().saturating_sub(n);
system_positions[..cut].to_vec()
}
};
if duplicate_indices.is_empty() {
eprintln!(
"mobkit-repair: session {session_id}: {system_before} System row(s), \
nothing to drop"
);
return Ok(SessionReport {
session_id: session_id.to_string(),
identity: Some(record.identity.to_string()),
outcome: "nothing_to_do".to_string(),
detail: None,
rows_before: Some(rows_before),
rows_after: Some(rows_before),
bytes_before: Some(bytes_before),
bytes_after: Some(bytes_before),
system_rows_before: Some(system_before),
system_rows_after: Some(system_before),
tool_result_messages: None,
tool_results_bounded: None,
tool_result_bytes_before: None,
tool_result_bytes_after: None,
});
}
let cleaned: Vec<meerkat_core::Message> = messages
.iter()
.enumerate()
.filter(|(index, _)| !duplicate_indices.contains(index))
.map(|(_, message)| message.clone())
.collect();
let system_after = system_before - duplicate_indices.len();
(cleaned, system_after, None)
}
RepairMode::ToolResults(bounds) => {
let plan = plan_tool_result_bounds(messages, bounds)?;
let accounting = ToolResultAccounting {
messages: plan.tool_result_messages,
bounded: plan.bounded,
bytes_before: plan.tool_bytes_before,
bytes_after: plan.tool_bytes_after,
};
if plan.bounded == 0 {
eprintln!(
"mobkit-repair: session {session_id}: {} tool-result row(s), \
nothing to bound",
accounting.messages
);
return Ok(SessionReport {
session_id: session_id.to_string(),
identity: Some(record.identity.to_string()),
outcome: "nothing_to_do".to_string(),
detail: None,
rows_before: Some(rows_before),
rows_after: Some(rows_before),
bytes_before: Some(bytes_before),
bytes_after: Some(bytes_before),
system_rows_before: Some(system_before),
system_rows_after: Some(system_before),
tool_result_messages: Some(accounting.messages),
tool_results_bounded: Some(0),
tool_result_bytes_before: Some(accounting.bytes_before),
tool_result_bytes_after: Some(accounting.bytes_after),
});
}
(plan.cleaned, system_before, Some(accounting))
}
};
let rows_after = cleaned.len();
let bytes_after: usize = cleaned
.iter()
.map(|m| serde_json::to_vec(m).map(|v| v.len()).unwrap_or(0))
.sum();
eprintln!(
"mobkit-repair: session {session_id}: plan {rows_before} -> {rows_after} rows, \
~{bytes_before} -> ~{bytes_after} bytes, System {system_before} -> {system_after}"
);
if let Some(tool) = tool_accounting {
eprintln!(
"mobkit-repair: session {session_id}: bounding {} tool-result entry(ies) \
across {} tool-result row(s) in place, ~{} -> ~{} tool-result bytes",
tool.bounded, tool.messages, tool.bytes_before, tool.bytes_after
);
}
if !apply {
return Ok(SessionReport {
session_id: session_id.to_string(),
identity: Some(record.identity.to_string()),
outcome: "dry_run".to_string(),
detail: None,
rows_before: Some(rows_before),
rows_after: Some(rows_after),
bytes_before: Some(bytes_before),
bytes_after: Some(bytes_after),
system_rows_before: Some(system_before),
system_rows_after: Some(system_after),
tool_result_messages: tool_accounting.map(|t| t.messages),
tool_results_bounded: tool_accounting.map(|t| t.bounded),
tool_result_bytes_before: tool_accounting.map(|t| t.bytes_before),
tool_result_bytes_after: tool_accounting.map(|t| t.bytes_after),
});
}
let (rewrite_reason, rewrite_provenance) = match mode {
RepairMode::System(_) => (
"operator prune of superseded System rows (mobkit-repair)",
"mobkit-repair/system-rows",
),
RepairMode::ToolResults(_) => (
"operator bound of tool-result bulk (mobkit-repair)",
"mobkit-repair/tool-results",
),
};
let parent_revision = session.transcript_revision()?;
let selection_end = messages.len();
let mut rewritten = session.clone();
let commit = rewritten.commit_transcript_rewrite(
meerkat_core::TranscriptRewriteSelection::MessageRange {
start: 0,
end: selection_end,
},
cleaned,
meerkat_core::TranscriptRewriteReason::new(rewrite_reason),
Some(rewrite_provenance.to_string()),
Some(parent_revision),
)?;
meerkat::SessionStore::save_transcript_rewrite(adapter.as_ref(), &rewritten, &commit).await?;
meerkat::SessionStore::save_authoritative_projection(adapter.as_ref(), &rewritten).await?;
eprintln!(
"mobkit-repair: session {session_id}: APPLIED at rewrite generation {}",
commit.rewrite_generation
);
Ok(SessionReport {
session_id: session_id.to_string(),
identity: Some(record.identity.to_string()),
outcome: "applied".to_string(),
detail: None,
rows_before: Some(rows_before),
rows_after: Some(rows_after),
bytes_before: Some(bytes_before),
bytes_after: Some(bytes_after),
system_rows_before: Some(system_before),
system_rows_after: Some(system_after),
tool_result_messages: tool_accounting.map(|t| t.messages),
tool_results_bounded: tool_accounting.map(|t| t.bounded),
tool_result_bytes_before: tool_accounting.map(|t| t.bytes_before),
tool_result_bytes_after: tool_accounting.map(|t| t.bytes_after),
})
}
struct ToolResultBoundsPlan {
cleaned: Vec<meerkat_core::Message>,
tool_result_messages: usize,
bounded: usize,
tool_bytes_before: usize,
tool_bytes_after: usize,
}
fn is_already_bounded(result: &meerkat_core::types::ToolResult) -> bool {
match result.content.as_slice() {
[meerkat_core::types::ContentBlock::Text { text }] => {
text.starts_with(TOOL_RESULT_MARKER_PREFIX)
}
_ => false,
}
}
fn plan_tool_result_bounds(
messages: &[meerkat_core::Message],
bounds: ToolResultBounds,
) -> Result<ToolResultBoundsPlan, serde_json::Error> {
let total = messages.len();
let mut cleaned = Vec::with_capacity(total);
let mut tool_result_messages = 0usize;
let mut bounded = 0usize;
let mut tool_bytes_before = 0usize;
let mut tool_bytes_after = 0usize;
for (index, message) in messages.iter().enumerate() {
let meerkat_core::Message::ToolResults {
results,
created_at,
} = message
else {
cleaned.push(message.clone());
continue;
};
tool_result_messages += 1;
let distance_from_tail = total - index;
let elide = bounds
.drop_older_than
.is_some_and(|n| distance_from_tail > n);
let mut bounded_results = Vec::with_capacity(results.len());
for result in results {
let entry_bytes = serde_json::to_vec(result)?.len();
tool_bytes_before += entry_bytes;
let replacement = if is_already_bounded(result) {
None
} else if elide {
Some(format!(
"{TOOL_RESULT_MARKER_PREFIX}elided original_bytes={entry_bytes}]"
))
} else {
match bounds.truncate_max_bytes {
Some(max_bytes) if entry_bytes > max_bytes => {
let text = result.text_content();
let mut cut = max_bytes.min(text.len());
while cut > 0 && !text.is_char_boundary(cut) {
cut -= 1;
}
Some(format!(
"{TOOL_RESULT_MARKER_PREFIX}truncated \
original_bytes={entry_bytes} kept_bytes={cut}]\n{}",
&text[..cut]
))
}
_ => None,
}
};
let entry = match replacement {
Some(marker_text) => {
bounded += 1;
let mut bounded_entry = result.clone();
bounded_entry.set_text_content(marker_text);
bounded_entry
}
None => result.clone(),
};
tool_bytes_after += serde_json::to_vec(&entry)?.len();
bounded_results.push(entry);
}
cleaned.push(meerkat_core::Message::ToolResults {
results: bounded_results,
created_at: *created_at,
});
}
Ok(ToolResultBoundsPlan {
cleaned,
tool_result_messages,
bounded,
tool_bytes_before,
tool_bytes_after,
})
}