1 files changed, 516 insertions, 0 deletions
diff --git a/Omni/Agent/Subagent.hs b/Omni/Agent/Subagent.hs
new file mode 100644
index 0000000..c8e56d5
--- /dev/null
+++ b/Omni/Agent/Subagent.hs
@@ -0,0 +1,516 @@
+{-# 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)
+--
+-- : out omni-agent-subagent
+-- : dep aeson
+-- : dep async
+module Omni.Agent.Subagent
+  ( -- * Types
+    SubagentRole (..),
+    SubagentConfig (..),
+    SubagentResult (..),
+    SubagentStatus (..),
+    SubagentCallbacks (..),
+
+    -- * Execution
+    runSubagent,
+    runSubagentWithCallbacks,
+
+    -- * Tool
+    spawnSubagentTool,
+
+    -- * Role-specific tools
+    SubagentApiKeys (..),
+    toolsForRole,
+    modelForRole,
+    systemPromptForRole,
+
+    -- * Defaults
+    defaultSubagentConfig,
+    defaultCallbacks,
+
+    -- * Testing
+    main,
+    test,
+  )
+where
+
+import Alpha
+import Data.Aeson ((.!=), (.:), (.:?), (.=))
+import qualified Data.Aeson as Aeson
+import qualified Data.Text as Text
+import qualified Data.Time.Clock as Clock
+import qualified Omni.Agent.Engine as Engine
+import qualified Omni.Agent.Provider as Provider
+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
+
+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"
+    ]
+
+data SubagentRole
+  = WebCrawler
+  | CodeReviewer
+  | DataExtractor
+  | Researcher
+  | 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 (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 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
+  }
+  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
+        ]
+
+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")
+
+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 ()
+    }
+
+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
+    }
+
+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 (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
+           ]
+toolsForRole (CustomRole _) keys = toolsForRole Researcher keys
+
+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\"",
+      "}",
+      "```"
+    ]
+  where
+    roleDescription :: SubagentRole -> Text
+    roleDescription WebCrawler = "web research"
+    roleDescription CodeReviewer = "code review"
+    roleDescription DataExtractor = "data extraction"
+    roleDescription Researcher = "research"
+    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
+  startTime <- Clock.getCurrentTime
+
+  let role = subagentRole config
+  let model = fromMaybe (modelForRole role) (subagentModel config)
+  let tools = toolsForRole role keys
+  let systemPrompt = systemPromptForRole 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 ()
+          }
+
+  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. "
+          <> "Use for tasks that benefit from deep exploration, parallel execution, "
+          <> "or specialized tools. The subagent will iterate until task completion "
+          <> "or resource limits are reached. "
+          <> "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 = executeSpawnSubagent keys
+    }
+
+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 config -> do
+      result <- runSubagent keys config
+      pure (Aeson.toJSON result)