task: t-rWa5yilwM.5 done (Omni.Agent.Log enhancements)

Amp-Thread-ID:
https://ampcode.com/threads/T-7109f8d0-feb4-4a24-bc4b-37743227e2cb
Co-authored-by: Amp <amp@ampcode.com>

2 files changed, 110 insertions, 6 deletions

diff --git a/Omni/Agent/Log.hs b/Omni/Agent/Log.hs
index c93479b..28dbab2 100644
--- a/Omni/Agent/Log.hs
+++ b/Omni/Agent/Log.hs
@@ -4,6 +4,10 @@
 
 module Omni.Agent.Log
   ( LogEntry (..),
+    Status (..),
+    initialStatus,
+    updateStatus,
+    renderStatus,
     parseLine,
     format,
   )
@@ -13,6 +17,7 @@ import Alpha
 import Data.Aeson (FromJSON (..), (.:), (.:?))
 import qualified Data.Aeson as Aeson
 import qualified Data.ByteString.Lazy as BSL
+import qualified Data.Set as Set
 
 data LogEntry = LogEntry
   { leMessage :: Text,
@@ -20,7 +25,8 @@ data LogEntry = LogEntry
     leToolName :: Maybe Text,
     leBatches :: Maybe [[Text]],
     leMethod :: Maybe Text,
-    lePath :: Maybe Text
+    lePath :: Maybe Text,
+    leTimestamp :: Maybe Text
   }
   deriving (Show, Eq, Generic)
 
@@ -40,6 +46,52 @@ instance FromJSON LogEntry where
         .:? "method"
         <*> v
         .:? "path"
+        <*> v
+        .:? "timestamp"
+
+data Status = Status
+  { sWorkerName :: Text,
+    sTaskId :: Maybe Text,
+    sFiles :: Set.Set Text,
+    sStartTime :: Maybe Text,
+    sLastActivity :: Text
+  }
+  deriving (Show, Eq, Generic)
+
+initialStatus :: Text -> Status
+initialStatus name =
+  Status
+    { sWorkerName = name,
+      sTaskId = Nothing,
+      sFiles = Set.empty,
+      sStartTime = Nothing,
+      sLastActivity = "Idle"
+    }
+
+updateStatus :: LogEntry -> Status -> Status
+updateStatus e s =
+  let s' = case format e of
+        Just msg -> s {sLastActivity = msg}
+        Nothing -> s
+      s'' = case leTimestamp e of
+        Just t -> if isNothing (sStartTime s) then s' {sStartTime = Just t} else s'
+        Nothing -> s'
+   in case (leMessage e, lePath e) of
+        ("ide-fs", Just p) -> s'' {sFiles = Set.insert p (sFiles s'')}
+        _ -> s''
+
+renderStatus :: Status -> Text
+renderStatus s =
+  let line1 =
+        "[Worker: "
+          <> sWorkerName s
+          <> "] "
+          <> "Task: "
+          <> fromMaybe "None" (sTaskId s)
+          <> " | Files: "
+          <> show (Set.size (sFiles s))
+      line2 = sLastActivity s
+   in line1 <> "\n" <> line2
 
 parseLine :: Text -> Maybe LogEntry
 parseLine line = Aeson.decode <| BSL.fromStrict <| encodeUtf8 line
diff --git a/Omni/Agent/LogTest.hs b/Omni/Agent/LogTest.hs
index 0d085b1..518147e 100644
--- a/Omni/Agent/LogTest.hs
+++ b/Omni/Agent/LogTest.hs
@@ -5,6 +5,7 @@
 module Omni.Agent.LogTest where
 
 import Alpha
+import qualified Data.Set as Set
 import Omni.Agent.Log
 import qualified Omni.Test as Test
 
@@ -16,7 +17,9 @@ tests =
   Test.group
     "Omni.Agent.Log"
     [ Test.unit "Parse LogEntry" testParse,
-      Test.unit "Format LogEntry" testFormat
+      Test.unit "Format LogEntry" testFormat,
+      Test.unit "Update Status" testUpdateStatus,
+      Test.unit "Render Status" testRenderStatus
     ]
 
 testParse :: IO ()
@@ -29,7 +32,8 @@ testParse = do
             leToolName = Nothing,
             leBatches = Just [["grep"]],
             leMethod = Nothing,
-            lePath = Nothing
+            lePath = Nothing,
+            leTimestamp = Nothing
           }
   parseLine json @?= Just expected
 
@@ -42,7 +46,8 @@ testFormat = do
             leToolName = Nothing,
             leBatches = Just [["grep"]],
             leMethod = Nothing,
-            lePath = Nothing
+            lePath = Nothing,
+            leTimestamp = Nothing
           }
   format entry @?= Just "🤖  THOUGHT: Planning tool execution (grep)"
 
@@ -53,7 +58,8 @@ testFormat = do
             leToolName = Nothing,
             leBatches = Nothing,
             leMethod = Nothing,
-            lePath = Nothing
+            lePath = Nothing,
+            leTimestamp = Nothing
           }
   format entry2 @?= Nothing
 
@@ -64,9 +70,55 @@ testFormat = do
             leToolName = Nothing,
             leBatches = Nothing,
             leMethod = Nothing,
-            lePath = Nothing
+            lePath = Nothing,
+            leTimestamp = Nothing
           }
   format entry3 @?= Just "❌  ERROR: some error"
 
+testUpdateStatus :: IO ()
+testUpdateStatus = do
+  let s0 = initialStatus "worker-1"
+  let e1 =
+        LogEntry
+          { leMessage = "executing 1 tools in 1 batch(es)",
+            leLevel = Nothing,
+            leToolName = Nothing,
+            leBatches = Just [["grep"]],
+            leMethod = Nothing,
+            lePath = Nothing,
+            leTimestamp = Just "12:00:00"
+          }
+  let s1 = updateStatus e1 s0
+  sLastActivity s1 @?= "🤖  THOUGHT: Planning tool execution (grep)"
+  sStartTime s1 @?= Just "12:00:00"
+
+  let e2 =
+        LogEntry
+          { leMessage = "ide-fs",
+            leLevel = Nothing,
+            leToolName = Nothing,
+            leBatches = Nothing,
+            leMethod = Just "readFile",
+            lePath = Just "/path/to/file",
+            leTimestamp = Just "12:00:01"
+          }
+  let s2 = updateStatus e2 s1
+  sLastActivity s2 @?= "📂  READ: /path/to/file"
+  Set.member "/path/to/file" (sFiles s2) @?= True
+  sStartTime s2 @?= Just "12:00:00" -- Should preserve start time
+
+testRenderStatus :: IO ()
+testRenderStatus = do
+  let s =
+        Status
+          { sWorkerName = "worker-1",
+            sTaskId = Just "t-123",
+            sFiles = Set.fromList ["file1", "file2"],
+            sStartTime = Just "12:00",
+            sLastActivity = "Running..."
+          }
+  let output = renderStatus s
+  output @?= "[Worker: worker-1] Task: t-123 | Files: 2\nRunning..."
+
 (@?=) :: (Eq a, Show a) => a -> a -> IO ()
 (@?=) = (Test.@?=)