path: root/Omni/Agent/Worker.hs

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

module Omni.Agent.Worker where

import Alpha
import qualified Data.Text as Text
import qualified Data.Text.IO as TIO
import qualified Omni.Agent.Core as Core
import qualified Omni.Agent.Log as AgentLog
import qualified Omni.Task.Core as TaskCore
import qualified System.Directory as Directory
import qualified System.Exit as Exit
import System.FilePath ((</>))
import qualified System.IO as IO
import qualified System.Process as Process

start :: Core.Worker -> Maybe Text -> IO ()
start worker maybeTaskId = do
  AgentLog.init (Core.workerName worker)
  AgentLog.log ("[worker] Starting for " <> Core.workerName worker)
  case maybeTaskId of
    Just tid -> AgentLog.log ("[worker] Target task: " <> tid)
    Nothing -> AgentLog.log "[worker] No specific task, will pick from ready queue"
  runOnce worker maybeTaskId

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
          AgentLog.updateActivity ("Task " <> tid <> " not found.")
          AgentLog.log ("[worker] Task " <> tid <> " not found.")
        Nothing -> do
          AgentLog.updateActivity "No work found."
          AgentLog.log "[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

  AgentLog.update (\s -> s {AgentLog.statusTask = Just tid})
  AgentLog.log ("[worker] Claiming task " <> tid)

  -- Claim task
  TaskCore.updateTaskStatus tid TaskCore.InProgress []
  AgentLog.log "[worker] Status -> InProgress"

  -- Run Amp
  AgentLog.log "[worker] Starting amp..."
  (exitCode, output) <- runAmp repo task
  AgentLog.log ("[worker] Amp exited with: " <> tshow exitCode)

  case exitCode of
    Exit.ExitSuccess -> do
      AgentLog.log "[worker] Running formatters..."
      _ <- runFormatters repo

      -- Try to commit (this runs git hooks which may fail)
      let commitMsg = formatCommitMessage task output
      AgentLog.log "[worker] Attempting commit..."
      commitResult <- tryCommit repo commitMsg

      case commitResult of
        Left commitErr -> do
          AgentLog.log ("[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
              AgentLog.log "[worker] Task failed 3 times, needs human intervention"
              TaskCore.updateTaskStatus tid TaskCore.Open []
            else do
              TaskCore.setRetryContext
                TaskCore.RetryContext
                  { TaskCore.retryTaskId = tid,
                    TaskCore.retryOriginalCommit = "",
                    TaskCore.retryConflictFiles = [],
                    TaskCore.retryAttempt = attempt,
                    TaskCore.retryReason = "commit_failed: " <> commitErr
                  }
              TaskCore.updateTaskStatus tid TaskCore.Open []
              AgentLog.log ("[worker] Task reopened (attempt " <> tshow attempt <> "/3)")
        Right () -> do
          -- Only set to Review after successful commit
          TaskCore.updateTaskStatus tid TaskCore.Review []
          AgentLog.log ("[worker] ✓ Task " <> tid <> " -> Review")
          AgentLog.update (\s -> s {AgentLog.statusTask = Nothing})
    Exit.ExitFailure code -> do
      AgentLog.log ("[worker] Amp failed with code " <> tshow code)
      -- Don't set back to Open here - leave in InProgress for debugging
      AgentLog.log "[worker] Task left in InProgress (amp 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

-- | Try to commit, returning error message on failure
tryCommit :: FilePath -> Text -> IO (Either Text ())
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 <| Left (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 (Right ()) -- Nothing to commit
        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 (Right ())
            Exit.ExitFailure _ -> pure <| Left (Text.pack commitErr)
        Exit.ExitFailure c -> pure <| Left ("git diff failed with code " <> tshow c)

runAmp :: FilePath -> TaskCore.Task -> IO (Exit.ExitCode, Text)
runAmp repo task = do
  -- Check for retry context
  maybeRetry <- TaskCore.getRetryContext (TaskCore.taskId task)

  let basePrompt =
        "You are a Worker Agent.\n"
          <> "Your goal is to implement the following task:\n\n"
          <> formatTask task
          <> "\n\nINSTRUCTIONS:\n"
          <> "1. Analyze the codebase (use finder/Grep) to understand where to make changes.\n"
          <> "2. Implement the changes by editing files.\n"
          <> "3. Run tests to verify your work (e.g., 'bild --test Omni/Namespace').\n"
          <> "4. Fix any errors found during testing.\n"
          <> "5. Do NOT update the task status or manage git branches.\n"
          <> "6. When finished and tested, exit.\n\n"
          <> "Context:\n"
          <> "- You are working in '"
          <> Text.pack repo
          <> "'.\n"
          <> "- The task is in namespace '"
          <> fromMaybe "root" (TaskCore.taskNamespace task)
          <> "'.\n"

  -- Add retry context if present
  let retryPrompt = case maybeRetry of
        Nothing -> ""
        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"
            <> "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"

  let prompt = basePrompt <> retryPrompt

  let logFile = repo </> "_/llm/amp.log"

  -- Read AGENTS.md
  agentsMd <-
    fmap (fromMaybe "") <| do
      exists <- Directory.doesFileExist (repo </> "AGENTS.md")
      if exists
        then Just </ readFile (repo </> "AGENTS.md")
        else pure Nothing

  let fullPrompt =
        prompt
          <> "\n\nREPOSITORY GUIDELINES (AGENTS.md):\n"
          <> agentsMd

  -- Remove old log file
  exists <- Directory.doesFileExist logFile
  when exists (Directory.removeFile logFile)

  Directory.createDirectoryIfMissing True (repo </> "_/llm")

  -- Assume amp is in PATH
  let args = ["--try-opus", "--log-level", "debug", "--log-file", "_/llm/amp.log", "--dangerously-allow-all", "-x", Text.unpack fullPrompt]

  let cp = (Process.proc "amp" args) {Process.cwd = Just repo, Process.std_out = Process.CreatePipe}
  (_, Just hOut, _, ph) <- Process.createProcess cp

  tid <- forkIO <| monitorLog logFile ph

  exitCode <- Process.waitForProcess ph
  output <- TIO.hGetContents hOut
  killThread tid
  pure (exitCode, output)

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"
    <> maybe "" (\d -> "Description:\n" <> d <> "\n\n") (TaskCore.taskDescription t)
    <> (if null (TaskCore.taskDependencies t) then "" else "\nDependencies:\n" <> Text.unlines (map formatDep (TaskCore.taskDependencies t)))
  where
    formatDep dep = "  - " <> TaskCore.depId dep <> " [" <> Text.pack (show (TaskCore.depType dep)) <> "]"

formatCommitMessage :: TaskCore.Task -> Text -> Text
formatCommitMessage task ampOutput =
  let tid = TaskCore.taskId task
      subject = cleanSubject (TaskCore.taskTitle task)
      body = cleanBody ampOutput
   in if Text.null body
        then subject <> "\n\nTask-Id: " <> tid
        else subject <> "\n\n" <> body <> "\n\nTask-Id: " <> tid
  where
    cleanSubject s =
      let stripped = Text.dropWhileEnd (`elem` ['.', ':', '!', '?', ' ']) s
          truncated = if Text.length stripped > 72 then Text.take 69 stripped <> "..." else stripped
          capitalized = case Text.uncons truncated of
            Just (c, rest) -> Text.cons (toUpper c) rest
            Nothing -> truncated
       in capitalized

    cleanBody :: Text -> Text
    cleanBody output =
      let stripped = Text.strip output
          lns = Text.lines stripped
          cleaned = map (Text.take 72) lns
       in Text.intercalate "\n" cleaned

monitorLog :: FilePath -> Process.ProcessHandle -> IO ()
monitorLog path ph = do
  waitForFile path
  IO.withFile path IO.ReadMode <| \h -> do
    IO.hSetBuffering h IO.LineBuffering
    go h
  where
    go h = do
      eof <- IO.hIsEOF h
      if eof
        then do
          mExit <- Process.getProcessExitCode ph
          case mExit of
            Nothing -> do
              threadDelay 100000 -- 0.1s
              go h
            Just _ -> pure ()
        else do
          line <- TIO.hGetLine h
          AgentLog.processLogLine line
          go h

waitForFile :: FilePath -> IO ()
waitForFile path = do
  exists <- Directory.doesFileExist path
  if exists
    then pure ()
    else do
      threadDelay 100000
      waitForFile path