Build Brief: Session State Management

Epic: Session State Management Project: Orchestrator Core Engine Status: Ready for implementation Depends on: None (first epic)

---

Context

This is the first epic of the Orchestrator service. The project is bare — no code exists yet. You are creating the project from scratch.

Tech Stack:

• Python 3.12+
• FastAPI (async web framework)
• Redis 7 (session cache, job queue)
• PostgreSQL 16 (persistent storage)
• Alembic (migrations)
• pytest + pytest-asyncio (testing)
• ruff (linting), mypy (type checking)
• prometheus-client (metrics)
• orjson (fast JSON serialization)

Project Root: orchestrator/ (create this directory)

Directory Layout:

orchestrator/
├── src/
│   ├── __init__.py
│   ├── main.py                  # FastAPI app, lifespan, health endpoints
│   ├── config.py                # Settings via pydantic-settings
│   ├── session/
│   │   ├── __init__.py
│   │   ├── state.py             # SessionState dataclass
│   │   ├── store.py             # Redis + Postgres session store
│   │   ├── memory.py            # In-memory LRU cache
│   │   └── reconstruction.py    # State reconstruction logic
│   ├── routing/
│   │   ├── __init__.py
│   │   └── sticky.py            # Consistent hash ring
│   └── metrics/
│       ├── __init__.py
│       └── session_metrics.py   # Prometheus metrics
├── tests/
│   ├── __init__.py
│   ├── conftest.py
│   ├── test_state.py
│   ├── test_memory.py
│   ├── test_store.py
│   ├── test_reconstruction.py
│   └── test_sticky.py
├── alembic/
│   ├── alembic.ini
│   └── versions/
├── pyproject.toml
├── docker-compose.yml
└── README.md

---

pyproject.toml

[project]
name = "orchestrator"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = [
    "fastapi>=0.115",
    "uvicorn[standard]>=0.34",
    "redis[hiredis]>=5.2",
    "asyncpg>=0.30",
    "pydantic>=2.10",
    "pydantic-settings>=2.7",
    "alembic>=1.14",
    "prometheus-client>=0.21",
    "orjson>=3.10",
]

[project.optional-dependencies]
dev = [
    "pytest>=8.3",
    "pytest-asyncio>=0.25",
    "pytest-cov>=6.0",
    "httpx>=0.28",
    "fakeredis>=2.26",
    "testcontainers[redis,postgres]>=4.9",
    "ruff>=0.9",
    "mypy>=1.14",
]

[tool.ruff]
target-version = "py312"
line-length = 100

[tool.mypy]
python_version = "3.12"
strict = true

[tool.pytest.ini_options]
asyncio_mode = "auto"
testpaths = ["tests"]

---

docker-compose.yml

services:
  redis:
    image: redis:7-alpine
    ports:
      - "6379:6379"
    volumes:
      - redis-data:/data

  postgres:
    image: postgres:16-alpine
    environment:
      POSTGRES_USER: orchestrator
      POSTGRES_PASSWORD: orchestrator
      POSTGRES_DB: orchestrator
    ports:
      - "5432:5432"
    volumes:
      - pg-data:/var/lib/postgresql/data

volumes:
  redis-data:
  pg-data:

---

Configuration (src/config.py)

Environment variables:

REDIS_URL=redis://localhost:6379/0
DATABASE_URL=postgresql+asyncpg://orchestrator:orchestrator@localhost:5432/orchestrator
SESSION_TTL_SECONDS=1800
SESSION_MEMORY_MAX_SIZE=10000
POSTGRES_BATCH_INTERVAL_SECONDS=5
INSTANCE_ID=orchestrator-1

---

Shared Schemas

Define these in src/session/state.py (inline, no shared module needed yet):

from __future__ import annotations
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any
from uuid import UUID, uuid4


class Role(str, Enum):
    USER = "user"
    ASSISTANT = "assistant"
    SYSTEM = "system"


@dataclass
class Turn:
    role: Role
    content: str
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


@dataclass
class TurnContext:
    user_message: str
    attached_context: dict[str, Any] = field(default_factory=dict)
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


@dataclass
class SessionState:
    session_id: str = field(default_factory=lambda: str(uuid4()))
    user_id: str = ""
    current_turn: TurnContext = field(default_factory=lambda: TurnContext(user_message=""))
    pending_jobs: dict[str, str] = field(default_factory=dict)  # job_id → agent_type
    intermediate_results: dict[str, Any] = field(default_factory=dict)  # job_id → result
    conversation_history: list[Turn] = field(default_factory=list)
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    updated_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    ttl_seconds: int = 1800

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dict for Redis storage via orjson."""
        return {
            "session_id": self.session_id,
            "user_id": self.user_id,
            "current_turn": {
                "user_message": self.current_turn.user_message,
                "attached_context": self.current_turn.attached_context,
                "timestamp": self.current_turn.timestamp.isoformat(),
            },
            "pending_jobs": self.pending_jobs,
            "intermediate_results": self.intermediate_results,
            "conversation_history": [
                {"role": t.role.value, "content": t.content, "timestamp": t.timestamp.isoformat()}
                for t in self.conversation_history
            ],
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "ttl_seconds": self.ttl_seconds,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> SessionState:
        """Deserialize from dict (Redis or Postgres)."""
        return cls(
            session_id=data["session_id"],
            user_id=data["user_id"],
            current_turn=TurnContext(
                user_message=data["current_turn"]["user_message"],
                attached_context=data["current_turn"].get("attached_context", {}),
                timestamp=datetime.fromisoformat(data["current_turn"]["timestamp"]),
            ),
            pending_jobs=data.get("pending_jobs", {}),
            intermediate_results=data.get("intermediate_results", {}),
            conversation_history=[
                Turn(
                    role=Role(t["role"]),
                    content=t["content"],
                    timestamp=datetime.fromisoformat(t["timestamp"]),
                )
                for t in data.get("conversation_history", [])
            ],
            created_at=datetime.fromisoformat(data["created_at"]),
            updated_at=datetime.fromisoformat(data["updated_at"]),
            ttl_seconds=data.get("ttl_seconds", 1800),
        )

---

Task Breakdown

Phase 1: Schema and Data Layer

T1.1 — Session state schema

• File: src/session/state.py
• Implement the SessionState dataclass above exactly as specified
• Implement to_dict() and from_dict() for serialization via orjson
• Unit tests: tests/test_state.py
  • Test roundtrip: SessionState → to_dict() → from_dict() produces identical object
  • Test defaults: new SessionState() has valid session_id, empty collections, ttl_seconds=1800
  • Test edge cases: empty conversation_history, empty pending_jobs, nested attached_context

Phase 2: Storage Layer

T2.1 — In-memory LRU cache

• File: src/session/memory.py
• Implement SessionMemoryCache class:
  • get(session_id: str) -> SessionState | None
  • put(session_id: str, state: SessionState) -> None
  • evict(session_id: str) -> None
  • size() -> int
• Use collections.OrderedDict with move_to_end for LRU behavior
• Max size configurable (default 10,000). On put when full, evict oldest.
• Unit tests: tests/test_memory.py
  • Test get/put roundtrip
  • Test LRU eviction: insert 101 items with max_size=100, verify oldest is evicted
  • Test evict removes specific entry
  • Test size tracks correctly

T2.2 — Redis session store

• File: src/session/store.py
• Implement RedisSessionStore class:
  • async get(session_id: str) -> SessionState | None — GET session:{id}, deserialize with orjson.loads
  • async set(session_id: str, state: SessionState) -> None — SETEX session:{id} {ttl} {orjson.dumps(state.to_dict())}
  • async delete(session_id: str) -> None — DEL session:{id}
  • async exists(session_id: str) -> bool — EXISTS session:{id}
  • async refresh_ttl(session_id: str) -> None — EXPIRE session:{id} {ttl} (sliding window)
• Connection pool: redis.asyncio.ConnectionPool.from_url(url, max_connections=20, min_connections=5)
• Unit tests: tests/test_store.py — use fakeredis.aioredis.FakeRedis() for unit tests

T2.3 — Postgres schema

• File: alembic/versions/001_session_tables.py
• Migration creates:

  CREATE TABLE sessions (
      id UUID PRIMARY KEY,
      user_id TEXT NOT NULL,
      created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
      updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
      expires_at TIMESTAMPTZ NOT NULL
  );
  
  CREATE TABLE conversation_history (
      id UUID PRIMARY KEY,
      session_id UUID REFERENCES sessions(id) ON DELETE CASCADE,
      role TEXT NOT NULL CHECK (role IN ('user', 'assistant', 'system')),
      content TEXT NOT NULL,
      timestamp TIMESTAMPTZ NOT NULL DEFAULT now()
  );
  
  CREATE INDEX idx_conv_session ON conversation_history(session_id, timestamp);

• Implement PostgresSessionStore class in src/session/store.py:
  • async save_session(state: SessionState) -> None — UPSERT into sessions (ON CONFLICT DO UPDATE)
  • async save_conversation_batch(session_id: str, turns: list[Turn]) -> None — batch INSERT
  • async load_session(session_id: str) -> SessionState | None — SELECT from sessions + conversation_history, reconstruct
  • async load_conversation(session_id: str, limit: int = 50) -> list[Turn] — SELECT with ORDER BY timestamp DESC
• Use asyncpg.create_pool() with same connection params as config

Phase 3: Write-Through and Reconstruction

T3.1 — Write-through persistence

• File: src/session/store.py (extend)
• Implement SessionWriteThrough class:
  • Constructor takes SessionMemoryCache, RedisSessionStore, PostgresSessionStore
  • async write(state: SessionState) -> None:
    1. memory.put(state.session_id, state) — synchronous
    2. await redis.set(state.session_id, state) — async, immediate
    3. Mark dirty in DirtyTracker
  • async flush_dirty() -> None — called every 5 seconds by background task:
    1. Collect all dirty session IDs
    2. Batch UPSERT to Postgres
    3. Clear dirty set
• Implement DirtyTracker:
  • mark_dirty(session_id: str) — add to set
  • get_dirty() -> set[str] — return and clear
  • Thread-safe via asyncio.Lock

T3.2 — Session reconstruction

• File: src/session/reconstruction.py
• Implement SessionReconstructor class:
  • Constructor takes SessionMemoryCache, RedisSessionStore, PostgresSessionStore
  • async reconstruct(session_id: str) -> SessionState:
    1. Check memory.get(session_id) — return if found (increment memory_hits counter)
    2. Check await redis.get(session_id) — if found, memory.put() and return (increment redis_hits counter)
    3. Check await postgres.load_session(session_id) — if found, redis.set() + memory.put() and return (increment postgres_hits counter)
    4. Return fresh SessionState(session_id=session_id) (not found anywhere)
  • Increment reconstructed_total counter on every call
  • Record reconstruction latency in histogram

Phase 4: TTL, Cleanup, and Routing

T4.1 — Session TTL

• Redis TTL: SETEX with configurable ttl_seconds (default 1800)
• Postgres: expires_at column = now() + interval '{ttl_seconds} seconds'
• Sliding window: refresh_ttl() called on every reconstruct() access
• Stale sessions: PostgresSessionStore.load_session() filters WHERE expires_at > now()

T4.2 — Session cleanup

• File: src/session/cleanup.py
• Implement SessionCleanup class:
  • async run_cleanup() -> None — called every 5 minutes by background task
  • Postgres: DELETE FROM conversation_history WHERE session_id IN (SELECT id FROM sessions WHERE expires_at < now())
  • Postgres: DELETE FROM sessions WHERE expires_at < now()
  • Log metrics: sessions_cleaned_total, conversations_cleaned_total

T4.3 — Sticky routing

• File: src/routing/sticky.py
• Implement ConsistentHashRing class:
  • __init__(self, nodes: list[str], virtual_nodes: int = 150)
  • get_node(self, session_id: str) -> str — consistent hash lookup
  • add_node(self, node: str) -> None — add node with virtual nodes
  • remove_node(self, node: str) -> None — remove node and its virtual nodes
• Use hashlib.md5 for hashing (deterministic, fast)
• Internal: sorted dict of hash_value → node_name
• Unit tests: tests/test_sticky.py
  • Test distribution: 1000 sessions across 3 nodes, each gets 28-38% (within 10%)
  • Test add_node: adding a node redistributes only nearby keys
  • Test remove_node: removing a node redistributes its keys to neighbors

T4.4 — Session migration

• File: src/session/migration.py
• Implement SessionMigrator class:
  • async on_instance_failure(failed_instance: str, ring: ConsistentHashRing) -> None
    1. Get all session IDs that hashed to failed instance (from Redis keys session:*)
    2. Remove failed instance from ring
    3. For each session: get new node from ring, trigger reconstruction on new instance
  • async graceful_shutdown(instance_id: str, ring: ConsistentHashRing) -> None
    1. Remove instance from ring
    2. Wait for in-flight requests to complete (drain)
    3. Mark sessions as migrating in Redis

Phase 5: Metrics and Observability

T5.1 — Session metrics

• File: src/metrics/session_metrics.py
• Define Prometheus metrics:

  from prometheus_client import Gauge, Counter, Histogram
  
  SESSION_ACTIVE = Gauge("orchestrator_session_active", "Active sessions per instance")
  SESSION_RECONSTRUCTED = Counter("orchestrator_session_reconstructed_total", "Session reconstructions")
  SESSION_RECONSTRUCTION_LATENCY = Histogram(
      "orchestrator_session_reconstruction_latency_seconds",
      "Reconstruction latency",
      buckets=[0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0],
  )
  SESSION_TTL_EXPIRED = Counter("orchestrator_session_ttl_expired_total", "TTL expirations")
  SESSION_MEMORY_HITS = Counter("orchestrator_session_memory_hits_total", "Memory cache hits")
  SESSION_REDIS_HITS = Counter("orchestrator_session_redis_hits_total", "Redis hits")
  SESSION_POSTGRES_HITS = Counter("orchestrator_session_postgres_hits_total", "Postgres hits")

• Expose at GET /metrics via prometheus_client.make_asgi_app()

T5.2 — FastAPI app

• File: src/main.py
• Implement FastAPI app with:
  • GET /healthz — returns {"status": "ok"}
  • GET /readyz — checks Redis ping + Postgres ping, returns 200 or 503
  • Mount Prometheus metrics at /metrics
  • Lifespan: startup creates Redis pool + Postgres pool, starts background tasks (flush, cleanup)
  • Lifespan: shutdown drains pools

---

Testing

Run

cd orchestrator
docker compose up -d          # Start Redis + Postgres
pip install -e ".[dev]"       # Install deps
alembic upgrade head          # Run migrations
pytest tests/ -v --cov=src --cov-report=term-missing
ruff check src/ tests/
mypy src/

Test Files

• tests/test_state.py — serialization roundtrip, defaults, edge cases
• tests/test_memory.py — LRU eviction, get/put/delete, max size
• tests/test_store.py — Redis get/set/delete with fakeredis, Postgres save/load with testcontainers
• tests/test_reconstruction.py — cache hit, Redis hit, Postgres fallback, fresh state
• tests/test_sticky.py — hash ring distribution, add/remove node, migration

Conftest Fixtures

import pytest
from uuid import uuid4
from src.session.state import SessionState, TurnContext, Turn, Role
from src.session.memory import SessionMemoryCache

@pytest.fixture
def memory_cache():
    return SessionMemoryCache(max_size=100)

@pytest.fixture
def sample_state():
    return SessionState(
        session_id=str(uuid4()),
        user_id="user-1",
        current_turn=TurnContext(user_message="Find sushi in Tokyo"),
        conversation_history=[
            Turn(role=Role.USER, content="Find sushi in Tokyo"),
            Turn(role=Role.ASSISTANT, content="Here are 3 options..."),
        ],
    )

---

Acceptance Criteria

• ☐ SessionState.to_dict() → from_dict() roundtrip produces identical object
• ☐ In-memory cache evicts LRU entries when max_size reached
• ☐ Redis store reads/writes with correct TTL (30 min default)
• ☐ Postgres migration creates sessions and conversation_history tables with indexes
• ☐ Write-through: Redis write is immediate, Postgres write batches every 5 seconds
• ☐ Reconstruction: cache hit <1ms, Redis hit <5ms, Postgres fallback <50ms
• ☐ Stale sessions (expired expires_at) return fresh state, not stale data
• ☐ Sticky routing distributes 1000 sessions across 3 nodes within 10% variance
• ☐ Prometheus metrics exposed at /metrics with all 7 counters/gauges/histograms
• ☐ GET /healthz returns 200, GET /readyz checks Redis + Postgres connectivity
• ☐ All unit tests pass with >80% coverage on src/session/
• ☐ ruff check and mypy pass with no errors
• ☐ README.md documents setup, configuration, schema, and usage