path: root/Omni/Agent/Subagent.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | Subagent system for spawning specialized agents.
--
-- Enables the orchestrator (Ava) to delegate focused tasks to specialized
-- subagents that run with their own tool sets and resource limits.
--
-- Key features:
-- - Role-based tool selection (WebCrawler, CodeReviewer, etc.)
-- - Per-subagent resource limits (timeout, cost, tokens)
-- - Structured result format with confidence scores
-- - No sub-subagent spawning (hierarchical control)
-- - Async execution with status polling
-- - Audit logging for all events
--
-- : out omni-agent-subagent
-- : dep aeson
-- : dep async
-- : dep directory
-- : dep mustache
-- : dep stm
-- : dep uuid
module Omni.Agent.Subagent
  ( -- * Types
    SubagentRole (..),
    SubagentConfig (..),
    SubagentResult (..),
    SubagentStatus (..),
    SubagentCallbacks (..),
    SubagentHandle (..),
    SubagentRunStatus (..),

    -- * Async Execution
    spawnSubagentAsync,
    querySubagentStatus,
    isSubagentDone,
    waitSubagent,
    cancelSubagent,

    -- * Sync Execution (legacy)
    runSubagent,
    runSubagentWithCallbacks,

    -- * 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,
    modelForRole,
    systemPromptForRole,

    -- * Defaults
    defaultSubagentConfig,
    defaultCallbacks,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Control.Concurrent.STM (TVar, newTVarIO, readTVar, readTVarIO, writeTVar)
import Data.Aeson ((.!=), (.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.KeyMap as KeyMap
import Data.IORef (IORef, modifyIORef', newIORef, readIORef)
import qualified Data.Map.Strict as Map
import qualified Data.Text as Text
import qualified Data.Time.Clock as Clock
import qualified Data.UUID
import qualified Data.UUID.V4
import qualified Omni.Agent.AuditLog as AuditLog
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Agent.Prompts as Prompts
import qualified Omni.Agent.Provider as Provider
import qualified Omni.Agent.Subagent.Coder as Coder
import qualified Omni.Agent.Tools as Tools
import qualified Omni.Agent.Tools.WebReader as WebReader
import qualified Omni.Agent.Tools.WebSearch as WebSearch
import qualified Omni.Test as Test
import System.IO.Unsafe (unsafePerformIO)
import Text.Printf (printf)

-- | Global registry of running subagents, keyed by SubagentId
subagentRegistry :: IORef (Map.Map Text SubagentHandle)
subagentRegistry = unsafePerformIO (newIORef Map.empty)
{-# NOINLINE subagentRegistry #-}

-- | Register a subagent handle
registerSubagent :: SubagentHandle -> IO ()
registerSubagent h = do
  let key = AuditLog.unSubagentId (handleId h)
  modifyIORef' subagentRegistry (Map.insert key h)

-- | Get a subagent handle by ID
getSubagentHandle :: Text -> IO (Maybe SubagentHandle)
getSubagentHandle sid = do
  registry <- readIORef subagentRegistry
  pure (Map.lookup sid registry)

-- | List all running subagent IDs with their status
listRunningSubagents :: IO [(Text, SubagentRunStatus)]
listRunningSubagents = do
  registry <- readIORef subagentRegistry
  forM (Map.toList registry) <| \(sid, h) -> do
    done <- isSubagentDone h
    if done
      then pure (sid, SubagentRunStatus 0 0 0 0 "Completed" Nothing)
      else do
        status <- querySubagentStatus h
        pure (sid, status)

-- | Remove completed subagents from registry
cleanupRegistry :: IO ()
cleanupRegistry = do
  registry <- readIORef subagentRegistry
  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,
    pendingSubagentId :: AuditLog.SubagentId,
    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
-- Returns (pendingId, subagentId) - the subagentId is pre-generated so agent can track it
createPendingSpawn :: SubagentConfig -> Int -> IO (Text, Text)
createPendingSpawn config chatId = do
  uuid <- Data.UUID.V4.nextRandom
  let pid = Text.take 8 (Data.UUID.toText uuid)
  subagentId <- AuditLog.newSubagentId
  now <- Clock.getCurrentTime
  let pending =
        PendingSpawn
          { pendingId = pid,
            pendingSubagentId = subagentId,
            pendingConfig = config,
            pendingChatId = chatId,
            pendingCreatedAt = now
          }
  modifyIORef' pendingSpawnRegistry (Map.insert pid pending)
  pure (pid, AuditLog.unSubagentId subagentId)

-- | 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

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Subagent"
    [ Test.unit "SubagentRole JSON roundtrip" <| do
        let roles = [WebCrawler, CodeReviewer, DataExtractor, Researcher]
        forM_ roles <| \role ->
          case Aeson.decode (Aeson.encode role) of
            Nothing -> Test.assertFailure ("Failed to decode role: " <> show role)
            Just decoded -> decoded Test.@=? role,
      Test.unit "SubagentConfig JSON roundtrip" <| do
        let cfg = defaultSubagentConfig WebCrawler "test task"
        case Aeson.decode (Aeson.encode cfg) of
          Nothing -> Test.assertFailure "Failed to decode SubagentConfig"
          Just decoded -> subagentTask decoded Test.@=? "test task",
      Test.unit "SubagentResult JSON roundtrip" <| do
        let result =
              SubagentResult
                { subagentOutput = Aeson.object ["data" .= ("test" :: Text)],
                  subagentSummary = "Test summary",
                  subagentConfidence = 0.85,
                  subagentTokensUsed = 1000,
                  subagentCostCents = 0.5,
                  subagentDuration = 30,
                  subagentIterations = 3,
                  subagentStatus = SubagentSuccess
                }
        case Aeson.decode (Aeson.encode result) of
          Nothing -> Test.assertFailure "Failed to decode SubagentResult"
          Just decoded -> subagentSummary decoded Test.@=? "Test summary",
      Test.unit "SubagentStatus JSON roundtrip" <| do
        let statuses =
              [ SubagentSuccess,
                SubagentTimeout,
                SubagentCostExceeded,
                SubagentError "test error"
              ]
        forM_ statuses <| \status ->
          case Aeson.decode (Aeson.encode status) of
            Nothing -> Test.assertFailure ("Failed to decode status: " <> show status)
            Just decoded -> decoded Test.@=? status,
      Test.unit "toolsForRole WebCrawler has web tools" <| do
        let keys = SubagentApiKeys "test-openrouter-key" (Just "test-kagi-key")
        let tools = toolsForRole WebCrawler keys
        let names = map Engine.toolName tools
        ("web_search" `elem` names) Test.@=? True
        ("read_webpages" `elem` names) Test.@=? True,
      Test.unit "toolsForRole CodeReviewer has code tools" <| do
        let keys = SubagentApiKeys "test-openrouter-key" Nothing
        let tools = toolsForRole CodeReviewer keys
        let names = map Engine.toolName tools
        ("read_file" `elem` names) Test.@=? True
        ("search_codebase" `elem` names) Test.@=? True,
      Test.unit "modelForRole returns appropriate models" <| do
        modelForRole WebCrawler Test.@=? "anthropic/claude-3-haiku"
        modelForRole CodeReviewer Test.@=? "anthropic/claude-sonnet-4"
        modelForRole Researcher Test.@=? "anthropic/claude-sonnet-4",
      Test.unit "defaultSubagentConfig has sensible defaults" <| do
        let cfg = defaultSubagentConfig WebCrawler "task"
        subagentTimeout cfg Test.@=? 600
        subagentMaxCost cfg Test.@=? 100.0
        subagentMaxTokens cfg Test.@=? 200000
        subagentMaxIterations cfg Test.@=? 20,
      Test.unit "spawnSubagentTool has correct name" <| do
        let keys = SubagentApiKeys "test-openrouter-key" (Just "test-kagi-key")
        let tool = spawnSubagentTool keys
        Engine.toolName tool Test.@=? "spawn_subagent",
      Test.unit "spawn_subagent returns approval request when not confirmed" <| do
        let keys = SubagentApiKeys "test-openrouter-key" (Just "test-kagi-key")
        let tool = spawnSubagentTool keys
        let args = Aeson.object ["role" .= ("web_crawler" :: Text), "task" .= ("test task" :: Text)]
        result <- Engine.toolExecute tool args
        case result of
          Aeson.Object obj -> do
            let status = KeyMap.lookup "status" obj
            status Test.@=? Just (Aeson.String "awaiting_approval")
          _ -> Test.assertFailure "Expected object response",
      Test.unit "pending spawn create and lookup works" <| do
        let config = defaultSubagentConfig WebCrawler "test pending task"
        (pid, sid) <- createPendingSpawn config 12345
        when (Text.null pid) <| Test.assertFailure "pending ID should not be empty"
        when (Text.null sid) <| Test.assertFailure "subagent ID should not be empty"
        maybePending <- getPendingSpawn pid
        case maybePending of
          Nothing -> Test.assertFailure "Pending spawn not found after creation"
          Just p -> do
            pendingChatId p Test.@=? 12345
            subagentTask (pendingConfig p) Test.@=? "test pending task"
            AuditLog.unSubagentId (pendingSubagentId p) Test.@=? sid
        removePendingSpawn pid
        afterRemove <- getPendingSpawn pid
        afterRemove Test.@=? Nothing,
      Test.unit "pending spawn registry is isolated" <| do
        let config = defaultSubagentConfig Researcher "isolated test"
        (pid1, _) <- createPendingSpawn config 111
        (pid2, _) <- createPendingSpawn config 222
        when (pid1 == pid2) <| Test.assertFailure "IDs should be different"
        p1 <- getPendingSpawn pid1
        p2 <- getPendingSpawn pid2
        when (isNothing p1 || isNothing p2) <| Test.assertFailure "Both should exist"
        removePendingSpawn pid1
        removePendingSpawn pid2
    ]

data SubagentRole
  = WebCrawler
  | CodeReviewer
  | DataExtractor
  | Researcher
  | Coder
  | CustomRole Text
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SubagentRole where
  toJSON WebCrawler = Aeson.String "web_crawler"
  toJSON CodeReviewer = Aeson.String "code_reviewer"
  toJSON DataExtractor = Aeson.String "data_extractor"
  toJSON Researcher = Aeson.String "researcher"
  toJSON Coder = Aeson.String "coder"
  toJSON (CustomRole name) = Aeson.String name

instance Aeson.FromJSON SubagentRole where
  parseJSON = Aeson.withText "SubagentRole" parseRole
    where
      parseRole "web_crawler" = pure WebCrawler
      parseRole "code_reviewer" = pure CodeReviewer
      parseRole "data_extractor" = pure DataExtractor
      parseRole "researcher" = pure Researcher
      parseRole "coder" = pure Coder
      parseRole name = pure (CustomRole name)

data SubagentConfig = SubagentConfig
  { subagentRole :: SubagentRole,
    subagentTask :: Text,
    subagentModel :: Maybe Text,
    subagentTimeout :: Int,
    subagentMaxCost :: Double,
    subagentMaxTokens :: Int,
    subagentMaxIterations :: Int,
    subagentExtendedThinking :: Bool,
    subagentContext :: Maybe Text,
    -- | Optional task ID for tracking (not used by Coder)
    subagentTaskId :: Maybe Text,
    -- | Namespace for Coder role - required (e.g., "Omni/Agent/Subagent")
    subagentNamespace :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SubagentConfig where
  toJSON c =
    Aeson.object
      <| catMaybes
        [ Just ("role" .= subagentRole c),
          Just ("task" .= subagentTask c),
          ("model" .=) </ subagentModel c,
          Just ("timeout" .= subagentTimeout c),
          Just ("max_cost_cents" .= subagentMaxCost c),
          Just ("max_tokens" .= subagentMaxTokens c),
          Just ("max_iterations" .= subagentMaxIterations c),
          Just ("extended_thinking" .= subagentExtendedThinking c),
          ("context" .=) </ subagentContext c,
          ("task_id" .=) </ subagentTaskId c,
          ("namespace" .=) </ subagentNamespace c
        ]

instance Aeson.FromJSON SubagentConfig where
  parseJSON =
    Aeson.withObject "SubagentConfig" <| \v ->
      (SubagentConfig </ (v .: "role"))
        <*> (v .: "task")
        <*> (v .:? "model")
        <*> (v .:? "timeout" .!= 600)
        <*> (v .:? "max_cost_cents" .!= 50.0)
        <*> (v .:? "max_tokens" .!= 100000)
        <*> (v .:? "max_iterations" .!= 20)
        <*> (v .:? "extended_thinking" .!= False)
        <*> (v .:? "context")
        <*> (v .:? "task_id")
        <*> (v .:? "namespace")

data SubagentResult = SubagentResult
  { subagentOutput :: Aeson.Value,
    subagentSummary :: Text,
    subagentConfidence :: Double,
    subagentTokensUsed :: Int,
    subagentCostCents :: Double,
    subagentDuration :: Int,
    subagentIterations :: Int,
    subagentStatus :: SubagentStatus
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SubagentResult

instance Aeson.FromJSON SubagentResult

data SubagentStatus
  = SubagentSuccess
  | SubagentTimeout
  | SubagentCostExceeded
  | SubagentError Text
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SubagentStatus where
  toJSON SubagentSuccess = Aeson.String "success"
  toJSON SubagentTimeout = Aeson.String "timeout"
  toJSON SubagentCostExceeded = Aeson.String "cost_exceeded"
  toJSON (SubagentError msg) = Aeson.object ["error" .= msg]

instance Aeson.FromJSON SubagentStatus where
  parseJSON (Aeson.String "success") = pure SubagentSuccess
  parseJSON (Aeson.String "timeout") = pure SubagentTimeout
  parseJSON (Aeson.String "cost_exceeded") = pure SubagentCostExceeded
  parseJSON (Aeson.Object v) = SubagentError </ (v .: "error")
  parseJSON _ = empty

data SubagentCallbacks = SubagentCallbacks
  { onSubagentStart :: Text -> IO (),
    onSubagentActivity :: Text -> IO (),
    onSubagentToolCall :: Text -> Text -> IO (),
    onSubagentComplete :: SubagentResult -> IO ()
  }

defaultCallbacks :: SubagentCallbacks
defaultCallbacks =
  SubagentCallbacks
    { onSubagentStart = \_ -> pure (),
      onSubagentActivity = \_ -> pure (),
      onSubagentToolCall = \_ _ -> pure (),
      onSubagentComplete = \_ -> pure ()
    }

data SubagentHandle = SubagentHandle
  { handleId :: AuditLog.SubagentId,
    handleAsync :: Async SubagentResult,
    handleStartTime :: Clock.UTCTime,
    handleConfig :: SubagentConfig,
    handleStatus :: TVar SubagentRunStatus
  }

data SubagentRunStatus = SubagentRunStatus
  { runIteration :: Int,
    runTokensUsed :: Int,
    runCostCents :: Double,
    runElapsedSeconds :: Int,
    runCurrentActivity :: Text,
    runLastToolCall :: Maybe (Text, Clock.UTCTime)
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SubagentRunStatus

initialRunStatus :: SubagentRunStatus
initialRunStatus =
  SubagentRunStatus
    { runIteration = 0,
      runTokensUsed = 0,
      runCostCents = 0.0,
      runElapsedSeconds = 0,
      runCurrentActivity = "Starting...",
      runLastToolCall = Nothing
    }

spawnSubagentAsync :: AuditLog.SessionId -> Maybe Text -> SubagentApiKeys -> SubagentConfig -> IO SubagentHandle
spawnSubagentAsync sessionId userId keys config =
  spawnSubagentAsyncWithId sessionId userId keys config Nothing

-- | Spawn subagent with optional pre-generated ID (for pending spawn flow)
spawnSubagentAsyncWithId :: AuditLog.SessionId -> Maybe Text -> SubagentApiKeys -> SubagentConfig -> Maybe AuditLog.SubagentId -> IO SubagentHandle
spawnSubagentAsyncWithId sessionId userId keys config maybePregenId = 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
                  ]
          }

  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
  let elapsed = round (Clock.diffUTCTime now (handleStartTime h))
  status <- readTVarIO (handleStatus h)
  pure <| status {runElapsedSeconds = elapsed}

isSubagentDone :: SubagentHandle -> IO Bool
isSubagentDone h = do
  result <- poll (handleAsync h)
  pure <| isJust result

waitSubagent :: SubagentHandle -> IO SubagentResult
waitSubagent h = wait (handleAsync h)

cancelSubagent :: SubagentHandle -> IO ()
cancelSubagent h = cancel (handleAsync h)

defaultSubagentConfig :: SubagentRole -> Text -> SubagentConfig
defaultSubagentConfig role task =
  SubagentConfig
    { subagentRole = role,
      subagentTask = task,
      subagentModel = Nothing,
      subagentTimeout = 600,
      subagentMaxCost = 100.0,
      subagentMaxTokens = 200000,
      subagentMaxIterations = 20,
      subagentExtendedThinking = False,
      subagentContext = Nothing,
      subagentTaskId = Nothing,
      subagentNamespace = Nothing
    }

modelForRole :: SubagentRole -> Text
modelForRole WebCrawler = "anthropic/claude-3-haiku"
modelForRole CodeReviewer = "anthropic/claude-sonnet-4"
modelForRole DataExtractor = "anthropic/claude-3-haiku"
modelForRole Researcher = "anthropic/claude-sonnet-4"
modelForRole Coder = "anthropic/claude-sonnet-4"
modelForRole (CustomRole _) = "anthropic/claude-sonnet-4"

data SubagentApiKeys = SubagentApiKeys
  { subagentOpenRouterKey :: Text,
    subagentKagiKey :: Maybe Text
  }
  deriving (Show, Eq)

toolsForRole :: SubagentRole -> SubagentApiKeys -> [Engine.Tool]
toolsForRole WebCrawler keys =
  let webSearchTools = case subagentKagiKey keys of
        Just kagiKey -> [WebSearch.webSearchTool kagiKey]
        Nothing -> []
   in webSearchTools
        <> [ WebReader.webReaderTool (subagentOpenRouterKey keys),
             Tools.searchCodebaseTool
           ]
toolsForRole CodeReviewer _keys =
  [ Tools.readFileTool,
    Tools.searchCodebaseTool,
    Tools.searchAndReadTool,
    Tools.runBashTool
  ]
toolsForRole DataExtractor keys =
  [ WebReader.webReaderTool (subagentOpenRouterKey keys),
    Tools.readFileTool,
    Tools.searchCodebaseTool
  ]
toolsForRole Researcher keys =
  let webSearchTools = case subagentKagiKey keys of
        Just kagiKey -> [WebSearch.webSearchTool kagiKey]
        Nothing -> []
   in webSearchTools
        <> [ WebReader.webReaderTool (subagentOpenRouterKey keys),
             Tools.readFileTool,
             Tools.searchCodebaseTool,
             Tools.searchAndReadTool
           ]
-- Coder uses the hardened Coder module, toolsForRole not used
toolsForRole Coder _keys = Coder.coderTools
toolsForRole (CustomRole _) keys = toolsForRole Researcher keys

-- | Load system prompt from template, falling back to hardcoded if unavailable
loadSystemPromptForRole :: SubagentRole -> Text -> Maybe Text -> IO Text
loadSystemPromptForRole role task maybeContext = do
  let ctx =
        Aeson.object
          [ "role_description" .= roleDescription role,
            "task" .= task,
            "context" .= maybeContext
          ]
  result <- Prompts.renderPrompt "subagents/generic/system" ctx
  case result of
    Right prompt -> pure prompt
    Left _err -> pure (systemPromptForRole role task maybeContext)

-- | Hardcoded fallback prompt for subagents
systemPromptForRole :: SubagentRole -> Text -> Maybe Text -> Text
systemPromptForRole role task maybeContext =
  Text.unlines
    [ "You are a specialized " <> roleDescription role <> " subagent working on a focused task.",
      "",
      "## Your Task",
      task,
      "",
      maybe "" (\ctx -> "## Context from Orchestrator\n" <> ctx <> "\n") maybeContext,
      "## Guidelines",
      "1. Be EFFICIENT with context - extract only key facts, don't save full page contents",
      "2. Summarize findings as you go rather than accumulating raw data",
      "3. Limit web page reads to 3-5 most relevant sources",
      "4. Work iteratively: search → skim results → read best 2-3 → synthesize",
      "5. ALWAYS cite sources - every claim needs a URL",
      "6. Stop when you have sufficient information - don't over-research",
      "",
      "## Output Format",
      "Return findings as a list of structured insights:",
      "",
      "```json",
      "{",
      "  \"summary\": \"Brief overall summary (1-2 sentences)\",",
      "  \"confidence\": 0.85,",
      "  \"findings\": [",
      "    {",
      "      \"claim\": \"The key insight or fact discovered\",",
      "      \"source_url\": \"https://example.com/page\",",
      "      \"quote\": \"Relevant excerpt supporting the claim\",",
      "      \"source_name\": \"Example Site\"",
      "    }",
      "  ],",
      "  \"caveats\": \"Any limitations or uncertainties\"",
      "}",
      "```"
    ]

roleDescription :: SubagentRole -> Text
roleDescription WebCrawler = "web research"
roleDescription CodeReviewer = "code review"
roleDescription DataExtractor = "data extraction"
roleDescription Researcher = "research"
roleDescription Coder = "coding"
roleDescription (CustomRole name) = name

runSubagent :: SubagentApiKeys -> SubagentConfig -> IO SubagentResult
runSubagent keys config = runSubagentWithCallbacks keys config defaultCallbacks

runSubagentWithCallbacks :: SubagentApiKeys -> SubagentConfig -> SubagentCallbacks -> IO SubagentResult
runSubagentWithCallbacks keys config callbacks = do
  let role = subagentRole config

  -- Coder role uses the hardened Coder module with init/verify/commit phases
  case role of
    Coder -> runCoderSubagentWrapper keys config callbacks
    _ -> runGenericSubagent keys config callbacks

-- | Run Coder subagent using the hardened Coder module
runCoderSubagentWrapper :: SubagentApiKeys -> SubagentConfig -> SubagentCallbacks -> IO SubagentResult
runCoderSubagentWrapper keys config callbacks = do
  startTime <- Clock.getCurrentTime

  -- Validate required namespace field for Coder role
  let namespace = fromMaybe "" (subagentNamespace config)

  if Text.null namespace
    then
      pure
        SubagentResult
          { subagentOutput = Aeson.object ["error" .= ("Coder role requires namespace field" :: Text)],
            subagentSummary = "Missing required field: namespace",
            subagentConfidence = 0.0,
            subagentTokensUsed = 0,
            subagentCostCents = 0.0,
            subagentDuration = 0,
            subagentIterations = 0,
            subagentStatus = SubagentError "Missing namespace"
          }
    else do
      onSubagentStart callbacks ("Starting Coder subagent for " <> namespace <> "...")

      -- Build CoderConfig from SubagentConfig
      let coderCfg =
            Coder.CoderConfig
              { Coder.coderNamespace = namespace,
                Coder.coderTask = subagentTask config,
                Coder.coderContext = subagentContext config,
                Coder.coderModel = fromMaybe "anthropic/claude-sonnet-4" (subagentModel config),
                Coder.coderTimeout = subagentTimeout config,
                Coder.coderMaxCost = subagentMaxCost config,
                Coder.coderMaxTokens = subagentMaxTokens config,
                Coder.coderMaxIterations = subagentMaxIterations config,
                Coder.coderMaxVerifyRetries = 3
              }

      result <- Coder.runCoderSubagent (subagentOpenRouterKey keys) coderCfg

      endTime <- Clock.getCurrentTime
      let durationSecs = round (Clock.diffUTCTime endTime startTime)

      case result of
        Left err -> do
          onSubagentComplete callbacks
            <| SubagentResult
              { subagentOutput = Aeson.object ["error" .= err],
                subagentSummary = "Coder failed: " <> Text.take 200 err,
                subagentConfidence = 0.0,
                subagentTokensUsed = 0,
                subagentCostCents = 0.0,
                subagentDuration = durationSecs,
                subagentIterations = 0,
                subagentStatus = SubagentError err
              }
          pure
            SubagentResult
              { subagentOutput = Aeson.object ["error" .= err],
                subagentSummary = "Coder failed: " <> Text.take 200 err,
                subagentConfidence = 0.0,
                subagentTokensUsed = 0,
                subagentCostCents = 0.0,
                subagentDuration = durationSecs,
                subagentIterations = 0,
                subagentStatus = SubagentError err
              }
        Right jsonResult -> do
          let summary = case jsonResult of
                Aeson.Object obj -> case KeyMap.lookup "summary" obj of
                  Just (Aeson.String s) -> Text.take 200 s
                  _ -> "Coder completed successfully"
                _ -> "Coder completed successfully"
          let tokens = case jsonResult of
                Aeson.Object obj -> case KeyMap.lookup "tokens_used" obj of
                  Just (Aeson.Number n) -> round n
                  _ -> 0
                _ -> 0
          let cost = case jsonResult of
                Aeson.Object obj -> case KeyMap.lookup "cost_cents" obj of
                  Just (Aeson.Number n) -> realToFrac n
                  _ -> 0.0
                _ -> 0.0
          let iters = case jsonResult of
                Aeson.Object obj -> case KeyMap.lookup "iterations" obj of
                  Just (Aeson.Number n) -> round n
                  _ -> 0
                _ -> 0
          let finalResult =
                SubagentResult
                  { subagentOutput = jsonResult,
                    subagentSummary = summary,
                    subagentConfidence = 0.9,
                    subagentTokensUsed = tokens,
                    subagentCostCents = cost,
                    subagentDuration = durationSecs,
                    subagentIterations = iters,
                    subagentStatus = SubagentSuccess
                  }
          onSubagentComplete callbacks finalResult
          pure finalResult

-- | Run generic (non-Coder) subagent
runGenericSubagent :: SubagentApiKeys -> SubagentConfig -> SubagentCallbacks -> IO SubagentResult
runGenericSubagent keys config callbacks = do
  startTime <- Clock.getCurrentTime

  let role = subagentRole config
  let model = fromMaybe (modelForRole role) (subagentModel config)
  let tools = toolsForRole role keys
  systemPrompt <- loadSystemPromptForRole role (subagentTask config) (subagentContext config)

  onSubagentStart callbacks ("Starting " <> tshow role <> " subagent...")

  let provider = Provider.defaultOpenRouter (subagentOpenRouterKey keys) model

  let guardrails =
        Engine.Guardrails
          { Engine.guardrailMaxCostCents = subagentMaxCost config,
            Engine.guardrailMaxTokens = subagentMaxTokens config,
            Engine.guardrailMaxDuplicateToolCalls = 20,
            Engine.guardrailMaxTestFailures = 3,
            Engine.guardrailMaxEditFailures = 5
          }

  let agentConfig =
        Engine.AgentConfig
          { Engine.agentModel = model,
            Engine.agentTools = tools,
            Engine.agentSystemPrompt = systemPrompt,
            Engine.agentMaxIterations = subagentMaxIterations config,
            Engine.agentGuardrails = guardrails
          }

  let engineConfig =
        Engine.EngineConfig
          { Engine.engineLLM = Engine.defaultLLM,
            Engine.engineOnCost = \_ _ -> pure (),
            Engine.engineOnActivity = onSubagentActivity callbacks,
            Engine.engineOnToolCall = onSubagentToolCall callbacks,
            Engine.engineOnAssistant = \_ -> pure (),
            Engine.engineOnToolResult = \_ _ _ -> pure (),
            Engine.engineOnComplete = pure (),
            Engine.engineOnError = \_ -> pure (),
            Engine.engineOnGuardrail = \_ -> pure (),
            Engine.engineOnToolTrace = \_ _ _ _ -> pure Nothing
          }

  let timeoutMicros = subagentTimeout config * 1000000

  resultOrTimeout <-
    race
      (threadDelay timeoutMicros)
      (Engine.runAgentWithProvider engineConfig provider agentConfig (subagentTask config))

  endTime <- Clock.getCurrentTime
  let durationSecs = round (Clock.diffUTCTime endTime startTime)

  let result = case resultOrTimeout of
        Left () ->
          SubagentResult
            { subagentOutput = Aeson.object ["error" .= ("Timeout after " <> tshow (subagentTimeout config) <> " seconds" :: Text)],
              subagentSummary = "Subagent timed out",
              subagentConfidence = 0.0,
              subagentTokensUsed = 0,
              subagentCostCents = 0.0,
              subagentDuration = durationSecs,
              subagentIterations = 0,
              subagentStatus = SubagentTimeout
            }
        Right (Left err) ->
          let status = if "cost" `Text.isInfixOf` Text.toLower err then SubagentCostExceeded else SubagentError err
           in SubagentResult
                { subagentOutput = Aeson.object ["error" .= err],
                  subagentSummary = "Subagent failed: " <> err,
                  subagentConfidence = 0.0,
                  subagentTokensUsed = 0,
                  subagentCostCents = 0.0,
                  subagentDuration = durationSecs,
                  subagentIterations = 0,
                  subagentStatus = status
                }
        Right (Right agentResult) ->
          SubagentResult
            { subagentOutput = Aeson.object ["response" .= Engine.resultFinalMessage agentResult],
              subagentSummary = truncateSummary (Engine.resultFinalMessage agentResult),
              subagentConfidence = 0.8,
              subagentTokensUsed = Engine.resultTotalTokens agentResult,
              subagentCostCents = Engine.resultTotalCost agentResult,
              subagentDuration = durationSecs,
              subagentIterations = Engine.resultIterations agentResult,
              subagentStatus = SubagentSuccess
            }

  onSubagentComplete callbacks result
  pure result
  where
    truncateSummary :: Text -> Text
    truncateSummary txt =
      let firstLine = Text.takeWhile (/= '\n') txt
       in if Text.length firstLine > 200
            then Text.take 197 firstLine <> "..."
            else firstLine

spawnSubagentTool :: SubagentApiKeys -> Engine.Tool
spawnSubagentTool keys =
  Engine.Tool
    { Engine.toolName = "spawn_subagent",
      Engine.toolDescription =
        "Spawn a specialized subagent for a focused task. "
          <> "IMPORTANT: First call with confirmed=false to get approval request, "
          <> "then present the approval to the user. Only call with confirmed=true "
          <> "after the user explicitly approves. "
          <> "Available roles: web_crawler (fast web research), code_reviewer (thorough code analysis), "
          <> "data_extractor (structured data extraction), researcher (general research), "
          <> "coder (hardened coding with init/verify/commit - requires namespace and context).",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "role"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "enum" .= (["web_crawler", "code_reviewer", "data_extractor", "researcher", "coder"] :: [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" .= ("Background context, related files, design decisions (required for coder)" :: 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)
                      ],
                  "namespace"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Code namespace like 'Omni/Agent/Subagent' (required for coder role)" :: Text)
                      ],
                  "confirmed"
                    .= Aeson.object
                      [ "type" .= ("boolean" :: Text),
                        "description" .= ("Set to true only after user approval. First call should use false." :: Text)
                      ]
                ],
            "required" .= (["role", "task"] :: [Text])
          ],
      Engine.toolExecute = executeSpawnSubagent keys
    }

data SpawnRequest = SpawnRequest
  { spawnConfig :: SubagentConfig,
    spawnConfirmed :: Bool
  }
  deriving (Show, Eq)

instance Aeson.FromJSON SpawnRequest where
  parseJSON =
    Aeson.withObject "SpawnRequest" <| \v -> do
      config <- Aeson.parseJSON (Aeson.Object v)
      confirmed <- v .:? "confirmed" .!= False
      pure SpawnRequest {spawnConfig = config, spawnConfirmed = confirmed}

formatApprovalRequest :: SubagentConfig -> Aeson.Value
formatApprovalRequest config =
  Aeson.object
    [ "status" .= ("awaiting_approval" :: Text),
      "message" .= approvalMessage,
      "estimated_time_minutes" .= estimatedTime,
      "max_cost_cents" .= subagentMaxCost config,
      "role" .= subagentRole config,
      "task" .= subagentTask config
    ]
  where
    approvalMessage :: Text
    approvalMessage =
      "I'd like to spawn a "
        <> roleText
        <> " subagent to: "
        <> subagentTask config
        <> "\n\nEstimated: "
        <> tshow estimatedTime
        <> " minutes, up to $"
        <> costStr
        <> "\n\nProceed? (yes/no)"
    roleText = case subagentRole config of
      WebCrawler -> "WebCrawler"
      CodeReviewer -> "CodeReviewer"
      DataExtractor -> "DataExtractor"
      Researcher -> "Researcher"
      Coder -> "Coder"
      CustomRole name -> name
    estimatedTime :: Int
    estimatedTime = subagentTimeout config `div` 60
    costStr = Text.pack (printf "%.2f" (subagentMaxCost config / 100))

executeSpawnSubagent :: SubagentApiKeys -> Aeson.Value -> IO Aeson.Value
executeSpawnSubagent keys v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure <| Aeson.object ["error" .= ("Invalid arguments: " <> Text.pack e)]
    Aeson.Success req ->
      if spawnConfirmed req
        then do
          uuid <- Data.UUID.V4.nextRandom
          let sessionId = AuditLog.SessionId ("subagent-" <> Text.take 8 (Data.UUID.toText uuid))
          subHandle <- spawnSubagentAsync sessionId Nothing keys (spawnConfig req)
          registerSubagent subHandle
          let sid = AuditLog.unSubagentId (handleId subHandle)
          pure
            <| Aeson.object
              [ "status" .= ("spawned" :: Text),
                "subagent_id" .= sid,
                "message"
                  .= ( "Subagent spawned in background. Use check_subagent with id '"
                         <> sid
                         <> "' to monitor progress."
                     )
              ]
        else pure (formatApprovalRequest (spawnConfig req))

-- | Tool for checking subagent status or getting results
checkSubagentTool :: Engine.Tool
checkSubagentTool =
  Engine.Tool
    { Engine.toolName = "check_subagent",
      Engine.toolDescription =
        "Check the status of a running subagent or retrieve its result if completed. "
          <> "Pass the subagent_id returned from spawn_subagent. "
          <> "If no id is given, lists all running subagents.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "subagent_id"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("The subagent ID to check (e.g. 'abc123')" :: Text)
                      ]
                ],
            "required" .= ([] :: [Text])
          ],
      Engine.toolExecute = executeCheckSubagent
    }

executeCheckSubagent :: Aeson.Value -> IO Aeson.Value
executeCheckSubagent v = do
  let maybeId = case v of
        Aeson.Object obj -> case KeyMap.lookup "subagent_id" obj of
          Just (Aeson.String sid) -> Just sid
          _ -> Nothing
        _ -> Nothing
  case maybeId of
    Nothing -> do
      running <- listRunningSubagents
      pure
        <| Aeson.object
          [ "status" .= ("list" :: Text),
            "subagents"
              .= [ Aeson.object
                     [ "id" .= sid,
                       "activity" .= runCurrentActivity status,
                       "elapsed_seconds" .= runElapsedSeconds status,
                       "tokens" .= runTokensUsed status
                     ]
                   | (sid, status) <- running
                 ]
          ]
    Just sid -> do
      maybeHandle <- getSubagentHandle sid
      case maybeHandle of
        Nothing ->
          pure <| Aeson.object ["error" .= ("No subagent found with id: " <> sid)]
        Just h -> do
          done <- isSubagentDone h
          if done
            then do
              result <- waitSubagent h
              pure (Aeson.toJSON result)
            else do
              status <- querySubagentStatus h
              pure
                <| Aeson.object
                  [ "status" .= ("running" :: Text),
                    "subagent_id" .= sid,
                    "activity" .= runCurrentActivity status,
                    "elapsed_seconds" .= runElapsedSeconds status,
                    "iteration" .= runIteration status,
                    "tokens_used" .= runTokensUsed status,
                    "cost_cents" .= runCostCents status
                  ]

-- | 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 =
        "Request to spawn a specialized subagent for a focused task. "
          <> "The user will receive a confirmation button to approve. "
          <> "IMPORTANT: The subagent does NOT start until the user clicks Approve - "
          <> "do NOT say 'spawned' or 'started', say 'requested' or 'awaiting approval'. "
          <> "Available roles: web_crawler (fast web research), code_reviewer (thorough code analysis), "
          <> "data_extractor (structured data extraction), researcher (general research), "
          <> "coder (hardened coding with init/verify/commit - requires namespace and context).",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "role"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "enum" .= (["web_crawler", "code_reviewer", "data_extractor", "researcher", "coder"] :: [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" .= ("Background context, related files, design decisions (required for coder)" :: 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)
                      ],
                  "namespace"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Code namespace like 'Omni/Agent/Subagent' (required for coder role)" :: 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
      (pid, subagentId) <- createPendingSpawn config chatId
      let roleText = case subagentRole config of
            WebCrawler -> "web_crawler"
            CodeReviewer -> "code_reviewer"
            DataExtractor -> "data_extractor"
            Researcher -> "researcher"
            Coder -> "coder"
            CustomRole name -> name
          estimatedMins = subagentTimeout config `div` 60
          maxCost = subagentMaxCost config
      onApprovalNeeded chatId pid roleText (subagentTask config) estimatedMins maxCost
      pure
        <| Aeson.object
          [ "status" .= ("pending_approval" :: Text),
            "subagent_id" .= subagentId,
            "message" .= ("Subagent requested. User must click Approve button before it starts. Do not say spawned yet." :: Text)
          ]

-- | Approve a pending spawn and start the subagent
approveAndSpawnSubagent :: SubagentApiKeys -> Text -> IO (Either Text Text)
approveAndSpawnSubagent keys pid = do
  maybePending <- getPendingSpawn pid
  case maybePending of
    Nothing -> pure (Left "Pending spawn not found or expired")
    Just pending -> 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))
      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]