commit a6dbe64442f99362eec2d26f0605e97352abc46a
Author: Ben Sima <ben@bensima.com>
Date: Wed Dec 31 00:22:14 2025
Omni/Agent/Tools.hs: Implement undo_edit tool for reverting file edits
Automated via pi-review.
Task-Id: t-242
diff --git a/Omni/Agent/Tools.hs b/Omni/Agent/Tools.hs
index ae822443..ea765968 100644
--- a/Omni/Agent/Tools.hs
+++ b/Omni/Agent/Tools.hs
@@ -20,14 +20,19 @@ module Omni.Agent.Tools
( readFileTool,
writeFileTool,
editFileTool,
+ editFileToolWithHistory,
+ undoEditTool,
runBashTool,
searchCodebaseTool,
searchAndReadTool,
globTool,
allTools,
+ allToolsWithHistory,
+ EditHistory,
ReadFileArgs (..),
WriteFileArgs (..),
EditFileArgs (..),
+ UndoEditArgs (..),
RunBashArgs (..),
SearchCodebaseArgs (..),
SearchAndReadArgs (..),
@@ -41,7 +46,9 @@ where
import Alpha
import Data.Aeson ((.!=), (.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
+import Data.IORef (IORef, modifyIORef', newIORef, readIORef)
import qualified Data.List as List
+import qualified Data.Map.Strict as Map
import qualified Data.Text as Text
import qualified Data.Text.IO as TextIO
import qualified Omni.Agent.Engine as Engine
@@ -50,6 +57,9 @@ import qualified System.Directory as Directory
import qualified System.Exit as Exit
import qualified System.Process as Process
+-- | Edit history maps file paths to their previous content (before last edit)
+type EditHistory = IORef (Map.Map FilePath Text)
+
main :: IO ()
main = Test.run test
@@ -82,8 +92,8 @@ test =
case schema of
Aeson.Object _ -> pure ()
_ -> Test.assertFailure "Schema should be an object",
- Test.unit "allTools contains 7 tools" <| do
- length allTools Test.@=? 7,
+ Test.unit "allTools contains 8 tools" <| do
+ length allTools Test.@=? 8,
Test.unit "ReadFileArgs parses correctly" <| do
let json = Aeson.object ["path" .= ("test.txt" :: Text)]
case Aeson.fromJSON json of
@@ -210,7 +220,31 @@ test =
toolResultSuccess tr Test.@=? True
-- Should find at least Tools.hs
not (Text.null (toolResultOutput tr)) Test.@=? True
- Aeson.Error e -> Test.assertFailure e
+ Aeson.Error e -> Test.assertFailure e,
+ Test.unit "undoEditTool schema is valid" <| do
+ historyRef <- Data.IORef.newIORef Map.empty
+ let schema = Engine.toolJsonSchema (undoEditTool historyRef)
+ case schema of
+ Aeson.Object _ -> pure ()
+ _ -> Test.assertFailure "Schema should be an object",
+ Test.unit "UndoEditArgs parses correctly" <| do
+ let json = Aeson.object ["path" .= ("/tmp/test.txt" :: Text)]
+ case Aeson.fromJSON json of
+ Aeson.Success (args :: UndoEditArgs) -> undoPath args Test.@=? "/tmp/test.txt"
+ Aeson.Error e -> Test.assertFailure e,
+ Test.unit "undoEditTool fails when no history" <| do
+ historyRef <- Data.IORef.newIORef Map.empty
+ let tool = undoEditTool historyRef
+ args = Aeson.object ["path" .= ("/tmp/nonexistent.txt" :: Text)]
+ result <- Engine.toolExecute tool args
+ case Aeson.fromJSON result of
+ Aeson.Success (tr :: ToolResult) -> do
+ toolResultSuccess tr Test.@=? False
+ isJust (toolResultError tr) Test.@=? True
+ Aeson.Error e -> Test.assertFailure e,
+ Test.unit "allToolsWithHistory contains 8 tools" <| do
+ historyRef <- Data.IORef.newIORef Map.empty
+ length (allToolsWithHistory historyRef) Test.@=? 8
]
data ToolResult = ToolResult
@@ -269,12 +303,36 @@ allTools =
[ readFileTool,
writeFileTool,
editFileTool,
+ undoEditToolStub,
runBashTool,
searchCodebaseTool,
searchAndReadTool,
globTool
]
+-- | Stub version of undo_edit that returns an error (for use without history tracking)
+undoEditToolStub :: Engine.Tool
+undoEditToolStub =
+ Engine.Tool
+ { Engine.toolName = "undo_edit",
+ Engine.toolDescription = "Undo the last edit made to a file. Only works for files edited in this session.",
+ Engine.toolJsonSchema =
+ Aeson.object
+ [ "type" .= ("object" :: Text),
+ "properties"
+ .= Aeson.object
+ [ "path"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("Absolute path to the file to undo" :: Text)
+ ]
+ ],
+ "required" .= (["path"] :: [Text])
+ ],
+ Engine.toolExecute = \_ ->
+ pure <| mkError "undo_edit requires edit history tracking. Use allToolsWithHistory instead of allTools."
+ }
+
data ReadFileArgs = ReadFileArgs
{ readFilePath :: Text,
readFileStartLine :: Maybe Int,
@@ -513,6 +571,135 @@ replaceFirst old new content =
then content
else before <> new <> Text.drop (Text.length old) after
+-- | Edit file tool that tracks history for undo support
+editFileToolWithHistory :: EditHistory -> Engine.Tool
+editFileToolWithHistory historyRef =
+ Engine.Tool
+ { Engine.toolName = "edit_file",
+ Engine.toolDescription = "Edit a file by replacing old_str with new_str. By default replaces only the first occurrence unless replace_all is true.",
+ Engine.toolJsonSchema =
+ Aeson.object
+ [ "type" .= ("object" :: Text),
+ "properties"
+ .= Aeson.object
+ [ "path"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("Absolute path to the file to edit" :: Text)
+ ],
+ "old_str"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("The text to search for and replace" :: Text)
+ ],
+ "new_str"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("The replacement text" :: Text)
+ ],
+ "replace_all"
+ .= Aeson.object
+ [ "type" .= ("boolean" :: Text),
+ "description" .= ("If true, replace all occurrences; otherwise replace only the first" :: Text)
+ ]
+ ],
+ "required" .= (["path", "old_str", "new_str"] :: [Text])
+ ],
+ Engine.toolExecute = executeEditFileWithHistory historyRef
+ }
+
+executeEditFileWithHistory :: EditHistory -> Aeson.Value -> IO Aeson.Value
+executeEditFileWithHistory historyRef v =
+ case Aeson.fromJSON v of
+ Aeson.Error e -> pure <| mkError (Text.pack e)
+ Aeson.Success args -> do
+ let path = Text.unpack (editFilePath args)
+ exists <- Directory.doesFileExist path
+ if exists
+ then do
+ content <- TextIO.readFile path
+ let oldStr = editFileOldStr args
+ newStr = editFileNewStr args
+ replaceAll = fromMaybe False (editFileReplaceAll args)
+ if Text.isInfixOf oldStr content
+ then do
+ -- Save original content to history before editing
+ modifyIORef' historyRef (Map.insert path content)
+ let newContent =
+ if replaceAll
+ then Text.replace oldStr newStr content
+ else replaceFirst oldStr newStr content
+ TextIO.writeFile path newContent
+ let count =
+ if replaceAll
+ then Text.count oldStr content
+ else 1
+ pure <| mkSuccess ("Replaced " <> tshow count <> " occurrence(s)")
+ else pure <| mkError ("old_str not found in file: " <> editFilePath args)
+ else pure <| mkError ("File not found: " <> editFilePath args)
+
+newtype UndoEditArgs = UndoEditArgs
+ { undoPath :: Text
+ }
+ deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON UndoEditArgs where
+ parseJSON =
+ Aeson.withObject "UndoEditArgs" <| \v ->
+ UndoEditArgs </ (v .: "path")
+
+-- | Tool to undo the last edit made to a file
+undoEditTool :: EditHistory -> Engine.Tool
+undoEditTool historyRef =
+ Engine.Tool
+ { Engine.toolName = "undo_edit",
+ Engine.toolDescription = "Undo the last edit made to a file. Only works for files edited in this session.",
+ Engine.toolJsonSchema =
+ Aeson.object
+ [ "type" .= ("object" :: Text),
+ "properties"
+ .= Aeson.object
+ [ "path"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("Absolute path to the file to undo" :: Text)
+ ]
+ ],
+ "required" .= (["path"] :: [Text])
+ ],
+ Engine.toolExecute = executeUndoEdit historyRef
+ }
+
+executeUndoEdit :: EditHistory -> Aeson.Value -> IO Aeson.Value
+executeUndoEdit historyRef v =
+ case Aeson.fromJSON v of
+ Aeson.Error e -> pure <| mkError (Text.pack e)
+ Aeson.Success args -> do
+ let path = Text.unpack (undoPath args)
+ history <- readIORef historyRef
+ case Map.lookup path history of
+ Nothing ->
+ pure <| mkError ("No edit history for file: " <> undoPath args)
+ Just previousContent -> do
+ -- Write previous content back to file
+ TextIO.writeFile path previousContent
+ -- Remove entry from history (cannot undo twice)
+ modifyIORef' historyRef (Map.delete path)
+ pure <| mkSuccess ("Reverted file to previous state: " <> undoPath args)
+
+-- | All tools with edit history support for undo functionality
+allToolsWithHistory :: EditHistory -> [Engine.Tool]
+allToolsWithHistory historyRef =
+ [ readFileTool,
+ writeFileTool,
+ editFileToolWithHistory historyRef,
+ undoEditTool historyRef,
+ runBashTool,
+ searchCodebaseTool,
+ searchAndReadTool,
+ globTool
+ ]
+
data RunBashArgs = RunBashArgs
{ runBashCommand :: Text,
runBashCwd :: Maybe Text,
@@ -767,7 +954,7 @@ globTool :: Engine.Tool
globTool =
Engine.Tool
{ Engine.toolName = "glob",
- Engine.toolDescription = "Find files matching a glob pattern. Respects .gitignore.",
+ Engine.toolDescription = "Find files matching a glob pattern. Returns matching paths. Respects .gitignore.",
Engine.toolJsonSchema =
Aeson.object
[ "type" .= ("object" :: Text),