From 2b58c127b1435d77fcb9089e44ea1cee66b417e5 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Tue, 3 Mar 2026 22:57:34 -0800 Subject: [PATCH] Add pluggable storage backend (SQLite + PostgreSQL) and deployment packaging (#20) * Add pluggable storage backend (SQLite + PostgreSQL) and deployment packaging Database abstraction: StorageBackend protocol with 21 methods, SQLAlchemy Core schema, SQLite backend (FTS5), PostgreSQL backend (tsvector/ILIKE), Alembic migrations, singleton registry. memory.py reduced to thin facade. Session.py open_db() calls replaced with generic KV methods. [database] config section with env var support. Deployment: Docker Compose production profile with PostgreSQL, Dockerfile with postgres extras and migration entrypoint, Helm chart with bitnami subcharts, Terraform AWS ECS/Fargate module with RDS + ElastiCache + ALB. 39 new storage tests (934 total). mypy strict clean. Docs and diagrams updated. * Address PR #20 review feedback (16 items) - Backends only call create_all() when Alembic migrations are disabled - Helm configmap uses correct TURNSTONE_DB_BACKEND env var; DB URL constructed via env expansion with secret reference instead of ConfigMap - Migration errors fail fast for PostgreSQL (only non-fatal for SQLite) - save_memory/delete_memory wrapped in exception handling like other facade fns - pool_size passed through from config/env to init_storage() in cli + server - Terraform: DB URL moved to Secrets Manager, auth enabled flag set, optional TLS listeners with certificate_arn, Redis transit encryption on - Docker entrypoint no longer suppresses migration output - Diagram fixes: removed StaticPool claim, removed non-existent migration ref - compose.yaml/README: clarified production profile requires DB env vars --- .dockerignore | 9 + .env.example | 97 +-- Dockerfile | 13 +- README.md | 27 +- compose.yaml | 45 +- deploy/helm/turnstone/Chart.yaml | 16 + deploy/helm/turnstone/templates/NOTES.txt | 42 ++ deploy/helm/turnstone/templates/_helpers.tpl | 163 +++++ .../helm/turnstone/templates/configmap.yaml | 23 + .../templates/deployment-bridge.yaml | 45 ++ .../templates/deployment-console.yaml | 62 ++ .../templates/deployment-server.yaml | 64 ++ deploy/helm/turnstone/templates/ingress.yaml | 47 ++ .../helm/turnstone/templates/job-migrate.yaml | 38 + deploy/helm/turnstone/templates/secret.yaml | 28 + .../turnstone/templates/service-console.yaml | 17 + .../turnstone/templates/service-server.yaml | 17 + .../turnstone/templates/serviceaccount.yaml | 6 + deploy/helm/turnstone/values.yaml | 103 +++ .../terraform/examples/aws-ecs-basic/main.tf | 36 + .../examples/aws-ecs-basic/outputs.tf | 29 + .../aws-ecs-basic/terraform.tfvars.example | 22 + .../examples/aws-ecs-basic/variables.tf | 62 ++ deploy/terraform/modules/aws-ecs/alb.tf | 150 ++++ .../terraform/modules/aws-ecs/elasticache.tf | 30 + deploy/terraform/modules/aws-ecs/iam.tf | 70 ++ deploy/terraform/modules/aws-ecs/main.tf | 313 ++++++++ deploy/terraform/modules/aws-ecs/outputs.tf | 29 + deploy/terraform/modules/aws-ecs/rds.tf | 42 ++ deploy/terraform/modules/aws-ecs/security.tf | 133 ++++ deploy/terraform/modules/aws-ecs/variables.tf | 130 ++++ docker/entrypoint.sh | 5 + docs/architecture.md | 94 ++- docs/diagrams/02-package-structure.puml | 4 +- docs/diagrams/14-storage-architecture.puml | 156 ++++ docs/diagrams/png/02-package-structure.png | 4 +- docs/diagrams/png/14-storage-architecture.png | 3 + pyproject.toml | 7 + tests/conftest.py | 12 +- tests/test_config.py | 16 +- tests/test_db.py | 29 +- tests/test_fts5.py | 30 +- tests/test_mcp_client.py | 10 +- tests/test_server_live.py | 10 +- tests/test_sessions.py | 66 +- tests/test_storage_registry.py | 54 ++ tests/test_storage_sqlite.py | 262 +++++++ turnstone/chat.py | 3 - turnstone/cli.py | 15 +- turnstone/core/config.py | 32 + turnstone/core/memory.py | 692 ++++-------------- turnstone/core/session.py | 88 +-- turnstone/core/storage/__init__.py | 14 + turnstone/core/storage/_migrate.py | 84 +++ turnstone/core/storage/_postgresql.py | 386 ++++++++++ turnstone/core/storage/_protocol.py | 117 +++ turnstone/core/storage/_registry.py | 86 +++ turnstone/core/storage/_schema.py | 56 ++ turnstone/core/storage/_sqlite.py | 546 ++++++++++++++ turnstone/core/storage/migrations/env.py | 27 + .../core/storage/migrations/script.py.mako | 23 + .../migrations/versions/001_initial_schema.py | 68 ++ turnstone/eval.py | 8 +- turnstone/server.py | 15 +- 64 files changed, 4083 insertions(+), 847 deletions(-) create mode 100644 deploy/helm/turnstone/Chart.yaml create mode 100644 deploy/helm/turnstone/templates/NOTES.txt create mode 100644 deploy/helm/turnstone/templates/_helpers.tpl create mode 100644 deploy/helm/turnstone/templates/configmap.yaml create mode 100644 deploy/helm/turnstone/templates/deployment-bridge.yaml create mode 100644 deploy/helm/turnstone/templates/deployment-console.yaml create mode 100644 deploy/helm/turnstone/templates/deployment-server.yaml create mode 100644 deploy/helm/turnstone/templates/ingress.yaml create mode 100644 deploy/helm/turnstone/templates/job-migrate.yaml create mode 100644 deploy/helm/turnstone/templates/secret.yaml create mode 100644 deploy/helm/turnstone/templates/service-console.yaml create mode 100644 deploy/helm/turnstone/templates/service-server.yaml create mode 100644 deploy/helm/turnstone/templates/serviceaccount.yaml create mode 100644 deploy/helm/turnstone/values.yaml create mode 100644 deploy/terraform/examples/aws-ecs-basic/main.tf create mode 100644 deploy/terraform/examples/aws-ecs-basic/outputs.tf create mode 100644 deploy/terraform/examples/aws-ecs-basic/terraform.tfvars.example create mode 100644 deploy/terraform/examples/aws-ecs-basic/variables.tf create mode 100644 deploy/terraform/modules/aws-ecs/alb.tf create mode 100644 deploy/terraform/modules/aws-ecs/elasticache.tf create mode 100644 deploy/terraform/modules/aws-ecs/iam.tf create mode 100644 deploy/terraform/modules/aws-ecs/main.tf create mode 100644 deploy/terraform/modules/aws-ecs/outputs.tf create mode 100644 deploy/terraform/modules/aws-ecs/rds.tf create mode 100644 deploy/terraform/modules/aws-ecs/security.tf create mode 100644 deploy/terraform/modules/aws-ecs/variables.tf create mode 100755 docker/entrypoint.sh create mode 100644 docs/diagrams/14-storage-architecture.puml create mode 100644 docs/diagrams/png/14-storage-architecture.png create mode 100644 tests/test_storage_registry.py create mode 100644 tests/test_storage_sqlite.py create mode 100644 turnstone/core/storage/__init__.py create mode 100644 turnstone/core/storage/_migrate.py create mode 100644 turnstone/core/storage/_postgresql.py create mode 100644 turnstone/core/storage/_protocol.py create mode 100644 turnstone/core/storage/_registry.py create mode 100644 turnstone/core/storage/_schema.py create mode 100644 turnstone/core/storage/_sqlite.py create mode 100644 turnstone/core/storage/migrations/env.py create mode 100644 turnstone/core/storage/migrations/script.py.mako create mode 100644 turnstone/core/storage/migrations/versions/001_initial_schema.py diff --git a/.dockerignore b/.dockerignore index 4e9bd020..73e53853 100644 --- a/.dockerignore +++ b/.dockerignore @@ -12,3 +12,12 @@ venv/ .mypy_cache/ .ruff_cache/ .hypothesis/ +deploy/ +docs/ +tests/ +sdk/ +*.md +!README.md +!LICENSE +.coverage +.swp diff --git a/.env.example b/.env.example index 1d04abd4..7660ddef 100644 --- a/.env.example +++ b/.env.example @@ -1,85 +1,28 @@ # ============================================================================= -# Turnstone Docker Compose — Environment Configuration -# Copy to .env and fill in your values: cp .env.example .env +# Turnstone Environment Variables +# Copy to .env and adjust values for your deployment # ============================================================================= -# --------------------------------------------------------------------------- -# LLM Backend -# --------------------------------------------------------------------------- -# OpenAI-compatible API URL (vLLM, llama.cpp, OpenAI, etc.) +# -- LLM Backend -------------------------------------------------------------- LLM_BASE_URL=http://host.docker.internal:8000/v1 +OPENAI_API_KEY=sk-... +# ANTHROPIC_API_KEY=sk-ant-... # Set instead for Anthropic provider +# TAVILY_API_KEY=tvly-... # For web search fallback (local models only) -# API key for the LLM backend ("dummy" for local servers without auth) -OPENAI_API_KEY=dummy +# -- Database (production profile) -------------------------------------------- +# DB_BACKEND=postgresql +# POSTGRES_USER=turnstone +# POSTGRES_PASSWORD=changeme +# DATABASE_URL=postgresql+psycopg://turnstone:changeme@postgres:5432/turnstone -# Tavily API key for web_search tool (optional) -TAVILY_API_KEY= +# -- Redis --------------------------------------------------------------------- +# REDIS_PASSWORD= +# REDIS_PORT=6379 -# --------------------------------------------------------------------------- -# Redis -# --------------------------------------------------------------------------- -# Redis password (leave empty for no authentication) -REDIS_PASSWORD= +# -- Authentication ------------------------------------------------------------ +# TURNSTONE_AUTH_ENABLED=true +# TURNSTONE_AUTH_TOKEN=your-secret-token -# Host port for Redis -REDIS_PORT=6379 - -# --------------------------------------------------------------------------- -# Server -# --------------------------------------------------------------------------- -# Host port for the turnstone web UI -SERVER_PORT=8080 - -# Set to any non-empty value to auto-approve all tool calls -SKIP_PERMISSIONS= - -# --------------------------------------------------------------------------- -# Bridge -# --------------------------------------------------------------------------- -# Heartbeat TTL in seconds -HEARTBEAT_TTL=60 - -# Seconds to wait for external approval responses -APPROVAL_TIMEOUT=300 - -# --------------------------------------------------------------------------- -# Console (Cluster Dashboard) -# --------------------------------------------------------------------------- -# Host port for the cluster dashboard -CONSOLE_PORT=8090 - -# Seconds between node polling cycles -CONSOLE_POLL_INTERVAL=10 - -# --------------------------------------------------------------------------- -# Auth (optional) -# --------------------------------------------------------------------------- -# Set to "1" to require Bearer token authentication -TURNSTONE_AUTH_ENABLED= - -# Bearer token for server/bridge/console authentication -TURNSTONE_AUTH_TOKEN= - -# --------------------------------------------------------------------------- -# Simulator (used with: docker compose --profile sim up) -# --------------------------------------------------------------------------- -# Number of simulated nodes -SIM_NODES=100 - -# Scenario: steady, burst, node_failure, directed, lifecycle -SIM_SCENARIO=steady - -# Scenario duration in seconds -SIM_DURATION=60 - -# Messages per second (steady scenario) -SIM_MPS=5.0 - -# Log level -SIM_LOG_LEVEL=INFO - -# Random seed for reproducibility (leave empty for random) -SIM_SEED= - -# Path to write JSON metrics report (leave empty to skip) -SIM_METRICS_FILE= +# -- Ports --------------------------------------------------------------------- +# SERVER_PORT=8080 +# CONSOLE_PORT=8090 diff --git a/Dockerfile b/Dockerfile index 0c57c901..f24335a5 100644 --- a/Dockerfile +++ b/Dockerfile @@ -25,22 +25,31 @@ FROM python:3.13-slim LABEL org.opencontainers.image.title="turnstone" \ org.opencontainers.image.description="Multi-node AI orchestration platform" +# System dependencies for psycopg (PostgreSQL client library) +RUN apt-get update && apt-get install -y --no-install-recommends libpq5 \ + && rm -rf /var/lib/apt/lists/* + # Non-root user RUN useradd --create-home --shell /bin/bash turnstone -# Install the wheel with all optional extras (redis for mq/console/sim) +# Install the wheel with all optional extras COPY --from=builder /build/wheels/*.whl /tmp/wheels/ -RUN pip install --no-cache-dir "$(ls /tmp/wheels/*.whl)[mq,console,sim]" \ +RUN pip install --no-cache-dir "$(ls /tmp/wheels/*.whl)[mq,console,sim,postgres]" \ && rm -rf /tmp/wheels # Health check script (stdlib only, no pip deps needed) COPY docker/healthcheck.py /usr/local/bin/healthcheck.py +# Entrypoint script — runs migrations before starting +COPY docker/entrypoint.sh /usr/local/bin/entrypoint.sh + # Data directory — SQLite DB is created in CWD WORKDIR /data RUN chown turnstone:turnstone /data USER turnstone +ENTRYPOINT ["entrypoint.sh"] + # Default command (overridden per service in compose.yaml) CMD ["turnstone-server", "--host", "0.0.0.0", "--port", "8080"] diff --git a/README.md b/README.md index c78a37b9..06eb61a4 100644 --- a/README.md +++ b/README.md @@ -78,7 +78,14 @@ Then open `http://localhost:8090` for the cluster-wide dashboard. Create workstr ```bash cp .env.example .env # edit LLM_BASE_URL, OPENAI_API_KEY, etc. -docker compose up # starts redis + server + bridge + console +docker compose up # starts redis + server + bridge + console (SQLite) +``` + +For production with PostgreSQL: + +```bash +# Requires POSTGRES_PASSWORD and DB_BACKEND=postgresql in .env (or exported) +docker compose --profile production up # adds PostgreSQL, uses it as database ``` Console dashboard at http://localhost:8090. See [docs/docker.md](docs/docker.md) for configuration, scaling, and profiles. @@ -117,7 +124,8 @@ turnstone/ │ ├── mcp_client.py # MCP client manager (external tool servers) │ ├── model_registry.py # ModelRegistry — named models, fallback routing, per-workstream selection │ ├── config.py # Unified TOML config (~/.config/turnstone/config.toml) -│ ├── memory.py # SQLite persistence (memories, conversations, FTS5) +│ ├── memory.py # Persistence facade (delegates to storage/) +│ ├── storage/ # Pluggable storage backend (SQLite + PostgreSQL) │ ├── metrics.py # Prometheus-compatible metrics collector │ ├── healthcheck.py # Backend health monitor + circuit breaker │ ├── ratelimit.py # Per-IP token-bucket rate limiter @@ -147,9 +155,12 @@ turnstone/ ├── cli.py # Terminal frontend (+ /cluster commands for console) ├── server.py # Web frontend (Starlette/ASGI + SSE) └── eval.py # Evaluation and prompt optimization harness +├── api/ # OpenAPI spec generation (Pydantic v2 models) +├── sdk/ # Client SDKs (sync + async, Python) docs/ ├── architecture.md # System architecture and threading model ├── api-reference.md # Web server API and SSE event reference +├── sdk.md # Client SDK reference (Python + TypeScript) ├── console.md # Cluster dashboard service (turnstone-console) ├── docker.md # Docker Compose deployment and configuration ├── simulator.md # Cluster simulator usage and scenarios @@ -157,6 +168,9 @@ docs/ ├── eval.md # Evaluation harness internals └── diagrams/ # UML architecture diagrams (PlantUML sources + PNGs) └── png/ # Pre-rendered diagram images +deploy/ +├── helm/turnstone/ # Helm chart for Kubernetes +└── terraform/ # Terraform modules (AWS ECS/Fargate) ``` ### Architecture Diagrams @@ -177,6 +191,8 @@ Detailed UML diagrams are available in [`docs/diagrams/`](docs/diagrams/): | [Simulator](docs/diagrams/png/10-simulator-architecture.png) | SimCluster, dispatchers, scenarios | | [Console Data Flow](docs/diagrams/png/11-console-data-flow.png) | Dashboard data collection threads | | [Deployment](docs/diagrams/png/12-deployment.png) | Docker Compose service topology | +| [SDK Architecture](docs/diagrams/png/13-sdk-architecture.png) | Python + TypeScript client libraries | +| [Storage Architecture](docs/diagrams/png/14-storage-architecture.png) | Pluggable database backends (SQLite + PostgreSQL) | ## Multi-node routing @@ -330,6 +346,12 @@ enabled = true requests_per_second = 10.0 burst = 20 +[database] +backend = "sqlite" # "sqlite" (default) or "postgresql" +path = ".turnstone.db" # SQLite file path (relative to working directory) +# url = "postgresql+psycopg://user:pass@host:5432/turnstone" # PostgreSQL +# pool_size = 5 # PostgreSQL connection pool size + [mcp] config_path = "" # path to MCP JSON config file (alternative to TOML sections) @@ -389,6 +411,7 @@ Per-workstream metrics are labeled by `ws_id` (bounded to 10 max workstreams). - An OpenAI-compatible API endpoint ([vLLM](https://github.com/vllm-project/vllm), [NVIDIA NIM](https://build.nvidia.com/), [llama.cpp](https://github.com/ggml-org/llama.cpp), etc.) or an Anthropic API key - Redis (for message queue bridge — `pip install turnstone[mq]`) - Anthropic provider (optional — `pip install turnstone[anthropic]`) +- PostgreSQL (optional, for production — `pip install turnstone[postgres]`) ## License diff --git a/compose.yaml b/compose.yaml index d8254aa1..97b5d5ad 100644 --- a/compose.yaml +++ b/compose.yaml @@ -2,10 +2,11 @@ # Turnstone Docker Compose Stack # # Usage: -# Full stack: docker compose up -# With simulator: docker compose --profile sim up -# Sim only: docker compose --profile sim up redis console sim -# Scale bridges: docker compose up --scale bridge=3 +# Default (SQLite): docker compose up +# Production (PG): DB_BACKEND=postgresql docker compose --profile production up +# (or set DB_BACKEND=postgresql in .env) +# With simulator: docker compose --profile sim up +# Scale bridges: docker compose up --scale bridge=3 # ============================================================================= name: turnstone @@ -17,8 +18,37 @@ networks: volumes: redis-data: turnstone-data: + postgres-data: services: + # ------------------------------------------------------------------- + # PostgreSQL — production database (profile: production) + # ------------------------------------------------------------------- + postgres: + image: postgres:17-alpine + profiles: + - production + environment: + POSTGRES_DB: turnstone + POSTGRES_USER: ${POSTGRES_USER:-turnstone} + POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:?POSTGRES_PASSWORD is required for production profile} + volumes: + - postgres-data:/var/lib/postgresql/data + networks: + - turnstone-net + healthcheck: + test: ["CMD-SHELL", "pg_isready -U ${POSTGRES_USER:-turnstone}"] + interval: 5s + timeout: 3s + retries: 5 + start_period: 5s + deploy: + resources: + limits: + memory: 512M + cpus: '1.0' + restart: unless-stopped + # ------------------------------------------------------------------- # Redis — message broker, pub/sub, node registry # ------------------------------------------------------------------- @@ -66,6 +96,7 @@ services: --port 8080 --base-url "$${LLM_BASE_URL}" --api-key "$${OPENAI_API_KEY}" + $${MODEL:+--model $$MODEL} $${SKIP_PERMISSIONS:+--skip-permissions} ports: - "${SERVER_PORT:-8080}:8080" @@ -78,6 +109,9 @@ services: - SKIP_PERMISSIONS=${SKIP_PERMISSIONS:-} - TURNSTONE_AUTH_ENABLED=${TURNSTONE_AUTH_ENABLED:-} - TURNSTONE_AUTH_TOKEN=${TURNSTONE_AUTH_TOKEN:-} + - MODEL=${MODEL:-} + - TURNSTONE_DB_BACKEND=${DB_BACKEND:-sqlite} + - TURNSTONE_DB_URL=${DATABASE_URL:-} extra_hosts: - "host.docker.internal:host-gateway" networks: @@ -85,6 +119,9 @@ services: depends_on: redis: condition: service_healthy + postgres: + condition: service_healthy + required: false healthcheck: test: ["CMD", "python", "/usr/local/bin/healthcheck.py", "http://127.0.0.1:8080/health"] interval: 10s diff --git a/deploy/helm/turnstone/Chart.yaml b/deploy/helm/turnstone/Chart.yaml new file mode 100644 index 00000000..403e65dd --- /dev/null +++ b/deploy/helm/turnstone/Chart.yaml @@ -0,0 +1,16 @@ +apiVersion: v2 +name: turnstone +description: Multi-node AI orchestration platform with tool use, agent routing, and cluster simulation +type: application +version: 0.1.0 +appVersion: "0.3.0" + +dependencies: + - name: postgresql + version: ~16.0 + repository: https://charts.bitnami.com/bitnami + condition: postgresql.enabled + - name: redis + version: ~20.0 + repository: https://charts.bitnami.com/bitnami + condition: redis.enabled diff --git a/deploy/helm/turnstone/templates/NOTES.txt b/deploy/helm/turnstone/templates/NOTES.txt new file mode 100644 index 00000000..c070d789 --- /dev/null +++ b/deploy/helm/turnstone/templates/NOTES.txt @@ -0,0 +1,42 @@ +Turnstone {{ .Chart.AppVersion }} has been deployed. + +{{- if .Values.ingress.enabled }} + +Access the application via your ingress: +{{- range .Values.ingress.hosts }} + http{{ if $.Values.ingress.tls }}s{{ end }}://{{ .host }} +{{- end }} + +{{- else }} + +To access the Turnstone server, run: + + kubectl port-forward svc/{{ include "turnstone.fullname" . }}-server {{ .Values.server.service.port }}:{{ .Values.server.service.port }} + +Then open: http://localhost:{{ .Values.server.service.port }} + +To access the Turnstone console (cluster dashboard), run: + + kubectl port-forward svc/{{ include "turnstone.fullname" . }}-console {{ .Values.console.service.port }}:{{ .Values.console.service.port }} + +Then open: http://localhost:{{ .Values.console.service.port }} + +{{- end }} + +Components deployed: + - Server: {{ include "turnstone.fullname" . }}-server ({{ .Values.server.replicas }} replica(s)) + - Bridge: {{ include "turnstone.fullname" . }}-bridge ({{ .Values.bridge.replicas }} replica(s)) + - Console: {{ include "turnstone.fullname" . }}-console ({{ .Values.console.replicas }} replica(s)) +{{- if .Values.postgresql.enabled }} + - PostgreSQL (bitnami subchart) +{{- end }} +{{- if .Values.redis.enabled }} + - Redis (bitnami subchart) +{{- end }} + +{{- if not .Values.llm.apiKey }} +{{- if not .Values.llm.existingSecret }} + +WARNING: No LLM API key configured. Set llm.apiKey or llm.existingSecret in your values. +{{- end }} +{{- end }} diff --git a/deploy/helm/turnstone/templates/_helpers.tpl b/deploy/helm/turnstone/templates/_helpers.tpl new file mode 100644 index 00000000..966c3f9e --- /dev/null +++ b/deploy/helm/turnstone/templates/_helpers.tpl @@ -0,0 +1,163 @@ +{{/* +Expand the name of the chart. +*/}} +{{- define "turnstone.name" -}} +{{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }} +{{- end }} + +{{/* +Create a default fully qualified app name. +We truncate at 63 chars because some Kubernetes name fields are limited to this +(by the DNS naming spec). If release name contains chart name it will be used +as a full name. +*/}} +{{- define "turnstone.fullname" -}} +{{- if .Values.fullnameOverride }} +{{- .Values.fullnameOverride | trunc 63 | trimSuffix "-" }} +{{- else }} +{{- $name := default .Chart.Name .Values.nameOverride }} +{{- if contains $name .Release.Name }} +{{- .Release.Name | trunc 63 | trimSuffix "-" }} +{{- else }} +{{- printf "%s-%s" .Release.Name $name | trunc 63 | trimSuffix "-" }} +{{- end }} +{{- end }} +{{- end }} + +{{/* +Create chart name and version as used by the chart label. +*/}} +{{- define "turnstone.chart" -}} +{{- printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" | trunc 63 | trimSuffix "-" }} +{{- end }} + +{{/* +Common labels. +*/}} +{{- define "turnstone.labels" -}} +helm.sh/chart: {{ include "turnstone.chart" . }} +{{ include "turnstone.selectorLabels" . }} +{{- if .Chart.AppVersion }} +app.kubernetes.io/version: {{ .Chart.AppVersion | quote }} +{{- end }} +app.kubernetes.io/managed-by: {{ .Release.Service }} +{{- end }} + +{{/* +Selector labels. +*/}} +{{- define "turnstone.selectorLabels" -}} +app.kubernetes.io/name: {{ include "turnstone.name" . }} +app.kubernetes.io/instance: {{ .Release.Name }} +{{- end }} + +{{/* +Create the name of the service account to use. +*/}} +{{- define "turnstone.serviceAccountName" -}} +{{- if .Values.serviceAccount }} +{{- if .Values.serviceAccount.name }} +{{- .Values.serviceAccount.name }} +{{- else }} +{{- include "turnstone.fullname" . }} +{{- end }} +{{- else }} +{{- include "turnstone.fullname" . }} +{{- end }} +{{- end }} + +{{/* +Determine the PostgreSQL host. +*/}} +{{- define "turnstone.postgresql.host" -}} +{{- if .Values.postgresql.enabled }} +{{- printf "%s-postgresql" .Release.Name }} +{{- else }} +{{- .Values.database.external.host }} +{{- end }} +{{- end }} + +{{/* +Determine the PostgreSQL port. +*/}} +{{- define "turnstone.postgresql.port" -}} +{{- if .Values.postgresql.enabled }} +{{- printf "5432" }} +{{- else }} +{{- .Values.database.external.port | toString }} +{{- end }} +{{- end }} + +{{/* +Determine the PostgreSQL database name. +*/}} +{{- define "turnstone.postgresql.database" -}} +{{- if .Values.postgresql.enabled }} +{{- .Values.postgresql.auth.database }} +{{- else }} +{{- .Values.database.external.database }} +{{- end }} +{{- end }} + +{{/* +Determine the PostgreSQL username. +*/}} +{{- define "turnstone.postgresql.username" -}} +{{- if .Values.postgresql.enabled }} +{{- .Values.postgresql.auth.username }} +{{- else }} +{{- .Values.database.external.username }} +{{- end }} +{{- end }} + +{{/* +Determine the Redis host. +*/}} +{{- define "turnstone.redis.host" -}} +{{- if .Values.redis.enabled }} +{{- printf "%s-redis-master" .Release.Name }} +{{- else }} +{{- .Values.redis.external.host }} +{{- end }} +{{- end }} + +{{/* +Determine the Redis port. +*/}} +{{- define "turnstone.redis.port" -}} +{{- if .Values.redis.enabled }} +{{- printf "6379" }} +{{- else }} +{{- .Values.redis.external.port | toString }} +{{- end }} +{{- end }} + +{{/* +Determine the secret name for LLM API keys. +*/}} +{{- define "turnstone.llm.secretName" -}} +{{- if .Values.llm.existingSecret }} +{{- .Values.llm.existingSecret }} +{{- else }} +{{- printf "%s-secrets" (include "turnstone.fullname" .) }} +{{- end }} +{{- end }} + +{{/* +Determine the secret name for auth tokens. +*/}} +{{- define "turnstone.auth.secretName" -}} +{{- if .Values.auth.existingSecret }} +{{- .Values.auth.existingSecret }} +{{- else }} +{{- printf "%s-secrets" (include "turnstone.fullname" .) }} +{{- end }} +{{- end }} + +{{/* +Container image reference. +*/}} +{{- define "turnstone.image" -}} +{{- $tag := .Values.image.tag | default .Chart.AppVersion }} +{{- printf "%s:%s" .Values.image.repository $tag }} +{{- end }} diff --git a/deploy/helm/turnstone/templates/configmap.yaml b/deploy/helm/turnstone/templates/configmap.yaml new file mode 100644 index 00000000..a54af3af --- /dev/null +++ b/deploy/helm/turnstone/templates/configmap.yaml @@ -0,0 +1,23 @@ +apiVersion: v1 +kind: ConfigMap +metadata: + name: {{ include "turnstone.fullname" . }}-config + labels: + {{- include "turnstone.labels" . | nindent 4 }} +data: + TURNSTONE_DB_BACKEND: {{ .Values.database.backend | quote }} + TURNSTONE_DB_HOST: {{ include "turnstone.postgresql.host" . | quote }} + TURNSTONE_DB_PORT: {{ include "turnstone.postgresql.port" . | quote }} + TURNSTONE_DB_NAME: {{ include "turnstone.postgresql.database" . | quote }} + TURNSTONE_DB_USER: {{ include "turnstone.postgresql.username" . | quote }} + TURNSTONE_SERVER_HOST: "0.0.0.0" + TURNSTONE_SERVER_PORT: {{ .Values.server.service.port | quote }} + TURNSTONE_CONSOLE_HOST: "0.0.0.0" + TURNSTONE_CONSOLE_PORT: {{ .Values.console.service.port | quote }} + TURNSTONE_REDIS_HOST: {{ include "turnstone.redis.host" . | quote }} + TURNSTONE_REDIS_PORT: {{ include "turnstone.redis.port" . | quote }} + TURNSTONE_POLL_INTERVAL: "5" + {{- if .Values.llm.baseUrl }} + TURNSTONE_LLM_BASE_URL: {{ .Values.llm.baseUrl | quote }} + {{- end }} + TURNSTONE_LLM_PROVIDER: {{ .Values.llm.provider | quote }} diff --git a/deploy/helm/turnstone/templates/deployment-bridge.yaml b/deploy/helm/turnstone/templates/deployment-bridge.yaml new file mode 100644 index 00000000..ec32d7d3 --- /dev/null +++ b/deploy/helm/turnstone/templates/deployment-bridge.yaml @@ -0,0 +1,45 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ include "turnstone.fullname" . }}-bridge + labels: + {{- include "turnstone.labels" . | nindent 4 }} + app.kubernetes.io/component: bridge +spec: + replicas: {{ .Values.bridge.replicas }} + selector: + matchLabels: + {{- include "turnstone.selectorLabels" . | nindent 6 }} + app.kubernetes.io/component: bridge + template: + metadata: + labels: + {{- include "turnstone.selectorLabels" . | nindent 8 }} + app.kubernetes.io/component: bridge + spec: + serviceAccountName: {{ include "turnstone.serviceAccountName" . }} + containers: + - name: bridge + image: {{ include "turnstone.image" . }} + imagePullPolicy: {{ .Values.image.pullPolicy }} + command: + - turnstone-bridge + - --server-url={{ printf "http://%s-server:%s" (include "turnstone.fullname" .) (.Values.server.service.port | toString) }} + - --redis-host={{ include "turnstone.redis.host" . }} + - --redis-port={{ include "turnstone.redis.port" . }} + envFrom: + - configMapRef: + name: {{ include "turnstone.fullname" . }}-config + - secretRef: + name: {{ include "turnstone.llm.secretName" . }} + optional: true + {{- if and .Values.auth.enabled .Values.auth.existingSecret }} + env: + - name: TURNSTONE_AUTH_TOKEN + valueFrom: + secretKeyRef: + name: {{ .Values.auth.existingSecret }} + key: TURNSTONE_AUTH_TOKEN + {{- end }} + resources: + {{- toYaml .Values.bridge.resources | nindent 12 }} diff --git a/deploy/helm/turnstone/templates/deployment-console.yaml b/deploy/helm/turnstone/templates/deployment-console.yaml new file mode 100644 index 00000000..d297210d --- /dev/null +++ b/deploy/helm/turnstone/templates/deployment-console.yaml @@ -0,0 +1,62 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ include "turnstone.fullname" . }}-console + labels: + {{- include "turnstone.labels" . | nindent 4 }} + app.kubernetes.io/component: console +spec: + replicas: {{ .Values.console.replicas }} + selector: + matchLabels: + {{- include "turnstone.selectorLabels" . | nindent 6 }} + app.kubernetes.io/component: console + template: + metadata: + labels: + {{- include "turnstone.selectorLabels" . | nindent 8 }} + app.kubernetes.io/component: console + spec: + serviceAccountName: {{ include "turnstone.serviceAccountName" . }} + containers: + - name: console + image: {{ include "turnstone.image" . }} + imagePullPolicy: {{ .Values.image.pullPolicy }} + command: + - turnstone-console + - --host=0.0.0.0 + - --port={{ .Values.console.service.port }} + - --redis-host={{ include "turnstone.redis.host" . }} + - --redis-port={{ include "turnstone.redis.port" . }} + ports: + - name: http + containerPort: {{ .Values.console.service.port }} + protocol: TCP + envFrom: + - configMapRef: + name: {{ include "turnstone.fullname" . }}-config + - secretRef: + name: {{ include "turnstone.llm.secretName" . }} + optional: true + {{- if and .Values.auth.enabled .Values.auth.existingSecret }} + env: + - name: TURNSTONE_AUTH_TOKEN + valueFrom: + secretKeyRef: + name: {{ .Values.auth.existingSecret }} + key: TURNSTONE_AUTH_TOKEN + {{- end }} + readinessProbe: + httpGet: + path: /health + port: http + initialDelaySeconds: 10 + periodSeconds: 10 + livenessProbe: + httpGet: + path: /health + port: http + initialDelaySeconds: 15 + periodSeconds: 20 + resources: + {{- toYaml .Values.console.resources | nindent 12 }} diff --git a/deploy/helm/turnstone/templates/deployment-server.yaml b/deploy/helm/turnstone/templates/deployment-server.yaml new file mode 100644 index 00000000..e1f1a39c --- /dev/null +++ b/deploy/helm/turnstone/templates/deployment-server.yaml @@ -0,0 +1,64 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ include "turnstone.fullname" . }}-server + labels: + {{- include "turnstone.labels" . | nindent 4 }} + app.kubernetes.io/component: server +spec: + replicas: {{ .Values.server.replicas }} + selector: + matchLabels: + {{- include "turnstone.selectorLabels" . | nindent 6 }} + app.kubernetes.io/component: server + template: + metadata: + labels: + {{- include "turnstone.selectorLabels" . | nindent 8 }} + app.kubernetes.io/component: server + spec: + serviceAccountName: {{ include "turnstone.serviceAccountName" . }} + containers: + - name: server + image: {{ include "turnstone.image" . }} + imagePullPolicy: {{ .Values.image.pullPolicy }} + command: + - turnstone-server + - --host + - "0.0.0.0" + - --port + - {{ .Values.server.service.port | quote }} + ports: + - name: http + containerPort: {{ .Values.server.service.port }} + protocol: TCP + envFrom: + - configMapRef: + name: {{ include "turnstone.fullname" . }}-config + - secretRef: + name: {{ include "turnstone.llm.secretName" . }} + optional: true + env: + - name: TURNSTONE_DB_URL + value: "postgresql+psycopg://$(TURNSTONE_DB_USER):$(POSTGRES_PASSWORD)@$(TURNSTONE_DB_HOST):$(TURNSTONE_DB_PORT)/$(TURNSTONE_DB_NAME)" + {{- if and .Values.auth.enabled .Values.auth.existingSecret }} + - name: TURNSTONE_AUTH_TOKEN + valueFrom: + secretKeyRef: + name: {{ .Values.auth.existingSecret }} + key: TURNSTONE_AUTH_TOKEN + {{- end }} + readinessProbe: + httpGet: + path: /health + port: http + initialDelaySeconds: 10 + periodSeconds: 10 + livenessProbe: + httpGet: + path: /health + port: http + initialDelaySeconds: 15 + periodSeconds: 20 + resources: + {{- toYaml .Values.server.resources | nindent 12 }} diff --git a/deploy/helm/turnstone/templates/ingress.yaml b/deploy/helm/turnstone/templates/ingress.yaml new file mode 100644 index 00000000..b1472d4d --- /dev/null +++ b/deploy/helm/turnstone/templates/ingress.yaml @@ -0,0 +1,47 @@ +{{- if .Values.ingress.enabled -}} +apiVersion: networking.k8s.io/v1 +kind: Ingress +metadata: + name: {{ include "turnstone.fullname" . }} + labels: + {{- include "turnstone.labels" . | nindent 4 }} + {{- with .Values.ingress.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +spec: + {{- if .Values.ingress.className }} + ingressClassName: {{ .Values.ingress.className }} + {{- end }} + {{- if .Values.ingress.tls }} + tls: + {{- range .Values.ingress.tls }} + - hosts: + {{- range .hosts }} + - {{ . | quote }} + {{- end }} + secretName: {{ .secretName }} + {{- end }} + {{- end }} + rules: + {{- range .Values.ingress.hosts }} + - host: {{ .host | quote }} + http: + paths: + {{- range .paths }} + - path: {{ .path }} + pathType: {{ .pathType | default "Prefix" }} + backend: + service: + {{- if eq (.service | default "server") "console" }} + name: {{ include "turnstone.fullname" $ }}-console + port: + number: {{ $.Values.console.service.port }} + {{- else }} + name: {{ include "turnstone.fullname" $ }}-server + port: + number: {{ $.Values.server.service.port }} + {{- end }} + {{- end }} + {{- end }} +{{- end }} diff --git a/deploy/helm/turnstone/templates/job-migrate.yaml b/deploy/helm/turnstone/templates/job-migrate.yaml new file mode 100644 index 00000000..9fb99eba --- /dev/null +++ b/deploy/helm/turnstone/templates/job-migrate.yaml @@ -0,0 +1,38 @@ +apiVersion: batch/v1 +kind: Job +metadata: + name: {{ include "turnstone.fullname" . }}-migrate + labels: + {{- include "turnstone.labels" . | nindent 4 }} + app.kubernetes.io/component: migrate + annotations: + "helm.sh/hook": pre-install,pre-upgrade + "helm.sh/hook-weight": "-1" + "helm.sh/hook-delete-policy": before-hook-creation,hook-succeeded +spec: + backoffLimit: 3 + template: + metadata: + labels: + {{- include "turnstone.selectorLabels" . | nindent 8 }} + app.kubernetes.io/component: migrate + spec: + serviceAccountName: {{ include "turnstone.serviceAccountName" . }} + restartPolicy: OnFailure + containers: + - name: migrate + image: {{ include "turnstone.image" . }} + imagePullPolicy: {{ .Values.image.pullPolicy }} + command: + - python + - -m + - turnstone.core.storage._migrate + envFrom: + - configMapRef: + name: {{ include "turnstone.fullname" . }}-config + - secretRef: + name: {{ include "turnstone.llm.secretName" . }} + optional: true + env: + - name: TURNSTONE_DB_URL + value: "postgresql+psycopg://$(TURNSTONE_DB_USER):$(POSTGRES_PASSWORD)@$(TURNSTONE_DB_HOST):$(TURNSTONE_DB_PORT)/$(TURNSTONE_DB_NAME)" diff --git a/deploy/helm/turnstone/templates/secret.yaml b/deploy/helm/turnstone/templates/secret.yaml new file mode 100644 index 00000000..6c786726 --- /dev/null +++ b/deploy/helm/turnstone/templates/secret.yaml @@ -0,0 +1,28 @@ +{{- if not .Values.llm.existingSecret }} +apiVersion: v1 +kind: Secret +metadata: + name: {{ include "turnstone.fullname" . }}-secrets + labels: + {{- include "turnstone.labels" . | nindent 4 }} +type: Opaque +data: + {{- if .Values.llm.apiKey }} + OPENAI_API_KEY: {{ .Values.llm.apiKey | b64enc | quote }} + {{- end }} + {{- if and .Values.postgresql.enabled .Values.postgresql.auth.password }} + POSTGRES_PASSWORD: {{ .Values.postgresql.auth.password | b64enc | quote }} + {{- else if and (not .Values.postgresql.enabled) .Values.database.external.password }} + POSTGRES_PASSWORD: {{ .Values.database.external.password | b64enc | quote }} + {{- end }} + {{- if and .Values.auth.enabled .Values.auth.token (not .Values.auth.existingSecret) }} + TURNSTONE_AUTH_TOKEN: {{ .Values.auth.token | b64enc | quote }} + {{- end }} + {{- if and .Values.redis.enabled .Values.redis.auth }} + {{- if .Values.redis.auth.password }} + REDIS_PASSWORD: {{ .Values.redis.auth.password | b64enc | quote }} + {{- end }} + {{- else if and (not .Values.redis.enabled) .Values.redis.external.password }} + REDIS_PASSWORD: {{ .Values.redis.external.password | b64enc | quote }} + {{- end }} +{{- end }} diff --git a/deploy/helm/turnstone/templates/service-console.yaml b/deploy/helm/turnstone/templates/service-console.yaml new file mode 100644 index 00000000..62381064 --- /dev/null +++ b/deploy/helm/turnstone/templates/service-console.yaml @@ -0,0 +1,17 @@ +apiVersion: v1 +kind: Service +metadata: + name: {{ include "turnstone.fullname" . }}-console + labels: + {{- include "turnstone.labels" . | nindent 4 }} + app.kubernetes.io/component: console +spec: + type: {{ .Values.console.service.type }} + ports: + - port: {{ .Values.console.service.port }} + targetPort: http + protocol: TCP + name: http + selector: + {{- include "turnstone.selectorLabels" . | nindent 4 }} + app.kubernetes.io/component: console diff --git a/deploy/helm/turnstone/templates/service-server.yaml b/deploy/helm/turnstone/templates/service-server.yaml new file mode 100644 index 00000000..46f4ddf9 --- /dev/null +++ b/deploy/helm/turnstone/templates/service-server.yaml @@ -0,0 +1,17 @@ +apiVersion: v1 +kind: Service +metadata: + name: {{ include "turnstone.fullname" . }}-server + labels: + {{- include "turnstone.labels" . | nindent 4 }} + app.kubernetes.io/component: server +spec: + type: {{ .Values.server.service.type }} + ports: + - port: {{ .Values.server.service.port }} + targetPort: http + protocol: TCP + name: http + selector: + {{- include "turnstone.selectorLabels" . | nindent 4 }} + app.kubernetes.io/component: server diff --git a/deploy/helm/turnstone/templates/serviceaccount.yaml b/deploy/helm/turnstone/templates/serviceaccount.yaml new file mode 100644 index 00000000..2622dad6 --- /dev/null +++ b/deploy/helm/turnstone/templates/serviceaccount.yaml @@ -0,0 +1,6 @@ +apiVersion: v1 +kind: ServiceAccount +metadata: + name: {{ include "turnstone.serviceAccountName" . }} + labels: + {{- include "turnstone.labels" . | nindent 4 }} diff --git a/deploy/helm/turnstone/values.yaml b/deploy/helm/turnstone/values.yaml new file mode 100644 index 00000000..915e0f40 --- /dev/null +++ b/deploy/helm/turnstone/values.yaml @@ -0,0 +1,103 @@ +# -- Container image settings +image: + repository: ghcr.io/turnstonelabs/turnstone + tag: "" + pullPolicy: IfNotPresent + +# -- Database configuration +database: + # Backend type (postgresql) + backend: postgresql + # External database settings (used when postgresql.enabled is false) + external: + host: "" + port: 5432 + database: turnstone + username: turnstone + existingSecret: "" + sslmode: prefer + +# -- Bitnami PostgreSQL subchart +postgresql: + enabled: true + auth: + database: turnstone + username: turnstone + +# -- Redis configuration +redis: + enabled: true + architecture: standalone + # External Redis settings (used when redis.enabled is false) + external: + host: "" + port: 6379 + existingSecret: "" + +# -- Turnstone server (main API + web UI) +server: + replicas: 1 + resources: + requests: + cpu: 500m + memory: 512Mi + limits: + cpu: "2" + memory: 1Gi + service: + type: ClusterIP + port: 8080 + +# -- Turnstone bridge (Redis MQ connector) +bridge: + replicas: 1 + resources: + requests: + cpu: 250m + memory: 256Mi + limits: + cpu: "1" + memory: 512Mi + +# -- Turnstone console (cluster dashboard) +console: + replicas: 1 + resources: + requests: + cpu: 250m + memory: 256Mi + limits: + cpu: "1" + memory: 512Mi + service: + type: ClusterIP + port: 8090 + +# -- LLM provider configuration +llm: + baseUrl: "" + provider: openai + apiKey: "" + existingSecret: "" + +# -- Authentication +auth: + enabled: false + token: "" + existingSecret: "" + +# -- Ingress configuration +ingress: + enabled: false + className: "" + annotations: {} + hosts: [] + # - host: turnstone.example.com + # paths: + # - path: / + # pathType: Prefix + # service: server + tls: [] + # - secretName: turnstone-tls + # hosts: + # - turnstone.example.com diff --git a/deploy/terraform/examples/aws-ecs-basic/main.tf b/deploy/terraform/examples/aws-ecs-basic/main.tf new file mode 100644 index 00000000..124d7d35 --- /dev/null +++ b/deploy/terraform/examples/aws-ecs-basic/main.tf @@ -0,0 +1,36 @@ +terraform { + required_version = ">= 1.5" + + required_providers { + aws = { + source = "hashicorp/aws" + version = ">= 5.0" + } + } +} + +provider "aws" { + region = var.aws_region +} + +module "turnstone" { + source = "../../modules/aws-ecs" + + vpc_id = var.vpc_id + private_subnet_ids = var.private_subnet_ids + public_subnet_ids = var.public_subnet_ids + + image_repository = var.image_repository + image_tag = var.image_tag + + llm_base_url = var.llm_base_url + openai_api_key = var.openai_api_key + + environment = var.environment + name_prefix = var.name_prefix + auth_token = var.auth_token + + tags = { + Example = "aws-ecs-basic" + } +} diff --git a/deploy/terraform/examples/aws-ecs-basic/outputs.tf b/deploy/terraform/examples/aws-ecs-basic/outputs.tf new file mode 100644 index 00000000..bb4f6a07 --- /dev/null +++ b/deploy/terraform/examples/aws-ecs-basic/outputs.tf @@ -0,0 +1,29 @@ +output "alb_dns_name" { + description = "DNS name of the Application Load Balancer." + value = module.turnstone.alb_dns_name +} + +output "server_url" { + description = "HTTP URL for the Turnstone server." + value = module.turnstone.server_url +} + +output "console_url" { + description = "HTTP URL for the Turnstone console." + value = module.turnstone.console_url +} + +output "cluster_arn" { + description = "ARN of the ECS cluster." + value = module.turnstone.cluster_arn +} + +output "rds_endpoint" { + description = "RDS PostgreSQL endpoint." + value = module.turnstone.rds_endpoint +} + +output "redis_endpoint" { + description = "ElastiCache Redis endpoint." + value = module.turnstone.redis_endpoint +} diff --git a/deploy/terraform/examples/aws-ecs-basic/terraform.tfvars.example b/deploy/terraform/examples/aws-ecs-basic/terraform.tfvars.example new file mode 100644 index 00000000..e25c494a --- /dev/null +++ b/deploy/terraform/examples/aws-ecs-basic/terraform.tfvars.example @@ -0,0 +1,22 @@ +# --- Required --- + +# VPC and subnet IDs from your existing AWS infrastructure. +# The VPC must have DNS support and DNS hostnames enabled. +vpc_id = "vpc-0123456789abcdef0" +private_subnet_ids = ["subnet-aaa111", "subnet-bbb222"] +public_subnet_ids = ["subnet-ccc333", "subnet-ddd444"] + +# LLM provider configuration. +# For OpenAI: https://api.openai.com/v1 +# For a self-hosted vLLM instance: http://your-vllm-host:8000/v1 +llm_base_url = "https://api.openai.com/v1" +openai_api_key = "sk-..." + +# --- Optional --- + +# aws_region = "us-east-1" +# image_repository = "ghcr.io/turnstonelabs/turnstone" +# image_tag = "0.3.0" +# environment = "production" +# name_prefix = "turnstone" +# auth_token = "my-secret-token" diff --git a/deploy/terraform/examples/aws-ecs-basic/variables.tf b/deploy/terraform/examples/aws-ecs-basic/variables.tf new file mode 100644 index 00000000..1d609039 --- /dev/null +++ b/deploy/terraform/examples/aws-ecs-basic/variables.tf @@ -0,0 +1,62 @@ +variable "aws_region" { + description = "AWS region to deploy into." + type = string + default = "us-east-1" +} + +variable "vpc_id" { + description = "ID of the VPC where all resources will be created." + type = string +} + +variable "private_subnet_ids" { + description = "List of private subnet IDs for ECS tasks, RDS, and ElastiCache." + type = list(string) +} + +variable "public_subnet_ids" { + description = "List of public subnet IDs for the Application Load Balancer." + type = list(string) +} + +variable "image_repository" { + description = "Container image repository." + type = string + default = "ghcr.io/turnstonelabs/turnstone" +} + +variable "image_tag" { + description = "Container image tag." + type = string + default = "latest" +} + +variable "llm_base_url" { + description = "Base URL for the LLM provider API." + type = string +} + +variable "openai_api_key" { + description = "API key for the LLM provider." + type = string + sensitive = true +} + +variable "environment" { + description = "Deployment environment name." + type = string + default = "production" +} + +variable "name_prefix" { + description = "Prefix for all resource names." + type = string + default = "turnstone" +} + +variable "auth_token" { + description = "Optional authentication token for the Turnstone API." + type = string + sensitive = true + default = "" +} diff --git a/deploy/terraform/modules/aws-ecs/alb.tf b/deploy/terraform/modules/aws-ecs/alb.tf new file mode 100644 index 00000000..bd39557b --- /dev/null +++ b/deploy/terraform/modules/aws-ecs/alb.tf @@ -0,0 +1,150 @@ +# ---------- Application Load Balancer ---------- +# +# HTTP listeners are provided as a starter baseline. For production, set +# var.certificate_arn to an ACM certificate ARN to enable HTTPS listeners +# that redirect HTTP traffic to TLS. + +resource "aws_lb" "this" { + name = "${var.name_prefix}-${var.environment}" + internal = false + load_balancer_type = "application" + security_groups = [aws_security_group.alb.id] + subnets = var.public_subnet_ids + tags = local.common_tags +} + +# ---------- Server Target Group + Listeners ---------- + +resource "aws_lb_target_group" "server" { + name = "${var.name_prefix}-server-${var.environment}" + port = 8080 + protocol = "HTTP" + vpc_id = var.vpc_id + target_type = "ip" + tags = local.common_tags + + health_check { + path = "/health" + port = "traffic-port" + protocol = "HTTP" + healthy_threshold = 2 + unhealthy_threshold = 3 + timeout = 5 + interval = 30 + matcher = "200" + } +} + +# HTTP listener: forwards directly when no certificate, redirects to HTTPS otherwise. +resource "aws_lb_listener" "server" { + count = var.certificate_arn == "" ? 1 : 0 + load_balancer_arn = aws_lb.this.arn + port = 80 + protocol = "HTTP" + tags = local.common_tags + + default_action { + type = "forward" + target_group_arn = aws_lb_target_group.server.arn + } +} + +resource "aws_lb_listener" "server_http_redirect" { + count = var.certificate_arn != "" ? 1 : 0 + load_balancer_arn = aws_lb.this.arn + port = 80 + protocol = "HTTP" + tags = local.common_tags + + default_action { + type = "redirect" + redirect { + port = "443" + protocol = "HTTPS" + status_code = "HTTP_301" + } + } +} + +resource "aws_lb_listener" "server_https" { + count = var.certificate_arn != "" ? 1 : 0 + load_balancer_arn = aws_lb.this.arn + port = 443 + protocol = "HTTPS" + ssl_policy = "ELBSecurityPolicy-TLS13-1-2-2021-06" + certificate_arn = var.certificate_arn + tags = local.common_tags + + default_action { + type = "forward" + target_group_arn = aws_lb_target_group.server.arn + } +} + +# ---------- Console Target Group + Listeners ---------- + +resource "aws_lb_target_group" "console" { + name = "${var.name_prefix}-console-${var.environment}" + port = 8090 + protocol = "HTTP" + vpc_id = var.vpc_id + target_type = "ip" + tags = local.common_tags + + health_check { + path = "/health" + port = "traffic-port" + protocol = "HTTP" + healthy_threshold = 2 + unhealthy_threshold = 3 + timeout = 5 + interval = 30 + matcher = "200" + } +} + +# HTTP listener: forwards directly when no certificate, redirects to HTTPS otherwise. +resource "aws_lb_listener" "console" { + count = var.certificate_arn == "" ? 1 : 0 + load_balancer_arn = aws_lb.this.arn + port = 8090 + protocol = "HTTP" + tags = local.common_tags + + default_action { + type = "forward" + target_group_arn = aws_lb_target_group.console.arn + } +} + +resource "aws_lb_listener" "console_http_redirect" { + count = var.certificate_arn != "" ? 1 : 0 + load_balancer_arn = aws_lb.this.arn + port = 8090 + protocol = "HTTP" + tags = local.common_tags + + default_action { + type = "redirect" + redirect { + port = "8443" + protocol = "HTTPS" + status_code = "HTTP_301" + } + } +} + +resource "aws_lb_listener" "console_https" { + count = var.certificate_arn != "" ? 1 : 0 + load_balancer_arn = aws_lb.this.arn + port = 8443 + protocol = "HTTPS" + ssl_policy = "ELBSecurityPolicy-TLS13-1-2-2021-06" + certificate_arn = var.certificate_arn + tags = local.common_tags + + default_action { + type = "forward" + target_group_arn = aws_lb_target_group.console.arn + } +} diff --git a/deploy/terraform/modules/aws-ecs/elasticache.tf b/deploy/terraform/modules/aws-ecs/elasticache.tf new file mode 100644 index 00000000..d6e68708 --- /dev/null +++ b/deploy/terraform/modules/aws-ecs/elasticache.tf @@ -0,0 +1,30 @@ +# ---------- ElastiCache Subnet Group ---------- + +resource "aws_elasticache_subnet_group" "this" { + name = "${var.name_prefix}-${var.environment}" + subnet_ids = var.private_subnet_ids + tags = local.common_tags +} + +# ---------- ElastiCache Redis Replication Group ---------- + +resource "aws_elasticache_replication_group" "this" { + replication_group_id = "${var.name_prefix}-${var.environment}" + description = "Turnstone Redis for MQ and session state" + + engine = "redis" + engine_version = "7.1" + node_type = var.redis_node_type + num_cache_clusters = 1 + port = 6379 + + subnet_group_name = aws_elasticache_subnet_group.this.name + security_group_ids = [aws_security_group.redis.id] + + at_rest_encryption_enabled = true + transit_encryption_enabled = true + + automatic_failover_enabled = false + + tags = local.common_tags +} diff --git a/deploy/terraform/modules/aws-ecs/iam.tf b/deploy/terraform/modules/aws-ecs/iam.tf new file mode 100644 index 00000000..c65a608c --- /dev/null +++ b/deploy/terraform/modules/aws-ecs/iam.tf @@ -0,0 +1,70 @@ +# ---------- ECS Task Execution Role ---------- +# Used by the ECS agent to pull images and retrieve secrets. + +resource "aws_iam_role" "ecs_execution" { + name = "${var.name_prefix}-ecs-execution-${var.environment}" + tags = local.common_tags + + assume_role_policy = jsonencode({ + Version = "2012-10-17" + Statement = [ + { + Effect = "Allow" + Principal = { + Service = "ecs-tasks.amazonaws.com" + } + Action = "sts:AssumeRole" + }, + ] + }) +} + +resource "aws_iam_role_policy_attachment" "ecs_execution_base" { + role = aws_iam_role.ecs_execution.name + policy_arn = "arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy" +} + +resource "aws_iam_role_policy" "ecs_execution_secrets" { + name = "${var.name_prefix}-secrets-read-${var.environment}" + role = aws_iam_role.ecs_execution.id + + policy = jsonencode({ + Version = "2012-10-17" + Statement = [ + { + Effect = "Allow" + Action = [ + "secretsmanager:GetSecretValue", + ] + Resource = concat( + [ + aws_secretsmanager_secret.openai_api_key.arn, + aws_secretsmanager_secret.db_password.arn, + ], + var.auth_token != "" ? [aws_secretsmanager_secret.auth_token[0].arn] : [], + ) + }, + ] + }) +} + +# ---------- ECS Task Role ---------- +# Assumed by the running container. Minimal permissions; extend as needed. + +resource "aws_iam_role" "ecs_task" { + name = "${var.name_prefix}-ecs-task-${var.environment}" + tags = local.common_tags + + assume_role_policy = jsonencode({ + Version = "2012-10-17" + Statement = [ + { + Effect = "Allow" + Principal = { + Service = "ecs-tasks.amazonaws.com" + } + Action = "sts:AssumeRole" + }, + ] + }) +} diff --git a/deploy/terraform/modules/aws-ecs/main.tf b/deploy/terraform/modules/aws-ecs/main.tf new file mode 100644 index 00000000..10b25a18 --- /dev/null +++ b/deploy/terraform/modules/aws-ecs/main.tf @@ -0,0 +1,313 @@ +terraform { + required_version = ">= 1.5" + + required_providers { + aws = { + source = "hashicorp/aws" + version = ">= 5.0" + } + random = { + source = "hashicorp/random" + version = ">= 3.5" + } + } +} + +locals { + full_image = "${var.image_repository}:${var.image_tag}" + + common_tags = merge(var.tags, { + Project = "turnstone" + Environment = var.environment + ManagedBy = "terraform" + }) + + # Shared environment variables injected into every container. + common_env = [ + { name = "TURNSTONE_ENV", value = var.environment }, + { name = "TURNSTONE_DB_BACKEND", value = "postgresql" }, + { name = "TURNSTONE_LLM_BASE_URL", value = var.llm_base_url }, + { name = "TURNSTONE_REDIS_URL", value = "redis://${aws_elasticache_replication_group.this.primary_endpoint_address}:6379/0" }, + ] + + # Secrets pulled from Secrets Manager at container start. + common_secrets = [ + { + name = "OPENAI_API_KEY" + valueFrom = aws_secretsmanager_secret_version.openai_api_key.arn + }, + { + name = "TURNSTONE_DB_URL" + valueFrom = aws_secretsmanager_secret_version.db_url.arn + }, + ] + + auth_env = var.auth_token != "" ? [ + { name = "TURNSTONE_AUTH_ENABLED", value = "true" }, + ] : [] + + auth_secrets = var.auth_token != "" ? [ + { + name = "TURNSTONE_AUTH_TOKEN" + valueFrom = aws_secretsmanager_secret_version.auth_token[0].arn + }, + ] : [] +} + +# ---------- Secrets Manager ---------- + +resource "aws_secretsmanager_secret" "openai_api_key" { + name = "${var.name_prefix}-${var.environment}-openai-api-key" + tags = local.common_tags +} + +resource "aws_secretsmanager_secret_version" "openai_api_key" { + secret_id = aws_secretsmanager_secret.openai_api_key.id + secret_string = var.openai_api_key +} + +resource "aws_secretsmanager_secret" "auth_token" { + count = var.auth_token != "" ? 1 : 0 + name = "${var.name_prefix}-${var.environment}-auth-token" + tags = local.common_tags +} + +resource "aws_secretsmanager_secret_version" "auth_token" { + count = var.auth_token != "" ? 1 : 0 + secret_id = aws_secretsmanager_secret.auth_token[0].id + secret_string = var.auth_token +} + +resource "aws_secretsmanager_secret" "db_password" { + name = "${var.name_prefix}-${var.environment}-db-password" + tags = local.common_tags +} + +resource "aws_secretsmanager_secret_version" "db_password" { + secret_id = aws_secretsmanager_secret.db_password.id + secret_string = random_password.db.result +} + +resource "aws_secretsmanager_secret" "db_url" { + name = "${var.name_prefix}-${var.environment}-db-url" + tags = local.common_tags +} + +resource "aws_secretsmanager_secret_version" "db_url" { + secret_id = aws_secretsmanager_secret.db_url.id + secret_string = "postgresql+psycopg://${aws_db_instance.this.username}:${random_password.db.result}@${aws_db_instance.this.endpoint}/turnstone" +} + +# ---------- ECS Cluster ---------- + +resource "aws_ecs_cluster" "this" { + name = "${var.name_prefix}-${var.environment}" + tags = local.common_tags + + setting { + name = "containerInsights" + value = "enabled" + } +} + +# ---------- CloudWatch Log Group ---------- + +resource "aws_cloudwatch_log_group" "this" { + name = "/ecs/${var.name_prefix}-${var.environment}" + retention_in_days = 30 + tags = local.common_tags +} + +# ---------- Server Task Definition + Service ---------- + +resource "aws_ecs_task_definition" "server" { + family = "${var.name_prefix}-server" + requires_compatibilities = ["FARGATE"] + network_mode = "awsvpc" + cpu = var.server_cpu + memory = var.server_memory + execution_role_arn = aws_iam_role.ecs_execution.arn + task_role_arn = aws_iam_role.ecs_task.arn + tags = local.common_tags + + container_definitions = jsonencode([ + { + name = "server" + image = local.full_image + essential = true + command = ["turnstone-server", "--host", "0.0.0.0", "--port", "8080"] + + portMappings = [ + { containerPort = 8080, protocol = "tcp" }, + ] + + environment = concat(local.common_env, local.auth_env) + secrets = concat(local.common_secrets, local.auth_secrets) + + logConfiguration = { + logDriver = "awslogs" + options = { + "awslogs-group" = aws_cloudwatch_log_group.this.name + "awslogs-region" = data.aws_region.current.name + "awslogs-stream-prefix" = "server" + } + } + + healthCheck = { + command = ["CMD-SHELL", "curl -f http://localhost:8080/health || exit 1"] + interval = 30 + timeout = 5 + retries = 3 + startPeriod = 10 + } + }, + ]) +} + +resource "aws_ecs_service" "server" { + name = "${var.name_prefix}-server" + cluster = aws_ecs_cluster.this.id + task_definition = aws_ecs_task_definition.server.arn + desired_count = 1 + launch_type = "FARGATE" + tags = local.common_tags + + network_configuration { + subnets = var.private_subnet_ids + security_groups = [aws_security_group.ecs_tasks.id] + assign_public_ip = false + } + + load_balancer { + target_group_arn = aws_lb_target_group.server.arn + container_name = "server" + container_port = 8080 + } + + depends_on = [aws_lb_target_group.server] +} + +# ---------- Bridge Task Definition + Service ---------- + +resource "aws_ecs_task_definition" "bridge" { + family = "${var.name_prefix}-bridge" + requires_compatibilities = ["FARGATE"] + network_mode = "awsvpc" + cpu = var.bridge_cpu + memory = var.bridge_memory + execution_role_arn = aws_iam_role.ecs_execution.arn + task_role_arn = aws_iam_role.ecs_task.arn + tags = local.common_tags + + container_definitions = jsonencode([ + { + name = "bridge" + image = local.full_image + essential = true + command = ["turnstone-bridge"] + + environment = concat(local.common_env, local.auth_env) + secrets = concat(local.common_secrets, local.auth_secrets) + + logConfiguration = { + logDriver = "awslogs" + options = { + "awslogs-group" = aws_cloudwatch_log_group.this.name + "awslogs-region" = data.aws_region.current.name + "awslogs-stream-prefix" = "bridge" + } + } + }, + ]) +} + +resource "aws_ecs_service" "bridge" { + name = "${var.name_prefix}-bridge" + cluster = aws_ecs_cluster.this.id + task_definition = aws_ecs_task_definition.bridge.arn + desired_count = 1 + launch_type = "FARGATE" + tags = local.common_tags + + network_configuration { + subnets = var.private_subnet_ids + security_groups = [aws_security_group.ecs_tasks.id] + assign_public_ip = false + } + + depends_on = [aws_ecs_service.server] +} + +# ---------- Console Task Definition + Service ---------- + +resource "aws_ecs_task_definition" "console" { + family = "${var.name_prefix}-console" + requires_compatibilities = ["FARGATE"] + network_mode = "awsvpc" + cpu = var.console_cpu + memory = var.console_memory + execution_role_arn = aws_iam_role.ecs_execution.arn + task_role_arn = aws_iam_role.ecs_task.arn + tags = local.common_tags + + container_definitions = jsonencode([ + { + name = "console" + image = local.full_image + essential = true + command = ["turnstone-console", "--host", "0.0.0.0", "--port", "8090"] + + portMappings = [ + { containerPort = 8090, protocol = "tcp" }, + ] + + environment = concat(local.common_env, local.auth_env) + secrets = concat(local.common_secrets, local.auth_secrets) + + logConfiguration = { + logDriver = "awslogs" + options = { + "awslogs-group" = aws_cloudwatch_log_group.this.name + "awslogs-region" = data.aws_region.current.name + "awslogs-stream-prefix" = "console" + } + } + + healthCheck = { + command = ["CMD-SHELL", "curl -f http://localhost:8090/health || exit 1"] + interval = 30 + timeout = 5 + retries = 3 + startPeriod = 10 + } + }, + ]) +} + +resource "aws_ecs_service" "console" { + name = "${var.name_prefix}-console" + cluster = aws_ecs_cluster.this.id + task_definition = aws_ecs_task_definition.console.arn + desired_count = 1 + launch_type = "FARGATE" + tags = local.common_tags + + network_configuration { + subnets = var.private_subnet_ids + security_groups = [aws_security_group.ecs_tasks.id] + assign_public_ip = false + } + + load_balancer { + target_group_arn = aws_lb_target_group.console.arn + container_name = "console" + container_port = 8090 + } + + depends_on = [aws_lb_target_group.console] +} + +# ---------- Data Sources ---------- + +data "aws_region" "current" {} +data "aws_caller_identity" "current" {} diff --git a/deploy/terraform/modules/aws-ecs/outputs.tf b/deploy/terraform/modules/aws-ecs/outputs.tf new file mode 100644 index 00000000..da47a6e3 --- /dev/null +++ b/deploy/terraform/modules/aws-ecs/outputs.tf @@ -0,0 +1,29 @@ +output "alb_dns_name" { + description = "DNS name of the Application Load Balancer." + value = aws_lb.this.dns_name +} + +output "server_url" { + description = "HTTP URL for the Turnstone server API and web UI." + value = "http://${aws_lb.this.dns_name}" +} + +output "console_url" { + description = "HTTP URL for the Turnstone console dashboard." + value = "http://${aws_lb.this.dns_name}:8090" +} + +output "cluster_arn" { + description = "ARN of the ECS cluster." + value = aws_ecs_cluster.this.arn +} + +output "rds_endpoint" { + description = "Endpoint of the RDS PostgreSQL instance (host:port)." + value = aws_db_instance.this.endpoint +} + +output "redis_endpoint" { + description = "Primary endpoint of the ElastiCache Redis replication group." + value = aws_elasticache_replication_group.this.primary_endpoint_address +} diff --git a/deploy/terraform/modules/aws-ecs/rds.tf b/deploy/terraform/modules/aws-ecs/rds.tf new file mode 100644 index 00000000..e3e746a5 --- /dev/null +++ b/deploy/terraform/modules/aws-ecs/rds.tf @@ -0,0 +1,42 @@ +# ---------- Random Password ---------- + +resource "random_password" "db" { + length = 32 + special = false +} + +# ---------- DB Subnet Group ---------- + +resource "aws_db_subnet_group" "this" { + name = "${var.name_prefix}-${var.environment}" + subnet_ids = var.private_subnet_ids + tags = local.common_tags +} + +# ---------- RDS PostgreSQL ---------- + +resource "aws_db_instance" "this" { + identifier = "${var.name_prefix}-${var.environment}" + + engine = "postgres" + engine_version = "17" + instance_class = var.db_instance_class + allocated_storage = 20 + storage_type = "gp3" + storage_encrypted = true + deletion_protection = true + skip_final_snapshot = false + final_snapshot_identifier = "${var.name_prefix}-${var.environment}-final" + + db_name = "turnstone" + username = "turnstone" + password = random_password.db.result + + db_subnet_group_name = aws_db_subnet_group.this.name + vpc_security_group_ids = [aws_security_group.rds.id] + + backup_retention_period = 7 + multi_az = false + + tags = local.common_tags +} diff --git a/deploy/terraform/modules/aws-ecs/security.tf b/deploy/terraform/modules/aws-ecs/security.tf new file mode 100644 index 00000000..a3bd2246 --- /dev/null +++ b/deploy/terraform/modules/aws-ecs/security.tf @@ -0,0 +1,133 @@ +# ---------- ALB Security Group ---------- + +resource "aws_security_group" "alb" { + name = "${var.name_prefix}-alb-${var.environment}" + description = "Allow inbound HTTP to ALB for server and console" + vpc_id = var.vpc_id + tags = local.common_tags +} + +resource "aws_vpc_security_group_ingress_rule" "alb_http" { + security_group_id = aws_security_group.alb.id + description = "HTTP traffic to server" + from_port = 80 + to_port = 80 + ip_protocol = "tcp" + cidr_ipv4 = "0.0.0.0/0" + tags = local.common_tags +} + +resource "aws_vpc_security_group_ingress_rule" "alb_https" { + count = var.certificate_arn != "" ? 1 : 0 + security_group_id = aws_security_group.alb.id + description = "HTTPS traffic to server" + from_port = 443 + to_port = 443 + ip_protocol = "tcp" + cidr_ipv4 = "0.0.0.0/0" + tags = local.common_tags +} + +resource "aws_vpc_security_group_ingress_rule" "alb_console" { + security_group_id = aws_security_group.alb.id + description = "HTTP traffic to console" + from_port = 8090 + to_port = 8090 + ip_protocol = "tcp" + cidr_ipv4 = "0.0.0.0/0" + tags = local.common_tags +} + +resource "aws_vpc_security_group_ingress_rule" "alb_console_https" { + count = var.certificate_arn != "" ? 1 : 0 + security_group_id = aws_security_group.alb.id + description = "HTTPS traffic to console" + from_port = 8443 + to_port = 8443 + ip_protocol = "tcp" + cidr_ipv4 = "0.0.0.0/0" + tags = local.common_tags +} + +resource "aws_vpc_security_group_egress_rule" "alb_all" { + security_group_id = aws_security_group.alb.id + description = "Allow all outbound" + ip_protocol = "-1" + cidr_ipv4 = "0.0.0.0/0" + tags = local.common_tags +} + +# ---------- ECS Tasks Security Group ---------- + +resource "aws_security_group" "ecs_tasks" { + name = "${var.name_prefix}-ecs-tasks-${var.environment}" + description = "Allow traffic from ALB to ECS tasks and outbound internet" + vpc_id = var.vpc_id + tags = local.common_tags +} + +resource "aws_vpc_security_group_ingress_rule" "ecs_from_alb_server" { + security_group_id = aws_security_group.ecs_tasks.id + description = "Server port from ALB" + from_port = 8080 + to_port = 8080 + ip_protocol = "tcp" + referenced_security_group_id = aws_security_group.alb.id + tags = local.common_tags +} + +resource "aws_vpc_security_group_ingress_rule" "ecs_from_alb_console" { + security_group_id = aws_security_group.ecs_tasks.id + description = "Console port from ALB" + from_port = 8090 + to_port = 8090 + ip_protocol = "tcp" + referenced_security_group_id = aws_security_group.alb.id + tags = local.common_tags +} + +resource "aws_vpc_security_group_egress_rule" "ecs_all" { + security_group_id = aws_security_group.ecs_tasks.id + description = "Allow all outbound (LLM APIs, ECR, Secrets Manager, etc.)" + ip_protocol = "-1" + cidr_ipv4 = "0.0.0.0/0" + tags = local.common_tags +} + +# ---------- RDS Security Group ---------- + +resource "aws_security_group" "rds" { + name = "${var.name_prefix}-rds-${var.environment}" + description = "Allow PostgreSQL access from ECS tasks" + vpc_id = var.vpc_id + tags = local.common_tags +} + +resource "aws_vpc_security_group_ingress_rule" "rds_from_ecs" { + security_group_id = aws_security_group.rds.id + description = "PostgreSQL from ECS tasks" + from_port = 5432 + to_port = 5432 + ip_protocol = "tcp" + referenced_security_group_id = aws_security_group.ecs_tasks.id + tags = local.common_tags +} + +# ---------- Redis Security Group ---------- + +resource "aws_security_group" "redis" { + name = "${var.name_prefix}-redis-${var.environment}" + description = "Allow Redis access from ECS tasks" + vpc_id = var.vpc_id + tags = local.common_tags +} + +resource "aws_vpc_security_group_ingress_rule" "redis_from_ecs" { + security_group_id = aws_security_group.redis.id + description = "Redis from ECS tasks" + from_port = 6379 + to_port = 6379 + ip_protocol = "tcp" + referenced_security_group_id = aws_security_group.ecs_tasks.id + tags = local.common_tags +} diff --git a/deploy/terraform/modules/aws-ecs/variables.tf b/deploy/terraform/modules/aws-ecs/variables.tf new file mode 100644 index 00000000..9e3bad22 --- /dev/null +++ b/deploy/terraform/modules/aws-ecs/variables.tf @@ -0,0 +1,130 @@ +# --- Networking --- + +variable "vpc_id" { + description = "ID of the VPC where all resources will be created." + type = string +} + +variable "private_subnet_ids" { + description = "List of private subnet IDs for ECS tasks, RDS, and ElastiCache." + type = list(string) +} + +variable "public_subnet_ids" { + description = "List of public subnet IDs for the Application Load Balancer." + type = list(string) +} + +# --- Container Image --- + +variable "image_repository" { + description = "Container image repository." + type = string + default = "ghcr.io/turnstonelabs/turnstone" +} + +variable "image_tag" { + description = "Container image tag." + type = string + default = "latest" +} + +# --- LLM Provider --- + +variable "llm_base_url" { + description = "Base URL for the LLM provider API (e.g. https://api.openai.com/v1)." + type = string +} + +variable "openai_api_key" { + description = "API key for the LLM provider. Stored in AWS Secrets Manager." + type = string + sensitive = true +} + +# --- RDS --- + +variable "db_instance_class" { + description = "RDS instance class for PostgreSQL." + type = string + default = "db.t4g.micro" +} + +# --- ElastiCache --- + +variable "redis_node_type" { + description = "ElastiCache node type for Redis." + type = string + default = "cache.t4g.micro" +} + +# --- ECS Task Sizing --- + +variable "server_cpu" { + description = "CPU units for the server task (1 vCPU = 1024)." + type = number + default = 512 +} + +variable "server_memory" { + description = "Memory (MiB) for the server task." + type = number + default = 1024 +} + +variable "bridge_cpu" { + description = "CPU units for the bridge task." + type = number + default = 256 +} + +variable "bridge_memory" { + description = "Memory (MiB) for the bridge task." + type = number + default = 512 +} + +variable "console_cpu" { + description = "CPU units for the console task." + type = number + default = 256 +} + +variable "console_memory" { + description = "Memory (MiB) for the console task." + type = number + default = 512 +} + +# --- General --- + +variable "environment" { + description = "Deployment environment name (e.g. production, staging)." + type = string + default = "production" +} + +variable "name_prefix" { + description = "Prefix for all resource names." + type = string + default = "turnstone" +} + +variable "auth_token" { + description = "Optional authentication token for the Turnstone API. Empty string disables auth." + type = string + sensitive = true + default = "" +} + +variable "certificate_arn" { + description = "ACM certificate ARN for HTTPS listeners. Leave empty for HTTP-only (not recommended for production)." + type = string + default = "" +} + +variable "tags" { + description = "Additional tags to apply to all resources." + type = map(string) + default = {} +} diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh new file mode 100755 index 00000000..5a9481b0 --- /dev/null +++ b/docker/entrypoint.sh @@ -0,0 +1,5 @@ +#!/bin/sh +# Run database migrations before starting the service +python -m turnstone.core.storage._migrate || true +# Execute the actual command +exec "$@" diff --git a/docs/architecture.md b/docs/architecture.md index 146065a8..e1268125 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -42,7 +42,8 @@ turnstone/ tools.py Tool schema loader (JSON -> OpenAI function-calling format) mcp_client.py MCPClientManager — MCP server connections, tool discovery, async-sync bridge model_registry.py ModelRegistry — named model configs, lazy client creation, fallback routing - memory.py SQLite persistence (conversations, memories, FTS5 search) + memory.py Persistence facade (delegates to storage backend) + storage/ Pluggable storage: StorageBackend protocol, SQLite + PostgreSQL metrics.py Prometheus-compatible metrics collector (MetricsCollector) healthcheck.py BackendHealthMonitor — periodic probe + circuit breaker ratelimit.py Per-IP token-bucket rate limiter (RateLimiter, TokenBucket) @@ -637,11 +638,37 @@ This truncation message is visible to the model, so it knows output was cut. ## Persistence -### Database +### Storage Architecture -SQLite via `turnstone.core.memory`. Database file: `.turnstone.db` in the -current working directory (overridable via `memory.db_override` for eval -isolation). +Persistence is managed by the `turnstone.core.storage` package — a pluggable +backend behind a `StorageBackend` protocol. The `memory.py` facade provides +backward-compatible module-level functions that delegate to the active backend. + +``` +session.py / server.py / cli.py + ↓ + memory.py (facade — silent-failure wrappers) + ↓ + storage._registry (singleton factory) + ↓ + ┌─────────────┐ ┌──────────────────┐ + │ SQLiteBackend │ │ PostgreSQLBackend │ + │ (FTS5 search) │ │ (tsvector/ILIKE) │ + └─────────────┘ └──────────────────┘ + ↓ ↓ + storage._schema (SQLAlchemy Core tables — single source of truth) + ↓ + storage._migrate (programmatic Alembic) +``` + +**SQLite** is the default (zero-config, single file at `.turnstone.db`). +**PostgreSQL** is the production backend (connection pooling, `tsvector` +full-text search). Select via `[database]` in `config.toml`, CLI flags, or +environment variables (`TURNSTONE_DB_BACKEND`, `TURNSTONE_DB_URL`). + +Schema migrations are managed by Alembic and run automatically on startup. +Existing SQLite databases created before the migration system are auto-stamped +at the baseline revision. ### Tables @@ -668,34 +695,53 @@ conversations tool_name TEXT tool_args TEXT tool_call_id TEXT -- links tool_call ↔ tool_result for resume + provider_data TEXT -- raw provider content (e.g. Anthropic encrypted) -conversations_fts -- FTS5 virtual table +session_config + session_id TEXT NOT NULL -- composite PK with key + key TEXT NOT NULL + value TEXT + +conversations_fts -- SQLite FTS5 virtual table (optional) content (content=conversations, content_rowid=id) ``` -The `tool_call_id` column was added via schema migration (`ALTER TABLE`) for -backwards compatibility with existing databases. +Table definitions live in `storage/_schema.py` (SQLAlchemy Core `Table` objects) +and are the single source of truth for both backends and Alembic migrations. -### Key Functions +### StorageBackend Protocol -| Function | Purpose | -|----------|---------| -| `open_db()` | Open/create database, run migrations, initialize tables | -| `load_memories()` | Return all `(key, value)` pairs sorted by key | -| `save_message(session_id, role, content, ...)` | Log a message to conversations (accepts `tool_call_id`) | -| `search_history(query, limit)` | Full-text search via FTS5 (falls back to LIKE) | -| `search_history_recent(limit)` | Return most recent messages | +| Method | Purpose | +|--------|---------| | `register_session(session_id, title)` | Create a sessions row (no-op if exists) | -| `update_session_title(session_id, title)` | Set/update LLM-generated title | +| `save_message(session_id, role, content, ...)` | Log a message to conversations | +| `load_session_messages(session_id)` | Reconstruct OpenAI message format from DB rows | +| `list_sessions(limit)` | List sessions with >=1 message, ordered by updated DESC | +| `delete_session(session_id)` | Delete session and all its messages | +| `prune_sessions(retention_days)` | Remove empty sessions and old unnamed sessions | +| `resolve_session(alias_or_id)` | Resolve alias, exact id, or id prefix to full session_id | +| `save_session_config(session_id, config)` | Persist session configuration key/value pairs | +| `load_session_config(session_id)` | Retrieve session configuration | | `set_session_alias(session_id, alias)` | Set user-friendly alias (returns False if taken) | | `get_session_name(session_id)` | Return alias if set, else title, else None | -| `resolve_session(alias_or_id)` | Resolve alias, exact id, or id prefix to full session_id | -| `list_sessions(limit)` | List sessions with ≥1 message, ordered by updated DESC | -| `load_session_messages(session_id)` | Reconstruct OpenAI message format from DB rows | -| `delete_session(session_id)` | Delete session and all its messages | -| `prune_sessions(retention_days, log_fn)` | Remove empty sessions and old unnamed sessions; called at startup | -| `normalize_key(key)` | Normalize memory keys (`lower`, replace `-`/` ` with `_`) | -| `fts5_query(query)` | Convert plain text to safe FTS5 query (quoted terms) | +| `update_session_title(session_id, title)` | Set/update LLM-generated title | +| `kv_get(key)` / `kv_set(key, value)` / `kv_delete(key)` | Generic key-value store (backs memories table) | +| `kv_list()` / `kv_search(query)` | List or search key-value pairs | +| `search_history(query, limit)` | Full-text search (FTS5 on SQLite, tsvector on PostgreSQL) | +| `search_history_recent(limit)` | Return most recent messages | +| `close()` | Release resources (connection pool, engine) | + +### Database Configuration + +```toml +[database] +backend = "sqlite" # "sqlite" | "postgresql" +path = ".turnstone.db" # SQLite file path +url = "" # PostgreSQL connection URL +pool_size = 5 # PostgreSQL connection pool size +``` + +Environment variables: `TURNSTONE_DB_BACKEND`, `TURNSTONE_DB_URL`, `TURNSTONE_DB_PATH`. ### Session Persistence and Resume diff --git a/docs/diagrams/02-package-structure.puml b/docs/diagrams/02-package-structure.puml index 537742db..088ad1ac 100644 --- a/docs/diagrams/02-package-structure.puml +++ b/docs/diagrams/02-package-structure.puml @@ -29,7 +29,8 @@ package "turnstone/core/" <> { component [providers/\nLLMProvider, OpenAI, Anthropic] as providers <> component [workstream.py\nWorkstreamManager] as workstream <> component [tools.py\nTool loader] as tools <> - component [memory.py\nSQLite + FTS5] as memory <> + component [memory.py\nPersistence facade] as memory <> + component [storage/\nStorageBackend protocol\nSQLite + PostgreSQL] as storage <> component [metrics.py\nPrometheus metrics] as metrics <> component [config.py\nTOML config] as config <> component [safety.py\nPath validation] as safety <> @@ -128,6 +129,7 @@ chat --> session session --> providers session --> tools session --> memory +memory --> storage session --> safety session --> sandbox session --> edit diff --git a/docs/diagrams/14-storage-architecture.puml b/docs/diagrams/14-storage-architecture.puml new file mode 100644 index 00000000..366a25e7 --- /dev/null +++ b/docs/diagrams/14-storage-architecture.puml @@ -0,0 +1,156 @@ +@startuml +!theme plain +title Turnstone — Storage Architecture + +skinparam class { + BackgroundColor<> #E8EAF6 + BackgroundColor<> #C8E6C9 + BackgroundColor<> #B3E5FC + BackgroundColor<> #FFF9C4 + BackgroundColor<> #FFE0B2 + BackgroundColor<> #F3E5F5 +} + +' -- Protocol -- +interface "StorageBackend" as SB <> { + +register_session(session_id, title) + +save_message(session_id, role, content, ...) + +load_session_messages(session_id) → list[dict] + +list_sessions(limit) → list + +delete_session(session_id) → bool + +prune_sessions(retention_days) → (int, int) + +resolve_session(alias_or_id) → str | None + +save_session_config(session_id, config) + +load_session_config(session_id) → dict + +set_session_alias(session_id, alias) → bool + +get_session_name(session_id) → str | None + +update_session_title(session_id, title) + +kv_get(key) → str | None + +kv_set(key, value) → str | None + +kv_delete(key) → bool + +kv_list() → list[(str, str)] + +kv_search(query) → list[(str, str)] + +search_history(query, limit) → list + +search_history_recent(limit) → list + +close() +} + +' -- Backends -- +class "SQLiteBackend" as SQLite <> { + -_engine: sa.Engine + -_fts5_available: bool + +__init__(path: str) + -- + FTS5 full-text search + Default pool, check_same_thread=False +} + +class "PostgreSQLBackend" as PG <> { + -_engine: sa.Engine + +__init__(url: str, pool_size: int) + -- + tsvector + ILIKE search + Connection pooling +} + +' -- Schema -- +class "_schema.py" as Schema <> { + +metadata: MetaData + +memories: Table + +conversations: Table + +sessions: Table + +session_config: Table + -- + SQLAlchemy Core + Single source of truth +} + +' -- Migration -- +class "_migrate.py" as Migrate <> { + +run_migrations(storage, backend) + -_bootstrap_existing_sqlite() + -- + Programmatic Alembic + Auto-bootstrap existing DBs +} + +class "migrations/" as Versions <> { + 001_initial_schema.py +} + +' -- Registry -- +class "_registry.py" as Registry { + -_storage: StorageBackend | None + +init_storage(backend, path, url) → StorageBackend + +get_storage() → StorageBackend + +reset_storage() + -- + Auto-initializes SQLite + if not configured +} + +' -- Facade -- +class "memory.py" as Facade <> { + +register_session() + +save_message() + +load_session_messages() + +save_memory() / delete_memory() + +search_memories() + +... (all 18 functions) + -- + Thin delegation to + get_storage() + Silent failure behavior +} + +' -- Consumers -- +class "session.py\nChatSession" as Session { +} + +class "server.py\nWeb UI" as Server { +} + +class "cli.py\nTerminal" as CLI { +} + +' -- Relationships -- +SQLite ..|> SB +PG ..|> SB + +SQLite --> Schema : uses +PG --> Schema : uses + +Registry --> SB : creates +Registry --> Migrate : calls + +Migrate --> Versions : applies +Migrate --> Schema : references + +Facade --> Registry : get_storage() + +Session --> Facade : imports +Server --> Facade : imports +CLI --> Facade : imports + +' -- Config -- +note right of Registry + [database] + backend = "sqlite" | "postgresql" + url = "postgresql+psycopg://..." + path = ".turnstone.db" + pool_size = 5 +end note + +note bottom of SQLite + Default backend. + Zero-config for + single-node / dev. +end note + +note bottom of PG + Production backend. + Multi-node / Docker + default. +end note + +@enduml diff --git a/docs/diagrams/png/02-package-structure.png b/docs/diagrams/png/02-package-structure.png index a4c1c786..f6e94944 100644 --- a/docs/diagrams/png/02-package-structure.png +++ b/docs/diagrams/png/02-package-structure.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:ab184b57aff615d64082434faf09ec444bd2f26643269c37d3b51fa6868b45da -size 326359 +oid sha256:7d75da92a657525bcbb7a425dc6c8d3cafe074c3ac3bff3cf4b1d44aea607b50 +size 330156 diff --git a/docs/diagrams/png/14-storage-architecture.png b/docs/diagrams/png/14-storage-architecture.png new file mode 100644 index 00000000..1bc7d2cb --- /dev/null +++ b/docs/diagrams/png/14-storage-architecture.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0c615984373b4893b6cc5755604f137541e9d391862122746a4fcbae63543563 +size 201041 diff --git a/pyproject.toml b/pyproject.toml index f03505e9..b46e785b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,8 @@ dependencies = [ "sse-starlette>=2.0", "httpx-sse>=0.4", "pydantic>=2.0", + "sqlalchemy>=2.0", + "alembic>=1.14", ] [project.urls] @@ -44,6 +46,7 @@ mq = ["redis>=7.2"] console = ["redis>=7.2"] sim = ["redis>=7.2"] anthropic = ["anthropic>=0.39"] +postgres = ["psycopg[binary]>=3.2"] [project.scripts] @@ -121,6 +124,10 @@ ignore_missing_imports = true module = ["sse_starlette", "sse_starlette.*", "uvicorn", "uvicorn.*", "httpx_sse", "httpx_sse.*"] ignore_missing_imports = true +[[tool.mypy.overrides]] +module = ["sqlalchemy", "sqlalchemy.*", "alembic", "alembic.*"] +ignore_missing_imports = true + [[tool.mypy.overrides]] module = ["anthropic", "anthropic.*"] ignore_missing_imports = true diff --git a/tests/conftest.py b/tests/conftest.py index 1413884a..3d539c80 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,15 +4,15 @@ import pytest @pytest.fixture -def tmp_db(tmp_path, monkeypatch): - """Provide a temporary SQLite database.""" - import turnstone.core.memory as memory +def tmp_db(tmp_path): + """Provide a temporary SQLite storage backend.""" + from turnstone.core.storage import init_storage, reset_storage db_path = str(tmp_path / "test.db") - monkeypatch.setattr(memory, "db_override", db_path) - memory.db_initialized.discard(db_path) + reset_storage() + init_storage("sqlite", path=db_path, run_migrations=False) yield db_path - memory.db_initialized.discard(db_path) + reset_storage() @pytest.fixture diff --git a/tests/test_config.py b/tests/test_config.py index 7c91e298..55b9cd3e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -153,27 +153,23 @@ def test_apply_config_model_section(tmp_path, monkeypatch): def test_tavily_key_from_config(tmp_path, monkeypatch): """get_tavily_key() reads from config.toml [api] tavily_key.""" _reset_cache() - import turnstone.core.memory as mem - - mem._tavily_key = None - mem._tavily_key_loaded = False + config_mod._tavily_key = None + config_mod._tavily_key_loaded = False cfg = tmp_path / "config.toml" cfg.write_text('[api]\ntavily_key = "tvly-from-config"\n') monkeypatch.setattr(config_mod, "CONFIG_PATH", cfg) monkeypatch.delenv("TAVILY_API_KEY", raising=False) - key = mem.get_tavily_key() + key = config_mod.get_tavily_key() assert key == "tvly-from-config" def test_tavily_key_fallback_to_env(tmp_path, monkeypatch): """get_tavily_key() falls back to $TAVILY_API_KEY env var.""" _reset_cache() - import turnstone.core.memory as mem - - mem._tavily_key = None - mem._tavily_key_loaded = False + config_mod._tavily_key = None + config_mod._tavily_key_loaded = False # Config exists but no tavily_key in it cfg = tmp_path / "config.toml" @@ -181,5 +177,5 @@ def test_tavily_key_fallback_to_env(tmp_path, monkeypatch): monkeypatch.setattr(config_mod, "CONFIG_PATH", cfg) monkeypatch.setenv("TAVILY_API_KEY", "tvly-from-env") - key = mem.get_tavily_key() + key = config_mod.get_tavily_key() assert key == "tvly-from-env" diff --git a/tests/test_db.py b/tests/test_db.py index 788cd0fb..5c3de26c 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -1,38 +1,30 @@ """Tests for turnstone.core.memory — database operations.""" +import sqlalchemy as sa + from turnstone.core.memory import ( normalize_key, - open_db, save_message, search_history, search_history_recent, ) +from turnstone.core.storage import get_storage -class TestOpenDb: +class TestSchemaCreation: def test_creates_tables(self, tmp_db): - conn = open_db() - try: - # Check memories table exists + engine = get_storage()._engine # noqa: SLF001 + with engine.connect() as conn: rows = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='memories'" + sa.text("SELECT name FROM sqlite_master WHERE type='table' AND name='memories'") ).fetchall() assert len(rows) == 1 - - # Check conversations table exists rows = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='conversations'" + sa.text( + "SELECT name FROM sqlite_master WHERE type='table' AND name='conversations'" + ) ).fetchall() assert len(rows) == 1 - finally: - conn.close() - - def test_idempotent_open(self, tmp_db): - # Opening twice should not raise - conn1 = open_db() - conn1.close() - conn2 = open_db() - conn2.close() class TestSaveAndSearchHistory: @@ -40,7 +32,6 @@ class TestSaveAndSearchHistory: save_message("sess1", "user", "hello world test message") results = search_history("hello") assert len(results) >= 1 - # Result tuple: (timestamp, session_id, role, content, tool_name) found = any(r[3] == "hello world test message" for r in results) assert found diff --git a/tests/test_fts5.py b/tests/test_fts5.py index eed2f4a5..eb495da7 100644 --- a/tests/test_fts5.py +++ b/tests/test_fts5.py @@ -1,50 +1,48 @@ -"""Tests for turnstone.core.memory — fts5_query and escape_like.""" +"""Tests for SQLite FTS5 query building and LIKE escaping.""" -from turnstone.core.memory import escape_like, fts5_query +from turnstone.core.storage._sqlite import _escape_like, _fts5_query class TestFts5Query: def test_single_word(self): - result = fts5_query("hello") + result = _fts5_query("hello") assert result == '"hello"' def test_multiple_words_joined_with_and(self): - result = fts5_query("hello world") - # Each word is quoted; space between = implicit AND + result = _fts5_query("hello world") assert result == '"hello" "world"' def test_special_chars_safely_quoted(self): - result = fts5_query("test*") + result = _fts5_query("test*") assert result == '"test*"' def test_dash_safely_quoted(self): - result = fts5_query("-negative") + result = _fts5_query("-negative") assert result == '"-negative"' def test_embedded_double_quotes(self): - # Double quotes inside a term are doubled per FTS5 convention - result = fts5_query('say"hello') + result = _fts5_query('say"hello') assert result == '"say""hello"' def test_empty_query(self): - assert fts5_query("") == "" + assert _fts5_query("") == "" def test_whitespace_only(self): - assert fts5_query(" ") == "" + assert _fts5_query(" ") == "" class TestEscapeLike: def test_percent_escaped(self): - assert escape_like("100%") == "100\\%" + assert _escape_like("100%") == "100\\%" def test_underscore_escaped(self): - assert escape_like("a_b") == "a\\_b" + assert _escape_like("a_b") == "a\\_b" def test_backslash_escaped(self): - assert escape_like("a\\b") == "a\\\\b" + assert _escape_like("a\\b") == "a\\\\b" def test_no_metacharacters(self): - assert escape_like("hello") == "hello" + assert _escape_like("hello") == "hello" def test_combined(self): - assert escape_like("50%_off\\sale") == "50\\%\\_off\\\\sale" + assert _escape_like("50%_off\\sale") == "50\\%\\_off\\\\sale" diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 3054091d..aaaad850 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -243,11 +243,13 @@ class TestMCPClientManager: class TestSessionIntegration: @pytest.fixture() - def tmp_db(self, tmp_path, monkeypatch): - monkeypatch.setenv("TURNSTONE_DB_PATH", str(tmp_path / "test.db")) - from turnstone.core.memory import open_db + def tmp_db(self, tmp_path): + from turnstone.core.storage import init_storage, reset_storage - open_db() + reset_storage() + init_storage("sqlite", path=str(tmp_path / "test.db"), run_migrations=False) + yield + reset_storage() def _make_session(self, mcp_client=None, **kwargs): from turnstone.core.session import ChatSession diff --git a/tests/test_server_live.py b/tests/test_server_live.py index b48c55ae..2a45b316 100644 --- a/tests/test_server_live.py +++ b/tests/test_server_live.py @@ -30,8 +30,8 @@ import httpx import pytest from openai import OpenAI -import turnstone.core.memory as _memory_module from turnstone.core.session import ChatSession +from turnstone.core.storage import init_storage, reset_storage # --------------------------------------------------------------------------- # Fixtures @@ -126,12 +126,10 @@ def tmp_db(): """Temp DB to avoid polluting real conversation history.""" with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: path = f.name - old = _memory_module.db_override - _memory_module.db_override = path - _memory_module.db_initialized.discard(path) + reset_storage() + init_storage("sqlite", path=path, run_migrations=False) yield path - _memory_module.db_override = old - _memory_module.db_initialized.discard(path) + reset_storage() os.unlink(path) diff --git a/tests/test_sessions.py b/tests/test_sessions.py index 727e85bf..ca2d0da7 100644 --- a/tests/test_sessions.py +++ b/tests/test_sessions.py @@ -2,12 +2,13 @@ from unittest.mock import MagicMock +import sqlalchemy as sa + from turnstone.core.memory import ( delete_session, list_sessions, load_session_config, load_session_messages, - open_db, prune_sessions, register_session, resolve_session, @@ -17,6 +18,7 @@ from turnstone.core.memory import ( update_session_title, ) from turnstone.core.session import ChatSession +from turnstone.core.storage import get_storage # ── Session registration ────────────────────────────────────────────── @@ -225,25 +227,21 @@ class TestDeleteSession: class TestSaveMessageToolCallId: def test_tool_call_id_stored(self, tmp_db): save_message("s1", "tool_call", None, "bash", '{"cmd":"ls"}', tool_call_id="call_xyz") - conn = open_db() - try: + engine = get_storage()._engine # noqa: SLF001 + with engine.connect() as conn: row = conn.execute( - "SELECT tool_call_id FROM conversations WHERE session_id = 's1'" + sa.text("SELECT tool_call_id FROM conversations WHERE session_id = 's1'") ).fetchone() assert row[0] == "call_xyz" - finally: - conn.close() def test_tool_call_id_none_by_default(self, tmp_db): save_message("s1", "user", "hello") - conn = open_db() - try: + engine = get_storage()._engine # noqa: SLF001 + with engine.connect() as conn: row = conn.execute( - "SELECT tool_call_id FROM conversations WHERE session_id = 's1'" + sa.text("SELECT tool_call_id FROM conversations WHERE session_id = 's1'") ).fetchone() assert row[0] is None - finally: - conn.close() # ── Sessions table creation ─────────────────────────────────────────── @@ -251,22 +249,18 @@ class TestSaveMessageToolCallId: class TestSessionsTable: def test_sessions_table_exists(self, tmp_db): - conn = open_db() - try: + engine = get_storage()._engine # noqa: SLF001 + with engine.connect() as conn: rows = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='sessions'" + sa.text("SELECT name FROM sqlite_master WHERE type='table' AND name='sessions'") ).fetchall() assert len(rows) == 1 - finally: - conn.close() def test_tool_call_id_column_exists(self, tmp_db): - conn = open_db() - try: + engine = get_storage()._engine # noqa: SLF001 + with engine.connect() as conn: # Should not raise - conn.execute("SELECT tool_call_id FROM conversations LIMIT 0") - finally: - conn.close() + conn.execute(sa.text("SELECT tool_call_id FROM conversations LIMIT 0")) # ── ChatSession.resume_session ──────────────────────────────────────── @@ -495,10 +489,12 @@ class TestPruneSessions: register_session("old1") save_message("old1", "user", "ancient message") # Force the updated timestamp to the past so it looks stale - conn = open_db() - conn.execute("UPDATE sessions SET updated = '2020-01-01' WHERE session_id = 'old1'") - conn.commit() - conn.close() + engine = get_storage()._engine # noqa: SLF001 + with engine.connect() as conn: + conn.execute( + sa.text("UPDATE sessions SET updated = '2020-01-01' WHERE session_id = 'old1'") + ) + conn.commit() _orphans, stale = prune_sessions(retention_days=30) assert stale == 1 @@ -508,10 +504,12 @@ class TestPruneSessions: set_session_alias("old2", "important") save_message("old2", "user", "old but named") # Force old timestamp - conn = open_db() - conn.execute("UPDATE sessions SET updated = '2020-01-01' WHERE session_id = 'old2'") - conn.commit() - conn.close() + engine = get_storage()._engine # noqa: SLF001 + with engine.connect() as conn: + conn.execute( + sa.text("UPDATE sessions SET updated = '2020-01-01' WHERE session_id = 'old2'") + ) + conn.commit() _orphans, stale = prune_sessions(retention_days=30) assert stale == 0 assert len(list_sessions()) == 1 @@ -532,10 +530,12 @@ class TestPruneSessions: register_session("stale_cfg") save_message("stale_cfg", "user", "old") save_session_config("stale_cfg", {"temperature": "0.9"}) - conn = open_db() - conn.execute("UPDATE sessions SET updated = '2020-01-01' WHERE session_id = 'stale_cfg'") - conn.commit() - conn.close() + engine = get_storage()._engine # noqa: SLF001 + with engine.connect() as conn: + conn.execute( + sa.text("UPDATE sessions SET updated = '2020-01-01' WHERE session_id = 'stale_cfg'") + ) + conn.commit() # Both should have config before prune assert load_session_config("orphan_cfg") == {"temperature": "0.5"} diff --git a/tests/test_storage_registry.py b/tests/test_storage_registry.py new file mode 100644 index 00000000..2e5df097 --- /dev/null +++ b/tests/test_storage_registry.py @@ -0,0 +1,54 @@ +"""Tests for the storage backend registry.""" + +import pytest + +from turnstone.core.storage import get_storage, init_storage, reset_storage +from turnstone.core.storage._sqlite import SQLiteBackend + + +@pytest.fixture(autouse=True) +def _clean_registry(): + """Reset the storage registry before and after each test.""" + reset_storage() + yield + reset_storage() + + +class TestInitStorage: + def test_sqlite_default(self, tmp_path): + backend = init_storage("sqlite", path=str(tmp_path / "test.db"), run_migrations=False) + assert isinstance(backend, SQLiteBackend) + + def test_unknown_backend_raises(self): + with pytest.raises(ValueError, match="Unknown storage backend"): + init_storage("mongodb") + + def test_postgresql_requires_url(self): + with pytest.raises(ValueError, match="requires a connection URL"): + init_storage("postgresql") + + +class TestGetStorage: + def test_auto_init(self, tmp_path, monkeypatch): + """get_storage() auto-initializes with SQLite if not yet initialized.""" + monkeypatch.chdir(tmp_path) + storage = get_storage() + assert isinstance(storage, SQLiteBackend) + + def test_returns_same_instance(self, tmp_path): + init_storage("sqlite", path=str(tmp_path / "test.db"), run_migrations=False) + s1 = get_storage() + s2 = get_storage() + assert s1 is s2 + + +class TestResetStorage: + def test_reset_clears_singleton(self, tmp_path): + init_storage("sqlite", path=str(tmp_path / "test.db"), run_migrations=False) + s1 = get_storage() + reset_storage() + # After reset, get_storage() auto-inits a new instance + monkeypatch_not_needed = True # noqa: F841 + init_storage("sqlite", path=str(tmp_path / "test2.db"), run_migrations=False) + s2 = get_storage() + assert s1 is not s2 diff --git a/tests/test_storage_sqlite.py b/tests/test_storage_sqlite.py new file mode 100644 index 00000000..6880484c --- /dev/null +++ b/tests/test_storage_sqlite.py @@ -0,0 +1,262 @@ +"""Tests for the SQLite storage backend.""" + +import pytest + +from turnstone.core.storage import init_storage, reset_storage + + +@pytest.fixture +def backend(tmp_path): + """Create a fresh SQLiteBackend for each test.""" + reset_storage() + b = init_storage("sqlite", path=str(tmp_path / "test.db"), run_migrations=False) + yield b + reset_storage() + + +# -- Session operations -------------------------------------------------------- + + +class TestRegisterSession: + def test_register_creates_session(self, backend): + backend.register_session("s1", title="Test") + name = backend.get_session_name("s1") + assert name == "Test" + + def test_register_idempotent(self, backend): + backend.register_session("s1", title="First") + backend.register_session("s1", title="Second") + name = backend.get_session_name("s1") + assert name == "First" # INSERT OR IGNORE preserves first + + +class TestSaveAndLoadMessages: + def test_roundtrip(self, backend): + backend.register_session("s1") + backend.save_message("s1", "user", "hello") + backend.save_message("s1", "assistant", "world") + msgs = backend.load_session_messages("s1") + assert len(msgs) == 2 + assert msgs[0]["role"] == "user" + assert msgs[0]["content"] == "hello" + assert msgs[1]["role"] == "assistant" + assert msgs[1]["content"] == "world" + + def test_tool_call_grouping(self, backend): + backend.register_session("s1") + backend.save_message("s1", "user", "do something") + backend.save_message("s1", "tool_call", None, "bash", '{"cmd":"ls"}', tool_call_id="c1") + backend.save_message("s1", "tool_result", "file.txt", tool_call_id="c1") + backend.save_message("s1", "assistant", "done") + msgs = backend.load_session_messages("s1") + assert len(msgs) == 4 + assert msgs[1]["role"] == "assistant" + assert len(msgs[1]["tool_calls"]) == 1 + assert msgs[1]["tool_calls"][0]["id"] == "c1" + assert msgs[2]["role"] == "tool" + assert msgs[2]["content"] == "file.txt" + + def test_incomplete_turn_repair(self, backend): + backend.register_session("s1") + backend.save_message("s1", "user", "do something") + backend.save_message("s1", "tool_call", None, "bash", '{"cmd":"ls"}', tool_call_id="c1") + backend.save_message("s1", "tool_call", None, "read", '{"path":"a"}', tool_call_id="c2") + # Only 1 result for 2 calls — incomplete turn + backend.save_message("s1", "tool_result", "ok", tool_call_id="c1") + msgs = backend.load_session_messages("s1") + # Incomplete turn should be stripped + assert len(msgs) == 1 # only the user message remains + + def test_provider_data_preserved(self, backend): + import json + + backend.register_session("s1") + pd = json.dumps({"encrypted": True}) + backend.save_message("s1", "assistant", "hi", provider_data=pd) + msgs = backend.load_session_messages("s1") + assert msgs[0].get("_provider_content") == {"encrypted": True} + + def test_empty_session_returns_empty(self, backend): + assert backend.load_session_messages("nonexistent") == [] + + +class TestListSessions: + def test_lists_sessions_with_messages(self, backend): + backend.register_session("s1") + backend.save_message("s1", "user", "hi") + backend.register_session("s2") # no messages + rows = backend.list_sessions() + assert len(rows) == 1 + assert rows[0][0] == "s1" + + def test_respects_limit(self, backend): + for i in range(5): + sid = f"s{i}" + backend.register_session(sid) + backend.save_message(sid, "user", f"msg {i}") + rows = backend.list_sessions(limit=3) + assert len(rows) == 3 + + +class TestDeleteSession: + def test_deletes_all_data(self, backend): + backend.register_session("s1") + backend.save_message("s1", "user", "hi") + backend.save_session_config("s1", {"temp": "0.5"}) + assert backend.delete_session("s1") + assert backend.load_session_messages("s1") == [] + assert backend.load_session_config("s1") == {} + assert backend.get_session_name("s1") is None + + +class TestPruneSessions: + def test_orphan_removed(self, backend): + backend.register_session("orphan") + orphans, stale = backend.prune_sessions() + assert orphans == 1 + + def test_stale_removed(self, backend): + import sqlalchemy as sa + + backend.register_session("old") + backend.save_message("old", "user", "hi") + # Force old timestamp + with backend._engine.connect() as conn: + conn.execute( + sa.text("UPDATE sessions SET updated = '2020-01-01' WHERE session_id = 'old'") + ) + conn.commit() + _, stale = backend.prune_sessions(retention_days=30) + assert stale == 1 + + +class TestResolveSession: + def test_exact_alias(self, backend): + backend.register_session("s1") + backend.set_session_alias("s1", "myalias") + assert backend.resolve_session("myalias") == "s1" + + def test_exact_id(self, backend): + backend.register_session("abc-123-def") + assert backend.resolve_session("abc-123-def") == "abc-123-def" + + def test_prefix_match(self, backend): + backend.register_session("abc-123-def") + assert backend.resolve_session("abc") == "abc-123-def" + + def test_not_found(self, backend): + assert backend.resolve_session("nonexistent") is None + + +# -- Session config ------------------------------------------------------------ + + +class TestSessionConfig: + def test_roundtrip(self, backend): + backend.register_session("s1") + backend.save_session_config("s1", {"temperature": "0.7", "effort": "high"}) + cfg = backend.load_session_config("s1") + assert cfg == {"temperature": "0.7", "effort": "high"} + + def test_empty_config(self, backend): + assert backend.load_session_config("nonexistent") == {} + + +# -- Session metadata ---------------------------------------------------------- + + +class TestSessionMetadata: + def test_alias(self, backend): + backend.register_session("s1") + assert backend.set_session_alias("s1", "my-session") + assert backend.get_session_name("s1") == "my-session" + + def test_alias_conflict(self, backend): + backend.register_session("s1") + backend.register_session("s2") + backend.set_session_alias("s1", "taken") + assert not backend.set_session_alias("s2", "taken") + + def test_title(self, backend): + backend.register_session("s1") + backend.update_session_title("s1", "My Title") + assert backend.get_session_name("s1") == "My Title" + + def test_alias_preferred_over_title(self, backend): + backend.register_session("s1") + backend.update_session_title("s1", "Title") + backend.set_session_alias("s1", "Alias") + assert backend.get_session_name("s1") == "Alias" + + +# -- Key-value store ----------------------------------------------------------- + + +class TestKVStore: + def test_set_and_get(self, backend): + assert backend.kv_set("key1", "value1") is None # no previous + assert backend.kv_get("key1") == "value1" + + def test_set_returns_old_value(self, backend): + backend.kv_set("key1", "v1") + old = backend.kv_set("key1", "v2") + assert old == "v1" + assert backend.kv_get("key1") == "v2" + + def test_delete(self, backend): + backend.kv_set("key1", "v1") + assert backend.kv_delete("key1") + assert backend.kv_get("key1") is None + + def test_delete_nonexistent(self, backend): + assert not backend.kv_delete("nope") + + def test_list(self, backend): + backend.kv_set("b", "2") + backend.kv_set("a", "1") + assert backend.kv_list() == [("a", "1"), ("b", "2")] + + def test_search(self, backend): + backend.kv_set("project_name", "turnstone") + backend.kv_set("version", "0.3") + results = backend.kv_search("turnstone") + assert len(results) == 1 + assert results[0] == ("project_name", "turnstone") + + def test_search_empty_lists_all(self, backend): + backend.kv_set("a", "1") + backend.kv_set("b", "2") + assert len(backend.kv_search("")) == 2 + + +# -- Conversation search ------------------------------------------------------- + + +class TestSearch: + def test_search_history(self, backend): + backend.register_session("s1") + backend.save_message("s1", "user", "hello world") + backend.save_message("s1", "user", "goodbye world") + results = backend.search_history("hello") + assert len(results) >= 1 + assert any("hello" in str(r[3]) for r in results) + + def test_search_recent(self, backend): + backend.register_session("s1") + backend.save_message("s1", "user", "msg1") + backend.save_message("s1", "user", "msg2") + results = backend.search_history_recent(limit=1) + assert len(results) == 1 + + +# -- Lifecycle ----------------------------------------------------------------- + + +class TestLifecycle: + def test_close(self, backend): + backend.close() # Should not raise + + def test_isinstance_check(self, backend): + from turnstone.core.storage._protocol import StorageBackend + + assert isinstance(backend, StorageBackend) diff --git a/turnstone/chat.py b/turnstone/chat.py index ab5cdf3e..9907bf60 100755 --- a/turnstone/chat.py +++ b/turnstone/chat.py @@ -16,9 +16,6 @@ All functionality has been moved to submodules: # Re-export public API for backward compatibility from turnstone.cli import detect_model, main # noqa: F401 -from turnstone.core.memory import ( # noqa: F401 - open_db as _open_db, -) from turnstone.core.session import ChatSession, SessionUI # noqa: F401 from turnstone.core.tools import AGENT_TOOLS, TASK_AGENT_TOOLS, TOOLS # noqa: F401 from turnstone.core.web import strip_html as _strip_html # noqa: F401 diff --git a/turnstone/cli.py b/turnstone/cli.py index d1fc1087..b3f5c408 100644 --- a/turnstone/cli.py +++ b/turnstone/cli.py @@ -825,9 +825,22 @@ def main() -> None: ) from turnstone.core.config import apply_config - apply_config(parser, ["api", "model", "session", "tools", "console", "auth", "mcp"]) + apply_config(parser, ["api", "model", "session", "tools", "console", "auth", "mcp", "database"]) args = parser.parse_args() + # Initialize storage backend + from turnstone.core.storage import init_storage + + db_backend = getattr(args, "db_backend", None) or os.environ.get( + "TURNSTONE_DB_BACKEND", "sqlite" + ) + db_url = getattr(args, "db_url", None) or os.environ.get("TURNSTONE_DB_URL", "") + db_path = getattr(args, "db_path", None) or os.environ.get("TURNSTONE_DB_PATH", "") + db_pool_size = int( + getattr(args, "db_pool_size", None) or os.environ.get("TURNSTONE_DB_POOL_SIZE", "5") + ) + init_storage(db_backend, path=db_path, url=db_url, pool_size=db_pool_size) + # Prune stale / empty sessions on startup from turnstone.core.memory import prune_sessions diff --git a/turnstone/core/config.py b/turnstone/core/config.py index c73dc760..6a761e05 100644 --- a/turnstone/core/config.py +++ b/turnstone/core/config.py @@ -115,8 +115,40 @@ _CONFIG_MAP: dict[str, dict[str, str]] = { "circuit_breaker_threshold": "circuit_breaker_threshold", "circuit_breaker_cooldown": "circuit_breaker_cooldown", }, + "database": { + "backend": "db_backend", + "url": "db_url", + "path": "db_path", + "pool_size": "db_pool_size", + }, } +# -- Tavily API key (cached) -------------------------------------------------- + +_tavily_key: str | None = None +_tavily_key_loaded: bool = False + + +def get_tavily_key() -> str | None: + """Load Tavily API key (cached after first call). + + Precedence: config.toml [api] tavily_key -> $TAVILY_API_KEY + """ + import os + + global _tavily_key, _tavily_key_loaded + if _tavily_key_loaded: + return _tavily_key + _tavily_key_loaded = True + cfg_key = load_config("api").get("tavily_key", "").strip() + if cfg_key: + _tavily_key = cfg_key + return _tavily_key + env_key = os.environ.get("TAVILY_API_KEY", "").strip() + if env_key: + _tavily_key = env_key + return _tavily_key + def apply_config(parser: argparse.ArgumentParser, sections: list[str]) -> None: """Set argparse defaults from config file. diff --git a/turnstone/core/memory.py b/turnstone/core/memory.py index a462274f..88844fc3 100644 --- a/turnstone/core/memory.py +++ b/turnstone/core/memory.py @@ -1,126 +1,33 @@ -"""SQLite database for persistent memories and conversation history.""" +"""Persistence facade — delegates to the pluggable storage backend. + +All functions maintain their existing signatures for consumers (session.py, +server.py, cli.py). The actual storage implementation lives in +``turnstone.core.storage``. +""" from __future__ import annotations import contextlib -import json -import os -import sqlite3 -from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any +from turnstone.core.storage import get_storage + if TYPE_CHECKING: from collections.abc import Callable -TURNSTONE_DB = os.path.join(os.getcwd(), ".turnstone.db") -db_override: str | None = None -db_initialized: set[str] = set() -_fts5_available: bool = False - -_tavily_key: str | None = None -_tavily_key_loaded: bool = False - - -def get_tavily_key() -> str | None: - """Load Tavily API key (cached after first call). - - Precedence: config.toml [api] tavily_key → $TAVILY_API_KEY - """ - global _tavily_key, _tavily_key_loaded - if _tavily_key_loaded: - return _tavily_key - _tavily_key_loaded = True - from turnstone.core.config import load_config - - cfg_key = load_config("api").get("tavily_key", "").strip() - if cfg_key: - _tavily_key = cfg_key - return _tavily_key - env_key = os.environ.get("TAVILY_API_KEY", "").strip() - if env_key: - _tavily_key = env_key - return _tavily_key - - -def open_db() -> sqlite3.Connection: - """Open the turnstone database, creating tables on first use per path.""" - global _fts5_available - path = db_override or TURNSTONE_DB - conn = sqlite3.connect(path) - if path not in db_initialized: - conn.execute( - "CREATE TABLE IF NOT EXISTS memories " - "(key TEXT PRIMARY KEY, value TEXT NOT NULL, " - "created TEXT NOT NULL, updated TEXT NOT NULL)" - ) - conn.execute( - "CREATE TABLE IF NOT EXISTS conversations " - "(id INTEGER PRIMARY KEY AUTOINCREMENT, " - "session_id TEXT NOT NULL, timestamp TEXT NOT NULL, " - "role TEXT NOT NULL, content TEXT, " - "tool_name TEXT, tool_args TEXT)" - ) - conn.execute("CREATE INDEX IF NOT EXISTS idx_conv_session ON conversations(session_id)") - # Migration: add tool_call_id column if missing (for session resume) - try: - conn.execute("SELECT tool_call_id FROM conversations LIMIT 0") - except sqlite3.OperationalError: - conn.execute("ALTER TABLE conversations ADD COLUMN tool_call_id TEXT") - conn.commit() - # Migration: add provider_data column for raw provider content blocks - try: - conn.execute("SELECT provider_data FROM conversations LIMIT 0") - except sqlite3.OperationalError: - conn.execute("ALTER TABLE conversations ADD COLUMN provider_data TEXT") - conn.commit() - # Sessions table — maps session_id to human-friendly alias/title - conn.execute( - "CREATE TABLE IF NOT EXISTS sessions " - "(session_id TEXT PRIMARY KEY, alias TEXT UNIQUE, " - "title TEXT, created TEXT NOT NULL, updated TEXT NOT NULL)" - ) - conn.execute("CREATE INDEX IF NOT EXISTS idx_sessions_alias ON sessions(alias)") - conn.execute("CREATE INDEX IF NOT EXISTS idx_sessions_updated ON sessions(updated)") - # Session config — persists LLM-affecting parameters across resume - conn.execute( - "CREATE TABLE IF NOT EXISTS session_config " - "(session_id TEXT NOT NULL, key TEXT NOT NULL, value TEXT, " - "PRIMARY KEY (session_id, key))" - ) - try: - # Check if FTS table already exists - fts_exists = conn.execute( - "SELECT 1 FROM sqlite_master WHERE type='table' AND name='conversations_fts'" - ).fetchone() - if not fts_exists: - conn.execute( - "CREATE VIRTUAL TABLE conversations_fts " - "USING fts5(content, content=conversations, content_rowid=id)" - ) - conn.execute("INSERT INTO conversations_fts(conversations_fts) VALUES('rebuild')") - conn.commit() - _fts5_available = True - except Exception: - _fts5_available = False - db_initialized.add(path) - return conn - def normalize_key(key: str) -> str: """Normalize a memory key for consistent lookup.""" return key.lower().replace("-", "_").replace(" ", "_") -def load_memories() -> list[tuple[str, str]]: - """Return all (key, value) pairs sorted by key.""" - try: - conn = open_db() - try: - return conn.execute("SELECT key, value FROM memories ORDER BY key").fetchall() - finally: - conn.close() - except Exception: - return [] +# -- Core session operations --------------------------------------------------- + + +def register_session(session_id: str, title: str | None = None) -> None: + """Create a sessions row for a new session (no-op if already exists).""" + with contextlib.suppress(Exception): + get_storage().register_session(session_id, title) def save_message( @@ -133,298 +40,46 @@ def save_message( provider_data: str | None = None, ) -> None: """Log a message to the conversations table.""" - global _fts5_available + with contextlib.suppress(Exception): + get_storage().save_message( + session_id, role, content, tool_name, tool_args, tool_call_id, provider_data + ) + + +def load_session_messages(session_id: str) -> list[dict[str, Any]]: + """Load messages for a session and reconstruct OpenAI message format.""" try: - conn = open_db() - try: - conn.execute( - "INSERT INTO conversations (session_id, timestamp, role, content, " - "tool_name, tool_args, tool_call_id, provider_data) " - "VALUES (?, datetime('now'), ?, ?, ?, ?, ?, ?)", - (session_id, role, content, tool_name, tool_args, tool_call_id, provider_data), - ) - if _fts5_available and content: - try: - rowid = conn.execute("SELECT last_insert_rowid()").fetchone()[0] - conn.execute( - "INSERT INTO conversations_fts(rowid, content) VALUES (?, ?)", - (rowid, content), - ) - except Exception: - _fts5_available = False # degrade to LIKE for rest of session - # Bump session updated timestamp - conn.execute( - "UPDATE sessions SET updated = datetime('now') WHERE session_id = ?", - (session_id,), - ) - conn.commit() - finally: - conn.close() - except Exception: - pass # Don't let logging failures break the session - - -def escape_like(s: str) -> str: - """Escape LIKE metacharacters for use with ESCAPE '\\\\'.""" - return s.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") - - -def fts5_query(query: str) -> str: - """Convert a plain search string into a safe FTS5 query. - - Quotes each term so FTS5 special characters (*, -, etc.) are treated - as literals, then joins with implicit AND. Embedded double quotes - are doubled per FTS5 quoting convention. - """ - terms = query.split() - safe = [] - for t in terms: - if t: - safe.append(f'"{t.replace(chr(34), chr(34) + chr(34))}"') - return " ".join(safe) - - -def search_history(query: str, limit: int = 20) -> list[tuple[Any, ...]]: - """Search conversation history. Returns (timestamp, session_id, role, content, tool_name).""" - if not query or not query.strip(): - return [] - try: - conn = open_db() - try: - if _fts5_available: - return conn.execute( - "SELECT c.timestamp, c.session_id, c.role, c.content, c.tool_name " - "FROM conversations_fts f " - "JOIN conversations c ON c.id = f.rowid " - "WHERE conversations_fts MATCH ? " - "ORDER BY f.rank ASC LIMIT ?", - (fts5_query(query), min(limit, 100)), - ).fetchall() - return conn.execute( - "SELECT timestamp, session_id, role, content, tool_name " - "FROM conversations WHERE content LIKE ? ESCAPE '\\' " - "ORDER BY timestamp DESC LIMIT ?", - (f"%{escape_like(query)}%", min(limit, 100)), - ).fetchall() - finally: - conn.close() + return get_storage().load_session_messages(session_id) except Exception: return [] -def search_history_recent(limit: int = 20) -> list[tuple[Any, ...]]: - """Return most recent conversation messages.""" +# -- Session management -------------------------------------------------------- + + +def list_sessions(limit: int = 20) -> list[Any]: + """List recent sessions with message counts.""" try: - conn = open_db() - try: - return conn.execute( - "SELECT timestamp, session_id, role, content, tool_name " - "FROM conversations ORDER BY timestamp DESC LIMIT ?", - (min(limit, 100),), - ).fetchall() - finally: - conn.close() + return get_storage().list_sessions(limit) except Exception: return [] -# ── Session management ──────────────────────────────────────────────── - - -def register_session(session_id: str, title: str | None = None) -> None: - """Create a sessions row for a new session (no-op if already exists).""" +def delete_session(session_id: str) -> bool: + """Delete a session and all its messages.""" try: - conn = open_db() - try: - conn.execute( - "INSERT OR IGNORE INTO sessions " - "(session_id, title, created, updated) " - "VALUES (?, ?, datetime('now'), datetime('now'))", - (session_id, title), - ) - conn.commit() - finally: - conn.close() - except Exception: - pass - - -def update_session_title(session_id: str, title: str) -> None: - """Set or update the auto-generated title for a session.""" - try: - conn = open_db() - try: - conn.execute( - "UPDATE sessions SET title = ? WHERE session_id = ?", - (title, session_id), - ) - conn.commit() - finally: - conn.close() - except Exception: - pass - - -def set_session_alias(session_id: str, alias: str) -> bool: - """Set a human-friendly alias for a session. Returns False if alias is taken.""" - try: - conn = open_db() - try: - existing = conn.execute( - "SELECT session_id FROM sessions WHERE alias = ?", (alias,) - ).fetchone() - if existing and existing[0] != session_id: - return False - conn.execute( - "UPDATE sessions SET alias = ? WHERE session_id = ?", - (alias, session_id), - ) - conn.commit() - return True - finally: - conn.close() + return get_storage().delete_session(session_id) except Exception: return False -def get_session_name(session_id: str) -> str | None: - """Return the alias (or title if no alias) for a session, or None if unset.""" - try: - conn = open_db() - try: - row = conn.execute( - "SELECT alias, title FROM sessions WHERE session_id = ?", - (session_id,), - ).fetchone() - if row: - value = row[0] or row[1] - return str(value) if value is not None else None - finally: - conn.close() - except Exception: - pass - return None - - -def resolve_session(alias_or_id: str) -> str | None: - """Resolve an alias or session_id (or prefix) to a full session_id.""" - try: - conn = open_db() - try: - # 1. Exact alias match - row = conn.execute( - "SELECT session_id FROM sessions WHERE alias = ?", - (alias_or_id,), - ).fetchone() - if row: - return str(row[0]) - # 2. Exact session_id match - row = conn.execute( - "SELECT session_id FROM sessions WHERE session_id = ?", - (alias_or_id,), - ).fetchone() - if row: - return str(row[0]) - # 3. Session_id prefix match - rows = conn.execute( - "SELECT session_id FROM sessions WHERE session_id LIKE ?", - (alias_or_id + "%",), - ).fetchall() - if len(rows) == 1: - return str(rows[0][0]) - # 4. Fallback: check conversations table for legacy sessions - row = conn.execute( - "SELECT DISTINCT session_id FROM conversations WHERE session_id = ? LIMIT 1", - (alias_or_id,), - ).fetchone() - if row: - # Auto-register legacy session - conn.execute( - "INSERT OR IGNORE INTO sessions " - "(session_id, created, updated) VALUES (" - "?, " - "(SELECT MIN(timestamp) FROM conversations WHERE session_id = ?), " - "(SELECT MAX(timestamp) FROM conversations WHERE session_id = ?))", - (row[0], row[0], row[0]), - ) - conn.commit() - return str(row[0]) - return None - finally: - conn.close() - except Exception: - return None - - def prune_sessions( retention_days: int = 90, log_fn: Callable[[str], None] | None = None, ) -> tuple[int, int]: - """Prune orphaned and stale sessions. - - Removes: - - Sessions with no messages (orphaned registrations from process startup). - - Sessions whose ``updated`` timestamp is older than ``retention_days`` - **and** that have no alias (named sessions are kept indefinitely). - - Pass ``retention_days=0`` to skip age-based pruning (only orphans removed). - - Returns: - (orphans_removed, stale_removed) - """ - orphans = stale = 0 + """Prune orphaned and stale sessions.""" try: - conn = open_db() - try: - # 1. Remove sessions that have no messages at all. - orphan_ids = [ - row[0] - for row in conn.execute( - "SELECT session_id FROM sessions " - "WHERE NOT EXISTS " - " (SELECT 1 FROM conversations c " - " WHERE c.session_id = sessions.session_id)" - ).fetchall() - ] - if orphan_ids: - placeholders = ",".join("?" * len(orphan_ids)) - conn.execute( - f"DELETE FROM session_config WHERE session_id IN ({placeholders})", - orphan_ids, - ) - cur = conn.execute( - f"DELETE FROM sessions WHERE session_id IN ({placeholders})", - orphan_ids, - ) - orphans = cur.rowcount - - # 2. Remove old unnamed sessions. - if retention_days > 0: - cutoff = (datetime.utcnow() - timedelta(days=retention_days)).strftime( - "%Y-%m-%dT%H:%M:%S" - ) - stale_ids = [ - row[0] - for row in conn.execute( - "SELECT session_id FROM sessions WHERE alias IS NULL AND updated < ?", - (cutoff,), - ).fetchall() - ] - if stale_ids: - placeholders = ",".join("?" * len(stale_ids)) - conn.execute( - f"DELETE FROM session_config WHERE session_id IN ({placeholders})", - stale_ids, - ) - cur = conn.execute( - f"DELETE FROM sessions WHERE session_id IN ({placeholders})", - stale_ids, - ) - stale = cur.rowcount - - conn.commit() - finally: - conn.close() + orphans, stale = get_storage().prune_sessions(retention_days) except Exception: return (0, 0) @@ -441,206 +96,105 @@ def prune_sessions( return (orphans, stale) -def list_sessions(limit: int = 20) -> list[tuple[Any, ...]]: - """List recent sessions. - - Returns (session_id, alias, title, created, updated, msg_count) - ordered by updated DESC. - """ +def resolve_session(alias_or_id: str) -> str | None: + """Resolve an alias or session_id (or prefix) to a full session_id.""" try: - conn = open_db() - try: - return conn.execute( - "SELECT s.session_id, s.alias, s.title, s.created, s.updated, " - "(SELECT COUNT(*) FROM conversations c " - " WHERE c.session_id = s.session_id) " - "FROM sessions s " - "WHERE EXISTS " - " (SELECT 1 FROM conversations c WHERE c.session_id = s.session_id) " - "ORDER BY s.updated DESC LIMIT ?", - (limit,), - ).fetchall() - finally: - conn.close() + return get_storage().resolve_session(alias_or_id) except Exception: - return [] + return None -def load_session_messages(session_id: str) -> list[dict[str, Any]]: - """Load messages for a session and reconstruct OpenAI message format. - - Handles tool_call / tool_result rows by grouping consecutive tool_call - rows into one assistant message with tool_calls, then pairing subsequent - tool_result rows as tool messages. - """ - try: - conn = open_db() - try: - rows = conn.execute( - "SELECT role, content, tool_name, tool_args, tool_call_id, provider_data " - "FROM conversations WHERE session_id = ? ORDER BY id", - (session_id,), - ).fetchall() - finally: - conn.close() - except Exception: - return [] - - messages: list[dict[str, Any]] = [] - i = 0 - while i < len(rows): - role, content, tool_name, tool_args, tc_id, provider_data = rows[i] - - if role == "user": - messages.append({"role": "user", "content": content or ""}) - i += 1 - - elif role == "assistant": - msg: dict[str, Any] = {"role": "assistant", "content": content} - if provider_data: - with contextlib.suppress(json.JSONDecodeError, TypeError): - msg["_provider_content"] = json.loads(provider_data) - messages.append(msg) - i += 1 - - elif role == "tool_call": - # Collect consecutive tool_call rows into one assistant message. - # If the previous message was an assistant with content (text + - # tool calls in the same turn), merge tool_calls into it. - assistant_msg: dict[str, Any] = { - "role": "assistant", - "content": None, - "tool_calls": [], - } - if ( - messages - and messages[-1]["role"] == "assistant" - and not messages[-1].get("tool_calls") - ): - assistant_msg = messages.pop() - assistant_msg["tool_calls"] = [] - - while i < len(rows) and rows[i][0] == "tool_call": - _, _, tn, ta, stored_tc_id, _ = rows[i] - call_id = stored_tc_id or f"call_{session_id}_{i}" - assistant_msg["tool_calls"].append( - { - "id": call_id, - "type": "function", - "function": {"name": tn or "", "arguments": ta or ""}, - } - ) - i += 1 - messages.append(assistant_msg) - - # Consume matching tool_result rows - result_idx = 0 - while i < len(rows) and rows[i][0] == "tool_result": - _, result_content, _, _, result_tc_id, _ = rows[i] - if result_tc_id: - tc_id_to_use = result_tc_id - elif result_idx < len(assistant_msg["tool_calls"]): - tc_id_to_use = assistant_msg["tool_calls"][result_idx]["id"] - else: - tc_id_to_use = f"call_orphan_{i}" - messages.append( - { - "role": "tool", - "tool_call_id": tc_id_to_use, - "content": result_content or "", - } - ) - result_idx += 1 - i += 1 - - elif role == "tool_result": - # Orphaned tool_result (no preceding tool_call) — skip - i += 1 - else: - i += 1 - - # Repair: strip trailing incomplete tool call turns. - # If an assistant message has tool_calls but fewer tool results follow - # than expected, the session was interrupted mid-execution. Remove - # the incomplete turn so the LLM can re-generate cleanly. - while messages: - # Count trailing tool messages - tail_tools = 0 - for j in range(len(messages) - 1, -1, -1): - if messages[j].get("role") == "tool": - tail_tools += 1 - else: - break - # Check the assistant message that should precede them - asst_idx = len(messages) - 1 - tail_tools - if asst_idx < 0: - break - asst = messages[asst_idx] - if asst.get("role") != "assistant" or not asst.get("tool_calls"): - break - if tail_tools >= len(asst["tool_calls"]): - break # complete turn, nothing to repair - # Incomplete: remove partial tool messages + the assistant message - del messages[asst_idx:] - # Loop to check for nested incomplete turns - - return messages - - -def delete_session(session_id: str) -> bool: - """Delete a session and all its messages. Returns True on success.""" - try: - conn = open_db() - try: - conn.execute( - "DELETE FROM conversations WHERE session_id = ?", - (session_id,), - ) - conn.execute( - "DELETE FROM session_config WHERE session_id = ?", - (session_id,), - ) - conn.execute( - "DELETE FROM sessions WHERE session_id = ?", - (session_id,), - ) - conn.commit() - return True - finally: - conn.close() - except Exception: - return False +# -- Session config ------------------------------------------------------------ def save_session_config(session_id: str, config: dict[str, str]) -> None: """Persist session configuration key/value pairs.""" - try: - conn = open_db() - try: - for key, value in config.items(): - conn.execute( - "INSERT OR REPLACE INTO session_config " - "(session_id, key, value) VALUES (?, ?, ?)", - (session_id, key, value), - ) - conn.commit() - finally: - conn.close() - except Exception: - pass + with contextlib.suppress(Exception): + get_storage().save_session_config(session_id, config) def load_session_config(session_id: str) -> dict[str, str]: - """Load session configuration. Returns empty dict if none stored.""" + """Load session configuration.""" try: - conn = open_db() - try: - rows = conn.execute( - "SELECT key, value FROM session_config WHERE session_id = ?", - (session_id,), - ).fetchall() - return {row[0]: row[1] for row in rows} - finally: - conn.close() + return get_storage().load_session_config(session_id) except Exception: return {} + + +# -- Session metadata ---------------------------------------------------------- + + +def set_session_alias(session_id: str, alias: str) -> bool: + """Set a human-friendly alias. Returns False if alias is taken.""" + try: + return get_storage().set_session_alias(session_id, alias) + except Exception: + return False + + +def get_session_name(session_id: str) -> str | None: + """Return the alias (or title) for a session, or None if unset.""" + try: + return get_storage().get_session_name(session_id) + except Exception: + return None + + +def update_session_title(session_id: str, title: str) -> None: + """Set or update the auto-generated title for a session.""" + with contextlib.suppress(Exception): + get_storage().update_session_title(session_id, title) + + +# -- Key-value store (memories) ------------------------------------------------ + + +def save_memory(key: str, value: str) -> str | None: + """Save a memory. Returns the previous value if it existed.""" + try: + return get_storage().kv_set(key, value) + except Exception: + return None + + +def delete_memory(key: str) -> bool: + """Delete a memory by key. Returns True if the key existed.""" + try: + return get_storage().kv_delete(key) + except Exception: + return False + + +def load_memories() -> list[tuple[str, str]]: + """Return all (key, value) memory pairs sorted by key.""" + try: + return get_storage().kv_list() + except Exception: + return [] + + +def search_memories(query: str) -> list[tuple[str, str]]: + """Search memories by query. Returns matching (key, value) pairs.""" + try: + return get_storage().kv_search(query) + except Exception: + return [] + + +# -- Conversation search ------------------------------------------------------- + + +def search_history(query: str, limit: int = 20) -> list[Any]: + """Search conversation history.""" + try: + return get_storage().search_history(query, limit) + except Exception: + return [] + + +def search_history_recent(limit: int = 20) -> list[Any]: + """Return most recent conversation messages.""" + try: + return get_storage().search_history_recent(limit) + except Exception: + return [] diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 5926b3d9..fdc012d4 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -24,24 +24,25 @@ from typing import TYPE_CHECKING, Any, Protocol import httpx +from turnstone.core.config import get_tavily_key from turnstone.core.edit import find_occurrences, pick_nearest from turnstone.core.memory import ( + delete_memory, delete_session, - escape_like, get_session_name, - get_tavily_key, list_sessions, load_memories, load_session_config, load_session_messages, normalize_key, - open_db, register_session, resolve_session, + save_memory, save_message, save_session_config, search_history, search_history_recent, + search_memories, set_session_alias, update_session_title, ) @@ -2330,29 +2331,14 @@ class ChatSession: """Save a persistent memory.""" call_id, key, value = item["call_id"], item["key"], item["value"] try: - conn = open_db() - try: - existing = conn.execute( - "SELECT value FROM memories WHERE key = ?", (key,) - ).fetchone() - conn.execute( - "INSERT OR REPLACE INTO memories (key, value, created, updated) " - "VALUES (?, ?, COALESCE(" - " (SELECT created FROM memories WHERE key = ?), " - " datetime('now')" - "), datetime('now'))", - (key, value, key), - ) - conn.commit() - self._init_system_messages() - if existing: - msg = f"Updated memory: {key} = {value} (was: {existing[0]})" - else: - msg = f"Saved memory: {key} = {value}" - self.ui.on_tool_result(call_id, "remember", msg) - return call_id, msg - finally: - conn.close() + old_value = save_memory(key, value) + self._init_system_messages() + if old_value is not None: + msg = f"Updated memory: {key} = {value} (was: {old_value})" + else: + msg = f"Saved memory: {key} = {value}" + self.ui.on_tool_result(call_id, "remember", msg) + return call_id, msg except Exception as e: return call_id, f"Error: {e}" @@ -2360,19 +2346,14 @@ class ChatSession: """Remove a persistent memory by key.""" call_id, key = item["call_id"], item["key"] try: - conn = open_db() - try: - cursor = conn.execute("DELETE FROM memories WHERE key = ?", (key,)) - conn.commit() - if cursor.rowcount == 0: - msg = f"Error: memory '{key}' not found" - else: - self._init_system_messages() - msg = f"Forgot: {key}" - self.ui.on_tool_result(call_id, "forget", msg) - return call_id, msg - finally: - conn.close() + deleted = delete_memory(key) + if not deleted: + msg = f"Error: memory '{key}' not found" + else: + self._init_system_messages() + msg = f"Forgot: {key}" + self.ui.on_tool_result(call_id, "forget", msg) + return call_id, msg except Exception as e: return call_id, f"Error: {e}" @@ -2384,30 +2365,11 @@ class ChatSession: # Memories: list all (no query) or search (with query) try: - conn = open_db() - try: - if not query: - rows = conn.execute("SELECT key, value FROM memories ORDER BY key").fetchall() - else: - terms = query.split() - clauses = [] - params: list[str] = [] - for t in terms: - escaped = escape_like(t) - clauses.append("(key LIKE ? ESCAPE '\\' OR value LIKE ? ESCAPE '\\')") - params.extend([f"%{escaped}%", f"%{escaped}%"]) - rows = conn.execute( - "SELECT key, value FROM memories WHERE " - + " AND ".join(clauses) - + " ORDER BY key", - params, - ).fetchall() - if rows: - parts.append("Memories:\n" + "\n".join(f" {k}={v}" for k, v in rows)) - elif not query: - parts.append("No memories stored.") - finally: - conn.close() + rows = search_memories(query) if query else load_memories() + if rows: + parts.append("Memories:\n" + "\n".join(f" {k}={v}" for k, v in rows)) + elif not query: + parts.append("No memories stored.") except Exception: pass diff --git a/turnstone/core/storage/__init__.py b/turnstone/core/storage/__init__.py new file mode 100644 index 00000000..728b81b6 --- /dev/null +++ b/turnstone/core/storage/__init__.py @@ -0,0 +1,14 @@ +"""Pluggable storage backend for turnstone persistence. + +Supports SQLite (default, zero-config) and PostgreSQL (multi-node, production). +""" + +from turnstone.core.storage._protocol import StorageBackend +from turnstone.core.storage._registry import get_storage, init_storage, reset_storage + +__all__ = [ + "StorageBackend", + "get_storage", + "init_storage", + "reset_storage", +] diff --git a/turnstone/core/storage/_migrate.py b/turnstone/core/storage/_migrate.py new file mode 100644 index 00000000..7eab8972 --- /dev/null +++ b/turnstone/core/storage/_migrate.py @@ -0,0 +1,84 @@ +"""Programmatic Alembic migration runner.""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any + +log = logging.getLogger(__name__) + +_MIGRATIONS_DIR = str(Path(__file__).parent / "migrations") + + +def run_migrations(storage: Any, backend: str) -> None: + """Run pending Alembic migrations. + + For SQLite backends, also handles bootstrapping existing databases + that were created before the migration system existed. + """ + from alembic import command + from alembic.config import Config + + engine = storage._engine # noqa: SLF001 + + # Build Alembic config programmatically (no alembic.ini needed) + cfg = Config() + cfg.set_main_option("script_location", _MIGRATIONS_DIR) + cfg.set_main_option("sqlalchemy.url", str(engine.url)) + + # Check if this is an existing database without alembic_version + if backend == "sqlite": + _bootstrap_existing_sqlite(engine, cfg) + + try: + command.upgrade(cfg, "head") + except Exception as exc: + if backend == "sqlite": + log.warning("Migration failed (non-fatal for SQLite): %s", exc) + else: + raise + + +def _bootstrap_existing_sqlite(engine: Any, cfg: Any) -> None: + """Stamp existing SQLite databases at the baseline revision. + + If the database has tables but no alembic_version, it was created + before the migration system. Stamp it so Alembic knows the schema + is already at the baseline. + """ + import sqlalchemy as sa + from alembic import command + + with engine.connect() as conn: + # Check if alembic_version table exists + has_alembic = conn.execute( + sa.text("SELECT 1 FROM sqlite_master WHERE type='table' AND name='alembic_version'") + ).fetchone() + if has_alembic: + return # Already managed by Alembic + + # Check if sessions table exists (indicates pre-existing database) + has_sessions = conn.execute( + sa.text("SELECT 1 FROM sqlite_master WHERE type='table' AND name='sessions'") + ).fetchone() + if has_sessions: + log.info("Bootstrapping existing database into Alembic (stamping at baseline)") + command.stamp(cfg, "001") + + +if __name__ == "__main__": + # Allow running as: python -m turnstone.core.storage._migrate + # Used by Docker entrypoint to apply migrations before starting services. + import os + + from turnstone.core.storage import get_storage, init_storage + + backend = os.environ.get("TURNSTONE_DB_BACKEND", "sqlite") + url = os.environ.get("TURNSTONE_DB_URL", "") + path = os.environ.get("TURNSTONE_DB_PATH", "") + + logging.basicConfig(level=logging.INFO) + init_storage(backend, path=path, url=url, run_migrations=True) + log.info("Migrations complete") + get_storage().close() diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py new file mode 100644 index 00000000..627ab98d --- /dev/null +++ b/turnstone/core/storage/_postgresql.py @@ -0,0 +1,386 @@ +"""PostgreSQL storage backend.""" + +from __future__ import annotations + +import logging +from datetime import UTC, datetime, timedelta +from typing import Any + +import sqlalchemy as sa + +from turnstone.core.storage._schema import ( + conversations, + memories, + metadata, + session_config, + sessions, +) +from turnstone.core.storage._sqlite import _reconstruct_messages + +log = logging.getLogger(__name__) + + +class PostgreSQLBackend: + """PostgreSQL implementation of the StorageBackend protocol.""" + + def __init__( + self, url: str, pool_size: int = 5, max_overflow: int = 10, *, create_tables: bool = True + ) -> None: + self._engine = sa.create_engine( + url, + pool_size=pool_size, + max_overflow=max_overflow, + pool_pre_ping=True, + ) + if create_tables: + metadata.create_all(self._engine) + + # -- Core session operations ----------------------------------------------- + + def register_session(self, session_id: str, title: str | None = None) -> None: + now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") + with self._engine.connect() as conn: + # Use dialect-neutral upsert pattern + existing = conn.execute( + sa.select(sessions.c.session_id).where(sessions.c.session_id == session_id) + ).fetchone() + if not existing: + conn.execute( + sa.insert(sessions), + {"session_id": session_id, "title": title, "created": now, "updated": now}, + ) + conn.commit() + + def save_message( + self, + session_id: str, + role: str, + content: str | None, + tool_name: str | None = None, + tool_args: str | None = None, + tool_call_id: str | None = None, + provider_data: str | None = None, + ) -> None: + now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") + with self._engine.connect() as conn: + conn.execute( + sa.insert(conversations), + { + "session_id": session_id, + "timestamp": now, + "role": role, + "content": content, + "tool_name": tool_name, + "tool_args": tool_args, + "tool_call_id": tool_call_id, + "provider_data": provider_data, + }, + ) + conn.execute( + sa.update(sessions).where(sessions.c.session_id == session_id).values(updated=now) + ) + conn.commit() + + def load_session_messages(self, session_id: str) -> list[dict[str, Any]]: + with self._engine.connect() as conn: + rows = conn.execute( + sa.select( + conversations.c.role, + conversations.c.content, + conversations.c.tool_name, + conversations.c.tool_args, + conversations.c.tool_call_id, + conversations.c.provider_data, + ) + .where(conversations.c.session_id == session_id) + .order_by(conversations.c.id) + ).fetchall() + return _reconstruct_messages(list(rows), session_id) + + # -- Session management ---------------------------------------------------- + + def list_sessions(self, limit: int = 20) -> list[Any]: + with self._engine.connect() as conn: + return list( + conn.execute( + sa.text( + "SELECT s.session_id, s.alias, s.title, s.created, s.updated, " + "(SELECT COUNT(*) FROM conversations c " + " WHERE c.session_id = s.session_id) " + "FROM sessions s " + "WHERE EXISTS " + " (SELECT 1 FROM conversations c WHERE c.session_id = s.session_id) " + "ORDER BY s.updated DESC LIMIT :limit" + ), + {"limit": limit}, + ).fetchall() + ) + + def delete_session(self, session_id: str) -> bool: + with self._engine.connect() as conn: + conn.execute(sa.delete(conversations).where(conversations.c.session_id == session_id)) + conn.execute(sa.delete(session_config).where(session_config.c.session_id == session_id)) + conn.execute(sa.delete(sessions).where(sessions.c.session_id == session_id)) + conn.commit() + return True + + def prune_sessions(self, retention_days: int = 90) -> tuple[int, int]: + orphans = stale = 0 + with self._engine.connect() as conn: + # 1. Remove sessions with no messages + orphan_rows = conn.execute( + sa.text( + "SELECT session_id FROM sessions " + "WHERE NOT EXISTS " + " (SELECT 1 FROM conversations c " + " WHERE c.session_id = sessions.session_id)" + ) + ).fetchall() + orphan_ids = [r[0] for r in orphan_rows] + if orphan_ids: + conn.execute( + sa.delete(session_config).where(session_config.c.session_id.in_(orphan_ids)) + ) + result = conn.execute( + sa.delete(sessions).where(sessions.c.session_id.in_(orphan_ids)) + ) + orphans = result.rowcount + + # 2. Remove old unnamed sessions + if retention_days > 0: + cutoff = (datetime.now(UTC) - timedelta(days=retention_days)).strftime( + "%Y-%m-%dT%H:%M:%S" + ) + stale_rows = conn.execute( + sa.select(sessions.c.session_id).where( + sessions.c.alias.is_(None), + sessions.c.updated < cutoff, + ) + ).fetchall() + stale_ids = [r[0] for r in stale_rows] + if stale_ids: + conn.execute( + sa.delete(session_config).where(session_config.c.session_id.in_(stale_ids)) + ) + result = conn.execute( + sa.delete(sessions).where(sessions.c.session_id.in_(stale_ids)) + ) + stale = result.rowcount + + conn.commit() + return (orphans, stale) + + def resolve_session(self, alias_or_id: str) -> str | None: + with self._engine.connect() as conn: + # 1. Exact alias + row = conn.execute( + sa.select(sessions.c.session_id).where(sessions.c.alias == alias_or_id) + ).fetchone() + if row: + return str(row[0]) + # 2. Exact session_id + row = conn.execute( + sa.select(sessions.c.session_id).where(sessions.c.session_id == alias_or_id) + ).fetchone() + if row: + return str(row[0]) + # 3. Prefix match + rows = conn.execute( + sa.select(sessions.c.session_id).where( + sessions.c.session_id.like(alias_or_id + "%") + ) + ).fetchall() + if len(rows) == 1: + return str(rows[0][0]) + # 4. Legacy: check conversations + row = conn.execute( + sa.select(sa.distinct(conversations.c.session_id)) + .where(conversations.c.session_id == alias_or_id) + .limit(1) + ).fetchone() + if row: + now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") + existing = conn.execute( + sa.select(sessions.c.session_id).where(sessions.c.session_id == row[0]) + ).fetchone() + if not existing: + conn.execute( + sa.insert(sessions), + {"session_id": row[0], "created": now, "updated": now}, + ) + conn.commit() + return str(row[0]) + return None + + # -- Session config -------------------------------------------------------- + + def save_session_config(self, session_id: str, config: dict[str, str]) -> None: + with self._engine.connect() as conn: + for key, value in config.items(): + # Upsert: delete + insert + conn.execute( + sa.delete(session_config).where( + session_config.c.session_id == session_id, + session_config.c.key == key, + ) + ) + conn.execute( + sa.insert(session_config), + {"session_id": session_id, "key": key, "value": value}, + ) + conn.commit() + + def load_session_config(self, session_id: str) -> dict[str, str]: + with self._engine.connect() as conn: + rows = conn.execute( + sa.select(session_config.c.key, session_config.c.value).where( + session_config.c.session_id == session_id + ) + ).fetchall() + return {row[0]: row[1] for row in rows} + + # -- Session metadata ------------------------------------------------------ + + def set_session_alias(self, session_id: str, alias: str) -> bool: + with self._engine.connect() as conn: + existing = conn.execute( + sa.select(sessions.c.session_id).where(sessions.c.alias == alias) + ).fetchone() + if existing and existing[0] != session_id: + return False + conn.execute( + sa.update(sessions).where(sessions.c.session_id == session_id).values(alias=alias) + ) + conn.commit() + return True + + def get_session_name(self, session_id: str) -> str | None: + with self._engine.connect() as conn: + row = conn.execute( + sa.select(sessions.c.alias, sessions.c.title).where( + sessions.c.session_id == session_id + ) + ).fetchone() + if row: + value = row[0] or row[1] + return str(value) if value is not None else None + return None + + def update_session_title(self, session_id: str, title: str) -> None: + with self._engine.connect() as conn: + conn.execute( + sa.update(sessions).where(sessions.c.session_id == session_id).values(title=title) + ) + conn.commit() + + # -- Generic key-value store ----------------------------------------------- + + def kv_get(self, key: str) -> str | None: + with self._engine.connect() as conn: + row = conn.execute(sa.select(memories.c.value).where(memories.c.key == key)).fetchone() + return str(row[0]) if row else None + + def kv_set(self, key: str, value: str) -> str | None: + now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") + with self._engine.connect() as conn: + existing = conn.execute( + sa.select(memories.c.value, memories.c.created).where(memories.c.key == key) + ).fetchone() + old_value = str(existing[0]) if existing else None + created = str(existing[1]) if existing else now + # Delete + insert for cross-dialect upsert + conn.execute(sa.delete(memories).where(memories.c.key == key)) + conn.execute( + sa.insert(memories), + {"key": key, "value": value, "created": created, "updated": now}, + ) + conn.commit() + return old_value + + def kv_delete(self, key: str) -> bool: + with self._engine.connect() as conn: + result = conn.execute(sa.delete(memories).where(memories.c.key == key)) + conn.commit() + return result.rowcount > 0 + + def kv_list(self) -> list[tuple[str, str]]: + with self._engine.connect() as conn: + rows = conn.execute( + sa.select(memories.c.key, memories.c.value).order_by(memories.c.key) + ).fetchall() + return [(str(r[0]), str(r[1])) for r in rows] + + def kv_search(self, query: str) -> list[tuple[str, str]]: + if not query or not query.strip(): + return self.kv_list() + terms = query.split() + with self._engine.connect() as conn: + clauses = [] + params: dict[str, str] = {} + for i, t in enumerate(terms): + clauses.append(f"(key ILIKE :k{i} OR value ILIKE :v{i})") + params[f"k{i}"] = f"%{t}%" + params[f"v{i}"] = f"%{t}%" + rows = conn.execute( + sa.text( + "SELECT key, value FROM memories WHERE " + + " AND ".join(clauses) + + " ORDER BY key" + ), + params, + ).fetchall() + return [(str(r[0]), str(r[1])) for r in rows] + + # -- Conversation search --------------------------------------------------- + + def search_history(self, query: str, limit: int = 20) -> list[Any]: + if not query or not query.strip(): + return [] + capped = min(limit, 100) + with self._engine.connect() as conn: + # Use PostgreSQL full-text search if search_vector column exists + try: + return list( + conn.execute( + sa.text( + "SELECT c.timestamp, c.session_id, c.role, c.content, c.tool_name " + "FROM conversations c " + "WHERE to_tsvector('english', COALESCE(c.content, '')) " + " @@ plainto_tsquery('english', :query) " + "ORDER BY ts_rank(to_tsvector('english', COALESCE(c.content, '')), " + " plainto_tsquery('english', :query)) DESC " + "LIMIT :limit" + ), + {"query": query, "limit": capped}, + ).fetchall() + ) + except Exception: + # Fallback to ILIKE + return list( + conn.execute( + sa.text( + "SELECT timestamp, session_id, role, content, tool_name " + "FROM conversations WHERE content ILIKE :pattern " + "ORDER BY timestamp DESC LIMIT :limit" + ), + {"pattern": f"%{query}%", "limit": capped}, + ).fetchall() + ) + + def search_history_recent(self, limit: int = 20) -> list[Any]: + capped = min(limit, 100) + with self._engine.connect() as conn: + return list( + conn.execute( + sa.text( + "SELECT timestamp, session_id, role, content, tool_name " + "FROM conversations ORDER BY timestamp DESC LIMIT :limit" + ), + {"limit": capped}, + ).fetchall() + ) + + # -- Lifecycle ------------------------------------------------------------- + + def close(self) -> None: + self._engine.dispose() diff --git a/turnstone/core/storage/_protocol.py b/turnstone/core/storage/_protocol.py new file mode 100644 index 00000000..a97a24f2 --- /dev/null +++ b/turnstone/core/storage/_protocol.py @@ -0,0 +1,117 @@ +"""Storage backend protocol — the contract every persistence adapter must implement.""" + +from __future__ import annotations + +from typing import Any, Protocol, runtime_checkable + + +@runtime_checkable +class StorageBackend(Protocol): + """Protocol that every storage backend adapter must implement. + + Provides session management, conversation persistence, key-value storage + (for memories), and full-text search. + """ + + # -- Core session operations ----------------------------------------------- + + def register_session(self, session_id: str, title: str | None = None) -> None: + """Create a sessions row for a new session (no-op if already exists).""" + ... + + def save_message( + self, + session_id: str, + role: str, + content: str | None, + tool_name: str | None = None, + tool_args: str | None = None, + tool_call_id: str | None = None, + provider_data: str | None = None, + ) -> None: + """Log a message to the conversations table.""" + ... + + def load_session_messages(self, session_id: str) -> list[dict[str, Any]]: + """Load messages for a session and reconstruct OpenAI message format.""" + ... + + # -- Session management ---------------------------------------------------- + + def list_sessions(self, limit: int = 20) -> list[Any]: + """List recent sessions with message counts, ordered by updated DESC.""" + ... + + def delete_session(self, session_id: str) -> bool: + """Delete a session and all its messages. Returns True on success.""" + ... + + def prune_sessions(self, retention_days: int = 90) -> tuple[int, int]: + """Remove orphaned + stale unnamed sessions. Returns (orphans, stale).""" + ... + + def resolve_session(self, alias_or_id: str) -> str | None: + """Resolve an alias or session_id (or prefix) to a full session_id.""" + ... + + # -- Session config -------------------------------------------------------- + + def save_session_config(self, session_id: str, config: dict[str, str]) -> None: + """Persist session configuration key/value pairs.""" + ... + + def load_session_config(self, session_id: str) -> dict[str, str]: + """Load session configuration. Returns empty dict if none stored.""" + ... + + # -- Session metadata ------------------------------------------------------ + + def set_session_alias(self, session_id: str, alias: str) -> bool: + """Set a human-friendly alias. Returns False if alias is taken.""" + ... + + def get_session_name(self, session_id: str) -> str | None: + """Return the alias (or title) for a session, or None if unset.""" + ... + + def update_session_title(self, session_id: str, title: str) -> None: + """Set or update the auto-generated title for a session.""" + ... + + # -- Generic key-value store (backs memories table) ------------------------ + + def kv_get(self, key: str) -> str | None: + """Get a value by key. Returns None if not found.""" + ... + + def kv_set(self, key: str, value: str) -> str | None: + """Set a key-value pair. Returns the previous value if it existed.""" + ... + + def kv_delete(self, key: str) -> bool: + """Delete a key. Returns True if the key existed.""" + ... + + def kv_list(self) -> list[tuple[str, str]]: + """Return all (key, value) pairs sorted by key.""" + ... + + def kv_search(self, query: str) -> list[tuple[str, str]]: + """Search key-value pairs by query. Returns matching (key, value) pairs.""" + ... + + # -- Conversation search --------------------------------------------------- + + def search_history(self, query: str, limit: int = 20) -> list[Any]: + """Search conversation history. Returns (timestamp, session_id, role, content, tool_name).""" + ... + + def search_history_recent(self, limit: int = 20) -> list[Any]: + """Return most recent conversation messages.""" + ... + + # -- Lifecycle ------------------------------------------------------------- + + def close(self) -> None: + """Release resources (connection pool, engine, etc.).""" + ... diff --git a/turnstone/core/storage/_registry.py b/turnstone/core/storage/_registry.py new file mode 100644 index 00000000..e9699cbc --- /dev/null +++ b/turnstone/core/storage/_registry.py @@ -0,0 +1,86 @@ +"""Storage backend singleton registry.""" + +from __future__ import annotations + +import logging +import os +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from turnstone.core.storage._protocol import StorageBackend + +log = logging.getLogger(__name__) + +_storage: StorageBackend | None = None + + +def init_storage( + backend: str = "sqlite", + *, + path: str = "", + url: str = "", + pool_size: int = 5, + run_migrations: bool = True, +) -> StorageBackend: + """Initialize the storage backend singleton. + + Args: + backend: "sqlite" or "postgresql" + path: SQLite database file path (default: .turnstone.db in cwd) + url: PostgreSQL connection URL (e.g. postgresql+psycopg://user:pass@host/db) + pool_size: Connection pool size (PostgreSQL only) + run_migrations: Whether to run Alembic migrations on init + """ + global _storage + + # When Alembic migrations will run, skip create_all() to avoid + # bypassing migration-managed DDL. Tests pass run_migrations=False + # and rely on create_all() instead. + create_tables = not run_migrations + + if backend == "sqlite": + from turnstone.core.storage._sqlite import SQLiteBackend + + db_path = path or os.path.join(os.getcwd(), ".turnstone.db") + _storage = SQLiteBackend(db_path, create_tables=create_tables) + log.info("Storage initialized: SQLite at %s", db_path) + + elif backend == "postgresql": + from turnstone.core.storage._postgresql import PostgreSQLBackend + + if not url: + msg = "PostgreSQL backend requires a connection URL (db_url)" + raise ValueError(msg) + _storage = PostgreSQLBackend(url, pool_size=pool_size, create_tables=create_tables) + log.info("Storage initialized: PostgreSQL") + + else: + msg = f"Unknown storage backend: {backend!r} (expected 'sqlite' or 'postgresql')" + raise ValueError(msg) + + if run_migrations: + from turnstone.core.storage._migrate import run_migrations as _run_migrations + + _run_migrations(_storage, backend) + + return _storage + + +def get_storage() -> StorageBackend: + """Return the initialized storage backend. + + Auto-initializes with SQLite defaults if not yet initialized. + """ + global _storage + if _storage is None: + init_storage("sqlite") + assert _storage is not None + return _storage + + +def reset_storage() -> None: + """Close and clear the storage backend singleton (for tests).""" + global _storage + if _storage is not None: + _storage.close() + _storage = None diff --git a/turnstone/core/storage/_schema.py b/turnstone/core/storage/_schema.py new file mode 100644 index 00000000..f1b39b0b --- /dev/null +++ b/turnstone/core/storage/_schema.py @@ -0,0 +1,56 @@ +"""SQLAlchemy Core schema — single source of truth for all table definitions. + +Used by both storage backends and Alembic migrations. +""" + +from __future__ import annotations + +import sqlalchemy as sa + +metadata = sa.MetaData() + +memories = sa.Table( + "memories", + metadata, + sa.Column("key", sa.Text, primary_key=True), + sa.Column("value", sa.Text, nullable=False), + sa.Column("created", sa.Text, nullable=False), + sa.Column("updated", sa.Text, nullable=False), +) + +conversations = sa.Table( + "conversations", + metadata, + sa.Column("id", sa.Integer, primary_key=True, autoincrement=True), + sa.Column("session_id", sa.Text, nullable=False, index=True), + sa.Column("timestamp", sa.Text, nullable=False), + sa.Column("role", sa.Text, nullable=False), + sa.Column("content", sa.Text), + sa.Column("tool_name", sa.Text), + sa.Column("tool_args", sa.Text), + sa.Column("tool_call_id", sa.Text), + sa.Column("provider_data", sa.Text), +) + +sessions = sa.Table( + "sessions", + metadata, + sa.Column("session_id", sa.Text, primary_key=True), + sa.Column("alias", sa.Text, unique=True), + sa.Column("title", sa.Text), + sa.Column("created", sa.Text, nullable=False), + sa.Column("updated", sa.Text, nullable=False), +) + +# Additional indexes on sessions (name-based to avoid duplication with SA's auto-index) +sa.Index("idx_sessions_alias", sessions.c.alias) +sa.Index("idx_sessions_updated", sessions.c.updated) + +session_config = sa.Table( + "session_config", + metadata, + sa.Column("session_id", sa.Text, nullable=False), + sa.Column("key", sa.Text, nullable=False), + sa.Column("value", sa.Text), + sa.PrimaryKeyConstraint("session_id", "key"), +) diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py new file mode 100644 index 00000000..42a172fa --- /dev/null +++ b/turnstone/core/storage/_sqlite.py @@ -0,0 +1,546 @@ +"""SQLite storage backend.""" + +from __future__ import annotations + +import contextlib +import json +import logging +from datetime import UTC, datetime, timedelta +from typing import Any + +import sqlalchemy as sa + +from turnstone.core.storage._schema import ( + conversations, + memories, + metadata, + session_config, + sessions, +) + +log = logging.getLogger(__name__) + + +def _escape_like(s: str) -> str: + """Escape LIKE metacharacters for use with ESCAPE '\\\\'.""" + return s.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + + +def _fts5_query(query: str) -> str: + """Convert a plain search string into a safe FTS5 query.""" + terms = query.split() + safe = [] + for t in terms: + if t: + safe.append(f'"{t.replace(chr(34), chr(34) + chr(34))}"') + return " ".join(safe) + + +class SQLiteBackend: + """SQLite implementation of the StorageBackend protocol.""" + + def __init__(self, path: str, *, create_tables: bool = True) -> None: + self._path = path + self._engine = sa.create_engine( + f"sqlite:///{path}", + pool_pre_ping=True, + connect_args={"check_same_thread": False}, + ) + self._fts5_available = False + if create_tables: + self._init_schema() + + def _init_schema(self) -> None: + """Create tables and FTS5 index.""" + metadata.create_all(self._engine) + # Try to set up FTS5 for full-text search + with self._engine.connect() as conn: + try: + fts_exists = conn.execute( + sa.text( + "SELECT 1 FROM sqlite_master " + "WHERE type='table' AND name='conversations_fts'" + ) + ).fetchone() + if not fts_exists: + conn.execute( + sa.text( + "CREATE VIRTUAL TABLE conversations_fts " + "USING fts5(content, content=conversations, content_rowid=id)" + ) + ) + conn.execute( + sa.text( + "INSERT INTO conversations_fts(conversations_fts) VALUES('rebuild')" + ) + ) + conn.commit() + self._fts5_available = True + except Exception: + self._fts5_available = False + + # -- Core session operations ----------------------------------------------- + + def register_session(self, session_id: str, title: str | None = None) -> None: + now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") + with self._engine.connect() as conn: + conn.execute( + sa.insert(sessions).prefix_with("OR IGNORE"), + {"session_id": session_id, "title": title, "created": now, "updated": now}, + ) + conn.commit() + + def save_message( + self, + session_id: str, + role: str, + content: str | None, + tool_name: str | None = None, + tool_args: str | None = None, + tool_call_id: str | None = None, + provider_data: str | None = None, + ) -> None: + now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") + with self._engine.connect() as conn: + result = conn.execute( + sa.insert(conversations), + { + "session_id": session_id, + "timestamp": now, + "role": role, + "content": content, + "tool_name": tool_name, + "tool_args": tool_args, + "tool_call_id": tool_call_id, + "provider_data": provider_data, + }, + ) + # FTS5 indexing + if self._fts5_available and content: + try: + rowid = result.lastrowid + conn.execute( + sa.text( + "INSERT INTO conversations_fts(rowid, content) VALUES (:rowid, :content)" + ), + {"rowid": rowid, "content": content}, + ) + except Exception: + self._fts5_available = False + # Bump session updated timestamp + conn.execute( + sa.update(sessions).where(sessions.c.session_id == session_id).values(updated=now) + ) + conn.commit() + + def load_session_messages(self, session_id: str) -> list[dict[str, Any]]: + with self._engine.connect() as conn: + rows = conn.execute( + sa.select( + conversations.c.role, + conversations.c.content, + conversations.c.tool_name, + conversations.c.tool_args, + conversations.c.tool_call_id, + conversations.c.provider_data, + ) + .where(conversations.c.session_id == session_id) + .order_by(conversations.c.id) + ).fetchall() + + return _reconstruct_messages(list(rows), session_id) + + # -- Session management ---------------------------------------------------- + + def list_sessions(self, limit: int = 20) -> list[Any]: + with self._engine.connect() as conn: + return list( + conn.execute( + sa.text( + "SELECT s.session_id, s.alias, s.title, s.created, s.updated, " + "(SELECT COUNT(*) FROM conversations c " + " WHERE c.session_id = s.session_id) " + "FROM sessions s " + "WHERE EXISTS " + " (SELECT 1 FROM conversations c WHERE c.session_id = s.session_id) " + "ORDER BY s.updated DESC LIMIT :limit" + ), + {"limit": limit}, + ).fetchall() + ) + + def delete_session(self, session_id: str) -> bool: + with self._engine.connect() as conn: + conn.execute(sa.delete(conversations).where(conversations.c.session_id == session_id)) + conn.execute(sa.delete(session_config).where(session_config.c.session_id == session_id)) + conn.execute(sa.delete(sessions).where(sessions.c.session_id == session_id)) + conn.commit() + return True + + def prune_sessions(self, retention_days: int = 90) -> tuple[int, int]: + orphans = stale = 0 + with self._engine.connect() as conn: + # 1. Remove sessions with no messages + orphan_ids = [ + row[0] + for row in conn.execute( + sa.text( + "SELECT session_id FROM sessions " + "WHERE NOT EXISTS " + " (SELECT 1 FROM conversations c " + " WHERE c.session_id = sessions.session_id)" + ) + ).fetchall() + ] + if orphan_ids: + placeholders = ",".join([":p" + str(i) for i in range(len(orphan_ids))]) + params = {f"p{i}": oid for i, oid in enumerate(orphan_ids)} + conn.execute( + sa.text(f"DELETE FROM session_config WHERE session_id IN ({placeholders})"), + params, + ) + result = conn.execute( + sa.text(f"DELETE FROM sessions WHERE session_id IN ({placeholders})"), + params, + ) + orphans = result.rowcount + + # 2. Remove old unnamed sessions + if retention_days > 0: + cutoff = (datetime.now(UTC) - timedelta(days=retention_days)).strftime( + "%Y-%m-%dT%H:%M:%S" + ) + stale_ids = [ + row[0] + for row in conn.execute( + sa.text( + "SELECT session_id FROM sessions " + "WHERE alias IS NULL AND updated < :cutoff" + ), + {"cutoff": cutoff}, + ).fetchall() + ] + if stale_ids: + placeholders = ",".join([":p" + str(i) for i in range(len(stale_ids))]) + params = {f"p{i}": sid for i, sid in enumerate(stale_ids)} + conn.execute( + sa.text(f"DELETE FROM session_config WHERE session_id IN ({placeholders})"), + params, + ) + result = conn.execute( + sa.text(f"DELETE FROM sessions WHERE session_id IN ({placeholders})"), + params, + ) + stale = result.rowcount + + conn.commit() + return (orphans, stale) + + def resolve_session(self, alias_or_id: str) -> str | None: + with self._engine.connect() as conn: + # 1. Exact alias match + row = conn.execute( + sa.select(sessions.c.session_id).where(sessions.c.alias == alias_or_id) + ).fetchone() + if row: + return str(row[0]) + # 2. Exact session_id match + row = conn.execute( + sa.select(sessions.c.session_id).where(sessions.c.session_id == alias_or_id) + ).fetchone() + if row: + return str(row[0]) + # 3. Session_id prefix match + rows = conn.execute( + sa.select(sessions.c.session_id).where( + sessions.c.session_id.like(alias_or_id + "%") + ) + ).fetchall() + if len(rows) == 1: + return str(rows[0][0]) + # 4. Legacy: check conversations table + row = conn.execute( + sa.text( + "SELECT DISTINCT session_id FROM conversations WHERE session_id = :sid LIMIT 1" + ), + {"sid": alias_or_id}, + ).fetchone() + if row: + # Auto-register legacy session + conn.execute( + sa.text( + "INSERT OR IGNORE INTO sessions " + "(session_id, created, updated) VALUES (" + ":sid, " + "(SELECT MIN(timestamp) FROM conversations WHERE session_id = :sid), " + "(SELECT MAX(timestamp) FROM conversations WHERE session_id = :sid))" + ), + {"sid": row[0]}, + ) + conn.commit() + return str(row[0]) + return None + + # -- Session config -------------------------------------------------------- + + def save_session_config(self, session_id: str, config: dict[str, str]) -> None: + with self._engine.connect() as conn: + for key, value in config.items(): + conn.execute( + sa.text( + "INSERT OR REPLACE INTO session_config " + "(session_id, key, value) VALUES (:sid, :key, :value)" + ), + {"sid": session_id, "key": key, "value": value}, + ) + conn.commit() + + def load_session_config(self, session_id: str) -> dict[str, str]: + with self._engine.connect() as conn: + rows = conn.execute( + sa.select(session_config.c.key, session_config.c.value).where( + session_config.c.session_id == session_id + ) + ).fetchall() + return {row[0]: row[1] for row in rows} + + # -- Session metadata ------------------------------------------------------ + + def set_session_alias(self, session_id: str, alias: str) -> bool: + with self._engine.connect() as conn: + existing = conn.execute( + sa.select(sessions.c.session_id).where(sessions.c.alias == alias) + ).fetchone() + if existing and existing[0] != session_id: + return False + conn.execute( + sa.update(sessions).where(sessions.c.session_id == session_id).values(alias=alias) + ) + conn.commit() + return True + + def get_session_name(self, session_id: str) -> str | None: + with self._engine.connect() as conn: + row = conn.execute( + sa.select(sessions.c.alias, sessions.c.title).where( + sessions.c.session_id == session_id + ) + ).fetchone() + if row: + value = row[0] or row[1] + return str(value) if value is not None else None + return None + + def update_session_title(self, session_id: str, title: str) -> None: + with self._engine.connect() as conn: + conn.execute( + sa.update(sessions).where(sessions.c.session_id == session_id).values(title=title) + ) + conn.commit() + + # -- Generic key-value store ----------------------------------------------- + + def kv_get(self, key: str) -> str | None: + with self._engine.connect() as conn: + row = conn.execute(sa.select(memories.c.value).where(memories.c.key == key)).fetchone() + return str(row[0]) if row else None + + def kv_set(self, key: str, value: str) -> str | None: + with self._engine.connect() as conn: + existing = conn.execute( + sa.select(memories.c.value).where(memories.c.key == key) + ).fetchone() + old_value = str(existing[0]) if existing else None + now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") + conn.execute( + sa.text( + "INSERT OR REPLACE INTO memories (key, value, created, updated) " + "VALUES (:key, :value, " + "COALESCE((SELECT created FROM memories WHERE key = :key), :now), " + ":now)" + ), + {"key": key, "value": value, "now": now}, + ) + conn.commit() + return old_value + + def kv_delete(self, key: str) -> bool: + with self._engine.connect() as conn: + result = conn.execute(sa.delete(memories).where(memories.c.key == key)) + conn.commit() + return result.rowcount > 0 + + def kv_list(self) -> list[tuple[str, str]]: + with self._engine.connect() as conn: + rows = conn.execute( + sa.select(memories.c.key, memories.c.value).order_by(memories.c.key) + ).fetchall() + return [(str(r[0]), str(r[1])) for r in rows] + + def kv_search(self, query: str) -> list[tuple[str, str]]: + if not query or not query.strip(): + return self.kv_list() + terms = query.split() + with self._engine.connect() as conn: + # Build WHERE clause: each term must match key OR value + clauses = [] + params: dict[str, str] = {} + for i, t in enumerate(terms): + escaped = _escape_like(t) + clauses.append(f"(key LIKE :k{i} ESCAPE '\\' OR value LIKE :v{i} ESCAPE '\\')") + params[f"k{i}"] = f"%{escaped}%" + params[f"v{i}"] = f"%{escaped}%" + rows = conn.execute( + sa.text( + "SELECT key, value FROM memories WHERE " + + " AND ".join(clauses) + + " ORDER BY key" + ), + params, + ).fetchall() + return [(str(r[0]), str(r[1])) for r in rows] + + # -- Conversation search --------------------------------------------------- + + def search_history(self, query: str, limit: int = 20) -> list[Any]: + if not query or not query.strip(): + return [] + capped = min(limit, 100) + with self._engine.connect() as conn: + if self._fts5_available: + return list( + conn.execute( + sa.text( + "SELECT c.timestamp, c.session_id, c.role, c.content, c.tool_name " + "FROM conversations_fts f " + "JOIN conversations c ON c.id = f.rowid " + "WHERE conversations_fts MATCH :query " + "ORDER BY f.rank ASC LIMIT :limit" + ), + {"query": _fts5_query(query), "limit": capped}, + ).fetchall() + ) + return list( + conn.execute( + sa.text( + "SELECT timestamp, session_id, role, content, tool_name " + "FROM conversations WHERE content LIKE :pattern ESCAPE '\\' " + "ORDER BY timestamp DESC LIMIT :limit" + ), + {"pattern": f"%{_escape_like(query)}%", "limit": capped}, + ).fetchall() + ) + + def search_history_recent(self, limit: int = 20) -> list[Any]: + capped = min(limit, 100) + with self._engine.connect() as conn: + return list( + conn.execute( + sa.text( + "SELECT timestamp, session_id, role, content, tool_name " + "FROM conversations ORDER BY timestamp DESC LIMIT :limit" + ), + {"limit": capped}, + ).fetchall() + ) + + # -- Lifecycle ------------------------------------------------------------- + + def close(self) -> None: + self._engine.dispose() + + +def _reconstruct_messages(rows: list[Any], session_id: str) -> list[dict[str, Any]]: + """Reconstruct OpenAI message format from stored conversation rows. + + Handles tool_call / tool_result grouping and incomplete turn repair. + """ + messages: list[dict[str, Any]] = [] + i = 0 + while i < len(rows): + role, content, tool_name, tool_args, tc_id, provider_data = rows[i] + + if role == "user": + messages.append({"role": "user", "content": content or ""}) + i += 1 + + elif role == "assistant": + msg: dict[str, Any] = {"role": "assistant", "content": content} + if provider_data: + with contextlib.suppress(json.JSONDecodeError, TypeError): + msg["_provider_content"] = json.loads(provider_data) + messages.append(msg) + i += 1 + + elif role == "tool_call": + assistant_msg: dict[str, Any] = { + "role": "assistant", + "content": None, + "tool_calls": [], + } + if ( + messages + and messages[-1]["role"] == "assistant" + and not messages[-1].get("tool_calls") + ): + assistant_msg = messages.pop() + assistant_msg["tool_calls"] = [] + + while i < len(rows) and rows[i][0] == "tool_call": + _, _, tn, ta, stored_tc_id, _ = rows[i] + call_id = stored_tc_id or f"call_{session_id}_{i}" + assistant_msg["tool_calls"].append( + { + "id": call_id, + "type": "function", + "function": {"name": tn or "", "arguments": ta or ""}, + } + ) + i += 1 + messages.append(assistant_msg) + + # Consume matching tool_result rows + result_idx = 0 + while i < len(rows) and rows[i][0] == "tool_result": + _, result_content, _, _, result_tc_id, _ = rows[i] + if result_tc_id: + tc_id_to_use = result_tc_id + elif result_idx < len(assistant_msg["tool_calls"]): + tc_id_to_use = assistant_msg["tool_calls"][result_idx]["id"] + else: + tc_id_to_use = f"call_orphan_{i}" + messages.append( + { + "role": "tool", + "tool_call_id": tc_id_to_use, + "content": result_content or "", + } + ) + result_idx += 1 + i += 1 + + elif role == "tool_result": + # Orphaned tool_result (no preceding tool_call) — skip + i += 1 + else: + i += 1 + + # Repair: strip trailing incomplete tool call turns + while messages: + tail_tools = 0 + for j in range(len(messages) - 1, -1, -1): + if messages[j].get("role") == "tool": + tail_tools += 1 + else: + break + asst_idx = len(messages) - 1 - tail_tools + if asst_idx < 0: + break + asst = messages[asst_idx] + if asst.get("role") != "assistant" or not asst.get("tool_calls"): + break + if tail_tools >= len(asst["tool_calls"]): + break + del messages[asst_idx:] + + return messages diff --git a/turnstone/core/storage/migrations/env.py b/turnstone/core/storage/migrations/env.py new file mode 100644 index 00000000..e316f8a8 --- /dev/null +++ b/turnstone/core/storage/migrations/env.py @@ -0,0 +1,27 @@ +"""Alembic environment for turnstone storage migrations.""" + +from alembic import context + +from turnstone.core.storage._schema import metadata + +target_metadata = metadata + + +def run_migrations_online() -> None: + """Run migrations using the engine from alembic config.""" + from sqlalchemy import engine_from_config, pool + + config = context.config + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure(connection=connection, target_metadata=target_metadata) + with context.begin_transaction(): + context.run_migrations() + + +run_migrations_online() diff --git a/turnstone/core/storage/migrations/script.py.mako b/turnstone/core/storage/migrations/script.py.mako new file mode 100644 index 00000000..02ec213b --- /dev/null +++ b/turnstone/core/storage/migrations/script.py.mako @@ -0,0 +1,23 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} +""" +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers +revision = ${repr(up_revision)} +down_revision = ${repr(down_revision)} +branch_labels = ${repr(branch_labels)} +depends_on = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/turnstone/core/storage/migrations/versions/001_initial_schema.py b/turnstone/core/storage/migrations/versions/001_initial_schema.py new file mode 100644 index 00000000..62552379 --- /dev/null +++ b/turnstone/core/storage/migrations/versions/001_initial_schema.py @@ -0,0 +1,68 @@ +"""Initial schema — baseline for all tables. + +Revision ID: 001 +Revises: +Create Date: 2026-03-03 +""" + +import sqlalchemy as sa +from alembic import op + +revision = "001" +down_revision = None +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # memories — persistent key-value store + op.create_table( + "memories", + sa.Column("key", sa.Text, primary_key=True), + sa.Column("value", sa.Text, nullable=False), + sa.Column("created", sa.Text, nullable=False), + sa.Column("updated", sa.Text, nullable=False), + ) + + # conversations — message history + op.create_table( + "conversations", + sa.Column("id", sa.Integer, primary_key=True, autoincrement=True), + sa.Column("session_id", sa.Text, nullable=False), + sa.Column("timestamp", sa.Text, nullable=False), + sa.Column("role", sa.Text, nullable=False), + sa.Column("content", sa.Text), + sa.Column("tool_name", sa.Text), + sa.Column("tool_args", sa.Text), + sa.Column("tool_call_id", sa.Text), + sa.Column("provider_data", sa.Text), + ) + op.create_index("idx_conv_session", "conversations", ["session_id"]) + + # sessions — session metadata + op.create_table( + "sessions", + sa.Column("session_id", sa.Text, primary_key=True), + sa.Column("alias", sa.Text, unique=True), + sa.Column("title", sa.Text), + sa.Column("created", sa.Text, nullable=False), + sa.Column("updated", sa.Text, nullable=False), + ) + op.create_index("idx_sessions_alias", "sessions", ["alias"]) + op.create_index("idx_sessions_updated", "sessions", ["updated"]) + + # session_config — per-session LLM parameters + op.create_table( + "session_config", + sa.Column("session_id", sa.Text, nullable=False), + sa.Column("key", sa.Text, nullable=False), + sa.Column("value", sa.Text), + sa.PrimaryKeyConstraint("session_id", "key"), + ) + + +def downgrade() -> None: + op.drop_table("session_config") + op.drop_table("sessions") + op.drop_table("conversations") + op.drop_table("memories") diff --git a/turnstone/eval.py b/turnstone/eval.py index 6da97fb8..b59cbd7f 100644 --- a/turnstone/eval.py +++ b/turnstone/eval.py @@ -29,9 +29,9 @@ from typing import Any from openai import OpenAI -import turnstone.core.memory as _memory_module from turnstone.core.providers import LLMProvider, create_provider from turnstone.core.session import ChatSession +from turnstone.core.storage import init_storage, reset_storage from turnstone.core.tools import PRIMARY_KEY_MAP, TOOLS # ─── ANSI & logging helpers ─────────────────────────────────────────────────── @@ -334,7 +334,8 @@ def _run_single_test( workdir = tempfile.mkdtemp(prefix="turnstone_eval_") original_cwd = os.getcwd() eval_db = os.path.join(workdir, ".turnstone_eval.db") - _memory_module.db_override = eval_db + reset_storage() + init_storage("sqlite", path=eval_db, run_migrations=False) t0 = time.monotonic() try: @@ -399,8 +400,7 @@ def _run_single_test( "elapsed": round(elapsed, 1), } finally: - _memory_module.db_override = None - _memory_module.db_initialized.discard(eval_db) + reset_storage() os.chdir(original_cwd) shutil.rmtree(workdir, ignore_errors=True) diff --git a/turnstone/server.py b/turnstone/server.py index 760d63ee..43944495 100644 --- a/turnstone/server.py +++ b/turnstone/server.py @@ -1261,10 +1261,23 @@ def main() -> None: apply_config( parser, - ["api", "model", "session", "tools", "server", "mcp", "ratelimit", "health"], + ["api", "model", "session", "tools", "server", "mcp", "ratelimit", "health", "database"], ) args = parser.parse_args() + # Initialize storage backend + from turnstone.core.storage import init_storage + + db_backend = getattr(args, "db_backend", None) or os.environ.get( + "TURNSTONE_DB_BACKEND", "sqlite" + ) + db_url = getattr(args, "db_url", None) or os.environ.get("TURNSTONE_DB_URL", "") + db_path = getattr(args, "db_path", None) or os.environ.get("TURNSTONE_DB_PATH", "") + db_pool_size = int( + getattr(args, "db_pool_size", None) or os.environ.get("TURNSTONE_DB_POOL_SIZE", "5") + ) + init_storage(db_backend, path=db_path, url=db_url, pool_size=db_pool_size) + # Prune stale / empty sessions on startup from turnstone.core.memory import prune_sessions