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>

2 files changed, 108 insertions, 11 deletions

diff --git a/Omni/Agent/Subagent.hs b/Omni/Agent/Subagent.hs
index cb8c090..29286a0 100644
--- a/Omni/Agent/Subagent.hs
+++ b/Omni/Agent/Subagent.hs
@@ -62,6 +62,8 @@ module Omni.Agent.Subagent
     getPendingSpawn,
     removePendingSpawn,
     approveAndSpawnSubagent,
+    approveAndSpawnSubagentWithCallback,
+    CompletionCallback,
     rejectPendingSpawn,
     cleanupExpiredPending,
 
@@ -532,6 +534,80 @@ spawnSubagentAsyncWithId sessionId userId keys config maybePregenId = do
         handleStatus = statusVar
       }
 
+-- | Spawn subagent with optional pre-generated ID and external completion callback
+spawnSubagentAsyncWithIdAndCallback :: AuditLog.SessionId -> Maybe Text -> SubagentApiKeys -> SubagentConfig -> Maybe AuditLog.SubagentId -> Maybe (Text -> SubagentResult -> IO ()) -> IO SubagentHandle
+spawnSubagentAsyncWithIdAndCallback sessionId userId keys config maybePregenId maybeExternalCallback = do
+  sid <- maybe AuditLog.newSubagentId pure maybePregenId
+  startTime <- Clock.getCurrentTime
+  statusVar <- newTVarIO initialRunStatus
+
+  let logEntry evType content = do
+        entry <-
+          AuditLog.mkLogEntry
+            sessionId
+            (AuditLog.AgentId ("subagent-" <> AuditLog.unSubagentId sid))
+            userId
+            evType
+            content
+            AuditLog.emptyMetadata
+        AuditLog.writeSubagentLog sid entry
+
+  logEntry AuditLog.SubagentSpawn
+    <| Aeson.object
+      [ "role" .= subagentRole config,
+        "task" .= subagentTask config,
+        "subagent_id" .= sid
+      ]
+
+  let callbacks =
+        SubagentCallbacks
+          { onSubagentStart = \msg -> do
+              logEntry AuditLog.AssistantMessage (Aeson.String msg)
+              atomically <| writeTVar statusVar <| initialRunStatus {runCurrentActivity = msg},
+            onSubagentActivity = \msg -> do
+              now <- Clock.getCurrentTime
+              let elapsed = round (Clock.diffUTCTime now startTime)
+              logEntry AuditLog.AssistantMessage (Aeson.String msg)
+              atomically <| do
+                status <- readTVar statusVar
+                writeTVar statusVar <| status {runCurrentActivity = msg, runElapsedSeconds = elapsed},
+            onSubagentToolCall = \tool args -> do
+              now <- Clock.getCurrentTime
+              let elapsed = round (Clock.diffUTCTime now startTime)
+              logEntry AuditLog.ToolCall (Aeson.object ["tool" .= tool, "args" .= args])
+              atomically <| do
+                status <- readTVar statusVar
+                writeTVar statusVar
+                  <| status
+                    { runCurrentActivity = "Calling " <> tool,
+                      runLastToolCall = Just (tool, now),
+                      runElapsedSeconds = elapsed
+                    },
+            onSubagentComplete = \result -> do
+              logEntry AuditLog.SubagentComplete
+                <| Aeson.object
+                  [ "status" .= subagentStatus result,
+                    "summary" .= subagentSummary result,
+                    "tokens" .= subagentTokensUsed result,
+                    "cost_cents" .= subagentCostCents result,
+                    "duration" .= subagentDuration result
+                  ]
+              case maybeExternalCallback of
+                Just cb -> cb (AuditLog.unSubagentId sid) result
+                Nothing -> pure ()
+          }
+
+  asyncHandle <- async (runSubagentWithCallbacks keys config callbacks)
+
+  pure
+    SubagentHandle
+      { handleId = sid,
+        handleAsync = asyncHandle,
+        handleStartTime = startTime,
+        handleConfig = config,
+        handleStatus = statusVar
+      }
+
 querySubagentStatus :: SubagentHandle -> IO SubagentRunStatus
 querySubagentStatus h = do
   now <- Clock.getCurrentTime
@@ -1189,7 +1265,15 @@ executeSpawnWithApproval _keys chatId onApprovalNeeded v =
 
 -- | Approve a pending spawn and start the subagent
 approveAndSpawnSubagent :: SubagentApiKeys -> Text -> IO (Either Text Text)
-approveAndSpawnSubagent keys pid = do
+approveAndSpawnSubagent keys pid =
+  approveAndSpawnSubagentWithCallback keys pid Nothing
+
+-- | Callback invoked when subagent completes
+type CompletionCallback = Text -> SubagentResult -> IO ()
+
+-- | Approve a pending spawn and start the subagent, with optional completion callback
+approveAndSpawnSubagentWithCallback :: SubagentApiKeys -> Text -> Maybe CompletionCallback -> IO (Either Text Text)
+approveAndSpawnSubagentWithCallback keys pid maybeOnComplete = do
   maybePending <- getPendingSpawn pid
   case maybePending of
     Nothing -> pure (Left "Pending spawn not found or expired")
@@ -1197,7 +1281,7 @@ approveAndSpawnSubagent keys pid = do
       removePendingSpawn pid
       uuid <- Data.UUID.V4.nextRandom
       let sessionId = AuditLog.SessionId ("subagent-" <> Text.take 8 (Data.UUID.toText uuid))
-      subHandle <- spawnSubagentAsyncWithId sessionId Nothing keys (pendingConfig pending) (Just (pendingSubagentId pending))
+      subHandle <- spawnSubagentAsyncWithIdAndCallback sessionId Nothing keys (pendingConfig pending) (Just (pendingSubagentId pending)) maybeOnComplete
       registerSubagent subHandle
       let sid = AuditLog.unSubagentId (handleId subHandle)
       pure (Right sid)
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 ()