From eec85a324745d8f06edb401ce28ecef7f6561ead Mon Sep 17 00:00:00 2001 From: Hamza Imran Date: Wed, 21 Oct 2026 02:51:36 +0500 Subject: [PATCH] fix(api): bound session_id at the request boundary (fixes #66) --- agentkit/api/schemas.py | 3 + agentkit/db/models.py | 5 +- docs/PROGRESS.md | 7 ++ tests/integration/test_api_runs.py | 7 ++ tests/unit/test_audit_regressions.py | 127 +++++++++++++++++++++++++++ 5 files changed, 148 insertions(+), 1 deletion(-) diff --git a/agentkit/api/schemas.py b/agentkit/api/schemas.py index 8d4d533..2e44dad 100644 --- a/agentkit/api/schemas.py +++ b/agentkit/api/schemas.py @@ -5,6 +5,8 @@ from pydantic import BaseModel, ConfigDict, Field +SESSION_ID_MAX_LENGTH: int = 255 + class RunCreateRequest(BaseModel): """Payload to initiate a new agent execution run.""" @@ -14,6 +16,7 @@ class RunCreateRequest(BaseModel): goal: str = Field(..., min_length=1, description="Goal or task description for the agent.") session_id: str | None = Field( default=None, + max_length=SESSION_ID_MAX_LENGTH, description="Optional conversation session ID to preserve context across multiple runs.", ) diff --git a/agentkit/db/models.py b/agentkit/db/models.py index 258d684..30545e8 100644 --- a/agentkit/db/models.py +++ b/agentkit/db/models.py @@ -12,6 +12,9 @@ class Base(DeclarativeBase): """Base class for all SQLAlchemy declarative models.""" +SESSION_ID_MAX_LENGTH: int = 255 + + class Run(Base): """Represents an agent execution run.""" @@ -23,7 +26,7 @@ class Run(Base): default=uuid.uuid4, ) session_id: Mapped[str | None] = mapped_column( - String(255), + String(SESSION_ID_MAX_LENGTH), nullable=True, index=True, ) diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 2c7b115..100fd4e 100644 --- a/docs/PROGRESS.md +++ b/docs/PROGRESS.md @@ -61,6 +61,7 @@ This document tracks execution progress across all backlog tasks, including pull | AUD-09 | AUD | Live End-to-End Infrastructure Verification & Documentation Proof | `docs/e2e-verification-report` | #53 | 3 | Done | Standardized model, verified live stack, and documented in E2E_REPORT.md | | AUD-10 | AUD | Rate Limiter Outage Fallback & Provider-Valid Session History | `fix/ratelimit-and-session-repair` | #56 | 3 | Done | In-memory fallback on Redis connection failures; synthetic tool results repair aborted sessions | | M13-T06 | M13 | Whole-Project Audit Remediation | `fix/audit-hardening` | #47 | 3 | Done | Fail-closed config and LLM resolution, SSRF post-connect check, atomic fallback | +| AUDIT-09 | AUD | Bound session_id Length at API Request Boundary | ix/session-id-length-validation | #80 | 1 | Done | Enforce max_length=255 matching VARCHAR(255) DB column, preventing persistence failures | --- @@ -407,3 +408,9 @@ This document tracks execution progress across all backlog tasks, including pull - **PR**: #47 - **Commits**: 3 - **Notes**: Remediated 17 findings from a full-codebase audit, rebased onto the AUD-01/AUD-04/AUD-05 follow-ups. Critical: removed the `ak_test_key_12345` default from `Settings.API_KEYS` and made the field required; replaced the `FakeLLMClient` fallback in `get_llm_client` with a `ServiceUnavailableError` (503) boundary around provider construction; kept the AUD-04 factory-level `RetryingLLMClient` wrapping and made it configurable through new `LLM_MAX_RETRIES`/`LLM_RETRY_BASE_DELAY_S`/`LLM_RETRY_MAX_DELAY_S` settings; made the per-request `Agent` assembly's override explicit as `app.state.agent_override`; forwarded `get_run_trace`/`list_runs` through `TieredMemoryStore`; moved `http_fetch` DNS off the event loop with `asyncio.to_thread`, added a post-connect `server_addr` check against the validated addresses to close the DNS-rebinding TOCTOU window, and replaced full-body buffering with a byte-budgeted stream; replaced the rate limiter's non-scripting fallback with `WATCH`/`MULTI`/`EXEC`, leaving the AUD-10 in-memory fallback for Redis outages intact. Moderate: corrected the `RedisMemoryStore.truncate_messages` off-by-one that returned `max_messages + 1` entries when a system prompt was present; switched generated `tool_call_id`s to full UUIDs; wrapped response parsing in both LLM adapters inside the error boundary so malformed responses and malformed tool-argument JSON raise `LLMProviderError`; guarded `RetryingLLMClient.__getattr__` against recursion on a partially constructed instance; sanitized `/health` details so the unauthenticated endpoint no longer echoes driver error strings; made `sql_readonly` mask string literals and quoted identifiers before its keyword scan, allow `#>`/`#>>`/`#-`, apply the row limit in SQL and push `statement_timeout` to PostgreSQL. Minor: stopped logging raw run goals at INFO (length only, content at DEBUG); replaced `split(";")` in `scripts/seed.py` with a statement-aware splitter handling literals, comments and dollar-quotes; bounded calculator exponentiation by result bit length. Added `tests/unit/test_audit_regressions.py` and a real-Redis concurrency regression test pinning the fallback fix. Recorded in ADR-009, ADR-010, ADR-011. Suite: 298 passed, 25 skipped; coverage 92.86% on core/tools/llm; `ruff check`, `ruff format --check` and `mypy agentkit tests` all clean. + +### AUDIT-09: Bound session_id Length at API Request Boundary +- **Status**: Done +- **PR**: #80 +- **Commits**: 1 +- **Notes**: In gentkit/api/schemas.py, added SESSION_ID_MAX_LENGTH: int = 255 and set max_length=SESSION_ID_MAX_LENGTH on RunCreateRequest.session_id. Aligned gentkit/db/models.py to use SESSION_ID_MAX_LENGTH for Run.session_id. Over-long session IDs are now rejected with HTTP 422 before agent execution starts, preventing data-length insertion failures and unbounded Redis key creation. Added unit, integration, and drift regression tests in ests/unit/test_audit_regressions.py and ests/integration/test_api_runs.py. diff --git a/tests/integration/test_api_runs.py b/tests/integration/test_api_runs.py index b79a639..fdf73c0 100644 --- a/tests/integration/test_api_runs.py +++ b/tests/integration/test_api_runs.py @@ -65,6 +65,13 @@ def test_runs_validation_error() -> None: assert resp.status_code == 422 assert resp.json()["error"]["code"] == "VALIDATION_ERROR" + # Over-long session_id (> 255 chars) + resp = client.post( + "/runs", json={"goal": "calculate 2+2", "session_id": "x" * 256}, headers=headers + ) + assert resp.status_code == 422 + assert resp.json()["error"]["code"] == "VALIDATION_ERROR" + @pytest.mark.asyncio async def test_create_and_query_run_lifecycle( diff --git a/tests/unit/test_audit_regressions.py b/tests/unit/test_audit_regressions.py index e7c0d9d..f4eb055 100644 --- a/tests/unit/test_audit_regressions.py +++ b/tests/unit/test_audit_regressions.py @@ -23,9 +23,22 @@ from agentkit.api.deps import get_llm_client from agentkit.api.main import create_app from agentkit.api.ratelimit import RateLimiter +from agentkit.api.schemas import ( + SESSION_ID_MAX_LENGTH as API_SESSION_ID_MAX_LENGTH, +) +from agentkit.api.schemas import ( + RunCreateRequest, +) from agentkit.config import Settings from agentkit.core.errors import AuthenticationError, ServiceUnavailableError from agentkit.core.trace import InMemoryTraceSink, StepTrace +from agentkit.db.models import ( + SESSION_ID_MAX_LENGTH as DB_SESSION_ID_MAX_LENGTH, +) +from agentkit.db.models import ( + Base, + Run, +) from agentkit.llm.base import LLMClient, LLMResponse, Message from agentkit.llm.retry import RetryingLLMClient from agentkit.memory.base import InMemoryMemoryStore, RunRecord @@ -590,3 +603,117 @@ async def test_unatomic_fallback_does_not_over_admit_under_concurrency() -> None assert not errors, f"unexpected exceptions: {errors}" assert len(admitted) == 5, f"expected exactly 5 admissions, got {len(admitted)}" assert len(denied) == 15 + + +# --- AUDIT-09: session_id bounded at API boundary and matching DB column ------------- + + +def test_session_id_db_and_api_max_length_aligned() -> None: + """AUDIT-09: DB column length and API request schema max_length must match exactly.""" + col_length = Run.__table__.c.session_id.type.length + assert col_length == 255 + assert DB_SESSION_ID_MAX_LENGTH == 255 + assert API_SESSION_ID_MAX_LENGTH == 255 + + schema_max_length = next( + m.max_length + for m in RunCreateRequest.model_fields["session_id"].metadata + if hasattr(m, "max_length") + ) + assert schema_max_length == col_length == 255 + + +@pytest.mark.asyncio +async def test_session_id_over_long_rejected_at_api_boundary() -> None: + """AUDIT-09: Over-long session_id (>255 chars) must return 422 before agent execution.""" + agent_called = False + + class _GuardedLLM(LLMClient): + async def chat(self, _messages: list[Message]) -> LLMResponse: + nonlocal agent_called + agent_called = True + return LLMResponse(text="should not be reached") + + engine = create_async_engine("sqlite+aiosqlite:///:memory:", echo=False) + factory = async_sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + mem_store = InMemoryMemoryStore() + app = create_app(settings=Settings(API_KEYS="secret", LLM_PROVIDER="fake")) + app.state.session_factory = factory + app.state.memory_store = mem_store + app.state.trace_sink = InMemoryTraceSink() + app.state.llm_client = _GuardedLLM() + + client = TestClient(app, raise_server_exceptions=False) + headers = {"X-API-Key": "secret"} + + # 256 characters session_id + overlong_id = "s" * 256 + resp = client.post( + "/runs", + json={"goal": "test goal", "session_id": overlong_id}, + headers=headers, + ) + + assert resp.status_code == 422 + data = resp.json() + assert data["error"]["code"] == "VALIDATION_ERROR" + assert not agent_called + assert len(mem_store._runs) == 0 + + await engine.dispose() + + +@pytest.mark.asyncio +async def test_session_id_max_boundary_accepted_and_roundtrips() -> None: + """AUDIT-09: A 255-character session_id is accepted and persisted across DB and memory.""" + from agentkit.memory.pg_store import PostgresMemoryStore + from agentkit.memory.redis_store import RedisMemoryStore + + engine = create_async_engine("sqlite+aiosqlite:///:memory:", echo=False) + factory = async_sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + pg_store = PostgresMemoryStore(factory) + fake_redis: Redis[Any] = fakeredis.aioredis.FakeRedis(decode_responses=True) + redis_store = RedisMemoryStore(fake_redis) + tiered = TieredMemoryStore(session_store=redis_store, run_store=pg_store) + + class _EchoLLM(LLMClient): + async def chat(self, _messages: list[Message]) -> LLMResponse: + return LLMResponse(text="Echo response") + + app = create_app(settings=Settings(API_KEYS="secret", LLM_PROVIDER="fake")) + app.state.session_factory = factory + app.state.memory_store = tiered + app.state.trace_sink = InMemoryTraceSink() + app.state.llm_client = _EchoLLM() + + client = TestClient(app, raise_server_exceptions=False) + headers = {"X-API-Key": "secret"} + + valid_boundary_id = "a" * 255 + resp = client.post( + "/runs", + json={"goal": "test goal", "session_id": valid_boundary_id}, + headers=headers, + ) + + assert resp.status_code == 200, resp.text + run_data = resp.json() + assert run_data["session_id"] == valid_boundary_id + + # Verify DB persistence + stored_run = await pg_store.get_run(run_data["run_id"]) + assert stored_run is not None + assert stored_run.session_id == valid_boundary_id + + # Verify Redis persistence + session_messages = await redis_store.get_messages(valid_boundary_id) + assert len(session_messages) > 0 + + await fake_redis.aclose() + await engine.dispose()