Grep.hs

module LLM.Tools.Grep
  ( grepToolTyped,
    GrepToolArgs (..),
  )
where

import Autodocodec qualified as AC
import Control.Exception (IOException, try)
import Control.Monad (unless)
import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef)
import Data.Aeson (FromJSON)
import Data.List (sort)
import Data.Text (Text)
import Data.ByteString qualified as BS
import Data.Maybe (catMaybes)
import Data.Text qualified as T
import Data.Text.Encoding (decodeUtf8With)
import Data.Text.Encoding.Error (lenientDecode)
import GHC.Generics (Generic)
import LLM.Core.Types (TypedTool (..))
import LLM.Tools.FsConfig (FsConfig, isFileHidden, isSymlink, sandboxPath, withWorkspacePaths)
import LLM.Tools.FsLimits (expensiveWalkDepth, maxOutputBytes, maxSearchBytes, maxVisitedEntries)
import LLM.Tools.Internal.TextFile (withTextHandle)
import LLM.Tools.Internal.WalkBudget (WalkBudget, budgetExhausted, budgetNote, capLines, newWalkBudget, takeEntries)
import System.Directory
  ( doesDirectoryExist,
    doesFileExist,
    getFileSize,
    listDirectory,
  )
import System.FilePath (takeExtension, takeFileName, (</>))
import System.IO
  ( IOMode (ReadMode),
    hFileSize,
    withBinaryFile,
  )

-- | Default cap on the number of match lines returned.
defaultMaxResults :: Int
defaultMaxResults = 50

-- | Hard cap on results to protect the model's context window.
maxResultsCap :: Int
maxResultsCap = 500

-- | Skip files larger than this (bytes). Big files are usually generated
-- artifacts (lock files, fixtures) that aren't useful to grep through.
maxFileSizeBytes :: Integer
maxFileSizeBytes = 1024 * 1024 -- 1 MB

-- | Maximum directory depth from the search root. Sourced from
-- 'LLM.Tools.FsLimits.expensiveWalkDepth' because @grep@ reads every
-- encountered file.
maxDepth :: Int
maxDepth = expensiveWalkDepth

-- | Truncate matched line text in output to keep results compact.
maxLineLength :: Int
maxLineLength = 240

data GrepToolArgs = GrepToolArgs
  { _grepPattern :: Text,
    _grepPath :: Text,
    _grepCaseInsensitive :: Bool,
    _grepExtensions :: [Text],
    _grepMaxResults :: Int
  }
  deriving (Generic)
  deriving (FromJSON) via (AC.Autodocodec GrepToolArgs)

instance AC.HasCodec GrepToolArgs where
  codec :: AC.JSONCodec GrepToolArgs
  codec =
    AC.object "search files for a literal substring" $
      GrepToolArgs
        <$> AC.requiredField "pattern" "Literal substring to search for (not a regex)"
          AC..= (._grepPattern)
        <*> AC.optionalFieldWithDefault
          "path"
          "."
          "Relative directory or file to search under (default '.')"
          AC..= (._grepPath)
        <*> AC.optionalFieldWithDefault
          "caseInsensitive"
          False
          "Match case-insensitively (default false)"
          AC..= (._grepCaseInsensitive)
        <*> AC.optionalFieldWithDefault
          "extensions"
          []
          "Restrict to files with these extensions, without the dot, e.g. [\"hs\", \"cabal\"]. Empty means all text files."
          AC..= (._grepExtensions)
        <*> AC.optionalFieldWithDefault
          "maxResults"
          defaultMaxResults
          "Maximum number of match lines to return (default 50, hard cap 500)"
          AC..= (._grepMaxResults)

grepToolTyped :: FsConfig -> TypedTool ctx GrepToolArgs
grepToolTyped cfg =
  TypedTool
    { ttoolName = "grep",
      ttoolDescription =
        "Recursively search files for a literal substring (not a regex). "
          <> "Returns matches as 'relative/path:line: snippet', sorted by file. "
          <> "Skips hidden files, binary files, and files larger than 1 MB. "
          <> "Use 'extensions' to restrict by file type (e.g. [\"hs\"]). "
          <> "Output is capped at 'maxResults' lines (default 50, hard cap 500); "
          <> "the header reports whether results were truncated and how many files were scanned. "
          <> "A search looks at most "
          <> T.pack (show maxVisitedEntries)
          <> " entries, reads at most "
          <> T.pack (show (maxSearchBytes `div` (1024 * 1024)))
          <> " MiB of files, and stops at maxResults; then the file count is shown as 'N+'.",
      ttoolReadonly = True,
      ttoolExecute = const (withWorkspacePaths cfg . execute cfg)
    }

execute :: FsConfig -> GrepToolArgs -> IO Text
execute cfg args = do
  let rel = T.unpack args._grepPath
      pat = args._grepPattern
      ci = args._grepCaseInsensitive
      exts = normalizeExtensions args._grepExtensions
      cap = min maxResultsCap args._grepMaxResults
  if T.null pat
    then pure "Error: pattern must be non-empty"
    else if args._grepMaxResults <= 0
    then pure $ "Error: maxResults must be positive, got " <> T.pack (show args._grepMaxResults)
    else do
      resolved <- sandboxPath cfg rel
      isDir <- doesDirectoryExist resolved
      isFile <- doesFileExist resolved
      -- An explicit file gets the same size cap as files found by the walk.
      sizeOk <- if isFile then isSizeOk resolved else pure True
      if not (isDir || isFile)
        then pure $ "Error: path does not exist: " <> args._grepPath
        else
          if not sizeOk
            then
              pure $
                "Error: "
                  <> args._grepPath
                  <> " exceeds the "
                  <> T.pack (show maxFileSizeBytes)
                  <> "-byte cap for grep."
            else do
              budget <- newWalkBudget maxVisitedEntries
              st <- newIORef (GrepState [] 0 0 False 0 False)
              let search = searchFile st pat ci cap
              if isFile
                then search resolved (takeFileName resolved)
                else walkDir budget st search exts resolved "" 0
              GrepState matches _ scanned capHit bytesRead bytesOut <- readIORef st
              stopped <- budgetExhausted budget
              let notes =
                    catMaybes
                      [ if stopped then Just (budgetNote budget) else Nothing,
                        if bytesOut then Just (bytesNote bytesRead) else Nothing
                      ]
              pure $
                renderResults
                  args._grepPath
                  pat
                  (reverse matches)
                  scanned
                  ((capHit && not isFile) || stopped || bytesOut)
                  capHit
                  (if null notes then Nothing else Just (T.intercalate " | " notes))

-- | Lowercase and strip leading dots so callers can pass "hs" or ".hs".
normalizeExtensions :: [Text] -> [Text]
normalizeExtensions = map (T.toLower . T.dropWhile (== '.'))

-- | Matches so far (newest first), their count, files searched, whether
-- the walk stopped because @maxResults@ was reached, the file bytes read,
-- and whether the walk stopped because 'maxSearchBytes' ran out.
data GrepState = GrepState [Match] Int Int Bool Int Bool

-- | True once the search should look at no more files.
searchDone :: GrepState -> Bool
searchDone (GrepState _ _ _ capHit _ bytesOut) = capHit || bytesOut

-- | The note added to the output when 'maxSearchBytes' ran out.
bytesNote :: Int -> Text
bytesNote n = "stopped after reading " <> T.pack (show n) <> " bytes; narrow the path or pattern"

-- | Walk a directory and search each candidate file as it is found,
-- respecting depth, hidden-file, extension and size limits and the entry
-- budget. Stops once @maxResults@ is reached or 'maxSearchBytes' runs
-- out. Order is stable (entries
-- sorted) so output is deterministic. Each file is shown by its path
-- relative to the search root, built from the names walked, so no host
-- path is shown.
walkDir ::
  WalkBudget ->
  IORef GrepState ->
  (FilePath -> FilePath -> IO ()) ->
  [Text] ->
  FilePath ->
  FilePath ->
  Int ->
  IO ()
walkDir budget st search exts = go
  where
    go dir display depth
      | depth > maxDepth = pure ()
      | otherwise = do
          entries <- sort <$> safeListDirectory dir
          allowed <- takeEntries budget (filter (not . isFileHidden) entries)
          mapM_ (visit dir display depth) allowed
    visit parent display depth name = do
      done <- searchDone <$> readIORef st
      unless done $ do
        let full = parent </> name
            fullDisplay = display </> name
        -- Skip symlinks to prevent traversal out of the sandbox.
        isLink <- isSymlink full
        unless isLink $ do
          isDir <- doesDirectoryExist full
          if isDir
            then go full fullDisplay (depth + 1)
            else do
              isFile <- doesFileExist full
              ok <-
                if isFile && extensionAllowed exts name
                  then isSizeOk full
                  else pure False
              if ok then search full fullDisplay else pure ()

extensionAllowed :: [Text] -> FilePath -> Bool
extensionAllowed [] _ = True
extensionAllowed exts name =
  let ext = T.toLower (T.pack (dropDot (takeExtension name)))
   in ext `elem` exts
  where
    dropDot ('.' : rest) = rest
    dropDot s = s

isSizeOk :: FilePath -> IO Bool
isSizeOk path = do
  r <- try (getFileSize path) :: IO (Either IOException Integer)
  case r of
    Right sz -> pure (sz <= maxFileSizeBytes)
    Left _ -> pure False

safeListDirectory :: FilePath -> IO [FilePath]
safeListDirectory path = do
  r <- try (listDirectory path) :: IO (Either IOException [FilePath])
  case r of
    Right xs -> pure xs
    Left _ -> pure []

-- | Search one file and add its matches, up to the cap. Binary and
-- unreadable files count as searched but give no matches. A text file is
-- charged to 'maxSearchBytes' before it is read; a file that does not fit
-- is not read or counted, and ends the search.
searchFile :: IORef GrepState -> Text -> Bool -> Int -> FilePath -> FilePath -> IO ()
searchFile st pat ci cap f display = do
  r <-
    try (withTextHandle (withBinaryFile f ReadMode) readCharged) ::
      IO (Either IOException (Maybe (Maybe Text)))
  case r of
    Right (Just Nothing) -> pure ()
    _ -> do
      let ms = case r of
            Right (Just (Just content)) -> matchLines display content needle ci
            _ -> []
      modifyIORef' st $ \(GrepState acc n scanned _ used out) ->
        let remaining = cap - n
            taken = take remaining ms
            n' = n + length taken
         in GrepState (reverse taken ++ acc) n' (scanned + 1) (n' >= cap) used out
  where
    needle = if ci then T.toLower pat else pat
    -- Charge the file's size (at most the per-file cap) and read that
    -- many bytes, or refuse it when the budget cannot cover it.
    readCharged h = do
      sz <- min maxFileSizeBytes <$> hFileSize h
      let n = fromIntegral sz
      fits <- charge n
      if fits
        then Just . decodeUtf8With lenientDecode <$> BS.hGet h n
        else pure Nothing
    charge n = atomicModifyIORef' st $ \(GrepState acc k scanned capHit used out) ->
      if used + n <= maxSearchBytes
        then (GrepState acc k scanned capHit (used + n) out, True)
        else (GrepState acc k scanned capHit used True, False)

data Match = Match
  { mPath :: FilePath,
    mLine :: Int,
    mText :: Text
  }

matchLines :: FilePath -> Text -> Text -> Bool -> [Match]
matchLines path content needle ci =
  [ Match path n (truncateLine line)
    | (n, line) <- zip [1 ..] (T.lines content),
      let hay = if ci then T.toLower line else line,
      needle `T.isInfixOf` hay
  ]

truncateLine :: Text -> Text
truncateLine t
  | T.length t <= maxLineLength = t
  | otherwise = T.take maxLineLength t <> "…"

-- | Render the results. @partial@: the walk stopped early, so the file
-- count is a lower bound. @capHit@: @maxResults@ was reached.
renderResults :: Text -> Text -> [Match] -> Int -> Bool -> Bool -> Maybe Text -> Text
renderResults rootDisplay pat allMatches scanned partial capHit note =
  let (body, cut) = capLines maxOutputBytes (map fmtMatch allMatches)
      header =
        "=== grep "
          <> T.pack (show pat)
          <> " in "
          <> rootDisplay
          <> " | "
          <> T.pack (show (length body))
          <> " match(es) across "
          <> T.pack (show scanned)
          <> (if partial then "+" else "")
          <> " file(s)"
          <> ( if cut
                 then " | output truncated at " <> T.pack (show maxOutputBytes) <> " bytes"
                 else if capHit then " | stopped at maxResults (raise it to see more)" else ""
             )
          <> maybe "" (" | " <>) note
          <> " ==="
   in T.unlines (header : body)
  where
    fmtMatch m =
      T.pack m.mPath
        <> ":"
        <> T.pack (show m.mLine)
        <> ": "
        <> m.mText