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.hs240
1 files changed, 215 insertions, 25 deletions
diff --git a/Omni/Agent/Subagent.hs b/Omni/Agent/Subagent.hs
index 3278e4c..06ef938 100644
--- a/Omni/Agent/Subagent.hs
+++ b/Omni/Agent/Subagent.hs
@@ -93,6 +93,7 @@ import qualified Data.UUID.V4
import qualified Omni.Agent.AuditLog as AuditLog
import qualified Omni.Agent.Engine as Engine
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
@@ -139,6 +140,7 @@ cleanupRegistry = do
-- | A pending spawn request awaiting user confirmation
data PendingSpawn = PendingSpawn
{ pendingId :: Text,
+ pendingSubagentId :: AuditLog.SubagentId,
pendingConfig :: SubagentConfig,
pendingChatId :: Int,
pendingCreatedAt :: Clock.UTCTime
@@ -151,20 +153,23 @@ pendingSpawnRegistry = unsafePerformIO (newIORef Map.empty)
{-# NOINLINE pendingSpawnRegistry #-}
-- | Create a new pending spawn request
-createPendingSpawn :: SubagentConfig -> Int -> IO Text
+-- 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 pendingId = Text.take 8 (Data.UUID.toText uuid)
+ let pid = Text.take 8 (Data.UUID.toText uuid)
+ subagentId <- AuditLog.newSubagentId
now <- Clock.getCurrentTime
let pending =
PendingSpawn
- { pendingId = pendingId,
+ { pendingId = pid,
+ pendingSubagentId = subagentId,
pendingConfig = config,
pendingChatId = chatId,
pendingCreatedAt = now
}
- modifyIORef' pendingSpawnRegistry (Map.insert pendingId pending)
- pure pendingId
+ modifyIORef' pendingSpawnRegistry (Map.insert pid pending)
+ pure (pid, AuditLog.unSubagentId subagentId)
-- | Get a pending spawn by ID
getPendingSpawn :: Text -> IO (Maybe PendingSpawn)
@@ -262,7 +267,32 @@ test =
Aeson.Object obj -> do
let status = KeyMap.lookup "status" obj
status Test.@=? Just (Aeson.String "awaiting_approval")
- _ -> Test.assertFailure "Expected object response"
+ _ -> 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
@@ -270,6 +300,7 @@ data SubagentRole
| CodeReviewer
| DataExtractor
| Researcher
+ | Coder
| CustomRole Text
deriving (Show, Eq, Generic)
@@ -278,6 +309,7 @@ instance Aeson.ToJSON SubagentRole where
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
@@ -287,6 +319,7 @@ instance Aeson.FromJSON SubagentRole where
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
@@ -298,7 +331,11 @@ data SubagentConfig = SubagentConfig
subagentMaxTokens :: Int,
subagentMaxIterations :: Int,
subagentExtendedThinking :: Bool,
- subagentContext :: Maybe Text
+ 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)
@@ -314,7 +351,9 @@ instance Aeson.ToJSON SubagentConfig where
Just ("max_tokens" .= subagentMaxTokens c),
Just ("max_iterations" .= subagentMaxIterations c),
Just ("extended_thinking" .= subagentExtendedThinking c),
- ("context" .=) </ subagentContext c
+ ("context" .=) </ subagentContext c,
+ ("task_id" .=) </ subagentTaskId c,
+ ("namespace" .=) </ subagentNamespace c
]
instance Aeson.FromJSON SubagentConfig where
@@ -329,6 +368,8 @@ instance Aeson.FromJSON SubagentConfig where
<*> (v .:? "max_iterations" .!= 20)
<*> (v .:? "extended_thinking" .!= False)
<*> (v .:? "context")
+ <*> (v .:? "task_id")
+ <*> (v .:? "namespace")
data SubagentResult = SubagentResult
{ subagentOutput :: Aeson.Value,
@@ -414,8 +455,13 @@ initialRunStatus =
}
spawnSubagentAsync :: AuditLog.SessionId -> Maybe Text -> SubagentApiKeys -> SubagentConfig -> IO SubagentHandle
-spawnSubagentAsync sessionId userId keys config = do
- sid <- AuditLog.newSubagentId
+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
@@ -512,7 +558,9 @@ defaultSubagentConfig role task =
subagentMaxTokens = 200000,
subagentMaxIterations = 20,
subagentExtendedThinking = False,
- subagentContext = Nothing
+ subagentContext = Nothing,
+ subagentTaskId = Nothing,
+ subagentNamespace = Nothing
}
modelForRole :: SubagentRole -> Text
@@ -520,6 +568,7 @@ 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
@@ -558,6 +607,8 @@ toolsForRole Researcher keys =
Tools.searchCodebaseTool,
Tools.searchAndReadTool
]
+-- Coder uses the hardened Coder module, toolsForRole not used
+toolsForRole Coder _keys = Coder.coderTools
toolsForRole (CustomRole _) keys = toolsForRole Researcher keys
systemPromptForRole :: SubagentRole -> Text -> Maybe Text -> Text
@@ -602,6 +653,7 @@ systemPromptForRole role task maybeContext =
roleDescription CodeReviewer = "code review"
roleDescription DataExtractor = "data extraction"
roleDescription Researcher = "research"
+ roleDescription Coder = "coding"
roleDescription (CustomRole name) = name
runSubagent :: SubagentApiKeys -> SubagentConfig -> IO SubagentResult
@@ -609,6 +661,118 @@ 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
@@ -717,7 +881,8 @@ spawnSubagentTool keys =
<> "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).",
+ <> "data_extractor (structured data extraction), researcher (general research), "
+ <> "coder (hardened coding with init/verify/commit phases - requires task_id and namespace).",
Engine.toolJsonSchema =
Aeson.object
[ "type" .= ("object" :: Text),
@@ -726,7 +891,7 @@ spawnSubagentTool keys =
[ "role"
.= Aeson.object
[ "type" .= ("string" :: Text),
- "enum" .= (["web_crawler", "code_reviewer", "data_extractor", "researcher"] :: [Text]),
+ "enum" .= (["web_crawler", "code_reviewer", "data_extractor", "researcher", "coder"] :: [Text]),
"description" .= ("Subagent role determining tools and model" :: Text)
],
"task"
@@ -754,6 +919,16 @@ spawnSubagentTool keys =
[ "type" .= ("number" :: Text),
"description" .= ("Maximum cost in cents (default: 50)" :: Text)
],
+ "task_id"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("Task ID from jr task (required for coder role)" :: 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),
@@ -805,6 +980,7 @@ formatApprovalRequest config =
CodeReviewer -> "CodeReviewer"
DataExtractor -> "DataExtractor"
Researcher -> "Researcher"
+ Coder -> "Coder"
CustomRole name -> name
estimatedTime :: Int
estimatedTime = subagentTimeout config `div` 60
@@ -920,10 +1096,13 @@ 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. "
+ "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).",
+ <> "data_extractor (structured data extraction), researcher (general research), "
+ <> "coder (hardened coding with init/verify/commit - requires task_id and namespace).",
Engine.toolJsonSchema =
Aeson.object
[ "type" .= ("object" :: Text),
@@ -932,7 +1111,7 @@ spawnSubagentToolWithApproval keys chatId onApprovalNeeded =
[ "role"
.= Aeson.object
[ "type" .= ("string" :: Text),
- "enum" .= (["web_crawler", "code_reviewer", "data_extractor", "researcher"] :: [Text]),
+ "enum" .= (["web_crawler", "code_reviewer", "data_extractor", "researcher", "coder"] :: [Text]),
"description" .= ("Subagent role determining tools and model" :: Text)
],
"task"
@@ -959,6 +1138,16 @@ spawnSubagentToolWithApproval keys chatId onApprovalNeeded =
.= Aeson.object
[ "type" .= ("number" :: Text),
"description" .= ("Maximum cost in cents (default: 50)" :: Text)
+ ],
+ "task_id"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("Task ID from jr task (required for coder role)" :: Text)
+ ],
+ "namespace"
+ .= Aeson.object
+ [ "type" .= ("string" :: Text),
+ "description" .= ("Code namespace like 'Omni/Agent/Subagent' (required for coder role)" :: Text)
]
],
"required" .= (["role", "task"] :: [Text])
@@ -971,34 +1160,35 @@ 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
+ (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 pendingId roleText (subagentTask config) estimatedMins maxCost
+ onApprovalNeeded chatId pid 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)
+ "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 pendingId = do
- maybePending <- getPendingSpawn pendingId
+approveAndSpawnSubagent keys pid = do
+ maybePending <- getPendingSpawn pid
case maybePending of
Nothing -> pure (Left "Pending spawn not found or expired")
Just pending -> do
- removePendingSpawn pendingId
+ removePendingSpawn pid
uuid <- Data.UUID.V4.nextRandom
let sessionId = AuditLog.SessionId ("subagent-" <> Text.take 8 (Data.UUID.toText uuid))
- subHandle <- spawnSubagentAsync sessionId Nothing keys (pendingConfig pending)
+ subHandle <- spawnSubagentAsyncWithId sessionId Nothing keys (pendingConfig pending) (Just (pendingSubagentId pending))
registerSubagent subHandle
let sid = AuditLog.unSubagentId (handleId subHandle)
pure (Right sid)