diff options
Diffstat (limited to 'Omni/Agent')
| -rw-r--r-- | Omni/Agent/Core.hs | 37 | ||||
| -rw-r--r-- | Omni/Agent/DESIGN.md | 117 | ||||
| -rw-r--r-- | Omni/Agent/Engine.hs | 582 | ||||
| -rw-r--r-- | Omni/Agent/Git.hs | 232 | ||||
| -rw-r--r-- | Omni/Agent/Log.hs | 154 | ||||
| -rw-r--r-- | Omni/Agent/Tools.hs | 582 | ||||
| -rw-r--r-- | Omni/Agent/Worker.hs | 446 |
7 files changed, 2150 insertions, 0 deletions
diff --git a/Omni/Agent/Core.hs b/Omni/Agent/Core.hs new file mode 100644 index 0000000..88f7237 --- /dev/null +++ b/Omni/Agent/Core.hs @@ -0,0 +1,37 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE NoImplicitPrelude #-} + +module Omni.Agent.Core where + +import Alpha +import Data.Aeson (FromJSON, ToJSON) + +-- | Status of a worker agent +data WorkerStatus + = Idle + | Syncing + | -- | Task ID + Working Text + | -- | Task ID + Submitting Text + | -- | Error message + Error Text + deriving (Show, Eq, Generic) + +instance ToJSON WorkerStatus + +instance FromJSON WorkerStatus + +-- | Representation of a worker agent +data Worker = Worker + { workerName :: Text, + workerPid :: Maybe Int, + workerStatus :: WorkerStatus, + workerPath :: FilePath, + workerQuiet :: Bool -- Disable ANSI status bar (for loop mode) + } + deriving (Show, Eq, Generic) + +instance ToJSON Worker + +instance FromJSON Worker diff --git a/Omni/Agent/DESIGN.md b/Omni/Agent/DESIGN.md new file mode 100644 index 0000000..ae1f6b3 --- /dev/null +++ b/Omni/Agent/DESIGN.md @@ -0,0 +1,117 @@ +# Multi-Agent System 2.0 Design + +**Goal:** Replace the current bash-script based worker system (`start-worker.sh`, etc.) with a robust, type-safe Haskell application `Omni/Agent.hs`. + +## 1. CLI Interface + +The `agent` command (compiled from `Omni/Agent.hs`) will provide a unified interface for managing workers. + +```bash +agent start <name> [--background] # Start a worker (foreground by default, background with flag) +agent stop <name> # Stop a background worker +agent status # List all workers and their status +agent log <name> [-f] # View/tail worker logs +agent sync # Sync local state with live (helper) +``` + +## 2. Module Structure (`Omni/Agent/`) + +We will refactor the bash logic into Haskell modules: + +- **Omni.Agent** (`Omni/Agent.hs`): Main entry point and CLI parsing (Docopt). +- **Omni.Agent.Core**: Core data types and state management. +- **Omni.Agent.Worker**: The worker loop logic (sync, claim, work, submit). +- **Omni.Agent.Git**: Git operations (worktree, branch, merge, commit). +- **Omni.Agent.Process**: Process management (PID files, signals). +- **Omni.Agent.Log**: Log streaming and filtering (the "monitor" logic). + +## 3. Data Types + +```haskell +data WorkerStatus + = Idle + | Syncing + | Working TaskId + | Submitting TaskId + | Error Text + deriving (Show, Eq, Generic) + +data Worker = Worker + { workerName :: Text + , workerPid :: Maybe Int + , workerStatus :: WorkerStatus + , workerPath :: FilePath + } +``` + +## 4. Implementation Details + +### 4.1 Worker Loop (`agent start`) +The Haskell implementation should replicate the logic of `start-worker.sh` but with better error handling and logging. + +1. **Setup**: Ensure worktree exists (or create it). +2. **Loop**: + - `Git.syncWithLive` (Rebase-based to preserve local history) + - `Task.sync` + - `task <- Task.findReady` + - If `task`: + - `Task.claim task` + - `baseBranch <- Git.determineBaseBranch task` (Check dependencies) + - `Git.checkoutTaskBranch task baseBranch` (Force checkout to clean untracked files) + - `Engine.runAgent prompt` (Native LLM agent via OpenRouter) + - `Git.commit` + - `Git.checkoutBase` + - `Task.submitReview task` + - Else: `sleep 60` + +### 4.2 Process Management +- Store PIDs in `.tasks/workers/<name>.pid`. +- `agent stop` sends SIGTERM to the PID. +- `agent status` checks if PID is alive. + +### 4.3 Logging +- The Engine module uses callbacks to report activity and costs in real-time. +- `agent log` displays the status bar with worker progress information. +- **UI Design**: + - **Two-line Status**: The CLI should maintain two reserved lines at the bottom (or top) of the output for each worker: + - **Line 1 (Meta)**: `[Worker: omni-worker-1] Task: t-123 | Files: 3 | Credits: $0.45 | Time: 05:23` + - **Line 2 (Activity)**: `[14:05:22] 🤖 Thinking...` (updates in place) + - **Task Details**: When claiming a task, print the full task description/details to the log/console so the user can see what is being worked on without looking it up. + - **Completion**: When a task finishes, print a summary line (e.g., `[✓] Task t-123 completed in 12m 30s`) and a hard line break before starting the next loop. + - **History**: Previous log lines (tool outputs, thoughts) scroll up above these two status lines. + +### 4.5 Git Robustness (Learnings) +- **Identity**: Configure `git config user.name "Omni Worker"` and `user.email` in the worktree to clearly distinguish worker commits from human commits. +- **Force Checkout**: The worker must use `git checkout -f` (or equivalent) when switching to task branches to ensure untracked files (like `.tasks/counters.jsonl`) don't block the switch. +- **Base Branch Logic**: + - If the task depends on another task that is *not* yet in `live` (e.g., in `Review`), the worker should branch off the dependency's branch (`task/<dep-id>`). + - Otherwise, branch off `live` directly. Do NOT use the local worker branch (`omni-worker-N`) as the base, as it may contain temporary sync commits that shouldn't be merged. +- **Commit Hygiene**: Bundle the task status update (marking as 'Review') *inside* the feature implementation commit. This keeps the history clean (one commit per feature) and avoids separate "sync" commits for status changes. +- **Clean State**: The worker should ensure the workspace is clean (no uncommitted changes) before starting a new loop iteration. +- **Rebase Safety**: Always check the exit code of `git rebase`. If it fails (conflicts), abort immediately (`git rebase --abort`) to avoid leaving the repo in a broken interactive rebase state. +- **Status Verification**: Verify that task status updates actually succeed. Check `task ready` output against `live` state to prevent "zombie" tasks (completed in live but stuck in local loop) from being re-claimed. +- **Binary Freshness**: Ensure the `task` binary used by the worker is rebuilt/updated when source code changes, otherwise logic fixes (like `task ready` filtering) won't take effect. + +## 5. Migration Strategy + +1. **Parallel Existence**: Keep bash scripts while developing Haskell version. +2. **Feature Parity**: Ensure `agent start` works exactly like `start-worker.sh`. +3. **Cutover**: Update `WORKER_AGENT_GUIDE.md` to use `agent` command. +4. **Cleanup**: Delete bash scripts. + +## 6. Testing Plan + +### 6.1 Unit Tests (`Omni/Agent/Test.hs`) +- Test `Git` module commands (mocked). +- Test `Log` filtering logic. +- Test CLI argument parsing. + +### 6.2 Integration Tests +- Create a temporary test repo. +- Spawn a worker. +- Mock the Engine LLM calls or use a test API key. +- Verify task moves from Open -> InProgress -> Review. + +## 7. References +- `Omni/Agent/start-worker.sh` (Current implementation) +- `Omni/Task.hs` (Task manager integration) diff --git a/Omni/Agent/Engine.hs b/Omni/Agent/Engine.hs new file mode 100644 index 0000000..e019341 --- /dev/null +++ b/Omni/Agent/Engine.hs @@ -0,0 +1,582 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | LLM Agent Engine - Tool protocol and LLM provider abstraction. +-- +-- This module provides the core abstractions for building LLM-powered agents: +-- - Tool: Defines tools that agents can use +-- - LLM: OpenAI-compatible chat completions API provider +-- - AgentConfig: Configuration for running agents +-- +-- : out omni-agent-engine +-- : dep http-conduit +-- : dep aeson +-- : dep case-insensitive +module Omni.Agent.Engine + ( Tool (..), + LLM (..), + EngineConfig (..), + AgentConfig (..), + AgentResult (..), + Message (..), + Role (..), + ToolCall (..), + FunctionCall (..), + ToolResult (..), + ChatCompletionRequest (..), + ChatCompletionResponse (..), + Choice (..), + Usage (..), + defaultLLM, + defaultEngineConfig, + defaultAgentConfig, + chat, + runAgent, + main, + test, + ) +where + +import Alpha +import Data.Aeson ((.!=), (.:), (.:?), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.ByteString.Lazy as BL +import qualified Data.CaseInsensitive as CI +import qualified Data.Map.Strict as Map +import qualified Data.Text as Text +import qualified Data.Text.Encoding as TE +import qualified Network.HTTP.Simple as HTTP +import qualified Omni.Test as Test + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Engine" + [ Test.unit "Tool JSON roundtrip" <| do + let tool = + Tool + { toolName = "get_weather", + toolDescription = "Get weather for a location", + toolJsonSchema = Aeson.object ["type" .= ("object" :: Text), "properties" .= Aeson.object []], + toolExecute = \_ -> pure (Aeson.String "sunny") + } + let encoded = encodeToolForApi tool + case Aeson.decode (Aeson.encode encoded) of + Nothing -> Test.assertFailure "Failed to decode tool" + Just decoded -> toolName tool Test.@=? toolApiName decoded, + Test.unit "Message JSON roundtrip" <| do + let msg = Message User "Hello" Nothing Nothing + case Aeson.decode (Aeson.encode msg) of + Nothing -> Test.assertFailure "Failed to decode message" + Just decoded -> msgContent msg Test.@=? msgContent decoded, + Test.unit "defaultLLM has correct endpoint" <| do + llmBaseUrl defaultLLM Test.@=? "https://openrouter.ai/api/v1", + Test.unit "defaultLLM has OpenRouter headers" <| do + length (llmExtraHeaders defaultLLM) Test.@=? 2 + llmModel defaultLLM Test.@=? "anthropic/claude-sonnet-4.5", + Test.unit "defaultAgentConfig has sensible defaults" <| do + agentMaxIterations defaultAgentConfig Test.@=? 10, + Test.unit "defaultEngineConfig has no-op callbacks" <| do + engineOnCost defaultEngineConfig 100 5 + engineOnActivity defaultEngineConfig "test" + engineOnToolCall defaultEngineConfig "tool" "result" + True Test.@=? True, + Test.unit "buildToolMap creates correct map" <| do + let tool1 = + Tool + { toolName = "tool1", + toolDescription = "First tool", + toolJsonSchema = Aeson.object [], + toolExecute = \_ -> pure Aeson.Null + } + tool2 = + Tool + { toolName = "tool2", + toolDescription = "Second tool", + toolJsonSchema = Aeson.object [], + toolExecute = \_ -> pure Aeson.Null + } + toolMap = buildToolMap [tool1, tool2] + Map.size toolMap Test.@=? 2 + Map.member "tool1" toolMap Test.@=? True + Map.member "tool2" toolMap Test.@=? True, + Test.unit "Usage JSON parsing" <| do + let json = "{\"prompt_tokens\":100,\"completion_tokens\":50,\"total_tokens\":150}" + case Aeson.decode json of + Nothing -> Test.assertFailure "Failed to decode usage" + Just usage -> do + usagePromptTokens usage Test.@=? 100 + usageCompletionTokens usage Test.@=? 50 + usageTotalTokens usage Test.@=? 150, + Test.unit "AgentResult JSON roundtrip" <| do + let result = + AgentResult + { resultFinalMessage = "Done", + resultToolCallCount = 3, + resultIterations = 2, + resultTotalCost = 50, + resultTotalTokens = 1500 + } + case Aeson.decode (Aeson.encode result) of + Nothing -> Test.assertFailure "Failed to decode AgentResult" + Just decoded -> do + resultFinalMessage decoded Test.@=? "Done" + resultToolCallCount decoded Test.@=? 3 + resultIterations decoded Test.@=? 2, + Test.unit "estimateCost calculates correctly" <| do + let gpt4oCost = estimateCost "gpt-4o" 1000 + gpt4oMiniCost = estimateCost "gpt-4o-mini" 1000 + (gpt4oCost >= gpt4oMiniCost) Test.@=? True, + Test.unit "ToolCall JSON roundtrip" <| do + let tc = + ToolCall + { tcId = "call_123", + tcType = "function", + tcFunction = FunctionCall "read_file" "{\"path\":\"/tmp/test\"}" + } + case Aeson.decode (Aeson.encode tc) of + Nothing -> Test.assertFailure "Failed to decode ToolCall" + Just decoded -> tcId decoded Test.@=? "call_123", + Test.unit "FunctionCall JSON roundtrip" <| do + let fc = FunctionCall "test_func" "{\"arg\":\"value\"}" + case Aeson.decode (Aeson.encode fc) of + Nothing -> Test.assertFailure "Failed to decode FunctionCall" + Just decoded -> do + fcName decoded Test.@=? "test_func" + fcArguments decoded Test.@=? "{\"arg\":\"value\"}", + Test.unit "Role JSON roundtrip for all roles" <| do + let roles = [System, User, Assistant, ToolRole] + forM_ roles <| \role -> + case Aeson.decode (Aeson.encode role) of + Nothing -> Test.assertFailure ("Failed to decode Role: " <> show role) + Just decoded -> decoded Test.@=? role + ] + +data Tool = Tool + { toolName :: Text, + toolDescription :: Text, + toolJsonSchema :: Aeson.Value, + toolExecute :: Aeson.Value -> IO Aeson.Value + } + +data ToolApi = ToolApi + { toolApiName :: Text, + toolApiDescription :: Text, + toolApiParameters :: Aeson.Value + } + deriving (Generic) + +instance Aeson.ToJSON ToolApi where + toJSON t = + Aeson.object + [ "type" .= ("function" :: Text), + "function" + .= Aeson.object + [ "name" .= toolApiName t, + "description" .= toolApiDescription t, + "parameters" .= toolApiParameters t + ] + ] + +instance Aeson.FromJSON ToolApi where + parseJSON = + Aeson.withObject "ToolApi" <| \v -> do + fn <- v .: "function" + (ToolApi </ (fn .: "name")) + <*> (fn .: "description") + <*> (fn .: "parameters") + +encodeToolForApi :: Tool -> ToolApi +encodeToolForApi t = + ToolApi + { toolApiName = toolName t, + toolApiDescription = toolDescription t, + toolApiParameters = toolJsonSchema t + } + +data LLM = LLM + { llmBaseUrl :: Text, + llmApiKey :: Text, + llmModel :: Text, + llmExtraHeaders :: [(ByteString, ByteString)] + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON LLM where + toJSON l = + Aeson.object + [ "llmBaseUrl" .= llmBaseUrl l, + "llmApiKey" .= llmApiKey l, + "llmModel" .= llmModel l + ] + +instance Aeson.FromJSON LLM where + parseJSON = + Aeson.withObject "LLM" <| \v -> + (LLM </ (v .: "llmBaseUrl")) + <*> (v .: "llmApiKey") + <*> (v .: "llmModel") + <*> pure [] + +defaultLLM :: LLM +defaultLLM = + LLM + { llmBaseUrl = "https://openrouter.ai/api/v1", + llmApiKey = "", + llmModel = "anthropic/claude-sonnet-4.5", + llmExtraHeaders = + [ ("HTTP-Referer", "https://omni.dev"), + ("X-Title", "Omni Agent") + ] + } + +data AgentConfig = AgentConfig + { agentModel :: Text, + agentTools :: [Tool], + agentSystemPrompt :: Text, + agentMaxIterations :: Int + } + +defaultAgentConfig :: AgentConfig +defaultAgentConfig = + AgentConfig + { agentModel = "gpt-4", + agentTools = [], + agentSystemPrompt = "You are a helpful assistant.", + agentMaxIterations = 10 + } + +data EngineConfig = EngineConfig + { engineLLM :: LLM, + engineOnCost :: Int -> Int -> IO (), + engineOnActivity :: Text -> IO (), + engineOnToolCall :: Text -> Text -> IO () + } + +defaultEngineConfig :: EngineConfig +defaultEngineConfig = + EngineConfig + { engineLLM = defaultLLM, + engineOnCost = \_ _ -> pure (), + engineOnActivity = \_ -> pure (), + engineOnToolCall = \_ _ -> pure () + } + +data AgentResult = AgentResult + { resultFinalMessage :: Text, + resultToolCallCount :: Int, + resultIterations :: Int, + resultTotalCost :: Int, + resultTotalTokens :: Int + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON AgentResult + +instance Aeson.FromJSON AgentResult + +data Role = System | User | Assistant | ToolRole + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON Role where + toJSON System = Aeson.String "system" + toJSON User = Aeson.String "user" + toJSON Assistant = Aeson.String "assistant" + toJSON ToolRole = Aeson.String "tool" + +instance Aeson.FromJSON Role where + parseJSON = Aeson.withText "Role" parseRole + where + parseRole "system" = pure System + parseRole "user" = pure User + parseRole "assistant" = pure Assistant + parseRole "tool" = pure ToolRole + parseRole _ = empty + +data Message = Message + { msgRole :: Role, + msgContent :: Text, + msgToolCalls :: Maybe [ToolCall], + msgToolCallId :: Maybe Text + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON Message where + toJSON m = + Aeson.object + <| catMaybes + [ Just ("role" .= msgRole m), + Just ("content" .= msgContent m), + ("tool_calls" .=) </ msgToolCalls m, + ("tool_call_id" .=) </ msgToolCallId m + ] + +instance Aeson.FromJSON Message where + parseJSON = + Aeson.withObject "Message" <| \v -> + (Message </ (v .: "role")) + <*> (v .:? "content" .!= "") + <*> (v .:? "tool_calls") + <*> (v .:? "tool_call_id") + +data ToolCall = ToolCall + { tcId :: Text, + tcType :: Text, + tcFunction :: FunctionCall + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON ToolCall where + toJSON tc = + Aeson.object + [ "id" .= tcId tc, + "type" .= tcType tc, + "function" .= tcFunction tc + ] + +instance Aeson.FromJSON ToolCall where + parseJSON = + Aeson.withObject "ToolCall" <| \v -> + (ToolCall </ (v .: "id")) + <*> (v .:? "type" .!= "function") + <*> (v .: "function") + +data FunctionCall = FunctionCall + { fcName :: Text, + fcArguments :: Text + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON FunctionCall where + toJSON fc = + Aeson.object + [ "name" .= fcName fc, + "arguments" .= fcArguments fc + ] + +instance Aeson.FromJSON FunctionCall where + parseJSON = + Aeson.withObject "FunctionCall" <| \v -> + (FunctionCall </ (v .: "name")) + <*> (v .: "arguments") + +data ToolResult = ToolResult + { trToolCallId :: Text, + trContent :: Text + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON ToolResult + +instance Aeson.FromJSON ToolResult + +data ChatCompletionRequest = ChatCompletionRequest + { reqModel :: Text, + reqMessages :: [Message], + reqTools :: Maybe [ToolApi] + } + deriving (Generic) + +instance Aeson.ToJSON ChatCompletionRequest where + toJSON r = + Aeson.object + <| catMaybes + [ Just ("model" .= reqModel r), + Just ("messages" .= reqMessages r), + ("tools" .=) </ reqTools r + ] + +data Choice = Choice + { choiceIndex :: Int, + choiceMessage :: Message, + choiceFinishReason :: Maybe Text + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON Choice where + parseJSON = + Aeson.withObject "Choice" <| \v -> + (Choice </ (v .: "index")) + <*> (v .: "message") + <*> (v .:? "finish_reason") + +data Usage = Usage + { usagePromptTokens :: Int, + usageCompletionTokens :: Int, + usageTotalTokens :: Int + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON Usage where + parseJSON = + Aeson.withObject "Usage" <| \v -> + (Usage </ (v .: "prompt_tokens")) + <*> (v .: "completion_tokens") + <*> (v .: "total_tokens") + +data ChatCompletionResponse = ChatCompletionResponse + { respId :: Text, + respChoices :: [Choice], + respModel :: Text, + respUsage :: Maybe Usage + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON ChatCompletionResponse where + parseJSON = + Aeson.withObject "ChatCompletionResponse" <| \v -> + (ChatCompletionResponse </ (v .: "id")) + <*> (v .: "choices") + <*> (v .: "model") + <*> (v .:? "usage") + +data ChatResult = ChatResult + { chatMessage :: Message, + chatUsage :: Maybe Usage + } + deriving (Show, Eq) + +chatWithUsage :: LLM -> [Tool] -> [Message] -> IO (Either Text ChatResult) +chatWithUsage llm tools messages = do + let url = Text.unpack (llmBaseUrl llm) <> "/chat/completions" + req0 <- HTTP.parseRequest url + let toolApis = [encodeToolForApi t | not (null tools), t <- tools] + body = + ChatCompletionRequest + { reqModel = llmModel llm, + reqMessages = messages, + reqTools = if null toolApis then Nothing else Just toolApis + } + baseReq = + HTTP.setRequestMethod "POST" + <| HTTP.setRequestHeader "Content-Type" ["application/json"] + <| HTTP.setRequestHeader "Authorization" ["Bearer " <> TE.encodeUtf8 (llmApiKey llm)] + <| HTTP.setRequestBodyLBS (Aeson.encode body) + <| req0 + req = foldr addHeader baseReq (llmExtraHeaders llm) + addHeader (name, value) = HTTP.addRequestHeader (CI.mk name) value + + response <- HTTP.httpLBS req + let status = HTTP.getResponseStatusCode response + if status >= 200 && status < 300 + then case Aeson.decode (HTTP.getResponseBody response) of + Just resp -> + case respChoices resp of + (c : _) -> pure (Right (ChatResult (choiceMessage c) (respUsage resp))) + [] -> pure (Left "No choices in response") + Nothing -> pure (Left "Failed to parse response") + else pure (Left ("HTTP error: " <> tshow status <> " - " <> TE.decodeUtf8 (BL.toStrict (HTTP.getResponseBody response)))) + +chat :: LLM -> [Tool] -> [Message] -> IO (Either Text Message) +chat llm tools messages = do + result <- chatWithUsage llm tools messages + pure (chatMessage </ result) + +runAgent :: EngineConfig -> AgentConfig -> Text -> IO (Either Text AgentResult) +runAgent engineCfg agentCfg userPrompt = do + let llm = + (engineLLM engineCfg) + { llmModel = agentModel agentCfg + } + tools = agentTools agentCfg + toolMap = buildToolMap tools + systemMsg = Message System (agentSystemPrompt agentCfg) Nothing Nothing + userMsg = Message User userPrompt Nothing Nothing + initialMessages = [systemMsg, userMsg] + + engineOnActivity engineCfg "Starting agent loop" + loop llm tools toolMap initialMessages 0 0 0 + where + maxIter = agentMaxIterations agentCfg + + loop :: LLM -> [Tool] -> Map.Map Text Tool -> [Message] -> Int -> Int -> Int -> IO (Either Text AgentResult) + loop llm tools' toolMap msgs iteration totalCalls totalTokens + | iteration >= maxIter = + pure + <| Left + <| "Max iterations (" + <> tshow maxIter + <> ") reached" + | otherwise = do + engineOnActivity engineCfg <| "Iteration " <> tshow (iteration + 1) + result <- chatWithUsage llm tools' msgs + case result of + Left err -> pure (Left err) + Right chatRes -> do + let msg = chatMessage chatRes + tokens = maybe 0 usageTotalTokens (chatUsage chatRes) + cost = estimateCost (llmModel llm) tokens + engineOnCost engineCfg tokens cost + let newTokens = totalTokens + tokens + case msgToolCalls msg of + Nothing -> do + engineOnActivity engineCfg "Agent completed" + pure + <| Right + <| AgentResult + { resultFinalMessage = msgContent msg, + resultToolCallCount = totalCalls, + resultIterations = iteration + 1, + resultTotalCost = estimateTotalCost (llmModel llm) newTokens, + resultTotalTokens = newTokens + } + Just [] -> do + engineOnActivity engineCfg "Agent completed (empty tool calls)" + pure + <| Right + <| AgentResult + { resultFinalMessage = msgContent msg, + resultToolCallCount = totalCalls, + resultIterations = iteration + 1, + resultTotalCost = estimateTotalCost (llmModel llm) newTokens, + resultTotalTokens = newTokens + } + Just tcs -> do + toolResults <- executeToolCalls engineCfg toolMap tcs + let newMsgs = msgs <> [msg] <> toolResults + newCalls = totalCalls + length tcs + loop llm tools' toolMap newMsgs (iteration + 1) newCalls newTokens + +buildToolMap :: [Tool] -> Map.Map Text Tool +buildToolMap = Map.fromList <. map (\t -> (toolName t, t)) + +executeToolCalls :: EngineConfig -> Map.Map Text Tool -> [ToolCall] -> IO [Message] +executeToolCalls engineCfg toolMap = traverse executeSingle + where + executeSingle tc = do + let name = fcName (tcFunction tc) + argsText = fcArguments (tcFunction tc) + callId = tcId tc + engineOnActivity engineCfg <| "Executing tool: " <> name + case Map.lookup name toolMap of + Nothing -> do + let errMsg = "Tool not found: " <> name + engineOnToolCall engineCfg name errMsg + pure <| Message ToolRole errMsg Nothing (Just callId) + Just tool -> do + case Aeson.decode (BL.fromStrict (TE.encodeUtf8 argsText)) of + Nothing -> do + let errMsg = "Invalid JSON arguments: " <> argsText + engineOnToolCall engineCfg name errMsg + pure <| Message ToolRole errMsg Nothing (Just callId) + Just args -> do + resultValue <- toolExecute tool args + let resultText = TE.decodeUtf8 (BL.toStrict (Aeson.encode resultValue)) + summary = Text.take 100 resultText + engineOnToolCall engineCfg name summary + pure <| Message ToolRole resultText Nothing (Just callId) + +estimateCost :: Text -> Int -> Int +estimateCost model tokens + | "gpt-4o-mini" `Text.isInfixOf` model = tokens * 15 `div` 1000000 + | "gpt-4o" `Text.isInfixOf` model = tokens * 250 `div` 100000 + | "gpt-4" `Text.isInfixOf` model = tokens * 3 `div` 100000 + | "claude" `Text.isInfixOf` model = tokens * 3 `div` 100000 + | otherwise = tokens `div` 100000 + +estimateTotalCost :: Text -> Int -> Int +estimateTotalCost = estimateCost diff --git a/Omni/Agent/Git.hs b/Omni/Agent/Git.hs new file mode 100644 index 0000000..4c06cf6 --- /dev/null +++ b/Omni/Agent/Git.hs @@ -0,0 +1,232 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Git operations for the agent. +-- +-- : out omni-agent-git +-- : dep temporary +module Omni.Agent.Git + ( checkout, + syncWithLive, + commit, + createBranch, + getCurrentBranch, + branchExists, + isMerged, + listBranches, + showFile, + getRepoRoot, + runGit, + main, + test, + ) +where + +import Alpha +import qualified Data.Text as Text +import qualified Omni.Log as Log +import Omni.Test ((@=?)) +import qualified Omni.Test as Test +import qualified System.Directory as Directory +import qualified System.Exit as Exit +import qualified System.IO.Temp as Temp +import qualified System.Process as Process + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Git" + [ Test.unit "checkout works" <| do + Temp.withSystemTempDirectory "omni-agent-git-test" <| \tmpDir -> do + let repo = tmpDir <> "/repo" + Directory.createDirectory repo + -- init repo + git repo ["init"] + git repo ["branch", "-m", "master"] + git repo ["config", "user.email", "you@example.com"] + git repo ["config", "user.name", "Your Name"] + + -- commit A + writeFile (repo <> "/a.txt") "A" + git repo ["add", "a.txt"] + git repo ["commit", "-m", "A"] + shaA <- getSha repo "HEAD" + + -- create branch dev + git repo ["checkout", "-b", "dev"] + + -- commit B + writeFile (repo <> "/b.txt") "B" + git repo ["add", "b.txt"] + git repo ["commit", "-m", "B"] + shaB <- getSha repo "HEAD" + + -- switch back to master + git repo ["checkout", "master"] + + -- Test 1: checkout dev + checkout repo "dev" + current <- getSha repo "HEAD" + shaB @=? current + + -- Test 2: checkout master + checkout repo "master" + current' <- getSha repo "HEAD" + shaA @=? current' + + -- Test 3: dirty state + writeFile (repo <> "/a.txt") "DIRTY" + checkout repo "dev" + current'' <- getSha repo "HEAD" + shaB @=? current'' + -- Verify dirty file is gone/overwritten (b.txt should exist, a.txt should be A from master? No, a.txt is in A and B) + -- Wait, in dev, a.txt is "A". + content <- readFile (repo <> "/a.txt") + "A" @=? content + + -- Test 4: untracked file + writeFile (repo <> "/untracked.txt") "DELETE ME" + checkout repo "master" + exists <- Directory.doesFileExist (repo <> "/untracked.txt") + False @=? exists + ] + +getSha :: FilePath -> String -> IO String +getSha dir ref = do + let cmd = (Process.proc "git" ["rev-parse", ref]) {Process.cwd = Just dir} + (code, out, _) <- Process.readCreateProcessWithExitCode cmd "" + case code of + Exit.ExitSuccess -> pure <| strip out + _ -> panic "getSha failed" + +-- | Checkout a specific ref (SHA, branch, tag) in the given repository path. +-- This function ensures the repository is in the correct state by: +-- 1. Fetching all updates +-- 2. Checking out the ref (forcing overwrites of local changes) +-- 3. Resetting hard to the ref (to ensure clean state) +-- 4. Cleaning untracked files +-- 5. Updating submodules +checkout :: FilePath -> Text -> IO () +checkout repoPath ref = do + let r = Text.unpack ref + + Log.info ["git", "checkout", ref, "in", Text.pack repoPath] + + -- Fetch all refs to ensure we have the target + git repoPath ["fetch", "--all", "--tags"] + + -- Checkout the ref, discarding local changes + git repoPath ["checkout", "--force", r] + + -- Reset hard to ensure we are exactly at the target state + git repoPath ["reset", "--hard", r] + + -- Remove untracked files and directories + git repoPath ["clean", "-fdx"] + + -- Update submodules + git repoPath ["submodule", "update", "--init", "--recursive"] + + Log.good ["git", "checkout", "complete"] + Log.br + +-- | Run a git command in the given directory. +git :: FilePath -> [String] -> IO () +git dir args = do + let cmd = (Process.proc "git" args) {Process.cwd = Just dir} + (exitCode, out, err) <- Process.readCreateProcessWithExitCode cmd "" + case exitCode of + Exit.ExitSuccess -> pure () + Exit.ExitFailure code -> do + Log.fail ["git command failed", Text.pack (show args), "code: " <> show code] + Log.info [Text.pack out] + Log.info [Text.pack err] + Log.br + panic <| "git command failed: git " <> show args + +syncWithLive :: FilePath -> IO () +syncWithLive repo = do + Log.info ["git", "syncing with live"] + -- git repo ["fetch", "origin", "live"] -- Optional + + -- Try sync (branchless sync), if fail, panic + -- This replaces manual rebase and handles stack movement + let cmd = (Process.proc "git" ["sync"]) {Process.cwd = Just repo} + (code, out, err) <- Process.readCreateProcessWithExitCode cmd "" + case code of + Exit.ExitSuccess -> pure () + Exit.ExitFailure _ -> do + Log.warn ["git sync failed", Text.pack err] + Log.info [Text.pack out] + panic "Sync with live failed (git sync)" + +commit :: FilePath -> Text -> IO () +commit repo msg = do + Log.info ["git", "commit", msg] + git repo ["add", "."] + + -- Check for changes before committing to avoid error + let checkCmd = (Process.proc "git" ["diff", "--cached", "--quiet"]) {Process.cwd = Just repo} + (code, _, _) <- Process.readCreateProcessWithExitCode checkCmd "" + + case code of + Exit.ExitSuccess -> Log.warn ["git", "nothing to commit", "skipping"] + Exit.ExitFailure 1 -> git repo ["commit", "-m", Text.unpack msg] + Exit.ExitFailure c -> panic <| "git diff failed with code " <> show c + +createBranch :: FilePath -> Text -> IO () +createBranch repo branch = do + Log.info ["git", "create branch", branch] + git repo ["checkout", "-b", Text.unpack branch] + +getCurrentBranch :: FilePath -> IO Text +getCurrentBranch repo = do + let cmd = (Process.proc "git" ["branch", "--show-current"]) {Process.cwd = Just repo} + (code, out, _) <- Process.readCreateProcessWithExitCode cmd "" + case code of + Exit.ExitSuccess -> pure <| Text.strip (Text.pack out) + _ -> panic "git branch failed" + +branchExists :: FilePath -> Text -> IO Bool +branchExists repo branch = do + let cmd = (Process.proc "git" ["show-ref", "--verify", "refs/heads/" <> Text.unpack branch]) {Process.cwd = Just repo} + (code, _, _) <- Process.readCreateProcessWithExitCode cmd "" + pure (code == Exit.ExitSuccess) + +isMerged :: FilePath -> Text -> Text -> IO Bool +isMerged repo branch target = do + -- Check if 'branch' is merged into 'target' + -- git merge-base --is-ancestor <branch> <target> + let cmd = (Process.proc "git" ["merge-base", "--is-ancestor", Text.unpack branch, Text.unpack target]) {Process.cwd = Just repo} + (code, _, _) <- Process.readCreateProcessWithExitCode cmd "" + pure (code == Exit.ExitSuccess) + +listBranches :: FilePath -> Text -> IO [Text] +listBranches repo pat = do + let cmd = (Process.proc "git" ["branch", "--list", Text.unpack pat, "--format=%(refname:short)"]) {Process.cwd = Just repo} + (code, out, _) <- Process.readCreateProcessWithExitCode cmd "" + case code of + Exit.ExitSuccess -> pure <| filter (not <. Text.null) (Text.lines (Text.pack out)) + _ -> panic "git branch list failed" + +showFile :: FilePath -> Text -> FilePath -> IO (Maybe Text) +showFile repo branch path = do + let cmd = (Process.proc "git" ["show", Text.unpack branch <> ":" <> path]) {Process.cwd = Just repo} + (code, out, _) <- Process.readCreateProcessWithExitCode cmd "" + case code of + Exit.ExitSuccess -> pure <| Just (Text.pack out) + _ -> pure Nothing + +getRepoRoot :: FilePath -> IO FilePath +getRepoRoot dir = do + let cmd = (Process.proc "git" ["rev-parse", "--show-toplevel"]) {Process.cwd = Just dir} + (code, out, _) <- Process.readCreateProcessWithExitCode cmd "" + case code of + Exit.ExitSuccess -> pure <| strip out + _ -> panic "git rev-parse failed" + +runGit :: FilePath -> [String] -> IO () +runGit = git diff --git a/Omni/Agent/Log.hs b/Omni/Agent/Log.hs new file mode 100644 index 0000000..46ea009 --- /dev/null +++ b/Omni/Agent/Log.hs @@ -0,0 +1,154 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Status of the agent for the UI +module Omni.Agent.Log where + +import Alpha +import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef) +import qualified Data.Text as Text +import qualified Data.Text.IO as TIO +import Data.Time.Clock (NominalDiffTime, UTCTime, diffUTCTime, getCurrentTime) +import Data.Time.Format (defaultTimeLocale, parseTimeOrError) +import qualified System.Console.ANSI as ANSI +import qualified System.IO as IO +import System.IO.Unsafe (unsafePerformIO) +import Text.Printf (printf) + +-- | Status of the agent for the UI +data Status = Status + { statusWorker :: Text, + statusTask :: Maybe Text, + statusThread :: Maybe Text, + statusFiles :: Int, + statusCredits :: Double, + statusStartTime :: UTCTime, + statusActivity :: Text + } + deriving (Show, Eq) + +emptyStatus :: Text -> UTCTime -> Status +emptyStatus workerName startTime = + Status + { statusWorker = workerName, + statusTask = Nothing, + statusThread = Nothing, + statusFiles = 0, + statusCredits = 0.0, + statusStartTime = startTime, + statusActivity = "Idle" + } + +-- | Global state for the status bar +{-# NOINLINE currentStatus #-} +currentStatus :: IORef Status +currentStatus = unsafePerformIO (newIORef (emptyStatus "Unknown" defaultStartTime)) + +defaultStartTime :: UTCTime +defaultStartTime = parseTimeOrError True defaultTimeLocale "%Y-%m-%d %H:%M:%S %Z" "2000-01-01 00:00:00 UTC" + +-- | Initialize the status bar system +init :: Text -> IO () +init workerName = do + IO.hSetBuffering IO.stderr IO.LineBuffering + startTime <- getCurrentTime + writeIORef currentStatus (emptyStatus workerName startTime) + -- Reserve 5 lines at bottom + IO.hPutStrLn IO.stderr "" + IO.hPutStrLn IO.stderr "" + IO.hPutStrLn IO.stderr "" + IO.hPutStrLn IO.stderr "" + IO.hPutStrLn IO.stderr "" + ANSI.hCursorUp IO.stderr 5 + +-- | Update the status +update :: (Status -> Status) -> IO () +update f = do + modifyIORef' currentStatus f + render + +-- | Get the current status +getStatus :: IO Status +getStatus = readIORef currentStatus + +-- | Set the activity message +updateActivity :: Text -> IO () +updateActivity msg = update (\s -> s {statusActivity = msg}) + +-- | Log a scrolling message (appears above status bars) +log :: Text -> IO () +log msg = do + -- Clear status bars + ANSI.hClearLine IO.stderr + ANSI.hCursorDown IO.stderr 1 + ANSI.hClearLine IO.stderr + ANSI.hCursorDown IO.stderr 1 + ANSI.hClearLine IO.stderr + ANSI.hCursorDown IO.stderr 1 + ANSI.hClearLine IO.stderr + ANSI.hCursorDown IO.stderr 1 + ANSI.hClearLine IO.stderr + ANSI.hCursorUp IO.stderr 4 + + -- Print message (scrolls screen) + TIO.hPutStrLn IO.stderr msg + + -- Re-render status bars at bottom + -- (Since we scrolled, we are now on the line above where the first status line should be) + render + +-- | Render the five status lines +render :: IO () +render = do + Status {..} <- readIORef currentStatus + now <- getCurrentTime + let taskStr = maybe "None" identity statusTask + threadStr = maybe "None" identity statusThread + elapsed = diffUTCTime now statusStartTime + elapsedStr = formatElapsed elapsed + + -- Line 1: Worker | Thread + ANSI.hSetCursorColumn IO.stderr 0 + ANSI.hClearLine IO.stderr + TIO.hPutStr IO.stderr ("[Worker: " <> statusWorker <> "] Thread: " <> threadStr) + + -- Line 2: Task + ANSI.hCursorDown IO.stderr 1 + ANSI.hSetCursorColumn IO.stderr 0 + ANSI.hClearLine IO.stderr + TIO.hPutStr IO.stderr ("Task: " <> taskStr) + + -- Line 3: Files | Credits + ANSI.hCursorDown IO.stderr 1 + ANSI.hSetCursorColumn IO.stderr 0 + ANSI.hClearLine IO.stderr + let creditsStr = Text.pack (printf "%.2f" statusCredits) + TIO.hPutStr IO.stderr ("Files: " <> tshow statusFiles <> " | Credits: $" <> creditsStr) + + -- Line 4: Time (elapsed duration) + ANSI.hCursorDown IO.stderr 1 + ANSI.hSetCursorColumn IO.stderr 0 + ANSI.hClearLine IO.stderr + TIO.hPutStr IO.stderr ("Time: " <> elapsedStr) + + -- Line 5: Activity + ANSI.hCursorDown IO.stderr 1 + ANSI.hSetCursorColumn IO.stderr 0 + ANSI.hClearLine IO.stderr + TIO.hPutStr IO.stderr ("> " <> statusActivity) + + -- Return cursor to line 1 + ANSI.hCursorUp IO.stderr 4 + IO.hFlush IO.stderr + +-- | Format elapsed time as MM:SS or HH:MM:SS +formatElapsed :: NominalDiffTime -> Text +formatElapsed elapsed = + let totalSecs = floor elapsed :: Int + hours = totalSecs `div` 3600 + mins = (totalSecs `mod` 3600) `div` 60 + secs = totalSecs `mod` 60 + in if hours > 0 + then Text.pack (printf "%02d:%02d:%02d" hours mins secs) + else Text.pack (printf "%02d:%02d" mins secs) diff --git a/Omni/Agent/Tools.hs b/Omni/Agent/Tools.hs new file mode 100644 index 0000000..0312924 --- /dev/null +++ b/Omni/Agent/Tools.hs @@ -0,0 +1,582 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE NoImplicitPrelude #-} + +-- | Core coding tools for Jr agent. +-- +-- Provides implementations for: +-- - readFile: Read file contents +-- - writeFile: Write file contents +-- - editFile: Search/replace edit +-- - runBash: Execute shell commands +-- - searchCodebase: Ripgrep wrapper for code search +-- +-- All tools return structured JSON results. +-- +-- : out omni-agent-tools +-- : dep aeson +-- : dep directory +module Omni.Agent.Tools + ( readFileTool, + writeFileTool, + editFileTool, + runBashTool, + searchCodebaseTool, + allTools, + ReadFileArgs (..), + WriteFileArgs (..), + EditFileArgs (..), + RunBashArgs (..), + SearchCodebaseArgs (..), + ToolResult (..), + main, + test, + ) +where + +import Alpha +import Data.Aeson ((.!=), (.:), (.:?), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.List as List +import qualified Data.Text as Text +import qualified Data.Text.IO as TextIO +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Test as Test +import qualified System.Directory as Directory +import qualified System.Exit as Exit +import qualified System.Process as Process + +main :: IO () +main = Test.run test + +test :: Test.Tree +test = + Test.group + "Omni.Agent.Tools" + [ Test.unit "readFileTool schema is valid" <| do + let schema = Engine.toolJsonSchema readFileTool + case schema of + Aeson.Object _ -> pure () + _ -> Test.assertFailure "Schema should be an object", + Test.unit "writeFileTool schema is valid" <| do + let schema = Engine.toolJsonSchema writeFileTool + case schema of + Aeson.Object _ -> pure () + _ -> Test.assertFailure "Schema should be an object", + Test.unit "editFileTool schema is valid" <| do + let schema = Engine.toolJsonSchema editFileTool + case schema of + Aeson.Object _ -> pure () + _ -> Test.assertFailure "Schema should be an object", + Test.unit "runBashTool schema is valid" <| do + let schema = Engine.toolJsonSchema runBashTool + case schema of + Aeson.Object _ -> pure () + _ -> Test.assertFailure "Schema should be an object", + Test.unit "searchCodebaseTool schema is valid" <| do + let schema = Engine.toolJsonSchema searchCodebaseTool + case schema of + Aeson.Object _ -> pure () + _ -> Test.assertFailure "Schema should be an object", + Test.unit "allTools contains 5 tools" <| do + length allTools Test.@=? 5, + Test.unit "ReadFileArgs parses correctly" <| do + let json = Aeson.object ["path" .= ("test.txt" :: Text)] + case Aeson.fromJSON json of + Aeson.Success (args :: ReadFileArgs) -> readFilePath args Test.@=? "test.txt" + Aeson.Error e -> Test.assertFailure e, + Test.unit "WriteFileArgs parses correctly" <| do + let json = Aeson.object ["path" .= ("test.txt" :: Text), "content" .= ("hello" :: Text)] + case Aeson.fromJSON json of + Aeson.Success (args :: WriteFileArgs) -> do + writeFilePath args Test.@=? "test.txt" + writeFileContent args Test.@=? "hello" + Aeson.Error e -> Test.assertFailure e, + Test.unit "EditFileArgs parses correctly" <| do + let json = + Aeson.object + [ "path" .= ("test.txt" :: Text), + "old_str" .= ("foo" :: Text), + "new_str" .= ("bar" :: Text) + ] + case Aeson.fromJSON json of + Aeson.Success (args :: EditFileArgs) -> do + editFilePath args Test.@=? "test.txt" + editFileOldStr args Test.@=? "foo" + editFileNewStr args Test.@=? "bar" + Aeson.Error e -> Test.assertFailure e, + Test.unit "RunBashArgs parses correctly" <| do + let json = Aeson.object ["command" .= ("ls -la" :: Text)] + case Aeson.fromJSON json of + Aeson.Success (args :: RunBashArgs) -> runBashCommand args Test.@=? "ls -la" + Aeson.Error e -> Test.assertFailure e, + Test.unit "SearchCodebaseArgs parses correctly" <| do + let json = Aeson.object ["pattern" .= ("TODO" :: Text)] + case Aeson.fromJSON json of + Aeson.Success (args :: SearchCodebaseArgs) -> searchPattern args Test.@=? "TODO" + Aeson.Error e -> Test.assertFailure e, + Test.unit "ToolResult success JSON roundtrip" <| do + let result = ToolResult True "done" Nothing + case Aeson.decode (Aeson.encode result) of + Nothing -> Test.assertFailure "Failed to decode ToolResult" + Just decoded -> toolResultSuccess decoded Test.@=? True, + Test.unit "ToolResult failure JSON roundtrip" <| do + let result = ToolResult False "" (Just "error occurred") + case Aeson.decode (Aeson.encode result) of + Nothing -> Test.assertFailure "Failed to decode ToolResult" + Just decoded -> toolResultError decoded Test.@=? Just "error occurred", + Test.unit "readFileTool handles missing files" <| do + let args = Aeson.object ["path" .= ("/nonexistent/path/to/file.txt" :: Text)] + result <- Engine.toolExecute readFileTool args + case Aeson.fromJSON result of + Aeson.Success (tr :: ToolResult) -> do + toolResultSuccess tr Test.@=? False + isJust (toolResultError tr) Test.@=? True + Aeson.Error e -> Test.assertFailure e, + Test.unit "editFileTool handles no-match case" <| do + let args = + Aeson.object + [ "path" .= ("/nonexistent/file.txt" :: Text), + "old_str" .= ("needle" :: Text), + "new_str" .= ("replacement" :: Text) + ] + result <- Engine.toolExecute editFileTool args + case Aeson.fromJSON result of + Aeson.Success (tr :: ToolResult) -> toolResultSuccess tr Test.@=? False + Aeson.Error e -> Test.assertFailure e, + Test.unit "runBashTool captures exit codes" <| do + let args = Aeson.object ["command" .= ("exit 42" :: Text)] + result <- Engine.toolExecute runBashTool args + case Aeson.fromJSON result of + Aeson.Success (tr :: ToolResult) -> do + toolResultSuccess tr Test.@=? False + toolResultError tr Test.@=? Just "Exit code: 42" + Aeson.Error e -> Test.assertFailure e, + Test.unit "runBashTool captures stdout" <| do + let args = Aeson.object ["command" .= ("echo hello" :: Text)] + result <- Engine.toolExecute runBashTool args + case Aeson.fromJSON result of + Aeson.Success (tr :: ToolResult) -> do + toolResultSuccess tr Test.@=? True + ("hello" `Text.isInfixOf` toolResultOutput tr) Test.@=? True + Aeson.Error e -> Test.assertFailure e, + Test.unit "searchCodebaseTool returns structured results" <| do + let args = + Aeson.object + [ "pattern" .= ("module" :: Text), + "path" .= ("." :: Text), + "max_results" .= (5 :: Int) + ] + result <- Engine.toolExecute searchCodebaseTool args + case Aeson.fromJSON result of + Aeson.Success (tr :: ToolResult) -> toolResultSuccess tr Test.@=? True + Aeson.Error e -> Test.assertFailure e + ] + +data ToolResult = ToolResult + { toolResultSuccess :: Bool, + toolResultOutput :: Text, + toolResultError :: Maybe Text + } + deriving (Show, Eq, Generic) + +instance Aeson.ToJSON ToolResult where + toJSON r = + Aeson.object + <| catMaybes + [ Just ("success" .= toolResultSuccess r), + Just ("output" .= toolResultOutput r), + ("error" .=) </ toolResultError r + ] + +instance Aeson.FromJSON ToolResult where + parseJSON = + Aeson.withObject "ToolResult" <| \v -> + (ToolResult </ (v .: "success")) + <*> (v .:? "output" .!= "") + <*> (v .:? "error") + +mkSuccess :: Text -> Aeson.Value +mkSuccess output = Aeson.toJSON <| ToolResult True output Nothing + +mkError :: Text -> Aeson.Value +mkError err = Aeson.toJSON <| ToolResult False "" (Just err) + +allTools :: [Engine.Tool] +allTools = + [ readFileTool, + writeFileTool, + editFileTool, + runBashTool, + searchCodebaseTool + ] + +data ReadFileArgs = ReadFileArgs + { readFilePath :: Text, + readFileStartLine :: Maybe Int, + readFileEndLine :: Maybe Int + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON ReadFileArgs where + parseJSON = + Aeson.withObject "ReadFileArgs" <| \v -> + (ReadFileArgs </ (v .: "path")) + <*> (v .:? "start_line") + <*> (v .:? "end_line") + +readFileTool :: Engine.Tool +readFileTool = + Engine.Tool + { Engine.toolName = "read_file", + Engine.toolDescription = "Read the contents of a file. Can optionally read a specific line range.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "path" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Absolute path to the file to read" :: Text) + ], + "start_line" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Optional: first line to read (1-indexed)" :: Text) + ], + "end_line" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Optional: last line to read (1-indexed)" :: Text) + ] + ], + "required" .= (["path"] :: [Text]) + ], + Engine.toolExecute = executeReadFile + } + +executeReadFile :: Aeson.Value -> IO Aeson.Value +executeReadFile v = + case Aeson.fromJSON v of + Aeson.Error e -> pure <| mkError (Text.pack e) + Aeson.Success args -> do + let path = Text.unpack (readFilePath args) + exists <- Directory.doesFileExist path + if exists + then do + content <- TextIO.readFile path + let allLines = Text.lines content + startIdx = maybe 0 (\n -> n - 1) (readFileStartLine args) + endIdx = maybe (length allLines) identity (readFileEndLine args) + selectedLines = take (endIdx - startIdx) (drop startIdx allLines) + numberedLines = zipWith formatLine [(startIdx + 1) ..] selectedLines + result = Text.unlines numberedLines + pure <| mkSuccess result + else pure <| mkError ("File not found: " <> readFilePath args) + where + formatLine :: Int -> Text -> Text + formatLine n line = Text.pack (show n) <> ": " <> line + +data WriteFileArgs = WriteFileArgs + { writeFilePath :: Text, + writeFileContent :: Text + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON WriteFileArgs where + parseJSON = + Aeson.withObject "WriteFileArgs" <| \v -> + (WriteFileArgs </ (v .: "path")) + <*> (v .: "content") + +writeFileTool :: Engine.Tool +writeFileTool = + Engine.Tool + { Engine.toolName = "write_file", + Engine.toolDescription = "Create or overwrite a file with the given content.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "path" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Absolute path to the file to write" :: Text) + ], + "content" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Content to write to the file" :: Text) + ] + ], + "required" .= (["path", "content"] :: [Text]) + ], + Engine.toolExecute = executeWriteFile + } + +executeWriteFile :: Aeson.Value -> IO Aeson.Value +executeWriteFile v = + case Aeson.fromJSON v of + Aeson.Error e -> pure <| mkError (Text.pack e) + Aeson.Success args -> do + let path = Text.unpack (writeFilePath args) + let dir = takeDirectory path + dirExists <- Directory.doesDirectoryExist dir + if dirExists + then do + TextIO.writeFile path (writeFileContent args) + pure <| mkSuccess ("File written: " <> writeFilePath args) + else pure <| mkError ("Parent directory does not exist: " <> Text.pack dir) + where + takeDirectory :: FilePath -> FilePath + takeDirectory p = + let parts = Text.splitOn "/" (Text.pack p) + in Text.unpack (Text.intercalate "/" (List.init parts)) + +data EditFileArgs = EditFileArgs + { editFilePath :: Text, + editFileOldStr :: Text, + editFileNewStr :: Text, + editFileReplaceAll :: Maybe Bool + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON EditFileArgs where + parseJSON = + Aeson.withObject "EditFileArgs" <| \v -> + (EditFileArgs </ (v .: "path")) + <*> (v .: "old_str") + <*> (v .: "new_str") + <*> (v .:? "replace_all") + +editFileTool :: Engine.Tool +editFileTool = + Engine.Tool + { Engine.toolName = "edit_file", + Engine.toolDescription = "Edit a file by replacing old_str with new_str. By default replaces only the first occurrence unless replace_all is true.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "path" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Absolute path to the file to edit" :: Text) + ], + "old_str" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("The text to search for and replace" :: Text) + ], + "new_str" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("The replacement text" :: Text) + ], + "replace_all" + .= Aeson.object + [ "type" .= ("boolean" :: Text), + "description" .= ("If true, replace all occurrences; otherwise replace only the first" :: Text) + ] + ], + "required" .= (["path", "old_str", "new_str"] :: [Text]) + ], + Engine.toolExecute = executeEditFile + } + +executeEditFile :: Aeson.Value -> IO Aeson.Value +executeEditFile v = + case Aeson.fromJSON v of + Aeson.Error e -> pure <| mkError (Text.pack e) + Aeson.Success args -> do + let path = Text.unpack (editFilePath args) + exists <- Directory.doesFileExist path + if exists + then do + content <- TextIO.readFile path + let oldStr = editFileOldStr args + newStr = editFileNewStr args + replaceAll = fromMaybe False (editFileReplaceAll args) + if Text.isInfixOf oldStr content + then do + let newContent = + if replaceAll + then Text.replace oldStr newStr content + else replaceFirst oldStr newStr content + TextIO.writeFile path newContent + let count = + if replaceAll + then Text.count oldStr content + else 1 + pure <| mkSuccess ("Replaced " <> tshow count <> " occurrence(s)") + else pure <| mkError ("old_str not found in file: " <> editFilePath args) + else pure <| mkError ("File not found: " <> editFilePath args) + +replaceFirst :: Text -> Text -> Text -> Text +replaceFirst old new content = + case Text.breakOn old content of + (before, after) -> + if Text.null after + then content + else before <> new <> Text.drop (Text.length old) after + +data RunBashArgs = RunBashArgs + { runBashCommand :: Text, + runBashCwd :: Maybe Text, + runBashTimeout :: Maybe Int + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON RunBashArgs where + parseJSON = + Aeson.withObject "RunBashArgs" <| \v -> + (RunBashArgs </ (v .: "command")) + <*> (v .:? "cwd") + <*> (v .:? "timeout") + +runBashTool :: Engine.Tool +runBashTool = + Engine.Tool + { Engine.toolName = "run_bash", + Engine.toolDescription = "Execute a shell command and return stdout/stderr.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "command" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("The shell command to execute" :: Text) + ], + "cwd" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Optional: working directory for the command" :: Text) + ], + "timeout" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Optional: timeout in seconds (default: 300)" :: Text) + ] + ], + "required" .= (["command"] :: [Text]) + ], + Engine.toolExecute = executeRunBash + } + +executeRunBash :: Aeson.Value -> IO Aeson.Value +executeRunBash v = + case Aeson.fromJSON v of + Aeson.Error e -> pure <| mkError (Text.pack e) + Aeson.Success args -> do + let cmd = Text.unpack (runBashCommand args) + proc = + (Process.shell cmd) + { Process.cwd = Text.unpack </ runBashCwd args + } + (exitCode, stdoutStr, stderrStr) <- Process.readCreateProcessWithExitCode proc "" + let output = Text.pack stdoutStr <> Text.pack stderrStr + case exitCode of + Exit.ExitSuccess -> + pure + <| Aeson.toJSON + <| ToolResult + { toolResultSuccess = True, + toolResultOutput = output, + toolResultError = Nothing + } + Exit.ExitFailure code -> + pure + <| Aeson.toJSON + <| ToolResult + { toolResultSuccess = False, + toolResultOutput = output, + toolResultError = Just ("Exit code: " <> tshow code) + } + +data SearchCodebaseArgs = SearchCodebaseArgs + { searchPattern :: Text, + searchPath :: Maybe Text, + searchGlob :: Maybe Text, + searchCaseSensitive :: Maybe Bool, + searchMaxResults :: Maybe Int + } + deriving (Show, Eq, Generic) + +instance Aeson.FromJSON SearchCodebaseArgs where + parseJSON = + Aeson.withObject "SearchCodebaseArgs" <| \v -> + (SearchCodebaseArgs </ (v .: "pattern")) + <*> (v .:? "path") + <*> (v .:? "glob") + <*> (v .:? "case_sensitive") + <*> (v .:? "max_results") + +searchCodebaseTool :: Engine.Tool +searchCodebaseTool = + Engine.Tool + { Engine.toolName = "search_codebase", + Engine.toolDescription = "Search the codebase using ripgrep. Returns matching lines with file paths and line numbers.", + Engine.toolJsonSchema = + Aeson.object + [ "type" .= ("object" :: Text), + "properties" + .= Aeson.object + [ "pattern" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("The regex pattern to search for" :: Text) + ], + "path" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Optional: directory or file path to search in" :: Text) + ], + "glob" + .= Aeson.object + [ "type" .= ("string" :: Text), + "description" .= ("Optional: file glob pattern (e.g., '*.hs')" :: Text) + ], + "case_sensitive" + .= Aeson.object + [ "type" .= ("boolean" :: Text), + "description" .= ("Optional: case sensitive search (default: false)" :: Text) + ], + "max_results" + .= Aeson.object + [ "type" .= ("integer" :: Text), + "description" .= ("Optional: maximum number of results (default: 100)" :: Text) + ] + ], + "required" .= (["pattern"] :: [Text]) + ], + Engine.toolExecute = executeSearchCodebase + } + +executeSearchCodebase :: Aeson.Value -> IO Aeson.Value +executeSearchCodebase v = + case Aeson.fromJSON v of + Aeson.Error e -> pure <| mkError (Text.pack e) + Aeson.Success args -> do + let pat = Text.unpack (searchPattern args) + maxRes = fromMaybe 100 (searchMaxResults args) + caseSensitive = fromMaybe False (searchCaseSensitive args) + baseArgs = + ["--line-number", "--no-heading", "--max-count=" <> show maxRes, pat] + caseArgs = ["--ignore-case" | not caseSensitive] + globArgs = maybe [] (\g -> ["--glob", Text.unpack g]) (searchGlob args) + pathArg = maybe ["."] (\p -> [Text.unpack p]) (searchPath args) + allArgs = caseArgs <> globArgs <> baseArgs <> pathArg + proc = Process.proc "rg" allArgs + (exitCode, stdoutStr, stderrStr) <- Process.readCreateProcessWithExitCode proc "" + case exitCode of + Exit.ExitSuccess -> + pure <| mkSuccess (Text.pack stdoutStr) + Exit.ExitFailure 1 -> + pure <| mkSuccess "No matches found" + Exit.ExitFailure code -> + pure <| mkError ("ripgrep failed with code " <> tshow code <> ": " <> Text.pack stderrStr) diff --git a/Omni/Agent/Worker.hs b/Omni/Agent/Worker.hs new file mode 100644 index 0000000..61c392b --- /dev/null +++ b/Omni/Agent/Worker.hs @@ -0,0 +1,446 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE NoImplicitPrelude #-} + +module Omni.Agent.Worker where + +import Alpha +import qualified Data.Aeson as Aeson +import qualified Data.Aeson.Key as AesonKey +import qualified Data.ByteString.Lazy as BSL +import Data.IORef (modifyIORef', newIORef, readIORef) +import qualified Data.List as List +import qualified Data.Text as Text +import qualified Data.Text.Encoding as TE +import qualified Data.Time +import qualified Omni.Agent.Core as Core +import qualified Omni.Agent.Engine as Engine +import qualified Omni.Agent.Log as AgentLog +import qualified Omni.Agent.Tools as Tools +import qualified Omni.Fact as Fact +import qualified Omni.Task.Core as TaskCore +import qualified System.Directory as Directory +import qualified System.Environment as Env +import qualified System.Exit as Exit +import System.FilePath ((</>)) +import qualified System.Process as Process + +start :: Core.Worker -> Maybe Text -> IO () +start worker maybeTaskId = do + if Core.workerQuiet worker + then putText ("[worker] Starting for " <> Core.workerName worker) + else do + AgentLog.init (Core.workerName worker) + AgentLog.log ("[worker] Starting for " <> Core.workerName worker) + case maybeTaskId of + Just tid -> logMsg worker ("[worker] Target task: " <> tid) + Nothing -> logMsg worker "[worker] No specific task, will pick from ready queue" + runOnce worker maybeTaskId + +-- | Log message respecting quiet mode +logMsg :: Core.Worker -> Text -> IO () +logMsg worker msg = + if Core.workerQuiet worker + then putText msg + else AgentLog.log msg + +-- | Convert key-value pairs to JSON metadata string +toMetadata :: [(Text, Text)] -> Text +toMetadata pairs = + let obj = Aeson.object [(AesonKey.fromText k, Aeson.String v) | (k, v) <- pairs] + in TE.decodeUtf8 (BSL.toStrict (Aeson.encode obj)) + +runOnce :: Core.Worker -> Maybe Text -> IO () +runOnce worker maybeTaskId = do + -- Find work + targetTask <- case maybeTaskId of + Just tid -> do + TaskCore.findTask tid </ TaskCore.loadTasks + Nothing -> do + readyTasks <- TaskCore.getReadyTasks + case readyTasks of + [] -> pure Nothing + (task : _) -> pure (Just task) + + case targetTask of + Nothing -> do + case maybeTaskId of + Just tid -> do + unless (Core.workerQuiet worker) <| AgentLog.updateActivity ("Task " <> tid <> " not found.") + logMsg worker ("[worker] Task " <> tid <> " not found.") + Nothing -> do + unless (Core.workerQuiet worker) <| AgentLog.updateActivity "No work found." + logMsg worker "[worker] No ready tasks found." + Just task -> do + processTask worker task + +processTask :: Core.Worker -> TaskCore.Task -> IO () +processTask worker task = do + let repo = Core.workerPath worker + let tid = TaskCore.taskId task + let quiet = Core.workerQuiet worker + let say = logMsg worker + + unless quiet <| AgentLog.update (\s -> s {AgentLog.statusTask = Just tid}) + say ("[worker] Claiming task " <> tid) + + -- Claim task + TaskCore.logActivity tid TaskCore.Claiming Nothing + TaskCore.updateTaskStatus tid TaskCore.InProgress [] + say "[worker] Status -> InProgress" + + -- Run agent with timing + startTime <- Data.Time.getCurrentTime + activityId <- TaskCore.logActivityWithMetrics tid TaskCore.Running Nothing Nothing (Just startTime) Nothing Nothing Nothing + + say "[worker] Starting engine..." + (exitCode, output, costCents) <- runWithEngine repo task + + endTime <- Data.Time.getCurrentTime + say ("[worker] Agent exited with: " <> tshow exitCode) + + -- Update the activity record with metrics + TaskCore.updateActivityMetrics activityId Nothing (Just endTime) (Just costCents) Nothing + + case exitCode of + Exit.ExitSuccess -> do + TaskCore.logActivity tid TaskCore.Reviewing Nothing + say "[worker] Running formatters..." + _ <- runFormatters repo + + -- Try to commit (this runs git hooks which may fail) + let commitMsg = formatCommitMessage task output + say "[worker] Attempting commit..." + commitResult <- tryCommit repo commitMsg + + case commitResult of + CommitFailed commitErr -> do + say ("[worker] Commit failed: " <> commitErr) + + -- Save failure context and reopen task for retry + maybeCtx <- TaskCore.getRetryContext tid + let attempt = maybe 1 (\c -> TaskCore.retryAttempt c + 1) maybeCtx + + if attempt > 3 + then do + say "[worker] Task failed 3 times, needs human intervention" + TaskCore.logActivity tid TaskCore.Failed (Just (toMetadata [("reason", "max_retries_exceeded")])) + TaskCore.updateTaskStatus tid TaskCore.Open [] + else do + let currentReason = "attempt " <> tshow attempt <> ": commit_failed: " <> commitErr + let accumulatedReason = case maybeCtx of + Nothing -> currentReason + Just ctx -> TaskCore.retryReason ctx <> "\n" <> currentReason + TaskCore.setRetryContext + TaskCore.RetryContext + { TaskCore.retryTaskId = tid, + TaskCore.retryOriginalCommit = "", + TaskCore.retryConflictFiles = [], + TaskCore.retryAttempt = attempt, + TaskCore.retryReason = accumulatedReason, + TaskCore.retryNotes = maybeCtx +> TaskCore.retryNotes + } + TaskCore.logActivity tid TaskCore.Retrying (Just (toMetadata [("attempt", tshow attempt)])) + TaskCore.updateTaskStatus tid TaskCore.Open [] + say ("[worker] Task reopened (attempt " <> tshow attempt <> "/3)") + NoChanges -> do + -- No changes = task already implemented, mark as Done + say "[worker] No changes to commit - task already done" + TaskCore.clearRetryContext tid + TaskCore.logActivity tid TaskCore.Completed (Just (toMetadata [("result", "no_changes")])) + TaskCore.updateTaskStatus tid TaskCore.Done [] + say ("[worker] ✓ Task " <> tid <> " -> Done (no changes)") + unless quiet <| AgentLog.update (\s -> s {AgentLog.statusTask = Nothing}) + CommitSuccess -> do + -- Commit succeeded, set to Review + TaskCore.logActivity tid TaskCore.Completed (Just (toMetadata [("result", "committed")])) + TaskCore.updateTaskStatus tid TaskCore.Review [] + say ("[worker] ✓ Task " <> tid <> " -> Review") + unless quiet <| AgentLog.update (\s -> s {AgentLog.statusTask = Nothing}) + Exit.ExitFailure code -> do + say ("[worker] Engine failed with code " <> tshow code) + TaskCore.logActivity tid TaskCore.Failed (Just (toMetadata [("exit_code", tshow code)])) + -- Don't set back to Open here - leave in InProgress for debugging + say "[worker] Task left in InProgress (engine failure)" + +-- | Run lint --fix to format and fix lint issues +runFormatters :: FilePath -> IO (Either Text ()) +runFormatters repo = do + let cmd = (Process.proc "lint" ["--fix"]) {Process.cwd = Just repo} + (code, _, _) <- Process.readCreateProcessWithExitCode cmd "" + case code of + Exit.ExitSuccess -> pure (Right ()) + Exit.ExitFailure _ -> pure (Right ()) -- lint --fix may exit non-zero but still fix things + +data CommitResult = CommitSuccess | NoChanges | CommitFailed Text + deriving (Show, Eq) + +-- | Try to commit, returning result +tryCommit :: FilePath -> Text -> IO CommitResult +tryCommit repo msg = do + -- Stage all changes + let addCmd = (Process.proc "git" ["add", "."]) {Process.cwd = Just repo} + (addCode, _, addErr) <- Process.readCreateProcessWithExitCode addCmd "" + case addCode of + Exit.ExitFailure _ -> pure <| CommitFailed (Text.pack addErr) + Exit.ExitSuccess -> do + -- Check for changes + let checkCmd = (Process.proc "git" ["diff", "--cached", "--quiet"]) {Process.cwd = Just repo} + (checkCode, _, _) <- Process.readCreateProcessWithExitCode checkCmd "" + case checkCode of + Exit.ExitSuccess -> pure NoChanges + Exit.ExitFailure 1 -> do + -- There are changes, commit them + let commitCmd = (Process.proc "git" ["commit", "-m", Text.unpack msg]) {Process.cwd = Just repo} + (commitCode, _, commitErr) <- Process.readCreateProcessWithExitCode commitCmd "" + case commitCode of + Exit.ExitSuccess -> pure CommitSuccess + Exit.ExitFailure _ -> pure <| CommitFailed (Text.pack commitErr) + Exit.ExitFailure c -> pure <| CommitFailed ("git diff failed with code " <> tshow c) + +-- | Run task using native Engine +-- Returns (ExitCode, output text, cost in cents) +runWithEngine :: FilePath -> TaskCore.Task -> IO (Exit.ExitCode, Text, Int) +runWithEngine repo task = do + -- Read API key from environment + maybeApiKey <- Env.lookupEnv "OPENROUTER_API_KEY" + case maybeApiKey of + Nothing -> pure (Exit.ExitFailure 1, "OPENROUTER_API_KEY not set", 0) + Just apiKey -> do + -- Check for retry context + maybeRetry <- TaskCore.getRetryContext (TaskCore.taskId task) + + -- Build the full prompt + let ns = fromMaybe "." (TaskCore.taskNamespace task) + let basePrompt = buildBasePrompt task ns repo + + -- Add retry context if present + let retryPrompt = buildRetryPrompt maybeRetry + + let prompt = basePrompt <> retryPrompt + + -- Read AGENTS.md + agentsMd <- + fmap (fromMaybe "") <| do + exists <- Directory.doesFileExist (repo </> "AGENTS.md") + if exists + then Just </ readFile (repo </> "AGENTS.md") + else pure Nothing + + -- Get relevant facts from the knowledge base + relevantFacts <- getRelevantFacts task + let factsSection = formatFacts relevantFacts + + -- Build system prompt + let systemPrompt = + prompt + <> "\n\nREPOSITORY GUIDELINES (AGENTS.md):\n" + <> agentsMd + <> factsSection + + -- Build user prompt from task comments + let userPrompt = formatTask task + + -- Select model based on task complexity (simple heuristic) + let model = selectModel task + + -- Build Engine config with callbacks + totalCostRef <- newIORef (0 :: Int) + let engineCfg = + Engine.EngineConfig + { Engine.engineLLM = + Engine.defaultLLM + { Engine.llmApiKey = Text.pack apiKey + }, + Engine.engineOnCost = \tokens cost -> do + modifyIORef' totalCostRef (+ cost) + AgentLog.log <| "Cost: " <> tshow cost <> " cents (" <> tshow tokens <> " tokens)", + Engine.engineOnActivity = \activity -> do + AgentLog.log <| "[engine] " <> activity, + Engine.engineOnToolCall = \toolName result -> do + AgentLog.log <| "[tool] " <> toolName <> ": " <> Text.take 100 result + } + + -- Build Agent config + let agentCfg = + Engine.AgentConfig + { Engine.agentModel = model, + Engine.agentTools = Tools.allTools, + Engine.agentSystemPrompt = systemPrompt, + Engine.agentMaxIterations = 20 + } + + -- Run the agent + result <- Engine.runAgent engineCfg agentCfg userPrompt + totalCost <- readIORef totalCostRef + + case result of + Left err -> pure (Exit.ExitFailure 1, "Engine error: " <> err, totalCost) + Right agentResult -> do + let output = Engine.resultFinalMessage agentResult + pure (Exit.ExitSuccess, output, totalCost) + +-- | Build the base prompt for the agent +buildBasePrompt :: TaskCore.Task -> Text -> FilePath -> Text +buildBasePrompt task ns repo = + "You are a Worker Agent.\n" + <> "Your goal is to implement the following task:\n\n" + <> formatTask task + <> "\n\nCRITICAL INSTRUCTIONS:\n" + <> "1. Analyze the codebase to understand where to make changes.\n" + <> "2. Implement the changes by editing files.\n" + <> "3. BEFORE finishing, you MUST run: bild --test " + <> ns + <> "\n" + <> "4. Fix ALL errors from bild --test (including hlint suggestions).\n" + <> "5. Keep running bild --test until it passes with no errors.\n" + <> "6. Do NOT update task status or manage git.\n" + <> "7. Only exit after bild --test passes.\n\n" + <> "IMPORTANT: The git commit will fail if hlint finds issues.\n" + <> "You must fix hlint suggestions like:\n" + <> "- 'Use list comprehension' -> use [x | cond] instead of if/else\n" + <> "- 'Avoid lambda' -> use function composition\n" + <> "- 'Redundant bracket' -> remove unnecessary parens\n\n" + <> "Context:\n" + <> "- Working directory: " + <> Text.pack repo + <> "\n" + <> "- Namespace: " + <> ns + <> "\n" + +-- | Build retry context prompt +buildRetryPrompt :: Maybe TaskCore.RetryContext -> Text +buildRetryPrompt Nothing = "" +buildRetryPrompt (Just ctx) = + "\n\n## RETRY CONTEXT (IMPORTANT)\n\n" + <> "This task was previously attempted but failed. Attempt: " + <> tshow (TaskCore.retryAttempt ctx) + <> "/3\n" + <> "Reason: " + <> TaskCore.retryReason ctx + <> "\n\n" + <> ( if null (TaskCore.retryConflictFiles ctx) + then "" + else + "Conflicting files from previous attempt:\n" + <> Text.unlines (map (" - " <>) (TaskCore.retryConflictFiles ctx)) + <> "\n" + ) + <> "Original commit: " + <> TaskCore.retryOriginalCommit ctx + <> "\n\n" + <> maybe "" (\notes -> "## HUMAN NOTES/GUIDANCE\n\n" <> notes <> "\n\n") (TaskCore.retryNotes ctx) + <> "INSTRUCTIONS FOR RETRY:\n" + <> "- The codebase has changed since your last attempt\n" + <> "- Re-implement this task on top of the CURRENT codebase\n" + <> "- If there were merge conflicts, the conflicting files may have been modified by others\n" + <> "- Review the current state of those files before making changes\n" + +-- | Select model based on task complexity (1-5 scale) +-- Uses OpenRouter model identifiers for Claude models +selectModel :: TaskCore.Task -> Text +selectModel task = selectModelByComplexity (TaskCore.taskComplexity task) + +-- | Select model based on complexity level +selectModelByComplexity :: Maybe Int -> Text +selectModelByComplexity Nothing = "anthropic/claude-sonnet-4.5" +selectModelByComplexity (Just 1) = "anthropic/claude-haiku-4.5" +selectModelByComplexity (Just 2) = "anthropic/claude-haiku-4.5" +selectModelByComplexity (Just 3) = "anthropic/claude-sonnet-4.5" +selectModelByComplexity (Just 4) = "anthropic/claude-sonnet-4.5" +selectModelByComplexity (Just 5) = "anthropic/claude-opus-4.5" +selectModelByComplexity (Just _) = "anthropic/claude-sonnet-4.5" + +formatTask :: TaskCore.Task -> Text +formatTask t = + "Task: " + <> TaskCore.taskId t + <> "\n" + <> "Title: " + <> TaskCore.taskTitle t + <> "\n" + <> "Type: " + <> Text.pack (show (TaskCore.taskType t)) + <> "\n" + <> "Status: " + <> Text.pack (show (TaskCore.taskStatus t)) + <> "\n" + <> "Priority: " + <> Text.pack (show (TaskCore.taskPriority t)) + <> "\n" + <> maybe "" (\p -> "Parent: " <> p <> "\n") (TaskCore.taskParent t) + <> maybe "" (\ns -> "Namespace: " <> ns <> "\n") (TaskCore.taskNamespace t) + <> "Created: " + <> Text.pack (show (TaskCore.taskCreatedAt t)) + <> "\n" + <> "Updated: " + <> Text.pack (show (TaskCore.taskUpdatedAt t)) + <> "\n" + <> (if Text.null (TaskCore.taskDescription t) then "" else "Description:\n" <> TaskCore.taskDescription t <> "\n\n") + <> formatDeps (TaskCore.taskDependencies t) + <> formatComments (TaskCore.taskComments t) + where + formatDeps [] = "" + formatDeps deps = "\nDependencies:\n" <> Text.unlines (map formatDep deps) + formatDep dep = " - " <> TaskCore.depId dep <> " [" <> Text.pack (show (TaskCore.depType dep)) <> "]" + formatComments [] = "" + formatComments cs = "\nComments/Notes:\n" <> Text.unlines (map formatComment cs) + formatComment c = " [" <> Text.pack (show (TaskCore.commentCreatedAt c)) <> "] " <> TaskCore.commentText c + +formatCommitMessage :: TaskCore.Task -> Text -> Text +formatCommitMessage task agentOutput = + let tid = TaskCore.taskId task + subject = cleanSubject (TaskCore.taskTitle task) + body = cleanBody agentOutput + in if Text.null body + then subject <> "\n\nTask-Id: " <> tid + else subject <> "\n\n" <> body <> "\n\nTask-Id: " <> tid + where + cleanSubject s = + let trailingPunct = ['.', ':', '!', '?', ',', ';', ' ', '-'] + stripped = Text.dropWhileEnd (`elem` trailingPunct) s + truncated = Text.take 72 stripped + noPunct = Text.dropWhileEnd (`elem` trailingPunct) truncated + capitalized = case Text.uncons noPunct of + Just (c, rest) -> Text.cons (toUpper c) rest + Nothing -> noPunct + in capitalized + + cleanBody :: Text -> Text + cleanBody output = + let stripped = Text.strip output + in if Text.null stripped + then "" + else + let lns = Text.lines stripped + cleaned = [Text.take 72 ln | ln <- lns] + in Text.intercalate "\n" cleaned + +-- | Get facts relevant to a task based on namespace/project +getRelevantFacts :: TaskCore.Task -> IO [TaskCore.Fact] +getRelevantFacts task = do + let namespace = fromMaybe "Omni" (TaskCore.taskNamespace task) + projectFacts <- Fact.getFactsByProject namespace + let sorted = List.sortBy (comparing (Down <. TaskCore.factConfidence)) projectFacts + pure (take 10 sorted) + +-- | Format facts for inclusion in the prompt +formatFacts :: [TaskCore.Fact] -> Text +formatFacts [] = "" +formatFacts facts = + Text.unlines + [ "\n\nKNOWLEDGE BASE FACTS:", + "(These are learned patterns/conventions from previous work)", + "" + ] + <> Text.unlines (map formatFact facts) + +-- | Format a single fact for the prompt +formatFact :: TaskCore.Fact -> Text +formatFact f = + "- " + <> TaskCore.factContent f + <> ( if null (TaskCore.factRelatedFiles f) + then "" + else " [" <> Text.intercalate ", " (TaskCore.factRelatedFiles f) <> "]" + ) |