summaryrefslogtreecommitdiff
path: root/Omni/Agent/Subagent.hs
diff options
context:
space:
mode:
Diffstat (limited to 'Omni/Agent/Subagent.hs')
-rw-r--r--Omni/Agent/Subagent.hs169
1 files changed, 168 insertions, 1 deletions
diff --git a/Omni/Agent/Subagent.hs b/Omni/Agent/Subagent.hs
index c251e9d..3278e4c 100644
--- a/Omni/Agent/Subagent.hs
+++ b/Omni/Agent/Subagent.hs
@@ -43,14 +43,26 @@ module Omni.Agent.Subagent
-- * Tool
spawnSubagentTool,
+ spawnSubagentToolWithApproval,
checkSubagentTool,
subagentTools,
+ subagentToolsWithApproval,
+ ApprovalCallback,
-- * Registry
getSubagentHandle,
listRunningSubagents,
cleanupRegistry,
+ -- * Pending Spawns (for Telegram button confirmation)
+ PendingSpawn (..),
+ createPendingSpawn,
+ getPendingSpawn,
+ removePendingSpawn,
+ approveAndSpawnSubagent,
+ rejectPendingSpawn,
+ cleanupExpiredPending,
+
-- * Role-specific tools
SubagentApiKeys (..),
toolsForRole,
@@ -124,6 +136,53 @@ cleanupRegistry = do
stillRunning <- filterM (\(_, h) -> fmap not (isSubagentDone h)) (Map.toList registry)
modifyIORef' subagentRegistry (const (Map.fromList stillRunning))
+-- | A pending spawn request awaiting user confirmation
+data PendingSpawn = PendingSpawn
+ { pendingId :: Text,
+ pendingConfig :: SubagentConfig,
+ pendingChatId :: Int,
+ pendingCreatedAt :: Clock.UTCTime
+ }
+ deriving (Show, Eq)
+
+-- | Global registry of pending spawn requests
+pendingSpawnRegistry :: IORef (Map.Map Text PendingSpawn)
+pendingSpawnRegistry = unsafePerformIO (newIORef Map.empty)
+{-# NOINLINE pendingSpawnRegistry #-}
+
+-- | Create a new pending spawn request
+createPendingSpawn :: SubagentConfig -> Int -> IO Text
+createPendingSpawn config chatId = do
+ uuid <- Data.UUID.V4.nextRandom
+ let pendingId = Text.take 8 (Data.UUID.toText uuid)
+ now <- Clock.getCurrentTime
+ let pending =
+ PendingSpawn
+ { pendingId = pendingId,
+ pendingConfig = config,
+ pendingChatId = chatId,
+ pendingCreatedAt = now
+ }
+ modifyIORef' pendingSpawnRegistry (Map.insert pendingId pending)
+ pure pendingId
+
+-- | Get a pending spawn by ID
+getPendingSpawn :: Text -> IO (Maybe PendingSpawn)
+getPendingSpawn pid = do
+ registry <- readIORef pendingSpawnRegistry
+ pure (Map.lookup pid registry)
+
+-- | Remove a pending spawn (after approval/rejection)
+removePendingSpawn :: Text -> IO ()
+removePendingSpawn pid = modifyIORef' pendingSpawnRegistry (Map.delete pid)
+
+-- | Clean up expired pending spawns (older than 10 minutes)
+cleanupExpiredPending :: IO ()
+cleanupExpiredPending = do
+ now <- Clock.getCurrentTime
+ let isExpired p = Clock.diffUTCTime now (pendingCreatedAt p) > 600
+ modifyIORef' pendingSpawnRegistry (Map.filter (not <. isExpired))
+
main :: IO ()
main = Test.run test
@@ -847,6 +906,114 @@ executeCheckSubagent v = do
"cost_cents" .= runCostCents status
]
--- | All subagent-related tools
+-- | All subagent-related tools (legacy - agent can bypass approval)
subagentTools :: SubagentApiKeys -> [Engine.Tool]
subagentTools keys = [spawnSubagentTool keys, checkSubagentTool]
+
+-- | Callback for sending approval buttons
+-- Args: chatId, pendingId, role, task, estimatedMinutes, maxCostCents
+type ApprovalCallback = Int -> Text -> Text -> Text -> Int -> Double -> IO ()
+
+-- | Spawn subagent tool that requires external approval via callback
+spawnSubagentToolWithApproval :: SubagentApiKeys -> Int -> ApprovalCallback -> Engine.Tool
+spawnSubagentToolWithApproval keys chatId onApprovalNeeded =
+ Engine.Tool
+ { Engine.toolName = "spawn_subagent",
+ Engine.toolDescription =
+ "Spawn a specialized subagent for a focused task. "
+ <> "The user will receive a confirmation button to approve the spawn. "
+ <> "Available roles: web_crawler (fast web research), code_reviewer (thorough code analysis), "
+ <> "data_extractor (structured data extraction), researcher (general research).",
+ Engine.toolJsonSchema =
+ Aeson.object
+ [ "type" .= ("object" :: Text),
+ "properties"
+ .= Aeson.object
+ [ "role"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "enum" .= (["web_crawler", "code_reviewer", "data_extractor", "researcher"] :: [Text]),
+ "description" .= ("Subagent role determining tools and model" :: Text)
+ ],
+ "task"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("The specific task for the subagent to accomplish" :: Text)
+ ],
+ "context"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("Additional context to help the subagent understand the goal" :: Text)
+ ],
+ "model"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("Override the default model for this role" :: Text)
+ ],
+ "timeout"
+ .= Aeson.object
+ [ "type" .= ("integer" :: Text),
+ "description" .= ("Timeout in seconds (default: 600)" :: Text)
+ ],
+ "max_cost_cents"
+ .= Aeson.object
+ [ "type" .= ("number" :: Text),
+ "description" .= ("Maximum cost in cents (default: 50)" :: Text)
+ ]
+ ],
+ "required" .= (["role", "task"] :: [Text])
+ ],
+ Engine.toolExecute = executeSpawnWithApproval keys chatId onApprovalNeeded
+ }
+
+executeSpawnWithApproval :: SubagentApiKeys -> Int -> ApprovalCallback -> Aeson.Value -> IO Aeson.Value
+executeSpawnWithApproval _keys chatId onApprovalNeeded v =
+ case Aeson.fromJSON v of
+ Aeson.Error e -> pure <| Aeson.object ["error" .= ("Invalid arguments: " <> Text.pack e)]
+ Aeson.Success config -> do
+ pendingId <- createPendingSpawn config chatId
+ let roleText = case subagentRole config of
+ WebCrawler -> "web_crawler"
+ CodeReviewer -> "code_reviewer"
+ DataExtractor -> "data_extractor"
+ Researcher -> "researcher"
+ CustomRole name -> name
+ estimatedMins = subagentTimeout config `div` 60
+ maxCost = subagentMaxCost config
+ onApprovalNeeded chatId pendingId roleText (subagentTask config) estimatedMins maxCost
+ pure
+ <| Aeson.object
+ [ "status" .= ("pending_approval" :: Text),
+ "pending_id" .= pendingId,
+ "message" .= ("Approval button sent to user. Waiting for confirmation." :: Text)
+ ]
+
+-- | Approve a pending spawn and start the subagent
+approveAndSpawnSubagent :: SubagentApiKeys -> Text -> IO (Either Text Text)
+approveAndSpawnSubagent keys pendingId = do
+ maybePending <- getPendingSpawn pendingId
+ case maybePending of
+ Nothing -> pure (Left "Pending spawn not found or expired")
+ Just pending -> do
+ removePendingSpawn pendingId
+ uuid <- Data.UUID.V4.nextRandom
+ let sessionId = AuditLog.SessionId ("subagent-" <> Text.take 8 (Data.UUID.toText uuid))
+ subHandle <- spawnSubagentAsync sessionId Nothing keys (pendingConfig pending)
+ registerSubagent subHandle
+ let sid = AuditLog.unSubagentId (handleId subHandle)
+ pure (Right sid)
+
+-- | Reject a pending spawn
+rejectPendingSpawn :: Text -> IO Bool
+rejectPendingSpawn pendingId = do
+ maybePending <- getPendingSpawn pendingId
+ case maybePending of
+ Nothing -> pure False
+ Just _ -> do
+ removePendingSpawn pendingId
+ pure True
+
+-- | All subagent-related tools with approval callback
+subagentToolsWithApproval :: SubagentApiKeys -> Int -> ApprovalCallback -> [Engine.Tool]
+subagentToolsWithApproval keys chatId onApproval =
+ [spawnSubagentToolWithApproval keys chatId onApproval, checkSubagentTool]