From 0c3b77c06028205aac0184973037355689fc3c9e Mon Sep 17 00:00:00 2001
From: Ben Sima <ben@bensima.com>
Date: Mon, 1 Dec 2025 22:28:30 -0500
Subject: Compact amp-style timeline rendering and targeted file reading
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Timeline tool display:
- Grep/search: ✓ Grep pattern in filepath
- Read file: ✓ Read filepath @start-end
- Edit file: ✓ Edit filepath
- Bash: ϟ command (lightning bolt prompt)
- Tool results only shown for meaningful output

New search_and_read tool:
- Combines search + read in one operation
- Uses ripgrep --context for surrounding lines
- More efficient than separate search then read

Worker prompt updated to prefer search_and_read over
separate search + read_file calls
---
 Omni/Agent/Tools.hs | 84 +++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 81 insertions(+), 3 deletions(-)

(limited to 'Omni/Agent/Tools.hs')

diff --git a/Omni/Agent/Tools.hs b/Omni/Agent/Tools.hs
index 9664c76..22cc8a1 100644
--- a/Omni/Agent/Tools.hs
+++ b/Omni/Agent/Tools.hs
@@ -22,12 +22,14 @@ module Omni.Agent.Tools
     editFileTool,
     runBashTool,
     searchCodebaseTool,
+    searchAndReadTool,
     allTools,
     ReadFileArgs (..),
     WriteFileArgs (..),
     EditFileArgs (..),
     RunBashArgs (..),
     SearchCodebaseArgs (..),
+    SearchAndReadArgs (..),
     ToolResult (..),
     main,
     test,
@@ -78,8 +80,8 @@ test =
         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 "allTools contains 6 tools" <| do
+        length allTools Test.@=? 6,
       Test.unit "ReadFileArgs parses correctly" <| do
         let json = Aeson.object ["path" .= ("test.txt" :: Text)]
         case Aeson.fromJSON json of
@@ -220,7 +222,8 @@ allTools =
     writeFileTool,
     editFileTool,
     runBashTool,
-    searchCodebaseTool
+    searchCodebaseTool,
+    searchAndReadTool
   ]
 
 data ReadFileArgs = ReadFileArgs
@@ -602,3 +605,78 @@ executeSearchCodebase v =
           pure <| mkSuccess "No matches found"
         Exit.ExitFailure code ->
           pure <| mkError ("ripgrep failed with code " <> tshow code <> ": " <> Text.pack stderrStr)
+
+data SearchAndReadArgs = SearchAndReadArgs
+  { sarPattern :: Text,
+    sarPath :: Maybe Text,
+    sarContextLines :: Maybe Int
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON SearchAndReadArgs where
+  parseJSON =
+    Aeson.withObject "SearchAndReadArgs" <| \v ->
+      (SearchAndReadArgs </ (v .: "pattern"))
+        <*> (v .:? "path")
+        <*> (v .:? "context_lines")
+
+searchAndReadTool :: Engine.Tool
+searchAndReadTool =
+  Engine.Tool
+    { Engine.toolName = "search_and_read",
+      Engine.toolDescription =
+        "Search for a pattern, then read the matching lines with context. "
+          <> "More efficient than search + read separately - returns file content "
+          <> "around each match. Use this to find and understand code in one step.",
+      Engine.toolJsonSchema =
+        Aeson.object
+          [ "type" .= ("object" :: Text),
+            "properties"
+              .= Aeson.object
+                [ "pattern"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("Regex pattern to search for" :: Text)
+                      ],
+                  "path"
+                    .= Aeson.object
+                      [ "type" .= ("string" :: Text),
+                        "description" .= ("Optional: directory or file to search in" :: Text)
+                      ],
+                  "context_lines"
+                    .= Aeson.object
+                      [ "type" .= ("integer" :: Text),
+                        "description" .= ("Lines of context around each match (default: 10)" :: Text)
+                      ]
+                ],
+            "required" .= (["pattern"] :: [Text])
+          ],
+      Engine.toolExecute = executeSearchAndRead
+    }
+
+executeSearchAndRead :: Aeson.Value -> IO Aeson.Value
+executeSearchAndRead v =
+  case Aeson.fromJSON v of
+    Aeson.Error e -> pure <| mkError (Text.pack e)
+    Aeson.Success args -> do
+      let pat = Text.unpack (sarPattern args)
+          ctx = fromMaybe 10 (sarContextLines args)
+          pathArg = maybe ["."] (\p -> [Text.unpack p]) (sarPath args)
+          rgArgs =
+            [ "--line-number",
+              "--no-heading",
+              "--context=" <> show ctx,
+              "--max-count=20",
+              "--ignore-case",
+              pat
+            ]
+              <> pathArg
+          proc = Process.proc "rg" rgArgs
+      (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: " <> tshow code <> ": " <> Text.pack stderrStr)
-- 
cgit v1.2.3