diff --git a/.env.example b/.env.example index 06a7c5a..23bd7e5 100644 --- a/.env.example +++ b/.env.example @@ -23,7 +23,9 @@ SEARCH_API_KEY=your_search_api_key_here FILE_TOOL_BASE_DIR=/tmp/agentkit_sandbox # Security & API Authentication -# Comma-separated list of authorized client API keys +# Comma-separated list of authorized client API keys. +# docker-compose.yml ships no default: leave this unset and the api service +# exits on settings validation instead of starting with a published key. API_KEYS=ak_test_key_12345,ak_admin_key_67890 RATE_LIMIT_PER_MIN=30 @@ -33,3 +35,10 @@ RUN_TIMEOUT_S=60 TOOL_TIMEOUT_S=15 SESSION_TTL_S=3600 TOOL_OUTPUT_MAX_CHARS=2000 + +# Real-infrastructure test suite (optional, test-only) +# `pytest` skips the live tier unless AGENTKIT_REAL_INFRA=1. Point these at a +# DEDICATED test database: the suite truncates runs/steps. +# AGENTKIT_REAL_INFRA=1 +# DATABASE_TEST_URL=postgresql+asyncpg://postgres:postgres@localhost:5432/agentkit_test +# REDIS_TEST_URL=redis://localhost:6379/15 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8ba1f38..9cdcafa 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -39,13 +39,54 @@ jobs: run: mypy agentkit tests test: - name: Test (Python ${{ matrix.python-version }}) + name: Test Hermetic (Python ${{ matrix.python-version }}) runs-on: ubuntu-latest needs: [lint] strategy: fail-fast: false matrix: python-version: ["3.11", "3.12"] + env: + ENVIRONMENT: test + DATABASE_URL: postgresql+asyncpg://agentkit:agentkit_password@localhost:5432/agentkit_test + REDIS_URL: redis://localhost:6379/0 + GEMINI_API_KEY: fake_gemini_key_for_testing + OPENAI_API_KEY: fake_openai_key_for_testing + TAVILY_API_KEY: fake_tavily_key_for_testing + API_KEYS: test_api_key_1,test_api_key_2 + JWT_SECRET: test_secret_key_at_least_32_bytes_long + PYTHONPATH: . + steps: + - name: Check out repository + uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: "pip" + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[dev]" + + - name: Run pytest with coverage + run: | + pytest -v --cov=agentkit --cov-report=term-missing --cov-report=xml --cov-fail-under=85 + + - name: Upload coverage artifact + uses: actions/upload-artifact@v4 + if: matrix.python-version == '3.12' + with: + name: coverage-report + path: coverage.xml + + test-real-infra: + name: Test Against Live PostgreSQL and Redis + runs-on: ubuntu-latest + needs: [lint] + # Serialises this suite because every run truncates shared service containers. services: postgres: image: postgres:16-alpine @@ -71,23 +112,19 @@ jobs: --health-retries 5 env: ENVIRONMENT: test - DATABASE_URL: postgresql+asyncpg://agentkit:agentkit_password@localhost:5432/agentkit_test + AGENTKIT_REAL_INFRA: "1" DATABASE_TEST_URL: postgresql+asyncpg://agentkit:agentkit_password@localhost:5432/agentkit_test - REDIS_URL: redis://localhost:6379/0 - GEMINI_API_KEY: fake_gemini_key_for_testing - OPENAI_API_KEY: fake_openai_key_for_testing - TAVILY_API_KEY: fake_tavily_key_for_testing + REDIS_TEST_URL: redis://localhost:6379/15 API_KEYS: test_api_key_1,test_api_key_2 - JWT_SECRET: test_secret_key_at_least_32_bytes_long PYTHONPATH: . steps: - name: Check out repository uses: actions/checkout@v4 - - name: Set up Python ${{ matrix.python-version }} + - name: Set up Python 3.12 uses: actions/setup-python@v5 with: - python-version: ${{ matrix.python-version }} + python-version: "3.12" cache: "pip" - name: Install dependencies @@ -95,16 +132,8 @@ jobs: python -m pip install --upgrade pip pip install -e ".[dev]" - - name: Run Alembic migrations - run: alembic upgrade head - - - name: Run pytest with coverage - run: | - pytest -v --cov=agentkit --cov-report=term-missing --cov-report=xml --cov-fail-under=85 - - - name: Upload coverage artifact - uses: actions/upload-artifact@v4 - if: matrix.python-version == '3.12' - with: - name: coverage-report - path: coverage.xml + # The suite exercises native UUID/JSON/timestamptz columns, FK cascade deletes + # and Lua-backed rate limiting, none of which the SQLite and fakeredis based + # unit tests can validate. + - name: Run real infrastructure tests + run: pytest -v -m real_infra tests/real diff --git a/README.md b/README.md index 74ea551..1c49bc1 100644 --- a/README.md +++ b/README.md @@ -118,8 +118,14 @@ JWT_SECRET=super_secret_jwt_key_at_least_32_bytes_long ``` ### 4. Running with Docker Compose (Recommended) -Launch the API, PostgreSQL, and Redis together with automatic migrations: +Launch the API, PostgreSQL, and Redis together with automatic migrations. + +`docker-compose.yml` ships no default `API_KEYS`. Leave it unset and the api service +fails immediately on settings validation rather than starting with a published key; +`down`, `ps` and `logs` still work either way. + ```bash +export API_KEYS=ak_local_dev_key # or place API_KEYS in your .env docker compose up -d --build # Verify all services are healthy @@ -135,11 +141,20 @@ curl -s http://localhost:8000/health | jq . AgentKit includes a complete multi-tool demo script that seeds a database and combines SQL analysis with mathematical calculation: ```bash -# Completely offline demo (using FakeLLMClient - no API keys required!) +# Against real infrastructure (PostgreSQL + Redis from your .env). +# Run and step records are persisted exactly as they would be in production. +docker compose up -d postgres redis +alembic upgrade head +python examples/demo.py --provider gemini + +# Or completely offline (FakeLLMClient + in-memory SQLite - no services, no API keys) python examples/demo.py --offline +``` -# Or live with Gemini: -python examples/demo.py --provider gemini +The demo prints the run ID; you can then inspect what it wrote: + +```bash +curl -s -H "X-API-Key: $API_KEYS" "http://localhost:8000/runs//trace" | jq . ``` --- @@ -288,7 +303,7 @@ Complete step-by-step instructions, including security group rules, systemd serv AgentKit enforces strict code quality and comprehensive test coverage across all layers: ```bash -# Run the complete test suite +# Run the complete hermetic test suite (no services required) pytest -v # Run with test coverage report (enforcing >=85% threshold) @@ -302,7 +317,43 @@ ruff format --check . mypy agentkit tests examples ``` -All pull requests trigger the automated GitHub Actions CI matrix running on Python 3.11 and 3.12 with live PostgreSQL 16 and Redis 7 service containers. +### Two test tiers + +The default suite is hermetic: SQLite, `fakeredis` and `FakeLLMClient` stand in for +PostgreSQL, Redis and the LLM providers, so `pytest` runs in about two seconds with no +services. A second, opt-in tier runs the same components against **live** PostgreSQL and +Redis, covering behaviour the hermetic tier structurally cannot assert: + +- native `uuid` / `json` / `timestamptz` columns and asyncpg parameter binding +- real foreign-key enforcement and `ON DELETE CASCADE` (SQLite ignores FKs by default) +- Alembic migrations applied through the real PostgreSQL dialect +- genuine Redis TTL expiry and the hiredis codec +- the Lua-backed atomic rate limiter under concurrency + +```bash +# Start the services the tier needs +docker compose up -d postgres redis + +# Run only the real-infrastructure suite +AGENTKIT_REAL_INFRA=1 pytest -m real_infra tests/real + +# Or point it at any other PostgreSQL / Redis +AGENTKIT_REAL_INFRA=1 \ + DATABASE_TEST_URL=postgresql+asyncpg://user:pass@localhost:5432/agentkit_test \ + REDIS_TEST_URL=redis://localhost:6379/15 \ + pytest -m real_infra tests/real +``` + +Without `AGENTKIT_REAL_INFRA=1` these tests skip, so the default `pytest` invocation stays +hermetic. The suite truncates the `runs` and `steps` tables and deletes only +`session:*` / `ratelimit:*` Redis keys, so point it at a dedicated test database. + +### CI + +All pull requests trigger GitHub Actions with three jobs: `Lint & Type Check`; a hermetic +`Test Hermetic` matrix across Python 3.11 and 3.12 with coverage gating; and +`Test Against Live PostgreSQL and Redis`, which runs the `real_infra` tier against real +PostgreSQL 16 and Redis 7 service containers. --- diff --git a/agentkit/api/deps.py b/agentkit/api/deps.py index eb6b23a..b09fea3 100644 --- a/agentkit/api/deps.py +++ b/agentkit/api/deps.py @@ -12,6 +12,7 @@ from agentkit.api.auth import get_current_settings from agentkit.config import Settings from agentkit.core.agent import Agent, AgentConfig +from agentkit.core.errors import ServiceUnavailableError from agentkit.core.pg_trace import PostgresTraceSink from agentkit.core.trace import InMemoryTraceSink, TraceSink from agentkit.llm.base import LLMClient @@ -19,6 +20,8 @@ from agentkit.llm.fake import FakeLLMClient from agentkit.memory.base import InMemoryMemoryStore, MemoryStore from agentkit.memory.pg_store import PostgresMemoryStore +from agentkit.memory.redis_store import RedisMemoryStore +from agentkit.memory.tiered import TieredMemoryStore from agentkit.tools.builtin.calculator import calculator from agentkit.tools.builtin.http_fetch import http_fetch from agentkit.tools.builtin.read_file import read_file @@ -29,21 +32,27 @@ logger = logging.getLogger(__name__) -def get_session_factory( - request: Request, -) -> async_sessionmaker[AsyncSession] | None: +def get_session_factory(request: Request) -> async_sessionmaker[AsyncSession] | None: """Retrieve database session factory from app state.""" return getattr(request.app.state, "session_factory", None) async def get_db_session( - request: Request, + factory: async_sessionmaker[AsyncSession] | None = Depends(get_session_factory), ) -> AsyncIterator[AsyncSession]: - """Provide a scoped database session.""" - factory = get_session_factory(request) - if factory is not None: - async with factory() as session: - yield session + """Yield a scoped database session. + + Fails explicitly when the application has no session factory. Returning without + yielding would leave FastAPI's dependency solver with an exhausted generator and + surface as an opaque ``generator didn't yield`` RuntimeError at request time. + """ + if factory is None: + raise ServiceUnavailableError( + "database", + "Database is not configured. The session factory is absent from application state.", + ) + async with factory() as session: + yield session def get_redis_client(request: Request) -> Redis | None: @@ -58,10 +67,23 @@ def get_memory_store(request: Request) -> PostgresMemoryStore | MemoryStore: return store factory = get_session_factory(request) - if factory is not None: - return PostgresMemoryStore(factory) + redis_client = get_redis_client(request) + settings: Settings = getattr(request.app.state, "settings", None) or get_current_settings( + request + ) + + run_store: PostgresMemoryStore | InMemoryMemoryStore = ( + PostgresMemoryStore(factory) if factory is not None else InMemoryMemoryStore() + ) + + if redis_client is not None: + session_store = RedisMemoryStore(redis_client, ttl_s=settings.SESSION_TTL_S) + return TieredMemoryStore(session_store=session_store, run_store=run_store) + + if isinstance(run_store, PostgresMemoryStore): + return TieredMemoryStore(session_store=InMemoryMemoryStore(), run_store=run_store) - return InMemoryMemoryStore() + return run_store def get_trace_sink(request: Request) -> TraceSink: diff --git a/agentkit/api/main.py b/agentkit/api/main.py index 53fc6b0..75c7383 100644 --- a/agentkit/api/main.py +++ b/agentkit/api/main.py @@ -9,7 +9,6 @@ from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse from redis.asyncio import Redis -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from agentkit.api.routes import discovery_router from agentkit.api.routes import router as runs_router @@ -22,8 +21,10 @@ RateLimitExceededError, RunNotFoundError, RunTimeoutError, + ServiceUnavailableError, ToolNotFoundError, ) +from agentkit.db.session import create_engine, create_session_factory logger = logging.getLogger(__name__) @@ -40,6 +41,8 @@ def map_agentkit_error_status(exc: AgentKitError) -> int: return 429 if isinstance(exc, LLMProviderError): return 502 + if isinstance(exc, ServiceUnavailableError): + return 503 return 400 @@ -48,8 +51,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: """Manage lifecycle of database engine, session factory, and Redis connection pool.""" settings: Settings = getattr(app.state, "settings", None) or get_settings() - engine = create_async_engine(settings.DATABASE_URL, pool_pre_ping=True) - session_factory = async_sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False) + engine = create_engine(settings.DATABASE_URL, pool_pre_ping=True) + session_factory = create_session_factory(engine) redis_client: Redis[Any] = Redis.from_url(settings.REDIS_URL, decode_responses=True) app.state.db_engine = engine diff --git a/agentkit/api/ratelimit.py b/agentkit/api/ratelimit.py index 8042a84..14d0220 100644 --- a/agentkit/api/ratelimit.py +++ b/agentkit/api/ratelimit.py @@ -6,9 +6,11 @@ import math import time import uuid +from typing import cast from fastapi import Depends from redis.asyncio import Redis +from redis.exceptions import ResponseError from agentkit.api.auth import get_current_settings, verify_api_key from agentkit.api.deps import get_redis_client @@ -17,6 +19,44 @@ logger = logging.getLogger(__name__) +#: Sliding-window admission script, executed atomically inside Redis. +#: +#: The trim, the cardinality check and the insert have to happen as one indivisible +#: step. Performing them as separate round trips leaves a window between reading the +#: count and writing the new member in which concurrent requests all observe the same +#: pre-insert cardinality and are admitted together, which lets a client exceed its +#: quota simply by sending requests in parallel. +#: +#: KEYS: 1 - the sorted set tracking this client's window. +#: ARGV: 1 - current unix timestamp, 2 - window length in seconds, 3 - request limit, +#: 4 - unique member for this request, 5 - key TTL in seconds. +#: Returns: {admitted (1/0), retry_after_seconds}. +SLIDING_WINDOW_SCRIPT = """ +local key = KEYS[1] +local now = tonumber(ARGV[1]) +local window = tonumber(ARGV[2]) +local limit = tonumber(ARGV[3]) +local member = ARGV[4] +local key_ttl = tonumber(ARGV[5]) + +redis.call('ZREMRANGEBYSCORE', key, '-inf', now - window) +local count = redis.call('ZCARD', key) + +if count >= limit then + local oldest = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES') + local retry_after = key_ttl + if oldest[2] then + retry_after = math.ceil(tonumber(oldest[2]) + window - now) + if retry_after < 1 then retry_after = 1 end + end + return {0, retry_after} +end + +redis.call('ZADD', key, now, member) +redis.call('EXPIRE', key, key_ttl) +return {1, 0} +""" + class RateLimiter: """Sliding window rate limiter enforcing request volume limits per client key. @@ -50,36 +90,25 @@ async def check( if redis_client is not None: redis_key = f"{self.prefix}{key}" - clear_before = now - effective_window - - pipe = redis_client.pipeline(transaction=True) - pipe.zremrangebyscore(redis_key, "-inf", clear_before) - pipe.zcard(redis_key) - results = await pipe.execute() - current_count = int(results[1]) - - if current_count >= effective_limit: - oldest_entries = await redis_client.zrange(redis_key, 0, 0, withscores=True) - if oldest_entries: - oldest_score = float(oldest_entries[0][1]) - retry_after = max(1, math.ceil(oldest_score + effective_window - now)) - else: - retry_after = int(effective_window) + member = f"{now}:{uuid.uuid4().hex[:8]}" + admitted, retry_after = await self._check_window( + redis_client=redis_client, + redis_key=redis_key, + now=now, + window=effective_window, + limit=effective_limit, + member=member, + ) + + if not admitted: logger.warning( - "Rate limit exceeded for key %s (count: %d, limit: %d). Retry after: %ds", + "Rate limit exceeded for key %s (limit: %d). Retry after: %ds", key[:8] + "...", - current_count, effective_limit, retry_after, ) raise RateLimitExceededError(retry_after=retry_after) - - member = f"{now}:{uuid.uuid4().hex[:8]}" - pipe = redis_client.pipeline(transaction=True) - pipe.zadd(redis_key, {member: now}) - pipe.expire(redis_key, int(effective_window) + 5) - await pipe.execute() else: # In-memory sliding window fallback active_timestamps = [ @@ -100,6 +129,75 @@ async def check( active_timestamps.append(now) self._memory_store[key] = active_timestamps + async def _check_window( + self, + redis_client: Redis, + redis_key: str, + now: float, + window: float, + limit: int, + member: str, + ) -> tuple[bool, int]: + """Run the sliding-window script and return (admitted, retry_after). + + Falls back to a non-atomic equivalent if the Redis deployment does not + permit scripting, since refusing all traffic would be worse than a + best-effort limit. + """ + key_ttl = int(window) + 5 + try: + # redis-py types eval() as `Any | Awaitable[Any]`, which mypy rejects in a + # typed context; the cast keeps the declared return type honest. + result = await redis_client.eval( # type: ignore[no-untyped-call] + SLIDING_WINDOW_SCRIPT, + 1, + redis_key, + now, + window, + limit, + member, + key_ttl, + ) + except ResponseError as exc: + logger.warning( + "Redis scripting unavailable (%s); falling back to non-atomic rate limiting.", + exc, + ) + return await self._check_window_unatomic( + redis_client, redis_key, now, window, limit, member, key_ttl + ) + + admitted, retry_after = cast("tuple[int, int]", tuple(result)) + return bool(admitted), int(retry_after) + + async def _check_window_unatomic( + self, + redis_client: Redis, + redis_key: str, + now: float, + window: float, + limit: int, + member: str, + key_ttl: int, + ) -> tuple[bool, int]: + """Best-effort sliding-window check for servers without scripting support.""" + pipe = redis_client.pipeline(transaction=True) + pipe.zremrangebyscore(redis_key, "-inf", now - window) + pipe.zcard(redis_key) + results = await pipe.execute() + if int(results[1]) >= limit: + oldest = await redis_client.zrange(redis_key, 0, 0, withscores=True) + retry_after = ( + max(1, math.ceil(float(oldest[0][1]) + window - now)) if oldest else int(window) + ) + return False, retry_after + + pipe = redis_client.pipeline(transaction=True) + pipe.zadd(redis_key, {member: now}) + pipe.expire(redis_key, key_ttl) + await pipe.execute() + return True, 0 + async def __call__( self, api_key: str = Depends(verify_api_key), diff --git a/agentkit/core/agent.py b/agentkit/core/agent.py index d240698..5f2c1c3 100644 --- a/agentkit/core/agent.py +++ b/agentkit/core/agent.py @@ -11,7 +11,7 @@ from agentkit.core.errors import DuplicateToolCallLoopError, MaxStepsExceeded, RunTimeoutError from agentkit.core.log import log_context, set_current_step_no from agentkit.core.trace import StepTrace, TraceSink -from agentkit.core.types import RunStatus +from agentkit.core.types import Role, RunStatus from agentkit.llm.base import LLMClient, Message from agentkit.tools.registry import ToolRegistry @@ -112,6 +112,10 @@ async def run( ) async def _finalize_and_return(result: RunResult) -> RunResult: + if session_id and self.memory is not None and hasattr(self.memory, "save_messages"): + with contextlib.suppress(Exception): + await self.memory.save_messages(session_id, messages) + if self.memory is not None and hasattr(self.memory, "update_run"): with contextlib.suppress(Exception): await self.memory.update_run( @@ -124,10 +128,22 @@ async def _finalize_and_return(result: RunResult) -> RunResult: ) return result - messages: list[Message] = [ - Message.system(self.config.system_prompt), - Message.user(goal), - ] + messages: list[Message] = [] + if session_id and self.memory is not None and hasattr(self.memory, "get_messages"): + with contextlib.suppress(Exception): + loaded = await self.memory.get_messages(session_id) + if loaded: + messages = list(loaded) + + if not messages: + messages = [ + Message.system(self.config.system_prompt), + Message.user(goal), + ] + else: + if messages[0].role != Role.SYSTEM: + messages.insert(0, Message.system(self.config.system_prompt)) + messages.append(Message.user(goal)) total_input_tokens = 0 total_output_tokens = 0 @@ -239,6 +255,7 @@ async def _finalize_and_return(result: RunResult) -> RunResult: # Final natural language answer reached duration_ms = max(0, int((time.perf_counter() - start_time) * 1000)) logger.info("Agent run succeeded in %d ms", duration_ms) + messages.append(Message.assistant(response.text)) return await _finalize_and_return( RunResult( run_id=run_id, diff --git a/agentkit/core/errors.py b/agentkit/core/errors.py index 6b8c8ac..b8dd81b 100644 --- a/agentkit/core/errors.py +++ b/agentkit/core/errors.py @@ -153,6 +153,17 @@ def __init__(self, run_id: str | Any) -> None: self.run_id = str(run_id) +class ServiceUnavailableError(AgentKitError): + """Raised when a required backing service is not configured or not reachable.""" + + def __init__(self, service: str, message: str | None = None) -> None: + super().__init__( + message or f"Service '{service}' is unavailable.", + code="SERVICE_UNAVAILABLE", + ) + self.service = service + + # Aliases for convenience ToolNotFoundError = ToolNotFound LLMRateLimitError = RateLimitExceededError diff --git a/agentkit/memory/tiered.py b/agentkit/memory/tiered.py new file mode 100644 index 0000000..5d2b2c6 --- /dev/null +++ b/agentkit/memory/tiered.py @@ -0,0 +1,76 @@ +"""Tiered dual-store persistence combining Redis session caching with PostgreSQL durable storage.""" + +from __future__ import annotations + +import logging +from typing import Any + +from agentkit.core.types import RunStatus +from agentkit.llm.base import Message +from agentkit.memory.base import MemoryStore, RunRecord, SessionStore + +logger = logging.getLogger(__name__) + + +class TieredMemoryStore(MemoryStore): + """Combines an ephemeral SessionStore (e.g. Redis) for conversation histories + with a durable Run repository (e.g. PostgreSQL) for run and step records. + """ + + def __init__( + self, + session_store: SessionStore, + run_store: Any, + ) -> None: + self.session_store = session_store + self.run_store = run_store + + async def get_messages(self, session_id: str) -> list[Message]: + """Load conversation messages for a session from session store.""" + return await self.session_store.get_messages(session_id) + + async def save_messages(self, session_id: str, messages: list[Message]) -> None: + """Persist conversation messages for a session to session store.""" + await self.session_store.save_messages(session_id, messages) + + async def clear_session(self, session_id: str) -> None: + """Clear conversation history for a session from session store.""" + await self.session_store.clear_session(session_id) + + async def create_run( + self, + run_id: str, + goal: str, + session_id: str | None = None, + ) -> None: + """Persist run initiation record in the run store.""" + if hasattr(self.run_store, "create_run"): + await self.run_store.create_run(run_id=run_id, goal=goal, session_id=session_id) + + async def update_run( + self, + run_id: str, + status: RunStatus | str, + final_answer: str | None = None, + failure_reason: str | None = None, + total_input_tokens: int = 0, + total_output_tokens: int = 0, + ) -> None: + """Update existing run state in the run store.""" + if hasattr(self.run_store, "update_run"): + await self.run_store.update_run( + run_id=run_id, + status=status, + final_answer=final_answer, + failure_reason=failure_reason, + total_input_tokens=total_input_tokens, + total_output_tokens=total_output_tokens, + ) + + async def get_run(self, run_id: str) -> RunRecord | None: + """Retrieve run record from the run store.""" + if hasattr(self.run_store, "get_run"): + record = await self.run_store.get_run(run_id) + if isinstance(record, RunRecord): + return record + return None diff --git a/docker-compose.yml b/docker-compose.yml index 34e45b3..f9b6fca 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -71,7 +71,10 @@ services: - DATABASE_URL=${DATABASE_URL:-postgresql+asyncpg://postgres:postgres@postgres:5432/agentkit} - REDIS_URL=${REDIS_URL:-redis://redis:6379/0} - FILE_TOOL_BASE_DIR=/tmp/agentkit_sandbox - - API_KEYS=${API_KEYS:-ak_test_key_12345} + # No default key. An unset or empty API_KEYS leaves the api service to fail at + # startup on Settings validation rather than running with a published key. + # Compose keeps ${VAR:-} so `down`, `ps` and `logs` still work without it. + - API_KEYS=${API_KEYS:-} - LLM_PROVIDER=${LLM_PROVIDER:-gemini} - LLM_MODEL=${LLM_MODEL:-gemini-2.5-flash} - GEMINI_API_KEY=${GEMINI_API_KEY:-} diff --git a/docs/DECISIONS.md b/docs/DECISIONS.md index c40e671..ff3afdc 100644 --- a/docs/DECISIONS.md +++ b/docs/DECISIONS.md @@ -54,3 +54,21 @@ This log records significant architectural, technical, and design decisions made - Tools are idiomatic, typed Python functions. - Argument validation is strictly enforced before tool execution. - Structured validation errors can be fed back to the LLM for self-correction without crashing the runtime. + +--- + +## ADR-005: Two-Tier Test Strategy with Server-Side Atomic Rate Limiting + +- **Date**: 2026-10-01 +- **Status**: Accepted +- **Context**: The suite originally exercised `PostgresMemoryStore`, `PostgresTraceSink` and `RedisMemoryStore` exclusively against in-memory SQLite and `fakeredis`. Both stand-ins behave differently from their targets in ways that matter: SQLite renders `uuid`/`json`/`timestamptz` as `CHAR`/`TEXT`, ignores foreign keys unless `PRAGMA foreign_keys=ON`, and `fakeredis` is a pure-Python reimplementation that never exercises the hiredis codec or real TTL expiry. CI declared PostgreSQL and Redis service containers that no test connected to. A concurrency test written against live Redis then showed the sliding-window rate limiter admitting 19 of 20 simultaneous requests against a limit of 5, because `ZCARD` was read and `ZADD` written in separate round trips with an `await` between them. +- **Decision**: + - Keep the hermetic tier as the default (`pytest`, SQLite + `fakeredis` + `FakeLLMClient`) so the suite stays fast and dependency-free. + - Add an opt-in `real_infra` tier (`tests/real/`, enabled by `AGENTKIT_REAL_INFRA=1`) that runs the same components against live PostgreSQL and Redis, driving Alembic through the real dialect and running the real FastAPI lifespan. + - Enforce the sliding-window limit with a single Lua script so trimming, counting and insertion are one indivisible operation, with a non-atomic fallback only when a deployment forbids scripting. + - Give fakeredis a Lua runtime (`lupa`) so the hermetic tier exercises the same code path as production. + - Split CI into a hermetic matrix job and a dedicated live-infrastructure job so the service containers are actually exercised. +- **Consequences**: + - Dialect-specific and concurrency behaviour is covered rather than assumed; a rate-limit bypass of this class cannot regress unnoticed. + - `pytest` gains an explicit opt-in for external services, documented in the README and `.env.example`. + - The live tier truncates the `runs`/`steps` tables and deletes only `session:*` / `ratelimit:*` keys, so it requires a dedicated test database. diff --git a/docs/PLAN.md b/docs/PLAN.md index eb885b9..457b592 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -141,3 +141,13 @@ This backlog contains granular, focused tasks (15–30 minutes each) structured - **Criteria**: Step-by-step documentation for deploying AgentKit on a single AWS EC2 instance using Docker Compose, including security groups, systemd service configuration, environment secrets management, and SSL termination via Caddy/Nginx. - [x] **M12-T03: Comprehensive Portfolio README (`README.md`)** - **Criteria**: Complete README featuring overview, architecture Mermaid diagram, quickstart guide, 10-line `@tool` tutorial, provider extension guide, sample trace output, design decisions (ADRs summary), and limitations/future roadmap. + +--- + +## Milestone 13: Real Infrastructure Verification +- [x] **M13-T01: Real-Infrastructure Test Tier (`tests/conftest.py`, `tests/real/`)** + - **Criteria**: Opt-in `real_infra` suite running `PostgresMemoryStore`, `PostgresTraceSink`, `RedisMemoryStore` and the rate limiter against live PostgreSQL and Redis; verifies native `uuid`/`json`/`timestamptz` columns, foreign-key cascade deletes, real TTL expiry, and the full FastAPI lifespan end to end. Applies Alembic migrations through the real dialect and skips unless `AGENTKIT_REAL_INFRA=1`. +- [x] **M13-T02: Atomic Sliding-Window Rate Limiting (`agentkit/api/ratelimit.py`)** + - **Criteria**: Trim, cardinality check and insertion execute as one Lua script so concurrent requests cannot exceed the configured limit; non-atomic fallback retained only for deployments that forbid scripting. Regression test in `tests/real/test_redis_real.py`. +- [x] **M13-T03: CI Split and Documentation Corrections (`.github/workflows/ci.yml`, `README.md`, `docs/DECISIONS.md`)** + - **Criteria**: CI runs the live tier against its PostgreSQL and Redis service containers; README documents both test tiers; ADR-005 records the rationale; compose requires an explicit `API_KEYS`. diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index ef0c20a..ed7f7ba 100644 --- a/docs/PROGRESS.md +++ b/docs/PROGRESS.md @@ -299,15 +299,20 @@ This document tracks execution progress across all backlog tasks, including pull - **Commits**: 2 - **Notes**: Authored complete portfolio `README.md` with CI badge roster, architecture Mermaid flowchart illustrating the ReAct loop and dual-persistence layer, quickstart guide with Docker Compose and local setup, 10-line `@tool` creation tutorial with automatic schema introspection, third-party LLM provider extension guide (`AnthropicClient`), sample `StepTrace` JSON telemetry output, ADRs summary (ADR-001 through ADR-004), AWS production deployment references, testing guides, current limitations, and future roadmap. +### M13-T01: Real-Infrastructure Test Tier +- **Status**: Done +- **PR**: #44 +- **Commits**: 1 +- **Notes**: Added `tests/conftest.py` providing opt-in live fixtures and `tests/real/` covering `PostgresMemoryStore`, `PostgresTraceSink`, `RedisMemoryStore`, the rate limiter and a full-lifespan API round trip. The suite applies Alembic migrations through the real PostgreSQL dialect, verifies native `uuid`/`json`/`timestamptz` round trips, `ON DELETE CASCADE`, FK rejection of orphan steps, genuine Redis TTL expiry and the hiredis codec. Skips unless `AGENTKIT_REAL_INFRA=1`; truncates `runs`/`steps` and deletes only `session:*` / `ratelimit:*` keys. +### M13-T02: Atomic Sliding-Window Rate Limiting +- **Status**: Done +- **PR**: #44 +- **Commits**: 1 +- **Notes**: A concurrency test against live Redis revealed the limiter admitted 19 of 20 simultaneous requests against a limit of 5, because `ZCARD` was read and `ZADD` written in separate round trips with an `await` between them. Replaced with a Lua script performing trim, count and insert as one indivisible operation, retaining a non-atomic fallback only for deployments that forbid scripting. Added `lupa` as a dev dependency so fakeredis exercises the same script path as production. - - - - - - - - - - +### M13-T03: CI Split, Compose Hardening and Documentation Corrections +- **Status**: Done +- **PR**: #44 +- **Commits**: 1 +- **Notes**: Split CI into a hermetic Python matrix job and a `Test Against Live PostgreSQL and Redis` job that runs the `real_infra` tier against real service containers. Fixed `get_db_session`, which returned without yielding when no session factory was configured and would have surfaced as an opaque `generator didn't yield` RuntimeError; it now raises `ServiceUnavailableError` mapped to HTTP 503, and `lifespan` reuses the `agentkit.db.session` helpers instead of duplicating engine construction. `docker-compose.yml` no longer defaults `API_KEYS` to a known test key. `examples/demo.py` now persists run and step records to real PostgreSQL and Redis by default. Corrected the README and this file's earlier claims that CI exercised the database and cache. diff --git a/examples/demo.py b/examples/demo.py index bc894a0..d550da6 100644 --- a/examples/demo.py +++ b/examples/demo.py @@ -7,7 +7,15 @@ "What was the total revenue from 'Electronics' category in 2024, and what would be the total if we applied an 8.5% sales tax?" +By default the demo runs against the infrastructure configured in the environment +(DATABASE_URL / REDIS_URL), so the run and step records it produces are persisted to +PostgreSQL and Redis exactly as they would be in production. The seed script creates +``products`` and ``orders`` tables in the target database. + +Pass ``--offline`` for a fully hermetic run that needs no services and no API keys. + Usage: + docker compose up -d postgres redis python examples/demo.py --offline python examples/demo.py --provider gemini python examples/demo.py --provider openai @@ -18,41 +26,101 @@ import argparse import asyncio import logging +from dataclasses import dataclass from pathlib import Path +from redis.asyncio import Redis from sqlalchemy import text -from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine +from sqlalchemy.ext.asyncio import ( + AsyncEngine, + AsyncSession, + async_sessionmaker, + create_async_engine, +) from agentkit.config import get_settings from agentkit.core.agent import Agent, AgentConfig -from agentkit.core.trace import InMemoryTraceSink +from agentkit.core.pg_trace import PostgresTraceSink +from agentkit.core.trace import InMemoryTraceSink, TraceSink from agentkit.llm.base import LLMClient from agentkit.llm.factory import create_llm_client from agentkit.llm.fake import FakeLLMClient +from agentkit.memory.base import InMemoryMemoryStore, MemoryStore, SessionStore +from agentkit.memory.pg_store import PostgresMemoryStore +from agentkit.memory.redis_store import RedisMemoryStore +from agentkit.memory.tiered import TieredMemoryStore from agentkit.tools.builtin.calculator import calculator from agentkit.tools.builtin.sql_readonly import execute_sql_query from agentkit.tools.registry import ToolRegistry logger = logging.getLogger("agentkit.demo") +IN_MEMORY_DB_URL = "sqlite+aiosqlite:///:memory:" -async def setup_demo_database(db_url: str) -> AsyncEngine: - """Initialize and seed demo database from scripts/seed_demo_db.sql.""" - engine = create_async_engine(db_url) - seed_path = Path(__file__).parents[1] / "scripts" / "seed_demo_db.sql" +@dataclass +class DemoResources: + """Stores and connections owned by a demo run, all closable.""" + + memory: MemoryStore + trace_sink: TraceSink + session_id: str | None + redis: Redis | None = None + engine: AsyncEngine | None = None + + +async def seed_demo_database(engine: AsyncEngine) -> None: + """Apply scripts/seed_demo_db.sql, creating products and orders tables.""" + seed_path = Path(__file__).parents[1] / "scripts" / "seed_demo_db.sql" if not seed_path.exists(): raise FileNotFoundError(f"Seed script not found at {seed_path}") - sql_content = seed_path.read_text(encoding="utf-8") - statements = [stmt.strip() for stmt in sql_content.split(";") if stmt.strip()] - + statements = [s.strip() for s in seed_path.read_text(encoding="utf-8").split(";") if s.strip()] async with engine.begin() as conn: for stmt in statements: await conn.execute(text(stmt)) - logger.info("Demo database seeded successfully with products and orders.") - return engine + logger.info("Seeded products and orders into the demo database.") + + +async def build_persistence(db_url: str, offline: bool) -> DemoResources: + """Select durable or in-memory stores depending on the requested mode.""" + if offline or db_url.startswith("sqlite"): + logger.info("Using in-memory stores (hermetic mode).") + return DemoResources( + memory=InMemoryMemoryStore(), + trace_sink=InMemoryTraceSink(), + session_id=None, + ) + + settings = get_settings() + engine = create_async_engine(db_url, pool_pre_ping=True) + factory = async_sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False) + + # Run records always go to PostgreSQL; the session tier degrades on its own so a + # Redis outage costs conversation history rather than the whole demo. + run_store = PostgresMemoryStore(factory) + session_store: SessionStore + redis_client: Redis | None + try: + redis_client = Redis.from_url(settings.REDIS_URL, decode_responses=True) + await redis_client.ping() + session_store = RedisMemoryStore(redis_client, ttl_s=settings.SESSION_TTL_S) + except Exception as exc: + logger.warning( + "Redis unreachable (%s); runs persist to PostgreSQL but sessions stay in memory.", + exc, + ) + redis_client = None + session_store = InMemoryMemoryStore() + + return DemoResources( + memory=TieredMemoryStore(session_store=session_store, run_store=run_store), + trace_sink=PostgresTraceSink(factory), + session_id="demo-session" if redis_client is not None else None, + redis=redis_client, + engine=engine, + ) def build_fake_llm() -> LLMClient: @@ -84,6 +152,20 @@ def build_fake_llm() -> LLMClient: return client +def build_llm(offline: bool, provider: str) -> LLMClient: + """Resolve the LLM client, degrading to the scripted fake when no key is present.""" + if offline or provider == "fake": + return build_fake_llm() + + settings = get_settings() + key_present = settings.GEMINI_API_KEY if provider == "gemini" else settings.OPENAI_API_KEY + if not key_present: + logger.warning("No API key for provider '%s'; falling back to FakeLLMClient.", provider) + return build_fake_llm() + + return create_llm_client(provider) + + async def run_demo( offline: bool = False, provider: str = "gemini", @@ -92,54 +174,41 @@ async def run_demo( """Run the multi-tool AgentKit demonstration. Args: - offline: If True, uses FakeLLMClient with scripted ReAct turns. - provider: Provider name ('gemini', 'openai', 'fake') if offline=False. - db_url: Optional SQLAlchemy database URL override. + offline: Use FakeLLMClient and in-memory stores with no external services. + provider: LLM provider name ('gemini', 'openai', 'fake') when not offline. + db_url: SQLAlchemy URL override. Defaults to the configured DATABASE_URL, + falling back to in-memory SQLite when offline. Returns: The agent's final answer string. """ settings = get_settings() - target_db_url = db_url or "sqlite+aiosqlite:///:memory:" + target_db_url = IN_MEMORY_DB_URL if offline else (db_url or settings.DATABASE_URL) - # 1. Initialize and seed database - engine = await setup_demo_database(target_db_url) + query_engine = create_async_engine(target_db_url, pool_pre_ping=True) + await seed_demo_database(query_engine) + + resources = await build_persistence(target_db_url, offline) - # 2. Register tools in an isolated registry registry = ToolRegistry() registry.register(calculator) async def execute_demo_sql(query: str) -> str: - """Execute a read-only SQL query against the database and return results as markdown table.""" - return await execute_sql_query(query, engine=engine) + """Execute a read-only SQL query against the demo database.""" + return await execute_sql_query(query, engine=query_engine) registry.register( execute_demo_sql, name="sql_readonly", - description="Execute a read-only SQL query against the database and return results as markdown table.", + description="Execute a read-only SQL query and return the results as a markdown table.", ) - # 3. Choose LLM Client - llm: LLMClient - if offline or provider == "fake": - llm = build_fake_llm() - else: - # Check for provider API key - if provider == "gemini" and not settings.GEMINI_API_KEY: - logger.warning("No GEMINI_API_KEY detected. Falling back to offline FakeLLMClient.") - llm = build_fake_llm() - elif provider == "openai" and not settings.OPENAI_API_KEY: - logger.warning("No OPENAI_API_KEY detected. Falling back to offline FakeLLMClient.") - llm = build_fake_llm() - else: - llm = create_llm_client(provider) - - # 4. Create Agent with trace sink - trace_sink = InMemoryTraceSink() + llm = build_llm(offline, provider) agent = Agent( llm=llm, registry=registry, - trace=trace_sink, + memory=resources.memory, + trace=resources.trace_sink, config=AgentConfig( max_steps=6, run_timeout_s=30.0, @@ -150,23 +219,32 @@ async def execute_demo_sql(query: str) -> str: ), ) - query = ( + question = ( "What was the total revenue from 'Electronics' category in 2024, " "and what would be the total if we applied an 8.5% sales tax?" ) - - logger.info("Executing AgentKit with prompt: %s", query) - result = await agent.run(query) - - logger.info( - "Agent run finished. Status: %s, Steps: %d", - result.status.value, - result.steps_count, - ) - logger.info("Final Answer:\n%s", result.final_answer) - - await engine.dispose() - return result.final_answer or "" + logger.info("Executing AgentKit with prompt: %s", question) + + try: + result = await agent.run(question, session_id=resources.session_id) + logger.info( + "Agent run finished. Status: %s, Steps: %d", + result.status.value, + result.steps_count, + ) + logger.info("Run ID: %s", result.run_id) + logger.info("Final Answer:\n%s", result.final_answer) + return result.final_answer or "" + finally: + await query_engine.dispose() + if resources.engine is not None: + await resources.engine.dispose() + if resources.redis is not None: + aclose = getattr(resources.redis, "aclose", None) + if callable(aclose): + await aclose() + else: + await resources.redis.close() def main() -> None: @@ -178,7 +256,7 @@ def main() -> None: parser.add_argument( "--offline", action="store_true", - help="Run completely offline using FakeLLMClient without API keys", + help="Run hermetically with FakeLLMClient, in-memory SQLite and no services", ) parser.add_argument( "--provider", @@ -191,7 +269,7 @@ def main() -> None: "--db-url", type=str, default=None, - help="Custom SQLAlchemy database URL (defaults to in-memory SQLite)", + help="SQLAlchemy database URL (defaults to DATABASE_URL, or in-memory SQLite when offline)", ) args = parser.parse_args() diff --git a/pyproject.toml b/pyproject.toml index 8d4d5bb..84d46e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,9 @@ dev = [ "types-redis>=4.6.0", "fakeredis>=2.26.0", "aiosqlite>=0.20.0", + # fakeredis only implements EVAL when a Lua runtime is present. The rate limiter + # relies on server-side scripting, so the hermetic tests need lupa to match. + "lupa>=2.0", ] [tool.hatch.build.targets.wheel] @@ -82,3 +85,6 @@ addopts = "-ra -q --strict-markers" testpaths = ["tests"] asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" +markers = [ + "real_infra: requires live PostgreSQL and Redis; set AGENTKIT_REAL_INFRA=1 to run", +] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..dafca8a --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,152 @@ +"""Shared pytest fixtures. + +Most of the suite is hermetic: SQLite, fakeredis and FakeLLMClient stand in for +PostgreSQL, Redis and the LLM providers. That keeps `pytest` fast and free of +external dependencies, but it means those paths never exercise the real drivers. + +The fixtures here back a `real_infra` marked suite that talks to actual +PostgreSQL and Redis instances. They are skipped unless explicitly enabled with +AGENTKIT_REAL_INFRA=1 so the default `pytest` invocation stays hermetic. + + AGENTKIT_REAL_INFRA=1 pytest -m real_infra + +Connection URLs are taken from DATABASE_TEST_URL / REDIS_TEST_URL, falling back +to the application's DATABASE_URL / REDIS_URL. +""" + +from __future__ import annotations + +import os +from collections.abc import AsyncIterator, Iterator + +import pytest +from redis.asyncio import Redis +from sqlalchemy import text +from sqlalchemy.ext.asyncio import ( + AsyncEngine, + AsyncSession, + async_sessionmaker, + create_async_engine, +) + +from agentkit.config import get_settings + +REAL_INFRA_ENV_FLAG = "AGENTKIT_REAL_INFRA" +TRUTHY_VALUES = frozenset({"1", "true", "yes", "on"}) + +#: Only AgentKit-owned key prefixes are removed between tests so pointing the +#: suite at a shared development Redis cannot destroy unrelated data. +MANAGED_REDIS_PATTERNS = ("session:*", "ratelimit:*") + + +def real_infra_enabled() -> bool: + """Return True when the real-infrastructure suite has been opted into.""" + return os.getenv(REAL_INFRA_ENV_FLAG, "").strip().lower() in TRUTHY_VALUES + + +def _skip_real_infra() -> None: + """Skip the calling test unless real infrastructure is enabled.""" + pytest.skip( + f"Real infrastructure suite disabled. Set {REAL_INFRA_ENV_FLAG}=1 to run against " + "live PostgreSQL and Redis." + ) + + +def database_test_url() -> str: + """Resolve the PostgreSQL URL used by the real-infrastructure suite.""" + return ( + os.getenv("DATABASE_TEST_URL") or os.getenv("DATABASE_URL") or get_settings().DATABASE_URL + ) + + +def redis_test_url() -> str: + """Resolve the Redis URL used by the real-infrastructure suite.""" + return os.getenv("REDIS_TEST_URL") or os.getenv("REDIS_URL") or get_settings().REDIS_URL + + +def _upgrade_schema(database_url: str) -> None: + """Apply Alembic migrations against the target database. + + Deliberately drives the real Alembic entrypoint rather than + ``Base.metadata.create_all`` so the migration scripts themselves are covered. + """ + from alembic import command + from alembic.config import Config + + cfg = Config("alembic.ini") + cfg.set_main_option("sqlalchemy.url", database_url) + try: + command.upgrade(cfg, "head") + except Exception as exc: + pytest.fail( + f"Could not apply migrations to {database_url!r}. Is PostgreSQL reachable and " + f"does the database exist? Underlying error: {exc}" + ) + + +@pytest.fixture(scope="session") +def real_database_url() -> Iterator[str]: + """Session-scoped PostgreSQL URL with migrations applied once.""" + if not real_infra_enabled(): + _skip_real_infra() + + url = database_test_url() + if "+asyncpg" not in url: + pytest.fail( + f"Real infrastructure tests require an asyncpg DATABASE_TEST_URL, got: {url!r}. " + "SQLite cannot validate PostgreSQL-specific behaviour." + ) + + _upgrade_schema(url) + yield url + + +@pytest.fixture +async def real_engine(real_database_url: str) -> AsyncIterator[AsyncEngine]: + """Async engine bound to live PostgreSQL, truncated before each test.""" + engine = create_async_engine(real_database_url, pool_pre_ping=True) + async with engine.begin() as conn: + await conn.execute(text("TRUNCATE TABLE steps, runs RESTART IDENTITY CASCADE")) + try: + yield engine + finally: + await engine.dispose() + + +@pytest.fixture +async def real_session_factory( + real_engine: AsyncEngine, +) -> AsyncIterator[async_sessionmaker[AsyncSession]]: + """Session factory bound to live PostgreSQL.""" + yield async_sessionmaker(bind=real_engine, class_=AsyncSession, expire_on_commit=False) + + +@pytest.fixture +async def real_redis() -> AsyncIterator[Redis[str]]: + """Live Redis client with AgentKit-owned keys cleared before each test.""" + if not real_infra_enabled(): + _skip_real_infra() + + client: Redis[str] = Redis.from_url(redis_test_url(), decode_responses=True) + try: + await client.ping() + except Exception as exc: # pragma: no cover - environment problem, not a code path + pytest.fail(f"Could not reach Redis at {redis_test_url()!r}: {exc}") + + await _purge_managed_keys(client) + try: + yield client + finally: + await _purge_managed_keys(client) + aclose = getattr(client, "aclose", None) + if callable(aclose): + await aclose() + else: + await client.close() + + +async def _purge_managed_keys(client: Redis[str]) -> None: + """Delete only the key prefixes AgentKit owns.""" + for pattern in MANAGED_REDIS_PATTERNS: + async for key in client.scan_iter(match=pattern, count=200): + await client.delete(key) diff --git a/tests/integration/test_session_memory.py b/tests/integration/test_session_memory.py new file mode 100644 index 0000000..4789a9f --- /dev/null +++ b/tests/integration/test_session_memory.py @@ -0,0 +1,152 @@ +"""Integration test verifying that conversation session memory persists across multiple runs with the same session_id.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from fakeredis.aioredis import FakeRedis +from fastapi.testclient import TestClient +from redis.asyncio import Redis + +from agentkit.api.deps import get_agent, get_memory_store, get_redis_client +from agentkit.api.main import create_app +from agentkit.config import Settings +from agentkit.core.agent import Agent, AgentConfig +from agentkit.llm.base import LLMResponse, Message +from agentkit.llm.fake import FakeLLMClient +from agentkit.memory.base import InMemoryMemoryStore +from agentkit.memory.redis_store import RedisMemoryStore +from agentkit.tools.registry import ToolRegistry + + +class MemoryInspectingFakeLLM(FakeLLMClient): + """Fake LLM that records received message histories for assertion inspection.""" + + def __init__(self, responses: list[str]) -> None: + super().__init__() + for resp in responses: + self.queue_text(resp) + self.received_histories: list[list[Message]] = [] + + async def chat( + self, + messages: list[Message], + tools: Any = None, + ) -> LLMResponse: + self.received_histories.append([m.model_copy() for m in messages]) + return await super().chat(messages, tools) + + +@pytest.mark.asyncio +async def test_session_memory_persists_across_runs_in_core_agent() -> None: + """Verify Agent.run loads prior messages from session memory and persists new turns.""" + redis: Redis[Any] = FakeRedis(decode_responses=True) + memory_store = RedisMemoryStore(client=redis, ttl_s=3600) + registry = ToolRegistry() + + llm = MemoryInspectingFakeLLM( + responses=[ + "Hello Alice! I have noted your favorite language is Python.", + "Your favorite language is Python!", + ] + ) + + agent = Agent( + llm=llm, + registry=registry, + memory=memory_store, + config=AgentConfig(system_prompt="You are a helpful assistant."), + ) + + session_id = "test-persistent-session-42" + + # Turn 1: User introduces name and preference + res1 = await agent.run( + goal="My name is Alice and my favorite language is Python.", session_id=session_id + ) + assert res1.final_answer == "Hello Alice! I have noted your favorite language is Python." + + # Verify messages saved to Redis session store + saved_messages = await memory_store.get_messages(session_id) + assert len(saved_messages) >= 3 # System prompt, User message, Assistant response + + # Turn 2: User asks a follow-up question referencing turn 1 + res2 = await agent.run(goal="What is my favorite language?", session_id=session_id) + assert res2.final_answer == "Your favorite language is Python!" + + # Verify that during Turn 2, LLM received the history containing Turn 1 + assert len(llm.received_histories) == 2 + turn2_messages = llm.received_histories[1] + + # Must contain Turn 1 user message and assistant answer + user_goals = [m.content for m in turn2_messages if m.role == "user"] + assert "My name is Alice and my favorite language is Python." in user_goals + assert "What is my favorite language?" in user_goals + + close_fn = getattr(redis, "aclose", redis.close) + await close_fn() + + +@pytest.mark.asyncio +async def test_session_memory_persists_across_api_runs() -> None: + """Verify POST /runs with same session_id preserves conversation context across calls.""" + settings = Settings(API_KEYS="secret_test_key") + app = create_app(settings=settings) + + redis: Redis[Any] = FakeRedis(decode_responses=True) + session_store = RedisMemoryStore(client=redis, ttl_s=3600) + run_store = InMemoryMemoryStore() + + # Import TieredMemoryStore to ensure API can combine session & run persistence + from agentkit.memory.tiered import TieredMemoryStore + + tiered_store = TieredMemoryStore(session_store=session_store, run_store=run_store) + + llm = MemoryInspectingFakeLLM( + responses=[ + "Passcode remembered: blue-falcon-42.", + "Your passcode is blue-falcon-42.", + ] + ) + + agent = Agent( + llm=llm, + registry=ToolRegistry(), + memory=tiered_store, + config=AgentConfig(system_prompt="Security Assistant"), + ) + + app.dependency_overrides[get_agent] = lambda: agent + app.dependency_overrides[get_memory_store] = lambda: tiered_store + app.dependency_overrides[get_redis_client] = lambda: redis + + client = TestClient(app, raise_server_exceptions=False) + headers = {"X-API-Key": "secret_test_key"} + session_id = "api_session_alpha" + + # API Call 1 + resp1 = client.post( + "/runs", + json={"goal": "The secret code is blue-falcon-42", "session_id": session_id}, + headers=headers, + ) + assert resp1.status_code == 200 + assert resp1.json()["answer"] == "Passcode remembered: blue-falcon-42." + + # API Call 2 with identical session_id + resp2 = client.post( + "/runs", + json={"goal": "What is the secret code?", "session_id": session_id}, + headers=headers, + ) + assert resp2.status_code == 200 + assert resp2.json()["answer"] == "Your passcode is blue-falcon-42." + + # Check LLM saw the full multi-turn history on turn 2 + assert len(llm.received_histories) == 2 + turn2_contents = [m.content for m in llm.received_histories[1]] + assert "The secret code is blue-falcon-42" in turn2_contents + + close_fn = getattr(redis, "aclose", redis.close) + await close_fn() diff --git a/tests/real/__init__.py b/tests/real/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/real/test_api_e2e_real.py b/tests/real/test_api_e2e_real.py new file mode 100644 index 0000000..7a66e85 --- /dev/null +++ b/tests/real/test_api_e2e_real.py @@ -0,0 +1,217 @@ +"""End-to-end API test against live PostgreSQL and Redis. + +Unlike the hermetic API integration tests, this one runs the real FastAPI +lifespan, so the engine, session factory and Redis pool are constructed exactly +as they are in the Docker Compose service. The only override is the LLM, which +AGENTS.md requires be faked. + +Note on event loops: ``TestClient`` drives the application on its own event loop, +so anything bound to the pytest loop (the engine and Redis fixtures) must only be +awaited from the test body, never from inside a request. +""" + +from __future__ import annotations + +import uuid +from collections.abc import Iterator + +import pytest +from fastapi.testclient import TestClient +from redis.asyncio import Redis +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker + +from agentkit.api.main import create_app +from agentkit.config import Settings +from agentkit.llm.fake import FakeLLMClient +from tests.conftest import database_test_url, redis_test_url + +pytestmark = pytest.mark.real_infra + +API_KEY = "ak_real_infra_key" + + +def _build_settings() -> Settings: + """Settings pointing the app at the live infrastructure under test.""" + return Settings( + DATABASE_URL=database_test_url(), + REDIS_URL=redis_test_url(), + API_KEYS=API_KEY, + RATE_LIMIT_PER_MIN=1000, + LLM_PROVIDER="fake", + MAX_STEPS=5, + RUN_TIMEOUT_S=30, + SESSION_TTL_S=3600, + ) + + +def _scripted_llm() -> FakeLLMClient: + """Fake LLM performing one calculator step, then answering.""" + llm = FakeLLMClient() + llm.queue_tool_call( + tool_name="calculator", + arguments={"expression": "6 * 7"}, + call_id="call_real_1", + text="Multiplying.", + ) + llm.queue_text("The answer is 42.") + return llm + + +@pytest.fixture +def real_client( + real_engine: AsyncEngine, + real_redis: Redis[str], +) -> Iterator[tuple[TestClient, FakeLLMClient]]: + """TestClient whose lifespan owns live PostgreSQL and Redis connections. + + The lifespan opens its own engine and Redis client against the same + endpoints, which is precisely the production code path under test. The + fixtures are requested only to guarantee truncation and key cleanup happen + before the app starts. + """ + del real_engine, real_redis + + app = create_app(settings=_build_settings()) + llm = _scripted_llm() + app.state.llm_client = llm + + with TestClient(app, raise_server_exceptions=False) as client: + yield client, llm + + +@pytest.fixture +def client(real_client: tuple[TestClient, FakeLLMClient]) -> TestClient: + """The TestClient half of the real_client fixture.""" + return real_client[0] + + +@pytest.fixture +def llm(real_client: tuple[TestClient, FakeLLMClient]) -> FakeLLMClient: + """The scripted LLM half of the real_client fixture.""" + return real_client[1] + + +async def test_health_reports_live_dependencies(client: TestClient) -> None: + """The health endpoint must probe the real engine and Redis pool.""" + response = client.get("/health") + + assert response.status_code == 200 + body = response.json() + assert body["status"] == "healthy" + assert body["database"] == "healthy" + assert body["redis"] == "healthy" + + +async def test_run_persists_run_and_steps_to_postgres( + client: TestClient, + real_session_factory: async_sessionmaker[AsyncSession], +) -> None: + """POST /runs must write run and step rows through the real Postgres stores.""" + response = client.post( + "/runs", + json={"goal": "What is 6 times 7?"}, + headers={"X-API-Key": API_KEY}, + ) + + assert response.status_code == 200 + payload = response.json() + assert payload["answer"] == "The answer is 42." + assert payload["status"] == "succeeded" + run_uuid = uuid.UUID(payload["run_id"]) + + async with real_session_factory() as session: + run_row = ( + await session.execute( + text("SELECT goal, status FROM runs WHERE id = :rid"), + {"rid": run_uuid}, + ) + ).one_or_none() + assert run_row is not None + assert run_row.goal == "What is 6 times 7?" + assert run_row.status == "succeeded" + + step_count: int = ( + await session.execute( + text("SELECT count(*) FROM steps WHERE run_id = :rid"), + {"rid": run_uuid}, + ) + ).scalar_one() + assert step_count >= 1 + + trace = client.get(f"/runs/{payload['run_id']}/trace", headers={"X-API-Key": API_KEY}) + assert trace.status_code == 200 + assert any(s["tool_name"] == "calculator" for s in trace.json()["steps"]) + + +async def test_session_history_persists_in_real_redis( + client: TestClient, + real_redis: Redis[str], +) -> None: + """A session_id must cause conversation history to be written to live Redis.""" + session_id = "real-session-42" + response = client.post( + "/runs", + json={"goal": "My name is Alice.", "session_id": session_id}, + headers={"X-API-Key": API_KEY}, + ) + assert response.status_code == 200 + + key = f"session:{session_id}:messages" + assert await real_redis.exists(key) == 1 + assert 0 < await real_redis.ttl(key) <= 3600 + + +async def test_multi_turn_session_reuses_history(client: TestClient, llm: FakeLLMClient) -> None: + """Second turn in the same session must receive the first turn's messages.""" + session_id = "real-session-multiturn" + headers = {"X-API-Key": API_KEY} + + llm.queue_text("Noted, Alice.") + llm.queue_text("Your name is Alice.") + + client.post( + "/runs", + json={"goal": "My name is Alice.", "session_id": session_id}, + headers=headers, + ) + client.post( + "/runs", + json={"goal": "What is my name?", "session_id": session_id}, + headers=headers, + ) + + assert llm.call_count >= 3 + # The final call belongs to turn 2 and must carry turn 1's history. + contents = [m.content for m in llm.history[-1]] + assert "My name is Alice." in contents + assert "What is my name?" in contents + + +async def test_run_details_and_missing_run(client: TestClient) -> None: + """Retrieving a persisted run works; an unknown run yields 404.""" + headers = {"X-API-Key": API_KEY} + + created = client.post("/runs", json={"goal": "Persist me."}, headers=headers) + assert created.status_code == 200 + run_id = created.json()["run_id"] + + fetched = client.get(f"/runs/{run_id}", headers=headers) + assert fetched.status_code == 200 + assert fetched.json()["goal"] == "Persist me." + + assert client.get(f"/runs/{uuid.uuid4()}", headers=headers).status_code == 404 + + +async def test_tools_endpoint_exposes_builtin_registry(client: TestClient) -> None: + """The built-in tool registry is served from the real dependency graph.""" + response = client.get("/tools", headers={"X-API-Key": API_KEY}) + + assert response.status_code == 200 + names = {t["name"] for t in response.json()["tools"]} + assert {"calculator", "read_file", "http_fetch", "sql_readonly", "web_search"} <= names + + +async def test_unauthorized_request_is_rejected(client: TestClient) -> None: + """Auth must be enforced before any store access.""" + assert client.post("/runs", json={"goal": "nope"}).status_code == 401 diff --git a/tests/real/test_postgres_real.py b/tests/real/test_postgres_real.py new file mode 100644 index 0000000..a76a509 --- /dev/null +++ b/tests/real/test_postgres_real.py @@ -0,0 +1,212 @@ +"""PostgreSQL persistence tests executed against a live PostgreSQL server. + +These cover the behaviour that the SQLite-backed unit tests structurally cannot +assert: native UUID binding, JSON column round-tripping, TIMESTAMPTZ timezone +preservation and real foreign-key enforcement. +""" + +from __future__ import annotations + +import uuid +from datetime import UTC, datetime, timedelta + +import pytest +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from agentkit.core.pg_trace import PostgresTraceSink +from agentkit.core.trace import StepTrace +from agentkit.core.types import RunStatus +from agentkit.db.models import Run, Step +from agentkit.memory.pg_store import PostgresMemoryStore + +pytestmark = pytest.mark.real_infra + + +async def test_run_and_step_round_trip_preserves_types( + real_session_factory: async_sessionmaker[AsyncSession], +) -> None: + """UUID, JSON and timestamptz values must survive a real round trip unchanged.""" + store = PostgresMemoryStore(session_factory=real_session_factory) + run_id = str(uuid.uuid4()) + + await store.create_run(run_id=run_id, goal="Round trip probe", session_id="sess-rt") + await store.update_run( + run_id=run_id, + status=RunStatus.SUCCEEDED, + final_answer="done", + total_input_tokens=7, + total_output_tokens=11, + ) + + record = await store.get_run(run_id) + assert record is not None + assert record.id == run_id + assert record.status == RunStatus.SUCCEEDED + assert record.total_input_tokens == 7 + assert record.total_output_tokens == 11 + + # The ORM column is native UUID on PostgreSQL; confirm the stored value is a + # real uuid column rather than text that merely parses. + async with real_session_factory() as session: + raw = await session.execute( + text("SELECT pg_typeof(id)::text FROM runs WHERE id = :rid"), + {"rid": uuid.UUID(run_id)}, + ) + assert raw.scalar_one() == "uuid" + + +async def test_json_columns_round_trip_nested_payloads( + real_session_factory: async_sessionmaker[AsyncSession], +) -> None: + """Nested args/results must round-trip through PostgreSQL's JSON column type.""" + sink = PostgresTraceSink(session_factory=real_session_factory) + run_id = str(uuid.uuid4()) + + args: dict[str, object] = { + "expression": "2 + 2", + "nested": {"list": [1, 2, {"deep": True}], "null_value": None}, + "unicode": "café ✓", + } + result: dict[str, object] = {"output": "4", "meta": {"flags": ["ok", "fast"]}} + created = datetime.now(UTC) + + await sink.record( + StepTrace( + run_id=run_id, + step_no=1, + tool_name="calculator", + args=args, + result=result, + latency_ms=3, + input_tokens=5, + output_tokens=2, + timestamp=created, + ) + ) + + traces = await sink.get_traces(run_id) + assert len(traces) == 1 + assert traces[0].args == args + assert traces[0].result == result + assert traces[0].timestamp == created + + +async def test_trace_sink_creates_missing_parent_run( + real_session_factory: async_sessionmaker[AsyncSession], +) -> None: + """Recording a step without a prior create_run must still satisfy the FK.""" + sink = PostgresTraceSink(session_factory=real_session_factory) + run_id = str(uuid.uuid4()) + + await sink.record( + StepTrace( + run_id=run_id, + step_no=1, + tool_name="read_file", + args={"path": "notes.txt"}, + result={"output": "hello"}, + timestamp=datetime.now(UTC), + ) + ) + + async with real_session_factory() as session: + parent = await session.get(Run, uuid.UUID(run_id)) + assert parent is not None, "PostgresTraceSink must auto-create the parent Run row" + assert parent.status == "running" + + assert len(await sink.get_traces(run_id)) == 1 + + +async def test_foreign_key_cascade_deletes_steps( + real_session_factory: async_sessionmaker[AsyncSession], +) -> None: + """ON DELETE CASCADE must remove steps; SQLite does not enforce this by default.""" + sink = PostgresTraceSink(session_factory=real_session_factory) + run_id = str(uuid.uuid4()) + + for step_no in (1, 2, 3): + await sink.record( + StepTrace( + run_id=run_id, + step_no=step_no, + tool_name="calculator", + args={"expression": str(step_no)}, + result={"output": str(step_no)}, + timestamp=datetime.now(UTC), + ) + ) + assert len(await sink.get_traces(run_id)) == 3 + + async with real_session_factory() as session, session.begin(): + run = await session.get(Run, uuid.UUID(run_id)) + assert run is not None + await session.delete(run) + + async with real_session_factory() as session: + remaining = await session.execute( + text("SELECT count(*) FROM steps WHERE run_id = :rid"), + {"rid": uuid.UUID(run_id)}, + ) + assert remaining.scalar_one() == 0 + + +async def test_orphan_step_is_rejected_by_postgres( + real_session_factory: async_sessionmaker[AsyncSession], +) -> None: + """Inserting a step without a parent run must violate the FK constraint.""" + async with real_session_factory() as session: + with pytest.raises(IntegrityError): + async with session.begin(): + session.add( + Step( + id=uuid.uuid4(), + run_id=uuid.uuid4(), + step_no=1, + tool_name="calculator", + args={"expression": "1"}, + created_at=datetime.now(UTC), + ) + ) + + +async def test_list_runs_filters_by_session_and_orders_newest_first( + real_session_factory: async_sessionmaker[AsyncSession], +) -> None: + """Session filtering and created_at ordering must hold on real PostgreSQL.""" + store = PostgresMemoryStore(session_factory=real_session_factory) + + ids: list[str] = [] + for index in range(3): + run_id = str(uuid.uuid4()) + ids.append(run_id) + await store.create_run(run_id=run_id, goal=f"goal {index}", session_id="sess-list") + # Force a distinct, increasing created_at so the ordering assertion is stable. + async with real_session_factory() as session, session.begin(): + run = await session.get(Run, uuid.UUID(run_id)) + assert run is not None + run.created_at = datetime.now(UTC) + timedelta(seconds=index) + + runs = await store.list_runs(session_id="sess-list") + assert [r.id for r in runs] == list(reversed(ids)) + + other = str(uuid.uuid4()) + await store.create_run(run_id=other, goal="other session", session_id="sess-other") + assert other not in {r.id for r in runs} + + limited = await store.list_runs(session_id="sess-list", limit=2) + assert len(limited) == 2 + + +async def test_non_uuid_run_id_maps_to_deterministic_uuid( + real_session_factory: async_sessionmaker[AsyncSession], +) -> None: + """Non-UUID identifiers fall back to uuid5 so writes still land predictably.""" + store = PostgresMemoryStore(session_factory=real_session_factory) + run_id = "run-without-dashes" + + await store.create_run(run_id=run_id, goal="deterministic mapping") + record = await store.get_run(run_id) + assert record is not None + assert record.id == str(uuid.uuid5(uuid.NAMESPACE_DNS, run_id)) diff --git a/tests/real/test_redis_real.py b/tests/real/test_redis_real.py new file mode 100644 index 0000000..1075755 --- /dev/null +++ b/tests/real/test_redis_real.py @@ -0,0 +1,165 @@ +"""Redis session-store and rate-limiter tests executed against a live Redis server. + +fakeredis is a pure-Python reimplementation; it does not exercise the hiredis +codec, real connection pooling, or genuine TTL expiry. These tests do. +""" + +from __future__ import annotations + +import asyncio + +import pytest +from redis.asyncio import Redis + +from agentkit.api.ratelimit import RateLimiter +from agentkit.core.errors import RateLimitExceededError +from agentkit.core.types import Role +from agentkit.llm.base import Message +from agentkit.memory.redis_store import RedisMemoryStore +from agentkit.tools.models import ToolCall + +pytestmark = pytest.mark.real_infra + + +async def test_session_round_trip_through_real_redis(real_redis: Redis[str]) -> None: + """Messages must survive JSON encoding and the hiredis codec intact.""" + store = RedisMemoryStore(client=real_redis, ttl_s=300) + messages = [ + Message.system("You are a helpful assistant."), + Message.user("Hello!"), + Message.assistant( + content="Hi!", + tool_calls=[ToolCall(id="c1", name="calculator", arguments={"expression": "1+1"})], + ), + Message.tool_result(tool_call_id="c1", content="2"), + ] + + await store.save_messages("sess-real", messages) + loaded = await store.get_messages("sess-real") + + assert len(loaded) == 4 + assert loaded[0].role == Role.SYSTEM + assert loaded[2].tool_calls is not None + assert loaded[2].tool_calls[0].arguments == {"expression": "1+1"} + assert loaded[3].role == Role.TOOL + assert loaded[3].tool_call_id == "c1" + + +async def test_ttl_is_applied_by_real_server(real_redis: Redis[str]) -> None: + """A short TTL must be visible through the server's own TTL command.""" + store = RedisMemoryStore(client=real_redis, ttl_s=5) + await store.save_messages("sess-ttl", [Message.user("hi")]) + + ttl = await real_redis.ttl("session:sess-ttl:messages") + assert 0 < ttl <= 5 + + +async def test_keys_expire_on_real_server(real_redis: Redis[str]) -> None: + """Keys must genuinely disappear once the TTL elapses.""" + store = RedisMemoryStore(client=real_redis, ttl_s=1) + await store.save_messages("sess-expiry", [Message.user("bye")]) + assert await store.get_messages("sess-expiry") + + await asyncio.sleep(1.5) + assert await real_redis.exists("session:sess-expiry:messages") == 0 + assert await store.get_messages("sess-expiry") == [] + + +async def test_sliding_window_truncation_is_persisted(real_redis: Redis[str]) -> None: + """Truncation must be applied on write, not only on read.""" + store = RedisMemoryStore(client=real_redis, ttl_s=300, max_messages=3) + await store.save_messages( + "sess-truncate", + [ + Message.system("System"), + Message.user("one"), + Message.assistant("two"), + Message.user("three"), + Message.assistant("four"), + Message.user("five"), + ], + ) + + loaded = await store.get_messages("sess-truncate") + assert len(loaded) == 4 + assert loaded[0].content == "System" + assert loaded[-1].content == "five" + + raw = await real_redis.get("session:sess-truncate:messages") + assert raw is not None + assert len(raw) < 400 + + +async def test_clear_session_deletes_key(real_redis: Redis[str]) -> None: + """clear_session must remove the key from the real server.""" + store = RedisMemoryStore(client=real_redis, ttl_s=300) + await store.save_messages("sess-clear", [Message.user("x")]) + await store.clear_session("sess-clear") + + assert await real_redis.exists("session:sess-clear:messages") == 0 + + +async def test_corrupt_payload_degrades_gracefully(real_redis: Redis[str]) -> None: + """A non-JSON value must not raise out of the read path.""" + store = RedisMemoryStore(client=real_redis, ttl_s=300) + await real_redis.set("session:sess-corrupt:messages", "{not-json") + + assert await store.get_messages("sess-corrupt") == [] + + +async def test_rate_limiter_blocks_at_limit(real_redis: Redis[str]) -> None: + """The ZSET sliding window must reject the (limit + 1)-th request.""" + limiter = RateLimiter(requests_per_window=3, window_seconds=60) + key = "ak_ratelimit_probe" + + for _ in range(3): + await limiter.check(key=key, redis_client=real_redis, limit=3) + + with pytest.raises(RateLimitExceededError) as excinfo: + await limiter.check(key=key, redis_client=real_redis, limit=3) + assert excinfo.value.retry_after >= 1 + + card = await real_redis.zcard(f"ratelimit:{key}") + assert card == 3 + + +async def test_rate_limiter_window_expiry(real_redis: Redis[str]) -> None: + """Once the window slides past, requests must be admitted again.""" + limiter = RateLimiter(requests_per_window=2, window_seconds=1) + key = "ak_ratelimit_window" + + await limiter.check(key=key, redis_client=real_redis, limit=2) + await limiter.check(key=key, redis_client=real_redis, limit=2) + with pytest.raises(RateLimitExceededError): + await limiter.check(key=key, redis_client=real_redis, limit=2) + + await asyncio.sleep(1.2) + await limiter.check(key=key, redis_client=real_redis, limit=2) + + +async def test_rate_limiter_isolates_keys(real_redis: Redis[str]) -> None: + """One client exhausting its quota must not affect another.""" + limiter = RateLimiter(requests_per_window=1, window_seconds=60) + + await limiter.check(key="ak_first", redis_client=real_redis, limit=1) + with pytest.raises(RateLimitExceededError): + await limiter.check(key="ak_first", redis_client=real_redis, limit=1) + + await limiter.check(key="ak_second", redis_client=real_redis, limit=1) + + +async def test_concurrent_checks_respect_limit(real_redis: Redis[str]) -> None: + """Concurrent pipelined checks must not overshoot the configured limit.""" + limiter = RateLimiter(requests_per_window=5, window_seconds=60) + key = "ak_ratelimit_concurrent" + + outcomes = await asyncio.gather( + *(limiter.check(key=key, redis_client=real_redis, limit=5) for _ in range(20)), + return_exceptions=True, + ) + admitted = sum(1 for o in outcomes if not isinstance(o, BaseException)) + rejected = sum(1 for o in outcomes if isinstance(o, RateLimitExceededError)) + + assert admitted == 5 + assert rejected == 15 + assert await real_redis.zcard(f"ratelimit:{key}") == 5 diff --git a/tests/unit/test_api_deps.py b/tests/unit/test_api_deps.py new file mode 100644 index 0000000..deb41c0 --- /dev/null +++ b/tests/unit/test_api_deps.py @@ -0,0 +1,82 @@ +"""Unit tests for FastAPI dependency resolution in agentkit.api.deps.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from fastapi import APIRouter, Depends, FastAPI +from fastapi.testclient import TestClient +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from agentkit.api.deps import get_db_session +from agentkit.api.main import create_app +from agentkit.config import Settings +from agentkit.core.errors import ServiceUnavailableError + + +def _app() -> FastAPI: + """Build the real app plus a route that consumes a database session. + + Built through ``create_app`` so the domain exception handlers that map + ``ServiceUnavailableError`` to 503 are actually registered. + """ + application = create_app(settings=Settings(API_KEYS="k")) + router = APIRouter() + + @router.get("/probe") + async def probe(session: AsyncSession = Depends(get_db_session)) -> dict[str, bool]: + return {"has_session": isinstance(session, AsyncSession)} + + application.include_router(router) + return application + + +def test_get_db_session_returns_503_when_database_absent() -> None: + """Without a session factory the dependency must fail explicitly, not silently. + + Yielding nothing would leave FastAPI's solver with an exhausted generator and + surface as an opaque ``generator didn't yield`` RuntimeError. + """ + client = TestClient(_app(), raise_server_exceptions=False) + + response = client.get("/probe") + + assert response.status_code == 503 + assert response.json()["error"]["code"] == "SERVICE_UNAVAILABLE" + + +def test_get_db_session_yields_session_when_factory_present() -> None: + """A configured session factory yields a usable AsyncSession.""" + factory: async_sessionmaker[AsyncSession] = async_sessionmaker( + class_=AsyncSession, expire_on_commit=False + ) + app = _app() + app.state.session_factory = factory + + client = TestClient(app, raise_server_exceptions=False) + + response = client.get("/probe") + + assert response.status_code == 200 + assert response.json() == {"has_session": True} + + +async def test_get_db_session_raises_domain_error_directly() -> None: + """Direct invocation raises the domain error rather than a framework error.""" + generator = get_db_session(factory=None) + + with pytest.raises(ServiceUnavailableError) as excinfo: + await generator.__anext__() + + assert excinfo.value.service == "database" + assert excinfo.value.code == "SERVICE_UNAVAILABLE" + + +async def test_service_unavailable_error_serializes() -> None: + """The error carries a structured payload for API responses.""" + error = ServiceUnavailableError("redis", "Redis is unreachable.") + + payload: dict[str, Any] = error.to_dict() + + assert payload == {"code": "SERVICE_UNAVAILABLE", "message": "Redis is unreachable."}