diff --git a/CHANGELOG.md b/CHANGELOG.md index b58da68..d1b7eb7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- **Bounded, tenant/principal-aware proxy conversation correlation store** (P3): + the proxy `ConversationCorrelationStore` lineage map is now keyed by + `(principal, canonical_history_hash)` instead of the bare history hash, so + byte-identical conversation histories from different tenants no longer collide + onto one `conversation_id`. The two hash functions + (`canonical_history_hash` / `canonical_parent_history_hash`) are **byte-for-byte + unchanged** — isolation lives entirely in the map key, and an empty principal + reproduces the previous behavior exactly (`("", hash)` keyspace), so existing + single-turn and history-based multi-turn correlation is preserved. The store is + now bounded (TTL lazy-expiry on read + a max-entries FIFO cap, mirroring + `InMemorySessionStore`; defaults 3600 s / 20 000 entries, overridable via + `MORALSTACK_CORRELATION_TTL_SECONDS` / `MORALSTACK_CORRELATION_MAX_ENTRIES`), + removing the previous unbounded-growth/OOM risk. Principal is derived per + request as `X-Moralstack-Tenant-Id` header → HMAC-SHA256 of an + `Authorization: Bearer` token (secret from `MORALSTACK_PRINCIPAL_HMAC_SECRET`, + read per request, raw token/digest never logged) → empty-string sentinel. + `create_app` gains an optional `correlation_store=` parameter. `resolve()` is + best-effort on the TTL/eviction path (PROJECT_SPEC §5 invariant #6): a helper + failure never propagates into the request handler and a valid `msconv-*` id is + always minted. `ConversationLockManager._locks` bounding is deferred + (`TODO(P3-followup)`). No P0 decision/governance invariant is affected. - **Constitution retrieval unified to a single pass**: `get_relevant_principles` now runs exactly once per request, owned inside the risk thread (`models/risk/estimator.py:_get_principles_context`) at the unified @@ -54,6 +75,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Development +- **memory-guard trace-doc path casing fixed**: `scripts/check_memory_updated.py` + referenced `docs/TRACES/*` while the repository directory is `docs/traces/*`, making + the guard structurally unsatisfiable for `moralstack/server/` and + `moralstack/compliance/` changes (their only mapped docs were the mis-cased trace + files). Aligned the four `BEHAVIOR_DOC_MAP` entries to the actual `docs/traces/` + paths so the gate enforces the intended doc updates again. - **AI harness**: added a self-maintaining memory contract and evolved the Claude Code hook mechanics (Stop verify dedup, PreCompact snapshot, SessionEnd diary, docs-gate) with offline harness tests under `tests/harness/`. Test-suite isolation hardened diff --git a/README.md b/README.md index ca06fd6..f41f2f8 100644 --- a/README.md +++ b/README.md @@ -408,6 +408,9 @@ Environment is loaded via `moralstack/utils/env_loader.py`. | `MORALSTACK_OBSERVABILITY_JSONL_DIR` | `logs/observability` | JSONL output directory | | `MORALSTACK_DB_PATH` / `MORALSTACK_PERSIST_MODE` | — | Deprecated aliases — still work | | `MORALSTACK_ORCHESTRATOR_BORDERLINE_REFUSE_UPPER` | `0.95` | Upper boundary for the borderline-refuse zone | +| `MORALSTACK_CORRELATION_TTL_SECONDS` | `3600` | TTL (seconds) for the server proxy's lineage correlation store; must exceed the max expected inter-turn delay, and should stay aligned with the session-store TTL so a lineage entry doesn't outlive (or expire well before) its governance session | +| `MORALSTACK_CORRELATION_MAX_ENTRIES` | `20000` | Max entries in the server proxy's lineage correlation store (FIFO eviction beyond this cap) | +| `MORALSTACK_PRINCIPAL_HMAC_SECRET` | *(unset)* | Secret used to HMAC the bearer token into a tenant/principal identifier for conversation correlation; when unset, the HMAC principal path is disabled (falls back to the empty-string sentinel) | ### Default models by component diff --git a/docs/CODEBASE_FACTS.md b/docs/CODEBASE_FACTS.md index 95c41ef..95cbeb7 100644 --- a/docs/CODEBASE_FACTS.md +++ b/docs/CODEBASE_FACTS.md @@ -54,8 +54,9 @@ only when the `[ui]`/`[server]` extras are absent. | OpenAI-compatible HTTP serving is handled by the production proxy created by `moralstack.server.proxy.create_app` and launched by `examples/server_quickstart.py` | `server/proxy.py:create_app`; `examples/server_quickstart.py:build_app` | High | This is the path to use for COMPL-AI / IFBench-style proxy runs | | Production proxy serializes same-conversation requests with per-conversation locks | `server/proxy.py:72-119,234-242` | High | `ConversationLockManager`, 30s acquire timeout | | Proxy `turn_index` = user-message count − 1 (stateless) | `server/proxy.py:526-541` | High | | -| Proxy resolves conversation_id: header > extra_body > lineage correlation | `server/proxy.py:218-219,121-136` | High | header alias `X-Moralstack-Conversation-Id` | -| Lineage correlation maps history hashes → conversation_id; identical histories collide | `server/conversation_correlation.py:88-130` + module docstring | High | central COMPL-AI risk | +| Proxy resolves conversation_id: header > extra_body > lineage correlation (principal only affects the third path's map key) | `server/proxy.py:134-152` (resolver), `:654-709` (route handler) | High | header alias `X-Moralstack-Conversation-Id` | +| Lineage correlation maps `(principal, history hash)` → conversation_id; identical histories from the SAME principal still collide, but different principals never collide (P3 / P0-3 / A3, 2026-07-09) | `server/conversation_correlation.py:112-215` + module docstring | High | central COMPL-AI risk, narrowed to same-principal scope; see new row below | +| `ConversationCorrelationStore.resolve()` is best-effort on the TTL/eviction lookup/insert paths: a `_get_id`/`_put_id` failure is caught and logged at debug, never raised into the request; a valid `msconv-*` id is always returned (invariant #6, fix-pass 2026-07-09) | `server/conversation_correlation.py:143-188` | High | `resolve()` runs on the request path before the proxy's fail-closed `try` block (`server/proxy.py:356`), so this fix closes a real fail-open gap found by diff review | | Proxy quickstart recommended uvicorn command uses port 8080; `examples/server_quickstart.py:main()` defaults to port 8787 via `MORALSTACK_OPENAI_COMPATIBLE_API_PORT` | `examples/server_quickstart.py:14,74-79` | High | `python examples/server_quickstart.py` and `uvicorn examples.server_quickstart:app ...` both serve the production proxy | | Proxy `stream=True` replays the governed text as OpenAI-compatible synthetic SSE chunks (D2); it never forwards live upstream tokens | `server/proxy.py` (`_build_synthetic_sse_response`, `_handle_chat_completion_sync`) | High | Plan 1 governed delivery; governance runs first, then the full governed answer is replayed as SSE | | `moralstack-server` console entry raises NotImplementedError | `server/proxy.py:777-783` | High | use `create_app` | @@ -105,7 +106,7 @@ only when the `[ui]`/`[server]` extras are absent. | JSONL sink writes one file **per event_type** (`{event_type}.jsonl`), each line = `envelope.to_dict()`; SQLite normalizes the same `EventEnvelope` into typed columns. Same source, different shape | `observability/sinks/jsonl_sink.py:77-95`; `observability/router.py:37-54` | High | | | Per-module context-shape telemetry is emitted as `CONTEXT_SHAPE_RECORDED` orchestration events and folded into LLM-call/request metadata where available | `orchestration/orchestration_event_taxonomy.py`; `orchestration/controller.py:_emit_context_shape`; `observability/governance_audit.py:build_request_meta_from_result` | High | Queryable from JSONL and SQLite `orchestration_events`/`proxy_request_events` payload JSON; no dedicated SQLite migration or UI rendering in v1 | | Proxy `stream=True` returns a `StreamingResponse` of synthetic SSE chunks built from the governed answer (`_build_synthetic_sse_response`), terminated with `data: [DONE]`; no upstream `Stream` is created | `server/proxy.py` (`_build_synthetic_sse_response`, `_handle_chat_completion_sync`); `tests/test_server_proxy.py::test_streaming_replays_governed_text_as_sse_without_upstream` | High | Plan 1 / D2 governed synthetic SSE replay | -| Lineage correlation: identical canonicalized histories produce the same `conversation_id`; that id keys the per-conversation lock, the session store entry, and the ledger key — so colliding requests serialize and share governance state | `server/conversation_correlation.py:61-114`; `server/proxy.py:87-110,256,303-304`; `orchestration/ledger.py:254` | High | benchmark impact requires the dataset to actually contain identical-history samples | +| Lineage correlation: identical canonicalized histories, for the SAME principal, produce the same `conversation_id`; that id keys the per-conversation lock and the session store entry. **The ledger is NOT keyed by `conversation_id`** — its key is `(contract_hash, posture, domain)` (`LedgerKey`) | `server/conversation_correlation.py:112-159`; `server/proxy.py:78-131,381,438`; `orchestration/ledger.py:262` | High | Corrects a prior version of this row that cited `ledger.py:254` (a `LedgerResult` return statement, not a key) as a collision-shared surface — code wins (§9). Benchmark impact requires the dataset to actually contain identical-history samples under the same principal | | Client retry creates a new `requests` row at the same turn: `ProcessedRequest.request_id` is a fresh uuid4 per instance and proxy `turn_index` is recomputed statelessly | `orchestration/types.py:196`; `server/proxy.py:526-541` | High | retries are not deduplicated | | Full test suite: **2150 passed / 0 failed** with the `venv` (rerun 2026-07-07; includes the new `tests/harness/` hook tests). Additional skips only without `[ui]`/`[server]` extras | `./venv/Scripts/python.exe -m pytest -q` (run 2026-07-07) | High | Supersedes the earlier "1673, not rerun" note. A count is a point-in-time run outcome — re-run on later HEADs. | | DCCL calls the INJECTED policy via `generate_messages`/`generate` with a `GenerationConfig(response_format={"type":"json_object"})`, NOT a direct OpenAI SDK client call | `compliance/dccl.py:468-493` | High | Corrects an earlier cartography note that described DCCL as a "direct SDK path" | @@ -116,6 +117,7 @@ only when the `[ui]`/`[server]` extras are absent. | Prompt-caching reorder (Part A, 2026-07-06): every deliberative module/path system prompt is now a byte-stable, path-specific static prefix; user messages carry only per-request dynamic data. `response_format` stays `json_object` everywhere; retries/parsers unchanged. Full suite: 2024 passed / 0 failed with the project `venv` | `tests/test_static_prefix_stability.py`; `./venv/Scripts/python.exe -m pytest -q` | High | See `docs/MORALSTACK_CODEBASE_INDEX.md` §5.1 and `docs/modules/{critic,simulator,hindsight,perspectives,risk_estimator}.md` for the per-module split | | Unify-constitution-retrieval-single-pass (2026-07-07): `get_relevant_principles` runs exactly ONCE per request (owned in the risk thread, `estimator.py:_get_principles_context`), reused by deliberation's critic, the FAST_PATH `quick_check` (filtered to HARD), and the quick-check-failed deliberative fallback, via a controller-built `RequestAnalysisContext` gated on `RiskEstimation.retrieval_succeeded` (never on emptiness — an empty successful retrieval is authoritative, not degraded). `COMPLIANCE_FAST_PATH` consumes no principles and adds no retrieval. The 5 SEMANTIC ANALYSIS GUIDELINES moved from the intent mini's user message into `INTENT_CONTEXT_SYSTEM_PROMPT` (static/cacheable); the signals/operational minis remain principle-free (§5.3 unaffected). Full suite (after this change, including all new/adapted tests): 2083 passed / 0 failed with the project `venv` (`./venv/Scripts/python.exe -m pytest -q`, run 2026-07-07). 12 pre-existing tests required adaptation to the new `_get_principles_context` return type (`_PrinciplesContextResult` instead of a 2-tuple) and the new agent `retrieval_phase` keyword-only param — assertions unchanged, only unpacking/signatures adapted — see `tests/test_risk_estimator_runtime_domain.py`, `tests/test_estimator_developer_contract_interpretation.py`, `tests/test_multiturn_context_propagation.py`, `tests/test_prompt8_contract_priority.py`, `tests/test_constitution_retrieval.py`, `tests/test_constitution_retrieval_context_propagation.py`, `tests/test_controller_risk_context_propagation.py`) | `moralstack/models/risk/estimator.py:_get_principles_context,_PrinciplesContextResult`; `moralstack/orchestration/controller.py:_build_request_analysis_from_risk,_estimate_risk`; `moralstack/orchestration/deliberation_runner.py:run_deliberative_path,run_fast_path,_critique`; `moralstack/runtime/modules/critic_module.py:quick_check`; `moralstack/constitution/retriever.py` (`retrieval_phase` threaded through both agent kinds) | High | Release-time noise-floor gate (`scripts/ai/noise_floor_compare.py`) built and validated against the 4 existing HEAD benchmark runs already in `moralstack.db` (`81319498`, `79830edc`, `dceb24f8`, `60124777`; 6 pairs, n=83): reproduces the measured null band exactly — `final_action` divergence mean 4.22%/max 8.43%, all transitions confined to `NORMAL_COMPLETE<->SAFE_COMPLETE`, route flips 0%, hard-signal-code divergence 0% on the hard-signal-tripping subset. The post-change (branch-vs-HEAD) gate was run 2026-07-07 against the new branch benchmark `f803f5cf` (gpt-4o, 84 questions; 2 upstream rate-limited) vs the 4 HEAD baselines (n=83 matched each): **accepted as a conditional pass**. All hard criteria clean on all 4 pairs — route flips 0%, REFUSE-set identical, hard-signal codes byte-identical (0% divergence on 20-25 hard-signal pairs), every `final_action` transition confined to `NORMAL_COMPLETE<->SAFE_COMPLETE` (0 touching REFUSE). 3/4 pairs pass the `<=8%` `final_action`-divergence threshold (all 2.41%); the 4th (`81319498`) reads 8.43%, which is exactly the documented HEAD-vs-HEAD null-band `max` — i.e. stochastic gpt-4o noise on that pair, not a branch regression (the 8% threshold sits below the intrinsic noise ceiling). Thresholds were **not** adjusted (per plan). | | `GovernanceConfig.observability_mode` (default `"off"`) has NO wired runtime effect: it is read nowhere in `moralstack/` (grep `.observability_mode` → 0 non-test hits) and `get_observability_mode()` honors only `db_only`/`dual`/`file_only` (falling back to `db_only` if a DB path is set, else `file_only`), never `"off"`. Observability mode is controlled exclusively by `MORALSTACK_OBSERVABILITY_MODE` | `moralstack/sdk/config.py:58-59`; `moralstack/observability/config.py:64-77` | High | Re-verified 2026-07-07 by re-reading + grep; moved up from "conditionally verified" — fully code-determined, no external dependency | +| Bounded-tenant-aware-conversation-correlation (P3 / P0-3 / A3, 2026-07-09): `ConversationCorrelationStore`'s internal map is now `OrderedDict[tuple[str, str], _Entry]` keyed by `(principal, canonical_history_hash)` instead of a bare `dict[str, str]`, and bounded via TTL (lazy expiry on read, strict `>`, hit path does not refresh `inserted_at`) + a max-entries FIFO cap (`popitem(last=False)`), mirroring `InMemorySessionStore`. `canonical_history_hash`/`canonical_parent_history_hash` are **byte-for-byte unchanged** (no `salt` param) — isolation lives entirely in the map key, verified by a golden-digest test. Defaults `DEFAULT_CORRELATION_TTL_SECONDS=3600`, `DEFAULT_MAX_CORRELATION_ENTRIES=20_000`; constructor gains `ttl_seconds`/`max_entries`/`time_fn` (clock seam) and a `size()` helper. `resolve`/`observe_completed_turn` gain a keyword-only `principal: str = ""` param (default preserves pre-change behavior exactly: `("", hash)` keyspace). `proxy.py` adds `_extract_principal(request)`: (A) `X-Moralstack-Tenant-Id` header verbatim → (B) HMAC-SHA256 of an `Authorization: Bearer` token via `MORALSTACK_PRINCIPAL_HMAC_SECRET` (read per-request via `os.environ.get`, never captured once at `create_app`; unset silently skips B; raw token/digest never logged) → (C) empty-string sentinel. `create_app` gains `correlation_store: ConversationCorrelationStore | None = None` (mirrors `session_store=`) and, when not injected, wires `MORALSTACK_CORRELATION_TTL_SECONDS`/`MORALSTACK_CORRELATION_MAX_ENTRIES` (best-effort parse + range-validate, `ttl>0`/`max_entries>=1`, never raises at `create_app` on garbage/out-of-range values). `ConversationLockManager._locks` bounding is explicitly DEFERRED (`# TODO(P3-followup)` at `proxy.py:90-96`; a naive idle-prune races `acquire()`'s `_meta_lock` release before the blocking `lock.acquire()`) — the correlation-store bound already removes the dominant growth source. Full suite: 2214 passed / 1 deselected (`slow`) / 0 failed with the project `venv` (re-run 2026-07-09 after the best-effort `resolve()` fix-pass; the earlier pre-fix-pass count was 2208); `pre-commit run -a` green (ruff, black, mypy strict on `moralstack.orchestration.*`, changelog guard, memory guard) | `moralstack/server/conversation_correlation.py:105-215`; `moralstack/server/proxy.py:78-215,544-611,684` (`_extract_principal`, `_resolve_correlation_store_env_config`, `create_app`, `chat_completions`); `./venv/Scripts/python.exe -m pytest -q` (run 2026-07-09) | High | See `docs/traces/openai_compatible_multiturn.md` (conversation_id generation / propagation, Collision risks, Worker/concurrency sections) and `docs/MORALSTACK_CODEBASE_INDEX.md` §2/§17; new tests in `tests/test_conversation_correlation.py` (`TestPrincipalIsolation`, `TestBoundedStoreTTL`, `TestMaxsizeEviction`, `TestConcurrency`) and `tests/test_server_proxy.py` (`TestPrincipalAwareCorrelation`, `TestExtractPrincipal`, `TestCorrelationEnvWiring`) | | Optimize-domain-prefilter-caching-and-parallelism (2026-07-08): `DomainPrefilter.filter_domains` now sends a SYSTEM message (`_build_prefilter_system_prompt(domain_list)` — classifier role, `AVAILABLE DOMAINS`, procedure, falsification checks, confidence scale, JSON schema, byte-stable per domain config) + a query-only USER message (`f"USER QUERY:\n{query}"`); `_call_openai(prompt, *, system_prompt, retrieval_phase=...)` is now keyword-only on `system_prompt` and threads the SAME builder output into both the OpenAI system message and the persisted `system_prompt` (single source; old hardcoded 11-word `sys_msg` removed). Persisted-shape change is intentional: new `llm_calls` rows for `action="domain_prefilter"` carry the large static block under `system_prompt` and only the query under `prompt`; old rows unaffected, no migration. The local `_cache` key (`md5(f"{query}_{','.join(sorted(available_domains))}")`, `retriever.py:446`) is unchanged — it never hashed keyword/description content. Follow-up (2026-07-08): the `_build_prefilter_system_prompt` template was de-indented (flush-left instead of inheriting the method's 8-space body indent) to strip ~384 chars / ~96 tokens of leading-whitespace waste from the cached prefix; the real rendered prompt over the 21 bundled overlays is ~16,906 chars (~4,226 tokens, well above OpenAI's 1,024-token caching floor). Purely a whitespace change — verbatim instructions/schema unchanged, byte-stability tests substring-based so unaffected. `core` still never appears in the `AVAILABLE DOMAINS` section (`ALWAYS_EVALUATE`/`domains_to_check` exclusion unchanged, §5.5 P0 intact). The effective `max_parallel_agents` default is now **4** (was 2) at all five runtime sources — `ConstitutionRetrieverConfig` (`retriever.py:1099`), `ConstitutionStoreConfig` + `ConstitutionStore.__init__` kwarg (`store.py:462,498`), `CLIConfig` (`cli/models.py:492`), and `resolve_constitution_max_parallel_agents`'s env fallback (`pipeline/deliberation_stack.py:64`) — plus the `--max-parallel-agents` CLI help string (`cli/shell.py:1123`); `MORALSTACK_CONSTITUTION_MAX_PARALLEL_AGENTS` still overrides when set; the parallel runners (`_run_enhanced_agents_parallel`/`_run_agents_parallel`) were not otherwise touched. Full suite: 2173 passed / 0 failed (1 deselected `slow`) with the project `venv` | `moralstack/constitution/retriever.py:DomainPrefilter._build_prefilter_system_prompt,filter_domains,_call_openai`; `moralstack/constitution/store.py`; `moralstack/pipeline/deliberation_stack.py:resolve_constitution_max_parallel_agents`; `moralstack/cli/models.py`; `moralstack/cli/shell.py`; `./venv/Scripts/python.exe -m pytest -q` (run 2026-07-08) | High | See `docs/MORALSTACK_CODEBASE_INDEX.md` §5.1/§6 and `docs/modules/constitution_store.md`; `tests/test_static_prefix_stability.py::TestDomainPrefilterStaticPrefixStability`, `tests/test_domain_prefilter_cache.py`, `tests/test_constitution_retrieval_context_propagation.py`, `tests/test_constitution_max_parallel_agents_default.py` | --- @@ -128,7 +130,7 @@ These items involve external systems, deployment configuration, or runtime behav |---|---|---| | COMPL-AI uses the production proxy | Repo contains proxy mechanics and accommodation code (lineage correlation, history propagation, per-turn lock, risk-estimator COMPL-AI comment at `controller.py:797-799`) | Whether an actual external COMPL-AI runner points at the proxy, the exact request format it sends, and the benchmark dataset's collision prevalence are external facts not verifiable from this repo | | Single-uvicorn-worker requirement | `examples/server_quickstart.py:16-21` documents the requirement and explains why | Not runtime-enforced: an external load balancer routing turns of one conversation to different workers silently breaks continuity without an error | -| Benchmark collision prevalence | Hash-collision mechanism verified (`conversation_correlation.py:61-114`) | Whether a given run encounters identical-history samples depends on the external dataset | +| Benchmark collision prevalence | Hash-collision mechanism verified (`conversation_correlation.py:140-159`), now scoped to same-principal collisions only (P3 / P0-3 / A3) | Whether a given run encounters identical-history samples under the same principal depends on the external dataset | | Full test-suite pass count | Rerun 2026-07-07: **2150 passed / 0 failed** with the `venv` (`python -m pytest -q`, incl. new `tests/harness/`) | A pass count is a point-in-time run outcome, not a code fact; re-run to confirm on later HEADs | ## Future work / known gaps diff --git a/docs/MORALSTACK_CODEBASE_INDEX.md b/docs/MORALSTACK_CODEBASE_INDEX.md index a98f7b7..e8d20e4 100644 --- a/docs/MORALSTACK_CODEBASE_INDEX.md +++ b/docs/MORALSTACK_CODEBASE_INDEX.md @@ -270,10 +270,11 @@ Python `>=3.11` (`pyproject.toml:11`). Runtime deps: `openai>=2.24`, `pydantic>= - `emit_helpers.py` — `persist_*` / `async_persist_*` telemetry wrappers. ### Server proxy — `moralstack/server/` -- `proxy.py` — `create_app(openai_client, orchestrator, config, session_store)` - returns a FastAPI app exposing `POST /v1/chat/completions`, `/chat/completions`, - `GET /healthz`. `ConversationLockManager` (per-conversation locks), - `_handle_chat_completion_sync` (runs in a threadpool). Reads client +- `proxy.py` — `create_app(openai_client, orchestrator, config, session_store, + correlation_store)` returns a FastAPI app exposing `POST /v1/chat/completions`, + `/chat/completions`, `GET /healthz`. `ConversationLockManager` (per-conversation + locks; `_locks` growth-bounding is deferred — see FACTS), `_handle_chat_completion_sync` + (runs in a threadpool). Reads client `max_tokens`/`max_completion_tokens`/`temperature`/`top_p` into `ProcessedRequest.generation_overrides` via `GenerationOverrides.from_mapping(body, passthrough_unset=True)`; the SDK @@ -283,11 +284,22 @@ Python `>=3.11` (`pyproject.toml:11`). Runtime deps: `openai>=2.24`, `pydantic>= defaults do not apply there; on the **SDK/CLI** an unset field falls back to the env default (precedence override > `GenerationConfig` > env defaults `OPENAI_MAX_TOKENS`/`OPENAI_TEMPERATURE`/`OPENAI_TOP_P`). REFUSE wording is excluded - on every path. See `docs/modules/policy.md`. -- `conversation_correlation.py` — `ConversationCorrelationStore` (lineage hashing - → conversation_id) + `canonical_history_hash`, `canonical_parent_history_hash`. + on every path. See `docs/modules/policy.md`. Per request, `_extract_principal(request)` + derives a tenant/principal string (A: `X-Moralstack-Tenant-Id` header → B: HMAC-SHA256 + of an `Authorization: Bearer` token via `MORALSTACK_PRINCIPAL_HMAC_SECRET`, read + per-request → C: empty-string sentinel) that keys the correlation store's lineage map + (P3 / P0-3 / A3; see `docs/traces/openai_compatible_multiturn.md`). +- `conversation_correlation.py` — `ConversationCorrelationStore` (bounded, TTL + + max-entries FIFO eviction; lineage hashing → conversation_id, internal map keyed by + `(principal, history_hash)`) + `canonical_history_hash`, `canonical_parent_history_hash` + (unchanged by the principal-keying design — isolation lives entirely in the map key). + Constructor: `ttl_seconds` (default `DEFAULT_CORRELATION_TTL_SECONDS=3600`), + `max_entries` (default `DEFAULT_MAX_CORRELATION_ENTRIES=20_000`), `time_fn` (clock seam + for tests). `create_app` wires `MORALSTACK_CORRELATION_TTL_SECONDS` / + `MORALSTACK_CORRELATION_MAX_ENTRIES` (best-effort parse + range-validate, never raises). - `headers.py` — `build_governance_headers` (X-Moralstack-* response headers). -- `fingerprint.py` — request fingerprinting. +- `fingerprint.py` — request fingerprinting (calls the unchanged `canonical_history_hash`; + verified unaffected by the principal-keying change). ### UI — `moralstack/ui/` - `app.py` — FastAPI dashboard (`moralstack-ui`). Reads exclusively from the @@ -718,10 +730,19 @@ See `docs/traces/complai_llm_rules_flow.md`. ## 17. Known fragile areas -- **Lineage-based conversation correlation can collide.** Two samples with - byte-identical histories (and identical assistant outputs) map to the same - `conversation_id` (`conversation_correlation.py` docstring + `resolve`). This - is the central COMPL-AI risk — see `docs/traces/complai_llm_rules_flow.md`. +- **Lineage-based conversation correlation can still collide within a principal.** + Two samples with byte-identical histories (and identical assistant outputs) **for + the same principal** map to the same `conversation_id` (`conversation_correlation.py` + docstring + `resolve`) — this is the central COMPL-AI risk (see + `docs/traces/complai_llm_rules_flow.md`). As of P3 / P0-3 / A3, this no longer + crosses tenant/principal boundaries: the internal lineage map is keyed by + `(principal, history_hash)`, so identical histories from *different* principals + never collide (see `docs/traces/openai_compatible_multiturn.md`). +- **`ConversationLockManager._locks` growth is still unbounded (deferred).** One + `threading.Lock` is created per distinct `conversation_id` and never removed + (`proxy.py:83` area, `# TODO(P3-followup)`); the correlation store's TTL/max-entries + bound already removes the dominant growth source (~2 entries/turn vs 1 lock/conversation). + A safe idle-prune has a real race (see the plan for P3 / P0-3 / A3) and is a follow-up. - **UI requires SQLite.** `file_only` runs never appear in the dashboard; the UI reads only the DB. - **Cache governance.** Ledger fast-path can reuse a prior decision; the diff --git a/docs/traces/openai_compatible_multiturn.md b/docs/traces/openai_compatible_multiturn.md index d0d8237..cab1ba8 100644 --- a/docs/traces/openai_compatible_multiturn.md +++ b/docs/traces/openai_compatible_multiturn.md @@ -65,54 +65,127 @@ Empty governed content fails closed to a governed refusal guard fields are retained as audit metadata, but they no longer route delivery to an upstream call. The SDK uses the same governed delivery invariant. -### conversation_id generation / propagation (`proxy.py:218-219`, `121-136`) -Resolution precedence: +### conversation_id generation / propagation (`proxy.py:134-152`, resolver; +`proxy.py:654-709`, route handler) +Resolution precedence (principal keying only affects step 3): 1. HTTP header `X-Moralstack-Conversation-Id` (highest). 2. `extra_body.moralstack_conversation_id`. -3. **Lineage correlation** (`ConversationCorrelationStore.resolve`). +3. **Lineage correlation** (`ConversationCorrelationStore.resolve`), keyed + additionally by `principal` (see below). + +**Principal derivation (P3 / P0-3 / A3, `proxy.py:178-215`, `_extract_principal`)** +— layered, first match wins: +- **(A)** Trusted internal header `X-Moralstack-Tenant-Id` — used verbatim + (whitespace-stripped) when present and non-empty. Set by a trusted fronting + layer; this is a **documented trust boundary, not enforced by MoralStack** + (a hostile client could forge the header unless the fronting layer strips + client-supplied copies). Note this header starts with `x-` and therefore + passes `_collect_safe_headers` (`proxy.py:164-175`) and appears in the + correlation debug line (`proxy.py:683-691`) — acceptable for an identifier, + not for a secret. +- **(B)** HMAC-SHA256 of an `Authorization: Bearer ` value, keyed by + `MORALSTACK_PRINCIPAL_HMAC_SECRET` (read **per-request** via `os.environ.get`, + not captured once at `create_app`, so secret rotation takes effect + immediately). Only a `Bearer` scheme triggers this path; a non-Bearer scheme + (e.g. `Basic`) or a malformed/empty value skips to (C). If the secret is + unset, (B) is silently skipped (no hardcoded fallback). Neither the raw + token nor the derived digest is ever logged (the token stays filtered by + `_SENSITIVE_HEADER_MARKERS`, `proxy.py:155-161`). **Token rotation mid-conversation + changes the principal and silently splits lineage into a new conversation_id** + (expected, not a bug). +- **(C)** Empty-string sentinel `""` — the pre-change, no-principal behavior. + All existing single-tenant deployments and the byte-equality-locked tests + take this path. + +Composite keying is **anti-collision only, not an authentication/authorization +boundary**: it prevents two tenants' identical histories from sharing one +`conversation_id`, but it does not verify tenant identity beyond the (A)/(B) +trust assumptions above. Lineage correlation (`server/conversation_correlation.py`): -- `canonical_history_hash(messages)` = SHA-256 over canonicalized role+content. +- `canonical_history_hash(messages)` = SHA-256 over canonicalized role+content + (`:67-75`). **Unchanged by the principal-keying design** — no `salt` param, no + envelope; the isolation lives entirely in the store's internal map key. - `canonical_parent_history_hash` = hash of `messages[:-1]` when the last message - is `user`. -- `resolve`: if the request hash is known → return its conversation_id; else if - the parent hash is known → inherit that conversation_id and record the request - hash; else mint a new `msconv-` id (`conversation_correlation.py:99-114`). -- After a turn completes, `observe_completed_turn` records the history *including* - the assistant reply so the next request's parent hash links back - (`conversation_correlation.py:116-129`). - -### turn_index handling (`proxy.py:526-541`) + is `user` (`:78-91`). +- The store's internal map is `OrderedDict[tuple[str, str], _Entry]` keyed by + `(principal, hash)` (`:112-138`) — bounded via TTL (default + `DEFAULT_CORRELATION_TTL_SECONDS=3600`) and a max-entries FIFO cap (default + `DEFAULT_MAX_CORRELATION_ENTRIES=20_000`, `:41-42`), overridable via + `create_app(correlation_store=...)` or the env vars below. `resolve`: if the + `(principal, request_hash)` key is known and not expired → return its + conversation_id (hit path does **not** refresh the entry's TTL); else if + `(principal, parent_hash)` is known and not expired → inherit that + conversation_id and record the request hash under the same principal; else + mint a new `msconv-` id (`:140-159`). +- After a turn completes, `observe_completed_turn` records the completed history + (including the assistant reply) under `(principal, completed_hash)` so the next + request's parent hash links back within the same principal + (`:161-176`). +- Env wiring (`create_app`, `proxy.py:544-574`, `608-611`): + `MORALSTACK_CORRELATION_TTL_SECONDS` / `MORALSTACK_CORRELATION_MAX_ENTRIES`. + Best-effort parsed **and** range-validated (`ttl_seconds > 0`, `max_entries >= 1`); + a missing var, a parse error, or an out-of-range value falls back to the + constructor defaults — `create_app` never raises on a malformed value. + **Caveat**: the correlation TTL is not automatically aligned with the + `SessionStore` TTL; if an operator overrides only one, a lineage entry can + outlive (or expire before) its governance session state. + +### turn_index handling (`proxy.py:717-732`) Stateless: `turn_index = max(0, user_message_count - 1)`. Turn 0 = first request with one user message; turn 1 = two user messages, etc. Chosen so a server restart or multiple clients sharing a conversation_id don't desync from the client's view. ### Collision risks -- **Identical histories collide.** The module docstring states it explicitly: - two distinct samples whose histories (and assistant outputs) are byte-identical - cannot be distinguished without an external id - (`conversation_correlation.py:10-12`). For benchmarks that reuse the same - opening user message across many samples, all of them hash-collide to **one** - conversation_id. +- **Identical histories from the same principal still collide.** The module + docstring states it explicitly: two distinct samples whose histories (and + assistant outputs) are byte-identical, **for the same principal**, cannot be + distinguished without an external id (`conversation_correlation.py:17-18`). + For benchmarks that reuse the same opening user message across many samples + under one principal (the common no-tenant-header COMPL-AI case, principal + `""`), all of them hash-collide to **one** `conversation_id`. As of P3 / P0-3 / + A3, this collision **no longer crosses principal/tenant boundaries**: the + internal lineage map is keyed by `(principal, hash)`, so two tenants with + byte-identical histories resolve to different `conversation_id`s. - **Consequence (verified mechanism)**: the resolved conversation_id is the key - for the per-conversation lock (`ConversationLockManager`, `proxy.py:87-110`), - the `SessionStore` entry (`proxy.py:256,303-304`), and the ledger key - (`ledger.py:254`). So colliding requests are serialized under one lock and - share one governance-state/ledger entry — i.e. a decision or posture from one - sample can be read by another. (Whether a given run actually collides depends - on the benchmark data containing identical-history samples.) + for the per-conversation lock (`ConversationLockManager`, `proxy.py:78-131`), + the `SessionStore` entry (`proxy.py:381`, `:438`). **The ledger is NOT keyed by + `conversation_id`** — its key is `(contract_hash, posture, domain)` + (`orchestration/ledger.py:262`, `LedgerKey`); an earlier version of this trace + incorrectly cited `ledger.py:254` (a `LedgerResult` return statement, not a + key) as a collision-shared surface — corrected here (code wins, PROJECT_SPEC + §9). So same-principal colliding requests are serialized under one lock and + share one governance-state entry — i.e. a decision or posture from one sample + can be read by another sample under the *same* principal. (Whether a given + run actually collides depends on the benchmark data containing + identical-history samples under the same principal.) - **Mitigation available**: send a unique `X-Moralstack-Conversation-Id` header (or `extra_body.moralstack_conversation_id`) per logical conversation to bypass - lineage hashing entirely. - -### Worker / concurrency implications (`proxy.py:72-119`, `234-373`) + lineage hashing entirely, or a per-tenant `X-Moralstack-Tenant-Id` / + `Authorization` to isolate by principal. +- **Cross-tenant availability effect (new, P3 / P0-3 / A3)**: all principals + share one `max_entries` eviction pool, so a noisy tenant can evict another + tenant's lineage entries under load, causing a mid-conversation + `conversation_id` split for the affected tenant. This is an availability + effect, not a confidentiality leak (composite keying still prevents any + content/state sharing across principals). A per-principal cap is a possible + follow-up. + +### Worker / concurrency implications (`proxy.py:78-131`) - `ConversationLockManager` hands out one `threading.Lock` per conversation_id; empty conversation_id → no lock (independent request). Acquire timeout 30s → - HTTP 503 with `Retry-After: 10` (`ConversationLockTimeout`). + HTTP 503 with `Retry-After: 10` (`ConversationLockTimeout`). **`_locks` growth + is unbounded and its bounding is deferred** (`proxy.py:90-96`, + `# TODO(P3-followup)`): a naive idle-prune races `acquire()` releasing + `_meta_lock` before the blocking `lock.acquire()` call. The correlation + store's TTL/max-entries bound (above) already removes the dominant growth + source (~2 entries/turn vs 1 lock/conversation); a refcounted-waiter design is + the follow-up. - Must run a **single uvicorn worker** for multi-turn: each `--workers N` process - has its own pipeline, session store, and lock namespace; routing turns of one - conversation to different workers breaks continuity + has its own pipeline, session store, and lock/correlation-store namespace; + routing turns of one conversation to different workers breaks continuity, and + principal isolation and bounding are per-process, not cluster-wide (`examples/server_quickstart.py:16-21`). - Different conversation_ids run concurrently (threadpool); same conversation_id is serialized. diff --git a/moralstack/server/conversation_correlation.py b/moralstack/server/conversation_correlation.py index d1d18c0..2495a6c 100644 --- a/moralstack/server/conversation_correlation.py +++ b/moralstack/server/conversation_correlation.py @@ -7,6 +7,13 @@ records completed turns (including assistant replies) so the next request can link back to the same conversation. +The internal lineage map is keyed by ``(principal, history_hash)`` so that two +byte-identical histories from different principals (tenants) resolve to +different ``conversation_id``s (P3 / P0-3 / A3). The hash functions themselves +are unchanged; isolation lives entirely in the map key. The store is also +bounded (TTL + max entries, FIFO eviction) to avoid unbounded memory growth +under long-running benchmarks. + Limitation (informatic): two distinct samples whose histories and assistant outputs are byte-identical cannot be distinguished without an external id. """ @@ -15,9 +22,15 @@ import hashlib import json +import logging import threading +import time import uuid -from typing import Any +from collections import OrderedDict +from dataclasses import dataclass +from typing import Any, Callable + +_LOG = logging.getLogger(__name__) _MAX_CONTENT_PER_MESSAGE = 4096 @@ -25,6 +38,12 @@ # before the first user message in a user-only opening turn. _EMPTY_HISTORY_CANONICAL_HASH = hashlib.sha256(b"[]").hexdigest() +# Bounded-store defaults (P3 / P0-3 / A3): mirror InMemorySessionStore +# (moralstack/sdk/session_store.py) so a lineage entry and its governance +# session state expire on roughly the same horizon. +DEFAULT_CORRELATION_TTL_SECONDS = 3600 +DEFAULT_MAX_CORRELATION_ENTRIES = 20_000 + def _truncate(s: str) -> str: if len(s) <= _MAX_CONTENT_PER_MESSAGE: @@ -85,33 +104,88 @@ def canonical_parent_history_hash(messages: list[dict[str, Any]]) -> str | None: return canonical_history_hash(messages[:-1]) +@dataclass +class _Entry: + """Internal wrapper combining a resolved conversation_id with insertion time.""" + + conversation_id: str + inserted_at: float + + class ConversationCorrelationStore: """ - Thread-safe in-memory mapping from canonical history hashes to conversation_id. + Thread-safe, bounded in-memory mapping from ``(principal, canonical history + hash)`` to ``conversation_id``. - See module docstring for the resolution algorithm and limitations. + Bounded via TTL (lazy expiry on read) and a max-entries FIFO cap, mirroring + ``moralstack.sdk.session_store.InMemorySessionStore``. See module docstring + for the resolution algorithm and limitations. """ - def __init__(self) -> None: + def __init__( + self, + *, + ttl_seconds: float = DEFAULT_CORRELATION_TTL_SECONDS, + max_entries: int = DEFAULT_MAX_CORRELATION_ENTRIES, + time_fn: Callable[[], float] = time.time, + ) -> None: + if ttl_seconds <= 0: + raise ValueError(f"ttl_seconds must be > 0, got {ttl_seconds}") + if max_entries < 1: + raise ValueError(f"max_entries must be >= 1, got {max_entries}") + self._ttl_seconds = ttl_seconds + self._max_entries = max_entries + self._time_fn = time_fn self._lock = threading.RLock() - self._history_to_conversation: dict[str, str] = {} - - def resolve(self, messages: list[dict[str, Any]]) -> str: + # OrderedDict preserves insertion order for FIFO eviction. + self._history_to_conversation: OrderedDict[tuple[str, str], _Entry] = OrderedDict() + + def resolve(self, messages: list[dict[str, Any]], *, principal: str = "") -> str: + """ + Resolve (or mint) the ``conversation_id`` for a request history. + + Best-effort on the TTL/eviction lookup and insert paths (PROJECT_SPEC + §5 invariant #6): this runs on the request path, before the proxy's + fail-closed ``try`` block, so a failure in ``_get_id``/``_put_id`` + (e.g. from the injected clock or ``popitem`` during eviction) must + never propagate out of ``resolve()``. On any such failure we log at + debug and fall through to minting a fresh id, which is always + returned regardless of whether the store lookup/insert succeeded. + """ request_hash = canonical_history_hash(messages) parent_hash = canonical_parent_history_hash(messages) - - with self._lock: - if request_hash in self._history_to_conversation: - return self._history_to_conversation[request_hash] - - if parent_hash is not None and parent_hash in self._history_to_conversation: - conversation_id = self._history_to_conversation[parent_hash] - self._history_to_conversation[request_hash] = conversation_id - return conversation_id - - conversation_id = f"msconv-{uuid.uuid4().hex[:16]}" - self._history_to_conversation[request_hash] = conversation_id - return conversation_id + request_key = (principal, request_hash) + + try: + with self._lock: + existing = self._get_id(request_key) + if existing is not None: + return existing + + if parent_hash is not None: + parent_key = (principal, parent_hash) + parent_conversation_id = self._get_id(parent_key) + if parent_conversation_id is not None: + self._put_id(request_key, parent_conversation_id) + return parent_conversation_id + except Exception: + _LOG.debug( + "Lookup lineage nello store di correlazione fallito (TTL/eviction); " + "verrà generato un nuovo conversation_id.", + exc_info=True, + ) + + conversation_id = f"msconv-{uuid.uuid4().hex[:16]}" + try: + with self._lock: + self._put_id(request_key, conversation_id) + except Exception: + _LOG.debug( + "Inserimento del nuovo conversation_id nello store di correlazione fallito; " + "l'id generato viene comunque restituito.", + exc_info=True, + ) + return conversation_id def observe_completed_turn( self, @@ -119,6 +193,7 @@ def observe_completed_turn( messages: list[dict[str, Any]], assistant_content: str, conversation_id: str, + principal: str = "", ) -> None: """Register the completed history including the assistant reply for future turns.""" if not conversation_id: @@ -126,4 +201,32 @@ def observe_completed_turn( completed_messages = list(messages) + [{"role": "assistant", "content": assistant_content or ""}] completed_hash = canonical_history_hash(completed_messages) with self._lock: - self._history_to_conversation[completed_hash] = conversation_id + self._put_id((principal, completed_hash), conversation_id) + + def size(self) -> int: + """Number of stored entries, including expired ones (lazy eviction).""" + with self._lock: + return len(self._history_to_conversation) + + def _get_id(self, key: tuple[str, str]) -> str | None: + """Read under the lock, evicting the entry if expired. Hit path does not refresh TTL.""" + with self._lock: + entry = self._history_to_conversation.get(key) + if entry is None: + return None + if self._time_fn() - entry.inserted_at > self._ttl_seconds: + self._history_to_conversation.pop(key, None) + return None + return entry.conversation_id + + def _put_id(self, key: tuple[str, str], conversation_id: str) -> None: + """Insert/refresh an entry under the lock and enforce the max-entries FIFO cap.""" + with self._lock: + if key in self._history_to_conversation: + del self._history_to_conversation[key] + self._history_to_conversation[key] = _Entry( + conversation_id=conversation_id, + inserted_at=self._time_fn(), + ) + while len(self._history_to_conversation) > self._max_entries: + self._history_to_conversation.popitem(last=False) diff --git a/moralstack/server/proxy.py b/moralstack/server/proxy.py index 0a39441..8ab55fa 100644 --- a/moralstack/server/proxy.py +++ b/moralstack/server/proxy.py @@ -15,8 +15,11 @@ from __future__ import annotations +import hashlib +import hmac import json import logging +import os import threading import time import uuid @@ -49,7 +52,11 @@ from moralstack.sdk.bootstrap import _resolve_model from moralstack.sdk.config import GovernanceConfig from moralstack.sdk.session_store import InMemorySessionStore, SessionStoreProtocol -from moralstack.server.conversation_correlation import ConversationCorrelationStore +from moralstack.server.conversation_correlation import ( + DEFAULT_CORRELATION_TTL_SECONDS, + DEFAULT_MAX_CORRELATION_ENTRIES, + ConversationCorrelationStore, +) from moralstack.server.headers import build_governance_headers logger = logging.getLogger("moralstack.server.proxy") @@ -80,6 +87,13 @@ class ConversationLockManager: """ def __init__(self) -> None: + # TODO(P3-followup): bound _locks via refcounted-waiter design + # (increment waiters under _meta_lock BEFORE the blocking acquire; + # prune only entries with waiters==0 AND unlocked). Deferred: the + # naive idle-prune sketch races with acquire() releasing _meta_lock + # before blocking on lock.acquire() (see plan §5). The correlation + # store's TTL/max-entries bound already removes the dominant growth + # source (~2 entries/turn vs 1 lock/conversation). self._locks: dict[str, threading.Lock] = {} self._meta_lock = threading.Lock() @@ -121,18 +135,21 @@ def _resolve_conversation_id_from_body_and_correlation( messages: list[dict[str, Any]], extra_body: dict[str, Any] | None, correlation_store: ConversationCorrelationStore, + principal: str = "", ) -> str: """ Resolve ``conversation_id`` from ``extra_body`` or lineage correlation. The HTTP ``X-Moralstack-Conversation-Id`` header is handled in the route - handler and takes precedence over this function. + handler and takes precedence over this function. ``principal`` (empty by + default) keys the correlation store's internal lineage map so identical + histories from different principals do not collide (P3 / P0-3 / A3). """ if extra_body: explicit = extra_body.get("moralstack_conversation_id") if explicit: return str(explicit) - return correlation_store.resolve(messages) + return correlation_store.resolve(messages, principal=principal) _SENSITIVE_HEADER_MARKERS = ( @@ -158,6 +175,49 @@ def _collect_safe_headers(request: Request) -> dict[str, str]: return out +def _extract_principal(request: Request) -> str: + """ + Derive a per-request principal for conversation-correlation isolation + (P3 / P0-3 / A3). Layered derivation, first match wins: + + (A) Trusted internal header ``X-Moralstack-Tenant-Id`` \u2014 used verbatim + (after stripping whitespace) when present and non-empty. Set by a + trusted fronting layer; a hostile client can forge this header unless + that layer strips client-supplied copies (documented trust boundary, + not enforced here). + (B) HMAC-SHA256 of an ``Authorization: Bearer `` value, keyed by + the env secret ``MORALSTACK_PRINCIPAL_HMAC_SECRET`` (read per-request, + rotation-friendly \u2014 never captured once at ``create_app``). Only a + Bearer scheme triggers this path; a non-Bearer scheme (e.g. ``Basic``) + or a malformed/empty value skips straight to (C). If the secret is + unset, (B) is silently skipped (no hardcoded fallback, no per-request + log spam). + (C) Empty-string sentinel \u2014 the ``("", hash)`` keyspace, identical to the + pre-change (no-principal) resolution behavior. + + Composite keying built on this principal is anti-collision only, not an + authentication/authorization boundary (documented, not enforced here). + The raw ``Authorization`` value is never logged nor stored: this function + reads ``request.headers`` directly and returns only the resulting principal + string. The derived HMAC digest is likewise never logged and never + persisted to disk; it is used solely as an in-memory component of the + correlation-store map key (``(principal, history_hash)``). + """ + tenant_header = (request.headers.get("X-Moralstack-Tenant-Id") or "").strip() + if tenant_header: + return tenant_header + + auth_header = request.headers.get("Authorization") or "" + scheme, _, token = auth_header.partition(" ") + token = token.strip() + if scheme.strip().lower() == "bearer" and token: + secret = os.environ.get("MORALSTACK_PRINCIPAL_HMAC_SECRET") + if secret: + return hmac.new(secret.encode("utf-8"), token.encode("utf-8"), hashlib.sha256).hexdigest() + + return "" + + def _build_synthetic_chat_completion( content: str, *, @@ -265,15 +325,20 @@ def _handle_chat_completion_sync( openai_client: Any, orchestrator: Any, cfg: GovernanceConfig, + principal: str = "", ) -> Response: """ Synchronous request handler: governance, session store, upstream OpenAI call. Intended to run inside ``run_in_threadpool`` so blocking SDK calls do not stall - the ASGI event loop. + the ASGI event loop. ``principal`` (empty by default) is threaded into the + correlation store's map key so identical histories from different tenants + resolve to different conversation_ids (P3 / P0-3 / A3). """ hdr = (moralstack_conversation_id_header or "").strip() - conversation_id = hdr or _resolve_conversation_id_from_body_and_correlation(messages, extra_body, correlation_store) + conversation_id = hdr or _resolve_conversation_id_from_body_and_correlation( + messages, extra_body, correlation_store, principal=principal + ) request_id_for_audit: str = "" final_response_text: str = "" @@ -420,6 +485,7 @@ def _handle_chat_completion_sync( messages=messages, assistant_content=final_response_text, conversation_id=conversation_id, + principal=principal, ) except Exception: logger.debug("observe_completed_turn failed (non-fatal)", exc_info=True) @@ -479,12 +545,46 @@ def _handle_chat_completion_sync( return out_response +def _resolve_correlation_store_env_config() -> tuple[float, int]: + """ + Read the process-default correlation store bounds from the environment. + + Best-effort parse **and** range-validate ``MORALSTACK_CORRELATION_TTL_SECONDS`` + (> 0) and ``MORALSTACK_CORRELATION_MAX_ENTRIES`` (>= 1); a missing var, a + parse error, or an out-of-range value falls back to the constructor + defaults. This function never raises — ``create_app`` must never fail to + start on a malformed ``MORALSTACK_CORRELATION_*`` env value. + """ + ttl_seconds: float = DEFAULT_CORRELATION_TTL_SECONDS + raw_ttl = os.environ.get("MORALSTACK_CORRELATION_TTL_SECONDS") + if raw_ttl is not None: + try: + parsed_ttl = float(raw_ttl) + except ValueError: + parsed_ttl = None + if parsed_ttl is not None and parsed_ttl > 0: + ttl_seconds = parsed_ttl + + max_entries: int = DEFAULT_MAX_CORRELATION_ENTRIES + raw_max = os.environ.get("MORALSTACK_CORRELATION_MAX_ENTRIES") + if raw_max is not None: + try: + parsed_max = int(raw_max) + except ValueError: + parsed_max = None + if parsed_max is not None and parsed_max >= 1: + max_entries = parsed_max + + return ttl_seconds, max_entries + + def create_app( *, openai_client: Any, orchestrator: OrchestrationController, config: GovernanceConfig | None = None, session_store: SessionStoreProtocol | None = None, + correlation_store: ConversationCorrelationStore | None = None, ) -> FastAPI: """ Build a FastAPI app instance. @@ -495,6 +595,9 @@ def create_app( orchestrator: a configured OrchestrationController. config: optional GovernanceConfig. Defaults to GovernanceConfig(). session_store: optional SessionStoreProtocol (defaults to InMemorySessionStore). + correlation_store: optional ConversationCorrelationStore (defaults to a + store bounded via MORALSTACK_CORRELATION_TTL_SECONDS / + MORALSTACK_CORRELATION_MAX_ENTRIES, or their constructor defaults). Returns: A FastAPI app ready to be served by uvicorn. @@ -504,7 +607,9 @@ def create_app( logger.info("MoralStack proxy upstream generation model: %s", upstream_generation_model) store: SessionStoreProtocol = session_store if session_store is not None else InMemorySessionStore() lock_manager = ConversationLockManager() - correlation_store = ConversationCorrelationStore() + if correlation_store is None: + _ttl_seconds, _max_entries = _resolve_correlation_store_env_config() + correlation_store = ConversationCorrelationStore(ttl_seconds=_ttl_seconds, max_entries=_max_entries) # Initialize observability for the proxy lifetime so all governance events # are persisted to DB / JSONL per MORALSTACK_OBSERVABILITY_* env vars. # Pattern parallels cli/shell.py: init_db + create_run + set_current_run_id. @@ -576,6 +681,7 @@ async def chat_completions( raise HTTPException(status_code=400, detail="`messages` must not be empty") extra_body = body.get("extra_body") if isinstance(body.get("extra_body"), dict) else None + principal = _extract_principal(request) safe_headers = _collect_safe_headers(request) has_extra_conv = bool(extra_body and extra_body.get("moralstack_conversation_id")) @@ -601,6 +707,7 @@ async def chat_completions( openai_client=openai_client, orchestrator=orchestrator, cfg=cfg, + principal=principal, ) return app diff --git a/scripts/check_memory_updated.py b/scripts/check_memory_updated.py index 52f9ff9..128549d 100644 --- a/scripts/check_memory_updated.py +++ b/scripts/check_memory_updated.py @@ -26,19 +26,19 @@ # Tieni i prefissi allineati a .claude/hooks/stop_gate.py:BEHAVIOR_PREFIXES. BEHAVIOR_DOC_MAP: tuple[tuple[str, tuple[str, ...]], ...] = ( ("moralstack/runtime/decision/", ("docs/decision_policy.md",)), - ("moralstack/compliance/", ("docs/TRACES/complai_llm_rules_flow.md",)), + ("moralstack/compliance/", ("docs/traces/complai_llm_rules_flow.md",)), ( "moralstack/orchestration/", - ("docs/modules/orchestrator.md", "docs/TRACES/governance_decision_flow.md"), + ("docs/modules/orchestrator.md", "docs/traces/governance_decision_flow.md"), ), ("moralstack/prompts/", ("docs/decision_policy.md",)), ( "moralstack/observability/", - ("docs/modules/observability.md", "docs/TRACES/observability_db_to_ui.md"), + ("docs/modules/observability.md", "docs/traces/observability_db_to_ui.md"), ), ( "moralstack/server/", - ("docs/TRACES/openai_compatible_multiturn.md", "docs/TRACES/observability_db_to_ui.md"), + ("docs/traces/openai_compatible_multiturn.md", "docs/traces/observability_db_to_ui.md"), ), ("moralstack/constitution/", ("docs/modules/constitution_store.md", "docs/constitution.md")), ) diff --git a/tests/test_conversation_correlation.py b/tests/test_conversation_correlation.py index e52a28c..0f13bd4 100644 --- a/tests/test_conversation_correlation.py +++ b/tests/test_conversation_correlation.py @@ -2,17 +2,133 @@ from __future__ import annotations +import hashlib +import json +import threading +from unittest.mock import Mock + +import pytest + from moralstack.server.conversation_correlation import ( + _EMPTY_HISTORY_CANONICAL_HASH, ConversationCorrelationStore, canonical_history_hash, + canonical_parent_history_hash, ) +def _reference_canonical_history_hash(messages: list[dict]) -> str: + """ + Independent reimplementation of ``canonical_history_hash``'s algorithm + (normalization + hashing reproduced, the module function is NOT called). + + Used as a golden-digest lock: the composite-key isolation design (P3 / + P0-3 / A3) must not incidentally change the hash output, since the + isolation lives entirely in the store's internal map key, never in the + digest. + """ + max_len = 4096 + + def truncate(s: str) -> str: + return s if len(s) <= max_len else s[:max_len] + + def normalize_content(content) -> str: + if content is None: + return "" + if isinstance(content, str): + return truncate(content) + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if not isinstance(block, dict): + parts.append(json.dumps(block, sort_keys=True, ensure_ascii=False, separators=(",", ":"))) + continue + btype = str(block.get("type") or "") + if btype == "text": + parts.append(truncate(str(block.get("text") or ""))) + else: + parts.append(json.dumps(block, sort_keys=True, ensure_ascii=False, separators=(",", ":"))) + return truncate("".join(parts)) + return truncate(json.dumps(content, sort_keys=True, ensure_ascii=False, separators=(",", ":"))) + + normalized = [{"role": str(m.get("role", "") or ""), "content": normalize_content(m.get("content"))} for m in messages] + blob = json.dumps(normalized, sort_keys=True, ensure_ascii=False, separators=(",", ":")) + return hashlib.sha256(blob.encode("utf-8")).hexdigest() + + class TestCanonicalHistoryHash: def test_deterministic(self): msgs = [{"role": "system", "content": "S"}, {"role": "user", "content": "U"}] assert canonical_history_hash(msgs) == canonical_history_hash(list(msgs)) + def test_hash_functions_unchanged_golden_digest(self): + """ + Golden-digest lock (P3 design point, "do not deviate" in the handoff): + canonical_history_hash / canonical_parent_history_hash stay byte-for-byte + unchanged. Composite (principal, hash) keying lives entirely inside the + store's internal map key, never in the digest itself. + """ + cases = [ + ([], "4f53cda18c2baa0c0354bb5f9a3ecbe5ed12ab4d8e11ba873c2f11161202b945"), + ( + [{"role": "user", "content": "hi"}], + "4e79873118cd9be7a1f0308b9cd772950c5410c74ca3fe1ba2626cba009a9237", + ), + ( + [ + {"role": "system", "content": "You are helpful. 你好"}, + {"role": "user", "content": "First Q café"}, + {"role": "assistant", "content": "Reply 🌍"}, + {"role": "user", "content": "Second Q"}, + ], + "4666af65b1114ebc48fa2e9f3e8fa41494d995ad413b38f764159bf5d9ed5714", + ), + ] + for messages, golden_digest in cases: + actual = canonical_history_hash(messages) + assert actual == golden_digest, f"digest moved for {messages!r}: {actual!r} != {golden_digest!r}" + assert actual == _reference_canonical_history_hash(messages) + + # _MAX_CONTENT_PER_MESSAGE (4096) truncation boundary: content at exactly + # the boundary and content beyond it must hash identically (both truncate + # to the same 4096-char prefix). + at_boundary = [{"role": "user", "content": "a" * 4096}] + beyond_boundary = [{"role": "user", "content": "a" * 5000}] + assert canonical_history_hash(at_boundary) == canonical_history_hash(beyond_boundary) + assert canonical_history_hash(beyond_boundary) == _reference_canonical_history_hash(beyond_boundary) + + def test_empty_list_equals_root_constant(self): + """The parent-lookup root path (canonical_parent_history_hash for a + single opening user turn) depends on this identity holding.""" + assert canonical_history_hash([]) == _EMPTY_HISTORY_CANONICAL_HASH + + def test_canonical_parent_history_hash_golden_digest(self): + """ + Golden-digest lock for ``canonical_parent_history_hash`` directly + (FIX 4, diff review 2026-07-09): the prior golden test only pinned + ``canonical_history_hash``; this pins the parent-hash function itself + so composite-key isolation cannot incidentally change either digest. + """ + # Empty-history-root case: a single opening user turn has no prior + # history, so the parent hash is the frozen empty-list root constant. + assert canonical_parent_history_hash([{"role": "user", "content": "hi"}]) == _EMPTY_HISTORY_CANONICAL_HASH + + # Fixed multi-turn input (same case pinned above for + # canonical_history_hash): the parent hash of the trailing user turn + # must equal canonical_history_hash(messages[:-1]) and a hard-coded + # pre-change hex digest. + multi_turn = [ + {"role": "system", "content": "You are helpful. 你好"}, + {"role": "user", "content": "First Q café"}, + {"role": "assistant", "content": "Reply 🌍"}, + {"role": "user", "content": "Second Q"}, + ] + golden_parent_digest = "86d67e09fd83cf50b53927afacfff5af76e124ebea65be1f3820a31524f0beae" + actual_parent = canonical_parent_history_hash(multi_turn) + assert actual_parent == golden_parent_digest + assert actual_parent == canonical_history_hash(multi_turn[:-1]) + assert actual_parent == _reference_canonical_history_hash(multi_turn[:-1]) + class TestConversationCorrelationStore: def test_resolve_new_conversation_id_format(self): @@ -40,3 +156,261 @@ def test_exact_request_replay_returns_same_id(self): store = ConversationCorrelationStore() msgs = [{"role": "user", "content": "x"}] assert store.resolve(msgs) == store.resolve(list(msgs)) + + def test_resolve_default_principal_matches_pre_change_behavior(self): + """resolve(msgs) (no principal kwarg) must be identical to + resolve(msgs, principal="") — the (\"\", hash) keyspace matches the + pre-change, no-principal path exactly.""" + store = ConversationCorrelationStore() + msgs = [{"role": "user", "content": "same content"}] + cid_default = store.resolve(msgs) + cid_explicit_empty = store.resolve(list(msgs), principal="") + assert cid_default == cid_explicit_empty + + +class TestPrincipalIsolation: + def test_resolve_isolates_principals(self): + """Core bug fix: identical histories, different principals -> different + conversation_ids; the same principal reused -> the same id.""" + store = ConversationCorrelationStore() + msgs = [{"role": "user", "content": "identical history"}] + + cid_a1 = store.resolve(msgs, principal="A") + cid_b = store.resolve(msgs, principal="B") + cid_a2 = store.resolve(msgs, principal="A") + + assert cid_a1 != cid_b + assert cid_a1 == cid_a2 + + def test_observe_completed_turn_lineage_within_principal(self): + store = ConversationCorrelationStore() + t0 = [ + {"role": "system", "content": "Scenario rules"}, + {"role": "user", "content": "First test message"}, + ] + cid0 = store.resolve(t0, principal="A") + store.observe_completed_turn( + messages=t0, + assistant_content="Model reply", + conversation_id=cid0, + principal="A", + ) + t1 = t0 + [{"role": "assistant", "content": "Model reply"}, {"role": "user", "content": "Second"}] + cid1 = store.resolve(t1, principal="A") + assert cid1 == cid0 + + def test_observe_completed_turn_does_not_leak_lineage_across_principal(self): + """ + Primary P3/P0-3 regression lock: byte-identical follow-up history resolved + under a DIFFERENT principal must NOT chain into principal A's lineage — + it must mint a fresh conversation_id. + """ + store = ConversationCorrelationStore() + t0 = [ + {"role": "system", "content": "Scenario rules"}, + {"role": "user", "content": "First test message"}, + ] + cid0 = store.resolve(t0, principal="A") + store.observe_completed_turn( + messages=t0, + assistant_content="Model reply", + conversation_id=cid0, + principal="A", + ) + t1 = t0 + [{"role": "assistant", "content": "Model reply"}, {"role": "user", "content": "Second"}] + cid1_other_principal = store.resolve(t1, principal="B") + assert cid1_other_principal != cid0 + + def test_odd_character_principal_isolation_stable(self): + """Principals containing odd characters (no serialization to spoof with + tuple keying) still isolate and remain stable.""" + store = ConversationCorrelationStore() + msgs = [{"role": "user", "content": "same"}] + cid_brace = store.resolve(msgs, principal="}") + cid_colon = store.resolve(msgs, principal=":") + cid_brace_again = store.resolve(list(msgs), principal="}") + assert cid_brace != cid_colon + assert cid_brace == cid_brace_again + + +class TestBoundedStoreTTL: + def test_ttl_expiry_mints_new_id(self): + clock = {"t": 1000.0} + store = ConversationCorrelationStore(ttl_seconds=10.0, time_fn=lambda: clock["t"]) + msgs = [{"role": "user", "content": "hi"}] + + cid0 = store.resolve(msgs) + + # Boundary: age == ttl_seconds is NOT expired (strict '>'). + clock["t"] = 1010.0 + assert store.resolve(list(msgs)) == cid0 + + # Past the TTL: stale entry evicted-on-read, a new id is minted. + clock["t"] = 1010.001 + cid1 = store.resolve(list(msgs)) + assert cid1 != cid0 + + def test_ttl_hit_path_does_not_refresh_inserted_at(self): + clock = {"t": 0.0} + store = ConversationCorrelationStore(ttl_seconds=5.0, time_fn=lambda: clock["t"]) + msgs = [{"role": "user", "content": "hi"}] + + cid0 = store.resolve(msgs) + clock["t"] = 4.0 + # A read at t=4 must NOT push expiry out to t=9; expiry stays anchored + # to the original insertion at t=0. + assert store.resolve(list(msgs)) == cid0 + clock["t"] = 5.001 + assert store.resolve(list(msgs)) != cid0 + + def test_ttl_lineage_survives_when_only_inter_turn_delay_elapses(self): + """3-turn conversation where turn-1 entries have individually expired by + the time of turn 3, but chaining still works because each turn writes + fresh request/completed entries. Locks: TTL only needs to exceed the + inter-turn delay, not the whole conversation duration.""" + clock = {"t": 0.0} + ttl = 100.0 + store = ConversationCorrelationStore(ttl_seconds=ttl, time_fn=lambda: clock["t"]) + + t0 = [{"role": "user", "content": "Q1"}] + cid0 = store.resolve(t0) + store.observe_completed_turn(messages=t0, assistant_content="A1", conversation_id=cid0) + + clock["t"] = 60.0 # inter-turn delay < ttl + t1 = t0 + [{"role": "assistant", "content": "A1"}, {"role": "user", "content": "Q2"}] + cid1 = store.resolve(t1) + store.observe_completed_turn(messages=t1, assistant_content="A2", conversation_id=cid1) + assert cid1 == cid0 + + clock["t"] = 120.0 # turn-0 entries (inserted at t=0) are now expired, + # but turn-1 entries (inserted at t=60) are still within ttl. + t2 = t1 + [{"role": "assistant", "content": "A2"}, {"role": "user", "content": "Q3"}] + cid2 = store.resolve(t2) + assert cid2 == cid0 + + +class TestMaxsizeEviction: + def test_maxsize_eviction_oldest_first(self): + store = ConversationCorrelationStore(max_entries=2) + cid_a = store.resolve([{"role": "user", "content": "history A"}]) + store.resolve([{"role": "user", "content": "history B"}]) + store.resolve([{"role": "user", "content": "history C"}]) + + assert store.size() <= 2 + # 'history A' was the oldest insertion -> evicted -> re-resolving mints + # a new id (rather than returning the original cid_a). + cid_a_again = store.resolve([{"role": "user", "content": "history A"}]) + assert cid_a_again != cid_a + + def test_maxsize_one_observe_after_eviction_no_crash(self): + store = ConversationCorrelationStore(max_entries=1) + t0 = [{"role": "user", "content": "Q1"}] + cid0 = store.resolve(t0) + # Every new history evicts the prior single entry. + store.resolve([{"role": "user", "content": "Q-other"}]) + # observe_completed_turn's target hash was already evicted; must not + # raise (KeyError or otherwise) — it simply (re)inserts. + store.observe_completed_turn(messages=t0, assistant_content="A1", conversation_id=cid0) + assert store.size() <= 1 + + +class TestConcurrency: + def test_concurrent_resolve_thread_safety(self): + store = ConversationCorrelationStore(max_entries=200) + errors: list[BaseException] = [] + + def worker(worker_id: int) -> None: + try: + for i in range(50): + msgs = [{"role": "user", "content": f"w{worker_id}-{i}"}] + store.resolve(msgs) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + threads = [threading.Thread(target=worker, args=(w,)) for w in range(20)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors + assert store.size() <= 200 + + def test_concurrent_mixed_principals_identical_history(self): + """Threads racing with different principals on an identical history: + no exceptions, and ids partition cleanly by principal.""" + store = ConversationCorrelationStore(max_entries=1000) + msgs = [{"role": "user", "content": "shared identical history"}] + errors: list[BaseException] = [] + results: dict[str, set[str]] = {"A": set(), "B": set(), "C": set()} + lock = threading.Lock() + + def worker(principal: str) -> None: + try: + for _ in range(30): + cid = store.resolve(list(msgs), principal=principal) + with lock: + results[principal].add(cid) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + threads = [threading.Thread(target=worker, args=(p,)) for p in ("A", "B", "C") for _ in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors + # Each principal converges to exactly one id (races resolve to a single + # winner id per principal since resolve() is lock-protected). + assert all(len(ids) == 1 for ids in results.values()) + assert results["A"] != results["B"] != results["C"] + assert results["A"].isdisjoint(results["B"]) + assert results["A"].isdisjoint(results["C"]) + assert results["B"].isdisjoint(results["C"]) + + +class TestResolveBestEffort: + """ + Blocking regression lock (diff review 2026-07-09): ``resolve()`` must never + raise from the TTL/eviction helper paths (``_get_id``/``_put_id``), and must + always return a valid ``msconv-*`` id even when a helper fails. This runs on + the request path, ahead of the proxy's fail-closed ``try`` block, so an + unguarded raise here would break the request (PROJECT_SPEC §5 invariant #6). + """ + + def test_resolve_survives_lookup_failure_and_mints_fresh_id(self): + """The hit-lookup branch (`_get_id`) raises -> `resolve()` swallows it + and still returns a valid minted id, never propagating the exception.""" + store = ConversationCorrelationStore() + store._get_id = lambda key: (_ for _ in ()).throw(RuntimeError("boom")) + cid = store.resolve([{"role": "user", "content": "hi"}]) + assert isinstance(cid, str) + assert cid.startswith("msconv-") + + def test_resolve_survives_insert_failure_and_still_returns_id(self): + """The insert branch (`_put_id`) raises on every call -> `resolve()` + swallows it and still returns the freshly minted id.""" + store = ConversationCorrelationStore() + store._put_id = lambda key, conversation_id: (_ for _ in ()).throw(RuntimeError("boom")) + cid = store.resolve([{"role": "user", "content": "hi"}]) + assert isinstance(cid, str) + assert cid.startswith("msconv-") + + def test_resolve_survives_time_fn_failure_during_eviction(self): + """A failing injected clock (used internally by `_get_id`/`_put_id` for + TTL/eviction bookkeeping) must not escape `resolve()`.""" + store = ConversationCorrelationStore(time_fn=Mock(side_effect=RuntimeError("clock failed"))) + cid = store.resolve([{"role": "user", "content": "hi"}]) + assert isinstance(cid, str) + assert cid.startswith("msconv-") + + +class TestStoreConstructorValidation: + def test_invalid_ttl_raises(self): + with pytest.raises(ValueError, match="ttl_seconds"): + ConversationCorrelationStore(ttl_seconds=0) + + def test_invalid_max_entries_raises(self): + with pytest.raises(ValueError, match="max_entries"): + ConversationCorrelationStore(max_entries=0) diff --git a/tests/test_server_proxy.py b/tests/test_server_proxy.py index 8740127..7148a24 100644 --- a/tests/test_server_proxy.py +++ b/tests/test_server_proxy.py @@ -5,7 +5,10 @@ from __future__ import annotations import asyncio +import hashlib +import hmac import json +import logging import random import threading import time @@ -20,6 +23,7 @@ import httpx # noqa: E402 from fastapi.testclient import TestClient # noqa: E402 from httpx import ASGITransport # noqa: E402 +from starlette.datastructures import Headers # noqa: E402 from moralstack.orchestration.types import ( # noqa: E402 FinalResponse, @@ -29,6 +33,7 @@ ) from moralstack.runtime.orchestrator import create_minimal_orchestrator # noqa: E402 from moralstack.sdk.config import GovernanceConfig # noqa: E402 +from moralstack.server.conversation_correlation import ConversationCorrelationStore # noqa: E402 from moralstack.server.proxy import create_app # noqa: E402 from tests.test_orchestrator import MockPolicyLLM, MockRiskEstimator # noqa: E402 @@ -1275,3 +1280,364 @@ def test_proxy_output_finalized_event_persisted(self, tmp_path, monkeypatch): assert payload["governed_delivery"] is True assert payload["wrapped_client_delivery_call"] is False assert payload["finish_reason"] == "stop" + + +# ─── Bounded, tenant/principal-aware conversation correlation (P3 / P0-3 / A3) ─── + + +def _build_correlation_client( + final_action: str = "NORMAL_COMPLETE", + *, + correlation_store: ConversationCorrelationStore | None = None, + refuse_content: str = "Cannot help with that.", +): + """Standalone variant of the ``client_factory`` fixture that additionally + supports injecting a ``correlation_store`` (mirrors ``session_store=``).""" + mock_orchestrator = MagicMock() + if final_action == "REFUSE": + mock_orchestrator.process = MagicMock(return_value=_make_result("REFUSE", content=refuse_content)) + else: + mock_orchestrator.process = MagicMock(return_value=_make_result(final_action, content="guidance")) + + mock_openai = MagicMock() + mock_openai.chat.completions.create = MagicMock(return_value=_make_upstream_chat_completion()) + + kwargs: dict[str, Any] = {} + if correlation_store is not None: + kwargs["correlation_store"] = correlation_store + app = create_app(openai_client=mock_openai, orchestrator=mock_orchestrator, config=GovernanceConfig(), **kwargs) + return TestClient(app), mock_openai, mock_orchestrator + + +def _extract_default_correlation_store(app) -> ConversationCorrelationStore: + """ + Recover the ``ConversationCorrelationStore`` that ``create_app`` built + internally (when no ``correlation_store=`` is injected) via the + ``chat_completions`` closure — the store is not exposed on ``app.state``, + so this is the only way to assert env-parsed values actually reached it + without modifying ``proxy.py``. + """ + for route in app.routes: + if getattr(route, "path", None) == "/v1/chat/completions": + endpoint = route.endpoint + idx = endpoint.__code__.co_freevars.index("correlation_store") + return endpoint.__closure__[idx].cell_contents + raise AssertionError("/v1/chat/completions route not found on app") + + +class _FakeRequest: + """Minimal ``Request``-like stand-in exposing only ``.headers`` (Starlette + ``Headers``, case-insensitive) — enough for ``_extract_principal`` unit tests.""" + + def __init__(self, headers: dict[str, str]): + self.headers = Headers(headers) + + +class TestPrincipalAwareCorrelation: + def test_tenant_header_isolates_conversation_id(self): + client, _, _ = _build_correlation_client("NORMAL_COMPLETE") + body = {"model": "gpt-4o", "messages": [{"role": "user", "content": "Identical body"}]} + + r_a = client.post("/v1/chat/completions", json=body, headers={"X-Moralstack-Tenant-Id": "tenant-A"}) + r_b = client.post("/v1/chat/completions", json=body, headers={"X-Moralstack-Tenant-Id": "tenant-B"}) + + cid_a = r_a.headers.get("X-Moralstack-Conversation-Id") + cid_b = r_b.headers.get("X-Moralstack-Conversation-Id") + assert cid_a is not None and cid_b is not None + assert cid_a != cid_b + + def test_anonymous_when_bearer_present_but_no_tenant_and_no_secret(self, monkeypatch): + """Strengthened re-run: an Authorization Bearer is present with no + tenant header and no HMAC secret configured — lineage stays stable + across turns via the anonymous ("", hash) path (asserts B is skipped, + not merely that the header is ignored).""" + monkeypatch.delenv("MORALSTACK_PRINCIPAL_HMAC_SECRET", raising=False) + client, _, _ = _build_correlation_client("NORMAL_COMPLETE") + + r1 = client.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o", + "messages": [ + {"role": "system", "content": "Sys"}, + {"role": "user", "content": "Q1"}, + ], + }, + headers={"Authorization": "Bearer key-X"}, + ) + cid1 = r1.headers.get("X-Moralstack-Conversation-Id") + assert cid1 is not None and cid1.startswith("msconv-") + + r2 = client.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o", + "messages": [ + {"role": "system", "content": "Sys"}, + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "guidance"}, + {"role": "user", "content": "Q2"}, + ], + }, + headers={"Authorization": "Bearer key-X"}, + ) + cid2 = r2.headers.get("X-Moralstack-Conversation-Id") + assert cid2 == cid1 + + def test_authorization_hmac_derives_distinct_principal(self, monkeypatch, caplog): + """Path B: identical body, different bearer tokens -> different + conversation_ids. Security regression guard: neither the raw + Authorization value nor the derived HMAC digest may appear in + ``_collect_safe_headers`` output or captured log records.""" + from moralstack.server.proxy import _collect_safe_headers + + monkeypatch.setenv("MORALSTACK_PRINCIPAL_HMAC_SECRET", "dummy-test-secret") + client, _, _ = _build_correlation_client("NORMAL_COMPLETE") + body = {"model": "gpt-4o", "messages": [{"role": "user", "content": "same body"}]} + + with caplog.at_level(logging.DEBUG, logger="moralstack.server.proxy"): + r_a = client.post("/v1/chat/completions", json=body, headers={"Authorization": "Bearer key-A"}) + r_b = client.post("/v1/chat/completions", json=body, headers={"Authorization": "Bearer key-B"}) + + cid_a = r_a.headers.get("X-Moralstack-Conversation-Id") + cid_b = r_b.headers.get("X-Moralstack-Conversation-Id") + assert cid_a is not None and cid_b is not None + assert cid_a != cid_b + + expected_digest_a = hmac.new(b"dummy-test-secret", b"key-A", hashlib.sha256).hexdigest() + expected_digest_b = hmac.new(b"dummy-test-secret", b"key-B", hashlib.sha256).hexdigest() + log_text = caplog.text + assert "key-A" not in log_text + assert "key-B" not in log_text + assert "Bearer" not in log_text + assert expected_digest_a not in log_text + assert expected_digest_b not in log_text + + safe_headers = _collect_safe_headers(_FakeRequest({"Authorization": "Bearer key-A"})) + assert "key-A" not in str(safe_headers) + assert expected_digest_a not in str(safe_headers) + + def test_hmac_secret_unset_falls_back_to_anonymous(self, monkeypatch): + """Path B/C: secret unset -> B disabled, no hardcoded fallback. Two + different bearer tokens with an identical body collide onto the same + anonymous ("", hash) id.""" + monkeypatch.delenv("MORALSTACK_PRINCIPAL_HMAC_SECRET", raising=False) + client, _, _ = _build_correlation_client("NORMAL_COMPLETE") + body = {"model": "gpt-4o", "messages": [{"role": "user", "content": "collide via anonymous path"}]} + + r_a = client.post("/v1/chat/completions", json=body, headers={"Authorization": "Bearer key-A"}) + r_b = client.post("/v1/chat/completions", json=body, headers={"Authorization": "Bearer key-B"}) + + cid_a = r_a.headers.get("X-Moralstack-Conversation-Id") + cid_b = r_b.headers.get("X-Moralstack-Conversation-Id") + assert cid_a == cid_b + + def test_tenant_header_precedence_over_authorization(self, monkeypatch): + """(A) wins over (B): both headers present -> same id as tenant-header-only.""" + monkeypatch.setenv("MORALSTACK_PRINCIPAL_HMAC_SECRET", "dummy-test-secret") + client, _, _ = _build_correlation_client("NORMAL_COMPLETE") + body = {"model": "gpt-4o", "messages": [{"role": "user", "content": "tenant precedence body"}]} + + r_tenant_only = client.post("/v1/chat/completions", json=body, headers={"X-Moralstack-Tenant-Id": "tenant-Z"}) + r_tenant_and_auth = client.post( + "/v1/chat/completions", + json=body, + headers={"X-Moralstack-Tenant-Id": "tenant-Z", "Authorization": "Bearer some-token"}, + ) + + cid1 = r_tenant_only.headers.get("X-Moralstack-Conversation-Id") + cid2 = r_tenant_and_auth.headers.get("X-Moralstack-Conversation-Id") + assert cid1 == cid2 + + def test_conversation_id_header_still_overrides_correlation(self): + """Explicit header wins regardless of principal keying: differing tenant + headers across two calls with the same explicit conversation_id header + must not change the resolved id.""" + client, _, _ = _build_correlation_client("NORMAL_COMPLETE") + r1 = client.post( + "/v1/chat/completions", + json={"model": "gpt-4o", "messages": [{"role": "user", "content": "Q"}]}, + headers={"X-Moralstack-Conversation-Id": "explicit-conv-1", "X-Moralstack-Tenant-Id": "tenant-A"}, + ) + r2 = client.post( + "/v1/chat/completions", + json={"model": "gpt-4o", "messages": [{"role": "user", "content": "Q"}]}, + headers={"X-Moralstack-Conversation-Id": "explicit-conv-1", "X-Moralstack-Tenant-Id": "tenant-B"}, + ) + assert r1.headers.get("X-Moralstack-Conversation-Id") == "explicit-conv-1" + assert r2.headers.get("X-Moralstack-Conversation-Id") == "explicit-conv-1" + + def test_correlation_store_bound_does_not_break_request_on_eviction(self): + """Invariant #6: a tiny-bound store forcing eviction mid-run must never + break the request — every response is still 200.""" + tiny_store = ConversationCorrelationStore(max_entries=1) + client, _, _ = _build_correlation_client("NORMAL_COMPLETE", correlation_store=tiny_store) + + for i in range(5): + response = client.post( + "/v1/chat/completions", + json={"model": "gpt-4o", "messages": [{"role": "user", "content": f"Q{i}"}]}, + ) + assert response.status_code == 200 + assert tiny_store.size() <= 1 + + def test_correlation_store_raising_helper_does_not_break_request(self): + """Invariant #6 / blocking fix (diff review 2026-07-09): a correlation + store whose internal TTL/eviction helper raises must NOT escape into + the endpoint. Unlike the eviction test above (which exercises normal, + successful eviction), this forces `_get_id`/`_put_id` to raise so the + request path exercises `resolve()`'s best-effort swallow-and-mint + fallback ahead of the proxy's fail-closed `try` block.""" + raising_store = ConversationCorrelationStore() + raising_store._get_id = MagicMock(side_effect=RuntimeError("lineage lookup boom")) + raising_store._put_id = MagicMock(side_effect=RuntimeError("lineage insert boom")) + client, _, _ = _build_correlation_client("NORMAL_COMPLETE", correlation_store=raising_store) + + response = client.post( + "/v1/chat/completions", + json={"model": "gpt-4o", "messages": [{"role": "user", "content": "Q"}]}, + ) + assert response.status_code == 200 + cid = response.headers.get("X-Moralstack-Conversation-Id") + assert cid is not None and cid.startswith("msconv-") + + def test_explicit_header_then_correlation_chains_under_principal(self): + """Turn 1 uses the explicit conversation_id header under tenant A; turn 2 + drops the header (same tenant, extended history) and must chain back to + the header-supplied id: observe_completed_turn records (principal, + completed_hash) -> cid even when the cid came from the explicit header, + for a non-empty principal.""" + client, _, _ = _build_correlation_client("NORMAL_COMPLETE") + + r1 = client.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o", + "messages": [ + {"role": "system", "content": "Sys"}, + {"role": "user", "content": "First Q"}, + ], + }, + headers={"X-Moralstack-Conversation-Id": "explicit-chain-1", "X-Moralstack-Tenant-Id": "tenant-chain"}, + ) + assert r1.headers.get("X-Moralstack-Conversation-Id") == "explicit-chain-1" + + r2 = client.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o", + "messages": [ + {"role": "system", "content": "Sys"}, + {"role": "user", "content": "First Q"}, + {"role": "assistant", "content": "guidance"}, + {"role": "user", "content": "Second Q"}, + ], + }, + headers={"X-Moralstack-Tenant-Id": "tenant-chain"}, + ) + assert r2.headers.get("X-Moralstack-Conversation-Id") == "explicit-chain-1" + + +class TestExtractPrincipal: + """Direct unit tests of ``_extract_principal`` layering (A -> B -> C).""" + + def test_tenant_header_used_verbatim(self): + from moralstack.server.proxy import _extract_principal + + req = _FakeRequest({"X-Moralstack-Tenant-Id": "tenant-x"}) + assert _extract_principal(req) == "tenant-x" + + def test_whitespace_only_tenant_header_treated_as_absent(self, monkeypatch): + from moralstack.server.proxy import _extract_principal + + monkeypatch.delenv("MORALSTACK_PRINCIPAL_HMAC_SECRET", raising=False) + req = _FakeRequest({"X-Moralstack-Tenant-Id": " "}) + assert _extract_principal(req) == "" + + def test_no_tenant_bearer_with_secret_returns_hmac_digest(self, monkeypatch): + from moralstack.server.proxy import _extract_principal + + monkeypatch.setenv("MORALSTACK_PRINCIPAL_HMAC_SECRET", "test-secret") + req = _FakeRequest({"Authorization": "Bearer tok-123"}) + principal = _extract_principal(req) + assert principal + assert principal != "tok-123" + expected = hmac.new(b"test-secret", b"tok-123", hashlib.sha256).hexdigest() + assert principal == expected + + def test_no_tenant_bearer_without_secret_returns_empty(self, monkeypatch): + from moralstack.server.proxy import _extract_principal + + monkeypatch.delenv("MORALSTACK_PRINCIPAL_HMAC_SECRET", raising=False) + req = _FakeRequest({"Authorization": "Bearer tok-123"}) + assert _extract_principal(req) == "" + + def test_non_bearer_scheme_returns_empty(self, monkeypatch): + from moralstack.server.proxy import _extract_principal + + monkeypatch.setenv("MORALSTACK_PRINCIPAL_HMAC_SECRET", "test-secret") + req = _FakeRequest({"Authorization": "Basic dXNlcjpwYXNz"}) + assert _extract_principal(req) == "" + + def test_no_headers_returns_empty(self): + from moralstack.server.proxy import _extract_principal + + req = _FakeRequest({}) + assert _extract_principal(req) == "" + + +class TestCorrelationEnvWiring: + def test_env_wiring_valid_values_parsed(self, monkeypatch): + from moralstack.server.proxy import _resolve_correlation_store_env_config + + monkeypatch.setenv("MORALSTACK_CORRELATION_TTL_SECONDS", "42") + monkeypatch.setenv("MORALSTACK_CORRELATION_MAX_ENTRIES", "7") + ttl_seconds, max_entries = _resolve_correlation_store_env_config() + assert ttl_seconds == 42.0 + assert max_entries == 7 + + def test_env_wiring_garbage_and_out_of_range_falls_back_to_defaults(self, monkeypatch): + from moralstack.server.proxy import ( + DEFAULT_CORRELATION_TTL_SECONDS, + DEFAULT_MAX_CORRELATION_ENTRIES, + _resolve_correlation_store_env_config, + ) + + for bad_ttl in ("abc", "-5", "0"): + monkeypatch.setenv("MORALSTACK_CORRELATION_TTL_SECONDS", bad_ttl) + monkeypatch.delenv("MORALSTACK_CORRELATION_MAX_ENTRIES", raising=False) + ttl_seconds, max_entries = _resolve_correlation_store_env_config() + assert ttl_seconds == DEFAULT_CORRELATION_TTL_SECONDS + assert max_entries == DEFAULT_MAX_CORRELATION_ENTRIES + + monkeypatch.delenv("MORALSTACK_CORRELATION_TTL_SECONDS", raising=False) + for bad_max in ("abc", "-5", "0"): + monkeypatch.setenv("MORALSTACK_CORRELATION_MAX_ENTRIES", bad_max) + ttl_seconds, max_entries = _resolve_correlation_store_env_config() + assert ttl_seconds == DEFAULT_CORRELATION_TTL_SECONDS + assert max_entries == DEFAULT_MAX_CORRELATION_ENTRIES + + def test_create_app_never_raises_on_garbage_correlation_env(self, monkeypatch): + monkeypatch.setenv("MORALSTACK_CORRELATION_TTL_SECONDS", "-5") + monkeypatch.setenv("MORALSTACK_CORRELATION_MAX_ENTRIES", "0") + client, _, _ = _build_correlation_client("NORMAL_COMPLETE") + response = client.post( + "/v1/chat/completions", + json={"model": "gpt-4o", "messages": [{"role": "user", "content": "Q"}]}, + ) + assert response.status_code == 200 + + def test_env_wiring_valid_values_reach_default_store(self, monkeypatch): + """ + FIX 4 (diff review 2026-07-09): prove parsed valid env values actually + reach the ``ConversationCorrelationStore`` that ``create_app`` builds + by default (no ``correlation_store=`` injection) — not only that + ``_resolve_correlation_store_env_config()`` parses them in isolation. + """ + monkeypatch.setenv("MORALSTACK_CORRELATION_TTL_SECONDS", "42") + monkeypatch.setenv("MORALSTACK_CORRELATION_MAX_ENTRIES", "7") + client, _, _ = _build_correlation_client("NORMAL_COMPLETE") + + store = _extract_default_correlation_store(client.app) + assert store._ttl_seconds == 42.0 + assert store._max_entries == 7