1 files changed, 69 insertions, 0 deletions
diff --git a/Omni/Task/Core.hs b/Omni/Task/Core.hs
index ffecd60..5af4ce4 100644
--- a/Omni/Task/Core.hs
+++ b/Omni/Task/Core.hs
@@ -374,6 +374,75 @@ initTaskDb = do
       \ tokens_used INTEGER, \
       \ FOREIGN KEY (task_id) REFERENCES tasks(id) \
       \)"
+    runMigrations conn
+
+-- | Run schema migrations to add missing columns to existing tables
+runMigrations :: SQL.Connection -> IO ()
+runMigrations conn = do
+  migrateTable conn "task_activity" taskActivityColumns
+  migrateTable conn "tasks" tasksColumns
+  migrateTable conn "retry_context" retryContextColumns
+
+-- | Expected columns for task_activity table (name, type, nullable)
+taskActivityColumns :: [(Text, Text)]
+taskActivityColumns =
+  [ ("id", "INTEGER"),
+    ("task_id", "TEXT"),
+    ("timestamp", "DATETIME"),
+    ("stage", "TEXT"),
+    ("message", "TEXT"),
+    ("metadata", "TEXT"),
+    ("amp_thread_url", "TEXT"),
+    ("started_at", "DATETIME"),
+    ("completed_at", "DATETIME"),
+    ("cost_cents", "INTEGER"),
+    ("tokens_used", "INTEGER")
+  ]
+
+-- | Expected columns for tasks table
+tasksColumns :: [(Text, Text)]
+tasksColumns =
+  [ ("id", "TEXT"),
+    ("title", "TEXT"),
+    ("type", "TEXT"),
+    ("parent", "TEXT"),
+    ("namespace", "TEXT"),
+    ("status", "TEXT"),
+    ("priority", "TEXT"),
+    ("dependencies", "TEXT"),
+    ("description", "TEXT"),
+    ("created_at", "TIMESTAMP"),
+    ("updated_at", "TIMESTAMP")
+  ]
+
+-- | Expected columns for retry_context table
+retryContextColumns :: [(Text, Text)]
+retryContextColumns =
+  [ ("task_id", "TEXT"),
+    ("original_commit", "TEXT"),
+    ("conflict_files", "TEXT"),
+    ("attempt", "INTEGER"),
+    ("reason", "TEXT")
+  ]
+
+-- | Migrate a table by adding any missing columns
+migrateTable :: SQL.Connection -> Text -> [(Text, Text)] -> IO ()
+migrateTable conn tableName expectedCols = do
+  existingCols <- getTableColumns conn tableName
+  let missingCols = filter (\(name, _) -> name `notElem` existingCols) expectedCols
+  traverse_ (addColumn conn tableName) missingCols
+
+-- | Get list of column names for a table using PRAGMA table_info
+getTableColumns :: SQL.Connection -> Text -> IO [Text]
+getTableColumns conn tableName = do
+  rows <- SQL.query conn "PRAGMA table_info(?)" (SQL.Only tableName) :: IO [(Int, Text, Text, Int, Maybe Text, Int)]
+  pure [colName | (_, colName, _, _, _, _) <- rows]
+
+-- | Add a column to a table
+addColumn :: SQL.Connection -> Text -> (Text, Text) -> IO ()
+addColumn conn tableName (colName, colType) = do
+  let sql = "ALTER TABLE " <> tableName <> " ADD COLUMN " <> colName <> " " <> colType
+  SQL.execute_ conn (SQL.Query sql)
 
 -- Generate a sequential task ID (t-1, t-2, t-3, ...)
 generateId :: IO Text