path: root/Omni/Agent/Tools.hs

{-# 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)