kairo-code 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

kairo/backend/services/task_service.py

@@ -0,0 +1,119 @@
+import logging
+from datetime import datetime, UTC
+
+from sqlalchemy import select, text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from backend.models.agent import Agent
+from backend.models.task import Task
+
+logger = logging.getLogger(__name__)
+
+
+class TaskService:
+    def __init__(self, db: AsyncSession):
+        self.db = db
+
+    async def create_task(self, user_id: str, input_data: dict, agent_id: str | None = None) -> Task:
+        """Create a new task for a user, optionally targeting a specific agent."""
+        # If agent_id specified, verify the agent belongs to this user
+        if agent_id:
+            stmt = select(Agent).where(Agent.id == agent_id, Agent.user_id == user_id)
+            result = await self.db.execute(stmt)
+            if not result.scalar_one_or_none():
+                raise ValueError("Agent not found or not owned by user")
+
+        task = Task(
+            user_id=user_id,
+            agent_id=agent_id,
+            input=input_data,
+            status="pending",
+        )
+        self.db.add(task)
+        await self.db.commit()
+        await self.db.refresh(task)
+        logger.info("Task created: id=%s user=%s agent=%s", task.id, user_id, agent_id)
+        return task
+
+    async def claim_task(self, agent_id: str, user_id: str) -> Task | None:
+        """Claim the oldest pending task for this user's agent.
+
+        Uses FOR UPDATE SKIP LOCKED to prevent race conditions
+        when multiple agents poll simultaneously.
+        """
+        # Verify agent belongs to this user
+        agent_stmt = select(Agent).where(Agent.id == agent_id, Agent.user_id == user_id)
+        agent_result = await self.db.execute(agent_stmt)
+        agent = agent_result.scalar_one_or_none()
+        if not agent:
+            return None
+
+        # Find oldest pending task: either targeted to this agent or unassigned
+        stmt = (
+            select(Task)
+            .where(
+                Task.user_id == user_id,
+                Task.status == "pending",
+                (Task.agent_id == agent_id) | (Task.agent_id.is_(None)),
+            )
+            .order_by(Task.created_at.asc())
+            .limit(1)
+            .with_for_update(skip_locked=True)
+        )
+        result = await self.db.execute(stmt)
+        task = result.scalar_one_or_none()
+
+        if not task:
+            return None
+
+        task.agent_id = agent_id
+        task.status = "running"
+        task.claimed_at = datetime.now(UTC)
+        task.updated_at = datetime.now(UTC)
+        await self.db.commit()
+        await self.db.refresh(task)
+        logger.info("Task claimed: id=%s agent=%s", task.id, agent_id)
+        return task
+
+    async def update_task(
+        self, task_id: str, user_id: str, status: str,
+        output: dict | None = None, error: str | None = None,
+    ) -> Task | None:
+        """Update a task's status, output, or error."""
+        stmt = select(Task).where(Task.id == task_id, Task.user_id == user_id)
+        result = await self.db.execute(stmt)
+        task = result.scalar_one_or_none()
+        if not task:
+            return None
+
+        task.status = status
+        task.updated_at = datetime.now(UTC)
+
+        if output is not None:
+            task.output = output
+        if error is not None:
+            task.error = error
+        if status in ("completed", "failed"):
+            task.completed_at = datetime.now(UTC)
+
+        await self.db.commit()
+        await self.db.refresh(task)
+        logger.info("Task updated: id=%s status=%s", task_id, status)
+        return task
+
+    async def list_tasks(
+        self, user_id: str, status: str | None = None, limit: int = 50,
+    ) -> list[Task]:
+        """List tasks for a user, optionally filtered by status."""
+        stmt = select(Task).where(Task.user_id == user_id)
+        if status:
+            stmt = stmt.where(Task.status == status)
+        stmt = stmt.order_by(Task.created_at.desc()).limit(limit)
+        result = await self.db.execute(stmt)
+        return list(result.scalars().all())
+
+    async def get_task(self, task_id: str, user_id: str) -> Task | None:
+        """Get a single task by ID, scoped to user."""
+        stmt = select(Task).where(Task.id == task_id, Task.user_id == user_id)
+        result = await self.db.execute(stmt)
+        return result.scalar_one_or_none()

kairo/backend/tests/__init__.py

@@ -0,0 +1 @@
+# Tests package

kairo/backend/tests/e2e/__init__.py

@@ -0,0 +1 @@
+# End-to-end tests package

kairo/backend/tests/e2e/agents/__init__.py

@@ -0,0 +1 @@
+# Agent end-to-end tests package

kairo/backend/tests/e2e/agents/conftest.py

@@ -0,0 +1,389 @@
+"""
+Shared fixtures for agent end-to-end tests.
+
+Provides test database setup, authenticated clients, and factory functions
+for creating test data.
+"""
+
+import asyncio
+import hashlib
+import secrets
+import uuid
+from datetime import datetime, timedelta, UTC
+from typing import AsyncGenerator, Generator
+
+import pytest
+import pytest_asyncio
+from argon2 import PasswordHasher
+from httpx import ASGITransport, AsyncClient
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+
+from backend.app import create_app
+from backend.config import settings
+from backend.core.database import Base, get_db
+from backend.models.agent import (
+    Agent,
+    AgentAlertConfig,
+    AgentAlertHistory,
+    AgentCommand,
+    AgentEvent,
+    AgentMetrics1m,
+    AgentSetupToken,
+)
+from backend.models.api_key import ApiKey
+from backend.models.user import User
+
+
+# Use a test database URL (modify for actual test environment)
+TEST_DATABASE_URL = settings.DATABASE_URL.replace("/kairo", "/kairo_test")
+
+# Create test engine and session
+test_engine = create_async_engine(TEST_DATABASE_URL, echo=False)
+TestAsyncSession = async_sessionmaker(test_engine, class_=AsyncSession, expire_on_commit=False)
+
+
+@pytest.fixture(scope="session")
+def event_loop() -> Generator:
+    """Create an instance of the default event loop for the test session."""
+    loop = asyncio.get_event_loop_policy().new_event_loop()
+    yield loop
+    loop.close()
+
+
+@pytest_asyncio.fixture(scope="function")
+async def db_session() -> AsyncGenerator[AsyncSession, None]:
+    """Create a fresh database session for each test."""
+    async with test_engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+
+    async with TestAsyncSession() as session:
+        yield session
+        # Rollback any uncommitted changes
+        await session.rollback()
+
+    # Clean up tables after test
+    async with test_engine.begin() as conn:
+        await conn.run_sync(Base.metadata.drop_all)
+
+
+@pytest_asyncio.fixture
+async def app(db_session: AsyncSession):
+    """Create a test application with dependency overrides."""
+    application = create_app()
+
+    # Override the database dependency
+    async def override_get_db():
+        yield db_session
+
+    application.dependency_overrides[get_db] = override_get_db
+
+    # Enable agents feature for testing
+    original_setting = settings.FEATURE_KAIRO_AGENTS_ENABLED
+    settings.FEATURE_KAIRO_AGENTS_ENABLED = True
+
+    yield application
+
+    # Restore original setting
+    settings.FEATURE_KAIRO_AGENTS_ENABLED = original_setting
+    application.dependency_overrides.clear()
+
+
+@pytest_asyncio.fixture
+async def client(app) -> AsyncGenerator[AsyncClient, None]:
+    """Create an async HTTP client for testing."""
+    transport = ASGITransport(app=app)
+    async with AsyncClient(transport=transport, base_url="http://test") as ac:
+        yield ac
+
+
+@pytest_asyncio.fixture
+async def test_user(db_session: AsyncSession) -> User:
+    """Create a test user."""
+    ph = PasswordHasher()
+    user = User(
+        id=str(uuid.uuid4()),
+        email=f"test_{uuid.uuid4().hex[:8]}@example.com",
+        hashed_password=ph.hash("testpassword123"),
+        email_verified=True,
+        status="active",
+        plan="pro",
+    )
+    db_session.add(user)
+    await db_session.commit()
+    await db_session.refresh(user)
+    return user
+
+
+@pytest_asyncio.fixture
+async def auth_token(test_user: User) -> str:
+    """Generate a valid JWT token for the test user."""
+    from backend.core.security import create_access_token
+    return create_access_token(test_user.id)
+
+
+@pytest_asyncio.fixture
+async def auth_headers(auth_token: str) -> dict:
+    """Return authorization headers for authenticated requests."""
+    return {"Authorization": f"Bearer {auth_token}"}
+
+
+@pytest_asyncio.fixture
+async def test_agent(db_session: AsyncSession, test_user: User) -> Agent:
+    """Create a test agent."""
+    agent = Agent(
+        id=str(uuid.uuid4()),
+        user_id=test_user.id,
+        name="Test Agent",
+        description="A test agent for e2e tests",
+        model_preference="nyx",
+        agent_type="sdk",
+        state="created",
+        status="offline",
+    )
+    db_session.add(agent)
+    await db_session.commit()
+    await db_session.refresh(agent)
+    return agent
+
+
+@pytest_asyncio.fixture
+async def online_agent(db_session: AsyncSession, test_user: User) -> Agent:
+    """Create an online test agent with connection info."""
+    now = datetime.now(UTC)
+    agent = Agent(
+        id=str(uuid.uuid4()),
+        user_id=test_user.id,
+        name="Online Test Agent",
+        description="An online agent for testing",
+        model_preference="nyx",
+        agent_type="sdk",
+        state="online",
+        status="online",
+        sdk_version="1.0.0",
+        host_info={"hostname": "test-host", "os": "linux"},
+        first_connected_at=now - timedelta(hours=1),
+        last_online_at=now,
+        last_heartbeat_at=now,
+    )
+    db_session.add(agent)
+    await db_session.commit()
+    await db_session.refresh(agent)
+    return agent
+
+
+@pytest_asyncio.fixture
+async def test_api_key(db_session: AsyncSession, test_user: User, test_agent: Agent) -> tuple[str, ApiKey]:
+    """Create an API key for agent authentication."""
+    key_raw = secrets.token_urlsafe(32)
+    key = f"sk-kairo-{key_raw}"
+    key_prefix = key[:20]
+    key_hash = hashlib.sha256(key.encode()).hexdigest()
+
+    api_key = ApiKey(
+        id=str(uuid.uuid4()),
+        user_id=test_user.id,
+        agent_id=test_agent.id,
+        name=f"Agent: {test_agent.name}",
+        key_prefix=key_prefix,
+        key_hash=key_hash,
+        key_type="agent",
+        is_active=True,
+    )
+    db_session.add(api_key)
+    await db_session.commit()
+    await db_session.refresh(api_key)
+    return key, api_key
+
+
+@pytest_asyncio.fixture
+async def api_key_headers(test_api_key: tuple[str, ApiKey]) -> dict:
+    """Return authorization headers for API key authentication."""
+    key, _ = test_api_key
+    return {"Authorization": f"Bearer {key}"}
+
+
+class AgentFactory:
+    """Factory for creating agent test data."""
+
+    def __init__(self, db_session: AsyncSession, user_id: str):
+        self.db = db_session
+        self.user_id = user_id
+
+    async def create_agent(
+        self,
+        name: str = "Test Agent",
+        state: str = "created",
+        **kwargs,
+    ) -> Agent:
+        """Create an agent with specified attributes."""
+        agent = Agent(
+            id=str(uuid.uuid4()),
+            user_id=self.user_id,
+            name=name,
+            state=state,
+            status=kwargs.get("status", "offline"),
+            model_preference=kwargs.get("model_preference", "nyx"),
+            agent_type=kwargs.get("agent_type", "sdk"),
+            description=kwargs.get("description"),
+            system_prompt=kwargs.get("system_prompt"),
+            sdk_version=kwargs.get("sdk_version"),
+            host_info=kwargs.get("host_info"),
+            last_heartbeat_at=kwargs.get("last_heartbeat_at"),
+            last_online_at=kwargs.get("last_online_at"),
+            first_connected_at=kwargs.get("first_connected_at"),
+        )
+        self.db.add(agent)
+        await self.db.commit()
+        await self.db.refresh(agent)
+        return agent
+
+    async def create_setup_token(
+        self,
+        agent_id: str,
+        ttl_minutes: int = 15,
+        used: bool = False,
+    ) -> tuple[str, AgentSetupToken]:
+        """Create a setup token for an agent."""
+        token = f"kairo_setup_{secrets.token_urlsafe(32)}"
+        token_hash = hashlib.sha256(token.encode()).hexdigest()
+        expires_at = datetime.now(UTC) + timedelta(minutes=ttl_minutes)
+
+        setup_token = AgentSetupToken(
+            id=str(uuid.uuid4()),
+            agent_id=agent_id,
+            user_id=self.user_id,
+            token_hash=token_hash,
+            expires_at=expires_at,
+            used=used,
+        )
+        self.db.add(setup_token)
+        await self.db.commit()
+        await self.db.refresh(setup_token)
+        return token, setup_token
+
+    async def create_command(
+        self,
+        agent_id: str,
+        command_type: str = "restart",
+        status: str = "pending",
+        **kwargs,
+    ) -> AgentCommand:
+        """Create a command for an agent."""
+        now = datetime.now(UTC)
+        command = AgentCommand(
+            id=str(uuid.uuid4()),
+            agent_id=agent_id,
+            command_type=command_type,
+            status=status,
+            issued_by=kwargs.get("issued_by", self.user_id),
+            payload=kwargs.get("payload"),
+            expires_at=kwargs.get("expires_at", now + timedelta(minutes=5)),
+        )
+        self.db.add(command)
+        await self.db.commit()
+        await self.db.refresh(command)
+        return command
+
+    async def create_event(
+        self,
+        agent_id: str,
+        event_type: str = "heartbeat",
+        **kwargs,
+    ) -> AgentEvent:
+        """Create an event for an agent."""
+        event = AgentEvent(
+            id=str(uuid.uuid4()),
+            agent_id=agent_id,
+            event_type=event_type,
+            event_data=kwargs.get("event_data"),
+            error_type=kwargs.get("error_type"),
+            error_message=kwargs.get("error_message"),
+            client_ip=kwargs.get("client_ip"),
+        )
+        self.db.add(event)
+        await self.db.commit()
+        await self.db.refresh(event)
+        return event
+
+    async def create_metrics(
+        self,
+        agent_id: str,
+        bucket_time: datetime | None = None,
+        **kwargs,
+    ) -> AgentMetrics1m:
+        """Create metrics for an agent."""
+        if bucket_time is None:
+            bucket_time = datetime.now(UTC).replace(second=0, microsecond=0)
+
+        metrics = AgentMetrics1m(
+            id=str(uuid.uuid4()),
+            agent_id=agent_id,
+            bucket_time=bucket_time,
+            request_count=kwargs.get("request_count", 10),
+            error_count=kwargs.get("error_count", 1),
+            input_tokens=kwargs.get("input_tokens", 1000),
+            output_tokens=kwargs.get("output_tokens", 500),
+            total_latency_ms=kwargs.get("total_latency_ms", 5000),
+            tool_calls=kwargs.get("tool_calls", 2),
+        )
+        self.db.add(metrics)
+        await self.db.commit()
+        await self.db.refresh(metrics)
+        return metrics
+
+    async def create_alert_config(
+        self,
+        agent_id: str,
+        name: str = "Test Alert",
+        **kwargs,
+    ) -> AgentAlertConfig:
+        """Create an alert configuration for an agent."""
+        config = AgentAlertConfig(
+            id=str(uuid.uuid4()),
+            agent_id=agent_id,
+            user_id=self.user_id,
+            name=name,
+            alert_type=kwargs.get("alert_type", "error_rate"),
+            metric=kwargs.get("metric", "error_rate"),
+            condition=kwargs.get("condition", "gt"),
+            threshold=kwargs.get("threshold", 10.0),
+            window_seconds=kwargs.get("window_seconds", 300),
+            cooldown_seconds=kwargs.get("cooldown_seconds", 3600),
+            severity=kwargs.get("severity", "warning"),
+            channels=kwargs.get("channels", ["email"]),
+            is_enabled=kwargs.get("is_enabled", True),
+        )
+        self.db.add(config)
+        await self.db.commit()
+        await self.db.refresh(config)
+        return config
+
+    async def create_alert_history(
+        self,
+        config_id: str,
+        agent_id: str,
+        **kwargs,
+    ) -> AgentAlertHistory:
+        """Create an alert history entry."""
+        history = AgentAlertHistory(
+            id=str(uuid.uuid4()),
+            config_id=config_id,
+            agent_id=agent_id,
+            alert_type=kwargs.get("alert_type", "error_rate"),
+            severity=kwargs.get("severity", "warning"),
+            status=kwargs.get("status", "triggered"),
+            message=kwargs.get("message", "Error rate exceeded threshold"),
+            trigger_value=kwargs.get("trigger_value", 15.0),
+            threshold_value=kwargs.get("threshold_value", 10.0),
+        )
+        self.db.add(history)
+        await self.db.commit()
+        await self.db.refresh(history)
+        return history
+
+
+@pytest_asyncio.fixture
+async def agent_factory(db_session: AsyncSession, test_user: User) -> AgentFactory:
+    """Provide an agent factory for creating test data."""
+    return AgentFactory(db_session, test_user.id)