1 files changed, 582 insertions, 0 deletions
diff --git a/Omni/Agent/Tools.hs b/Omni/Agent/Tools.hs
new file mode 100644
index 0000000..0312924
--- /dev/null
+++ b/Omni/Agent/Tools.hs
@@ -0,0 +1,582 @@
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+
+-- | Core coding tools for Jr agent.
+--
+-- Provides implementations for:
+-- - readFile: Read file contents
+-- - writeFile: Write file contents
+-- - editFile: Search/replace edit
+-- - runBash: Execute shell commands
+-- - searchCodebase: Ripgrep wrapper for code search
+--
+-- All tools return structured JSON results.
+--
+-- : out omni-agent-tools
+-- : dep aeson
+-- : dep directory
+module Omni.Agent.Tools
+  ( readFileTool,
+    writeFileTool,
+    editFileTool,
+    runBashTool,
+    searchCodebaseTool,
+    allTools,
+    ReadFileArgs (..),
+    WriteFileArgs (..),
+    EditFileArgs (..),
+    RunBashArgs (..),
+    SearchCodebaseArgs (..),
+    ToolResult (..),
+    main,
+    test,
+  )
+where
+
+import Alpha
+import Data.Aeson ((.!=), (.:), (.:?), (.=))
+import qualified Data.Aeson as Aeson
+import qualified Data.List as List
+import qualified Data.Text as Text
+import qualified Data.Text.IO as TextIO
+import qualified Omni.Agent.Engine as Engine
+import qualified Omni.Test as Test
+import qualified System.Directory as Directory
+import qualified System.Exit as Exit
+import qualified System.Process as Process
+
+main :: IO ()
+main = Test.run test
+
+test :: Test.Tree
+test =
+  Test.group
+    "Omni.Agent.Tools"
+    [ Test.unit "readFileTool schema is valid" <| do
+        let schema = Engine.toolJsonSchema readFileTool
+        case schema of
+          Aeson.Object _ -> pure ()
+          _ -> Test.assertFailure "Schema should be an object",
+      Test.unit "writeFileTool schema is valid" <| do
+        let schema = Engine.toolJsonSchema writeFileTool
+        case schema of
+          Aeson.Object _ -> pure ()
+          _ -> Test.assertFailure "Schema should be an object",
+      Test.unit "editFileTool schema is valid" <| do
+        let schema = Engine.toolJsonSchema editFileTool
+        case schema of
+          Aeson.Object _ -> pure ()
+          _ -> Test.assertFailure "Schema should be an object",
+      Test.unit "runBashTool schema is valid" <| do
+        let schema = Engine.toolJsonSchema runBashTool
+        case schema of
+          Aeson.Object _ -> pure ()
+          _ -> Test.assertFailure "Schema should be an object",
+      Test.unit "searchCodebaseTool schema is valid" <| do
+        let schema = Engine.toolJsonSchema searchCodebaseTool
+        case schema of
+          Aeson.Object _ -> pure ()
+          _ -> Test.assertFailure "Schema should be an object",
+      Test.unit "allTools contains 5 tools" <| do
+        length allTools Test.@=? 5,
+      Test.unit "ReadFileArgs parses correctly" <| do
+        let json = Aeson.object ["path" .= ("test.txt" :: Text)]
+        case Aeson.fromJSON json of
+          Aeson.Success (args :: ReadFileArgs) -> readFilePath args Test.@=? "test.txt"
+          Aeson.Error e -> Test.assertFailure e,
+      Test.unit "WriteFileArgs parses correctly" <| do
+        let json = Aeson.object ["path" .= ("test.txt" :: Text), "content" .= ("hello" :: Text)]
+        case Aeson.fromJSON json of
+          Aeson.Success (args :: WriteFileArgs) -> do
+            writeFilePath args Test.@=? "test.txt"
+            writeFileContent args Test.@=? "hello"
+          Aeson.Error e -> Test.assertFailure e,
+      Test.unit "EditFileArgs parses correctly" <| do
+        let json =
+              Aeson.object
+                [ "path" .= ("test.txt" :: Text),
+                  "old_str" .= ("foo" :: Text),
+                  "new_str" .= ("bar" :: Text)
+                ]
+        case Aeson.fromJSON json of
+          Aeson.Success (args :: EditFileArgs) -> do
+            editFilePath args Test.@=? "test.txt"
+            editFileOldStr args Test.@=? "foo"
+            editFileNewStr args Test.@=? "bar"
+          Aeson.Error e -> Test.assertFailure e,
+      Test.unit "RunBashArgs parses correctly" <| do
+        let json = Aeson.object ["command" .= ("ls -la" :: Text)]
+        case Aeson.fromJSON json of
+          Aeson.Success (args :: RunBashArgs) -> runBashCommand args Test.@=? "ls -la"
+          Aeson.Error e -> Test.assertFailure e,
+      Test.unit "SearchCodebaseArgs parses correctly" <| do
+        let json = Aeson.object ["pattern" .= ("TODO" :: Text)]
+        case Aeson.fromJSON json of
+          Aeson.Success (args :: SearchCodebaseArgs) -> searchPattern args Test.@=? "TODO"
+          Aeson.Error e -> Test.assertFailure e,
+      Test.unit "ToolResult success JSON roundtrip" <| do
+        let result = ToolResult True "done" Nothing
+        case Aeson.decode (Aeson.encode result) of
+          Nothing -> Test.assertFailure "Failed to decode ToolResult"
+          Just decoded -> toolResultSuccess decoded Test.@=? True,
+      Test.unit "ToolResult failure JSON roundtrip" <| do
+        let result = ToolResult False "" (Just "error occurred")
+        case Aeson.decode (Aeson.encode result) of
+          Nothing -> Test.assertFailure "Failed to decode ToolResult"
+          Just decoded -> toolResultError decoded Test.@=? Just "error occurred",
+      Test.unit "readFileTool handles missing files" <| do
+        let args = Aeson.object ["path" .= ("/nonexistent/path/to/file.txt" :: Text)]
+        result <- Engine.toolExecute readFileTool 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 "editFileTool handles no-match case" <| do
+        let args =
+              Aeson.object
+                [ "path" .= ("/nonexistent/file.txt" :: Text),
+                  "old_str" .= ("needle" :: Text),
+                  "new_str" .= ("replacement" :: Text)
+                ]
+        result <- Engine.toolExecute editFileTool args
+        case Aeson.fromJSON result of
+          Aeson.Success (tr :: ToolResult) -> toolResultSuccess tr Test.@=? False
+          Aeson.Error e -> Test.assertFailure e,
+      Test.unit "runBashTool captures exit codes" <| do
+        let args = Aeson.object ["command" .= ("exit 42" :: Text)]
+        result <- Engine.toolExecute runBashTool args
+        case Aeson.fromJSON result of
+          Aeson.Success (tr :: ToolResult) -> do
+            toolResultSuccess tr Test.@=? False
+            toolResultError tr Test.@=? Just "Exit code: 42"
+          Aeson.Error e -> Test.assertFailure e,
+      Test.unit "runBashTool captures stdout" <| do
+        let args = Aeson.object ["command" .= ("echo hello" :: Text)]
+        result <- Engine.toolExecute runBashTool args
+        case Aeson.fromJSON result of
+          Aeson.Success (tr :: ToolResult) -> do
+            toolResultSuccess tr Test.@=? True
+            ("hello" `Text.isInfixOf` toolResultOutput tr) Test.@=? True
+          Aeson.Error e -> Test.assertFailure e,
+      Test.unit "searchCodebaseTool returns structured results" <| do
+        let args =
+              Aeson.object
+                [ "pattern" .= ("module" :: Text),
+                  "path" .= ("." :: Text),
+                  "max_results" .= (5 :: Int)
+                ]
+        result <- Engine.toolExecute searchCodebaseTool args
+        case Aeson.fromJSON result of
+          Aeson.Success (tr :: ToolResult) -> toolResultSuccess tr Test.@=? True
+          Aeson.Error e -> Test.assertFailure e
+    ]
+
+data ToolResult = ToolResult
+  { toolResultSuccess :: Bool,
+    toolResultOutput :: Text,
+    toolResultError :: Maybe Text
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON ToolResult where
+  toJSON r =
+    Aeson.object
+      <| catMaybes
+        [ Just ("success" .= toolResultSuccess r),
+          Just ("output" .= toolResultOutput r),
+          ("error" .=) </ toolResultError r
+        ]
+
+instance Aeson.FromJSON ToolResult where
+  parseJSON =
+    Aeson.withObject "ToolResult" <| \v ->
+      (ToolResult </ (v .: "success"))
+        <*> (v .:? "output" .!= "")
+        <*> (v .:? "error")
+
+mkSuccess :: Text -> Aeson.Value
+mkSuccess output = Aeson.toJSON <| ToolResult True output Nothing
+
+mkError :: Text -> Aeson.Value
+mkError err = Aeson.toJSON <| ToolResult False "" (Just err)
+
+allTools :: [Engine.Tool]
+allTools =
+  [ readFileTool,
+    writeFileTool,
+    editFileTool,
+    runBashTool,
+    searchCodebaseTool
+  ]
+
+data ReadFileArgs = ReadFileArgs
+  { readFilePath :: Text,
+    readFileStartLine :: Maybe Int,
+    readFileEndLine :: Maybe Int
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON ReadFileArgs where
+  parseJSON =
+    Aeson.withObject "ReadFileArgs" <| \v ->
+      (ReadFileArgs </ (v .: "path"))
+        <*> (v .:? "start_line")
+        <*> (v .:? "end_line")
+
+readFileTool :: Engine.Tool
+readFileTool =
+  Engine.Tool
+    { Engine.toolName = "read_file",
+      Engine.toolDescription = "Read the contents of a file. Can optionally read a specific line range.",
+      Engine.toolJsonSchema =
+        Aeson.object
+          [ "type" .= ("object" :: Text),
+            "properties"
+              .= Aeson.object
+                [ "path"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("Absolute path to the file to read" :: Text)
+                      ],
+                  "start_line"
+                    .= Aeson.object
+                      [ "type" .= ("integer" :: Text),
+                        "description" .= ("Optional: first line to read (1-indexed)" :: Text)
+                      ],
+                  "end_line"
+                    .= Aeson.object
+                      [ "type" .= ("integer" :: Text),
+                        "description" .= ("Optional: last line to read (1-indexed)" :: Text)
+                      ]
+                ],
+            "required" .= (["path"] :: [Text])
+          ],
+      Engine.toolExecute = executeReadFile
+    }
+
+executeReadFile :: Aeson.Value -> IO Aeson.Value
+executeReadFile v =
+  case Aeson.fromJSON v of
+    Aeson.Error e -> pure <| mkError (Text.pack e)
+    Aeson.Success args -> do
+      let path = Text.unpack (readFilePath args)
+      exists <- Directory.doesFileExist path
+      if exists
+        then do
+          content <- TextIO.readFile path
+          let allLines = Text.lines content
+              startIdx = maybe 0 (\n -> n - 1) (readFileStartLine args)
+              endIdx = maybe (length allLines) identity (readFileEndLine args)
+              selectedLines = take (endIdx - startIdx) (drop startIdx allLines)
+              numberedLines = zipWith formatLine [(startIdx + 1) ..] selectedLines
+              result = Text.unlines numberedLines
+          pure <| mkSuccess result
+        else pure <| mkError ("File not found: " <> readFilePath args)
+  where
+    formatLine :: Int -> Text -> Text
+    formatLine n line = Text.pack (show n) <> ": " <> line
+
+data WriteFileArgs = WriteFileArgs
+  { writeFilePath :: Text,
+    writeFileContent :: Text
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON WriteFileArgs where
+  parseJSON =
+    Aeson.withObject "WriteFileArgs" <| \v ->
+      (WriteFileArgs </ (v .: "path"))
+        <*> (v .: "content")
+
+writeFileTool :: Engine.Tool
+writeFileTool =
+  Engine.Tool
+    { Engine.toolName = "write_file",
+      Engine.toolDescription = "Create or overwrite a file with the given content.",
+      Engine.toolJsonSchema =
+        Aeson.object
+          [ "type" .= ("object" :: Text),
+            "properties"
+              .= Aeson.object
+                [ "path"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("Absolute path to the file to write" :: Text)
+                      ],
+                  "content"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("Content to write to the file" :: Text)
+                      ]
+                ],
+            "required" .= (["path", "content"] :: [Text])
+          ],
+      Engine.toolExecute = executeWriteFile
+    }
+
+executeWriteFile :: Aeson.Value -> IO Aeson.Value
+executeWriteFile v =
+  case Aeson.fromJSON v of
+    Aeson.Error e -> pure <| mkError (Text.pack e)
+    Aeson.Success args -> do
+      let path = Text.unpack (writeFilePath args)
+      let dir = takeDirectory path
+      dirExists <- Directory.doesDirectoryExist dir
+      if dirExists
+        then do
+          TextIO.writeFile path (writeFileContent args)
+          pure <| mkSuccess ("File written: " <> writeFilePath args)
+        else pure <| mkError ("Parent directory does not exist: " <> Text.pack dir)
+  where
+    takeDirectory :: FilePath -> FilePath
+    takeDirectory p =
+      let parts = Text.splitOn "/" (Text.pack p)
+       in Text.unpack (Text.intercalate "/" (List.init parts))
+
+data EditFileArgs = EditFileArgs
+  { editFilePath :: Text,
+    editFileOldStr :: Text,
+    editFileNewStr :: Text,
+    editFileReplaceAll :: Maybe Bool
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON EditFileArgs where
+  parseJSON =
+    Aeson.withObject "EditFileArgs" <| \v ->
+      (EditFileArgs </ (v .: "path"))
+        <*> (v .: "old_str")
+        <*> (v .: "new_str")
+        <*> (v .:? "replace_all")
+
+editFileTool :: Engine.Tool
+editFileTool =
+  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 = executeEditFile
+    }
+
+executeEditFile :: Aeson.Value -> IO Aeson.Value
+executeEditFile 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
+              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)
+
+replaceFirst :: Text -> Text -> Text -> Text
+replaceFirst old new content =
+  case Text.breakOn old content of
+    (before, after) ->
+      if Text.null after
+        then content
+        else before <> new <> Text.drop (Text.length old) after
+
+data RunBashArgs = RunBashArgs
+  { runBashCommand :: Text,
+    runBashCwd :: Maybe Text,
+    runBashTimeout :: Maybe Int
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON RunBashArgs where
+  parseJSON =
+    Aeson.withObject "RunBashArgs" <| \v ->
+      (RunBashArgs </ (v .: "command"))
+        <*> (v .:? "cwd")
+        <*> (v .:? "timeout")
+
+runBashTool :: Engine.Tool
+runBashTool =
+  Engine.Tool
+    { Engine.toolName = "run_bash",
+      Engine.toolDescription = "Execute a shell command and return stdout/stderr.",
+      Engine.toolJsonSchema =
+        Aeson.object
+          [ "type" .= ("object" :: Text),
+            "properties"
+              .= Aeson.object
+                [ "command"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("The shell command to execute" :: Text)
+                      ],
+                  "cwd"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("Optional: working directory for the command" :: Text)
+                      ],
+                  "timeout"
+                    .= Aeson.object
+                      [ "type" .= ("integer" :: Text),
+                        "description" .= ("Optional: timeout in seconds (default: 300)" :: Text)
+                      ]
+                ],
+            "required" .= (["command"] :: [Text])
+          ],
+      Engine.toolExecute = executeRunBash
+    }
+
+executeRunBash :: Aeson.Value -> IO Aeson.Value
+executeRunBash v =
+  case Aeson.fromJSON v of
+    Aeson.Error e -> pure <| mkError (Text.pack e)
+    Aeson.Success args -> do
+      let cmd = Text.unpack (runBashCommand args)
+          proc =
+            (Process.shell cmd)
+              { Process.cwd = Text.unpack </ runBashCwd args
+              }
+      (exitCode, stdoutStr, stderrStr) <- Process.readCreateProcessWithExitCode proc ""
+      let output = Text.pack stdoutStr <> Text.pack stderrStr
+      case exitCode of
+        Exit.ExitSuccess ->
+          pure
+            <| Aeson.toJSON
+            <| ToolResult
+              { toolResultSuccess = True,
+                toolResultOutput = output,
+                toolResultError = Nothing
+              }
+        Exit.ExitFailure code ->
+          pure
+            <| Aeson.toJSON
+            <| ToolResult
+              { toolResultSuccess = False,
+                toolResultOutput = output,
+                toolResultError = Just ("Exit code: " <> tshow code)
+              }
+
+data SearchCodebaseArgs = SearchCodebaseArgs
+  { searchPattern :: Text,
+    searchPath :: Maybe Text,
+    searchGlob :: Maybe Text,
+    searchCaseSensitive :: Maybe Bool,
+    searchMaxResults :: Maybe Int
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON SearchCodebaseArgs where
+  parseJSON =
+    Aeson.withObject "SearchCodebaseArgs" <| \v ->
+      (SearchCodebaseArgs </ (v .: "pattern"))
+        <*> (v .:? "path")
+        <*> (v .:? "glob")
+        <*> (v .:? "case_sensitive")
+        <*> (v .:? "max_results")
+
+searchCodebaseTool :: Engine.Tool
+searchCodebaseTool =
+  Engine.Tool
+    { Engine.toolName = "search_codebase",
+      Engine.toolDescription = "Search the codebase using ripgrep. Returns matching lines with file paths and line numbers.",
+      Engine.toolJsonSchema =
+        Aeson.object
+          [ "type" .= ("object" :: Text),
+            "properties"
+              .= Aeson.object
+                [ "pattern"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("The regex pattern to search for" :: Text)
+                      ],
+                  "path"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("Optional: directory or file path to search in" :: Text)
+                      ],
+                  "glob"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("Optional: file glob pattern (e.g., '*.hs')" :: Text)
+                      ],
+                  "case_sensitive"
+                    .= Aeson.object
+                      [ "type" .= ("boolean" :: Text),
+                        "description" .= ("Optional: case sensitive search (default: false)" :: Text)
+                      ],
+                  "max_results"
+                    .= Aeson.object
+                      [ "type" .= ("integer" :: Text),
+                        "description" .= ("Optional: maximum number of results (default: 100)" :: Text)
+                      ]
+                ],
+            "required" .= (["pattern"] :: [Text])
+          ],
+      Engine.toolExecute = executeSearchCodebase
+    }
+
+executeSearchCodebase :: Aeson.Value -> IO Aeson.Value
+executeSearchCodebase v =
+  case Aeson.fromJSON v of
+    Aeson.Error e -> pure <| mkError (Text.pack e)
+    Aeson.Success args -> do
+      let pat = Text.unpack (searchPattern args)
+          maxRes = fromMaybe 100 (searchMaxResults args)
+          caseSensitive = fromMaybe False (searchCaseSensitive args)
+          baseArgs =
+            ["--line-number", "--no-heading", "--max-count=" <> show maxRes, pat]
+          caseArgs = ["--ignore-case" | not caseSensitive]
+          globArgs = maybe [] (\g -> ["--glob", Text.unpack g]) (searchGlob args)
+          pathArg = maybe ["."] (\p -> [Text.unpack p]) (searchPath args)
+          allArgs = caseArgs <> globArgs <> baseArgs <> pathArg
+          proc = Process.proc "rg" allArgs
+      (exitCode, stdoutStr, stderrStr) <- Process.readCreateProcessWithExitCode proc ""
+      case exitCode of
+        Exit.ExitSuccess ->
+          pure <| mkSuccess (Text.pack stdoutStr)
+        Exit.ExitFailure 1 ->
+          pure <| mkSuccess "No matches found"
+        Exit.ExitFailure code ->
+          pure <| mkError ("ripgrep failed with code " <> tshow code <> ": " <> Text.pack stderrStr)