XCClaw/app/services/session_manager.py

import uuid
from datetime import datetime
from app.models.session import SessionType, Task, TaskStatus, CreateSessionRequest
from app.services.opencode_client import opencode_client
from app.services.storage import storage
from app.services.history import history_service
from app.core.logging import logger

TASKS_KEY = "tasks"


class SessionManager:
    def __init__(self):
        self.tasks: dict[str, Task] = {}

    def load(self):
        data = storage.load(TASKS_KEY, {})
        if isinstance(data, dict):
            for k, v in data.items():
                self.tasks[k] = Task(**v)
        logger.info(f"Loaded {len(self.tasks)} tasks")

    def save(self):
        data = {k: v.model_dump() for k, v in self.tasks.items()}
        storage.save(TASKS_KEY, data)

    async def create_task(self, request: CreateSessionRequest) -> Task:
        task_id = str(uuid.uuid4())
        task = Task(
            id=task_id,
            type=request.type,
            prompt=request.prompt or "",
            created_at=datetime.now(),
        )
        self.tasks[task_id] = task
        self.save()
        logger.info(f"Created task: {task_id}")
        return task

    async def execute_task(self, task_id: str) -> dict:
        task = self.tasks.get(task_id)
        if not task:
            raise ValueError(f"Task {task_id} not found")

        task.status = TaskStatus.RUNNING
        task.started_at = datetime.now()
        self.save()

        session = await opencode_client.create_session(
            title=task.prompt[:50] if task.prompt else None
        )
        session_id = session["id"]
        task.session_id = session_id
        self.save()

        logger.info(f"Executing task {task_id} with session {session_id}")

        try:
            result = await opencode_client.send_message(session_id, task.prompt)
            task.status = TaskStatus.COMPLETED
            task.finished_at = datetime.now()
            self.save()
            history_service.add_task_history(task)
            return result
        except Exception as e:
            task.status = TaskStatus.FAILED
            task.finished_at = datetime.now()
            task.error = str(e)
            self.save()
            history_service.add_task_history(task)
            logger.error(f"Task {task_id} failed: {e}")
            raise

    async def execute_task_async(self, task_id: str) -> None:
        task = self.tasks.get(task_id)
        if not task:
            raise ValueError(f"Task {task_id} not found")

        task.status = TaskStatus.RUNNING
        task.started_at = datetime.now()
        self.save()

        session = await opencode_client.create_session(
            title=task.prompt[:50] if task.prompt else None
        )
        session_id = session["id"]
        task.session_id = session_id
        self.save()

        logger.info(f"Executing async task {task_id} with session {session_id}")

        await opencode_client.send_message_async(session_id, task.prompt)

        from app.services.websocket_manager import ws_manager

        await ws_manager.start_poll_task(task_id, session_id)

    async def update_task_status(
        self, task_id: str, status: TaskStatus, error: str | None = None
    ):
        task = self.tasks.get(task_id)
        if task:
            task.status = status
            if status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.ABORTED):
                task.finished_at = datetime.now()
            if error:
                task.error = error
            self.save()
            history_service.add_task_history(task)

    async def abort_task(self, task_id: str) -> bool:
        task = self.tasks.get(task_id)
        if not task or not task.session_id:
            return False

        result = await opencode_client.abort_session(task.session_id)
        task.status = TaskStatus.ABORTED
        task.finished_at = datetime.now()
        self.save()
        history_service.add_task_history(task)
        return result

    async def get_task(self, task_id: str) -> Task | None:
        return self.tasks.get(task_id)

    async def list_tasks(self) -> list[Task]:
        return list(self.tasks.values())


session_manager = SessionManager()