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, 86 insertions, 2 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)