116 lines
4.1 KiB
Python
116 lines
4.1 KiB
Python
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
import aiosqlite
|
|
|
|
from duck_core.tasks.state import TaskState
|
|
|
|
|
|
def utc_now() -> str:
|
|
return datetime.now(UTC).isoformat()
|
|
|
|
|
|
class TaskStore:
|
|
def __init__(self, db_path: str):
|
|
self.db_path = Path(db_path)
|
|
|
|
async def init(self) -> None:
|
|
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
await db.execute(
|
|
"""
|
|
create table if not exists tasks (
|
|
task_id text primary key,
|
|
status text not null,
|
|
user_message text not null,
|
|
workspace text,
|
|
debug integer not null default 0,
|
|
final_response text,
|
|
created_at text not null,
|
|
updated_at text not null
|
|
)
|
|
"""
|
|
)
|
|
await db.commit()
|
|
|
|
async def create_task(self, user_message: str, workspace: str | None, debug: bool) -> TaskState:
|
|
await self.init()
|
|
now = utc_now()
|
|
task_id = f"task_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}_{uuid4().hex[:8]}"
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
await db.execute(
|
|
"""
|
|
insert into tasks(task_id, status, user_message, workspace, debug, created_at, updated_at)
|
|
values (?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(task_id, "running", user_message, workspace, int(debug), now, now),
|
|
)
|
|
await db.commit()
|
|
return TaskState(
|
|
task_id=task_id,
|
|
status="running",
|
|
user_message=user_message,
|
|
workspace=workspace,
|
|
debug=debug,
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
|
|
async def update_status(
|
|
self, task_id: str, status: str, final_response: str | None = None
|
|
) -> None:
|
|
await self.init()
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
await db.execute(
|
|
"""
|
|
update tasks
|
|
set status = ?, final_response = coalesce(?, final_response), updated_at = ?
|
|
where task_id = ?
|
|
""",
|
|
(status, final_response, utc_now(), task_id),
|
|
)
|
|
await db.commit()
|
|
|
|
async def complete_task(self, task_id: str, final_response: str) -> None:
|
|
await self.update_status(task_id, "completed", final_response)
|
|
|
|
async def fail_task(self, task_id: str, message: str) -> None:
|
|
await self.update_status(task_id, "failed", message)
|
|
|
|
async def cancel_task(self, task_id: str) -> None:
|
|
await self.update_status(task_id, "cancelled")
|
|
|
|
async def waiting_for_approval(self, task_id: str) -> None:
|
|
await self.update_status(task_id, "waiting_for_approval")
|
|
|
|
async def get_task(self, task_id: str) -> TaskState | None:
|
|
await self.init()
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
cursor = await db.execute("select * from tasks where task_id = ?", (task_id,))
|
|
row = await cursor.fetchone()
|
|
return self._row_to_task(row) if row else None
|
|
|
|
async def list_tasks(self, limit: int = 50) -> list[TaskState]:
|
|
await self.init()
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
cursor = await db.execute(
|
|
"select * from tasks order by created_at desc limit ?", (limit,)
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return [self._row_to_task(row) for row in rows]
|
|
|
|
def _row_to_task(self, row: aiosqlite.Row) -> TaskState:
|
|
return TaskState(
|
|
task_id=row["task_id"],
|
|
status=row["status"],
|
|
user_message=row["user_message"],
|
|
workspace=row["workspace"],
|
|
debug=bool(row["debug"]),
|
|
final_response=row["final_response"],
|
|
created_at=row["created_at"],
|
|
updated_at=row["updated_at"],
|
|
)
|