Files
crank/__REFACTORING.md
T

31 KiB
Raw Blame History

Refactoring Plan

Purpose

This document defines the refactoring track required to improve maintainability, safety, and production-readiness in the current codebase. It is intentionally implementation-oriented: each area is split into concrete slices, target modules, function ownership, migration order, and acceptance criteria.

The refactoring scope covers:

  1. splitting crates/crank-registry/src/postgres.rs into domain modules;
  2. migrating static SQL to sqlx::query! / sqlx::query_as! where possible;
  3. replacing SHA-256 key derivation in SecretCrypto with HKDF;
  4. replacing string timestamps in domain models with typed timestamps;
  5. adding Display support to define_id!;
  6. introducing explicit PostgreSQL pool configuration;
  7. removing duplicated SQL timestamp formatting;
  8. making runtime and API error payloads more structured;
  9. introducing correlation IDs across transport, runtime, and adapters;
  10. adding rate limiting and backpressure controls;
  11. replacing in-memory MCP session state with a multi-instance-safe store.

Guiding Rules

  • Keep domain logic separate from SQL storage concerns.
  • Preserve current behavior while shrinking file size and coupling.
  • Move in vertical slices: storage module split first, then SQL typing, then timestamp migration.
  • Every slice must be independently reviewable and mergeable.
  • Prefer compatibility shims during migration instead of breaking multiple crates at once.
  • Add or update tests in the same slice as the refactoring.

Refactoring Tracks

Track A: Registry Modularization

Problem

crates/crank-registry/src/postgres.rs currently concentrates almost all database logic:

  • bootstrap and auth;
  • workspaces and memberships;
  • invitations;
  • operations and versions;
  • agents and bindings;
  • samples and descriptors;
  • API keys;
  • secrets;
  • stream sessions and async jobs;
  • invocation logs and usage;
  • YAML import jobs.

This creates:

  • a permanent git conflict hotspot;
  • slow code navigation;
  • oversized review diffs;
  • accidental coupling between unrelated domains.

Target Layout

crates/crank-registry/src/

error.rs
lib.rs
migrations.rs
model.rs
postgres/
  mod.rs
  connect.rs
  auth.rs
  workspace.rs
  invitation.rs
  secret.rs
  stream.rs
  api_key.rs
  operation.rs
  agent.rs
  observability.rs
  yaml_import.rs
  shared.rs

Ownership by Module

postgres/mod.rs

Responsibilities:

  • define PostgresRegistry;
  • wire submodules;
  • expose the public impl surface;
  • keep only high-level constructor entry points.

Functions staying here:

  • new
  • connect
  • connect_with_options
  • migrate

Functions moving out:

  • everything domain-specific.
postgres/connect.rs

Responsibilities:

  • connection setup;
  • pool initialization;
  • migration invocation;
  • pool option translation from config to PgPoolOptions.

Functions:

  • connect
  • connect_with_options
  • migrate
  • internal helpers:
    • build_pool_options
    • build_connect_options
postgres/auth.rs

Responsibilities:

  • auth users;
  • user sessions;
  • membership access checks;
  • auth profiles.

Functions:

  • upsert_bootstrap_user
  • get_auth_user_by_email
  • get_auth_user_by_id
  • update_user_profile
  • update_user_password
  • create_user_session
  • get_user_session
  • touch_user_session
  • revoke_user_session
  • set_user_session_current_workspace
  • user_has_workspace_access
  • save_auth_profile
  • get_auth_profile
  • list_auth_profiles
  • list_auth_profiles_referencing_secret

Internal helpers to extract:

  • map_auth_user_row
  • map_session_row
  • map_auth_profile_row
postgres/workspace.rs

Responsibilities:

  • workspaces;
  • memberships;
  • workspace CRUD.

Functions:

  • list_workspaces
  • list_workspaces_for_user
  • ensure_membership
  • list_memberships
  • update_membership_role
  • delete_membership
  • get_workspace
  • create_workspace
  • update_workspace
  • delete_workspace

Internal helpers:

  • workspace_exists
  • map_workspace_row
  • map_membership_row
postgres/invitation.rs

Responsibilities:

  • invitations lifecycle.

Functions:

  • list_invitations
  • create_invitation
  • delete_invitation
postgres/secret.rs

Responsibilities:

  • secrets;
  • secret versions;
  • secret touch/rotation.

Functions:

  • list_secrets
  • get_secret
  • get_current_secret_version
  • create_secret
  • rotate_secret
  • delete_secret
  • touch_secret
postgres/stream.rs

Responsibilities:

  • stream sessions;
  • async jobs;
  • state transitions;
  • cleanup.

Functions:

  • create_stream_session
  • get_stream_session
  • update_stream_session_state
  • close_stream_session
  • list_stream_sessions
  • delete_expired_stream_sessions
  • create_async_job
  • get_async_job
  • update_async_job_status
  • cancel_async_job
  • list_async_jobs
  • delete_expired_async_jobs

Internal helpers:

  • assert_stream_session_transition
  • assert_async_job_transition
  • map_stream_session_row
  • map_async_job_row
postgres/api_key.rs

Responsibilities:

  • platform API key lifecycle.

Functions:

  • list_platform_api_keys
  • get_platform_api_key_by_secret_for_workspace_slug
  • create_platform_api_key
  • revoke_platform_api_key
  • delete_platform_api_key
  • touch_platform_api_key
postgres/operation.rs

Responsibilities:

  • operations;
  • versions;
  • published operations;
  • samples;
  • descriptors.

Functions:

  • create_operation
  • create_version
  • update_operation_draft
  • archive_operation
  • delete_operation
  • has_published_agent_bindings_for_operation
  • list_operations
  • list_operation_usage_summaries
  • list_operation_agent_refs
  • get_operation_summary
  • get_operation_version
  • list_operation_versions
  • publish_operation
  • get_published_operation
  • list_published_operations
  • save_sample_metadata
  • list_sample_metadata
  • save_descriptor_metadata
  • list_descriptor_metadata
  • insert_version_row

Internal helpers:

  • map_operation_summary_row
  • map_operation_version_row
  • map_published_operation_row
  • serialize_optional_json
postgres/agent.rs

Responsibilities:

  • agents;
  • agent versions;
  • bindings;
  • publish lifecycle.

Functions:

  • list_agents
  • get_agent_summary
  • create_agent
  • get_agent_version
  • save_agent_bindings
  • update_agent_summary
  • delete_agent
  • publish_agent
  • unpublish_agent
  • archive_agent
  • get_published_agent_tools_by_slug
  • list_agent_bindings
  • insert_agent_version_row
  • replace_agent_bindings_rows
postgres/observability.rs

Responsibilities:

  • invocation logs;
  • usage rollups and breakdowns.

Functions:

  • create_invocation_log
  • list_invocation_logs
  • get_invocation_log
  • summarize_usage
  • list_usage_timeline
  • list_usage_by_operation
  • get_usage_for_operation
  • list_usage_by_agent
  • get_usage_for_agent
postgres/yaml_import.rs

Responsibilities:

  • YAML import jobs.

Functions:

  • create_yaml_import_job
  • finish_yaml_import_job
  • get_yaml_import_job
postgres/shared.rs

Responsibilities:

  • SQL snippets reused by several modules;
  • common timestamp select fragments during the migration period;
  • small row-mapping helpers;
  • pagination helpers.

Planned content:

  • Page<T> helpers if they remain storage-local;
  • constants for table-independent fragments;
  • transitional helpers until typed timestamps remove string formatting entirely.

Slice Plan for Track A

Slice A1: File layout and constructors
  • create postgres/ directory;
  • move connect, connect_with_options, migrate, and struct wiring;
  • keep all public signatures unchanged.
Slice A2: Workspace/Auth/API-key split
  • move auth.rs, workspace.rs, invitation.rs, api_key.rs;
  • leave behavior unchanged;
  • add targeted module tests if not present.
Slice A3: Operation/Agent split
  • move operation.rs, agent.rs, yaml_import.rs;
  • keep SQL unchanged for this slice;
  • ensure published-operation and published-agent behavior is unchanged.
Slice A4: Secret/Stream/Observability split
  • move secret.rs, stream.rs, observability.rs;
  • keep transition logic and cleanup behavior unchanged.
Slice A5: Remove legacy monolith
  • remove remaining methods from root postgres.rs;
  • rename module structure so postgres/mod.rs is the only entry point.

Done Criteria for Track A

  • no postgres.rs monolith remains;
  • no single registry source file exceeds roughly 8001200 lines;
  • all existing registry tests pass unchanged;
  • lib.rs still re-exports PostgresRegistry without downstream breakage.

Track B: Compile-Time SQL Verification

Problem

Most static SQL currently uses sqlx::query(...) / query_as::<_, T>(...) with runtime validation only. That leaves:

  • column name typos;
  • mismatched types;
  • bad nullability assumptions;
  • schema drift

to runtime instead of compile time.

Goal

Use sqlx::query! and sqlx::query_as! for every static query with stable shape.

Keep runtime builders only where the query shape is truly dynamic.

Query Classification

Must migrate to query! / query_as!

Use macros for:

  • single-row fetches by ID;
  • list queries with static columns;
  • insert/update/delete with static statement shape;
  • existence checks;
  • count queries;
  • cleanup queries with fixed predicates.

Expected modules:

  • auth.rs
  • workspace.rs
  • invitation.rs
  • secret.rs
  • stream.rs
  • api_key.rs
  • most of agent.rs
  • most of yaml_import.rs

May stay runtime-based

Keep dynamic APIs only when necessary:

  • optional filters that materially change SQL shape;
  • optional sort columns chosen from a whitelist;
  • bulk dynamic IN lists where macro usage becomes unreadable;
  • QueryBuilder-based multi-row inserts/updates.

Expected modules:

  • operation.rs list queries with optional filters;
  • observability.rs usage aggregations and breakdowns;
  • some search/filter endpoints.

Required Infrastructure

Add SQLx offline support

Repository-level requirements:

  • use sqlx prepare to generate checked metadata;
  • keep checked metadata in version control if the team wants deterministic CI;
  • wire verification into CI or just verify.

Concrete additions:

  • sqlx-data.json or the current SQLx metadata equivalent for the chosen version;
  • DATABASE_URL or test DB wiring for cargo sqlx prepare.

Introduce row structs where needed

For query_as!, add storage-local row structs for cases where domain structs should not be coupled directly to SQL column shape.

Examples:

  • WorkspaceRow
  • AuthUserRow
  • OperationSummaryRow
  • InvocationLogRow
  • UsageTimelineRow

Rule:

  • domain structs remain domain-oriented;
  • row structs remain SQL-oriented.

Function-Level Migration Rules

For query_as!

Use query_as! when:

  • the query result maps cleanly to a typed row struct;
  • all returned columns are explicit;
  • nullability is known.

For query!

Use query! when:

  • no reusable row struct is needed;
  • the query is a simple insert, update, delete, or exists fetch.

For dynamic list queries

Use:

  • sqlx::QueryBuilder<Postgres> with explicit allowed sort/filter sets;
  • storage-local mapper functions from row to domain object.

Migration Order

Slice B1: auth, workspace, invitation

Low-risk, fixed-shape queries.

Slice B2: secret, stream, api_key, yaml_import

Still mostly fixed-shape.

Slice B3: operation, agent

Introduce mixed strategy: macros for fixed queries, builder for dynamic lists.

Slice B4: observability

Likely the most dynamic area. Convert only static pieces first, leave complex aggregations to QueryBuilder.

Done Criteria for Track B

  • all fixed-shape SQL uses compile-time checked macros;
  • dynamic SQL is explicitly documented and isolated;
  • CI fails on SQL drift before runtime tests;
  • no new static query is accepted through sqlx::query(...) without justification.

Track C: SecretCrypto KDF Correction

Problem

Current code in crates/crank-runtime/src/secret_crypto.rs derives the AES key by hashing the master key with SHA-256:

  • no KDF semantics;
  • no labeled derivation context;
  • no future rotation/versioning path.

Goal

Replace direct hashing with HKDF-SHA256.

Target Design

SecretCrypto::new(master_key: &str) will:

  1. validate non-empty master key;
  2. derive a 32-byte AES key using Hkdf<Sha256>;
  3. use explicit info context, for example:
    • b"crank.secret-envelope.v1"
  4. initialize Aes256Gcm from the derived bytes.

Module Changes

File:

  • crates/crank-runtime/src/secret_crypto.rs

Changes:

  • replace Sha256::digest(trimmed.as_bytes()) with HKDF expansion;
  • preserve the public constructor signature;
  • keep ciphertext format unchanged if possible;
  • if ciphertext format must change later, introduce envelope version before doing so.

Required Tests

  • encrypt_decrypt_roundtrip_still_works
  • empty_master_key_is_rejected
  • same_master_key_derives_stable_key
  • different_master_keys_produce_different_ciphertexts

Done Criteria for Track C

  • no direct SHA-256 KDF remains;
  • SecretCrypto uses HKDF-SHA256 with explicit context;
  • all existing secret encrypt/decrypt flows remain compatible.

Track D: Typed Timestamps

Problem

Domain and storage structs use String timestamps. SQL converts timestamps via repeated to_char(...).

That causes:

  • repeated formatting logic in SQL;
  • weak type guarantees;
  • inability to do time arithmetic without reparsing;
  • broad surface for formatting drift.

Goal

Use time::OffsetDateTime in domain and storage models, with RFC 3339 serialization at the transport boundary.

Target Type Strategy

Domain layer

Use:

#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime

and the same for updated_at, expires_at, used_at, and similar fields.

Storage row layer

Also use OffsetDateTime.

Do not format timestamps in SQL except for transitional compatibility shims.

Transport layer

JSON still emits RFC 3339 strings through Serde.

Files Affected

Domain crates

  • crates/crank-core/src/workspace.rs
  • crates/crank-core/src/access.rs
  • crates/crank-core/src/auth.rs
  • crates/crank-core/src/agent.rs
  • crates/crank-core/src/operation.rs
  • crates/crank-core/src/observability.rs
  • crates/crank-core/src/secret.rs
  • crates/crank-core/src/stream_session.rs

Registry models

  • crates/crank-registry/src/model.rs

Admin API transport/service DTOs

  • apps/admin-api/src/service.rs
  • any route DTO files that mirror time fields directly

Migration Strategy

Slice D1: storage row types
  • change storage row structs to OffsetDateTime;
  • remove to_char(...) from fixed-shape SQL;
  • keep service DTOs as strings temporarily if needed.
Slice D2: domain models
  • convert crank-core models to typed timestamps with Serde helpers;
  • update constructors and tests.
Slice D3: API/service transport
  • remove string timestamp shims from service DTOs where possible;
  • ensure JSON remains stable.
Slice D4: cleanup
  • remove transitional formatting helpers;
  • remove now_string() helpers in favor of typed timestamps.

Acceptance Criteria

  • no created_at: String or updated_at: String remains in domain/storage models;
  • no to_char(created_at...) / to_char(updated_at...) remains in registry SQL except perhaps in legacy compatibility code slated for removal in the same slice;
  • JSON contract stays RFC 3339.

Track E: Display for ID Newtypes

Problem

define_id! currently requires .as_str() everywhere for formatting and messages.

Goal

Add std::fmt::Display support to every generated ID newtype.

File

  • crates/crank-core/src/ids.rs

Change

Add inside define_id!:

impl std::fmt::Display for $name {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

Follow-Up Cleanup

After macro support lands:

  • replace formatting sites like workspace_id.as_str() inside messages;
  • keep .as_str() only where an actual &str API requires it.

Acceptance Criteria

  • all ID newtypes implement Display;
  • formatting code no longer uses .as_str() inside format!, println!, or error messages unless &str is actually required.

Track F: PostgreSQL Pool Configuration

Problem

The pool is currently created via default PgPoolOptions, which means the runtime behavior is implicit and under-configured.

Goal

Make pool behavior explicit and configurable from environment.

Target Config Model

Add a typed config struct, for example:

pub struct PostgresPoolConfig {
    pub max_connections: u32,
    pub min_connections: u32,
    pub acquire_timeout_ms: u64,
    pub idle_timeout_ms: u64,
    pub max_lifetime_ms: u64,
}

Runtime Environment

Add env variables:

  • POSTGRES_MAX_CONNECTIONS
  • POSTGRES_MIN_CONNECTIONS
  • POSTGRES_ACQUIRE_TIMEOUT_MS
  • POSTGRES_IDLE_TIMEOUT_MS
  • POSTGRES_MAX_LIFETIME_MS

Files

  • crates/crank-registry/src/postgres/connect.rs
  • apps/admin-api/src/main.rs
  • apps/mcp-server/src/main.rs
  • config docs and examples

Behavior

connect_with_options should accept either:

  • direct PgConnectOptions plus explicit pool config;
  • or an internal default config built from env.

Defaults should be conservative and explicit, for example:

  • max_connections = 20
  • min_connections = 2
  • acquire_timeout = 5s
  • idle_timeout = 10m
  • max_lifetime = 30m

Exact values can be tuned later, but they must no longer be hidden defaults.

Acceptance Criteria

  • pool limits and timeouts are explicit;
  • docs reflect the new variables;
  • startup logs or config inspection make active pool settings visible.

Track G: Timestamp Formatting Deduplication

Problem

The same to_char(... at time zone 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS"Z"') fragment is copied across many queries.

Goal

Eliminate duplicated formatting logic.

Strategy

This track is mostly resolved by Track D, but there is a transition window.

During transition:

  • centralize any remaining formatting in postgres/shared.rs;
  • define named SQL fragments only once;
  • remove inline duplicated snippets.

After Track D:

  • remove formatting fragments entirely.

Acceptance Criteria

  • duplicated inline timestamp formatting is no longer scattered across query strings;
  • final state has typed timestamps instead of SQL formatting.

Track H: Error Structure Normalization

Problem

Some runtime and API error variants carry only a details: String, which makes programmatic handling and consistent observability harder.

Goal

Introduce structured error payloads where the current model is too string-heavy.

Scope

Priority types:

  • RuntimeError
  • registry storage errors where context is stringly typed;
  • API error mapping around secret/runtime/streaming faults.

Design Rules

  • keep thiserror;
  • use structured fields in variants whenever stable context exists;
  • reserve free-form details only for opaque third-party messages.

Examples:

Instead of:

RuntimeError::SecretCrypto { details: String }

prefer:

RuntimeError::SecretCrypto {
    operation: &'static str,
    details: String,
}

or dedicated variants such as:

  • InvalidSecretEnvelope
  • SecretDecryptFailed
  • SecretSerializeFailed

Slice Plan

Slice H1
  • normalize SecretCrypto and streaming/session/job errors.
Slice H2
  • normalize adapter errors for REST/gRPC/WebSocket/SOAP where stable categories exist.
Slice H3
  • normalize API mapping and error codes.

Acceptance Criteria

  • core runtime error variants expose structured context;
  • API can map key failure classes deterministically;
  • logs and telemetry can group errors by category instead of by ad hoc message text.

Track I: Correlation IDs and Cross-Layer Tracing

Problem

Today, a request can enter through:

  • admin-api HTTP routes;
  • mcp-server JSON-RPC over Streamable HTTP;
  • internal runtime execution;
  • protocol adapters (REST, gRPC, WebSocket, SOAP)

without a single stable correlation key that survives the entire path.

There are already partial pieces in the codebase:

  • request_id exists in observability models and storage;
  • MCP mappings can inject $.mcp.correlation_id;
  • some protocol fixtures mention correlation_id.

But there is no guaranteed end-to-end propagation model. When an adapter fails, logs cannot always be tied back to:

  • the incoming HTTP request;
  • the MCP session;
  • the runtime operation execution;
  • the invocation log row.

Goal

Introduce a first-class correlation model that is generated once at ingress and propagated across:

  • transport;
  • service/runtime orchestration;
  • adapter calls;
  • persistence and logs.

Target Model

Introduce two related identifiers:

  • request_id
    • unique per incoming HTTP or MCP request;
    • used for ingress/egress tracing and log correlation.
  • correlation_id
    • stable business-level identifier propagated into upstream payloads when configured;
    • may equal request_id by default, but remains distinct conceptually.

Flow Rules

admin-api

At HTTP ingress:

  • read X-Request-Id if supplied and valid;
  • otherwise generate a UUID v7 request ID;
  • store it in request extensions;
  • include it in structured logs;
  • pass it into runtime execution context;
  • persist it into invocation_logs.request_id when a test run or runtime execution creates a log row.

mcp-server

At JSON-RPC ingress:

  • derive a request-scoped request_id;
  • bind it to the current transport session and JSON-RPC request;
  • propagate it into runtime execution and adapter context;
  • return it in logs and diagnostics.

Runtime

Add an execution context struct, for example:

pub struct ExecutionContext {
    pub request_id: String,
    pub correlation_id: String,
    pub workspace_id: WorkspaceId,
    pub agent_id: Option<AgentId>,
    pub operation_id: OperationId,
    pub mcp_session_id: Option<String>,
}

This context must be passed through:

  • unary execution;
  • window execution;
  • session tool start/poll/stop;
  • async job start/status/result/cancel.

Adapters

Each adapter must receive execution context and apply it consistently:

  • REST
    • add X-Request-Id
    • add X-Correlation-Id
    • allow explicit mapping into headers/query/body where configured
  • gRPC
    • attach metadata keys such as x-request-id, x-correlation-id
  • WebSocket
    • inject correlation metadata into subscribe payloads when configured
  • SOAP
    • inject correlation value into configured header/body mappings

Files and Functions

New or expanded files

  • apps/admin-api/src/request_context.rs
  • apps/mcp-server/src/request_context.rs
  • crates/crank-runtime/src/context.rs
  • adapter request builders in each protocol crate

Functions to add or refactor

admin-api

  • extract_or_generate_request_id(...)
  • request_context_middleware(...)
  • runtime_request_context(...)

mcp-server

  • extract_mcp_request_context(...)
  • runtime_request_context_from_message(...)

crank-runtime

  • ExecutionContext
  • PreparedRequest / RuntimeOperation updates to carry context
  • execute_* functions accept &ExecutionContext

Acceptance Criteria

  • every runtime invocation has a request ID;
  • every adapter call can log the request ID;
  • invocation logs persist the same request ID;
  • operator can trace a failing adapter call back to the original HTTP or MCP request.

Track J: Rate Limiting and Backpressure

Problem

There is currently no explicit protection layer for:

  • excessive API request rates;
  • excessive tool invocation rates;
  • excessive stream/session polling;
  • unbounded adapter output pressure.

This is a production-readiness gap, especially now that the platform supports:

  • unary execution;
  • windowed streaming;
  • session tools;
  • async jobs;
  • WebSocket and SSE upstream adapters.

Goal

Introduce explicit controls for request rate, concurrency, and bounded output pressure.

Scope

API-level controls

Apply to:

  • admin-api routes;
  • mcp-server transport endpoints.

Dimensions:

  • per-IP
  • per-session
  • per-workspace
  • per-platform-key

Runtime-level controls

Apply to:

  • concurrent unary executions;
  • concurrent window executions;
  • concurrent live stream sessions;
  • concurrent async jobs;
  • polling frequency for session/job tools.

Adapter-level backpressure

Apply to:

  • REST SSE
  • gRPC server-streaming
  • WebSocket windows/sessions

Controls:

  • max buffered items
  • max buffered bytes
  • max window duration
  • max session polls
  • max per-message bytes

Target Design

admin-api

Add middleware with a keyed limiter, for example:

  • by authenticated user ID;
  • by workspace ID;
  • by remote IP for unauthenticated endpoints.

mcp-server

Add a limiter keyed by:

  • platform key ID;
  • workspace slug;
  • MCP session ID;
  • optionally agent slug.

crank-runtime

Add a runtime guard layer, for example:

pub struct RuntimeLimits {
    pub max_concurrent_unary: usize,
    pub max_concurrent_window: usize,
    pub max_concurrent_sessions: usize,
    pub max_concurrent_jobs: usize,
    pub max_polls_per_session: u32,
}

And an execution gate:

  • acquire permit before execution;
  • release permit on completion;
  • reject with stable error code if limits are exceeded.

Implementation Slices

Slice J1: API request throttling
  • add middleware and config for admin-api;
  • add middleware and config for mcp-server.
Slice J2: Runtime concurrency limits
  • add semaphores/permits around execution paths.
Slice J3: Streaming/session pressure limits
  • cap poll frequency and active session count;
  • cap pending buffer sizes in adapters.

Acceptance Criteria

  • platform can reject overload deterministically instead of degrading unpredictably;
  • stream/session polling is bounded;
  • slow clients cannot force unbounded buffering in streaming adapters.

Track K: Distributed MCP Session Store

Problem

apps/mcp-server/src/session.rs currently keeps transport/session state in process memory through Arc<RwLock<HashMap<...>>>.

This is acceptable only for a single instance. It breaks when:

  • multiple mcp-server instances run behind a load balancer;
  • a container restarts and loses in-memory session state;
  • GET/POST/DELETE transport requests land on different instances.

Goal

Replace the in-memory session store with a multi-instance-safe shared store.

Target Design

Introduce a SessionStore trait:

#[async_trait]
pub trait SessionStore {
    async fn create_transport_session(&self, session: TransportSessionRecord) -> Result<...>;
    async fn get_transport_session(&self, session_id: &str) -> Result<...>;
    async fn touch_transport_session(&self, session_id: &str, now: OffsetDateTime) -> Result<...>;
    async fn delete_transport_session(&self, session_id: &str) -> Result<...>;
    async fn append_pending_message(&self, session_id: &str, message: Value) -> Result<...>;
    async fn drain_pending_messages(&self, session_id: &str) -> Result<Vec<Value>, ...>;
    async fn cleanup_expired_transport_sessions(&self, now: OffsetDateTime) -> Result<u64, ...>;
}

Provide two implementations:

  • InMemorySessionStore
    • test-only and local dev fallback
  • PostgresSessionStore
    • production default

Storage Model

Add tables such as:

  • mcp_transport_sessions
  • mcp_transport_outbox

Suggested fields:

mcp_transport_sessions

  • id
  • workspace_id
  • agent_id
  • platform_api_key_id
  • transport_kind
  • status
  • created_at
  • updated_at
  • expires_at

mcp_transport_outbox

  • id
  • session_id
  • sequence_no
  • payload_json
  • created_at

Semantics

POST initialize

  • create transport session record in shared store.

GET SSE stream

  • attach to existing session record;
  • stream messages from outbox or polling bridge;
  • support reconnect without requiring process-local state.

POST JSON-RPC

  • resolve session from shared store;
  • append outbound messages to outbox.

DELETE

  • mark session closed in shared store;
  • best-effort cleanup of outbox rows.

Files and Modules

Add:

  • apps/mcp-server/src/session/mod.rs
  • apps/mcp-server/src/session/store.rs
  • apps/mcp-server/src/session/in_memory.rs
  • apps/mcp-server/src/session/postgres.rs

Potential shared data types:

  • crates/crank-core/src/transport_session.rs

Migration Order

Slice K1: trait and in-memory implementation
  • preserve current behavior but hide it behind a trait.
Slice K2: persistent schema and Postgres store
  • add migrations and registry/storage support.
Slice K3: transport handlers on shared store
  • wire GET/POST/DELETE through the trait-backed store.
Slice K4: cleanup and failover semantics
  • cleanup job for expired transport sessions;
  • reconnect tests across store-backed state.

Acceptance Criteria

  • transport session state survives process restart if store is persistent;
  • GET/POST/DELETE requests can hit different instances safely;
  • in-memory session state is no longer the production bottleneck.

Delivery Order

Recommended order:

  1. feat/refactoring-plan
  2. feat/id-display-support
  3. feat/postgres-pool-config
  4. feat/postgres-registry-modularization
  5. feat/sqlx-compile-time-verification
  6. feat/secret-crypto-hkdf
  7. feat/typed-timestamps
  8. feat/error-structure-normalization
  9. feat/correlation-id-propagation
  10. feat/runtime-rate-limiting-and-backpressure
  11. feat/distributed-mcp-session-store
  12. feat/post-refactor-cleanup

Reasoning:

  • Display is tiny and immediately reduces noise.
  • pool config is operational and isolated.
  • module split should happen before large query migrations, otherwise the macro migration lands into the monolith and becomes harder to review.
  • HKDF is isolated and low-risk.
  • typed timestamps should land after storage modules are split and query migration is underway.

Per-Slice Verification

Every slice must run:

  • just fmt-check
  • just check
  • just clippy
  • targeted tests for the touched area

Additionally:

  • SQL macro migration slices must run SQLx prepare/verification;
  • timestamp slices must run serialization regression tests;
  • secret crypto slice must run compatibility and roundtrip tests.

Success Definition

The refactoring track is complete when:

  • registry storage code is split into domain modules;
  • static SQL is compile-time validated;
  • SecretCrypto uses HKDF instead of SHA-256 as a KDF;
  • domain/storage timestamps are typed;
  • ID newtypes implement Display;
  • PostgreSQL pool settings are explicit and configurable;
  • timestamp formatting duplication is removed;
  • error variants carry structured context where it matters.

At that point, maintainability improves materially without changing the product model or feature scope.