summaryrefslogtreecommitdiff
path: root/Omni/Agent
diff options
context:
space:
mode:
Diffstat (limited to 'Omni/Agent')
-rw-r--r--Omni/Agent/Subagent.hs88
-rw-r--r--Omni/Agent/Telegram.hs31
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 ()
generated by cgit v1.2.3 (git 2.25.1) at 2025-12-30 12:08:55 +0000