forge_guardrails/tool_output/state/
dedup.rs

1use indexmap::IndexMap;
2use std::collections::{hash_map::DefaultHasher, VecDeque};
3use std::hash::{Hash, Hasher};
4use std::sync::Mutex;
5
6use super::memo::{MemoLookup, MemoRecord, MemoState};
7
8/// Bounded in-memory dedup and memo state for compressed tool outputs.
9#[derive(Debug, Default)]
10pub struct ToolOutputCompressionState {
11    inner: Mutex<DedupState>,
12    memo: Mutex<MemoState>,
13}
14
15#[derive(Debug, Default)]
16struct DedupState {
17    sessions: IndexMap<String, VecDeque<DedupRecord>>,
18    session_order: VecDeque<String>,
19}
20
21#[derive(Debug, Clone)]
22struct DedupRecord {
23    hash: u64,
24    output_len: usize,
25    tool_name: String,
26    tool_call_id: String,
27}
28
29impl ToolOutputCompressionState {
30    /// Create empty bounded compression state.
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    pub(in crate::tool_output) fn deduplicate(
36        &self,
37        session_id: &str,
38        tool_call_id: &str,
39        tool_name: &str,
40        output: &str,
41        max_sessions: usize,
42        max_entries_per_session: usize,
43    ) -> Option<String> {
44        if session_id.is_empty() || tool_call_id.is_empty() || output.is_empty() {
45            return None;
46        }
47        let max_sessions = max_sessions.max(1);
48        let max_entries_per_session = max_entries_per_session.max(1);
49        let hash = hash_output(output);
50        let mut state = self.inner.lock().expect("tool output dedup lock");
51
52        let output_len = output.len();
53        if let Some(records) = state.sessions.get(session_id) {
54            if let Some(record) = records.iter().find(|record| {
55                record.tool_name == tool_name
56                    && record.hash == hash
57                    && record.output_len == output_len
58            }) {
59                // The same call re-sent in a later request must keep its
60                // content; only a different call with identical output is a
61                // true duplicate.
62                if record.tool_call_id == tool_call_id {
63                    return None;
64                }
65                // Collapse the LATER duplicate, not the earlier one. The
66                // earlier message anchors the prompt-cache prefix; rewriting
67                // it on every resend would bust the prefix on every subsequent
68                // request. The later message is cheapest to replace: it has
69                // not yet been cached. Near-duplicate delta encoding was
70                // considered and rejected because it breaks determinism when
71                // the base entry is evicted under FIFO pressure.
72                return Some(format!(
73                    "[Duplicate of {} ({}); see earlier result]",
74                    record.tool_call_id, record.tool_name
75                ));
76            }
77        }
78
79        if !state.sessions.contains_key(session_id) {
80            state.session_order.push_back(session_id.to_string());
81            state
82                .sessions
83                .insert(session_id.to_string(), VecDeque::new());
84        }
85
86        let records = state
87            .sessions
88            .get_mut(session_id)
89            .expect("session inserted above");
90        records.push_back(DedupRecord {
91            hash,
92            output_len,
93            tool_name: tool_name.to_string(),
94            tool_call_id: tool_call_id.to_string(),
95        });
96        while records.len() > max_entries_per_session {
97            records.pop_front();
98        }
99
100        while state.sessions.len() > max_sessions {
101            let Some(oldest) = state.session_order.pop_front() else {
102                break;
103            };
104            if oldest != session_id {
105                state.sessions.shift_remove(&oldest);
106            } else {
107                state.session_order.push_back(oldest);
108                break;
109            }
110        }
111
112        None
113    }
114
115    pub(in crate::tool_output) fn lookup_memo(
116        &self,
117        session_id: &str,
118        tool_call_id: &str,
119        input_hash: u64,
120        input_len: usize,
121        config_fp: u64,
122    ) -> MemoLookup {
123        self.memo.lock().expect("tool output memo lock").lookup(
124            session_id,
125            tool_call_id,
126            input_hash,
127            input_len,
128            config_fp,
129        )
130    }
131
132    pub(in crate::tool_output) fn store_memo(
133        &self,
134        session_id: &str,
135        tool_call_id: &str,
136        record: MemoRecord,
137        max_sessions: usize,
138    ) {
139        self.memo.lock().expect("tool output memo lock").store(
140            session_id,
141            tool_call_id,
142            record,
143            max_sessions,
144        );
145    }
146}
147
148fn hash_output(output: &str) -> u64 {
149    let mut hasher = DefaultHasher::new();
150    output.hash(&mut hasher);
151    hasher.finish()
152}
forge_guardrails/tool_output/state/dedup.rs

forge_guardrails/tool_output/state/
dedup.rs
