From 0f20ee4eb871907defe7848f0d3e2203cfff057e Mon Sep 17 00:00:00 2001
From: Michelle Tilley <michelle@michelletilley.net>
Date: Tue, 21 Apr 2026 10:32:54 -0700
Subject: feat: AI tool rendering overhaul + edit_file tool (#3423)

Overhaul of how AI tool calls are modeled, rendered, and displayed in
the Atuin AI TUI. Fixes bugs in shell command output capture, implements
the `edit_file` tool with full safety infrastructure, and adds a diff
preview for edits.
---
 crates/atuin-ai/src/session.rs | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

(limited to 'crates/atuin-ai/src/session.rs')

diff --git a/crates/atuin-ai/src/session.rs b/crates/atuin-ai/src/session.rs
index d8314343..848330fc 100644
--- a/crates/atuin-ai/src/session.rs
+++ b/crates/atuin-ai/src/session.rs
@@ -51,6 +51,9 @@ pub(crate) trait SessionService: Send + Sync {
     ) -> Result<()>;
 
     async fn archive(&self, session_id: &str) -> Result<()>;
+
+    async fn get_metadata(&self, session_id: &str, key: &str) -> Result<Option<String>>;
+    async fn set_metadata(&self, session_id: &str, key: &str, value: &str) -> Result<()>;
 }
 
 // ---------------------------------------------------------------------------
@@ -128,6 +131,14 @@ impl SessionService for LocalSessionService {
     async fn archive(&self, session_id: &str) -> Result<()> {
         self.store.archive_session(session_id).await
     }
+
+    async fn get_metadata(&self, session_id: &str, key: &str) -> Result<Option<String>> {
+        self.store.get_metadata(session_id, key).await
+    }
+
+    async fn set_metadata(&self, session_id: &str, key: &str, value: &str) -> Result<()> {
+        self.store.set_metadata(session_id, key, value).await
+    }
 }
 
 // ---------------------------------------------------------------------------
@@ -310,6 +321,22 @@ impl SessionManager {
     pub fn invocation_id(&self) -> &str {
         &self.invocation_id
     }
+
+    /// Read a metadata value for the current session.
+    pub async fn get_metadata(&self, key: &str) -> Result<Option<String>> {
+        if !self.persisted_to_db {
+            return Ok(None);
+        }
+        self.service.get_metadata(&self.session_id, key).await
+    }
+
+    /// Write a metadata value for the current session.
+    pub async fn set_metadata(&mut self, key: &str, value: &str) -> Result<()> {
+        self.ensure_persisted().await?;
+        self.service
+            .set_metadata(&self.session_id, key, value)
+            .await
+    }
 }
 
 #[cfg(test)]
-- 
cgit v1.3.1