path: root/Omni/Agent/Worker.hs

{-# 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 qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
import qualified Data.Text.IO as TIO
import qualified Data.Time
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
  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 Amp with timing
  say "[worker] Starting amp..."
  startTime <- Data.Time.getCurrentTime
  activityId <- TaskCore.logActivityWithMetrics tid TaskCore.Running Nothing Nothing (Just startTime) Nothing Nothing Nothing
  (exitCode, output) <- runAmp repo task
  endTime <- Data.Time.getCurrentTime
  say ("[worker] Amp exited with: " <> tshow exitCode)

  -- Capture metrics from agent log (thread URL, credits)
  status <- AgentLog.getStatus
  let threadUrl = ("https://ampcode.com/threads/" <>) </ AgentLog.statusThread status
  let costCents = Just <| floor (AgentLog.statusCredits status * 100)

  -- Update the activity record with metrics
  TaskCore.updateActivityMetrics activityId threadUrl (Just endTime) 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] Amp 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 (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

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)

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

  let ns = fromMaybe "." (TaskCore.taskNamespace task)
  let basePrompt =
        "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"

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

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

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