From fcb8629182fa1552e4a840ccd4ec0aa2b8042cc0 Mon Sep 17 00:00:00 2001
From: Ben Sima <ben@bensima.com>
Date: Fri, 19 Dec 2025 21:54:54 -0500
Subject: feat(ava): add tool trace viewer mini-app

- Add SQLite storage for tool traces (Omni/Ava/Trace.hs)
- Add web server to serve trace viewer (Omni/Ava/Web.hs)
- Add HTML/CSS/JS trace viewer UI (Omni/Ava/Web/trace.html)
- Integrate trace storage into Engine.hs tool execution callback
- Add trace links to Telegram responses when AVA_WEB_URL is set
- Configure Tailscale Funnel for public access
- Fix pre-push hook variable scope bug
- Add direnv, bash, nix to Ava service PATH
- Add mustache dep to Ava.hs for template rendering

Epic: t-272
---
 Omni/Agent/Telegram.hs | 54 +++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 51 insertions(+), 3 deletions(-)

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

diff --git a/Omni/Agent/Telegram.hs b/Omni/Agent/Telegram.hs
index 7b2beaa..7183592 100644
--- a/Omni/Agent/Telegram.hs
+++ b/Omni/Agent/Telegram.hs
@@ -76,6 +76,7 @@ import Data.Aeson ((.=))
 import qualified Data.Aeson as Aeson
 import qualified Data.Aeson.KeyMap as KeyMap
 import qualified Data.ByteString.Lazy as BL
+import Data.IORef (newIORef, readIORef, writeIORef)
 import qualified Data.Text as Text
 import qualified Data.Text.Encoding as TE
 import Data.Time (getCurrentTime, utcToLocalTime)
@@ -110,6 +111,7 @@ import qualified Omni.Agent.Tools.Python as Python
 import qualified Omni.Agent.Tools.Todos as Todos
 import qualified Omni.Agent.Tools.WebReader as WebReader
 import qualified Omni.Agent.Tools.WebSearch as WebSearch
+import qualified Omni.Ava.Trace as Trace
 import qualified Omni.Test as Test
 import System.Environment (lookupEnv)
 import Text.Printf (printf)
@@ -524,6 +526,14 @@ leaveChat cfg chatId = do
 runTelegramBot :: Types.TelegramConfig -> Provider.Provider -> IO ()
 runTelegramBot tgConfig provider = do
   putText "Starting Telegram bot..."
+
+  cleanedCount <- Memory.withMemoryDb Trace.cleanupOldTraces
+  when (cleanedCount > 0)
+    <| putText
+    <| "Cleaned up "
+    <> tshow cleanedCount
+    <> " old tool traces"
+
   offsetVar <- newTVarIO 0
 
   botUsername <- getBotUsername tgConfig
@@ -551,7 +561,23 @@ runTelegramBot tgConfig provider = do
             Engine.engineOnToolResult = \toolName success result ->
               putText <| "Tool result: " <> toolName <> " " <> (if success then "ok" else "err") <> " " <> Text.take 200 result,
             Engine.engineOnActivity = \activity ->
-              putText <| "Agent: " <> activity
+              putText <| "Agent: " <> activity,
+            Engine.engineOnToolTrace = \toolName input output durationMs -> do
+              now <- getCurrentTime
+              let truncatedOutput = Text.take 1000000 output
+                  traceRecord =
+                    Trace.TraceRecord
+                      { Trace.trcId = "",
+                        Trace.trcCreatedAt = tshow now,
+                        Trace.trcToolName = toolName,
+                        Trace.trcInput = input,
+                        Trace.trcOutput = truncatedOutput,
+                        Trace.trcDurationMs = durationMs,
+                        Trace.trcUserId = Nothing,
+                        Trace.trcChatId = Nothing
+                      }
+              tid <- Memory.withMemoryDb <| \conn -> Trace.insertTrace conn traceRecord
+              pure (Just tid)
           }
 
   let processBatch = handleMessageBatch tgConfig provider engineCfg botName
@@ -1061,6 +1087,17 @@ processEngagedMessage ::
 processEngagedMessage tgConfig provider engineCfg msg uid userName chatId userMessage conversationContext = do
   let isGroup = Types.isGroupChat msg
 
+  lastTraceIdRef <- newIORef (Nothing :: Maybe Text)
+  let engineCfgWithTrace =
+        engineCfg
+          { Engine.engineOnToolTrace = \toolName input output durationMs -> do
+              maybeTid <- Engine.engineOnToolTrace engineCfg toolName input output durationMs
+              case maybeTid of
+                Just tid -> writeIORef lastTraceIdRef (Just tid)
+                Nothing -> pure ()
+              pure maybeTid
+          }
+
   personalMemories <- Memory.recallMemories uid userMessage 5
   groupMemories <-
     if isGroup
@@ -1245,7 +1282,11 @@ processEngagedMessage tgConfig provider engineCfg msg uid userName chatId userMe
 
   result <-
     withTypingIndicator tgConfig chatId
-      <| Engine.runAgentWithProvider engineCfg provider agentCfg userMessage
+      <| Engine.runAgentWithProvider engineCfgWithTrace provider agentCfg userMessage
+
+  lastTraceId <- readIORef lastTraceIdRef
+  maybeWebUrl <- lookupEnv "AVA_WEB_URL"
+  let traceLink = formatTraceLink lastTraceId (Text.pack </ maybeWebUrl)
 
   case result of
     Left err -> do
@@ -1253,7 +1294,8 @@ processEngagedMessage tgConfig provider engineCfg msg uid userName chatId userMe
       _ <- Messages.enqueueImmediate (Just uid) chatId (Types.tmThreadId msg) "sorry, i hit an error. please try again." (Just "agent_error") Nothing
       pure ()
     Right agentResult -> do
-      let response = Engine.resultFinalMessage agentResult
+      let baseResponse = Engine.resultFinalMessage agentResult
+          response = baseResponse <> traceLink
           threadId = Types.tmThreadId msg
       putText <| "Response text: " <> Text.take 200 response
 
@@ -1367,6 +1409,12 @@ mergeTooShort (x : y : rest)
   | Text.length x < 100 = mergeTooShort ((x <> "\n\n" <> y) : rest)
   | otherwise = x : mergeTooShort (y : rest)
 
+formatTraceLink :: Maybe Text -> Maybe Text -> Text
+formatTraceLink Nothing _ = ""
+formatTraceLink _ Nothing = ""
+formatTraceLink (Just tid) (Just baseUrl) =
+  "\n\n[view trace](" <> baseUrl <> "/trace/" <> tid <> ")"
+
 enqueueMultipart :: Maybe Text -> Int -> Maybe Int -> [Text] -> Maybe Text -> IO ()
 enqueueMultipart _ _ _ [] _ = pure ()
 enqueueMultipart mUid chatId mThreadId parts msgType = do
-- 
cgit v1.2.3