Omni/Ava: improve trace viewer and subagent notifications

- Add subagent completion callback to notify user when subagent finishes
- Show tool name in 'view trace' link (e.g. 'view web_search trace')
- Pretty-print JSON on trace web page using aeson-pretty

Amp-Thread-ID: https://ampcode.com/threads/T-019b3a13-bc75-7368-9ec9-362d462a022c
Co-authored-by: Amp <amp@ampcode.com>

1 files changed, 22 insertions, 9 deletions

diff --git a/Omni/Agent/Telegram.hs b/Omni/Agent/Telegram.hs
index 7183592..ae87c44 100644
--- a/Omni/Agent/Telegram.hs
+++ b/Omni/Agent/Telegram.hs
@@ -667,7 +667,20 @@ handleCallbackQueryInner tgConfig cq chatId userId callbackData = do
                 { Subagent.subagentOpenRouterKey = Types.tgOpenRouterApiKey tgConfig,
                   Subagent.subagentKagiKey = Types.tgKagiApiKey tgConfig
                 }
-        spawnResult <- Subagent.approveAndSpawnSubagent keys pendingId
+            onComplete :: Subagent.CompletionCallback
+            onComplete sid result = do
+              putText <| "Subagent " <> sid <> " completed"
+              let statusText = case Subagent.subagentStatus result of
+                    Subagent.SubagentSuccess -> "completed successfully"
+                    Subagent.SubagentTimeout -> "timed out"
+                    Subagent.SubagentCostExceeded -> "exceeded cost limit"
+                    Subagent.SubagentError err -> "failed: " <> err
+                  summary = Text.take 500 (Subagent.subagentSummary result)
+                  costStr = Text.pack (printf "$%.2f" (Subagent.subagentCostCents result / 100))
+                  durationStr = tshow (Subagent.subagentDuration result) <> "s"
+                  baseMsg = "subagent " <> sid <> " " <> statusText <> " (" <> durationStr <> ", " <> costStr <> ")\n\n" <> summary
+              sendMessage tgConfig chatId baseMsg
+        spawnResult <- Subagent.approveAndSpawnSubagentWithCallback keys pendingId (Just onComplete)
         putText <| "Spawn result: " <> tshow spawnResult
         case spawnResult of
           Left err -> do
@@ -675,7 +688,7 @@ handleCallbackQueryInner tgConfig cq chatId userId callbackData = do
             sendMessage tgConfig chatId ("Failed to spawn subagent: " <> err)
           Right subagentId -> do
             putText <| "Spawn succeeded, subagent ID: " <> subagentId
-            sendMessage tgConfig chatId ("Subagent " <> subagentId <> " spawned! Use check_subagent to monitor progress.")
+            sendMessage tgConfig chatId ("Subagent " <> subagentId <> " spawned! I'll notify you when it completes.")
       ["subagent_reject", pendingId] -> do
         putText <| "Rejecting subagent spawn: " <> pendingId
         rejected <- Subagent.rejectPendingSpawn pendingId
@@ -1087,13 +1100,13 @@ processEngagedMessage ::
 processEngagedMessage tgConfig provider engineCfg msg uid userName chatId userMessage conversationContext = do
   let isGroup = Types.isGroupChat msg
 
-  lastTraceIdRef <- newIORef (Nothing :: Maybe Text)
+  lastTraceRef <- newIORef (Nothing :: Maybe (Text, 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)
+                Just tid -> writeIORef lastTraceRef (Just (tid, toolName))
                 Nothing -> pure ()
               pure maybeTid
           }
@@ -1284,9 +1297,9 @@ processEngagedMessage tgConfig provider engineCfg msg uid userName chatId userMe
     withTypingIndicator tgConfig chatId
       <| Engine.runAgentWithProvider engineCfgWithTrace provider agentCfg userMessage
 
-  lastTraceId <- readIORef lastTraceIdRef
+  lastTrace <- readIORef lastTraceRef
   maybeWebUrl <- lookupEnv "AVA_WEB_URL"
-  let traceLink = formatTraceLink lastTraceId (Text.pack </ maybeWebUrl)
+  let traceLink = formatTraceLink lastTrace (Text.pack </ maybeWebUrl)
 
   case result of
     Left err -> do
@@ -1409,11 +1422,11 @@ 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 :: Maybe (Text, Text) -> Maybe Text -> Text
 formatTraceLink Nothing _ = ""
 formatTraceLink _ Nothing = ""
-formatTraceLink (Just tid) (Just baseUrl) =
-  "\n\n[view trace](" <> baseUrl <> "/trace/" <> tid <> ")"
+formatTraceLink (Just (tid, toolName)) (Just baseUrl) =
+  "\n\n[view " <> toolName <> " trace](" <> baseUrl <> "/trace/" <> tid <> ")"
 
 enqueueMultipart :: Maybe Text -> Int -> Maybe Int -> [Text] -> Maybe Text -> IO ()
 enqueueMultipart _ _ _ [] _ = pure ()