gemini_adk_rs/tool/
commit.rs

1//! Commit governance for a tool that acts on the world: idempotency and
2//! compensation.
3//!
4//! A voice agent that charges a card or books a table can call the same
5//! tool twice for one intent. The model retries after a timeout, the user
6//! repeats themselves, a resumed session replays. [`CommitGuard`] wraps such
7//! a tool:
8//!
9//! - **Idempotency.** An idempotency key is rendered from a template
10//!   (`"{user_id}:{amount}"`). A call whose key already succeeded returns the
11//!   first call's result without running the tool again. The result is
12//!   remembered in state under [`idempotency_key`]`(tool, key)`. Calls with
13//!   the same key run one at a time, so two concurrent retries (a duplicate
14//!   background call, say) still charge once.
15//! - **Compensation.** When the tool fails, a compensating tool (a refund, a
16//!   release) runs with the same arguments to undo any partial effect, and
17//!   `compensated:{tool}` is set in state. The failure is still reported to
18//!   the model.
19//!
20//! `Policy::commit(..)` in the fluent crate lowers to this.
21
22use std::collections::HashMap;
23use std::sync::Arc;
24
25use async_trait::async_trait;
26use serde_json::Value;
27
28use super::ToolFunction;
29use crate::error::ToolError;
30use crate::state::State;
31
32/// The state key that remembers a committed call's result.
33pub fn idempotency_key(tool: &str, key: &str) -> String {
34    format!("idempotency:{tool}:{key}")
35}
36
37/// The state key raised after a failed commit was compensated.
38pub fn compensated_key(tool: &str) -> String {
39    format!("compensated:{tool}")
40}
41
42/// A tool wrapped with idempotency and compensation. See the
43/// [module docs](self).
44pub struct CommitGuard {
45    inner: Arc<dyn ToolFunction>,
46    state: State,
47    key_template: Option<String>,
48    compensate: Option<Arc<dyn ToolFunction>>,
49    /// One lock per rendered key, held from the idempotency check to the
50    /// stored result.
51    in_flight: parking_lot::Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
52}
53
54impl CommitGuard {
55    /// Guard `inner`, remembering results in `state`.
56    pub fn new(inner: Arc<dyn ToolFunction>, state: State) -> Self {
57        Self {
58            inner,
59            state,
60            key_template: None,
61            compensate: None,
62            in_flight: parking_lot::Mutex::new(HashMap::new()),
63        }
64    }
65
66    /// Deduplicate calls by this key template. `{name}` takes the call's
67    /// argument `name`, else the state value at `name`. A call where any
68    /// placeholder has no value is not deduplicated (and says so in the
69    /// log), so an incomplete key never merges two different commits.
70    pub fn idempotency_key(mut self, template: impl Into<String>) -> Self {
71        self.key_template = Some(template.into());
72        self
73    }
74
75    /// Run `tool` with the same arguments when the guarded tool fails.
76    pub fn compensate_with(mut self, tool: Arc<dyn ToolFunction>) -> Self {
77        self.compensate = Some(tool);
78        self
79    }
80
81    fn render_key(&self, args: &Value) -> Option<String> {
82        let template = self.key_template.as_deref()?;
83        let mut out = String::with_capacity(template.len());
84        let mut rest = template;
85        while let Some(open) = rest.find('{') {
86            out.push_str(&rest[..open]);
87            let after = &rest[open + 1..];
88            let close = after.find('}')?;
89            let name = after[..close].trim();
90            let value = args
91                .get(name)
92                .cloned()
93                .or_else(|| self.state.get_raw(name))
94                .filter(|v| !v.is_null());
95            match value {
96                Some(Value::String(s)) => out.push_str(&s),
97                Some(v) => out.push_str(&v.to_string()),
98                None => {
99                    tracing::warn!(
100                        tool = self.inner.name(),
101                        placeholder = name,
102                        "idempotency key incomplete; this call is not deduplicated"
103                    );
104                    return None;
105                }
106            }
107            rest = &after[close + 1..];
108        }
109        out.push_str(rest);
110        Some(out)
111    }
112}
113
114#[async_trait]
115impl ToolFunction for CommitGuard {
116    fn name(&self) -> &str {
117        self.inner.name()
118    }
119
120    fn description(&self) -> &str {
121        self.inner.description()
122    }
123
124    fn parameters(&self) -> Option<Value> {
125        self.inner.parameters()
126    }
127
128    fn requires_confirmation(&self) -> bool {
129        self.inner.requires_confirmation()
130    }
131
132    fn confirmation_message(&self) -> Option<&str> {
133        self.inner.confirmation_message()
134    }
135
136    async fn call(&self, args: Value) -> Result<Value, ToolError> {
137        self.call_with_context(args, super::ToolContext::new(self.state.clone()))
138            .await
139    }
140
141    async fn call_with_context(
142        &self,
143        args: Value,
144        ctx: super::ToolContext,
145    ) -> Result<Value, ToolError> {
146        let tool = self.inner.name();
147        let key = self
148            .render_key(&args)
149            .map(|key| idempotency_key(tool, &key));
150        // Serialize calls with the same key: the second waits for the first
151        // to store its result, then returns it.
152        let lock = key.as_ref().map(|key| {
153            self.in_flight
154                .lock()
155                .entry(key.clone())
156                .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
157                .clone()
158        });
159        let _held = match &lock {
160            Some(lock) => Some(lock.lock().await),
161            None => None,
162        };
163        if let Some(key) = &key
164            && let Some(previous) = self.state.get_raw(key)
165        {
166            tracing::info!(tool, "commit already made; returning its result");
167            return Ok(previous);
168        }
169        match self
170            .inner
171            .call_with_context(args.clone(), ctx.clone())
172            .await
173        {
174            Ok(result) => {
175                if let Some(key) = key {
176                    let _ = self.state.set(key, &result);
177                }
178                Ok(result)
179            }
180            Err(error) => {
181                if let Some(compensate) = &self.compensate {
182                    match compensate.call_with_context(args, ctx).await {
183                        Ok(_) => {
184                            let _ = self.state.set(compensated_key(tool), true);
185                        }
186                        Err(e) => tracing::error!(
187                            tool,
188                            compensating = compensate.name(),
189                            "compensation failed: {e}"
190                        ),
191                    }
192                }
193                Err(error)
194            }
195        }
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use crate::tool::SimpleTool;
203    use serde_json::json;
204    use std::sync::atomic::{AtomicUsize, Ordering};
205
206    fn counting(name: &'static str, calls: Arc<AtomicUsize>, fail: bool) -> Arc<dyn ToolFunction> {
207        Arc::new(SimpleTool::new(name, name, None, move |args| {
208            let calls = calls.clone();
209            async move {
210                let n = calls.fetch_add(1, Ordering::SeqCst) + 1;
211                if fail {
212                    Err(ToolError::ExecutionFailed("card declined".into()))
213                } else {
214                    Ok(json!({ "charge": n, "amount": args["amount"] }))
215                }
216            }
217        }))
218    }
219
220    #[tokio::test]
221    async fn the_same_commit_runs_once() {
222        let calls = Arc::new(AtomicUsize::new(0));
223        let state = State::new();
224        let _ = state.set("user_id", "u1");
225        let guard = CommitGuard::new(counting("charge", calls.clone(), false), state.clone())
226            .idempotency_key("{user_id}:{amount}");
227
228        let first = guard.call(json!({ "amount": 40 })).await.unwrap();
229        let again = guard.call(json!({ "amount": 40 })).await.unwrap();
230        assert_eq!(first, again, "the retry gets the first charge back");
231        assert_eq!(calls.load(Ordering::SeqCst), 1, "the card was charged once");
232
233        guard.call(json!({ "amount": 55 })).await.unwrap();
234        assert_eq!(
235            calls.load(Ordering::SeqCst),
236            2,
237            "a different amount is a new commit"
238        );
239        assert!(state.get_raw(&idempotency_key("charge", "u1:40")).is_some());
240    }
241
242    #[tokio::test]
243    async fn concurrent_calls_with_one_key_commit_once() {
244        let calls = Arc::new(AtomicUsize::new(0));
245        let counter = calls.clone();
246        let slow: Arc<dyn ToolFunction> =
247            Arc::new(SimpleTool::new("charge", "charge", None, move |_| {
248                let counter = counter.clone();
249                async move {
250                    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
251                    Ok(json!({ "charge": counter.fetch_add(1, Ordering::SeqCst) + 1 }))
252                }
253            }));
254        let guard = Arc::new(CommitGuard::new(slow, State::new()).idempotency_key("{amount}"));
255        let (a, b) = tokio::join!(
256            guard.call(json!({ "amount": 40 })),
257            guard.call(json!({ "amount": 40 }))
258        );
259        assert_eq!(a.unwrap(), b.unwrap(), "both get the one charge");
260        assert_eq!(calls.load(Ordering::SeqCst), 1);
261    }
262
263    #[tokio::test]
264    async fn an_incomplete_key_never_merges_commits() {
265        let calls = Arc::new(AtomicUsize::new(0));
266        let guard = CommitGuard::new(counting("charge", calls.clone(), false), State::new())
267            .idempotency_key("{user_id}:{amount}");
268        guard.call(json!({ "amount": 40 })).await.unwrap();
269        guard.call(json!({ "amount": 40 })).await.unwrap();
270        assert_eq!(calls.load(Ordering::SeqCst), 2);
271    }
272
273    #[tokio::test]
274    async fn a_failed_commit_is_compensated_and_still_fails() {
275        let charges = Arc::new(AtomicUsize::new(0));
276        let refunds = Arc::new(AtomicUsize::new(0));
277        let state = State::new();
278        let guard = CommitGuard::new(counting("charge", charges, true), state.clone())
279            .compensate_with(counting("refund", refunds.clone(), false));
280        assert!(guard.call(json!({ "amount": 40 })).await.is_err());
281        assert_eq!(refunds.load(Ordering::SeqCst), 1);
282        assert_eq!(state.get::<bool>(&compensated_key("charge")), Some(true));
283    }
284}
gemini_adk_rs/tool/commit.rs

gemini_adk_rs/tool/
commit.rs
