task: complete t-rWcqsDZFM.2 (Merge)

Amp-Thread-ID:
https://ampcode.com/threads/T-ca3b086b-5a85-422a-b13d-256784c04221
Co-authored-by: Amp <amp@ampcode.com>

1 files changed, 100 insertions, 1 deletions

diff --git a/Omni/Agent.hs b/Omni/Agent.hs
index f3fc29d..bad2737 100644
--- a/Omni/Agent.hs
+++ b/Omni/Agent.hs
@@ -9,12 +9,21 @@ module Omni.Agent 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.Git as Git
 import qualified Omni.Agent.Log as Log
 import qualified Omni.Agent.Worker as Worker
 import qualified Omni.Cli as Cli
+import qualified Omni.Task.Core as TaskCore
 import qualified Omni.Test as Test
 import qualified System.Console.Docopt as Docopt
+import qualified System.Directory as Directory
+import qualified System.Environment as Env
+import qualified System.Exit as Exit
+import System.FilePath ((</>))
+import qualified System.IO.Temp as Temp
+import qualified System.Process as Process
 
 main :: IO ()
 main = Cli.main plan
@@ -35,6 +44,9 @@ agent
 
 Usage:
   agent start <name> [--path=<path>]
+  agent harvest [--path=<path>]
+  agent merge-driver <ours> <theirs>
+  agent setup <name>
   agent test
   agent --help
 
@@ -61,8 +73,85 @@ move args
               }
 
       Worker.start worker
+  | args `Cli.has` Cli.command "harvest" = harvest args
+  | args `Cli.has` Cli.command "merge-driver" = mergeDriver args
+  | args `Cli.has` Cli.command "setup" = setup args
   | otherwise = putStrLn (Cli.usage help)
 
+harvest :: Cli.Arguments -> IO ()
+harvest args = do
+  let path = Cli.getArgWithDefault args "." (Cli.longOption "path")
+  putText "Harvesting task updates from workers..."
+
+  branches <- Git.listBranches path "omni-worker-*"
+  if null branches
+    then putText "No worker branches found."
+    else do
+      updated <- foldlM (processBranch path) False branches
+      when updated <| do
+        -- Consolidate
+        Directory.setCurrentDirectory path
+        TaskCore.exportTasks
+
+        -- Commit if changed
+        Git.commit path "task: harvest updates from workers"
+        putText "Success: Task database updated and committed."
+
+processBranch :: FilePath -> Bool -> Text -> IO Bool
+processBranch repo updated branch = do
+  putText <| "Checking " <> branch <> "..."
+  maybeContent <- Git.showFile repo branch ".tasks/tasks.jsonl"
+  case maybeContent of
+    Nothing -> do
+      putText <| "  Warning: Could not read .tasks/tasks.jsonl from " <> branch
+      pure updated
+    Just content -> do
+      -- Write to temp file
+      Temp.withSystemTempFile "worker-tasks.jsonl" <| \tempPath h -> do
+        TIO.hPutStr h content
+        IO.hClose h
+        -- Import
+        -- We need to ensure we are in the repo directory for TaskCore to find .tasks/tasks.jsonl
+        Directory.setCurrentDirectory repo
+        TaskCore.importTasks tempPath
+        putText <| "  Imported tasks from " <> branch
+        pure True
+
+mergeDriver :: Cli.Arguments -> IO ()
+mergeDriver args = do
+  ours <- Cli.getArgOrExit args (Cli.argument "ours")
+  theirs <- Cli.getArgOrExit args (Cli.argument "theirs")
+
+  -- Set TASK_DB_PATH to ours (the file git provided as the current version)
+  Env.setEnv "TASK_DB_PATH" ours
+  TaskCore.importTasks theirs
+  Exit.exitSuccess
+
+setup :: Cli.Arguments -> IO ()
+setup args = do
+  nameStr <- Cli.getArgOrExit args (Cli.argument "name")
+  let name = Text.pack nameStr
+  root <- Git.getRepoRoot "."
+  let worktreePath = root <> "/../" <> nameStr
+
+  putText <| "Creating worktree '" <> Text.pack worktreePath <> "' on branch '" <> name <> "' (from live)..."
+
+  -- git worktree add -b <name> <path> live
+  Git.runGit root ["worktree", "add", "-b", nameStr, worktreePath, "live"]
+
+  -- Copy .envrc.local if exists
+  let envrc = root </> ".envrc.local"
+  exists <- Directory.doesFileExist envrc
+  when exists <| do
+    putText "Copying .envrc.local..."
+    Directory.copyFile envrc (worktreePath </> ".envrc.local")
+
+  -- Config git
+  Git.runGit worktreePath ["config", "user.name", "Omni Worker"]
+  Git.runGit worktreePath ["config", "user.email", "bot@omni.agent"]
+
+  putText <| "Worker setup complete at " <> Text.pack worktreePath
+
 test :: Test.Tree
 test = Test.group "Omni.Agent" [unitTests, logTests]
 
@@ -84,5 +173,15 @@ unitTests =
         let result = Docopt.parseArgs help ["start", "worker-1"]
         case result of
           Left err -> Test.assertFailure <| "Failed to parse 'start': " <> show err
-          Right args -> args `Cli.has` Cli.command "start" Test.@?= True
+          Right args -> args `Cli.has` Cli.command "start" Test.@?= True,
+      Test.unit "can parse harvest command" <| do
+        let result = Docopt.parseArgs help ["harvest"]
+        case result of
+          Left err -> Test.assertFailure <| "Failed to parse 'harvest': " <> show err
+          Right args -> args `Cli.has` Cli.command "harvest" Test.@?= True,
+      Test.unit "can parse setup command" <| do
+        let result = Docopt.parseArgs help ["setup", "worker-2"]
+        case result of
+          Left err -> Test.assertFailure <| "Failed to parse 'setup': " <> show err
+          Right args -> args `Cli.has` Cli.command "setup" Test.@?= True
     ]