diff --git a/CHANGELOG.md b/CHANGELOG.md index 6831d845d..1455615ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,38 @@ All notable changes to this project are documented in this file. Format follows For narrative release notes written for operators and product owners, see [RELEASE_NOTES.md](RELEASE_NOTES.md). +## [1.6.0] - 2026-07-15 + +Conversation-sharing and chat-reliability release. Large conversations can now be shared β€” their snapshots offload to a new S3 bucket instead of overflowing the 400 KB DynamoDB item limit. **Stop** now actually stops the server-side turn (distributed cancellation over the session lease), and a per-session single-flight lease plus restore-time history repair close a class of bugs where a tab switch or duplicate invocation could permanently brick a conversation. Web sources become removable and editor-manageable, and model RBAC grants written from the model admin page finally take effect. Requires a CDK deploy for the new shared-conversations S3 bucket and IAM grants. + +### πŸš€ Added + +- Share large conversations: snapshot bodies (messages + metadata) offload to a new private `shared-conversations` S3 bucket with a `body_ref` pointer in DynamoDB, so conversations over the 400 KB item limit can be shared. Legacy inline shares still read back with no migration; storage-unavailable surfaces as a friendly 503 (#658) +- Distributed turn cancellation β€” **Stop** now ends the running server-side turn instead of only the client stream. The app-api `user_stopped` endpoint stamps an owner-scoped `cancelRequestedFor` on the session lease; the inference-api heartbeat (tightened 30sβ†’10s) observes it and cooperatively tears down both the tool loop and the model stream, persisting the partial and releasing the lease. Shrinks the 409-on-resend window from a full turn to ~one heartbeat and halts wasted model/tool spend after Stop (#656) +- Remove a web source: `DELETE /assistants/{id}/web-sources/crawls/{crawl_id}` removes the crawl's sync policy, soft-deletes every page under its root URL (vector + S3 teardown via the existing background cleanup), then hard-deletes the crawl row. In-flight crawls are refused with a 409; a zombie 'running' crawl whose process died stays deletable. Edit-gated (owner or editor) (#648) + +### ✨ Improved + +- Editors (not just owners) can now start and view web crawls β€” `start_crawl`, `list_crawls`, and `get_crawl` route through the shared `_require_edit_permission` gate, so the "Add web content" button the SPA already renders for editors no longer 404s; viewers get a clean 403 (#650) + +### πŸ› Fixed + +- Model RBAC is now single-source-of-truth: the model admin page's role picker writes through to each role's `grantedModels` (mirroring tools/skills) instead of onto a dead `allowedAppRoles` field no access check read, so enabling a model for a role actually grants it. `allowedAppRoles` is derived on read; `can_access_model` and `filter_accessible_models` share one `_grants_access` predicate so a model can no longer be listed by the catalog yet denied on use. Removes the dead `POST /sync-roles` endpoint (#651) +- SSE chat stream stays open across tab switches β€” `@microsoft/fetch-event-source` `openWhenHidden` is now `true`, stopping the library from aborting and reopening the connection (a fresh `POST /invocations` for the same turn) on `visibilitychange`, which spawned a concurrent backend agent and corrupted tool-pairing history (#653) +- Restore-time tool-use/tool-result pairing repair β€” `TurnBasedSessionManager._repair_tool_pairing` unconditionally rebuilds a Bedrock-valid history on restore (one result turn per toolUse turn, duplicate/orphaned results dropped, same-role turns merged), recovering conversations already bricked by a "toolResult blocks exceed toolUse blocks" ValidationException. No-op on healthy history (#653) +- Reject duplicate concurrent turns β€” a per-session single-flight lease at the inference-api `/invocations` chokepoint (atomic conditional write on a `LEASE#{sid}` item) rejects a duplicate turn with 409 so two agent loops can't run against one Memory session; the SPA shows a soft "Already responding" notice instead of a hard error. Resume / max-tokens continuation force-acquire; fail-open on non-conflict DynamoDB errors (#655) +- Guard synthetic error persistence against role-alternation breaks β€” `persist_synthetic_messages` now drops a synthetic "⚠️ Something went wrong" turn that would land adjacent to a same-role turn, preventing the consecutive-assistant-message amplifier that turned one errored turn into a permanently bricked session (#654) +- Conversation-share operations no longer fail with a generic 500 β€” the app-api task role is granted DynamoDB access (including the `SessionShareIndex` GSI) on the shared-conversations table it was already wired to via env var, fixing `PutItem`/`Query` AccessDeniedException on share create/list (#657) + +### πŸ—οΈ Infrastructure + +- New private `shared-conversations` S3 bucket (SSE-S3, versioned) with an SSM param, a `PlatformComputeRefs` entry, the `SHARED_CONVERSATIONS_BUCKET_NAME` app-api env var, and an app-api-only `SharedConversationsBucketReadWrite` IAM grant (#658) +- `SharedConversationsAccess` added to the app-api `coreTables` grant list so the role gets the standard DynamoDB action set on the shared-conversations table and its GSIs (#657) + +### πŸ”§ CI/CD + +- Repointed `test_cache_savings.py` storage patch targets to `apis.shared.storage` (the accessor moved; `app_api.storage` is now an empty stub) and gated the compaction integration tests on an explicit `RUN_AGENTCORE_INTEGRATION_TESTS=1` opt-in so leaked env vars no longer make them run order-dependently against invalid credentials (#652) + ## [1.5.0] - 2026-07-13 Model-catalog and MCP-admin expansion, plus UI polish. Admins can now discover the tools behind OAuth-gated MCP servers (e.g. the GitHub remote MCP server) using their own vaulted 3LO token, two new curated model cards land (Claude Sonnet 5, GPT-5.4), and the max-output-tokens field goes optional so reasoning/Responses-API models without a fixed cap can be added. Alongside: sticky admin/settings sidebars, a redesigned 404 page, chat-scroll and sticky-nav fixes, a vitest flake fix, and a Docker curl security patch. No CDK deploy, no data migration, no breaking changes β€” ships via the backend and frontend pipelines. diff --git a/CLAUDE.MD b/CLAUDE.MD index be9f0ad89..8e1af5c83 100644 --- a/CLAUDE.MD +++ b/CLAUDE.MD @@ -97,6 +97,14 @@ The `inference-api` runs inside an AgentCore Runtime container. The runtime data If you're tempted to add to inference-api because of an existing route there (`converse_router`, `voice_router`), don't use them as templates without confirming their access path β€” they predate this rule and may rely on bypasses (API key, WebSocket upgrade, direct container reach in environments where the runtime isn't the only path). +### RBAC: the AppRole record is the source of truth + +A role grants a tool/model/skill by listing it in its own `grantedTools` / `grantedModels` / `grantedSkills` (or by granting the `*` wildcard, or inheriting from a parent). **That is the only thing access checks read.** The `allowedAppRoles` field on a resource (tool, model, skill) is a *derived, display-only* projection of those role records β€” never an independent grant. + +**Rule:** When an admin page offers a "which roles can use this?" picker on the resource, it must write *through* to each role's `granted*` list (`set_roles_for_tool`, `set_roles_for_model`) and derive the field back on read (`hydrate_model_roles`). Never persist a role list on the resource and expect it to grant anything: it will silently do nothing, and the resource page and the role page will disagree. + +Access decisions for models live in `apis/app_api/admin/services/model_access.py`. `can_access_model` and `filter_accessible_models` both delegate to a single `_grants_access` predicate β€” keep it that way. They previously diverged (one gated on `allowed_app_roles`, one didn't), so a model could be *listed* by the catalog and *denied* on use. + ### Auth dependency on app_api routes The SPA sends an httpOnly session cookie β€” not `Authorization: Bearer`. A route that declares a bare Bearer-only dependency on the SPA-facing surface causes a 401 β†’ centralized redirect loop the moment the SPA hits it. diff --git a/README.md b/README.md index 443604ebe..104f936e9 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ **An open-source, production-ready Generative AI platform for institutions** *Built by Boise State University, designed for everyone.* -[![Release](https://img.shields.io/badge/Release-v1.5.0-6366f1?style=flat&logo=github&logoColor=white)](RELEASE_NOTES.md) +[![Release](https://img.shields.io/badge/Release-v1.6.0-6366f1?style=flat&logo=github&logoColor=white)](RELEASE_NOTES.md) [![Nightly](https://github.com/Boise-State-Development/agentcore-public-stack/actions/workflows/nightly.yml/badge.svg)](https://github.com/Boise-State-Development/agentcore-public-stack/actions/workflows/nightly.yml) ![Python](https://img.shields.io/badge/Python-3.13+-3776AB?style=flat&logo=python&logoColor=white) @@ -296,7 +296,7 @@ agentcore-public-stack/ See [RELEASE_NOTES.md](RELEASE_NOTES.md) for the full changelog, including new features, bug fixes, platform upgrades, and deployment notes for each release. -**Current release:** v1.5.0 +**Current release:** v1.6.0 --- diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 859058fbf..73a013df0 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -1,7 +1,90 @@ -# Release Notes β€” v1.5.0 +# Release Notes β€” v1.6.0 -**Release Date:** July 13, 2026 -**Previous Release:** v1.4.0 (July 10, 2026) +**Release Date:** July 15, 2026 +**Previous Release:** v1.5.0 (July 13, 2026) + +--- + +> ⚠️ **Platform (CDK) deploy required.** This release adds a new `shared-conversations` S3 bucket and IAM grants, so it ships through `platform.yml` (CDK) **before** `backend.yml`. No data migration and no breaking changes β€” legacy inline shares keep working untouched. + +--- + +## Highlights + +v1.6.0 makes conversation sharing work for **large** conversations and makes chat **reliable under interruption**. Sharing a big conversation used to fail with a bare 500 because the whole snapshot was inlined into one DynamoDB item past the 400 KB limit; snapshots now offload to a new private S3 bucket, with legacy shares still readable and no migration. On the reliability side, **Stop now actually stops the server-side turn** β€” a distributed cancellation signal carried over the session lease tears down the running agent instead of letting it burn model and tool spend β€” and a **per-session single-flight lease** plus a **restore-time history repair** close a nasty class of bugs where a tab switch or duplicate invocation could permanently brick a conversation with a Bedrock tool-pairing error. Rounding it out: web sources are now **removable** and **editor-manageable**, and model RBAC grants written from the model admin page **finally take effect**. Operators must run a CDK deploy first for the new bucket and grants. + +## Share large conversations without hitting the DynamoDB item limit + +Sharing a large conversation failed with a generic 500 (observed in prod-ai as a `PutItem` ValidationException): `ShareService.create_share` inlined the full message list into a single DynamoDB item, exceeding the 400 KB item limit. Snapshot bodies now offload to a dedicated S3 bucket β€” mirroring the Memory Spaces / Artifacts / Skills offload pattern β€” while DynamoDB keeps only control fields plus a `body_ref` pointer. Reads fall back to inline for legacy shares, so existing shares keep working with no migration and the SPA contract is unchanged. + +### Backend + +- `shares/snapshot_store.py` β€” new `ShareSnapshotStore`: content-addressed S3 put/get/delete with SSE-S3 and dedupe. +- `shares/service.py` β€” `create_share` writes the body to S3 and stores a `body_ref` item; `_load_snapshot_body` reads from S3 or falls back to legacy inline items; revoke and session-cleanup best-effort delete the object. `ShareStorageUnavailableError` maps to a friendly 503 instead of a bare 500. + +### Infrastructure + +- `data/shared-conversations-construct.ts` β€” new private `shared-conversations` S3 bucket (SSE-S3, versioned) plus an SSM param; threaded to app-api via `PlatformComputeRefs` and the `SHARED_CONVERSATIONS_BUCKET_NAME` env var. +- `app-api/app-api-iam-grants.ts` β€” `SharedConversationsBucketReadWrite` grant (app-api only). + +### Test Coverage + +350+ lines: store round-trip / dedupe, a >400 KB regression, S3 and legacy reads, export-from-S3, revoke cleanup, and the storage-unavailable path. + +**Related fix (#657):** the shared-conversations *DynamoDB table* was wired into app-api by env var but never granted on the task role, so every share create/list already failed with `PutItem` / `Query` AccessDeniedException surfacing as a 500. Added `SharedConversationsAccess` to the app-api `coreTables` grant list (standard action set on the table and its `index/*` GSIs), with a synth regression test asserting the grant exists and carries no wildcard. + +## Stop actually stops the server turn + +A client abort β€” Stop, tab switch, dropped socket β€” does not propagate through the AgentCore Runtime data plane, so Stop was cosmetic server-side: the container ran the turn to completion, held the session lease, and burned model and tool spend. "Stop β†’ resend" then returned 409 until the prior turn finished on its own. This release reuses the single-flight lease (below) as a cross-container signalling channel so Stop ends the actual turn. + +### Backend + +- `apis/app_api/sessions/routes.py` β€” the `user_stopped` endpoint calls `request_session_cancel`, stamping `cancelRequestedFor=` on the lease item. Owner-scoped, so a stale Stop can't kill a later turn; best-effort, so it never fails the Stop. +- `apis/shared/sessions/session_lease.py` β€” the heartbeat (tightened 30sβ†’10s) renews with `ReturnValues=ALL_NEW` and, on `cancelRequestedFor == owner`, flips `session_manager.cancelled`. `acquire` clears any stale cancel marker on takeover. +- `main_agent/streaming/stream_coordinator.py` β€” two effects: the always-on `StopHook` cancels the next tool call; and a cooperative check at the top of the stream loop raises `_CooperativeStopSignal`, whose handler persists the partial via `_persist_interruption` (marked `user_stopped`), emits terminal SSE frames, and ends cleanly so the client closes and the lease releases. The cooperative arm is what ends a pure-chat turn, which has no tool boundary for `StopHook`. + +Net: the 409-on-resend window shrinks from a full turn to ~one heartbeat (10s), and post-Stop spend is halted. **Residual (documented):** an in-flight tool call finishes before cancel is seen, and already-generated Bedrock tokens are billed. + +## Duplicate-invocation hardening β€” no more bricked conversations + +Bedrock Converse rejects any history where a user turn's `toolResult` blocks don't exactly match the preceding assistant turn's `toolUse` blocks. A single such violation anywhere in a session's persisted history makes **every** subsequent turn fail, permanently bricking the conversation (this hit prod session `f761f59b`). The trigger was two agent loops running concurrently against one Memory session β€” spawned by a client reconnect or a duplicate `POST /invocations` the Runtime routed to a different container. This release closes the vector on three fronts. + +### Frontend + +- `chat-http.service.ts` / `preview-chat.service.ts` β€” set `openWhenHidden: true` on both `fetchEventSource` call sites so a single stream survives a tab switch instead of the library aborting and reopening it (a fresh `POST /invocations` for the same turn that bypassed the SPA's own double-submit guards). Also correct for long agentic turns (#653). + +### Backend + +- `apis/shared/sessions/session_lease.py` β€” new per-session single-flight lease: `acquire` / `renew` / `release` on a dedicated `PK=USER#{uid}, SK=LEASE#{sid}` item via an atomic conditional write. Owner-scoped renew/release; fail-open on any non-conflict DynamoDB error (#655). +- `apis/inference_api/chat/routes.py` β€” acquire the lease at turn-start and reject a duplicate with 409; resume and max-tokens continuation force-acquire (they re-enter an already-ended loop). Heartbeat renews while the turn streams; release in the generator finally and both except handlers (#655). +- `TurnBasedSessionManager._repair_tool_pairing` β€” an unconditional restore-time normalizer (sibling to `_strip_document_bytes`) that rebuilds a Bedrock-valid history: one matching result turn per `toolUse` turn (missing ones synthesized as errors), duplicate/orphaned result turns dropped, consecutive same-role turns merged. Identity no-op on healthy history, so it recovers already-corrupted sessions without touching clean ones (#653). +- `persist_synthetic_messages` β€” a centralized role-alternation guard drops a synthetic "⚠️ Something went wrong" turn that would land adjacent to a same-role turn, killing the consecutive-assistant-message amplifier that turned one errored turn into a permanent brick. Fixes the write side; `_repair_tool_pairing` masks it on the read side (#654). + +### Frontend (SPA) + +- The inference-api 409 is handled as a soft `AlreadyStreamingError` ("Already responding") notice rather than a hard "Chat Request Failed" toast; loading clears so the user can retry once the prior turn finishes (#655). + +## πŸ› Bug fixes + +- **Model RBAC grants from the model page now take effect (#651).** The model admin page and the role admin page wrote to two different, unlinked fields: enabling a model for a role on the model page wrote `allowedAppRoles` onto the model record β€” a field no access check ever read β€” so the grant silently did nothing and the role page still showed the model unchecked. The role record is now the single source of truth (matching tools and skills): the model form's picker writes through to each role's `grantedModels` (`set_roles_for_model`), `allowedAppRoles` is derived on read (`hydrate_model_roles`), and `can_access_model` / `filter_accessible_models` share one `_grants_access` predicate so a model can no longer be listed by the catalog yet denied on use. Removes the dead `POST /sync-roles` endpoint. +- **Editors can start and view web crawls (#650).** `start_crawl`, `list_crawls`, and `get_crawl` gated on the owner-keyed `get_assistant()`, which returns `None` for an editor-share holder β€” so the "Add web content" button the SPA already shows editors returned a 404. They now route through the shared `_require_edit_permission` gate (owner|editor), and a viewer gets a clean 403 instead of a misleading 404. + +## πŸ—οΈ Infrastructure + +- New private `shared-conversations` S3 bucket (SSE-S3, versioned), its SSM param, a `PlatformComputeRefs` entry, the `SHARED_CONVERSATIONS_BUCKET_NAME` app-api env var, and the app-api-only `SharedConversationsBucketReadWrite` grant (#658). +- `SharedConversationsAccess` added to the app-api `coreTables` DynamoDB grant list β€” standard action set on the shared-conversations table and its GSIs (#657). + +## πŸ”§ CI/CD + +- Fixed two pre-existing test failures on develop, both unrelated to the code under test: repointed `test_cache_savings.py`'s five `get_metadata_storage` patch targets to `apis.shared.storage` (the accessor moved; `app_api.storage` is now an empty stub), and re-gated the compaction integration tests on an explicit `RUN_AGENTCORE_INTEGRATION_TESTS=1` opt-in so a mid-suite env-var leak no longer makes them run order-dependently against invalid credentials. Full suite: 4771 passed, 6 skipped (#652). + +## πŸš€ Deployment notes + +1. **Deploy `platform.yml` (CDK) first.** This release adds the `shared-conversations` S3 bucket, its SSM param, and IAM grants (both the bucket read/write grant and the `SharedConversationsAccess` DynamoDB grant on the app-api role). The app-api container reads `SHARED_CONVERSATIONS_BUCKET_NAME` at runtime; deploying app-api before the platform stack would leave sharing large conversations broken. +2. **Then `backend.yml`** to ship the app-api / inference-api images, followed by **`frontend-deploy.yml`** for the SPA (the 409 "Already responding" handling, `openWhenHidden`, web-source delete UI, and model-role picker changes). +3. **No data migration.** Legacy inline shares read back unchanged; new shares offload to S3 automatically. No breaking API changes. + +--- --- diff --git a/VERSION b/VERSION index 3e1ad720b..dc1e644a1 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.5.0 \ No newline at end of file +1.6.0 diff --git a/backend/pyproject.toml b/backend/pyproject.toml index e6b648b19..df349ac98 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "agentcore-stack" -version = "1.5.0" +version = "1.6.0" requires-python = ">=3.10" description = "Multi-agent conversational AI system with AWS Bedrock AgentCore" readme = "README.md" diff --git a/backend/src/.env.example b/backend/src/.env.example index c8335ed14..623ec1192 100644 --- a/backend/src/.env.example +++ b/backend/src/.env.example @@ -323,6 +323,14 @@ DYNAMODB_USER_MENU_LINKS_TABLE_NAME= # Example: dev-boisestateai-v2-shared-conversations SHARED_CONVERSATIONS_TABLE_NAME= +# Shared-conversations snapshot-body S3 bucket. The share BODY (messages + +# metadata) is offloaded here because a long conversation exceeds DynamoDB's +# 400 KB item limit; only a pointer stays in the table. Leave empty locally +# (creating a share then returns a 503 "temporarily unavailable" rather than a +# 400 KB PutItem failure). Set to the CDK-created bucket name in deployed envs. +# Example: dev-boisestateai-v2-shared-conversations +SHARED_CONVERSATIONS_BUCKET_NAME= + # ============================================================================= # OAUTH PROVIDER MANAGEMENT (OPTIONAL) # ============================================================================= diff --git a/backend/src/agents/main_agent/config/constants.py b/backend/src/agents/main_agent/config/constants.py index d8f2f9a6e..c817f463f 100644 --- a/backend/src/agents/main_agent/config/constants.py +++ b/backend/src/agents/main_agent/config/constants.py @@ -28,6 +28,12 @@ class EnvVars: COMPACTION_PROTECTED_TURNS = "AGENTCORE_MEMORY_COMPACTION_PROTECTED_TURNS" COMPACTION_MAX_TOOL_CONTENT_LENGTH = "AGENTCORE_MEMORY_COMPACTION_MAX_TOOL_CONTENT_LENGTH" + # --- Restored-history repair --- + # Kill switch for the restore-time tool-pairing/alternation repair + # (`TurnBasedSessionManager._repair_tool_pairing`). Default ON; set to + # "false" to disable. See the method docstring for the failure it guards. + HISTORY_REPAIR_ENABLED = "AGENTCORE_MEMORY_HISTORY_REPAIR_ENABLED" + # --- DynamoDB Tables --- DYNAMODB_SESSIONS_METADATA_TABLE = "DYNAMODB_SESSIONS_METADATA_TABLE_NAME" DYNAMODB_QUOTA_TABLE = "DYNAMODB_QUOTA_TABLE" diff --git a/backend/src/agents/main_agent/session/persistence.py b/backend/src/agents/main_agent/session/persistence.py index 60170dd13..134366500 100644 --- a/backend/src/agents/main_agent/session/persistence.py +++ b/backend/src/agents/main_agent/session/persistence.py @@ -21,10 +21,17 @@ before the agent runs). Call sites in ``stream_coordinator`` and ``chat/routes.py`` repeat the high-level reasoning inline; if you need to revisit the invariant, start here. + +ROLE-ALTERNATION GUARD (also canonical here): +Passing ``last_persisted_role`` makes this helper drop any synthetic turn +that would create two consecutive same-role messages β€” the corruption that +bricks a session under Bedrock's strict alternation rule. Centralizing it +here protects every caller that writes an assistant turn after an error, +consistent with this module being the single source of truth. """ import logging -from typing import Any, List, Tuple +from typing import Any, List, Optional, Tuple from strands.types.content import Message from strands.types.session import SessionMessage @@ -40,6 +47,7 @@ def persist_synthetic_messages( messages: List[Tuple[str, str]], *, agent_id: str = "default", + last_persisted_role: Optional[str] = None, ) -> bool: """Write one or more synthetic ``(role, text)`` messages to a session. @@ -58,17 +66,69 @@ def persist_synthetic_messages( user turn has not been written yet. agent_id: AgentCore Memory agent_id. Defaults to ``"default"`` to match read paths in ``apis.shared.sessions.messages``. + last_persisted_role: Role of the message already at the tail of the + session (``"user"`` / ``"assistant"``), typically + ``agent.messages[-1]["role"]``. When provided, this helper drops + any synthetic turn that would land adjacent to a same-role turn, + preserving Bedrock's strict user/assistant alternation β€” see the + ROLE-ALTERNATION GUARD below. Pass ``None`` (the default) to + persist verbatim, e.g. the quota-exceeded path that writes a + fresh ``user`` + ``assistant`` pair onto an empty session. Returns: - ``True`` if all messages were written. ``False`` (with an ERROR - log) if the session manager has no ``create_message`` method β€” - the failure mode that previously went silent. + ``True`` if every message that survived alternation filtering was + written (including the case where the guard dropped all of them β€” + that is a successful no-op, not a failure). ``False`` (with an ERROR + log) if the session manager has no ``create_message`` method β€” the + failure mode that previously went silent. Raises: Whatever ``create_message`` raises is propagated. Callers wrap with their own try/except so the failure appears in logs at the call site rather than being swallowed here. """ + # ROLE-ALTERNATION GUARD (single source of truth for all callers). + # + # Bedrock Converse requires strict user/assistant alternation. TWO + # consecutive same-role turns ANYWHERE in stored history make EVERY + # subsequent turn on that session fail with a ValidationException, + # permanently bricking the conversation. The classic amplifier: a turn + # ends with a dangling assistant toolUse (or a prior synthetic error + # turn), an error fires, and the error handler appends ANOTHER assistant + # turn β€” assistant, assistant β€” which then fails the next turn, which + # persists yet another assistant error, and so on. + # + # When the caller knows the tail role (``last_persisted_role``), skip any + # synthetic turn that would sit next to a same-role turn. In practice + # this is the "error persisted after a dangling assistant turn" case: the + # synthetic assistant message is dropped and the error stays a live-only + # UI affordance for that turn β€” the same deliberate choice the max_tokens + # path already makes (see ``stream_coordinator``). Only the FIRST tuple can + # collide with the tail; the tuples within ``messages`` already alternate, + # so we thread ``prev_role`` through the loop to stay correct for any + # multi-message batch. + if last_persisted_role is not None: + filtered: List[Tuple[str, str]] = [] + prev_role = last_persisted_role + for role, text in messages: + if role == prev_role: + logger.warning( + f"Skipping synthetic {role} message for session " + f"{scrub_log(session_id)} to preserve role alternation " + f"(tail role is already {role})" + ) + continue + filtered.append((role, text)) + prev_role = role + messages = filtered + + if not messages: + logger.info( + f"No synthetic messages to persist to session {scrub_log(session_id)} " + f"after role-alternation filtering" + ) + return True + target_manager = next( ( m diff --git a/backend/src/agents/main_agent/session/tests/test_compaction_integration.py b/backend/src/agents/main_agent/session/tests/test_compaction_integration.py index 5421667b1..1bb2a93e9 100644 --- a/backend/src/agents/main_agent/session/tests/test_compaction_integration.py +++ b/backend/src/agents/main_agent/session/tests/test_compaction_integration.py @@ -20,10 +20,15 @@ import os import pytest -# Skip all tests in this module unless AWS integration env vars are set +# These hit real AWS services, so they must never run in the automated unit +# suite. Gate them behind an explicit opt-in rather than AGENTCORE_MEMORY_ID: +# that variable is present in any real environment (and a local backend/src/.env +# leaks it into the process mid-suite via the apis.app_api.main reload's +# load_dotenv(override=True)), which made these tests run order-dependently +# instead of skipping. Set RUN_AGENTCORE_INTEGRATION_TESTS=1 to run them. pytestmark = pytest.mark.skipif( - not os.environ.get("AGENTCORE_MEMORY_ID"), - reason="Integration test requires AGENTCORE_MEMORY_ID environment variable" + os.environ.get("RUN_AGENTCORE_INTEGRATION_TESTS", "").lower() not in ("1", "true"), + reason="Integration test requires RUN_AGENTCORE_INTEGRATION_TESTS=1 and real AWS services" ) import sys import uuid diff --git a/backend/src/agents/main_agent/session/tests/test_history_repair.py b/backend/src/agents/main_agent/session/tests/test_history_repair.py new file mode 100644 index 000000000..0b68b8f9d --- /dev/null +++ b/backend/src/agents/main_agent/session/tests/test_history_repair.py @@ -0,0 +1,246 @@ +"""Unit tests for restore-time tool-pairing / role-alternation repair. + +`TurnBasedSessionManager._repair_tool_pairing` is the safety net that keeps a +structurally-corrupt persisted history from bricking a session with a Bedrock +Converse ValidationException ("The number of toolResult blocks at messages.N +exceeds the number of toolUse blocks of previous turn."). + +The corruption shapes exercised here are the ones observed in a real bricked +production session (parallel tool calls interrupted by a user Stop): +duplicate tool-result turns, tool-result turns reordered away from their +tool-use turn (assistant/assistant/user/user), tool-results orphaned after a +synthetic error turn, and consecutive synthetic error turns. +""" + +from agents.main_agent.config.constants import EnvVars +from agents.main_agent.session.turn_based_session_manager import ( + TurnBasedSessionManager as T, +) + + +class _FakeConfig: + session_id = "test-session" + + +def _make_manager() -> T: + """A TurnBasedSessionManager built without the heavy AgentCore parent init. + + ``_repair_restored_history`` only reads ``self.config.session_id`` and calls + the ``_repair_tool_pairing`` classmethod, so an instance created via + ``__new__`` with just a fake config is sufficient β€” and avoids a subclass + whose ``__init__`` skips ``super().__init__``. + """ + manager = T.__new__(T) + manager.config = _FakeConfig() + return manager + + +class _FakeAgent: + def __init__(self, messages): + self.messages = messages + + +def _use(*ids): + return { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": t, "name": "search", "input": {}}} for t in ids], + } + + +def _res(*ids): + return { + "role": "user", + "content": [{"toolResult": {"toolUseId": t, "content": [{"text": "ok"}]}} for t in ids], + } + + +def _txt(role, text="hi"): + return {"role": role, "content": [{"text": text}]} + + +def _is_valid(messages): + """Assert the message list satisfies Bedrock's structural constraints: + strict role alternation, and every toolResult turn matches the toolUseIds + of the immediately-preceding turn exactly. + """ + for i, msg in enumerate(messages): + if i > 0 and messages[i - 1]["role"] == msg["role"]: + return False, f"consecutive role at {i}" + has_use, has_result = T._block_keys(msg) + if has_result: + prev_use = T._tool_use_ids(messages[i - 1]) if i > 0 else [] + if set(T._tool_result_ids(msg)) != set(prev_use): + return False, f"pairing mismatch at {i}" + return True, "valid" + + +class TestHealthyHistory: + def test_clean_history_is_noop_identity(self): + msgs = [_txt("user"), _use("a"), _res("a"), _txt("assistant")] + repaired, fixed = T._repair_tool_pairing(msgs) + assert fixed == 0 + assert repaired is msgs # identity: no rebuild for healthy history + + def test_clean_parallel_tools_is_noop(self): + msgs = [_txt("user"), _use("a", "b", "c"), _res("a", "b", "c"), _txt("assistant")] + _, fixed = T._repair_tool_pairing(msgs) + assert fixed == 0 + + +class TestDuplicateToolResults: + def test_duplicate_consecutive_result_turn_deduped(self): + # assistant USE x3, then the SAME results persisted twice. + msgs = [ + _txt("user"), + _use("a", "b", "c"), + _res("a", "b", "c"), + _res("a", "b", "c"), # duplicate + _txt("user", "next question"), + ] + repaired, fixed = T._repair_tool_pairing(msgs) + assert fixed > 0 + ok, why = _is_valid(repaired) + assert ok, why + # exactly one result turn survives + result_turns = [m for m in repaired if T._block_keys(m)[1]] + assert len(result_turns) == 1 + + +class TestReorderedParallelTools: + def test_assistant_assistant_user_user_reordered(self): + # Two tool-use turns back to back, then both result turns back to back. + msgs = [ + _txt("user"), + _use("a", "b", "c"), + _use("d", "e"), + _res("a", "b", "c"), + _res("d", "e"), + _txt("assistant", "done"), + ] + repaired, fixed = T._repair_tool_pairing(msgs) + assert fixed > 0 + ok, why = _is_valid(repaired) + assert ok, why + + +class TestOrphanedToolResults: + def test_orphan_result_after_assistant_text_dropped(self): + # A result turn re-injected after an assistant TEXT turn (no toolUse). + msgs = [ + _txt("user"), + _use("a", "b"), + _res("a", "b"), + _txt("assistant", "here is the answer"), + _res("a", "b"), # orphan: previous turn has no toolUse + _txt("user", "thanks"), + ] + repaired, fixed = T._repair_tool_pairing(msgs) + assert fixed > 0 + ok, why = _is_valid(repaired) + assert ok, why + + +class TestConsecutiveTextTurns: + def test_consecutive_synthetic_assistant_errors_merged(self): + msgs = [ + _txt("user"), + _txt("assistant", "⚠️ Something went wrong"), + _txt("assistant", "⚠️ Something went wrong"), + _txt("assistant", "final answer"), + ] + repaired, fixed = T._repair_tool_pairing(msgs) + assert fixed > 0 + ok, why = _is_valid(repaired) + assert ok, why + assert sum(1 for m in repaired if m["role"] == "assistant") == 1 + + def test_text_turn_merges_into_following_tooluse_turn(self): + # assistant text immediately followed by assistant tool-use. + msgs = [ + _txt("user"), + _txt("assistant", "let me look that up"), + _use("a", "b"), + _res("a", "b"), + _txt("assistant", "done"), + ] + repaired, fixed = T._repair_tool_pairing(msgs) + assert fixed > 0 + ok, why = _is_valid(repaired) + assert ok, why + # merged assistant turn keeps toolUse at the tail so its result follows + merged = repaired[1] + assert merged["role"] == "assistant" + assert "text" in merged["content"][0] + assert "toolUse" in merged["content"][-1] + + +class TestMissingResults: + def test_interrupted_tooluse_gets_synthesized_error_result(self): + # Mid-list tool-use whose results were never persisted. + msgs = [ + _txt("user"), + _use("a", "b"), # no result turn follows + _txt("assistant", "moving on"), + _txt("user", "ok"), + ] + repaired, fixed = T._repair_tool_pairing(msgs) + assert fixed > 0 + ok, why = _is_valid(repaired) + assert ok, why + # a synthesized error result exists for both ids + results = [b for m in repaired for b in m["content"] if "toolResult" in b] + assert {r["toolResult"]["toolUseId"] for r in results} == {"a", "b"} + assert all(r["toolResult"].get("status") == "error" for r in results) + + def test_trailing_tooluse_left_untouched(self): + # A tool-use turn as the LAST message is left for prompt-arrival + # handling (matching the SDK's own repair), not force-answered here. + msgs = [_txt("user"), _use("a")] + repaired, _ = T._repair_tool_pairing(msgs) + assert repaired[-1]["content"][-1].get("toolUse", {}).get("toolUseId") == "a" + + +class TestIdempotency: + def test_repair_is_idempotent(self): + msgs = [ + _txt("user"), + _use("a", "b", "c"), + _res("a", "b", "c"), + _res("a", "b", "c"), + _use("d", "e"), + _use("f"), + _res("d", "e"), + _res("f"), + _txt("assistant", "x"), + _txt("assistant", "y"), + ] + once, n1 = T._repair_tool_pairing(msgs) + assert n1 > 0 and _is_valid(once)[0] + twice, n2 = T._repair_tool_pairing(once) + assert n2 == 0 + assert _is_valid(twice)[0] + + +class TestRepairWrapper: + """Covers `_repair_restored_history`, the method wired into initialize().""" + + def test_wrapper_mutates_agent_messages(self, monkeypatch): + monkeypatch.delenv(EnvVars.HISTORY_REPAIR_ENABLED, raising=False) + agent = _FakeAgent([_txt("user"), _use("a"), _res("a"), _res("a")]) + _make_manager()._repair_restored_history(agent) + ok, why = _is_valid(agent.messages) + assert ok, why + + def test_kill_switch_disables_repair(self, monkeypatch): + monkeypatch.setenv(EnvVars.HISTORY_REPAIR_ENABLED, "false") + corrupt = [_txt("user"), _use("a"), _res("a"), _res("a")] + agent = _FakeAgent(corrupt) + _make_manager()._repair_restored_history(agent) + assert agent.messages is corrupt # untouched + + def test_wrapper_noop_on_healthy_history(self, monkeypatch): + monkeypatch.delenv(EnvVars.HISTORY_REPAIR_ENABLED, raising=False) + healthy = [_txt("user"), _use("a"), _res("a"), _txt("assistant")] + agent = _FakeAgent(healthy) + _make_manager()._repair_restored_history(agent) + assert agent.messages is healthy # identity preserved, no rebuild diff --git a/backend/src/agents/main_agent/session/tests/test_persistence.py b/backend/src/agents/main_agent/session/tests/test_persistence.py index 9d17921bb..18d02ba19 100644 --- a/backend/src/agents/main_agent/session/tests/test_persistence.py +++ b/backend/src/agents/main_agent/session/tests/test_persistence.py @@ -130,3 +130,101 @@ def create_message(self, *args: Any, **kwargs: Any) -> None: sm = _RaisingSessionManager() with pytest.raises(RuntimeError, match="rejected the write"): persist_synthetic_messages(sm, "sess-x", [("assistant", "boom")]) + + +# --------------------------------------------------------------------------- +# Role-alternation guard (last_persisted_role). +# +# Bedrock Converse requires strict user/assistant alternation. TWO consecutive +# same-role turns anywhere in stored history permanently brick the session. +# The synthetic-error paths append an assistant turn after a failure; if the +# session tail is ALREADY an assistant turn (a dangling toolUse or a prior +# synthetic error), that append is the corruption. The guard drops it. +# --------------------------------------------------------------------------- + + +def test_skips_assistant_write_when_tail_is_assistant(): + """The core fix: error persisted after a dangling assistant turn must NOT + create a second consecutive assistant message β€” it is dropped entirely.""" + sm = _RecordingSessionManager() + + ok = persist_synthetic_messages( + sm, + "sess-brick", + [("assistant", "⚠️ Something went wrong")], + last_persisted_role="assistant", + ) + + # No write happened, but it's a successful no-op (True), not a failure. + assert ok is True + assert sm.calls == [] + + +def test_persists_assistant_write_when_tail_is_user(): + """The common, healthy case: the user turn is the tail (persisted by the + hook at turn start), so appending the assistant error is valid.""" + sm = _RecordingSessionManager() + + ok = persist_synthetic_messages( + sm, + "sess-ok", + [("assistant", "the error explanation")], + last_persisted_role="user", + ) + + assert ok is True + assert len(sm.calls) == 1 + assert _extract(sm.calls[0]["message"]) == {"role": "assistant", "text": "the error explanation"} + + +def test_none_last_role_preserves_verbatim_behavior(): + """``last_persisted_role=None`` (the default) means "caller doesn't know the + tail" β†’ persist verbatim, matching pre-guard behavior. This is the + quota-exceeded path that writes a fresh user+assistant pair.""" + sm = _RecordingSessionManager() + + ok = persist_synthetic_messages( + sm, + "sess-quota", + [("user", "hi"), ("assistant", "quota exceeded")], + ) + + assert ok is True + assert len(sm.calls) == 2 + + +def test_guard_drops_leading_user_but_keeps_alternating_tail(): + """The guard tracks the running role through a multi-message batch: a + leading ``user`` that collides with a ``user`` tail is dropped, and the + following ``assistant`` (now valid after the user tail) is kept.""" + sm = _RecordingSessionManager() + + ok = persist_synthetic_messages( + sm, + "sess-multi", + [("user", "dup user"), ("assistant", "answer")], + last_persisted_role="user", + ) + + assert ok is True + assert len(sm.calls) == 1 + assert _extract(sm.calls[0]["message"]) == {"role": "assistant", "text": "answer"} + + +def test_guard_logs_when_skipping(caplog): + """A dropped synthetic turn is logged loudly so the skip is observable in + incident forensics, not silent.""" + sm = _RecordingSessionManager() + + with caplog.at_level("WARNING"): + persist_synthetic_messages( + sm, + "sess-log", + [("assistant", "dropped")], + last_persisted_role="assistant", + ) + + assert any( + "preserve role alternation" in rec.message and "sess-log" in rec.message + for rec in caplog.records + ), f"expected a skip warning, got: {[r.message for r in caplog.records]}" diff --git a/backend/src/agents/main_agent/session/turn_based_session_manager.py b/backend/src/agents/main_agent/session/turn_based_session_manager.py index d13a34cf5..c52238df1 100644 --- a/backend/src/agents/main_agent/session/turn_based_session_manager.py +++ b/backend/src/agents/main_agent/session/turn_based_session_manager.py @@ -214,6 +214,7 @@ def initialize(self, agent: "Agent", **kwargs: Any) -> None: logger.warning(f"Document byte stripping failed, continuing: {e}", exc_info=True) if not self.compaction_config or not self.compaction_config.enabled: + self._repair_restored_history(agent) return try: @@ -224,6 +225,14 @@ def initialize(self, agent: "Agent", **kwargs: Any) -> None: self._valid_cutoff_indices = [] self._all_messages_for_summary = [] + # Repair tool-use/tool-result pairing and role alternation on the FINAL + # restored list β€” after compaction slicing/truncation β€” so it is always + # the exact history sent to Bedrock. Runs unconditionally (independent of + # compaction), mirroring `_strip_document_bytes`, because the corruption + # it fixes originates on the write side and can exist regardless of + # whether compaction is enabled. + self._repair_restored_history(agent) + # ========================================================================= # Compaction β€” applied after SDK session restore # ========================================================================= @@ -784,6 +793,192 @@ def _strip_document_bytes(self, messages: List[Dict]) -> List[Dict]: return stripped_messages + # ========================================================================= + # Tool-pairing / role-alternation repair (restore-time safety net) + # ========================================================================= + + def _repair_restored_history(self, agent: "Agent") -> None: + """Repair tool-use/tool-result pairing on the restored history in place. + + Bedrock Converse rejects any request whose history violates its + structural rules β€” most relevantly: + "The number of toolResult blocks at messages.N.content exceeds the + number of toolUse blocks of previous turn." A single such violation + anywhere in stored history makes EVERY subsequent turn on that session + fail, permanently bricking the conversation. + + The corruption originates on the write side: a turn with parallel tool + calls in flight that is interrupted (user Stop / dropped connection) or + retried can persist duplicated tool-result messages, tool-result + messages reordered away from their tool-use turn (assistant, assistant, + user, user), or tool-result messages orphaned after a synthetic error + turn. The Strands SDK's own ``_fix_broken_tool_use`` only rebuilds the + single message immediately after each tool-use turn, so it does not + repair duplicates, reordering, or orphans β€” and it may not run at all on + the AgentCore Memory restore path. This is the unconditional safety net. + + Runs on the FINAL ``agent.messages`` (after compaction), mirroring + ``_strip_document_bytes``. Best-effort: any failure logs and leaves the + history untouched rather than breaking the turn. Kill switch: + ``AGENTCORE_MEMORY_HISTORY_REPAIR_ENABLED=false``. + """ + if os.environ.get(EnvVars.HISTORY_REPAIR_ENABLED, "").strip().lower() == "false": + return + try: + messages = agent.messages + if not messages: + return + repaired, fixed = self._repair_tool_pairing(messages) + if fixed > 0: + logger.warning( + "Restore repair: fixed %d tool-pairing/alternation " + "violation(s) in restored history (%d -> %d messages) " + "for session %s", + fixed, len(messages), len(repaired), self.config.session_id, + ) + agent.messages = repaired + except Exception as e: + logger.error(f"Restore history repair failed, using history as-is: {e}", exc_info=True) + + @staticmethod + def _block_keys(message: Dict) -> tuple: + """(has_tool_use, has_tool_result) for a message's content blocks.""" + has_use = has_result = False + for block in message.get("content", []) or []: + if isinstance(block, dict): + if "toolUse" in block: + has_use = True + elif "toolResult" in block: + has_result = True + return has_use, has_result + + @staticmethod + def _tool_use_ids(message: Dict) -> List[str]: + return [ + b["toolUse"]["toolUseId"] + for b in message.get("content", []) or [] + if isinstance(b, dict) and "toolUse" in b and "toolUseId" in b["toolUse"] + ] + + @staticmethod + def _tool_result_ids(message: Dict) -> List[str]: + return [ + b["toolResult"]["toolUseId"] + for b in message.get("content", []) or [] + if isinstance(b, dict) and "toolResult" in b and "toolUseId" in b["toolResult"] + ] + + @classmethod + def _count_pairing_violations(cls, messages: List[Dict]) -> int: + """Count Bedrock structural violations: consecutive same-role turns, + orphaned/mismatched toolResults, and tool-use turns not answered by the + next turn. Zero means the history is already valid β€” repair can no-op. + """ + violations = 0 + for i, msg in enumerate(messages): + role = msg.get("role") + prev = messages[i - 1] if i > 0 else None + if prev is not None and prev.get("role") == role: + violations += 1 + has_use, has_result = cls._block_keys(msg) + if has_result: + prev_use = cls._tool_use_ids(prev) if prev is not None else [] + if not prev_use or set(cls._tool_result_ids(msg)) != set(prev_use): + violations += 1 + if has_use and i + 1 < len(messages): + if set(cls._tool_use_ids(msg)) != set(cls._tool_result_ids(messages[i + 1])): + violations += 1 + return violations + + @classmethod + def _repair_tool_pairing(cls, messages: List[Dict]) -> tuple: + """Return ``(repaired_messages, violations_fixed)``. + + Rebuilds a Bedrock-valid history: every assistant tool-use turn is + immediately followed by exactly one user turn carrying one toolResult + per toolUseId (missing ones synthesized as errors), duplicate/orphaned + toolResult turns are dropped, and consecutive same-role turns are + merged. When the history is already valid the input list is returned + unchanged (identity) with a count of 0, so healthy sessions pay only a + scan. + """ + violations = cls._count_pairing_violations(messages) + if violations == 0: + return messages, 0 + + # Global toolUseId -> toolResult block (last occurrence wins: the most + # recent result for an id is the authoritative one). + result_by_id: Dict[str, Dict] = {} + for msg in messages: + for block in msg.get("content", []) or []: + if isinstance(block, dict) and "toolResult" in block: + tid = block["toolResult"].get("toolUseId") + if tid: + result_by_id[tid] = block + + def missing_result(tid: str) -> Dict: + return { + "toolResult": { + "toolUseId": tid, + "content": [{"text": "[tool result unavailable: the turn was interrupted]"}], + "status": "error", + } + } + + # Forward pass: emit each assistant tool-use turn followed by a freshly + # built result turn; drop standalone toolResult-only turns (their blocks + # are re-emitted from the map at the correct slot). + rebuilt: List[Dict] = [] + last_index = len(messages) - 1 + for i, msg in enumerate(messages): + has_use, has_result = cls._block_keys(msg) + + if has_result and not has_use: + # Standalone tool-result turn: keep only non-toolResult content + # (rare mixed text), drop the results themselves. + residual = [b for b in msg.get("content", []) if not (isinstance(b, dict) and "toolResult" in b)] + if residual: + rebuilt.append({"role": msg.get("role", "user"), "content": residual}) + continue + + if has_use and i < last_index: + rebuilt.append(msg) + use_ids = cls._tool_use_ids(msg) + rebuilt.append({ + "role": "user", + "content": [result_by_id.get(t, missing_result(t)) for t in use_ids], + }) + continue + + # Trailing tool-use turn (last message) or any non-tool turn: emit + # as-is. The trailing case is deliberately left for prompt-arrival + # handling, matching the SDK's own repair. + rebuilt.append(msg) + + # Merge consecutive same-role turns. Guard only on the PREVIOUS turn + # having no toolUse: a tool-use turn must stay immediately adjacent to + # its result turn, so it can never absorb a following turn β€” but a + # text turn may legitimately merge into a following tool-use turn + # (assistant text + toolUse in one turn is valid), keeping the toolUse + # at the tail so its result turn still follows. In a merged user turn, + # toolResult blocks come first. + merged: List[Dict] = [] + for msg in rebuilt: + prev = merged[-1] if merged else None + if ( + prev is not None + and prev.get("role") == msg.get("role") + and not cls._block_keys(prev)[0] # prev has no toolUse + ): + combined = list(prev.get("content", [])) + list(msg.get("content", [])) + if msg.get("role") == "user": + combined = sorted(combined, key=lambda b: 0 if isinstance(b, dict) and "toolResult" in b else 1) + prev["content"] = combined + else: + merged.append({"role": msg.get("role"), "content": list(msg.get("content", []))}) + + return merged, violations + def _truncate_tool_contents( self, messages: List[Dict], diff --git a/backend/src/agents/main_agent/streaming/stream_coordinator.py b/backend/src/agents/main_agent/streaming/stream_coordinator.py index a79bd5a70..4f4174f9b 100644 --- a/backend/src/agents/main_agent/streaming/stream_coordinator.py +++ b/backend/src/agents/main_agent/streaming/stream_coordinator.py @@ -23,6 +23,18 @@ logger = logging.getLogger(__name__) +class _CooperativeStopSignal(Exception): + """Internal: a user Stop was observed mid-stream (cooperative cancellation). + + Raised from the stream loop when ``session_manager.cancelled`` is set, and + caught in ``stream_response`` to persist the partial + end the stream + cleanly. A plain ``Exception`` (not ``CancelledError``) so it never reaches + the ASGI server as a spurious task cancellation when the client is still + connected, and so the generic ``except Exception`` error arm can't mistake + a deliberate stop for a failure. + """ + + class StreamCoordinator: """Coordinates streaming lifecycle for agent responses""" @@ -165,6 +177,26 @@ async def stream_response( # Process through new stream processor and format as SSE async for event in process_agent_stream(agent_stream): + # Cooperative stop. A user Stop arms a cancel on the session's + # single-flight lease; the inference-api heartbeat observes it + # and flips ``session_manager.cancelled``. Because a client + # abort does not propagate through the AgentCore Runtime data + # plane, this in-loop check is what actually ends a token-only + # turn β€” StopHook only fires at tool boundaries, of which a + # pure-chat turn has none. Raise into the dedicated stop arm + # below, which persists the partial the user already saw, marks + # the turn interrupted, and ends the stream cleanly so the lease + # releases and the user's resend isn't rejected. Checked before + # yielding so no further tokens are emitted once stop is seen; + # we then abandon the (suspended) agent stream, which cannot + # progress or write to Memory without a consumer. + if getattr(session_manager, "cancelled", False): + logger.info( + "Cooperative stop observed mid-stream for session %s; ending turn", + session_id, + ) + raise _CooperativeStopSignal() + # Track when new assistant messages start (to associate metadata with them) if event.get("type") == "message_start": role = event.get("data", {}).get("role") @@ -669,6 +701,16 @@ async def stream_response( # content_block_delta below yields to the SSE # stream. Persisting it keeps live and # refresh-hydrated views in sync. + # + # Alternation guard: if the turn already ended with a + # dangling assistant turn (an in-flight toolUse, or a + # prior synthetic error), appending another assistant + # message would create two consecutive assistant turns + # and brick the session under Bedrock's strict + # alternation rule. Passing the tail role makes + # persist_synthetic_messages drop the write in that case + # (the error stays a live-only UI affordance for this + # turn, mirroring the max_tokens path above). try: from agents.main_agent.session.persistence import persist_synthetic_messages from agents.main_agent.session.session_factory import SessionFactory @@ -678,6 +720,7 @@ async def stream_response( persist_session_manager, session_id, [("assistant", conv_error_event.message)], + last_persisted_role=self._last_persisted_role(agent), ) except Exception as persist_error: logger.error(f"Failed to persist intercepted error to session: {persist_error}", exc_info=True) @@ -919,6 +962,34 @@ async def stream_response( except Exception as e: logger.error(f"Failed to store user displayText: {e}", exc_info=True) + except _CooperativeStopSignal: + # Deliberate user Stop observed mid-stream (see the in-loop check). + # Unlike the CancelledError/GeneratorExit backstop below β€” which + # fires only when a client disconnect actually reaches this process + # β€” this path runs even when the client is still connected, because + # the stop was signalled out-of-band via the lease. Persist the + # partial the user already saw and mark the turn interrupted + # (user_stopped), then end the SSE stream cleanly rather than + # re-raising, so a still-connected client sees a proper close and + # the route's finally releases the lease. + await self._persist_interruption( + agent=agent, + session_id=session_id, + user_id=user_id, + partial_text="".join(assistant_text_acc), + main_agent_wrapper=main_agent_wrapper, + accumulated_metadata=accumulated_metadata, + initial_message_count=initial_message_count, + current_assistant_message_index=current_assistant_message_index, + stream_start_time=stream_start_time, + first_token_time=first_token_time, + reason="user_stopped", + ) + # Terminal frames so any still-connected client ends cleanly; the + # SPA already stopped rendering on Stop, so this is belt-and-braces. + yield 'event: message_stop\ndata: {"stopReason": "stopped"}\n\n' + yield "event: done\ndata: {}\n\n" + # No re-raise: the turn ended on purpose. except (asyncio.CancelledError, GeneratorExit): # Client interruption: Stop click, page refresh, or dropped # socket. Depending on where the generator is suspended when the @@ -983,7 +1054,10 @@ async def stream_response( # Persist ONLY the assistant turn. Same reasoning as the # AGENT_ERROR path above β€” the user turn was already persisted - # at turn start by Strands' MessageAddedEvent hook. + # at turn start by Strands' MessageAddedEvent hook. The same + # alternation guard applies: skip the write if the tail is + # already an assistant turn so we never append a second + # consecutive assistant message. try: from agents.main_agent.session.persistence import persist_synthetic_messages from agents.main_agent.session.session_factory import SessionFactory @@ -993,6 +1067,7 @@ async def stream_response( persist_session_manager, session_id, [("assistant", error_event.message)], + last_persisted_role=self._last_persisted_role(agent), ) except Exception as persist_error: logger.error(f"Failed to persist stream error to session: {persist_error}") @@ -1014,6 +1089,29 @@ async def stream_response( "MCP Apps broker unsubscribe failed", exc_info=True ) + @staticmethod + def _last_persisted_role(agent: Any) -> Optional[str]: + """Role of the message at the tail of ``agent.messages``, or ``None``. + + Strands' ``MessageAddedEvent``/``append_message`` hook persists each + completed message to AgentCore Memory as it lands, so the live + ``agent.messages`` list mirrors the persisted session tail. The + synthetic-error persistence paths pass this to + ``persist_synthetic_messages`` so it can drop a write that would create + two consecutive same-role turns (the corruption that bricks a session + under Bedrock's strict alternation rule). Reads the in-memory list + rather than round-tripping AgentCore Memory β€” the same 80-250ms query + the coordinator deliberately avoids elsewhere. Best-effort: any failure + returns ``None`` so the caller persists verbatim (prior behavior). + """ + try: + messages = getattr(agent, "messages", None) or [] + if messages: + return messages[-1].get("role") + except Exception: # noqa: BLE001 - guard is best-effort + logger.debug("Could not read agent.messages tail for alternation guard", exc_info=True) + return None + async def _persist_interruption( self, agent: Any, @@ -1026,10 +1124,17 @@ async def _persist_interruption( current_assistant_message_index: int = -1, stream_start_time: Optional[float] = None, first_token_time: Optional[float] = None, + reason: str = "connection_lost", ) -> None: """Persist the in-flight partial assistant turn + an interrupted marker when a turn is torn down mid-stream. + ``reason`` records why: the default ``connection_lost`` is the + disconnect backstop (conditional, never downgrades a stronger + ``user_stopped`` the client beacon may have written); the cooperative + Stop arm passes ``user_stopped`` so the marker is correct even if that + beacon never landed. + Runs from inside the ``except (CancelledError, GeneratorExit)`` arm, i.e. while the request task is being torn down. Any bare ``await`` here would itself be cancelled before it completes, so the work is @@ -1044,13 +1149,19 @@ async def _persist_interruption( reasoning as the error paths; canonical reference in ``session/persistence.py``). - Role-alternation repair: when the interruption lands before any token - of the in-flight message streamed (empty partial), a minimal - placeholder assistant turn is persisted ONLY if the last committed - message is a user turn β€” that is the dangling userβ†’user case the - placeholder exists to fix. On continuation/resume turns (history tail - is an assistant message) or a pre-turn cancellation nothing needs - repair, so no synthetic write happens and only the marker is set. + Role-alternation repair: a synthetic assistant turn is persisted only + when it keeps user/assistant alternation valid β€” i.e. the last + committed message is NOT itself an assistant turn. This covers two + cases at once: + * empty partial + dangling user tail β†’ persist a minimal placeholder + so the orphan user turn is answered (the userβ†’user repair); + * non-empty partial + assistant tail (an interrupted continuation/ + resume, where the tail is the message being extended) β†’ SKIP, so we + don't append a second consecutive assistant turn and brick the + session. The partial stays a live-only affordance for that turn. + Whenever nothing is persisted, only the marker is set. The write itself + also passes ``last_persisted_role`` to ``persist_synthetic_messages`` so + the centralized alternation guard is the single enforcement point. """ async def _do() -> None: text = partial_text.strip() @@ -1063,7 +1174,13 @@ async def _do() -> None: logger.debug("Could not read agent.messages tail", exc_info=True) message = text if text else "[Response interrupted before any content was generated]" - should_persist = bool(text) or last_role == "user" + # Persist the synthetic assistant turn only when it preserves + # alternation: never when the tail is already an assistant turn + # (an interrupted continuation/resume β€” appending here would create + # consecutive assistant turns and brick the session). Otherwise a + # non-empty partial is worth persisting, and an empty partial is + # persisted only to answer a dangling user turn. + should_persist = (bool(text) or last_role == "user") and last_role != "assistant" if should_persist: try: from agents.main_agent.session.persistence import persist_synthetic_messages @@ -1076,6 +1193,7 @@ async def _do() -> None: persist_session_manager, session_id, [("assistant", message)], + last_persisted_role=last_role, ) except Exception as persist_error: logger.error( @@ -1084,16 +1202,16 @@ async def _do() -> None: ) else: logger.info( - "Interruption with no in-flight partial and non-user history tail " - "(role=%s) for session %s β€” marker only, no synthetic write", - last_role, session_id, + "Interruption for session %s β€” marker only, no synthetic write " + "(in-flight partial present=%s, history tail role=%s)", + session_id, bool(text), last_role, ) try: from apis.shared.sessions.metadata import set_interrupted_turn await set_interrupted_turn( - session_id, user_id, reason="connection_lost", source="cancellation" + session_id, user_id, reason=reason, source="cancellation" ) except Exception as marker_error: logger.error( diff --git a/backend/src/apis/app_api/admin/routes.py b/backend/src/apis/app_api/admin/routes.py index 0b451f3a5..5b9244f88 100644 --- a/backend/src/apis/app_api/admin/routes.py +++ b/backend/src/apis/app_api/admin/routes.py @@ -5,7 +5,7 @@ """ from fastapi import APIRouter, HTTPException, Depends, Query, status -from typing import Literal, Optional +from typing import List, Literal, Optional import logging import os import re @@ -27,6 +27,7 @@ ManagedModelCreate, ManagedModelUpdate, ManagedModel, + ModelRoleAssignment, ) from apis.shared.auth import User, require_admin from apis.shared.feature_flags import skills_enabled @@ -37,6 +38,7 @@ update_managed_model, delete_managed_model, ) +from .services.model_roles import get_model_role_service logger = logging.getLogger(__name__) @@ -531,9 +533,13 @@ async def list_managed_models_endpoint( try: models = await list_managed_models(user_roles=None) # None = no role filtering + # Derive each model's role access from the AppRole records (one role query + # for the whole catalog), so the admin list reflects real grants. + await get_model_role_service().hydrate_model_roles(models) + # Convert ManagedModel instances to dicts for Pydantic v2 validation models_dict = [model.model_dump(by_alias=True) for model in models] - + return ManagedModelsListResponse( models=models_dict, total_count=len(models), @@ -576,10 +582,20 @@ async def create_managed_model_endpoint( try: model = await create_managed_model(model_data) + + # Grant the model to the requested roles. This is the write that actually + # controls access β€” the role record is the source of truth β€” after which + # we derive the role fields back onto the response. + role_service = get_model_role_service() + await role_service.set_roles_for_model( + model.model_id, model_data.allowed_app_roles, admin_user + ) + await role_service.hydrate_model_roles([model]) + return model except ValueError as e: - # Model already exists + # Model already exists, or an unknown AppRole was named logger.warning("Model creation failed") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -626,6 +642,10 @@ async def get_managed_model_endpoint( detail=f"Model with ID '{model_id}' not found" ) + # Derive role access from the AppRole records so the edit form shows the + # grants that are actually in effect. + await get_model_role_service().hydrate_model_roles([model]) + return model except HTTPException: @@ -668,6 +688,16 @@ async def update_managed_model_endpoint( logger.info("Admin updating enabled model") try: + # Capture the provider model id before the update: roles key their grants + # on it, so a rename has to move those grants rather than orphan them. + existing = await get_managed_model(model_id) + if not existing: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Model with ID '{model_id}' not found" + ) + previous_model_id = existing.model_id + model = await update_managed_model(model_id, updates) if not model: @@ -676,10 +706,30 @@ async def update_managed_model_endpoint( detail=f"Model with ID '{model_id}' not found" ) + role_service = get_model_role_service() + + # `None` means the caller didn't touch role access β€” but a rename still + # has to carry the existing direct grants over to the new model id. + requested_roles = updates.allowed_app_roles + renamed = model.model_id != previous_model_id + if requested_roles is None and renamed: + current = await role_service.get_roles_for_model(previous_model_id) + requested_roles = [a.role_id for a in current if a.grant_type == "direct"] + + if requested_roles is not None: + await role_service.set_roles_for_model( + model.model_id, + requested_roles, + admin_user, + previous_model_id=previous_model_id, + ) + + await role_service.hydrate_model_roles([model]) + return model except ValueError as e: - # Duplicate modelId or other validation error + # Duplicate modelId, or an unknown AppRole was named logger.warning("Model update failed") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -719,6 +769,10 @@ async def delete_managed_model_endpoint( logger.info("Admin deleting enabled model") try: + # Read the provider model id before deleting β€” roles key their grants on + # it, and we need to strip those so no role keeps granting a dead model. + existing = await get_managed_model(model_id) + deleted = await delete_managed_model(model_id) if not deleted: @@ -727,6 +781,11 @@ async def delete_managed_model_endpoint( detail=f"Model with ID '{model_id}' not found" ) + if existing: + await get_model_role_service().revoke_model_from_all_roles( + existing.model_id, admin_user + ) + return None except HTTPException: @@ -739,38 +798,39 @@ async def delete_managed_model_endpoint( ) -@router.post("/managed-models/{model_id}/sync-roles", response_model=ManagedModel) -async def sync_model_roles( +@router.get("/managed-models/{model_id}/roles", response_model=List[ModelRoleAssignment]) +async def get_managed_model_roles( model_id: str, admin_user: User = Depends(require_admin), ): """ - Sync a model's allowedAppRoles with the AppRole system. + List every AppRole that grants access to a model, and how. - This endpoint updates the model's allowedAppRoles based on which AppRoles - have this model in their granted_models list. It ensures bidirectional - consistency between models and roles. + Each assignment is tagged `direct` (the role lists the model in its + grantedModels), `wildcard` (the role grants '*'), or `inherited` (a parent + role grants it). Only `direct` grants are editable from the model form β€” + the others are changed by editing the role. - Requires system administrator access. + Replaces the old POST /sync-roles endpoint: role records are now the single + source of truth, so there is nothing left to reconcile. Args: - model_id: Model identifier - admin_user: Authenticated system admin user (injected by dependency) + model_id: Model identifier (the record's internal id) + admin_user: Authenticated admin user (injected by dependency) Returns: - ManagedModel: Updated model with synced allowedAppRoles + List[ModelRoleAssignment] Raises: HTTPException: - 401 if not authenticated - - 403 if user lacks system admin role + - 403 if user lacks admin role - 404 if model not found - 500 if server error """ - logger.info("Admin syncing roles for model") + logger.info("Admin listing roles for model") try: - # Get the model model = await get_managed_model(model_id) if not model: raise HTTPException( @@ -778,47 +838,15 @@ async def sync_model_roles( detail=f"Model with ID '{model_id}' not found" ) - # Import here to avoid circular imports - from apis.shared.rbac.admin_service import get_app_role_admin_service - - # Get all roles and find which ones grant access to this model - admin_service = get_app_role_admin_service() - all_roles = await admin_service.list_roles(enabled_only=False) - - # Find roles that grant access to this model - granting_roles = [] - for role in all_roles: - if model.model_id in role.granted_models: - granting_roles.append(role.role_id) - # Also check effective_permissions in case inheritance grants access - if role.effective_permissions and model.model_id in role.effective_permissions.models: - if role.role_id not in granting_roles: - granting_roles.append(role.role_id) - - # Update the model's allowed_app_roles - from apis.shared.models.models import ManagedModelUpdate - updates = ManagedModelUpdate(allowed_app_roles=granting_roles) - updated_model = await update_managed_model(model_id, updates) - - if not updated_model: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail="Failed to update model after computing roles" - ) - - logger.info( - "βœ… Synced model allowedAppRoles" - ) - - return updated_model + return await get_model_role_service().get_roles_for_model(model.model_id) except HTTPException: raise except Exception as e: - logger.error("Unexpected error syncing model roles", exc_info=True) + logger.error("Unexpected error listing model roles", exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Error syncing model roles: {str(e)}" + detail=f"Error listing model roles: {str(e)}" ) diff --git a/backend/src/apis/app_api/admin/services/model_access.py b/backend/src/apis/app_api/admin/services/model_access.py index 5a8d7c49e..f5dcacf50 100644 --- a/backend/src/apis/app_api/admin/services/model_access.py +++ b/backend/src/apis/app_api/admin/services/model_access.py @@ -4,14 +4,21 @@ supporting both the new AppRole system and legacy JWT role-based access. During the transition period, access is granted if the user matches EITHER: -1. AppRole-based access (via allowed_app_roles) -2. Legacy JWT role-based access (via available_to_roles) +1. AppRole-based access β€” the model is in the user's resolved ``permissions.models`` + (i.e. some role of theirs lists it in ``grantedModels``, or grants ``*``) +2. Legacy JWT role-based access (via ``available_to_roles``) Once migration is complete, the legacy JWT role check can be removed. + +Note ``ManagedModel.allowed_app_roles`` is NOT an input to either check: it is a +field derived *from* the role records for display (see :mod:`.model_roles`). +Access is decided by the roles alone. Both entry points below funnel through +``_grants_access`` so a single-model check and a catalog filter can never +disagree about the same model. """ import logging -from typing import List, Optional +from typing import List, Optional, Set from apis.shared.auth.models import User from apis.shared.rbac.service import AppRoleService, get_app_role_service @@ -40,15 +47,54 @@ def app_role_service(self) -> AppRoleService: self._app_role_service = get_app_role_service() return self._app_role_service + async def _resolve_model_permissions(self, user: User) -> Set[str]: + """ + Resolve the set of model ids the user's AppRoles grant (may contain '*'). + + Returns an empty set if permission resolution fails, so the caller falls + through to the legacy JWT check rather than erroring. + """ + try: + permissions = await self.app_role_service.resolve_user_permissions(user) + return set(permissions.models) + except Exception as e: + # JUSTIFICATION: AppRole permission resolution failures should not block access checks. + # We fall back to legacy JWT role checking to maintain system availability during + # the AppRole migration period. This ensures users can still access models even if + # the AppRole system has issues. We log the error for monitoring. + logger.warning( + f"Error resolving AppRole permissions for {user.email} (falling back to JWT roles): {e}", + exc_info=True + ) + return set() + + @staticmethod + def _grants_access( + model: ManagedModel, model_permissions: Set[str], user_roles: Set[str] + ) -> bool: + """ + Decide whether a user may use a model. The single access rule. + + Both ``can_access_model`` and ``filter_accessible_models`` delegate here, + so a model is never listed by one and denied by the other. + """ + if not model.enabled: + return False + + # AppRole-based access: a role of the user's grants this model (or all models). + if "*" in model_permissions or model.model_id in model_permissions: + return True + + # Legacy JWT role-based access (deprecated). + if model.available_to_roles and user_roles.intersection(model.available_to_roles): + return True + + return False + async def can_access_model(self, user: User, model: ManagedModel) -> bool: """ Check if a user can access a specific model. - Access is granted if ANY of the following is true: - 1. Model has allowed_app_roles AND user has matching AppRole permissions - 2. Model has available_to_roles AND user has matching JWT role - 3. Neither field is set (model is unrestricted - should not happen in practice) - Args: user: Authenticated user model: ManagedModel to check access for @@ -56,49 +102,17 @@ async def can_access_model(self, user: User, model: ManagedModel) -> bool: Returns: True if user can access the model, False otherwise """ + # Short-circuit before paying for a permission resolve. if not model.enabled: return False - # Check AppRole-based access first (new system) - if model.allowed_app_roles: - try: - permissions = await self.app_role_service.resolve_user_permissions(user) - - # Wildcard grants access to all models - if "*" in permissions.models: - logger.debug( - f"User {user.email} has wildcard model access via AppRole" - ) - return True - - # Check if user has access to this specific model - if model.model_id in permissions.models: - logger.debug( - f"User {user.email} has AppRole access to model {model.model_id}" - ) - return True - - except Exception as e: - # JUSTIFICATION: AppRole permission resolution failures should not block access checks. - # We fall back to legacy JWT role checking to maintain system availability during - # the AppRole migration period. This ensures users can still access models even if - # the AppRole system has issues. We log the error for monitoring. - logger.warning( - f"Error checking AppRole permissions for {user.email} (falling back to JWT roles): {e}", - exc_info=True - ) - - # Check legacy JWT role-based access (deprecated) - if model.available_to_roles: - user_roles = user.roles or [] - if any(role in model.available_to_roles for role in user_roles): - logger.debug( - f"User {user.email} has legacy JWT role access to model {model.model_id}" - ) - return True - - # No access configured or user doesn't match any access rules - return False + model_permissions = await self._resolve_model_permissions(user) + allowed = self._grants_access(model, model_permissions, set(user.roles or [])) + + logger.debug( + f"User {user.email} access to model {model.model_id}: {allowed}" + ) + return allowed async def filter_accessible_models( self, user: User, models: List[ManagedModel] @@ -113,48 +127,15 @@ async def filter_accessible_models( Returns: List of ManagedModel objects the user can access """ - accessible = [] - - # Get user's AppRole permissions once (cached) - try: - permissions = await self.app_role_service.resolve_user_permissions(user) - has_wildcard = "*" in permissions.models - model_permissions = set(permissions.models) - except Exception as e: - # JUSTIFICATION: AppRole permission resolution failures should not block model filtering. - # We fall back to legacy JWT role checking to maintain system availability during - # the AppRole migration period. This ensures users can still see accessible models - # even if the AppRole system has issues. We log the error for monitoring. - logger.warning( - f"Error resolving AppRole permissions for {user.email} (using JWT roles only): {e}", - exc_info=True - ) - permissions = None - has_wildcard = False - model_permissions = set() - + # Resolve the user's AppRole permissions once (cached) for the whole catalog. + model_permissions = await self._resolve_model_permissions(user) user_roles = set(user.roles or []) - for model in models: - if not model.enabled: - continue - - # Check AppRole-based access (wildcard or specific model) - if has_wildcard or model.model_id in model_permissions: - accessible.append(model) - continue - - # Check if model has allowed_app_roles and user matches - if model.allowed_app_roles and permissions: - # This model requires AppRole access, but user didn't have it - # Still check legacy JWT roles as fallback - pass - - # Check legacy JWT role-based access - if model.available_to_roles: - if user_roles.intersection(model.available_to_roles): - accessible.append(model) - continue + accessible = [ + model + for model in models + if self._grants_access(model, model_permissions, user_roles) + ] logger.debug( f"Filtered {len(models)} models to {len(accessible)} accessible for {user.email}" diff --git a/backend/src/apis/app_api/admin/services/model_roles.py b/backend/src/apis/app_api/admin/services/model_roles.py new file mode 100644 index 000000000..c87c2ae11 --- /dev/null +++ b/backend/src/apis/app_api/admin/services/model_roles.py @@ -0,0 +1,259 @@ +"""Model Role Service + +Keeps model↔role permissions in one place: the AppRole record. + +A role grants a model by listing it in its own ``grantedModels`` (or by granting +the ``*`` wildcard, or by inheriting from a parent that does). That is the only +thing the access checks in :mod:`.model_access` ever read. + +The admin model form still presents a "which roles can use this model?" picker. +This service is what makes that picker honest: + +* :meth:`set_roles_for_model` writes the picker's selection *through* to each + role's ``grantedModels`` β€” the same shape as ``set_roles_for_tool`` in + ``app_api/tools/service.py``. +* :meth:`hydrate_model_roles` recomputes ``allowed_app_roles`` / + ``inherited_app_roles`` from the role records on read, so the model record + never stores a role list that can drift out of date. + +Historically ``allowedAppRoles`` was a stored, client-writable field that no +access check consulted, so editing it on the model page silently did nothing. +""" + +import logging +from typing import Dict, List, Optional, Set + +from apis.shared.auth.models import User +from apis.shared.models.models import ManagedModel, ModelRoleAssignment +from apis.shared.rbac.admin_service import ( + AppRoleAdminService, + get_app_role_admin_service, +) +from apis.shared.rbac.models import AppRole, AppRoleUpdate + +logger = logging.getLogger(__name__) + +WILDCARD = "*" + + +class ModelRoleService: + """Reads and writes model grants against the AppRole records.""" + + def __init__(self, app_role_admin_service: Optional[AppRoleAdminService] = None): + self._admin_service = app_role_admin_service + + @property + def admin_service(self) -> AppRoleAdminService: + """Lazy-load AppRoleAdminService to avoid circular imports.""" + if self._admin_service is None: + self._admin_service = get_app_role_admin_service() + return self._admin_service + + # ========================================================================= + # Read β€” derive a model's roles from the role records + # ========================================================================= + + async def get_roles_for_model(self, model_id: str) -> List[ModelRoleAssignment]: + """ + Get every AppRole that grants access to a model. + + Args: + model_id: The *provider* model id (e.g. ``anthropic.claude-...``), + which is what roles store in ``grantedModels`` β€” not the model + record's internal UUID. + + Returns: + One assignment per granting role, tagged direct / wildcard / inherited. + """ + roles = await self.admin_service.list_roles(enabled_only=False) + return self._assignments_for(model_id, roles) + + async def hydrate_model_roles(self, models: List[ManagedModel]) -> List[ManagedModel]: + """ + Populate the derived ``allowed_app_roles`` / ``inherited_app_roles`` on + each model from the role records. + + Loads the role list once and resolves every model against it in memory, + so hydrating a full catalog costs a single role query. + + Args: + models: Models to hydrate (mutated in place and returned). + + Returns: + The same list, with derived role fields populated. + """ + if not models: + return models + + roles = await self.admin_service.list_roles(enabled_only=False) + + for model in models: + assignments = self._assignments_for(model.model_id, roles) + model.allowed_app_roles = [ + a.role_id for a in assignments if a.grant_type == "direct" + ] + model.inherited_app_roles = [ + a.role_id for a in assignments if a.grant_type != "direct" + ] + + return models + + def _assignments_for( + self, model_id: str, roles: List[AppRole] + ) -> List[ModelRoleAssignment]: + """Classify how (if at all) each role grants ``model_id``.""" + by_id: Dict[str, AppRole] = {r.role_id: r for r in roles} + assignments: List[ModelRoleAssignment] = [] + + for role in roles: + granted = set(role.granted_models) + + if model_id in granted: + grant_type, inherited_from = "direct", None + elif WILDCARD in granted: + grant_type, inherited_from = "wildcard", None + else: + # Not granted directly β€” see whether a parent supplies it. + inherited_from = self._parent_granting(model_id, role, by_id) + if inherited_from is None: + continue + grant_type = "inherited" + + assignments.append( + ModelRoleAssignment( + role_id=role.role_id, + display_name=role.display_name, + grant_type=grant_type, + inherited_from=inherited_from, + enabled=role.enabled, + ) + ) + + return assignments + + @staticmethod + def _parent_granting( + model_id: str, role: AppRole, by_id: Dict[str, AppRole] + ) -> Optional[str]: + """Return the id of the first enabled parent role granting the model.""" + for parent_id in role.inherits_from: + parent = by_id.get(parent_id) + if not parent or not parent.enabled: + continue + parent_grants = set(parent.granted_models) + if model_id in parent_grants or WILDCARD in parent_grants: + return parent_id + return None + + # ========================================================================= + # Write β€” push the model form's role picker into the role records + # ========================================================================= + + async def set_roles_for_model( + self, + model_id: str, + app_role_ids: List[str], + admin: User, + previous_model_id: Optional[str] = None, + ) -> None: + """ + Make exactly ``app_role_ids`` grant this model directly. + + Adds the model to each named role's ``grantedModels`` and removes it from + any role that grants it directly but isn't named. Roles that grant the + model only via wildcard or inheritance are left alone β€” they aren't + "direct" grants, so they're neither added to nor stripped by this call. + + Args: + model_id: The provider model id roles store in ``grantedModels``. + app_role_ids: Roles that should grant the model directly. + admin: Admin performing the change (for the audit log). + previous_model_id: Set when an update renamed the model id, so grants + pointing at the old id are migrated rather than orphaned. + + Raises: + ValueError: If any named role does not exist. + """ + roles = await self.admin_service.list_roles(enabled_only=False) + by_id: Dict[str, AppRole] = {r.role_id: r for r in roles} + + requested: Set[str] = set(app_role_ids) + unknown = requested - by_id.keys() + if unknown: + raise ValueError(f"Unknown AppRole(s): {sorted(unknown)}") + + # A rename leaves grants pointing at the old id; drop them so the role's + # grantedModels doesn't accumulate a dangling entry. + stale_id = ( + previous_model_id + if previous_model_id and previous_model_id != model_id + else None + ) + + for role in roles: + granted = list(role.granted_models) + updated = [m for m in granted if m != stale_id] if stale_id else list(granted) + + should_grant = role.role_id in requested + grants_directly = model_id in updated + + if should_grant and not grants_directly: + # A wildcard role already covers every model; adding the explicit + # id would be redundant noise in its grantedModels. + if WILDCARD not in updated: + updated.append(model_id) + elif not should_grant and grants_directly: + updated = [m for m in updated if m != model_id] + + if updated != granted: + await self.admin_service.update_role( + role.role_id, AppRoleUpdate(granted_models=updated), admin + ) + + logger.info( + f"Admin {admin.email} set roles for model {model_id}", + extra={ + "event": "model_roles_updated", + "model_id": model_id, + "admin_user_id": admin.user_id, + "roles": sorted(requested), + }, + ) + + async def revoke_model_from_all_roles(self, model_id: str, admin: User) -> None: + """ + Strip a model from every role's ``grantedModels``. + + Called when a model is deleted so roles don't retain grants for a model + that no longer exists. + """ + roles = await self.admin_service.list_roles(enabled_only=False) + + for role in roles: + if model_id not in role.granted_models: + continue + remaining = [m for m in role.granted_models if m != model_id] + await self.admin_service.update_role( + role.role_id, AppRoleUpdate(granted_models=remaining), admin + ) + + logger.info( + f"Admin {admin.email} revoked model {model_id} from all roles", + extra={ + "event": "model_roles_revoked", + "model_id": model_id, + "admin_user_id": admin.user_id, + }, + ) + + +# Global service instance +_service_instance: Optional[ModelRoleService] = None + + +def get_model_role_service() -> ModelRoleService: + """Get or create the global ModelRoleService instance.""" + global _service_instance + if _service_instance is None: + _service_instance = ModelRoleService() + return _service_instance diff --git a/backend/src/apis/app_api/admin/services/tests/test_model_access.py b/backend/src/apis/app_api/admin/services/tests/test_model_access.py index 6dd803d0f..2cb576c21 100644 --- a/backend/src/apis/app_api/admin/services/tests/test_model_access.py +++ b/backend/src/apis/app_api/admin/services/tests/test_model_access.py @@ -392,3 +392,85 @@ async def test_filter_permissions_error_fallback( # Should still work via JWT role fallback assert len(result) == 1 assert result[0].model_id == "model-1" + + +class TestModelAccessIgnoresAllowedAppRoles: + """ + ``allowed_app_roles`` is a display-only field derived from the role records. + It must never influence an access decision, and the single-model check must + always agree with the catalog filter. + + Regression: a model granted via a role's ``grantedModels`` but with an empty + ``allowed_app_roles`` was *listed* by ``filter_accessible_models`` and yet + *denied* by ``can_access_model``, because the latter gated on the field. + """ + + @pytest.fixture + def mock_app_role_service(self): + return AsyncMock() + + @pytest.fixture + def service(self, mock_app_role_service): + return ModelAccessService(app_role_service=mock_app_role_service) + + def _grant(self, user, mock_app_role_service, models): + mock_app_role_service.resolve_user_permissions.return_value = ( + UserEffectivePermissions( + user_id=user.user_id, + app_roles=["staff"], + tools=[], + models=models, + quota_tier=None, + resolved_at=datetime.now(timezone.utc).isoformat() + "Z", + ) + ) + + @pytest.mark.asyncio + async def test_role_grant_with_empty_allowed_app_roles_is_accessible( + self, service, mock_app_role_service + ): + """A role grant alone is sufficient β€” the model's role list is irrelevant.""" + user = create_test_user(roles=["Staff"]) + model = create_test_model(model_id="claude-sonnet-5", allowed_app_roles=[]) + self._grant(user, mock_app_role_service, ["claude-sonnet-5"]) + + assert await service.can_access_model(user, model) is True + + @pytest.mark.asyncio + async def test_allowed_app_roles_alone_grants_nothing( + self, service, mock_app_role_service + ): + """ + Listing a role on the model must NOT grant access on its own. Only the + role's own grantedModels does β€” otherwise the model record becomes a + second, competing source of truth. + """ + user = create_test_user(roles=["Staff"]) + model = create_test_model( + model_id="claude-sonnet-5", + allowed_app_roles=["staff"], # says "staff" but no role actually grants it + ) + self._grant(user, mock_app_role_service, []) + + assert await service.can_access_model(user, model) is False + + @pytest.mark.asyncio + async def test_single_check_agrees_with_filter( + self, service, mock_app_role_service + ): + """can_access_model and filter_accessible_models must never disagree.""" + user = create_test_user(roles=["Staff"]) + models = [ + create_test_model(model_id="granted", allowed_app_roles=[]), + create_test_model(model_id="not-granted", allowed_app_roles=["staff"]), + ] + self._grant(user, mock_app_role_service, ["granted"]) + + listed = {m.model_id for m in await service.filter_accessible_models(user, models)} + + for model in models: + assert ( + await service.can_access_model(user, model) + ) is (model.model_id in listed) + + assert listed == {"granted"} diff --git a/backend/src/apis/app_api/admin/services/tests/test_model_roles.py b/backend/src/apis/app_api/admin/services/tests/test_model_roles.py new file mode 100644 index 000000000..1e39a55e7 --- /dev/null +++ b/backend/src/apis/app_api/admin/services/tests/test_model_roles.py @@ -0,0 +1,240 @@ +"""Unit tests for ModelRoleService. + +The AppRole record is the single source of truth for model access. These tests +pin the two directions of that contract: + +* the model form's role picker writes THROUGH to each role's ``grantedModels`` + (previously it wrote a model-side field that no access check ever read, so + enabling a model for a role from the model page silently did nothing); +* the model's role fields are derived back FROM the roles on read, so the model + page and the role page can never show different answers. +""" + +import pytest +from datetime import datetime, timezone +from unittest.mock import AsyncMock + +from apis.app_api.admin.services.model_roles import ModelRoleService +from apis.shared.auth.models import User +from apis.shared.models.models import ManagedModel +from apis.shared.rbac.models import AppRole + + +def make_admin() -> User: + return User( + user_id="admin-1", email="admin@example.com", name="Admin", roles=["Admin"] + ) + + +def make_role( + role_id: str, + granted_models: list = None, + inherits_from: list = None, + enabled: bool = True, +) -> AppRole: + return AppRole( + role_id=role_id, + display_name=role_id.title(), + description=f"{role_id} role", + granted_models=granted_models or [], + inherits_from=inherits_from or [], + enabled=enabled, + ) + + +def make_model(model_id: str = "claude-sonnet-5") -> ManagedModel: + now = datetime.now(timezone.utc) + return ManagedModel( + id="uuid-1", + model_id=model_id, + model_name="Claude Sonnet 5", + provider="bedrock", + provider_name="AWS Bedrock", + input_modalities=["TEXT"], + output_modalities=["TEXT"], + max_input_tokens=200000, + enabled=True, + input_price_per_million_tokens=3.0, + output_price_per_million_tokens=15.0, + created_at=now, + updated_at=now, + ) + + +def make_service(roles: list) -> tuple: + """Build a ModelRoleService over a fixed role list. Returns (service, admin_mock).""" + admin_service = AsyncMock() + admin_service.list_roles.return_value = roles + return ModelRoleService(app_role_admin_service=admin_service), admin_service + + +def granted_models_written(admin_service, role_id: str) -> list: + """The grantedModels passed to update_role for a given role.""" + for call in admin_service.update_role.await_args_list: + if call.args[0] == role_id: + return call.args[1].granted_models + raise AssertionError(f"update_role was never called for {role_id}") + + +class TestSetRolesForModel: + """The write-through: the picker's selection lands on the ROLE records.""" + + @pytest.mark.asyncio + async def test_grant_adds_model_to_role_granted_models(self): + """ + The bug this fixes: enabling Sonnet 5 for Staff on the model page must + add the model to the Staff role's grantedModels β€” that is the only field + the chat model list reads. + """ + service, admin_service = make_service([make_role("staff")]) + + await service.set_roles_for_model("claude-sonnet-5", ["staff"], make_admin()) + + assert granted_models_written(admin_service, "staff") == ["claude-sonnet-5"] + + @pytest.mark.asyncio + async def test_deselecting_a_role_revokes_the_grant(self): + service, admin_service = make_service( + [make_role("staff", granted_models=["claude-sonnet-5", "other"])] + ) + + await service.set_roles_for_model("claude-sonnet-5", [], make_admin()) + + assert granted_models_written(admin_service, "staff") == ["other"] + + @pytest.mark.asyncio + async def test_untouched_roles_are_not_rewritten(self): + """Only roles whose grants actually change should be written.""" + service, admin_service = make_service( + [ + make_role("staff", granted_models=["claude-sonnet-5"]), + make_role("student", granted_models=["haiku"]), + ] + ) + + await service.set_roles_for_model("claude-sonnet-5", ["staff"], make_admin()) + + admin_service.update_role.assert_not_awaited() + + @pytest.mark.asyncio + async def test_wildcard_role_is_not_given_a_redundant_grant(self): + """A role granting '*' already covers every model; don't add the id too.""" + service, admin_service = make_service( + [make_role("system_admin", granted_models=["*"])] + ) + + await service.set_roles_for_model( + "claude-sonnet-5", ["system_admin"], make_admin() + ) + + admin_service.update_role.assert_not_awaited() + + @pytest.mark.asyncio + async def test_unknown_role_is_rejected(self): + service, admin_service = make_service([make_role("staff")]) + + with pytest.raises(ValueError, match="Unknown AppRole"): + await service.set_roles_for_model("claude-sonnet-5", ["ghost"], make_admin()) + + admin_service.update_role.assert_not_awaited() + + @pytest.mark.asyncio + async def test_rename_migrates_grants_to_the_new_model_id(self): + """Roles key grants on the provider model id, so a rename must move them.""" + service, admin_service = make_service( + [make_role("staff", granted_models=["old-id", "other"])] + ) + + await service.set_roles_for_model( + "new-id", ["staff"], make_admin(), previous_model_id="old-id" + ) + + written = granted_models_written(admin_service, "staff") + assert "old-id" not in written + assert written == ["other", "new-id"] + + +class TestRevokeModelFromAllRoles: + @pytest.mark.asyncio + async def test_delete_strips_the_model_from_every_role(self): + service, admin_service = make_service( + [ + make_role("staff", granted_models=["claude-sonnet-5", "haiku"]), + make_role("student", granted_models=["claude-sonnet-5"]), + make_role("guest", granted_models=["haiku"]), + ] + ) + + await service.revoke_model_from_all_roles("claude-sonnet-5", make_admin()) + + assert granted_models_written(admin_service, "staff") == ["haiku"] + assert granted_models_written(admin_service, "student") == [] + # 'guest' never granted it, so it should not be rewritten. + assert admin_service.update_role.await_count == 2 + + +class TestHydrateModelRoles: + """The derived read: role records -> the model's displayed role fields.""" + + @pytest.mark.asyncio + async def test_direct_grants_populate_allowed_app_roles(self): + service, _ = make_service( + [ + make_role("staff", granted_models=["claude-sonnet-5"]), + make_role("student", granted_models=["haiku"]), + ] + ) + model = make_model() + + await service.hydrate_model_roles([model]) + + assert model.allowed_app_roles == ["staff"] + assert model.inherited_app_roles == [] + + @pytest.mark.asyncio + async def test_wildcard_and_inherited_grants_are_reported_separately(self): + """ + A wildcard or inherited grant is real access, but it isn't a direct grant + β€” it must not show up as a checked box the admin could 'uncheck'. + """ + service, _ = make_service( + [ + make_role("system_admin", granted_models=["*"]), + make_role("staff", granted_models=["claude-sonnet-5"]), + make_role("ta", inherits_from=["staff"]), + ] + ) + model = make_model() + + await service.hydrate_model_roles([model]) + + assert model.allowed_app_roles == ["staff"] + assert sorted(model.inherited_app_roles) == ["system_admin", "ta"] + + @pytest.mark.asyncio + async def test_inheritance_from_a_disabled_parent_does_not_grant(self): + service, _ = make_service( + [ + make_role("staff", granted_models=["claude-sonnet-5"], enabled=False), + make_role("ta", inherits_from=["staff"]), + ] + ) + model = make_model() + + await service.hydrate_model_roles([model]) + + assert model.inherited_app_roles == [] + + @pytest.mark.asyncio + async def test_hydrating_a_catalog_queries_roles_once(self): + """Cost guard: one role query for the whole model list, not one per model.""" + service, admin_service = make_service( + [make_role("staff", granted_models=["m1"])] + ) + models = [make_model("m1"), make_model("m2"), make_model("m3")] + + await service.hydrate_model_roles(models) + + admin_service.list_roles.assert_awaited_once() + assert models[0].allowed_app_roles == ["staff"] + assert models[1].allowed_app_roles == [] diff --git a/backend/src/apis/app_api/agent_designer/services/binding_validation.py b/backend/src/apis/app_api/agent_designer/services/binding_validation.py index df8a13ba0..123eb1f01 100644 --- a/backend/src/apis/app_api/agent_designer/services/binding_validation.py +++ b/backend/src/apis/app_api/agent_designer/services/binding_validation.py @@ -92,13 +92,11 @@ async def _validate_model(user: User, cfg: AgentModelConfig, svc: ModelAccessSer model = next((m for m in await list_all_managed_models() if m.model_id == cfg.model_id), None) if model is None: raise BindingValidationError(f"Model '{cfg.model_id}' is not available.", status_code=400) - # Use the SAME predicate the ``/agents/bindable`` catalog uses (``filter_accessible_models``), - # not ``can_access_model``: the latter only honors an AppRole model grant when the model - # record also carries a non-empty ``allowed_app_roles``, so a model granted purely via the - # user's AppRole ``permissions.models`` (empty ``allowed_app_roles``) is *listed* by the - # palette but would be *rejected* here β€” the picker shows it, saving 403s. Filtering the - # single model guarantees "if the palette offers it, the write accepts it", and matches the - # runtime's membership grant. + # Use the SAME predicate the ``/agents/bindable`` catalog uses, so "if the palette + # offers it, the write accepts it". ``can_access_model`` is now equivalent (both + # delegate to one ``_grants_access`` rule); this once had to avoid it because it + # additionally gated on the model's ``allowed_app_roles``, rejecting models the + # palette had just listed. if not await svc.filter_accessible_models(user, [model]): raise BindingValidationError( f"You do not have access to model '{cfg.model_id}'.", status_code=403 diff --git a/backend/src/apis/app_api/models/routes.py b/backend/src/apis/app_api/models/routes.py index 1dfc0cab1..3f1058c5c 100644 --- a/backend/src/apis/app_api/models/routes.py +++ b/backend/src/apis/app_api/models/routes.py @@ -30,14 +30,16 @@ async def list_models_for_user( This endpoint returns models filtered by the user's permissions. Only models that are: - 1. Enabled - 2. Accessible via AppRole permissions (allowedAppRoles) OR + 1. Enabled, AND + 2. Granted by one of the user's AppRoles β€” the role lists the model in its + grantedModels, or grants the '*' wildcard β€” OR 3. Available via legacy JWT role matching (availableToRoles) will be returned. Access Control: - - AppRole-based access is checked first (via allowedAppRoles field) + - The AppRole record is the source of truth. The model's own allowedAppRoles + field is DERIVED from those roles for display and is not consulted here. - Legacy JWT role-based access is checked as fallback (via availableToRoles field) - During the transition period, access is granted if EITHER method matches diff --git a/backend/src/apis/app_api/sessions/routes.py b/backend/src/apis/app_api/sessions/routes.py index e0f6298c8..4ccfb5dd1 100644 --- a/backend/src/apis/app_api/sessions/routes.py +++ b/backend/src/apis/app_api/sessions/routes.py @@ -678,6 +678,18 @@ async def signal_turn_interrupted_endpoint( reason=body.reason, source="client_signal", ) + # Distributed turn cancellation: a client abort doesn't propagate + # through the AgentCore Runtime data plane, so arm a cancel on the + # session's single-flight lease. The container running the turn + # observes it on its next heartbeat and unwinds β€” releasing the lease + # so the user's resend isn't rejected with 409 and stopping wasted + # model/tool work. Owner-scoped, so a stale Stop can't kill a later + # turn. Best-effort: never fail the Stop signal on this. + try: + from apis.shared.sessions.session_lease import request_session_cancel + await request_session_cancel(session_id, user_id) + except Exception: + logger.warning("Failed to arm session cancel on stop", exc_info=True) return Response(status_code=204) except Exception: logger.error("Error recording turn interruption", exc_info=True) diff --git a/backend/src/apis/app_api/sessions/tests/test_cache_savings.py b/backend/src/apis/app_api/sessions/tests/test_cache_savings.py index dcb8b16d0..54c0e747a 100644 --- a/backend/src/apis/app_api/sessions/tests/test_cache_savings.py +++ b/backend/src/apis/app_api/sessions/tests/test_cache_savings.py @@ -73,7 +73,7 @@ def sample_message_metadata(self, sample_token_usage_with_cache, sample_model_in async def test_cache_savings_calculation(self, mock_storage, sample_message_metadata): """Test that cache savings are calculated correctly""" with patch( - 'apis.app_api.storage.get_metadata_storage', + 'apis.shared.storage.get_metadata_storage', return_value=mock_storage ): from apis.app_api.sessions.services.metadata import _update_cost_summary_async @@ -138,7 +138,7 @@ async def test_cache_savings_zero_when_no_cache_reads(self, mock_storage): ) with patch( - 'apis.app_api.storage.get_metadata_storage', + 'apis.shared.storage.get_metadata_storage', return_value=mock_storage ): from apis.app_api.sessions.services.metadata import _update_cost_summary_async @@ -180,7 +180,7 @@ async def test_cache_savings_zero_when_no_pricing_snapshot(self, mock_storage): ) with patch( - 'apis.app_api.storage.get_metadata_storage', + 'apis.shared.storage.get_metadata_storage', return_value=mock_storage ): from apis.app_api.sessions.services.metadata import _update_cost_summary_async @@ -230,7 +230,7 @@ async def test_cache_savings_large_cache_hit(self, mock_storage): ) with patch( - 'apis.app_api.storage.get_metadata_storage', + 'apis.shared.storage.get_metadata_storage', return_value=mock_storage ): from apis.app_api.sessions.services.metadata import _update_cost_summary_async @@ -287,7 +287,7 @@ async def test_cache_savings_with_haiku_pricing(self, mock_storage): ) with patch( - 'apis.app_api.storage.get_metadata_storage', + 'apis.shared.storage.get_metadata_storage', return_value=mock_storage ): from apis.app_api.sessions.services.metadata import _update_cost_summary_async diff --git a/backend/src/apis/app_api/shares/routes.py b/backend/src/apis/app_api/shares/routes.py index 530ae7c86..6d4629cf2 100644 --- a/backend/src/apis/app_api/shares/routes.py +++ b/backend/src/apis/app_api/shares/routes.py @@ -26,6 +26,7 @@ NotOwnerError, SessionNotFoundError, ShareNotFoundError, + ShareStorageUnavailableError, ShareTableNotFoundError, get_share_service, ) @@ -71,6 +72,11 @@ async def create_share( raise HTTPException(status_code=403, detail="You do not have permission to share this session") except ShareTableNotFoundError: raise HTTPException(status_code=503, detail="Share feature unavailable - table not deployed") + except ShareStorageUnavailableError: + raise HTTPException( + status_code=503, + detail="Sharing is temporarily unavailable. Please try again later.", + ) except Exception as e: safe_session_id = session_id.replace("\r", "").replace("\n", "") logger.error(f"Error creating share for session {safe_session_id}: {e}", exc_info=True) diff --git a/backend/src/apis/app_api/shares/service.py b/backend/src/apis/app_api/shares/service.py index f9a12f91a..9df891e71 100644 --- a/backend/src/apis/app_api/shares/service.py +++ b/backend/src/apis/app_api/shares/service.py @@ -4,13 +4,14 @@ conversation share snapshots. Supports multiple shares per session. """ +import json import logging import os import re import uuid from decimal import Decimal from datetime import datetime, timezone -from typing import Any, List, Optional +from typing import Any, List, Optional, Tuple import boto3 from boto3.dynamodb.conditions import Key @@ -27,6 +28,15 @@ SharedConversationResponse, UpdateShareRequest, ) +from .snapshot_store import ( + ShareSnapshotStore, + ShareSnapshotStoreError, + get_share_snapshot_store, +) + +# Snapshot schema version stamped on the S3 body pointer, for forward +# migration if the body shape ever changes. +_SNAPSHOT_SCHEMA_VERSION = 1 logger = logging.getLogger(__name__) @@ -34,10 +44,13 @@ class ShareService: """Handles share CRUD operations against the shared-conversations DynamoDB table.""" - def __init__(self) -> None: + def __init__(self, snapshot_store: Optional[ShareSnapshotStore] = None) -> None: table_name = os.environ.get("SHARED_CONVERSATIONS_TABLE_NAME", "") self._table_name = table_name self._enabled = bool(table_name) + # S3-backed snapshot body store. Injectable for tests; otherwise the + # process-global store (bucket from SHARED_CONVERSATIONS_BUCKET_NAME). + self._snapshot_store = snapshot_store or get_share_snapshot_store() if self._enabled: self._dynamodb = boto3.resource("dynamodb") @@ -78,17 +91,33 @@ async def create_share( metadata_snapshot = metadata.model_dump(by_alias=True, exclude_none=True) - # Convert floats to Decimal for DynamoDB compatibility - messages_snapshot = self._convert_floats_to_decimal(messages_snapshot) - metadata_snapshot = self._convert_floats_to_decimal(metadata_snapshot) - - # Build item share_id = str(uuid.uuid4()) now = datetime.now(timezone.utc).isoformat() allowed_emails = self._resolve_allowed_emails( request.access_level, request.allowed_emails, user.email ) + # Offload the snapshot BODY (messages + metadata) to S3 β€” a long + # conversation exceeds DynamoDB's 400 KB item limit if inlined. The + # DynamoDB row keeps only the small control fields plus a pointer. + # The body is opaque JSON bytes in S3, so we serialize the raw + # model_dump directly and skip the floatβ†’Decimal dance (that only + # exists to satisfy DynamoDB's boto3 resource, which rejects floats). + if not self._snapshot_store.enabled: + raise ShareStorageUnavailableError() + + body_bytes = json.dumps( + {"metadata": metadata_snapshot, "messages": messages_snapshot} + ).encode("utf-8") + + try: + bucket_key = self._snapshot_store.put(share_id=share_id, body=body_bytes) + except ShareSnapshotStoreError as e: + logger.error( + f"Failed to store snapshot body for share {self._sanitize_id(share_id)}: {e}" + ) + raise ShareStorageUnavailableError() from e + item = { "share_id": share_id, "session_id": session_id, @@ -96,8 +125,12 @@ async def create_share( "owner_email": user.email, "access_level": request.access_level, "created_at": now, - "metadata": metadata_snapshot, - "messages": messages_snapshot, + "body_ref": { + "bucket_key": bucket_key, + "format": "json", + "schema_version": _SNAPSHOT_SCHEMA_VERSION, + "byte_size": len(body_bytes), + }, } if allowed_emails is not None: item["allowed_emails"] = allowed_emails @@ -194,6 +227,7 @@ async def revoke_share(self, share_id: str, user: User) -> None: raise NotOwnerError() self._table.delete_item(Key={"share_id": item["share_id"]}) + self._delete_snapshot_body(item) logger.info(f"Revoked share {item['share_id']}") async def delete_shares_for_session(self, session_id: str) -> int: @@ -220,6 +254,12 @@ async def delete_shares_for_session(self, session_id: str) -> int: for item in items: batch.delete_item(Key={"share_id": item["share_id"]}) + # Best-effort cleanup of each share's S3 snapshot body. The + # DynamoDB delete above is what makes the share link stop working; + # an S3 miss here only leaves an orphan object, never a live share. + for item in items: + self._delete_snapshot_body(item) + logger.info( f"Deleted {len(items)} share(s) for session " f"{self._sanitize_id(session_id)}" @@ -264,8 +304,7 @@ async def export_shared_conversation( self._check_access(item, requester) - snapshot_messages = item.get("messages", []) - metadata = item.get("metadata", {}) + metadata, snapshot_messages = self._load_snapshot_body(item) original_title = metadata.get("title", "Untitled Conversation") new_title = f"{original_title} (shared)" @@ -427,6 +466,23 @@ def _convert_floats_to_decimal(obj: Any) -> Any: return [ShareService._convert_floats_to_decimal(item) for item in obj] return obj + @staticmethod + def _convert_decimals_to_float(obj: Any) -> Any: + """Recursively convert DynamoDB ``Decimal`` values back to native types. + + Legacy inline shares were written with ``_convert_floats_to_decimal``, + so their bodies come back off DynamoDB as ``Decimal``. Convert them + back β€” to ``int`` when integral, else ``float`` β€” so the legacy read + path yields the same plain-JSON shape as the S3-backed path. + """ + if isinstance(obj, Decimal): + return int(obj) if obj % 1 == 0 else float(obj) + elif isinstance(obj, dict): + return {k: ShareService._convert_decimals_to_float(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [ShareService._convert_decimals_to_float(item) for item in obj] + return obj + @staticmethod def _sanitize_id(value: str, max_length: int = 128) -> str: """Return a log-safe version of an ID string. @@ -507,11 +563,64 @@ def _build_share_response(self, item: dict) -> ShareResponse: share_url=f"/shared/{item['share_id']}", ) + def _load_snapshot_body(self, item: dict) -> Tuple[dict, list]: + """Return ``(metadata, messages)`` for a share item. + + Handles three item shapes for backward compatibility: + + - **New** (``body_ref`` present): fetch the JSON body from S3. + - **Legacy inline** (``messages`` present, no ``body_ref``): read the + body straight off the DynamoDB item, exactly as before the S3 + offload. Existing shares predate the offload and stay readable + with no migration. + - **Malformed** (neither): unreadable β†’ ``ShareNotFoundError``. + """ + body_ref = item.get("body_ref") + if body_ref: + key = body_ref.get("bucket_key") + try: + raw = self._snapshot_store.get(key) + body = json.loads(raw) + except (ShareSnapshotStoreError, ValueError) as e: + logger.error( + f"Failed to load snapshot body for share " + f"{self._sanitize_id(str(item.get('share_id', '')))} " + f"key={key}: {e}" + ) + raise ShareNotFoundError() from e + return body.get("metadata", {}) or {}, body.get("messages", []) or [] + + if item.get("messages") is not None: + # Legacy inline share β€” DynamoDB stored floats as Decimal; convert + # back so downstream JSON/Pydantic handling matches the S3 path. + metadata = self._convert_decimals_to_float(item.get("metadata", {}) or {}) + messages = self._convert_decimals_to_float(item.get("messages", [])) + return metadata, messages + + logger.warning( + f"Share {self._sanitize_id(str(item.get('share_id', '')))} has neither " + "body_ref nor inline messages β€” treating as unreadable" + ) + raise ShareNotFoundError() + + def _delete_snapshot_body(self, item: dict) -> None: + """Best-effort delete of a share's S3 snapshot body. + + No-op for legacy inline shares (no ``body_ref``). Never raises β€” the + store's delete swallows storage misses; a failure here logs but never + blocks the DynamoDB delete that actually revokes the share. + """ + body_ref = item.get("body_ref") + if not body_ref: + return + key = body_ref.get("bucket_key") + if key: + self._snapshot_store.delete(key) + def _build_shared_conversation_response(self, item: dict) -> SharedConversationResponse: from apis.shared.sessions.models import MessageResponse - metadata = item.get("metadata", {}) - raw_messages = item.get("messages", []) + metadata, raw_messages = self._load_snapshot_body(item) messages = [] for msg_data in raw_messages: @@ -557,6 +666,18 @@ class ShareTableNotFoundError(Exception): pass +class ShareStorageUnavailableError(Exception): + """Raised when the S3 snapshot-body store is unconfigured or unreachable. + + The share body is offloaded to S3; if the bucket is unset (misconfigured + deploy / local dev without AWS) or the write fails, creating a share can't + proceed. Surfaced to the client as a 503 with a friendly message rather + than silently falling back to inline (which would reintroduce the 400 KB + item-size failure). + """ + pass + + # Global service instance (singleton) _service_instance: Optional[ShareService] = None diff --git a/backend/src/apis/app_api/shares/snapshot_store.py b/backend/src/apis/app_api/shares/snapshot_store.py new file mode 100644 index 000000000..d26ad6c60 --- /dev/null +++ b/backend/src/apis/app_api/shares/snapshot_store.py @@ -0,0 +1,217 @@ +"""S3-backed store for a conversation share's snapshot body. + +A share is a point-in-time snapshot of a conversation. The body β€” the full +message list plus the session-metadata snapshot β€” is too large to inline in a +single DynamoDB item (400 KB limit; a long conversation with tool results, +images, or documents blows past it and PutItem fails with a +``ValidationException``). So the body lives in the ``shared-conversations`` S3 +bucket and the DynamoDB row carries only the small control fields plus a +pointer (``body_ref``) to the object. + +This store is a faithful sibling of ``apis/shared/memory/store.py`` (the Memory +Spaces S3 offload) and ``apis/shared/skills/resource_store.py``: + + - Objects are **content-addressed**: the key is + ``shares/{share_id}/{content_hash}`` where ``content_hash`` is the sha256 + hex of the bytes. A share is immutable once created (``update_share`` only + touches access-control fields, never the body), so there is exactly one + object per share; content-addressing makes a retried ``create_share`` + idempotent. + - The DynamoDB row references the object by key; the bytes never travel + through DynamoDB. + +It lives under ``app_api/shares/`` rather than ``apis/shared/`` because app-api +is the only consumer (create writes; view/export read; revoke deletes). If a +second consumer ever appears it moves to ``apis.shared`` per the import +boundary rule. + +Configuration: the bucket name comes from ``SHARED_CONVERSATIONS_BUCKET_NAME`` +(set on the app-api role by the CDK ``SharedConversationsConstruct`` wiring). +When boto3 or the bucket name is absent (local dev without AWS), the store is +``enabled == False`` and every write raises ``ShareSnapshotStoreError`` so a +misconfigured deploy surfaces loudly rather than silently reintroducing the +400 KB inline cliff. +""" + +from __future__ import annotations + +import hashlib +import logging +import os +from typing import Optional + +try: # boto3 is absent in some local-dev setups + import boto3 + from botocore.exceptions import ClientError +except ImportError: # pragma: no cover - exercised only without boto3 + boto3 = None + ClientError = Exception # type: ignore[assignment, misc] + +logger = logging.getLogger(__name__) + +# AWS-managed (SSE-S3 / AES256) encryption, matching the bucket default and the +# memory-spaces / skills / artifacts / file-upload buckets. +_SSE_ALGORITHM = "AES256" + + +class ShareSnapshotStoreError(RuntimeError): + """Raised when the store cannot complete a requested operation. + + Covers both "storage not configured" (no bucket / no boto3) and an + unexpected S3 failure, so callers have one error type to translate. + """ + + +def content_key(share_id: str, content_hash: str) -> str: + """Return the content-addressed object key for a share's snapshot body.""" + return f"shares/{share_id}/{content_hash}" + + +def compute_content_hash(content: bytes) -> str: + """Return the sha256 hex digest used as the content address.""" + return hashlib.sha256(content).hexdigest() + + +class ShareSnapshotStore: + """Put / get / delete a share's snapshot-body bytes in S3.""" + + def __init__( + self, + bucket_name: Optional[str] = None, + s3_client: Optional[object] = None, + ) -> None: + self.bucket_name = bucket_name or os.environ.get( + "SHARED_CONVERSATIONS_BUCKET_NAME" + ) + # Allow an explicit client (tests inject a moto client); otherwise it is + # created lazily on first use so importing the module never needs creds. + self._s3 = s3_client + + @property + def enabled(self) -> bool: + """True when a bucket is configured and boto3 is importable.""" + return bool(self.bucket_name) and boto3 is not None + + def _client(self): + if self._s3 is None: + if boto3 is None: # pragma: no cover - import-guarded above + raise ShareSnapshotStoreError( + "share snapshot storage unavailable: boto3 is not installed" + ) + self._s3 = boto3.client("s3") + return self._s3 + + def _require_enabled(self) -> None: + if not self.enabled: + raise ShareSnapshotStoreError( + "share snapshot storage is not configured " + "(SHARED_CONVERSATIONS_BUCKET_NAME is unset)" + ) + + def put(self, *, share_id: str, body: bytes) -> str: + """Persist snapshot-body bytes content-addressed; return the object key. + + Computes the sha256 of ``body``, derives the + ``shares/{share_id}/{content_hash}`` key, and uploads. If an object + already exists at that key (same content β€” a retried create), the + upload is skipped (dedupe); the key is returned either way. + """ + self._require_enabled() + digest = compute_content_hash(body) + key = content_key(share_id, digest) + client = self._client() + + if self._object_exists(key): + logger.info( + "shared-conversations: dedupe hit for share=%s key=%s (%d bytes)", + share_id, + key, + len(body), + ) + return key + + try: + client.put_object( + Bucket=self.bucket_name, + Key=key, + Body=body, + ContentType="application/json", + ServerSideEncryption=_SSE_ALGORITHM, + ) + except ClientError as e: # pragma: no cover - network/permission path + logger.error( + "shared-conversations: put failed for share=%s key=%s: %s", + share_id, + key, + e, + ) + raise ShareSnapshotStoreError( + f"failed to store snapshot body for share '{share_id}'" + ) from e + + logger.info( + "shared-conversations: stored share=%s key=%s (%d bytes)", + share_id, + key, + len(body), + ) + return key + + def get(self, bucket_key: str) -> bytes: + """Return the bytes for an object key. Raises if missing/unavailable.""" + self._require_enabled() + client = self._client() + try: + response = client.get_object(Bucket=self.bucket_name, Key=bucket_key) + return response["Body"].read() + except ClientError as e: + code = e.response.get("Error", {}).get("Code", "") + if code in ("NoSuchKey", "404"): + raise ShareSnapshotStoreError( + f"snapshot body not found at key '{bucket_key}'" + ) from e + logger.error( + "shared-conversations: get failed for key=%s: %s", bucket_key, e + ) + raise ShareSnapshotStoreError( + f"failed to read snapshot body at key '{bucket_key}'" + ) from e + + def delete(self, bucket_key: str) -> None: + """Delete an object key. Best-effort β€” never raises on the storage + miss path (deleting an already-absent object is a no-op in S3).""" + if not self.enabled: + return + client = self._client() + try: + client.delete_object(Bucket=self.bucket_name, Key=bucket_key) + except ClientError: # pragma: no cover - best-effort cleanup + logger.warning( + "shared-conversations: delete failed for key=%s", + bucket_key, + exc_info=True, + ) + + def _object_exists(self, key: str) -> bool: + client = self._client() + try: + client.head_object(Bucket=self.bucket_name, Key=key) + return True + except ClientError as e: + code = e.response.get("Error", {}).get("Code", "") + if code in ("404", "NoSuchKey", "NotFound"): + return False + # Any other error (permissions, throttling) is real β€” surface it + # rather than masquerading as "absent" and double-uploading. + raise + + +_store: Optional[ShareSnapshotStore] = None + + +def get_share_snapshot_store() -> ShareSnapshotStore: + """Get or create the process-global share snapshot store.""" + global _store + if _store is None: + _store = ShareSnapshotStore() + return _store diff --git a/backend/src/apis/app_api/web_sources/crawl_repository.py b/backend/src/apis/app_api/web_sources/crawl_repository.py index c95145846..6db068eec 100644 --- a/backend/src/apis/app_api/web_sources/crawl_repository.py +++ b/backend/src/apis/app_api/web_sources/crawl_repository.py @@ -175,12 +175,16 @@ async def hard_delete_crawl_job(assistant_id: str, crawl_id: str) -> bool: return False -def _is_crawl_stale(job: CrawlJob) -> bool: +def is_crawl_stale(job: CrawlJob) -> bool: """A `running` crawl past the crawler's own budget is dead. The crawler always finalizes in a `finally` block; the only way a `running` row outlives the budget is the owning process died. Mirrors the same staleness pattern documents use in `document_service`. + + Also read by the delete path: a live crawl refuses deletion (its pages + are still being written), but a dead-but-`running` row must stay + removable or a zombie web source could never be cleared from the UI. """ try: started_str = job.started_at.rstrip("Z") @@ -226,7 +230,7 @@ async def list_active_crawls(assistant_id: str) -> List[CrawlJob]: logger.warning("Skipping unparseable CrawlJob row: %s", e) continue - if _is_crawl_stale(job): + if is_crawl_stale(job): logger.info( "Auto-reaping stale crawl %s/%s (started_at=%s)", assistant_id, diff --git a/backend/src/apis/app_api/web_sources/deletion_service.py b/backend/src/apis/app_api/web_sources/deletion_service.py new file mode 100644 index 000000000..54fb6e1ce --- /dev/null +++ b/backend/src/apis/app_api/web_sources/deletion_service.py @@ -0,0 +1,132 @@ +"""Removal of a whole web source β€” the crawl record plus every page it ingested. + +A "web source" is not a single stored entity: it is a `CrawlJob` row, the +fan-out of `Document` rows the crawler wrote beneath it, and (optionally) a +`web_crawl` sync policy keyed on the crawl id. Removing one therefore has to +walk that fan-out itself. + +The reverse relationship already exists: deleting the *last* page of a crawl +lets `documents.services.cleanup_service._cascade_delete_orphaned_crawl_jobs` +drop the now-orphaned `CrawlJob`. This module drives the same graph from the +other end β€” delete the source, and the pages go with it. + +Pages are matched to their crawl by URL prefix (`source_file_id` starts with +the crawl's `root_url`), the same rule the orphan cascade uses and sound for +the same reason: the crawler only enqueues URLs that already passed its +same-domain / same-root filter. +""" + +import asyncio +import logging +from typing import List, Set + +from apis.app_api.documents.models import Document +from apis.app_api.documents.services.cleanup_service import cleanup_assistant_documents +from apis.app_api.documents.services.document_service import ( + list_assistant_documents, + soft_delete_document, +) +from apis.app_api.web_sources.crawl_repository import hard_delete_crawl_job +from apis.shared.security.log_sanitize import scrub_log +from apis.shared.sync_policies.service import delete_sync_policies_for_source + +logger = logging.getLogger(__name__) + +# Strong refs to the fire-and-forget cleanup task. The event loop only holds +# weak references to tasks, so a bare `create_task` can be collected mid-run β€” +# the same trap the route documents for its background crawls. +_BACKGROUND_CLEANUPS: Set[asyncio.Task] = set() + + +class WebSourceDeletionError(Exception): + """The crawl record itself could not be removed.""" + + +async def delete_web_source( + *, + assistant_id: str, + crawl_id: str, + root_url: str, + owner_id: str, +) -> int: + """Remove a web source and everything it owns. Returns the page count removed. + + Order matters: + + 1. The sync policy goes first. It keys on the crawl id, so a dispatcher + sweep landing mid-delete would otherwise re-crawl the very source we + are removing and resurrect its pages. + 2. Pages are soft-deleted (status `deleting` + TTL) so they leave the KB + immediately, with the slow half β€” vectors and S3 objects β€” handed to a + background task, exactly as the single-document delete path does. + 3. The crawl row is hard-deleted last. If a soft-delete fails partway, the + surviving row still lists the source and the user can retry; deleting + the row first would strand its pages with nothing to remove them from. + + Raises `WebSourceDeletionError` if the crawl row survives β€” leaving it + behind would show the caller a source that is now empty but still listed. + """ + pages = await _list_crawl_pages(assistant_id, owner_id, root_url) + + policies_removed = await delete_sync_policies_for_source(assistant_id, crawl_id) + + deleted: List[Document] = [] + for page in pages: + document = await soft_delete_document(assistant_id, page.document_id, owner_id) + if document: + deleted.append(document) + + if not await hard_delete_crawl_job(assistant_id, crawl_id): + raise WebSourceDeletionError( + "Removed the pages but could not remove the web source record. " + "Try again in a moment." + ) + + # Vectors + S3, off the request path. `cleanup_assistant_documents` gathers + # per-document cleanup and never raises; it also hard-deletes each DynamoDB + # row on success. It is deliberately *not* given the docs' `web` connector + # id β€” that would fire the orphaned-CrawlJob cascade once per page (a full + # document scan each), and we just deleted the crawl row ourselves. + if deleted: + task = asyncio.create_task(cleanup_assistant_documents(assistant_id, deleted)) + _BACKGROUND_CLEANUPS.add(task) + task.add_done_callback(_BACKGROUND_CLEANUPS.discard) + + logger.info( + "Removed web source %s from assistant %s (root=%s, pages=%d, sync_policies=%d)", + scrub_log(crawl_id), + scrub_log(assistant_id), + scrub_log(root_url), + len(deleted), + policies_removed, + ) + return len(deleted) + + +async def _list_crawl_pages( + assistant_id: str, owner_id: str, root_url: str +) -> List[Document]: + """Every live document this crawl produced, across all pages of results. + + Rows already in `deleting` are skipped β€” their cleanup is in flight and + re-soft-deleting them would just reset the TTL. + """ + pages: List[Document] = [] + next_token = None + while True: + batch, next_token = await list_assistant_documents( + assistant_id, owner_id, next_token=next_token + ) + for document in batch: + if document.source_connector_id != "web": + continue + if not document.source_file_id: + continue + if not document.source_file_id.startswith(root_url): + continue + if document.status == "deleting": + continue + pages.append(document) + if not next_token: + break + return pages diff --git a/backend/src/apis/app_api/web_sources/routes.py b/backend/src/apis/app_api/web_sources/routes.py index f7c8a439a..3cf703b3a 100644 --- a/backend/src/apis/app_api/web_sources/routes.py +++ b/backend/src/apis/app_api/web_sources/routes.py @@ -32,10 +32,15 @@ from apis.app_api.web_sources.crawl_repository import ( create_crawl_job, get_crawl_job, + is_crawl_stale, list_active_crawls, list_all_crawls, ) from apis.app_api.web_sources.crawler import run_crawl +from apis.app_api.web_sources.deletion_service import ( + WebSourceDeletionError, + delete_web_source, +) from apis.app_api.web_sources.models import ( ActiveCrawlsResponse, CrawlJob, @@ -48,7 +53,7 @@ assert_url_is_public, url_extension_hint, ) -from apis.shared.assistants.service import get_assistant +from apis.shared.assistants.service import resolve_assistant_permission from apis.shared.auth import User, get_current_user_from_session from apis.shared.security.log_sanitize import scrub_log @@ -59,6 +64,31 @@ prefix="/assistants/{assistant_id}/web-sources", tags=["web-sources"] ) + +async def _require_edit_permission(assistant_id: str, current_user: User) -> str: + """Owner or editor share required β€” the same gate the documents surface uses. + + Returns the assistant's real owner_id, which the owner-keyed document + services below need in order to see an editor's assistant at all. + """ + assistant, permission = await resolve_assistant_permission( + assistant_id=assistant_id, + user_id=current_user.user_id, + user_email=current_user.email, + ) + if not assistant: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Assistant not found: {assistant_id}", + ) + if permission not in ("owner", "editor"): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="You do not have permission to manage web sources for this assistant", + ) + return assistant.owner_id + + # Strong refs to in-flight crawl tasks. Python's event loop tracks tasks # with weak references, so a bare `asyncio.ensure_future(run_crawl(...))` # can be garbage-collected mid-execution β€” leaving the root document @@ -83,13 +113,13 @@ async def start_crawl( with `max_depth=0` β€” the BFS visits only the root and terminates. The same async pipeline is used either way so there is no separate code path to keep in sync. + + Owner or editor may crawl. Note the document writes below are keyed on + the assistant (`PK=AST#`), not on its owner, so no owner_id needs + threading through β€” and `imported_by_user_id`/`started_by_user_id` stay + the *acting* user, which is the point of recording them. """ - assistant = await get_assistant(assistant_id, current_user.user_id) - if not assistant: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail=f"Assistant not found: {assistant_id}", - ) + await _require_edit_permission(assistant_id, current_user) try: normalized = assert_url_is_public(request.url) @@ -164,12 +194,7 @@ async def list_crawls( `?active=true` is the only filter currently honored β€” drives the SPA's "should I keep polling for new docs" decision. """ - assistant = await get_assistant(assistant_id, current_user.user_id) - if not assistant: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail=f"Assistant not found: {assistant_id}", - ) + await _require_edit_permission(assistant_id, current_user) if active: crawls = await list_active_crawls(assistant_id) else: @@ -186,16 +211,66 @@ async def get_crawl( current_user: User = Depends(get_current_user_from_session), ) -> CrawlJob: """Return a single crawl's current status + counters.""" - assistant = await get_assistant(assistant_id, current_user.user_id) - if not assistant: + await _require_edit_permission(assistant_id, current_user) + job = await get_crawl_job(assistant_id, crawl_id) + if not job: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, - detail=f"Assistant not found: {assistant_id}", + detail=f"Crawl not found: {crawl_id}", ) + return job + + +@router.delete("/crawls/{crawl_id}", status_code=status.HTTP_204_NO_CONTENT) +async def delete_crawl( + assistant_id: str, + crawl_id: str, + current_user: User = Depends(get_current_user_from_session), +) -> None: + """Remove a web source: the crawl record, its pages, and its sync policy. + + A crawl that is genuinely still in flight is refused (409) rather than + raced β€” the crawler would keep writing pages we just enumerated, stranding + them under a deleted parent. A crawl stuck at `running` because its owning + process died is *not* in flight, and stays deletable (`is_crawl_stale`). + """ + owner_id = await _require_edit_permission(assistant_id, current_user) + job = await get_crawl_job(assistant_id, crawl_id) if not job: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Crawl not found: {crawl_id}", ) - return job + + if job.status == "running" and not is_crawl_stale(job): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="This crawl is still running. Wait for it to finish, then remove it.", + ) + + try: + removed = await delete_web_source( + assistant_id=assistant_id, + crawl_id=crawl_id, + root_url=job.root_url, + owner_id=owner_id, + ) + except WebSourceDeletionError as e: + logger.error( + "Failed to remove web source %s from assistant %s: %s", + scrub_log(crawl_id), + scrub_log(assistant_id), + e, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e) + ) + + logger.info( + "Deleted web source %s (%d pages) from assistant %s", + scrub_log(crawl_id), + removed, + scrub_log(assistant_id), + ) + return None diff --git a/backend/src/apis/inference_api/chat/routes.py b/backend/src/apis/inference_api/chat/routes.py index 7e5281954..e9d1dd852 100644 --- a/backend/src/apis/inference_api/chat/routes.py +++ b/backend/src/apis/inference_api/chat/routes.py @@ -8,6 +8,7 @@ """ import asyncio +import contextlib import json import logging import os @@ -94,6 +95,44 @@ DEFAULT_AGENT_TYPE = DEFAULT_CHAT_MODE +def _mark_session_cancelled(agent) -> None: + """Flip the agent's session-manager ``cancelled`` flag (cooperative stop). + + Both StopHook (tool boundaries) and the stream coordinator (mid-generation) + read this flag to unwind the turn. Defensive: a nonstandard agent without a + session manager is simply a no-op. + """ + session_manager = getattr(agent, "session_manager", None) + if session_manager is not None: + session_manager.cancelled = True + logger.info("Cooperative stop: cancel observed for the running turn") + + +async def _lease_heartbeat_loop(lease, agent) -> None: + """Renew the single-flight session lease and observe cancel requests. + + Runs as a background task for the life of the SSE stream. Renewing on a wall + clock (rather than piggybacking on SSE-event cadence) keeps the lease alive + across a long silent tool call β€” code-interpreter / browser can run past the + lease window between yielded events β€” and bounds Stopβ†’resend latency to one + interval. Each renew also reports whether a cancel has been armed for this + lease owner; on the first such observation we flip the agent's ``cancelled`` + flag and stop renewing (the turn is unwinding). Best-effort and owner-scoped; + cancelled in the stream generator's ``finally``. + """ + from apis.shared.sessions.session_lease import ( + LEASE_HEARTBEAT_SECONDS, + renew_session_lease, + ) + + while True: + await asyncio.sleep(LEASE_HEARTBEAT_SECONDS) + cancel_requested = await renew_session_lease(lease) + if cancel_requested: + _mark_session_cancelled(agent) + return + + def is_preview_session(session_id: str) -> bool: """Check if a session ID is a preview session (should skip persistence). @@ -1503,6 +1542,42 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g system_prompt = append_active_prompt(system_prompt, prompt_name, prompt_text) logger.info(f"Appended custom system prompt: {prompt_name!r}") + # Per-session single-flight guard (docs/specs/session-single-flight-guard.md, + # follow-up to PR #653). A client abort doesn't propagate through the + # AgentCore Runtime data plane and the Runtime can route a duplicate + # invocation to a *different* container, so two agent loops could otherwise + # run concurrently against one AgentCore Memory session and corrupt + # tool-pairing history. Acquire a distributed lease at turn-start; reject a + # duplicate with 409. Resume / max-tokens continuation re-enter a loop that + # already ended, so they take the lease with force=True (never blocked, but + # still install it so a fresh duplicate during them is rejected). Preview + # sessions and the local no-DynamoDB path (lease None) skip the guard. + session_lease = None + if not is_preview_session(input_data.session_id): + from apis.shared.sessions.session_lease import ( + acquire_session_lease, + SessionBusyError, + ) + + try: + session_lease = await acquire_session_lease( + input_data.session_id, + user_id, + force=is_resume or is_continuation, + ) + except SessionBusyError: + logger.warning( + "Rejected duplicate concurrent invocation for session %s (409)", + scrub_log(input_data.session_id), + ) + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=( + "A response is already streaming for this conversation. " + "Wait for it to finish before sending another message." + ), + ) + try: # Resume requests rebuild the agent from the persisted PausedTurnSnapshot # so a refresh / cache eviction / pod restart between pause and resume @@ -1912,20 +1987,52 @@ def _session_title_sse() -> Optional[str]: except Exception as cleanup_err: logger.error("Failed to clear resolved pending_interrupts: %s", cleanup_err, exc_info=True) + # Wrap the agent stream so the single-flight session lease is heartbeat- + # renewed while the turn runs and released when the stream ends. FastAPI + # runs this generator *after* the handler returns, so the lease can't be + # released in the handler body without ending it prematurely β€” the + # generator's finally is the release site for the happy path (the two + # except handlers below cover pre-stream failures). + async def _guarded_stream() -> AsyncGenerator[str, None]: + heartbeat_task = ( + asyncio.create_task(_lease_heartbeat_loop(session_lease, agent)) + if session_lease is not None + else None + ) + try: + async for chunk in stream_with_quota_warning(): + yield chunk + finally: + if heartbeat_task is not None: + heartbeat_task.cancel() + # Await the cancelled task so its CancelledError is retrieved + # (never re-raised) before the lease is released. + await asyncio.gather(heartbeat_task, return_exceptions=True) + from apis.shared.sessions.session_lease import release_session_lease + await release_session_lease(session_lease) + # Stream response from agent as SSE (with optional files) # Note: Compression is handled by GZipMiddleware if configured in main.py return StreamingResponse( - stream_with_quota_warning(), + _guarded_stream(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no", "X-Session-ID": input_data.session_id}, ) except HTTPException: - # Re-raise HTTP exceptions as-is (e.g., from auth) + # Re-raise HTTP exceptions as-is (e.g., from auth). Release the lease + # first β€” a failure raised after acquire (e.g. resume/interrupt 400s) + # means the turn won't stream, so its generator finally never runs. + from apis.shared.sessions.session_lease import release_session_lease + await release_session_lease(session_lease) raise except Exception as e: - # Stream error as a conversational assistant message for better UX + # Stream error as a conversational assistant message for better UX. + # The agent turn won't run, so release the lease here (the error stream + # is a canned single message, not an agent loop). logger.error("Error in invocations", exc_info=True) + from apis.shared.sessions.session_lease import release_session_lease + await release_session_lease(session_lease) error_event = build_conversational_error_event(code=ErrorCode.AGENT_ERROR, error=e, session_id=input_data.session_id, recoverable=True) diff --git a/backend/src/apis/shared/models/managed_models.py b/backend/src/apis/shared/models/managed_models.py index a50f56973..26a73668d 100644 --- a/backend/src/apis/shared/models/managed_models.py +++ b/backend/src/apis/shared/models/managed_models.py @@ -234,7 +234,9 @@ async def _create_managed_model_cloud(model_data: ManagedModelCreate, table_name output_modalities=model_data.output_modalities, max_input_tokens=model_data.max_input_tokens, max_output_tokens=model_data.max_output_tokens, - allowed_app_roles=model_data.allowed_app_roles, + # allowed_app_roles is intentionally not set here: it is derived from the + # AppRole records on read (see app_api/admin/services/model_roles.py). + # The caller writes the requested roles through to those records. available_to_roles=model_data.available_to_roles, enabled=model_data.enabled, input_price_per_million_tokens=model_data.input_price_per_million_tokens, @@ -265,7 +267,6 @@ async def _create_managed_model_cloud(model_data: ManagedModelCreate, table_name 'inputModalities': model_data.input_modalities, 'outputModalities': model_data.output_modalities, 'maxInputTokens': model_data.max_input_tokens, - 'allowedAppRoles': model_data.allowed_app_roles, 'availableToRoles': model_data.available_to_roles, 'enabled': model_data.enabled, 'inputPricePerMillionTokens': model_data.input_price_per_million_tokens, @@ -537,6 +538,12 @@ async def _update_managed_model_cloud(model_id: str, updates: ManagedModelUpdate # Get update data update_data = updates.model_dump(exclude_none=True, by_alias=True) + # Role access is owned by the AppRole records, not the model item. The admin + # route writes allowedAppRoles through to each role's grantedModels and the + # value is derived back on read, so persisting it here would only create a + # copy that drifts out of date. + update_data.pop('allowedAppRoles', None) + if not update_data: return existing_model # No updates to apply diff --git a/backend/src/apis/shared/models/models.py b/backend/src/apis/shared/models/models.py index 1234aee6b..3457e195c 100644 --- a/backend/src/apis/shared/models/models.py +++ b/backend/src/apis/shared/models/models.py @@ -161,11 +161,13 @@ class ManagedModelCreate(BaseModel): # value is only a ceiling for the admin-configured max_tokens inference # param β€” it is never sent to the provider β€” so leaving it unset is safe. max_output_tokens: Optional[int] = Field(None, alias="maxOutputTokens", ge=1) - # Access control: AppRoles (preferred) or legacy JWT roles + # Access control. Not stored on the model item: the admin routes write this + # through to each named role's ``grantedModels``, which is the source of truth. allowed_app_roles: List[str] = Field( default_factory=list, alias="allowedAppRoles", - description="AppRole IDs that can access this model (preferred over availableToRoles)" + description="AppRole IDs that should grant this model. Written through to each " + "role's grantedModels; not persisted on the model record." ) available_to_roles: List[str] = Field( default_factory=list, @@ -247,11 +249,14 @@ class ManagedModelUpdate(BaseModel): output_modalities: Optional[List[str]] = Field(None, alias="outputModalities") max_input_tokens: Optional[int] = Field(None, alias="maxInputTokens", ge=1) max_output_tokens: Optional[int] = Field(None, alias="maxOutputTokens", ge=1) - # Access control: AppRoles (preferred) or legacy JWT roles + # Access control. Not stored on the model item: the admin routes write this + # through to each named role's ``grantedModels``, which is the source of truth. + # None means "leave role grants alone"; [] means "revoke every direct grant". allowed_app_roles: Optional[List[str]] = Field( None, alias="allowedAppRoles", - description="AppRole IDs that can access this model (preferred over availableToRoles)" + description="AppRole IDs that should grant this model. Written through to each " + "role's grantedModels; not persisted on the model record." ) available_to_roles: Optional[List[str]] = Field( None, @@ -327,11 +332,23 @@ class ManagedModel(BaseModel): output_modalities: List[str] = Field(..., alias="outputModalities") max_input_tokens: int = Field(..., alias="maxInputTokens") max_output_tokens: Optional[int] = Field(None, alias="maxOutputTokens") - # Access control: AppRoles (preferred) or legacy JWT roles + # Access control. The AppRole record is the single source of truth: a role + # grants a model via its own ``grantedModels``. The two fields below are + # DERIVED from those role records on read (see ModelRoleService) and are not + # persisted on the model item β€” access checks never read them. allowed_app_roles: List[str] = Field( default_factory=list, alias="allowedAppRoles", - description="AppRole IDs that can access this model (preferred over availableToRoles)" + description="[DERIVED] AppRole IDs that grant this model DIRECTLY (the role lists " + "it in grantedModels). Editable via the admin model form, which writes " + "through to each role's grantedModels." + ) + inherited_app_roles: List[str] = Field( + default_factory=list, + alias="inheritedAppRoles", + description="[DERIVED, read-only] AppRole IDs that grant this model indirectly β€” " + "via a wildcard ('*') grant or via inheritance from a parent role. " + "Cannot be toggled from the model form; edit the role instead." ) available_to_roles: List[str] = Field( default_factory=list, @@ -387,3 +404,20 @@ class ManagedModel(BaseModel): ) created_at: datetime = Field(..., alias="createdAt") updated_at: datetime = Field(..., alias="updatedAt") + + +class ModelRoleAssignment(BaseModel): + """Role assignment info for a model. Mirrors ToolRoleAssignment.""" + + role_id: str = Field(..., alias="roleId") + display_name: str = Field(..., alias="displayName") + grant_type: str = Field( + ..., + alias="grantType", + description="'direct' (role lists the model in grantedModels), " + "'wildcard' (role grants '*'), or 'inherited' (a parent role grants it)", + ) + inherited_from: Optional[str] = Field(None, alias="inheritedFrom") + enabled: bool + + model_config = {"populate_by_name": True} diff --git a/backend/src/apis/shared/sessions/session_lease.py b/backend/src/apis/shared/sessions/session_lease.py new file mode 100644 index 000000000..5bfb3acd4 --- /dev/null +++ b/backend/src/apis/shared/sessions/session_lease.py @@ -0,0 +1,315 @@ +"""Per-session single-flight lease (distributed concurrency guard). + +Closes the server-side race where two concurrent ``POST /invocations`` for the +same session run two agent loops against one AgentCore Memory session and +corrupt tool-pairing history (see docs/specs/session-single-flight-guard.md and +PR #653). A client-side abort does not propagate through the AgentCore Runtime +data plane, and the Runtime can route the duplicate to a *different* container, +so an in-process lock is insufficient β€” we need a distributed lease. + +Design: +- A dedicated item on the existing ``sessions-metadata`` table + (``PK=USER#{user_id}``, ``SK=LEASE#{session_id}``). The deterministic key + means acquisition is one atomic conditional write with no GSI read first. +- ``leaseExpiresAt`` (epoch seconds) is the application-level validity check + used in the acquire ``ConditionExpression``; the ``ttl`` attribute is only a + coarse DynamoDB auto-reap backstop for crashed-container orphans (TTL delete + lags up to 48h and must never be the correctness mechanism). +- Renewal (heartbeat) and release are owner-scoped so a container that already + lost the lease can neither extend nor delete the new owner's lease. + +Fail-open: every operation except a genuine lock *conflict* proceeds/degrades +silently. A throttle or transient DynamoDB error must never block a legitimate +turn β€” the guard is a safety net, not a gate. +""" + +from __future__ import annotations + +import logging +import os +import time +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Optional + +logger = logging.getLogger(__name__) + +# Lease validity window. A turn renews (heartbeats) well inside this; after a +# genuine container crash the lease self-expires within the window and the +# session becomes usable again. 90s tolerates one missed 30s heartbeat before +# expiry while keeping crash-recovery fast (well under the 600s stream timeout, +# so a normal long turn never self-evicts). +LEASE_WINDOW_SECONDS = 90 + +# Heartbeat cadence β€” how often the live turn renews ``leaseExpiresAt`` and +# observes a cancel request. Also bounds worst-case Stopβ†’resend latency: a +# cooperative stop is seen within one interval. Kept well under the window so a +# single missed tick never expires the lease. +LEASE_HEARTBEAT_SECONDS = 10 + +# DynamoDB ``ttl`` backstop: how long an orphaned lease item lingers before +# DynamoDB reaps it. Only prevents unbounded accumulation; correctness rides on +# ``leaseExpiresAt``. +LEASE_TTL_BACKSTOP_SECONDS = 3600 + + +class SessionBusyError(Exception): + """Raised on acquire when an unexpired lease is already held for the session. + + The caller (``/invocations``) maps this to HTTP 409. + """ + + +@dataclass(frozen=True) +class SessionLease: + """Handle for a held lease β€” carries the owner token used by renew/release.""" + + session_id: str + user_id: str + owner: str + + @property + def pk(self) -> str: + return f"USER#{self.user_id}" + + @property + def sk(self) -> str: + return f"LEASE#{self.session_id}" + + +def _table(): + """Return the sessions-metadata DynamoDB table, or ``None`` if unconfigured. + + Unconfigured (local dev without DynamoDB) is a valid state: the guard simply + doesn't run there, matching the best-effort posture of the other + session-metadata helpers. + """ + table_name = os.environ.get("DYNAMODB_SESSIONS_METADATA_TABLE_NAME") + if not table_name: + return None + import boto3 + + return boto3.resource("dynamodb").Table(table_name) + + +async def acquire_session_lease( + session_id: str, + user_id: str, + *, + force: bool = False, +) -> Optional[SessionLease]: + """Acquire the single-flight lease for a session's turn. + + Args: + session_id / user_id: the turn's session and its owning user. + force: when True (resume / max-tokens continuation), take the lease + unconditionally β€” those turns re-enter a loop that already ended and + must never be blocked, but still install a lease so a fresh duplicate + arriving *during* them is rejected. + + Returns: + A ``SessionLease`` on success, or ``None`` when the guard is inactive + (table unconfigured) or a non-conflict DynamoDB error occurred + (fail-open β€” the turn proceeds unguarded rather than being blocked on + lock-infra failure). + + Raises: + SessionBusyError: a non-forced acquire found an unexpired lease held by + another in-flight turn. + """ + table = _table() + if table is None: + return None + + from botocore.exceptions import ClientError + + owner = uuid.uuid4().hex + now = int(time.time()) + expires_at = now + LEASE_WINDOW_SECONDS + ttl = now + LEASE_TTL_BACKSTOP_SECONDS + + update_kwargs = { + "Key": {"PK": f"USER#{user_id}", "SK": f"LEASE#{session_id}"}, + # REMOVE clears any stale cancel marker when taking over an expired + # lease item, so a prior turn's cancel request can't bleed into this + # one. (Owner-scoping already protects β€” the new owner token won't match + # the old cancelRequestedFor β€” but clearing keeps the row honest.) + "UpdateExpression": ( + "SET leaseOwner = :owner, leaseExpiresAt = :exp, " + "#ttl = :ttl, updatedAt = :updated " + "REMOVE cancelRequestedFor, cancelRequestedAt" + ), + "ExpressionAttributeNames": {"#ttl": "ttl"}, + "ExpressionAttributeValues": { + ":owner": owner, + ":exp": expires_at, + ":ttl": ttl, + ":updated": datetime.now(timezone.utc).isoformat(), + }, + } + if not force: + # Win iff no lease exists or the existing one has lapsed. Two racing + # duplicates both evaluate this against the same item; DynamoDB + # serializes them so exactly one satisfies the condition. + update_kwargs["ConditionExpression"] = ( + "attribute_not_exists(PK) OR leaseExpiresAt < :now" + ) + update_kwargs["ExpressionAttributeValues"][":now"] = now + + try: + table.update_item(**update_kwargs) + except ClientError as e: + if e.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + raise SessionBusyError(session_id) from e + # Any other DynamoDB failure: fail open. Log and let the turn run + # unguarded β€” never block a legitimate turn on lock-infra trouble. + logger.error( + "Session lease acquire failed for %s (fail-open, proceeding unguarded): %s", + session_id, + e, + exc_info=True, + ) + return None + + logger.info( + "Acquired session lease for %s (owner=%s, force=%s)", + session_id, + owner, + force, + ) + return SessionLease(session_id=session_id, user_id=user_id, owner=owner) + + +async def renew_session_lease(lease: Optional[SessionLease]) -> bool: + """Extend our lease's window and observe a cancel request in one round-trip. + + Owner-conditional so a container that already lost the lease (its window + lapsed and another turn took over) can't clobber the new owner's window. + Returns ``True`` iff a cancel has been requested for *this* lease owner β€” + the caller flips the session manager's ``cancelled`` flag to stop the turn. + Best-effort: any error (including having lost ownership) returns ``False``. + """ + if lease is None: + return False + table = _table() + if table is None: + return False + + from botocore.exceptions import ClientError + + now = int(time.time()) + try: + resp = table.update_item( + Key={"PK": lease.pk, "SK": lease.sk}, + UpdateExpression="SET leaseExpiresAt = :exp, #ttl = :ttl, updatedAt = :updated", + ConditionExpression="leaseOwner = :owner", + ExpressionAttributeNames={"#ttl": "ttl"}, + ExpressionAttributeValues={ + ":exp": now + LEASE_WINDOW_SECONDS, + ":ttl": now + LEASE_TTL_BACKSTOP_SECONDS, + ":owner": lease.owner, + ":updated": datetime.now(timezone.utc).isoformat(), + }, + ReturnValues="ALL_NEW", + ) + attrs = resp.get("Attributes", {}) + # Owner-scoped: only honor a cancel aimed at *our* turn, so a stale + # marker from a prior turn on the same row can't stop us. + return attrs.get("cancelRequestedFor") == lease.owner + except ClientError as e: + if e.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + # We no longer own the lease (took too long, another turn took over). + # Nothing to renew; the live loop will still finish, but the session + # is no longer reserved for it. + logger.warning("Session lease renew skipped for %s β€” no longer owner", lease.session_id) + return False + logger.warning("Session lease renew failed for %s: %s", lease.session_id, e) + return False + + +async def release_session_lease(lease: Optional[SessionLease]) -> None: + """Release our lease at turn end. Best-effort, owner-scoped, idempotent. + + Owner-conditional delete so we never remove a lease a later turn legitimately + took over after our window lapsed. Safe to call twice (the second is a no-op + conditional miss) and safe to call with ``None``. + """ + if lease is None: + return + table = _table() + if table is None: + return + + from botocore.exceptions import ClientError + + try: + table.delete_item( + Key={"PK": lease.pk, "SK": lease.sk}, + ConditionExpression="leaseOwner = :owner", + ExpressionAttributeValues={":owner": lease.owner}, + ) + logger.info("Released session lease for %s", lease.session_id) + except ClientError as e: + if e.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + # Already released, expired-and-retaken, or never persisted β€” fine. + return + logger.warning("Session lease release failed for %s: %s", lease.session_id, e) + + +async def request_session_cancel(session_id: str, user_id: str) -> bool: + """Ask the turn currently holding this session's lease to stop. + + Called from the app-api ``user_stopped`` path (any container). Reads the + lease's current ``leaseOwner`` and stamps ``cancelRequestedFor = `` + so the running container observes it on its next heartbeat and unwinds the + turn. Owner-scoping is the safety property: the request names the *current* + owner, so if the turn has already ended and a new one started (new owner + token), the new turn ignores it β€” a stale Stop can never kill a later turn. + + Returns ``True`` if a cancel was armed against an active lease. ``False`` + when there is no active turn (no lease item) or the guard is inactive + (table unconfigured) β€” both mean "nothing running to stop." Best-effort: + any DynamoDB error returns ``False`` rather than raising. + """ + table = _table() + if table is None: + return False + + from botocore.exceptions import ClientError + + try: + resp = table.get_item( + Key={"PK": f"USER#{user_id}", "SK": f"LEASE#{session_id}"} + ) + except ClientError as e: + logger.warning("Session cancel lookup failed for %s: %s", session_id, e) + return False + + item = resp.get("Item") + owner = item.get("leaseOwner") if item else None + if not owner: + # No lease β†’ no turn is streaming server-side for this session. + return False + + try: + table.update_item( + Key={"PK": f"USER#{user_id}", "SK": f"LEASE#{session_id}"}, + UpdateExpression="SET cancelRequestedFor = :owner, cancelRequestedAt = :ts", + # Only arm if that same owner still holds the lease β€” otherwise the + # turn already ended/rotated and there is nothing to cancel. + ConditionExpression="leaseOwner = :owner", + ExpressionAttributeValues={ + ":owner": owner, + ":ts": datetime.now(timezone.utc).isoformat(), + }, + ) + logger.info("Armed cancel for session %s (owner=%s)", session_id, owner) + return True + except ClientError as e: + if e.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + # The turn ended or rotated between the read and the write β€” nothing + # to cancel. + return False + logger.warning("Session cancel arm failed for %s: %s", session_id, e) + return False diff --git a/backend/tests/agents/main_agent/session/test_turn_based_session_manager.py b/backend/tests/agents/main_agent/session/test_turn_based_session_manager.py index a0586a548..efd908e57 100644 --- a/backend/tests/agents/main_agent/session/test_turn_based_session_manager.py +++ b/backend/tests/agents/main_agent/session/test_turn_based_session_manager.py @@ -593,10 +593,14 @@ def test_existing_agent_compaction_no_checkpoint(self, make_session_manager, com mgr = make_session_manager(compaction_config=compaction_config) mgr._load_compaction_state = MagicMock(return_value=CompactionState()) + # A valid Converse history: the truncatable toolResult is preceded by + # its matching toolUse turn, so the restore-time pairing repair no-ops + # and this test exercises compaction/truncation in isolation. messages = [ make_user_message("q1"), make_assistant_message("a1"), make_user_message("q2"), + make_tool_use_message("t1", "search", {"q": "x"}), make_tool_result_message("t1", "x" * 200), # will be truncated ] session_agent = self._make_mock_session_agent() @@ -607,9 +611,10 @@ def test_existing_agent_compaction_no_checkpoint(self, make_session_manager, com agent = self._make_mock_agent() mgr.initialize(agent) - # All 4 messages kept (checkpoint=0), but truncation applied - assert len(agent.messages) == 4 - # Valid cutoffs cached for user text messages (indices 0, 2) + # All 5 messages kept (checkpoint=0), but truncation applied + assert len(agent.messages) == 5 + # Valid cutoffs cached for user text messages (indices 0, 2); the + # toolResult user turn at index 4 is not a valid cutoff. assert mgr._valid_cutoff_indices == [0, 2] def test_existing_agent_compaction_with_checkpoint_slices_messages(self, make_session_manager, compaction_config): @@ -642,10 +647,14 @@ def test_existing_agent_compaction_checkpoint_plus_truncation(self, make_session return_value=CompactionState(checkpoint=2) ) + # Valid Converse history: the truncatable toolResult follows its + # matching toolUse turn, so the pairing repair no-ops and the slice + # boundary lands on a clean turn. messages = [ make_user_message("old1"), make_assistant_message("old2"), make_user_message("new1"), + make_tool_use_message("t1", "search", {"q": "r"}), make_tool_result_message("t1", "r" * 200), # truncatable ] session_agent = self._make_mock_session_agent() @@ -656,8 +665,8 @@ def test_existing_agent_compaction_checkpoint_plus_truncation(self, make_session agent = self._make_mock_agent() mgr.initialize(agent) - # Sliced from index 2: 2 messages remain - assert len(agent.messages) == 2 + # Sliced from index 2: [new1, toolUse, toolResult] = 3 messages remain + assert len(agent.messages) == 3 def test_duplicate_agent_id_raises(self, make_session_manager): """Second initialize with same agent_id should raise SessionException.""" diff --git a/backend/tests/agents/main_agent/streaming/test_force_stop_persistence.py b/backend/tests/agents/main_agent/streaming/test_force_stop_persistence.py index 350d757b9..30e151632 100644 --- a/backend/tests/agents/main_agent/streaming/test_force_stop_persistence.py +++ b/backend/tests/agents/main_agent/streaming/test_force_stop_persistence.py @@ -310,3 +310,111 @@ def _extract_text(session_message: Any) -> str: content = msg.get("content", []) if isinstance(msg, dict) else [] parts = [block.get("text", "") for block in content if isinstance(block, dict)] return "".join(parts) + + +# --------------------------------------------------------------------------- +# Role-alternation guard: an error persisted after a DANGLING ASSISTANT turn +# must not create two consecutive assistant messages. +# +# This is the amplifier that permanently bricked prod session f761f59b: a turn +# ended with a dangling assistant toolUse (a duplicate concurrent invocation +# double-wrote tool results), an error fired, and the handler appended ANOTHER +# assistant turn β€” assistant, assistant β€” which then failed every subsequent +# turn, each failure persisting yet another consecutive assistant error. +# +# The write-side fix: both synthetic-error paths pass the tail role +# (agent.messages[-1]["role"]) to persist_synthetic_messages, which drops the +# write when it would land next to a same-role turn. The error stays a +# live-only UI affordance for that turn (mirroring the max_tokens path). +# --------------------------------------------------------------------------- + + +def _assistant_tail_agent_force_stop(reason: str) -> _FakeAgent: + """A force_stop agent whose history tail is a dangling assistant toolUse.""" + agent = _FakeAgent([_force_stop_event(reason)]) + agent.messages = [ + {"role": "user", "content": [{"text": "run the tool"}]}, + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t1", "name": "search", "input": {}}}], + }, + ] + return agent + + +@pytest.mark.asyncio +async def test_agent_error_skips_persist_when_tail_is_assistant(): + """AGENT_ERROR (force_stop) path: history tail is a dangling assistant + turn, so the synthetic assistant error must NOT be persisted β€” that would + create consecutive assistant messages and brick the session.""" + raw_reason = ( + "An error occurred (ValidationException) when calling the " + "ConverseStream operation: This model doesn't support documents." + ) + persist_sm = _RecordingPersistSessionManager() + + with patch( + "agents.main_agent.session.session_factory.SessionFactory.create_session_manager", + return_value=persist_sm, + ): + await _collect(_assistant_tail_agent_force_stop(raw_reason), _NoopSessionManager()) + + assert persist_sm.calls == [], ( + f"expected NO create_message call (assistant tail β†’ skip to preserve " + f"alternation), got {len(persist_sm.calls)}: {persist_sm.calls}" + ) + + +@pytest.mark.asyncio +async def test_stream_error_skips_persist_when_tail_is_assistant(): + """Emergency path (raw exception out of stream_async β†’ STREAM_ERROR): + history tail is a dangling assistant turn, so the synthetic error is + dropped rather than appended as a second consecutive assistant message.""" + exc = Exception( + "An error occurred (ValidationException) when calling the " + "ConverseStream operation: This model doesn't support documents." + ) + agent = _RaisingAgent(exc) + agent.messages = [ + {"role": "user", "content": [{"text": "run the tool"}]}, + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t1", "name": "search", "input": {}}}], + }, + ] + persist_sm = _RecordingPersistSessionManager() + + with patch( + "agents.main_agent.session.session_factory.SessionFactory.create_session_manager", + return_value=persist_sm, + ): + await _collect(agent, _NoopSessionManager()) + + assert persist_sm.calls == [], ( + f"expected NO create_message call (assistant tail β†’ skip), got " + f"{len(persist_sm.calls)}: {persist_sm.calls}" + ) + + +@pytest.mark.asyncio +async def test_agent_error_still_persists_when_tail_is_user(): + """Control: the healthy case (user turn is the tail, as at turn start) still + persists the assistant error exactly once β€” the guard only suppresses the + consecutive-assistant case, not normal error persistence.""" + raw_reason = ( + "An error occurred (ValidationException) when calling the " + "ConverseStream operation: This model doesn't support documents." + ) + # _FakeAgent's default tail is the user turn. + persist_sm = _RecordingPersistSessionManager() + + with patch( + "agents.main_agent.session.session_factory.SessionFactory.create_session_manager", + return_value=persist_sm, + ): + await _collect(_FakeAgent([_force_stop_event(raw_reason)]), _NoopSessionManager()) + + assert len(persist_sm.calls) == 1 + inner = getattr(persist_sm.calls[0]["message"], "message", None) + role = inner.get("role") if isinstance(inner, dict) else None + assert role == "assistant" diff --git a/backend/tests/agents/main_agent/streaming/test_interrupted_turn_persistence.py b/backend/tests/agents/main_agent/streaming/test_interrupted_turn_persistence.py index 01cc974d5..744288b47 100644 --- a/backend/tests/agents/main_agent/streaming/test_interrupted_turn_persistence.py +++ b/backend/tests/agents/main_agent/streaming/test_interrupted_turn_persistence.py @@ -249,6 +249,40 @@ async def _fake_set_interrupted(session_id, user_id, reason="unknown", source="c assert marker_calls == [{"reason": "connection_lost"}] +@pytest.mark.asyncio +async def test_nonempty_partial_skips_write_when_tail_is_assistant(): + """An interrupted continuation/resume: the tail is already an assistant + message (the one being extended) and a partial streamed before teardown. + Persisting the partial as a NEW assistant turn would create consecutive + assistant messages and brick the session, so it is SKIPPED β€” the partial + stays a live-only affordance and only the marker is set.""" + persist_sm = _RecordingPersistSessionManager() + marker_calls: List[Dict[str, Any]] = [] + + async def _fake_set_interrupted(session_id, user_id, reason="unknown", source="cancellation"): + marker_calls.append({"reason": reason}) + + agent = _InterruptingAgent() + agent.messages = [ + {"role": "user", "content": [{"text": "hi"}]}, + {"role": "assistant", "content": [{"text": "resuming the truncated answer"}]}, + ] + coordinator = StreamCoordinator() + with patch( + "agents.main_agent.session.session_factory.SessionFactory.create_session_manager", + return_value=persist_sm, + ), patch("apis.shared.sessions.metadata.set_interrupted_turn", _fake_set_interrupted): + await coordinator._persist_interruption( + agent=agent, + session_id="sess-interrupt", + user_id="user-1", + partial_text="…and here is the continuation", + ) + + assert persist_sm.calls == [] + assert marker_calls == [{"reason": "connection_lost"}] + + @pytest.mark.asyncio async def test_interruption_persists_partial_turn_metadata(): """A Stop with a partial persists per-message metadata (which also bumps @@ -334,3 +368,97 @@ async def _fake_store_metadata(self, **kwargs): assert store_calls == [] assert len(persist_sm.calls) == 1 # partial still persisted + + +# --------------------------------------------------------------------------- +# Cooperative stop (distributed cancellation): a user Stop is observed via the +# session lease and flips session_manager.cancelled while the turn is still +# streaming server-side (the client abort never reached this process). Unlike +# the CancelledError/GeneratorExit backstop, this ends the stream CLEANLY +# (persist + terminal frames, no re-raise) so the lease releases and resend works. +# --------------------------------------------------------------------------- + + +class _CancellingSessionManager: + """Session manager whose ``cancelled`` flag the running agent flips mid-stream.""" + + def __init__(self) -> None: + self.cancelled = False + + async def update_after_turn(self, input_tokens: int, current_messages=None): + return None + + +@pytest.mark.asyncio +async def test_cooperative_stop_persists_partial_and_ends_cleanly(): + sm = _CancellingSessionManager() + + class _Agent: + def __init__(self) -> None: + self.messages = [{"role": "user", "content": [{"text": "hi"}]}] + + def stream_async(self, prompt: Any) -> AsyncIterator[Dict[str, Any]]: + async def _gen() -> AsyncIterator[Dict[str, Any]]: + yield _raw_message_start() + yield _raw_text_delta("Hello") + yield _raw_text_delta(" world") + # User clicks Stop; the heartbeat flips the flag. The NEXT + # loop-top check must end the turn before "!" is accumulated. + sm.cancelled = True + yield _raw_text_delta("!") + yield _raw_message_stop() + + return _gen() + + captured: Dict[str, Any] = {} + + async def _fake_persist(self, *, agent, session_id, user_id, partial_text, reason="connection_lost", **kwargs): + captured.update(partial_text=partial_text, reason=reason, session_id=session_id) + + sse: List[str] = [] + coordinator = StreamCoordinator() + with patch.object(StreamCoordinator, "_persist_interruption", _fake_persist): + # No exception escapes β€” the stream ends cleanly (contrast with the + # CancelledError arm, which re-raises). + async for chunk in coordinator.stream_response( + agent=_Agent(), + prompt="write an essay", + session_manager=sm, + session_id="sess-coop", + user_id="user-1", + main_agent_wrapper=None, + ): + sse.append(chunk) + + # Partial covers what streamed before Stop, not the post-stop delta. + assert captured["partial_text"] == "Hello world" + # Marked as a deliberate stop, not the connection_lost fallback. + assert captured["reason"] == "user_stopped" + assert captured["session_id"] == "sess-coop" + # A clean terminal frame was emitted for any still-connected client. + assert "event: done" in "".join(sse) + + +@pytest.mark.asyncio +async def test_persist_interruption_threads_user_stopped_reason(): + """The cooperative arm passes reason=user_stopped through to the marker.""" + persist_sm = _RecordingPersistSessionManager() + marker_calls: List[Dict[str, Any]] = [] + + async def _fake_set_interrupted(session_id, user_id, reason="unknown", source="cancellation"): + marker_calls.append({"reason": reason, "source": source}) + + coordinator = StreamCoordinator() + with patch( + "agents.main_agent.session.session_factory.SessionFactory.create_session_manager", + return_value=persist_sm, + ), patch("apis.shared.sessions.metadata.set_interrupted_turn", _fake_set_interrupted): + await coordinator._persist_interruption( + agent=_InterruptingAgent(), + session_id="sess-interrupt", + user_id="user-1", + partial_text="partial answer", + reason="user_stopped", + ) + + assert marker_calls == [{"reason": "user_stopped", "source": "cancellation"}] diff --git a/backend/tests/apis/app_api/shares/test_share_s3_offload.py b/backend/tests/apis/app_api/shares/test_share_s3_offload.py new file mode 100644 index 000000000..13a8efe5d --- /dev/null +++ b/backend/tests/apis/app_api/shares/test_share_s3_offload.py @@ -0,0 +1,254 @@ +"""Tests for the S3 snapshot-body offload in ShareService. + +Covers the regression this feature fixes β€” a conversation too large to inline +in a DynamoDB item (>400 KB) can now be shared β€” plus the write shape +(``body_ref``, no inline ``messages``), the S3-backed and legacy-inline read +paths, revoke cleanup, and the storage-unavailable guard. +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import boto3 +import pytest +from moto import mock_aws + +from apis.app_api.shares.service import ( + ShareService, + ShareStorageUnavailableError, +) +from apis.app_api.shares.snapshot_store import ShareSnapshotStore +from apis.app_api.shares.models import CreateShareRequest +from apis.shared.auth.models import User + +AWS_REGION = "us-east-1" +BUCKET = "test-shared-conversations" + + +@pytest.fixture() +def aws_env(monkeypatch): + monkeypatch.setenv("AWS_DEFAULT_REGION", AWS_REGION) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing") + monkeypatch.setenv("AWS_SESSION_TOKEN", "testing") + with mock_aws(): + yield + + +@pytest.fixture() +def s3_client(aws_env): + client = boto3.client("s3", region_name=AWS_REGION) + client.create_bucket(Bucket=BUCKET) + return client + + +@pytest.fixture() +def store(s3_client): + return ShareSnapshotStore(bucket_name=BUCKET, s3_client=s3_client) + + +@pytest.fixture() +def service(store, monkeypatch): + """A ShareService with a real moto S3 store and a mock DynamoDB table.""" + monkeypatch.setenv("SHARED_CONVERSATIONS_TABLE_NAME", "shares-table") + with patch("boto3.resource"): + svc = ShareService(snapshot_store=store) + svc._table = MagicMock() + return svc + + +def _owner() -> User: + return User(email="owner@example.com", user_id="owner-1", name="Owner", roles=["User"]) + + +def _message_dict(text: str, msg_id: str = "m0") -> dict: + return { + "id": msg_id, + "role": "user", + "content": [{"type": "text", "text": text}], + "createdAt": "2025-06-01T00:00:00Z", + } + + +def _patch_snapshot_sources(metadata_dump: dict, message_dumps: list): + """Patch get_session_metadata + get_messages used by create_share.""" + meta = MagicMock() + meta.model_dump.return_value = metadata_dump + + messages = [MagicMock(model_dump=MagicMock(return_value=m)) for m in message_dumps] + messages_response = MagicMock(messages=messages) + + return ( + patch( + "apis.app_api.shares.service.get_session_metadata", + new=AsyncMock(return_value=meta), + ), + patch( + "apis.app_api.shares.service.get_messages", + new=AsyncMock(return_value=messages_response), + ), + ) + + +class TestCreateOffloadsToS3: + @pytest.mark.asyncio + async def test_item_carries_body_ref_not_inline_messages(self, service, s3_client): + meta_patch, msgs_patch = _patch_snapshot_sources( + {"title": "Hello Chat"}, [_message_dict("hi")] + ) + with meta_patch, msgs_patch: + resp = await service.create_share( + "sess-1", _owner(), CreateShareRequest(accessLevel="public") + ) + + # The DynamoDB item was written with a body_ref and NO inline body. + item = service._table.put_item.call_args[1]["Item"] + assert "messages" not in item + assert "metadata" not in item + assert item["body_ref"]["bucket_key"].startswith("shares/") + assert item["body_ref"]["format"] == "json" + assert item["body_ref"]["byte_size"] > 0 + assert resp.session_id == "sess-1" + + # The S3 object decodes to the full snapshot body. + obj = s3_client.get_object(Bucket=BUCKET, Key=item["body_ref"]["bucket_key"]) + body = json.loads(obj["Body"].read()) + assert body["metadata"]["title"] == "Hello Chat" + assert body["messages"][0]["content"][0]["text"] == "hi" + + @pytest.mark.asyncio + async def test_large_conversation_succeeds(self, service): + """Regression: a >400 KB snapshot no longer fails on PutItem.""" + big = [_message_dict("x" * 5000, msg_id=f"m{i}") for i in range(120)] + # Sanity: the inline body would have blown DynamoDB's 400 KB item cap. + assert len(json.dumps({"metadata": {}, "messages": big})) > 400_000 + + meta_patch, msgs_patch = _patch_snapshot_sources({"title": "Big"}, big) + with meta_patch, msgs_patch: + resp = await service.create_share( + "sess-big", _owner(), CreateShareRequest(accessLevel="public") + ) + + assert resp.session_id == "sess-big" + item = service._table.put_item.call_args[1]["Item"] + assert item["body_ref"]["byte_size"] > 400_000 + + @pytest.mark.asyncio + async def test_storage_unavailable_raises(self, monkeypatch): + monkeypatch.setenv("SHARED_CONVERSATIONS_TABLE_NAME", "shares-table") + with patch("boto3.resource"): + svc = ShareService(snapshot_store=ShareSnapshotStore(bucket_name=None)) + svc._table = MagicMock() + + meta_patch, msgs_patch = _patch_snapshot_sources({"title": "T"}, [_message_dict("hi")]) + with meta_patch, msgs_patch: + with pytest.raises(ShareStorageUnavailableError): + await svc.create_share( + "sess-1", _owner(), CreateShareRequest(accessLevel="public") + ) + # Nothing was written to DynamoDB when storage is unavailable. + svc._table.put_item.assert_not_called() + + +class TestReadPaths: + @pytest.mark.asyncio + async def test_s3_backed_read_round_trips(self, service, store): + # Create, then read back through get_shared_conversation. + meta_patch, msgs_patch = _patch_snapshot_sources( + {"title": "Round Trip"}, [_message_dict("hello", "m1")] + ) + with meta_patch, msgs_patch: + await service.create_share( + "sess-1", _owner(), CreateShareRequest(accessLevel="public") + ) + written = service._table.put_item.call_args[1]["Item"] + + with patch.object(service, "_get_share_item", return_value=written): + result = await service.get_shared_conversation("share-x", _owner()) + + assert result.title == "Round Trip" + assert len(result.messages) == 1 + assert result.messages[0].content[0].text == "hello" + + @pytest.mark.asyncio + async def test_legacy_inline_read_still_works(self, service): + """A pre-offload item (inline messages, no body_ref) still resolves.""" + legacy_item = { + "share_id": "share-legacy", + "session_id": "sess-old", + "owner_id": "owner-1", + "access_level": "public", + "created_at": "2025-01-01T00:00:00Z", + "metadata": {"title": "Old One"}, + "messages": [_message_dict("legacy hi", "m0")], + } + with patch.object(service, "_get_share_item", return_value=legacy_item): + result = await service.get_shared_conversation("share-legacy", _owner()) + + assert result.title == "Old One" + assert result.messages[0].content[0].text == "legacy hi" + + @pytest.mark.asyncio + async def test_malformed_item_raises_not_found(self, service): + from apis.app_api.shares.service import ShareNotFoundError + + bad = { + "share_id": "share-bad", + "session_id": "s", + "owner_id": "owner-1", + "access_level": "public", + "created_at": "2025-01-01T00:00:00Z", + } + with patch.object(service, "_get_share_item", return_value=bad): + with pytest.raises(ShareNotFoundError): + await service.get_shared_conversation("share-bad", _owner()) + + + @pytest.mark.asyncio + async def test_export_reads_s3_backed_body(self, service): + """export_shared_conversation loads the body from S3, not inline.""" + meta_patch, msgs_patch = _patch_snapshot_sources( + {"title": "Export Me"}, + [_message_dict("one", "m0"), _message_dict("two", "m1")], + ) + with meta_patch, msgs_patch: + await service.create_share( + "sess-1", _owner(), CreateShareRequest(accessLevel="public") + ) + item = service._table.put_item.call_args[1]["Item"] + + with patch.object(service, "_get_share_item", return_value=item), \ + patch.object(service, "_check_access"), \ + patch.object( + service, "_copy_messages_to_memory", + new_callable=AsyncMock, return_value=2, + ) as mock_copy, \ + patch( + "apis.app_api.shares.service.store_session_metadata", + new_callable=AsyncMock, + ): + result = await service.export_shared_conversation("share-x", _owner()) + + assert result["title"] == "Export Me (shared)" + # The two S3-stored messages were the ones handed to the memory copy. + assert len(mock_copy.call_args[0][2]) == 2 + + +class TestRevokeCleansUpS3: + @pytest.mark.asyncio + async def test_revoke_deletes_s3_object(self, service, s3_client): + meta_patch, msgs_patch = _patch_snapshot_sources({"title": "T"}, [_message_dict("hi")]) + with meta_patch, msgs_patch: + await service.create_share( + "sess-1", _owner(), CreateShareRequest(accessLevel="public") + ) + item = service._table.put_item.call_args[1]["Item"] + key = item["body_ref"]["bucket_key"] + # Object exists before revoke. + assert s3_client.get_object(Bucket=BUCKET, Key=key)["Body"].read() + + with patch.object(service, "_get_share_item", return_value=item): + await service.revoke_share(item["share_id"], _owner()) + + with pytest.raises(s3_client.exceptions.NoSuchKey): + s3_client.get_object(Bucket=BUCKET, Key=key) diff --git a/backend/tests/apis/app_api/shares/test_snapshot_store.py b/backend/tests/apis/app_api/shares/test_snapshot_store.py new file mode 100644 index 000000000..78bcd87e6 --- /dev/null +++ b/backend/tests/apis/app_api/shares/test_snapshot_store.py @@ -0,0 +1,100 @@ +"""Tests for the S3-backed share snapshot-body store. + +moto-backed: exercises content-hash keying, dedupe (no second put for +identical bytes), get/delete round-trip, and the not-configured guard. +Mirrors ``tests/shared/test_memory_store.py``. +""" + +import boto3 +import pytest +from moto import mock_aws + +from apis.app_api.shares.snapshot_store import ( + ShareSnapshotStore, + ShareSnapshotStoreError, + compute_content_hash, + content_key, +) + +AWS_REGION = "us-east-1" +BUCKET = "test-shared-conversations" + + +@pytest.fixture() +def aws_env(monkeypatch): + monkeypatch.setenv("AWS_DEFAULT_REGION", AWS_REGION) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing") + monkeypatch.setenv("AWS_SESSION_TOKEN", "testing") + with mock_aws(): + yield + + +@pytest.fixture() +def s3_client(aws_env): + client = boto3.client("s3", region_name=AWS_REGION) + client.create_bucket(Bucket=BUCKET) + return client + + +@pytest.fixture() +def store(s3_client): + return ShareSnapshotStore(bucket_name=BUCKET, s3_client=s3_client) + + +class TestContentKey: + def test_key_is_content_addressed(self): + digest = compute_content_hash(b"body") + assert content_key("share-1", digest) == f"shares/share-1/{digest}" + + def test_hash_is_stable_and_distinct(self): + assert compute_content_hash(b"a") == compute_content_hash(b"a") + assert compute_content_hash(b"a") != compute_content_hash(b"b") + + +class TestPutGetDelete: + def test_put_returns_content_addressed_key(self, store): + key = store.put(share_id="share-1", body=b'{"messages":[]}') + assert key == content_key("share-1", compute_content_hash(b'{"messages":[]}')) + + def test_get_round_trip(self, store): + key = store.put(share_id="share-1", body=b"hello-body") + assert store.get(key) == b"hello-body" + + def test_put_is_idempotent_dedupe(self, store, s3_client): + k1 = store.put(share_id="s", body=b"same") + k2 = store.put(share_id="s", body=b"same") + assert k1 == k2 + # Exactly one object under the prefix. + listing = s3_client.list_objects_v2(Bucket=BUCKET, Prefix="shares/s/") + assert listing["KeyCount"] == 1 + + def test_get_missing_key_raises(self, store): + with pytest.raises(ShareSnapshotStoreError): + store.get("shares/nope/deadbeef") + + def test_delete_is_best_effort(self, store): + key = store.put(share_id="s", body=b"x") + store.delete(key) + # Deleting an already-absent object is a no-op (never raises). + store.delete(key) + with pytest.raises(ShareSnapshotStoreError): + store.get(key) + + +class TestEnabledGuard: + def test_disabled_when_bucket_unset(self, monkeypatch): + monkeypatch.delenv("SHARED_CONVERSATIONS_BUCKET_NAME", raising=False) + assert ShareSnapshotStore(bucket_name=None).enabled is False + + def test_put_raises_when_disabled(self, monkeypatch): + monkeypatch.delenv("SHARED_CONVERSATIONS_BUCKET_NAME", raising=False) + store = ShareSnapshotStore(bucket_name=None) + with pytest.raises(ShareSnapshotStoreError): + store.put(share_id="s", body=b"x") + + def test_delete_when_disabled_is_noop(self, monkeypatch): + monkeypatch.delenv("SHARED_CONVERSATIONS_BUCKET_NAME", raising=False) + store = ShareSnapshotStore(bucket_name=None) + # No exception even though storage is unconfigured. + store.delete("shares/s/abc") diff --git a/backend/tests/apis/app_api/web_sources/test_deletion_service.py b/backend/tests/apis/app_api/web_sources/test_deletion_service.py new file mode 100644 index 000000000..b3c177d79 --- /dev/null +++ b/backend/tests/apis/app_api/web_sources/test_deletion_service.py @@ -0,0 +1,209 @@ +"""Unit tests for the web-source deletion cascade. + +The cascade's whole job is deciding *which* documents belong to a crawl and +tearing them down in an order that can't strand state, so that's what these +assert: prefix scoping, page-by-page soft-delete, sync-policy removal, and the +crawl row going last. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, patch + +import pytest + +from apis.app_api.documents.models import Document +from apis.app_api.web_sources.deletion_service import ( + WebSourceDeletionError, + delete_web_source, +) + +ASSISTANT_ID = "ast-1" +OWNER_ID = "user-1" +CRAWL_ID = "CRAWL-1" +ROOT_URL = "https://example.com/docs/" + +MODULE = "apis.app_api.web_sources.deletion_service" + + +def _doc( + document_id: str, + *, + source_file_id: str | None = None, + connector: str | None = "web", + status: str = "complete", +) -> Document: + return Document.model_validate( + { + "documentId": document_id, + "assistantId": ASSISTANT_ID, + "filename": f"{document_id}.html", + "contentType": "text/html", + "sizeBytes": 10, + "s3Key": f"assistants/{ASSISTANT_ID}/documents/{document_id}/page.html", + "status": status, + "chunkCount": 2, + "sourceConnectorId": connector, + "sourceFileId": source_file_id, + "createdAt": "2026-07-14T00:00:00Z", + "updatedAt": "2026-07-14T00:00:00Z", + } + ) + + +@pytest.fixture +def patched(): + """Patch every collaborator the cascade reaches out to.""" + with patch( + f"{MODULE}.list_assistant_documents", new_callable=AsyncMock + ) as list_docs, patch( + f"{MODULE}.soft_delete_document", new_callable=AsyncMock + ) as soft_delete, patch( + f"{MODULE}.hard_delete_crawl_job", new_callable=AsyncMock, return_value=True + ) as hard_delete, patch( + f"{MODULE}.delete_sync_policies_for_source", + new_callable=AsyncMock, + return_value=0, + ) as delete_policies, patch( + f"{MODULE}.cleanup_assistant_documents", new_callable=AsyncMock + ) as cleanup: + soft_delete.side_effect = lambda assistant_id, document_id, owner_id: _doc( + document_id, source_file_id=f"{ROOT_URL}{document_id}" + ) + yield { + "list_docs": list_docs, + "soft_delete": soft_delete, + "hard_delete": hard_delete, + "delete_policies": delete_policies, + "cleanup": cleanup, + } + + +async def _run(patched) -> int: + removed = await delete_web_source( + assistant_id=ASSISTANT_ID, + crawl_id=CRAWL_ID, + root_url=ROOT_URL, + owner_id=OWNER_ID, + ) + # The cleanup fan-out is a background task; let it start so the assertion + # on `cleanup_assistant_documents` isn't racing the event loop. + await asyncio.sleep(0) + return removed + + +@pytest.mark.asyncio +async def test_deletes_only_pages_under_the_crawl_root(patched): + patched["list_docs"].return_value = ( + [ + _doc("DOC-1", source_file_id=f"{ROOT_URL}intro"), + _doc("DOC-2", source_file_id=f"{ROOT_URL}guide/setup"), + # Same site, different crawl root β€” must survive. + _doc("DOC-3", source_file_id="https://example.com/blog/post"), + # A device upload β€” no provenance at all. + _doc("DOC-4", connector=None, source_file_id=None), + # A Drive import that happens to sit under a similar path. + _doc("DOC-5", connector="google_drive", source_file_id=f"{ROOT_URL}sheet"), + ], + None, + ) + + removed = await _run(patched) + + assert removed == 2 + deleted_ids = {c.args[1] for c in patched["soft_delete"].await_args_list} + assert deleted_ids == {"DOC-1", "DOC-2"} + + +@pytest.mark.asyncio +async def test_skips_pages_already_being_deleted(patched): + patched["list_docs"].return_value = ( + [ + _doc("DOC-1", source_file_id=f"{ROOT_URL}intro"), + _doc("DOC-2", source_file_id=f"{ROOT_URL}gone", status="deleting"), + ], + None, + ) + + removed = await _run(patched) + + assert removed == 1 + patched["soft_delete"].assert_awaited_once() + + +@pytest.mark.asyncio +async def test_walks_every_page_of_results(patched): + patched["list_docs"].side_effect = [ + ([_doc("DOC-1", source_file_id=f"{ROOT_URL}a")], "token-2"), + ([_doc("DOC-2", source_file_id=f"{ROOT_URL}b")], None), + ] + + removed = await _run(patched) + + assert removed == 2 + assert patched["list_docs"].await_count == 2 + assert patched["list_docs"].await_args_list[1].kwargs["next_token"] == "token-2" + + +@pytest.mark.asyncio +async def test_removes_the_sync_policy_before_the_pages(patched): + """A dispatcher sweep landing mid-delete would re-crawl the source and + resurrect the very pages we're removing.""" + order: list[str] = [] + patched["list_docs"].return_value = ( + [_doc("DOC-1", source_file_id=f"{ROOT_URL}a")], + None, + ) + patched["delete_policies"].side_effect = lambda *a: order.append("policy") or 1 + patched["soft_delete"].side_effect = lambda assistant_id, document_id, owner_id: ( + order.append("page") or _doc(document_id, source_file_id=f"{ROOT_URL}a") + ) + patched["hard_delete"].side_effect = lambda *a: order.append("crawl") or True + + await _run(patched) + + assert order == ["policy", "page", "crawl"] + patched["delete_policies"].assert_awaited_once_with(ASSISTANT_ID, CRAWL_ID) + + +@pytest.mark.asyncio +async def test_hands_the_pages_to_background_cleanup(patched): + patched["list_docs"].return_value = ( + [_doc("DOC-1", source_file_id=f"{ROOT_URL}a")], + None, + ) + + await _run(patched) + + patched["cleanup"].assert_awaited_once() + assistant_id, documents = patched["cleanup"].await_args.args + assert assistant_id == ASSISTANT_ID + assert [d.document_id for d in documents] == ["DOC-1"] + + +@pytest.mark.asyncio +async def test_removes_a_crawl_that_produced_no_pages(patched): + """A crawl that failed on its root page has nothing to sweep, but the row + itself still has to go β€” otherwise it's undeletable from the UI.""" + patched["list_docs"].return_value = ([], None) + + removed = await _run(patched) + + assert removed == 0 + patched["hard_delete"].assert_awaited_once_with(ASSISTANT_ID, CRAWL_ID) + patched["cleanup"].assert_not_awaited() + + +@pytest.mark.asyncio +async def test_raises_when_the_crawl_row_survives(patched): + patched["list_docs"].return_value = ([], None) + patched["hard_delete"].return_value = False + + with pytest.raises(WebSourceDeletionError): + await delete_web_source( + assistant_id=ASSISTANT_ID, + crawl_id=CRAWL_ID, + root_url=ROOT_URL, + owner_id=OWNER_ID, + ) diff --git a/backend/tests/apis/app_api/web_sources/test_routes.py b/backend/tests/apis/app_api/web_sources/test_routes.py index 588321018..87e5f6d89 100644 --- a/backend/tests/apis/app_api/web_sources/test_routes.py +++ b/backend/tests/apis/app_api/web_sources/test_routes.py @@ -8,6 +8,7 @@ from __future__ import annotations +from types import SimpleNamespace from unittest.mock import AsyncMock, patch import pytest @@ -16,6 +17,7 @@ from apis.app_api.documents.models import Document from apis.app_api.web_sources import routes as web_routes +from apis.app_api.web_sources.deletion_service import WebSourceDeletionError from apis.app_api.web_sources.models import CrawlJob, CrawlSettings from apis.shared.auth.models import User from apis.shared.auth.dependencies import get_current_user_from_session @@ -66,6 +68,11 @@ def _stub_crawl(crawl_id: str = "CRAWL-1") -> CrawlJob: ) +def _stub_permission(permission: str = "owner"): + """(assistant, permission) as `resolve_assistant_permission` returns it.""" + return SimpleNamespace(owner_id=USER_ID), permission + + @pytest.fixture def app() -> FastAPI: _app = FastAPI() @@ -80,9 +87,9 @@ def test_returns_202_with_root_document_and_crawl(self, app: FastAPI): crawl = _stub_crawl() run_crawl_mock = AsyncMock(return_value=None) with patch( - "apis.app_api.web_sources.routes.get_assistant", + "apis.app_api.web_sources.routes.resolve_assistant_permission", new_callable=AsyncMock, - return_value={"assistantId": ASSISTANT_ID}, + return_value=_stub_permission("owner"), ), patch( "apis.app_api.web_sources.routes.create_document", new_callable=AsyncMock, @@ -113,12 +120,68 @@ def test_returns_202_with_root_document_and_crawl(self, app: FastAPI): # exercising the real crawler. run_crawl_mock.assert_called_once() - def test_returns_404_when_assistant_not_owned(self, app: FastAPI): + def test_editor_may_start_a_crawl(self, app: FastAPI): + """An editor share is enough to add web content β€” the SPA already + renders the "Add web content" button for anyone who isn't a viewer. + """ + editor = User( + email="editor@example.com", user_id="user-editor", name="E", roles=["User"] + ) + mock_auth_user(app, editor) + run_crawl_mock = AsyncMock(return_value=None) + create_doc_mock = AsyncMock(return_value=_stub_document()) + create_job_mock = AsyncMock(return_value=_stub_crawl()) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission("editor"), + ), patch( + "apis.app_api.web_sources.routes.create_document", create_doc_mock, + ), patch( + "apis.app_api.web_sources.routes.create_crawl_job", create_job_mock, + ), patch( + "apis.app_api.web_sources.routes.run_crawl", run_crawl_mock, + ), patch( + "apis.app_api.web_sources.routes.assert_url_is_public", + return_value="https://example.com/", + ): + client = TestClient(app) + resp = client.post( + f"/assistants/{ASSISTANT_ID}/web-sources/crawl", + json={"url": "https://example.com/"}, + ) + assert resp.status_code == 202 + + # The document row is keyed on the assistant, so the editor's crawl + # lands under the same assistant the owner's would... + assert create_doc_mock.await_args.kwargs["assistant_id"] == ASSISTANT_ID + # ...and the actor fields record the *editor*, not the owner β€” that + # is the whole point of tracking who imported/started. + provenance = create_doc_mock.await_args.kwargs["provenance"] + assert provenance.imported_by_user_id == "user-editor" + assert create_job_mock.await_args.kwargs["started_by_user_id"] == "user-editor" + assert run_crawl_mock.call_args.kwargs["user_id"] == "user-editor" + + def test_viewer_gets_403(self, app: FastAPI): mock_auth_user(app, _user()) with patch( - "apis.app_api.web_sources.routes.get_assistant", + "apis.app_api.web_sources.routes.resolve_assistant_permission", new_callable=AsyncMock, - return_value=None, + return_value=_stub_permission("viewer"), + ): + client = TestClient(app) + resp = client.post( + f"/assistants/{ASSISTANT_ID}/web-sources/crawl", + json={"url": "https://example.com/"}, + ) + assert resp.status_code == 403 + + def test_returns_404_when_assistant_not_visible(self, app: FastAPI): + mock_auth_user(app, _user()) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=(None, None), ): client = TestClient(app) resp = client.post( @@ -130,9 +193,9 @@ def test_returns_404_when_assistant_not_owned(self, app: FastAPI): def test_returns_422_on_invalid_url(self, app: FastAPI): mock_auth_user(app, _user()) with patch( - "apis.app_api.web_sources.routes.get_assistant", + "apis.app_api.web_sources.routes.resolve_assistant_permission", new_callable=AsyncMock, - return_value={"assistantId": ASSISTANT_ID}, + return_value=_stub_permission("owner"), ): client = TestClient(app) # Loopback URL β†’ SSRF guard rejects it. @@ -145,9 +208,9 @@ def test_returns_422_on_invalid_url(self, app: FastAPI): def test_returns_422_on_bad_settings_bounds(self, app: FastAPI): mock_auth_user(app, _user()) with patch( - "apis.app_api.web_sources.routes.get_assistant", + "apis.app_api.web_sources.routes.resolve_assistant_permission", new_callable=AsyncMock, - return_value={"assistantId": ASSISTANT_ID}, + return_value=_stub_permission("owner"), ): client = TestClient(app) # max_pages above the cap @@ -175,9 +238,9 @@ def test_returns_active_jobs(self, app: FastAPI): mock_auth_user(app, _user()) crawl = _stub_crawl() with patch( - "apis.app_api.web_sources.routes.get_assistant", + "apis.app_api.web_sources.routes.resolve_assistant_permission", new_callable=AsyncMock, - return_value={"assistantId": ASSISTANT_ID}, + return_value=_stub_permission("owner"), ), patch( "apis.app_api.web_sources.routes.list_active_crawls", new_callable=AsyncMock, @@ -193,15 +256,42 @@ def test_returns_active_jobs(self, app: FastAPI): assert len(body["crawls"]) == 1 assert body["crawls"][0]["crawlId"] == "CRAWL-1" + def test_editor_may_list_crawls(self, app: FastAPI): + mock_auth_user(app, _user()) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission("editor"), + ), patch( + "apis.app_api.web_sources.routes.list_all_crawls", + new_callable=AsyncMock, + return_value=[_stub_crawl()], + ): + client = TestClient(app) + resp = client.get(f"/assistants/{ASSISTANT_ID}/web-sources/crawls") + assert resp.status_code == 200 + assert len(resp.json()["crawls"]) == 1 + + def test_viewer_gets_403(self, app: FastAPI): + mock_auth_user(app, _user()) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission("viewer"), + ): + client = TestClient(app) + resp = client.get(f"/assistants/{ASSISTANT_ID}/web-sources/crawls") + assert resp.status_code == 403 + def test_returns_all_jobs_without_active_filter(self, app: FastAPI): """Default (no ?active=true) is full history β€” the sync-policy UI lists completed crawls as syncable sources.""" mock_auth_user(app, _user()) crawl = _stub_crawl() with patch( - "apis.app_api.web_sources.routes.get_assistant", + "apis.app_api.web_sources.routes.resolve_assistant_permission", new_callable=AsyncMock, - return_value={"assistantId": ASSISTANT_ID}, + return_value=_stub_permission("owner"), ), patch( "apis.app_api.web_sources.routes.list_all_crawls", new_callable=AsyncMock, @@ -220,9 +310,9 @@ def test_returns_single_crawl(self, app: FastAPI): mock_auth_user(app, _user()) crawl = _stub_crawl() with patch( - "apis.app_api.web_sources.routes.get_assistant", + "apis.app_api.web_sources.routes.resolve_assistant_permission", new_callable=AsyncMock, - return_value={"assistantId": ASSISTANT_ID}, + return_value=_stub_permission("owner"), ), patch( "apis.app_api.web_sources.routes.get_crawl_job", new_callable=AsyncMock, @@ -235,12 +325,42 @@ def test_returns_single_crawl(self, app: FastAPI): assert resp.status_code == 200 assert resp.json()["crawlId"] == "CRAWL-1" + def test_editor_may_get_a_crawl(self, app: FastAPI): + mock_auth_user(app, _user()) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission("editor"), + ), patch( + "apis.app_api.web_sources.routes.get_crawl_job", + new_callable=AsyncMock, + return_value=_stub_crawl(), + ): + client = TestClient(app) + resp = client.get( + f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-1" + ) + assert resp.status_code == 200 + + def test_viewer_gets_403(self, app: FastAPI): + mock_auth_user(app, _user()) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission("viewer"), + ): + client = TestClient(app) + resp = client.get( + f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-1" + ) + assert resp.status_code == 403 + def test_returns_404_when_missing(self, app: FastAPI): mock_auth_user(app, _user()) with patch( - "apis.app_api.web_sources.routes.get_assistant", + "apis.app_api.web_sources.routes.resolve_assistant_permission", new_callable=AsyncMock, - return_value={"assistantId": ASSISTANT_ID}, + return_value=_stub_permission("owner"), ), patch( "apis.app_api.web_sources.routes.get_crawl_job", new_callable=AsyncMock, @@ -251,3 +371,173 @@ def test_returns_404_when_missing(self, app: FastAPI): f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-X" ) assert resp.status_code == 404 + + +class TestDeleteCrawl: + def test_removes_web_source_and_returns_204(self, app: FastAPI): + mock_auth_user(app, _user()) + crawl = _stub_crawl() + crawl.status = "complete" + delete_mock = AsyncMock(return_value=3) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission("owner"), + ), patch( + "apis.app_api.web_sources.routes.get_crawl_job", + new_callable=AsyncMock, + return_value=crawl, + ), patch( + "apis.app_api.web_sources.routes.delete_web_source", delete_mock + ): + client = TestClient(app) + resp = client.delete( + f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-1" + ) + assert resp.status_code == 204 + # The crawl's own root_url is what scopes the page sweep β€” not a value + # the caller supplies, so a stale client can't widen the blast radius. + delete_mock.assert_awaited_once_with( + assistant_id=ASSISTANT_ID, + crawl_id="CRAWL-1", + root_url="https://example.com/", + owner_id=USER_ID, + ) + + def test_editor_may_delete(self, app: FastAPI): + """The web-sources list is rendered for editors, so the delete must + work for them β€” the owner-keyed services get the real owner_id.""" + mock_auth_user(app, _user()) + crawl = _stub_crawl() + crawl.status = "complete" + owner = SimpleNamespace(owner_id="someone-else") + delete_mock = AsyncMock(return_value=1) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=(owner, "editor"), + ), patch( + "apis.app_api.web_sources.routes.get_crawl_job", + new_callable=AsyncMock, + return_value=crawl, + ), patch( + "apis.app_api.web_sources.routes.delete_web_source", delete_mock + ): + client = TestClient(app) + resp = client.delete( + f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-1" + ) + assert resp.status_code == 204 + assert delete_mock.await_args.kwargs["owner_id"] == "someone-else" + + def test_viewer_gets_403(self, app: FastAPI): + mock_auth_user(app, _user()) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission("viewer"), + ): + client = TestClient(app) + resp = client.delete( + f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-1" + ) + assert resp.status_code == 403 + + def test_returns_404_when_crawl_missing(self, app: FastAPI): + mock_auth_user(app, _user()) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission(), + ), patch( + "apis.app_api.web_sources.routes.get_crawl_job", + new_callable=AsyncMock, + return_value=None, + ): + client = TestClient(app) + resp = client.delete( + f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-X" + ) + assert resp.status_code == 404 + + def test_returns_409_while_crawl_is_running(self, app: FastAPI): + """Deleting mid-crawl would strand pages the crawler is still writing.""" + mock_auth_user(app, _user()) + crawl = _stub_crawl() # status='running', started just now + delete_mock = AsyncMock(return_value=0) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission(), + ), patch( + "apis.app_api.web_sources.routes.get_crawl_job", + new_callable=AsyncMock, + return_value=crawl, + ), patch( + "apis.app_api.web_sources.routes.is_crawl_stale", return_value=False + ), patch( + "apis.app_api.web_sources.routes.delete_web_source", delete_mock + ): + client = TestClient(app) + resp = client.delete( + f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-1" + ) + assert resp.status_code == 409 + delete_mock.assert_not_awaited() + + def test_deletes_a_stale_running_crawl(self, app: FastAPI): + """A `running` row whose process died is a zombie β€” it must stay + removable, or the UI would show an undeletable 'Crawling…' source.""" + mock_auth_user(app, _user()) + crawl = _stub_crawl() # status='running' + delete_mock = AsyncMock(return_value=0) + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission(), + ), patch( + "apis.app_api.web_sources.routes.get_crawl_job", + new_callable=AsyncMock, + return_value=crawl, + ), patch( + "apis.app_api.web_sources.routes.is_crawl_stale", return_value=True + ), patch( + "apis.app_api.web_sources.routes.delete_web_source", delete_mock + ): + client = TestClient(app) + resp = client.delete( + f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-1" + ) + assert resp.status_code == 204 + delete_mock.assert_awaited_once() + + def test_returns_500_when_crawl_row_survives(self, app: FastAPI): + mock_auth_user(app, _user()) + crawl = _stub_crawl() + crawl.status = "complete" + with patch( + "apis.app_api.web_sources.routes.resolve_assistant_permission", + new_callable=AsyncMock, + return_value=_stub_permission(), + ), patch( + "apis.app_api.web_sources.routes.get_crawl_job", + new_callable=AsyncMock, + return_value=crawl, + ), patch( + "apis.app_api.web_sources.routes.delete_web_source", + new_callable=AsyncMock, + side_effect=WebSourceDeletionError("boom"), + ): + client = TestClient(app) + resp = client.delete( + f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-1" + ) + assert resp.status_code == 500 + + def test_returns_401_unauthenticated(self, app: FastAPI): + mock_no_auth(app) + client = TestClient(app) + resp = client.delete( + f"/assistants/{ASSISTANT_ID}/web-sources/crawls/CRAWL-1" + ) + assert resp.status_code == 401 diff --git a/backend/tests/routes/test_admin.py b/backend/tests/routes/test_admin.py index d187f4267..f81bf48fc 100644 --- a/backend/tests/routes/test_admin.py +++ b/backend/tests/routes/test_admin.py @@ -82,6 +82,25 @@ def _raise(): app.dependency_overrides[require_admin] = _raise +@pytest.fixture(autouse=True) +def stub_model_role_service(): + """ + Stub the ModelRoleService for every test in this module. + + The managed-model routes treat the AppRole records as the source of truth for + model access: reads derive allowedAppRoles from them, writes push the picker's + selection back into each role's grantedModels. These tests cover the model + storage layer only, so without this stub every request would hit the real + roles table. ``hydrate_model_roles`` passes the models straight through. + """ + service = AsyncMock() + service.hydrate_model_roles.side_effect = lambda models: models + with patch( + f"{MANAGED_MODELS_PATH}.get_model_role_service", return_value=service + ): + yield service + + # --------------------------------------------------------------------------- # Requirement 7.1: Admin endpoint returns 200 for user with Admin role # --------------------------------------------------------------------------- @@ -252,6 +271,46 @@ def test_create_returns_201(self, app, make_user): body = resp.json() assert body["modelId"] == "anthropic.claude-3-haiku" + def test_create_grants_the_model_to_the_selected_roles( + self, app, make_user, stub_model_role_service + ): + """ + The role picker must write through to the ROLE records β€” that is what + actually grants access. Previously it wrote a field on the model that no + access check read, so selecting a role here did nothing at all. + """ + admin = make_user(email="admin@example.com", user_id="admin-001", roles=["Admin"]) + _override_require_admin(app, admin) + + with patch( + f"{MANAGED_MODELS_PATH}.create_managed_model", + new_callable=AsyncMock, + return_value=SAMPLE_MODEL, + ): + client = TestClient(app) + resp = client.post( + "/admin/managed-models", + json={ + "modelId": "anthropic.claude-3-haiku", + "modelName": "Claude 3 Haiku", + "provider": "bedrock", + "providerName": "Anthropic", + "inputModalities": ["TEXT"], + "outputModalities": ["TEXT"], + "maxInputTokens": 200000, + "maxOutputTokens": 4096, + "inputPricePerMillionTokens": 0.25, + "outputPricePerMillionTokens": 1.25, + "allowedAppRoles": ["staff"], + }, + ) + + assert resp.status_code == 201 + stub_model_role_service.set_roles_for_model.assert_awaited_once() + model_id, role_ids, *_ = stub_model_role_service.set_roles_for_model.await_args.args + assert model_id == SAMPLE_MODEL.model_id + assert role_ids == ["staff"] + # --------------------------------------------------------------------------- # Requirement 7.7: DELETE managed model returns 204 for admin @@ -266,7 +325,13 @@ def test_delete_returns_204(self, app, make_user): admin = make_user(email="admin@example.com", user_id="admin-001", roles=["Admin"]) _override_require_admin(app, admin) + # The route reads the model before deleting it so it knows which provider + # model id to strip from the roles' grantedModels. with patch( + f"{MANAGED_MODELS_PATH}.get_managed_model", + new_callable=AsyncMock, + return_value=SAMPLE_MODEL, + ), patch( f"{MANAGED_MODELS_PATH}.delete_managed_model", new_callable=AsyncMock, return_value=True, @@ -275,3 +340,29 @@ def test_delete_returns_204(self, app, make_user): resp = client.delete("/admin/managed-models/model-001") assert resp.status_code == 204 + + def test_delete_revokes_the_model_from_every_role( + self, app, make_user, stub_model_role_service + ): + """Deleting a model must not leave roles granting a model that's gone.""" + admin = make_user(email="admin@example.com", user_id="admin-001", roles=["Admin"]) + _override_require_admin(app, admin) + + with patch( + f"{MANAGED_MODELS_PATH}.get_managed_model", + new_callable=AsyncMock, + return_value=SAMPLE_MODEL, + ), patch( + f"{MANAGED_MODELS_PATH}.delete_managed_model", + new_callable=AsyncMock, + return_value=True, + ): + client = TestClient(app) + resp = client.delete("/admin/managed-models/model-001") + + assert resp.status_code == 204 + stub_model_role_service.revoke_model_from_all_roles.assert_awaited_once() + assert ( + stub_model_role_service.revoke_model_from_all_roles.await_args.args[0] + == SAMPLE_MODEL.model_id + ) diff --git a/backend/tests/routes/test_inference.py b/backend/tests/routes/test_inference.py index 6342ec1c8..baace6847 100644 --- a/backend/tests/routes/test_inference.py +++ b/backend/tests/routes/test_inference.py @@ -335,3 +335,98 @@ def test_explicit_chat_opts_out(self, authed_app, authed_client): kwargs = get_agent_mock.call_args.kwargs assert kwargs["agent_type"] == "chat" assert kwargs["accessible_skill_ids"] is None + + +# --------------------------------------------------------------------------- +# Single-flight concurrency guard (follow-up to PR #653) +# See docs/specs/session-single-flight-guard.md +# --------------------------------------------------------------------------- + + +class TestInvocationsSingleFlight: + """POST /invocations acquires a per-session lease and rejects duplicates.""" + + def _mock_agent(self): + agent = MagicMock() + + async def fake_stream(*args, **kwargs): + yield "event: done\ndata: {}\n\n" + + agent.stream_async = fake_stream + return agent + + def test_duplicate_concurrent_invocation_returns_409(self, authed_app, authed_client): + """A second turn while the first holds the lease is rejected with 409.""" + from apis.shared.sessions.session_lease import SessionBusyError + + with patch( + "apis.inference_api.chat.routes.get_agent", + return_value=self._mock_agent(), + ), patch( + "apis.inference_api.chat.routes.is_quota_enforcement_enabled", + return_value=False, + ), patch( + # Patched at the source module β€” the route imports it locally at + # call time, so this binding is what it resolves. + "apis.shared.sessions.session_lease.acquire_session_lease", + AsyncMock(side_effect=SessionBusyError("sess-dup")), + ): + resp = authed_client.post( + "/invocations", + json={"session_id": "sess-dup", "message": "hi"}, + ) + + assert resp.status_code == 409 + assert "already streaming" in resp.text.lower() + + def test_lease_released_after_stream_completes(self, authed_app, authed_client): + """The happy path releases the lease when the SSE stream ends.""" + from apis.shared.sessions.session_lease import SessionLease + + sentinel = SessionLease(session_id="sess-ok", user_id="u1", owner="owner-xyz") + release_mock = AsyncMock() + + with patch( + "apis.inference_api.chat.routes.get_agent", + return_value=self._mock_agent(), + ), patch( + "apis.inference_api.chat.routes.is_quota_enforcement_enabled", + return_value=False, + ), patch( + "apis.shared.sessions.session_lease.acquire_session_lease", + AsyncMock(return_value=sentinel), + ), patch( + "apis.shared.sessions.session_lease.release_session_lease", + release_mock, + ): + resp = authed_client.post( + "/invocations", + json={"session_id": "sess-ok", "message": "hi"}, + ) + _ = resp.text # drive the streaming generator (and its finally) to completion + + assert resp.status_code == 200 + release_mock.assert_awaited_with(sentinel) + + def test_preview_session_skips_the_guard(self, authed_app, authed_client): + """Preview sessions never touch the lease (they don't persist).""" + acquire_mock = AsyncMock(return_value=None) + + with patch( + "apis.inference_api.chat.routes.get_agent", + return_value=self._mock_agent(), + ), patch( + "apis.inference_api.chat.routes.is_quota_enforcement_enabled", + return_value=False, + ), patch( + "apis.shared.sessions.session_lease.acquire_session_lease", + acquire_mock, + ): + resp = authed_client.post( + "/invocations", + json={"session_id": "preview-abc", "message": "hi"}, + ) + _ = resp.text + + assert resp.status_code == 200 + acquire_mock.assert_not_awaited() diff --git a/backend/tests/routes/test_sessions.py b/backend/tests/routes/test_sessions.py index 1900bc447..105f9d1d5 100644 --- a/backend/tests/routes/test_sessions.py +++ b/backend/tests/routes/test_sessions.py @@ -800,6 +800,49 @@ def test_returns_204_and_records_user_stopped(self, app, make_user, authenticate source="client_signal", ) + def test_arms_session_cancel_on_stop(self, app, make_user, authenticated_client): + """Stop also arms a cancel on the session lease so the container running + the turn can unwind it (distributed cancellation) β€” a client abort + doesn't propagate through the AgentCore Runtime data plane.""" + user = make_user() + client = authenticated_client(app, user) + + cancel = AsyncMock(return_value=True) + with patch( + "apis.app_api.sessions.routes.set_interrupted_turn", + AsyncMock(), + ), patch( + # Patched at the source module β€” the route imports it locally. + "apis.shared.sessions.session_lease.request_session_cancel", + cancel, + ): + resp = client.post( + "/sessions/sess-001/interrupt", + json={"reason": "user_stopped"}, + ) + + assert resp.status_code == 204 + cancel.assert_awaited_once_with("sess-001", user.user_id) + + def test_stop_succeeds_even_if_cancel_arm_fails(self, app, make_user, authenticated_client): + """Arming the cancel is best-effort β€” a failure never fails the Stop.""" + user = make_user() + client = authenticated_client(app, user) + + with patch( + "apis.app_api.sessions.routes.set_interrupted_turn", + AsyncMock(), + ), patch( + "apis.shared.sessions.session_lease.request_session_cancel", + AsyncMock(side_effect=RuntimeError("dynamo down")), + ): + resp = client.post( + "/sessions/sess-001/interrupt", + json={"reason": "user_stopped"}, + ) + + assert resp.status_code == 204 + def test_rejects_non_client_attested_reason(self, app, make_user, authenticated_client): """`connection_lost` is server-inferred only β€” a client must not be able to plant (or downgrade to) it through this endpoint.""" diff --git a/backend/tests/shared/test_session_lease.py b/backend/tests/shared/test_session_lease.py new file mode 100644 index 000000000..1f8779e12 --- /dev/null +++ b/backend/tests/shared/test_session_lease.py @@ -0,0 +1,235 @@ +"""Tests for the per-session single-flight lease (concurrency guard). + +Covers the server-side race PR #653's follow-up closes: two concurrent +/invocations for one session. See docs/specs/session-single-flight-guard.md and +apis/shared/sessions/session_lease.py. + +All tests run against a moto-backed ``sessions-metadata`` table via the shared +``sessions_metadata_table`` fixture. +""" + +import time + +import pytest + +from apis.shared.sessions.session_lease import ( + LEASE_WINDOW_SECONDS, + SessionBusyError, + SessionLease, + acquire_session_lease, + release_session_lease, + renew_session_lease, + request_session_cancel, +) + + +def _lease_item(table, session_id="s1", user_id="u1"): + resp = table.get_item(Key={"PK": f"USER#{user_id}", "SK": f"LEASE#{session_id}"}) + return resp.get("Item") + + +class TestAcquire: + @pytest.mark.asyncio + async def test_acquire_writes_lease_item(self, sessions_metadata_table): + lease = await acquire_session_lease("s1", "u1") + assert isinstance(lease, SessionLease) + assert lease.owner + + item = _lease_item(sessions_metadata_table) + assert item is not None + assert item["PK"] == "USER#u1" + assert item["SK"] == "LEASE#s1" + assert item["leaseOwner"] == lease.owner + # leaseExpiresAt is the app-level validity check; ttl is the auto-reap + # backstop set further out. + assert int(item["leaseExpiresAt"]) > int(time.time()) + assert int(item["ttl"]) > int(item["leaseExpiresAt"]) + + @pytest.mark.asyncio + async def test_second_concurrent_acquire_is_rejected(self, sessions_metadata_table): + first = await acquire_session_lease("s1", "u1") + assert first is not None + # A duplicate arriving while the first turn holds an unexpired lease. + with pytest.raises(SessionBusyError): + await acquire_session_lease("s1", "u1") + + @pytest.mark.asyncio + async def test_acquire_over_expired_lease_succeeds(self, sessions_metadata_table): + # Simulate a crashed turn that left a stale (expired) lease behind: + # write the item directly with a past leaseExpiresAt. + now = int(time.time()) + sessions_metadata_table.put_item( + Item={ + "PK": "USER#u1", + "SK": "LEASE#s1", + "leaseOwner": "dead-owner", + "leaseExpiresAt": now - 10, + "ttl": now + 3600, + } + ) + lease = await acquire_session_lease("s1", "u1") + assert lease is not None + assert lease.owner != "dead-owner" + assert _lease_item(sessions_metadata_table)["leaseOwner"] == lease.owner + + @pytest.mark.asyncio + async def test_distinct_sessions_do_not_conflict(self, sessions_metadata_table): + a = await acquire_session_lease("s1", "u1") + b = await acquire_session_lease("s2", "u1") + assert a is not None and b is not None + assert a.owner != b.owner + + @pytest.mark.asyncio + async def test_force_takes_over_active_lease(self, sessions_metadata_table): + # Resume / continuation: an active lease must NOT block them. + first = await acquire_session_lease("s1", "u1") + assert first is not None + resumed = await acquire_session_lease("s1", "u1", force=True) + assert resumed is not None + assert resumed.owner != first.owner + # The forced acquire installs its own lease so a *fresh* duplicate + # arriving during the resume is still rejected. + assert _lease_item(sessions_metadata_table)["leaseOwner"] == resumed.owner + with pytest.raises(SessionBusyError): + await acquire_session_lease("s1", "u1") + + @pytest.mark.asyncio + async def test_no_table_configured_returns_none(self, aws, monkeypatch): + # Local / no-DynamoDB path: guard is inactive, never blocks a turn. + monkeypatch.delenv("DYNAMODB_SESSIONS_METADATA_TABLE_NAME", raising=False) + assert await acquire_session_lease("s1", "u1") is None + + +class TestRenew: + @pytest.mark.asyncio + async def test_renew_extends_window_for_owner(self, sessions_metadata_table): + lease = await acquire_session_lease("s1", "u1") + # Backdate the stored window so a renewal is observably larger even + # within the same wall-clock second. + sessions_metadata_table.update_item( + Key={"PK": lease.pk, "SK": lease.sk}, + UpdateExpression="SET leaseExpiresAt = :old", + ExpressionAttributeValues={":old": int(time.time()) - 5}, + ) + await renew_session_lease(lease) + item = _lease_item(sessions_metadata_table) + assert int(item["leaseExpiresAt"]) >= int(time.time()) + LEASE_WINDOW_SECONDS - 1 + assert item["leaseOwner"] == lease.owner + + @pytest.mark.asyncio + async def test_renew_by_non_owner_is_noop(self, sessions_metadata_table): + await acquire_session_lease("s1", "u1") + original = _lease_item(sessions_metadata_table) + # A container that lost the lease tries to renew β€” owner-scoped, so the + # current owner's window is untouched and no error surfaces. + stale = SessionLease(session_id="s1", user_id="u1", owner="stale-owner") + await renew_session_lease(stale) + after = _lease_item(sessions_metadata_table) + assert after["leaseOwner"] == original["leaseOwner"] + assert int(after["leaseExpiresAt"]) == int(original["leaseExpiresAt"]) + + @pytest.mark.asyncio + async def test_renew_none_is_noop(self, sessions_metadata_table): + await renew_session_lease(None) # must not raise + + +class TestRelease: + @pytest.mark.asyncio + async def test_release_deletes_own_lease(self, sessions_metadata_table): + lease = await acquire_session_lease("s1", "u1") + await release_session_lease(lease) + assert _lease_item(sessions_metadata_table) is None + # Session is free again β€” a new turn can acquire. + again = await acquire_session_lease("s1", "u1") + assert again is not None + + @pytest.mark.asyncio + async def test_release_by_non_owner_leaves_lease(self, sessions_metadata_table): + real = await acquire_session_lease("s1", "u1") + stale = SessionLease(session_id="s1", user_id="u1", owner="stale-owner") + # A lapsed-then-retaken owner must not delete the new owner's lease. + await release_session_lease(stale) + item = _lease_item(sessions_metadata_table) + assert item is not None + assert item["leaseOwner"] == real.owner + + @pytest.mark.asyncio + async def test_release_is_idempotent(self, sessions_metadata_table): + lease = await acquire_session_lease("s1", "u1") + await release_session_lease(lease) + await release_session_lease(lease) # second is a no-op conditional miss + assert _lease_item(sessions_metadata_table) is None + + @pytest.mark.asyncio + async def test_release_none_is_noop(self, sessions_metadata_table): + await release_session_lease(None) # must not raise + + +class TestCancel: + @pytest.mark.asyncio + async def test_renew_reports_no_cancel_by_default(self, sessions_metadata_table): + lease = await acquire_session_lease("s1", "u1") + assert await renew_session_lease(lease) is False + + @pytest.mark.asyncio + async def test_request_cancel_is_observed_by_owner_renew(self, sessions_metadata_table): + lease = await acquire_session_lease("s1", "u1") + armed = await request_session_cancel("s1", "u1") + assert armed is True + # The container running the turn sees it on its next heartbeat renew. + assert await renew_session_lease(lease) is True + + @pytest.mark.asyncio + async def test_request_cancel_with_no_active_lease_is_noop(self, sessions_metadata_table): + # No turn streaming β†’ nothing to cancel. + assert await request_session_cancel("s1", "u1") is False + + @pytest.mark.asyncio + async def test_cancel_is_owner_scoped_across_takeover(self, sessions_metadata_table): + # A Stop arms a cancel against the current owner; then that turn ends + # and a new one force-acquires (resume). The new owner must NOT inherit + # the old cancel. + first = await acquire_session_lease("s1", "u1") + await request_session_cancel("s1", "u1") + assert await renew_session_lease(first) is True # armed for `first` + + resumed = await acquire_session_lease("s1", "u1", force=True) + assert resumed.owner != first.owner + # acquire cleared the stale marker; the new owner sees no cancel. + item = _lease_item(sessions_metadata_table) + assert "cancelRequestedFor" not in item + assert await renew_session_lease(resumed) is False + + @pytest.mark.asyncio + async def test_cancel_after_takeover_targets_only_new_owner(self, sessions_metadata_table): + first = await acquire_session_lease("s1", "u1") + resumed = await acquire_session_lease("s1", "u1", force=True) + # A fresh Stop now arms against the current (resumed) owner. + await request_session_cancel("s1", "u1") + assert await renew_session_lease(resumed) is True + # The superseded owner never sees it (it lost the lease anyway). + assert await renew_session_lease(first) is False + + @pytest.mark.asyncio + async def test_release_after_cancel_frees_session(self, sessions_metadata_table): + lease = await acquire_session_lease("s1", "u1") + await request_session_cancel("s1", "u1") + await release_session_lease(lease) + assert _lease_item(sessions_metadata_table) is None + # Resend acquires cleanly, with no leftover cancel marker. + again = await acquire_session_lease("s1", "u1") + assert again is not None + assert await renew_session_lease(again) is False + + +class TestLifecycle: + @pytest.mark.asyncio + async def test_acquire_release_reacquire_cycle(self, sessions_metadata_table): + """A normal turn: acquire, run, release; the next turn acquires cleanly.""" + for _ in range(3): + lease = await acquire_session_lease("s1", "u1") + assert lease is not None + # Duplicate during the turn is rejected. + with pytest.raises(SessionBusyError): + await acquire_session_lease("s1", "u1") + await release_session_lease(lease) diff --git a/backend/uv.lock b/backend/uv.lock index f9ce7188f..d06f3f5f3 100644 --- a/backend/uv.lock +++ b/backend/uv.lock @@ -12,7 +12,7 @@ resolution-markers = [ [[package]] name = "agentcore-stack" -version = "1.5.0" +version = "1.6.0" source = { editable = "." } dependencies = [ { name = "aiofiles" }, diff --git a/docs/specs/session-single-flight-guard.md b/docs/specs/session-single-flight-guard.md new file mode 100644 index 000000000..e8f8b4a84 --- /dev/null +++ b/docs/specs/session-single-flight-guard.md @@ -0,0 +1,331 @@ +# Session single-flight concurrency guard + +**Status:** implemented (branch `fix/session-single-flight-guard`, off `develop`) +**Follow-up to:** PR #653 (`fix/tool-pairing-restore-sanitizer`) + +## Problem + +A client-side abort (Stop button, tab switch, dropped socket, transport retry) +does **not** propagate through the AgentCore Runtime data plane to the backend +agent β€” the agent runs to completion server-side regardless. If a second +`POST /invocations` for the same session arrives while the first turn is still +running (genuine double-click, two tabs/devices, or an HTTP retry), the Runtime +can route the two invocations to **different containers**, so two agent loops +run concurrently against the same AgentCore Memory session. Both persist +`toolUse`/`toolResult` events, producing duplicate + interleaved +(assistant/assistant/user/user) history that Bedrock Converse rejects on every +subsequent turn: + +> The number of toolResult blocks at messages.N exceeds the number of toolUse +> blocks of previous turn. + +That was the prod incident that bricked session `f761f59b`. PR #653 closed the +frontend tab-switch vector (`openWhenHidden: true` so backgrounding a tab no +longer aborts + reopens the stream) and added a restore-time +`_repair_tool_pairing` sanitizer. This change closes the remaining **server-side +race**: two live invocations for one session. + +## Decision summary + +| Question | Decision | +|----------|----------| +| Reject-new vs supersede-old | **Reject-new** with HTTP **409** | +| Lock mechanism | **Distributed lease** β€” DynamoDB conditional write on `sessions-metadata` | +| Which service owns the lock | **inference-api `/invocations`** (the true turn-start chokepoint) | +| Lease validity window | **90 s**, renewed by heartbeat | +| Heartbeat | Background asyncio task, renews every **30 s** while the turn streams | +| TTL backstop | DynamoDB `ttl` attribute at **now + 1 h** (auto-reap orphans) | +| Resume / continuation carve-out | Bypass the conflict check; acquire the lease with `force=True` | +| Failure posture | **Fail-open** β€” any non-conflict DynamoDB error proceeds without a lease | + +### Reject-new vs supersede-old + +**Reject-new (chosen).** On an active lease, the duplicate gets `409 Conflict`; +the first turn finishes uninterrupted and the SPA surfaces "already streaming". + +Supersede-old matches the SPA's "newest stream wins" UX, but superseding +requires *signalling the in-flight agent loop to stop* β€” which is exactly the +thing that does **not** propagate through the Runtime data plane (the whole +premise of the bug). The session manager's `cancelled` flag and `StopHook` +(`turn_based_session_manager.py`, `session/hooks/stop.py`) only reach the agent +that shares the *same in-process* session manager instance; a loop on another +container never sees the flag. So supersede cannot actually stop the old writer +and both loops would still corrupt history. Reject-new is the only option that +holds without a cross-container stop signal. It is also strictly safer: the +worst case is a spurious 409 (recoverable by retrying after the first turn), +never a corrupted session. + +### Which service owns the lock + +The lease lives in **inference-api `/invocations`**, not the app-api +`/chat/stream` BFF proxy: + +- `/invocations` is the true turn-start chokepoint and the exact point where a + second container also begins. The harm (AgentCore Memory tool-pairing + corruption) is inference-api's domain β€” the agent loop it owns is the writer. +- An in-process lock is insufficient (two containers), so the guard must be + distributed regardless of which service holds it. +- Per the CLAUDE.md inference-api boundary rule and the "respect service + boundaries" principle, the guard belongs where the writes happen. +- The app-api proxy already relays a `>= 400` upstream status verbatim + (`proxy_routes.py`), so a 409 from inference-api reaches the SPA unchanged β€” + **no app-api change is required.** Putting a second lease in app-api would + duplicate state and couple the BFF to turn semantics it deliberately doesn't + own (the body is opaque bytes there). + +The lease **storage helper** lives in `apis/shared/sessions/` (like +`metadata.py`), since that package is the shared home for session-row access; +only inference-api *uses* it. + +### Lease shape & storage + +Dedicated item on the existing `boisestateai-v2-sessions-metadata` table β€” no +new table, no CDK change (the table already has `timeToLiveAttribute: 'ttl'` and +PAY_PER_REQUEST billing): + +``` +PK = USER#{user_id} +SK = LEASE#{session_id} +Attributes: + leaseOwner : uuid4 hex, unique per invocation + leaseExpiresAt : epoch seconds β€” the app-level validity check + ttl : epoch seconds (now + 3600) β€” DynamoDB auto-reap backstop + updatedAt : ISO8601 +``` + +A **deterministic key** (`LEASE#{session_id}`, no dynamic `lastMessageAt` +suffix) means acquisition is a single atomic conditional write with **no GSI +read first** β€” eliminating the read-before-write race that a marker on the META +row would have. The item is invisible to session listings (`list_user_sessions` +queries `SK begins_with 'S#ACTIVE#'`) and to `SessionLookupIndex` (it has no +`GSI_PK`/`GSI_SK`). + +Why both `leaseExpiresAt` *and* `ttl`: DynamoDB TTL deletion lags up to 48 h, so +it can't be the correctness mechanism. The **application** compares +`leaseExpiresAt < now` in the acquisition `ConditionExpression`; `ttl` only +stops crashed-container orphans from accumulating forever. + +**Acquisition (fresh turn):** + +``` +UpdateItem + ConditionExpression = attribute_not_exists(PK) OR leaseExpiresAt < :now + SET leaseOwner, leaseExpiresAt, ttl, updatedAt +``` + +`ConditionalCheckFailedException` β‡’ an unexpired lease is held β‡’ raise +`SessionBusyError` β‡’ `/invocations` returns 409. Any other `ClientError` +(throttle, transient) β‡’ log and **proceed without a lease** (fail-open: never +block a legitimate turn on lock-infra failure). + +**Renewal (heartbeat)** and **release** both carry +`ConditionExpression = leaseOwner = :owner` so a container that already lost the +lease (its window lapsed and another turn took over) can neither extend nor +delete the new owner's lease. Both are best-effort and swallow errors. + +### Heartbeat & window sizing + +- Window **90 s**, heartbeat every **30 s** β‡’ tolerates one missed renewal + (network blip) before expiry. +- A background asyncio task (not inline on SSE-event cadence) renews on a wall + clock, so it keeps the lease alive even across a **long silent tool call** + (code-interpreter / browser can run > 90 s between yielded events). +- After a genuine container crash the heartbeat stops; the lease self-expires in + ≀ 90 s and the session is usable again. This is the recovery time; it is well + under the 600 s stream timeout, so a normal long turn never self-evicts. + +### Resume / continuation carve-out + +Resume (`interrupt_responses` present, `is_resume`) and max-tokens continuation +(`continue_truncated`, `is_continuation`) re-enter a turn whose original agent +loop has **already ended** (it paused/truncated and its SSE stream closed). +There is no concurrent writer to guard against, and blocking them would strand +the user. They call `acquire_session_lease(..., force=True)`: + +- **`force=True`** does an *unconditional* write β€” it always succeeds, so an + authorized resume can never be rejected (even against a stale lease the paused + turn failed to release). +- It still *installs* a lease, so a fresh duplicate that arrives *during* the + resume is itself rejected β€” the session stays single-flight end to end. + +Two concurrent resumes for one session is not a real vector (resume is an +explicit post-OAuth user action), so `force` overwrite between them is +acceptable. + +### Placement & release wiring + +Acquire at the **top of the streaming `try:`** in `invocations` β€” after the +pre-flight checks (quota, model access, RAG validation, file resolution) and +immediately before agent construction. This keeps release management to three +contiguous sites with no early `return` in between: + +1. **Happy path:** the SSE generator (`stream_with_quota_warning`) starts the + heartbeat before `agent.stream_async`, and in its `finally` cancels the + heartbeat and releases the lease. +2. **`except HTTPException`** (e.g. resume/interrupt 400s raised after acquire): + release, then re-raise. +3. **`except Exception`** (agent build error β†’ conversational error stream): + release, then return the error stream. + +FastAPI runs the generator *after* the handler returns, so no single `finally` +can cover both the streaming and non-streaming exits β€” hence the three sites. +They are mutually exclusive (reaching the generator means no exception reached +the outer handlers), and release is idempotent (owner-conditional delete), so +there is no double-release hazard. + +Preview sessions (`is_preview_session`) and the local/no-DynamoDB path (table +env var unset) skip the guard entirely β€” they don't persist and can't corrupt +shared Memory. + +### Out of scope + +- **App-initiated invocations** (`app_tool_call`, `app_context_update`) are + short-circuited above the guard. They do write synthesized tool events, but + they are inert behind the MCP-Apps host flag (no live App), synchronous, and + fast. Guarding them is deferred. +- **Scheduled/headless runs** go through the same `/invocations` and get the + guard for free; a schedule that fires twice for one session is exactly the + kind of duplicate this rejects. + +## Known limitations of the lease alone + +The lease is a **safety floor** β€” it guarantees "never corrupt Memory," not +"Stop actually stops." Two rough edges remain, both stemming from the same root +cause (client aborts don't reach the running turn), and both are the motivation +for the distributed-cancel follow-on below: + +- **Best-effort, not a hard lock.** Fail-open on a DynamoDB brownout, a + heartbeat-failure window (~3 missed renewals β‰ˆ the lease window), and reliance + on NTP-level clock agreement between containers each leave a narrow window + where a duplicate could still slip through. Acceptable for a net; named so it + isn't mistaken for a hard guarantee. +- **Stop β†’ immediate resend returns 409.** Because Stop doesn't end the server + turn, the old turn keeps running and holds the lease (heartbeated, up to the + 600s stream timeout). The resend is a genuine second concurrent loop, so 409 + is *correct* β€” but it changes UX. This is strictly better than today (where + the same action corrupts the session), but the SPA must handle 409 as "prior + response still finishing," and the clean fix is to make Stop genuinely end the + turn (below). + +## Follow-on: distributed turn cancellation (make Stop real) + +**Status: IMPLEMENTED** (branch `feat/distributed-turn-cancellation`, off +`develop` after the lease landed). Complementary to the lease, not a +replacement β€” the lease prevents corruption; this makes Stop actually end the +server turn, which shrinks the 409-on-resend window from "up to a full turn" to +"one heartbeat interval (~10s)" and reclaims wasted model/tool spend. + +### What shipped + +- **Signal** β€” the app-api `POST /sessions/{id}/interrupt` (`user_stopped`) + handler calls `request_session_cancel`, which reads the lease's current + `leaseOwner` and stamps `cancelRequestedFor = ` (owner-scoped, so a + stale Stop can't kill a later turn). Best-effort β€” never fails the Stop. +- **Observe** β€” the inference-api lease heartbeat (`_lease_heartbeat_loop`, + tightened to a 10s cadence) renews with `ReturnValues=ALL_NEW`; when it sees + `cancelRequestedFor == our owner` it flips `agent.session_manager.cancelled`. +- **Effect A (tools)** β€” the always-on `StopHook` cancels the next tool call + (existing machinery; zero new code). +- **Effect B (model stream)** β€” a cooperative check at the top of the + `StreamCoordinator` loop raises `_CooperativeStopSignal` when `cancelled` is + set. A dedicated `except` arm persists the partial via `_persist_interruption` + (marked `user_stopped`), emits terminal SSE frames, and ends **cleanly** + (no re-raise) β€” so a still-connected client sees a proper close and the + route's `finally` releases the lease. This is what ends a pure-chat turn, + which has no tool boundary for `StopHook` to catch. +- **Release** β€” unchanged: the route's `_guarded_stream` `finally` releases the + lease as the turn unwinds. + +The spike that de-risked this: the teardown-and-persist-partial path already +existed (the `(CancelledError, GeneratorExit)` arm built for interrupted-turn / +connection-lost, hardened by #653's `_repair_tool_pairing`), so stopping +mid-stream never orphans or corrupts history. We abandon the *suspended* agent +stream, which cannot progress or write to Memory without a consumer. + +### Residual limitations (documented, not fixed) + +- **In-flight tool calls finish.** `StopHook` cancels at tool *boundaries*; a + browser / code-interpreter call already executing runs to completion before + the cancel is seen. Genuinely interrupting a running tool is a deeper change + (cooperative cancellation inside the tool executor) and is out of scope. +- **Bedrock token tail.** Aborting stops *further* generation, but tokens the + model already produced server-side are billed. The dominant saving is halting + the loop (no further model calls / tools), which both effects achieve. +- **Observe latency.** Stopβ†’release is bounded by the heartbeat cadence (~10s). + The client already stops rendering instantly on Stop; only *resend* waits, and + the SPA's 409 handling covers that window. + +### Why Stop couldn't propagate before this + +Two stacked reasons, both confirmed in code: + +### Why Stop can't propagate today + +Two stacked reasons, both confirmed in code: + +1. **The AgentCore Runtime data plane only proxies `/invocations` and `/ping`** + (`apis/app_api/sessions/routes.py` β€” *"a custom inference-api route would 404 + in cloud"*; `apis/shared/harness/runner.py` builds + `POST /runtimes/{ARN}/invocations?qualifier=DEFAULT`). There is no + "abort this invocation" operation, so closing the downstream HTTP connection + never signals the container to cancel its running coroutine β€” + `agent.stream_async` runs to completion, still writing to Memory. +2. **You can't address the container running the turn.** The Runtime routes + invocations to containers opaquely (the same reason a duplicate lands on a + *different* container). A "stop session X" request would likely hit the wrong + container. + +Consequently the existing `cancelled` flag + `StopHook` +(`session/turn_based_session_manager.py`, `session/hooks/stop.py`) are **dead +code**: nothing in `src/` ever sets `cancelled = True`. Today's Stop is a client +abort plus a `user_stopped` beacon that only writes an interrupted-turn *marker* +via app-api (`set_interrupted_turn`) for the "Continue" UX β€” it never reaches +the running turn. + +### Design: poll a distributed cancel flag from inside the loop + +Reuse the exact distributed-state pattern the lease already relies on. The stop +signal and the running turn coordinate through the session row, not through the +Runtime: + +1. **Signal (any container).** The `POST /sessions/{id}/interrupt` + `user_stopped` path *additionally* sets a `cancelRequested` marker on the + session's lease/META row β€” e.g. `cancelRequestedFor = ` (scope it + to the current lease owner so it can't cancel a later, unrelated turn), plus + a timestamp. Cheap conditional write; app-api already owns this endpoint and + already talks to this table. +2. **Observe (the running container).** Wire `StopHook` (and, ideally, a check + between agent steps / before each model call) to consult that flag instead of + the in-process boolean. The natural, low-latency read is to **piggyback on + the lease heartbeat** (`_lease_heartbeat_loop`, every 30s): when a renew sees + `cancelRequestedFor == our owner`, set `session_manager.cancelled = True`. + `StopHook.check_cancelled` (already registered on `BeforeToolCallEvent`) then + cancels the next tool call, and the loop unwinds. +3. **Release.** The turn ends β†’ the generator `finally` releases the lease as it + does now β†’ the user's resend acquires cleanly. No 409-on-resend. + +### Trade-offs & open questions + +- **Granularity.** Heartbeat-cadence polling (~30s) bounds worst-case + stop-to-effect latency. A tighter loop-level check (between steps / before each + Bedrock call) makes Stop feel instant but adds a DynamoDB read per step β€” likely + gate it behind "only read when a cheap in-process hint is unset," or accept the + heartbeat cadence for v1. +- **Mid-tool cancellation.** `StopHook` cancels at `BeforeToolCallEvent` + boundaries; a long-running tool already in flight (browser, code interpreter) + finishes before the cancel is seen. Genuinely interrupting an in-flight tool is + a deeper change (cooperative cancellation inside the tool executor) and + probably out of scope even for the follow-on. +- **Partial-write integrity.** When a turn is cancelled mid-stream, its partial + `toolUse`/`toolResult` must remain a valid pairing in Memory β€” this is exactly + what PR #653's `_repair_tool_pairing` restore-time sanitizer already guards, so + cancellation rides on machinery that already exists. +- **Cost upside.** Real cancellation stops burning Bedrock tokens and tool + invocations on output nobody will read β€” a side benefit beyond UX. + +### Relationship to the lease + +The lease prevents *corruption* (never two live writers); distributed +cancellation makes Stop *actually stop* (one writer, ended on demand). Ship the +lease first as the safety floor; the cancel follow-on removes the 409-on-resend +rough edge and reclaims wasted compute. Neither supersedes the other. diff --git a/docs/specs/share-large-conversations-s3-offload.md b/docs/specs/share-large-conversations-s3-offload.md new file mode 100644 index 000000000..df1100c51 --- /dev/null +++ b/docs/specs/share-large-conversations-s3-offload.md @@ -0,0 +1,329 @@ +# Sharing large conversations β€” S3 snapshot offload + +**Status:** Draft (spec-first; implementation gated on approval) +**Branch:** `feature/share-large-conversations-s3-offload` (off `develop`) +**Owner:** Phil Merrell +**Related:** Conversation share feature (`apis/app_api/shares/`); PR #657 (share IAM-grant fix β€” *separate* bug); Memory Spaces S3 store (`apis/shared/memory/store.py`); Artifacts / Skills / RAG S3-offload precedent. + +--- + +## 1. Problem + +Creating a share for a large conversation fails. `ShareService.create_share` inlines the full message list into a single DynamoDB item on the `shared-conversations` table: + +```python +item = { + "share_id": share_id, + ... + "metadata": metadata_snapshot, + "messages": messages_snapshot, # ← entire conversation, inline +} +self._table.put_item(Item=item) +``` + +DynamoDB caps an item at **400 KB**. A long conversation β€” especially one with tool results, images, documents, or reasoning blocks β€” blows past that. Observed in **prod-ai** (`897729136999`, `us-west-2`) app-api logs: + +``` +apis.app_api.shares.routes - ERROR - Error creating share for session 69ea19d9-...: +An error occurred (ValidationException) when calling the PutItem operation: +Item size has exceeded the maximum allowed size +``` + +The route's catch-all turns this into a bare `500 {"detail":"Failed to create share"}`. The user is told sharing failed with no reason and no recourse. + +This is **distinct** from the `AccessDeniedException` IAM-grant bug fixed in PR #657. This one is `ValidationException` / item-size. + +### Why this recurs + +The snapshot is a *copy* of the conversation taken at share time, so it grows monotonically with conversation length and with the richness of each turn (a single image or document block can be tens/hundreds of KB after base64). There is no upper bound on conversation size, so any "just trim it" mitigation only moves the cliff. + +--- + +## 2. Goal + +Support sharing conversations of any size, transparently, while: + +- Keeping the existing share read/write/update/revoke/export API contract unchanged (SPA needs no changes). +- Mirroring the **established S3-offload pattern** already used for Memory Spaces, Skills reference files, Artifacts, and RAG documents β€” not inventing a new one. +- Remaining backward-compatible with existing inline-item shares in prod/dev. +- Replacing the bare `500` with a specific, honest error if a share genuinely cannot be created. + +--- + +## 3. Design overview + +Offload the **snapshot body** (`messages`, and defensively `metadata`) to an S3 object. Keep in the DynamoDB item only: + +- the share's control-plane fields (`share_id`, `session_id`, `owner_id`, `owner_email`, `access_level`, `allowed_emails`, `created_at`) β€” these are small, queried by GSIs, and mutated by `update_share`; they **must** stay in DynamoDB, and +- a **pointer** to the S3 object holding the body. + +The DynamoDB item thus becomes small and bounded regardless of conversation size. The read path (`get_shared_conversation`, `export_shared_conversation`, `get_shares_for_session`) fetches the body from S3 when the pointer is present, and falls back to the inline `messages`/`metadata` fields when it is not (legacy items). + +``` +create_share + β”œβ”€ snapshot messages + metadata (unchanged) + β”œβ”€ serialize body β†’ JSON bytes + β”œβ”€ store.put(share_id, body) β†’ s3_key (NEW) + └─ put_item { control fields..., body_ref: {bucket_key, format, ...} } (no inline messages) + +get_shared_conversation / export_shared_conversation + β”œβ”€ get_item β†’ control fields + (body_ref | inline messages) + β”œβ”€ if body_ref: store.get(key) β†’ messages/metadata (NEW) + └─ else: read inline item["messages"]/["metadata"] (legacy fallback) +``` + +### Why S3, always (not a size threshold) + +The task raised "inline under N KB vs always-S3." **Recommendation: always offload the body to S3.** Rationale: + +- **One code path.** A size gate means two write paths and two read paths, each needing tests and each a place for the 400 KB cliff to hide (the gate has to account for DynamoDB's *item* overhead β€” attribute names, the `Decimal` re-encoding, the control fields β€” not just `len(json)`, so a naive threshold is itself a source of bugs). +- **The body is never queried.** Nothing in the share feature does a DynamoDB `Query`/`Scan` *on message content*; the GSIs key on `session_id` and `owner_id` only. So there is zero DynamoDB-side benefit to keeping messages inline. +- **The offloaded read is one `get_object`.** For a share view that already crosses the network and renders a whole conversation, a single S3 GET is negligible latency and removes a class of failure entirely. +- Precedent: Memory Spaces, Skills, and Artifacts all offload bytes to S3 unconditionally and keep only manifests/pointers in DynamoDB. This proposal is deliberately the *same shape*, for reviewer familiarity and code reuse. + +The one concession to "small items": we still **read** legacy inline items (they predate this change), but we never **write** new ones. + +--- + +## 4. Storage design + +### 4.1 New bucket: `shared-conversations` + +A dedicated private S3 bucket, created by extending `SharedConversationsConstruct` (co-locating bucket + table in the one construct that already owns this domain β€” same as `MemorySpacesConstruct` owning both its bucket and table, and `ArtifactsDataConstruct` owning both). + +```ts +// infrastructure/lib/constructs/data/shared-conversations-construct.ts +public readonly bucket: s3.Bucket; + +this.bucket = new s3.Bucket(this, 'SharedConversationsBucket', { + bucketName: getResourceName(config, 'shared-conversations'), + encryption: s3.BucketEncryption.S3_MANAGED, + blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL, + enforceSSL: true, + lifecycleRules: [ + { id: 'abort-stale-multipart', abortIncompleteMultipartUploadAfter: cdk.Duration.days(7) }, + ], + removalPolicy: getRemovalPolicy(config), + autoDeleteObjects: getAutoDeleteObjects(config), +}); +``` + +Private, server-side-only access (app-api reads/writes; never loaded cross-origin β€” the share view is JSON served by app-api, not an iframe). Byte-for-byte the Memory Spaces bucket recipe. + +**No expiration lifecycle rule.** A share's object lives exactly as long as the share row: it is deleted when the share is revoked or the session's shares are cleaned up (Β§5.4). Age-based reaping would break live shares β€” the same "reference recency β‰  object age" lesson from the MCP App UI-resource persistence work (`project_mcp_apps_uires_persistence`: never age-based deletes when a live row can still point at the object). + +> **Decision point for review:** dedicated bucket vs. reusing an existing one. A dedicated bucket keeps IAM scoping clean (its own ARN, its own grant sid) and lifecycle independent, at the cost of one more bucket. This is the pattern every sibling feature follows, so the spec assumes a dedicated bucket. (Reusing e.g. the file-upload bucket would muddy IAM and lifecycle for no real saving.) + +### 4.2 Object key layout + +Content-addressed, mirroring `memory/store.py` and the skills resource store: + +``` +shares/{share_id}/{content_hash} +``` + +- `content_hash` = `sha256(body_bytes)` hex. +- One object per share (a share is immutable once created β€” `update_share` only touches access-control fields, never the body). Content-addressing gives us free idempotency on retry: a re-run of the same `create_share` body writes the same key. +- Keyed under `share_id` so Β§5.4 revoke can delete the object by known key, and a stray-object sweep can list by `shares/{share_id}/` prefix. + +### 4.3 DynamoDB item shape (new writes) + +```python +item = { + "share_id": share_id, + "session_id": session_id, + "owner_id": user.user_id, + "owner_email": user.email, + "access_level": request.access_level, + "created_at": now, + "allowed_emails": [...], # when access_level == "specific" + "body_ref": { # ← NEW: pointer replaces inline body + "bucket_key": "shares/{share_id}/{hash}", + "format": "json", # serialization of the body object + "schema_version": 1, # snapshot schema, for forward migration + "byte_size": 812345, # observability / future gating + }, + # NOTE: no "messages" / "metadata" attributes on new items +} +``` + +The body object itself is the JSON serialization of: + +```json +{ "metadata": { ...session metadata snapshot... }, + "messages": [ ...MessageResponse dicts... ] } +``` + +Serialized with plain `json.dumps` (UTF-8 bytes). **The floatβ†’Decimal conversion is dropped for the offloaded body** β€” that conversion only exists to satisfy DynamoDB's boto3 resource, which rejects Python floats. S3 stores opaque bytes, so we serialize the raw Pydantic `model_dump` directly and skip `_convert_floats_to_decimal` for anything going to S3. (Control fields going into DynamoDB are all strings/lists, so no Decimal concern there.) This also sidesteps the read-side `Decimal`β†’float round-trip. + +### 4.4 Backward compatibility + +Reads must handle three item shapes: + +| Item shape | How it's read | +|---|---| +| **New** (`body_ref` present, no inline `messages`) | fetch body from S3 via `store.get(body_ref["bucket_key"])` | +| **Legacy inline** (`messages`/`metadata` present, no `body_ref`) | read inline exactly as today | +| **Malformed** (neither) | raise `ShareNotFoundError` / log; treat as unreadable | + +No data migration is required β€” existing inline shares keep working untouched. New shares are S3-backed. Optional one-time backfill is **out of scope** (existing inline shares are, by definition, already small enough to have been written). + +--- + +## 5. Backend changes + +All under `backend/src/apis/app_api/shares/`, plus one shared store. + +### 5.1 New: snapshot body store + +`apis/app_api/shares/snapshot_store.py` β€” a thin S3 put/get keyed by `share_id`, structurally identical to `memory/store.py` but domain-scoped to shares. (It lives under `app_api/shares/` because app-api is the only consumer; if a second consumer ever appears it moves to `apis/shared/` per the import-boundary rule. Memory's store is under `apis/shared/` precisely because both app-api and inference-api use it β€” shares are app-api-only.) + +```python +class ShareSnapshotStore: + def __init__(self, bucket_name: str | None = None, s3_client=None): ... + @property + def enabled(self) -> bool: ... # bucket configured AND boto3 present + def put(self, *, share_id: str, body: bytes) -> str: # β†’ bucket_key + def get(self, bucket_key: str) -> bytes: + def delete(self, bucket_key: str) -> None: # best-effort +``` + +- `put`: content-address (`sha256`), `head_object` dedupe, `put_object` with `ServerSideEncryption="AES256"`, `ContentType="application/json"`. +- `get`: `get_object`; `NoSuchKey`/`404` β†’ a typed `ShareSnapshotStoreError` the service maps to a friendly error. +- Errors raise `ShareSnapshotStoreError` (module-local), consistent with `MemorySpaceStoreError`. + +Env var: **`SHARED_CONVERSATIONS_BUCKET_NAME`** (naming parallels the existing `SHARED_CONVERSATIONS_TABLE_NAME`). + +### 5.2 `create_share` + +- Build `metadata_snapshot` + `messages_snapshot` as today (still via `model_dump(by_alias=True, exclude_none=True)`), **without** `_convert_floats_to_decimal` for the S3 body. +- Serialize `{"metadata": ..., "messages": ...}` β†’ UTF-8 JSON bytes. +- `bucket_key = store.put(share_id=share_id, body=body_bytes)`. +- Put the item with `body_ref` and **no** inline `messages`/`metadata`. +- If `store.enabled` is False (bucket unset) β†’ raise a new `ShareStorageUnavailableError` (maps to a specific message, Β§5.5), instead of silently falling back to inline (which would reintroduce the 400 KB cliff). + +### 5.3 Read paths + +- `_get_share_item` unchanged (still `get_item` by `share_id`). +- New helper `_load_snapshot_body(item) -> tuple[metadata, messages]`: + - if `item.get("body_ref")`: `json.loads(store.get(item["body_ref"]["bucket_key"]))` β†’ `(body["metadata"], body["messages"])`. + - elif `item.get("messages") is not None`: legacy inline β†’ `(item.get("metadata", {}), item["messages"])`. + - else: log + raise `ShareNotFoundError`. +- `_build_shared_conversation_response` and `export_shared_conversation._copy_messages_to_memory` consume `messages` from this helper rather than `item["messages"]` directly. +- `get_shares_for_session` / `_build_share_response` **do not** need the body β€” they only surface control fields β€” so listing stays a pure DynamoDB read with **no** S3 fetch. (Good: the list view is unaffected and fast.) + +### 5.4 Delete / revoke + +- `revoke_share`: after `delete_item`, best-effort `store.delete(body_ref["bucket_key"])` (guarded on `body_ref` present). +- `delete_shares_for_session`: for each item being batch-deleted, best-effort delete its S3 object. (These are cleanup paths; an S3 delete failure logs but never blocks the DynamoDB delete β€” matching the "revoked link stops working" guarantee, which is enforced by the DynamoDB row's absence, not the object's.) +- Legacy inline items have no `body_ref` β†’ nothing to delete in S3. + +### 5.5 Friendlier errors (route layer) + +Replace the bare `500` with specific handling in `routes.create_share`: + +- New `ShareStorageUnavailableError` β†’ `503 {"detail":"Sharing is temporarily unavailable. Please try again later."}` (bucket unset / boto3 missing β€” a config problem, not the user's fault). +- The 400 KB path essentially disappears for the body. If a *control* item still somehow exceeds limits (it can't in practice β€” it's a handful of strings), the generic `500` remains as the final catch-all, but with the item-size failure mode designed out. +- The catch-all `except Exception` stays as a backstop but the common failure now has a real message. + +> The SPA already renders `detail` from error responses, so a clearer `detail` string improves UX with zero frontend change. (A dedicated toast copy tweak is optional and out of scope here.) + +--- + +## 6. Infrastructure changes + +### 6.1 Construct + +`SharedConversationsConstruct` gains a `public readonly bucket: s3.Bucket` (Β§4.1) and an SSM publication for symmetry with the table: + +```ts +new ssm.StringParameter(this, 'SharedConversationsBucketNameParameter', { + parameterName: `/${config.projectPrefix}/shares/shared-conversations-bucket-name`, + stringValue: this.bucket.bucketName, + ... +}); +``` + +### 6.2 Platform wiring + +- `platform-stack.ts`: `this.sharedConversationsBucket = sharedConversations.bucket;` and add to the `PlatformComputeRefs` bundle passed to app-api (alongside the existing `sharedConversationsTable`). +- `platform-compute-refs.ts`: add `sharedConversationsBucket: s3.IBucket;`. +- `app-api-environment.ts`: thread `sharedConversationsBucketName` and set env `SHARED_CONVERSATIONS_BUCKET_NAME: params.sharedConversationsBucketName`. + +### 6.3 IAM grant (app-api task role) + +Add to `app-api-iam-grants.ts`, mirroring `MemorySpacesBucketReadWrite`: + +```ts +taskRole.addToPrincipalPolicy(new iam.PolicyStatement({ + sid: 'SharedConversationsBucketReadWrite', + effect: iam.Effect.ALLOW, + actions: ['s3:GetObject', 's3:PutObject', 's3:DeleteObject', 's3:ListBucket'], + resources: [ + props.refs.sharedConversationsBucket.bucketArn, + `${props.refs.sharedConversationsBucket.bucketArn}/*`, + ], +})); +``` + +app-api is the only reader/writer (create writes; view/export read; revoke deletes). Inference-api is **not** granted β€” it never touches shares. This respects the service boundary (`feedback_service_boundaries`). + +### 6.4 Local dev + +Add `SHARED_CONVERSATIONS_BUCKET_NAME=` to `backend/src/.env.example` next to the existing `SHARED_CONVERSATIONS_TABLE_NAME`. Locally unset β†’ `store.enabled == False` β†’ create returns the `503` "temporarily unavailable" message rather than a confusing 500, which is the honest state of a machine with no bucket. + +--- + +## 7. Backward-compat & rollout + +1. **Deploy order:** `platform.yml` (CDK β€” creates bucket, grant, env) **before** the `backend.yml` app-api image that writes `body_ref`. Because reads fall back to inline, an app-api that predates the bucket keeps working; an app-api that has the code but no bucket yet returns the `503` on *create* only (reads unaffected). So the CDK deploy must land first, but there is no hard coupling that breaks reads. +2. **Existing inline shares** keep resolving via the legacy read path indefinitely. No migration. +3. **Rollback:** reverting the app-api image returns to inline writes (large shares fail again, as before) but every S3-backed share created in the interim becomes unreadable by the old code (it doesn't know `body_ref`). This is the one rollback caveat β€” **note it in the PR.** Mitigation if that matters: land the *read* support (fallback + `body_ref` handling) in an earlier, separately-deployed change than the *write* switch, so a rollback target already understands `body_ref`. See Β§9 phasing. + +--- + +## 8. Testing + +Existing suites to extend (all present): `tests/routes/test_shares.py`, `test_share_export.py`, `test_share_properties.py`, `tests/apis/app_api/shares/`. + +- **Store unit tests** (moto S3): putβ†’get round-trip, dedupe on identical body, `NoSuchKey`β†’typed error, `enabled` False when bucket unset. +- **create_share** writes `body_ref` and **no** inline `messages`; body object in S3 decodes to the snapshot. +- **Large conversation:** a snapshot > 400 KB now succeeds (regression test for the actual bug β€” build a message list that exceeds 400 KB inline and assert `create_share` returns 201). +- **Legacy read:** an item with inline `messages` and no `body_ref` still resolves via `get_shared_conversation` and `export_shared_conversation`. +- **Revoke** deletes the S3 object; `delete_shares_for_session` cleans up objects. +- **Storage-unavailable:** bucket unset β†’ `create_share` β†’ `503` with the friendly detail. +- **Float handling:** a message with a float (e.g. a cost/score) round-trips through S3 JSON without the `Decimal` dance and validates back into `MessageResponse`. +- **Infra:** `infrastructure` jest snapshot updated for the new bucket + grant; `npx tsc --noEmit` + `npx cdk synth` clean. + +--- + +## 9. Phasing (PRs into `develop`) + +Small, reviewable, rollback-aware: + +- **PR-1 β€” Read support + infra (no behavior change to writes).** + CDK: bucket, SSM, compute-ref, env, IAM grant. Backend: `ShareSnapshotStore`, `_load_snapshot_body` with legacy fallback wired into read/export paths. **Writes still inline.** Deploying this makes every running app-api `body_ref`-aware *before* anything writes one β€” this is the rollback-safety anchor (Β§7.3). +- **PR-2 β€” Switch writes to S3.** + `create_share` serializes + `store.put` + `body_ref` item, drops inline body; revoke/session-cleanup delete objects; `ShareStorageUnavailableError` + friendly `503`. Tests: large-conversation regression, storage-unavailable. +- **PR-3 (optional) β€” SPA copy polish.** + Nicer toast for the `503`/failure states. Pure frontend; can be skipped if the `detail` passthrough is deemed sufficient. + +Each PR branches from `develop`, targets `develop`, conventional-commit titled (`feat(shares): ...`). + +--- + +## 10. Alternatives considered + +- **Size-gated inline vs. S3.** Rejected β€” two code paths, and the gate itself has to model DynamoDB item overhead correctly or it just relocates the cliff (Β§3). +- **Compress inline (gzip the body attribute).** Buys a ~3–5Γ— headroom but does not remove the ceiling; a big enough conversation with images still exceeds 400 KB compressed, and it adds an opaque binary blob to DynamoDB that PITR/console can't inspect. Rejected as a stopgap, not a fix. +- **Split across multiple DynamoDB items (chunking).** Reintroduces multi-item transactional complexity (partial writes, ordering) that S3 avoids entirely. Rejected. +- **Reference the live session instead of snapshotting.** Would eliminate the copy, but breaks the share's point-in-time semantics and its independence from later edits/deletes of the source session (a share deliberately survives the owner continuing or deleting the conversation β€” see `delete_shares_for_session`'s "exported conversations unaffected" note). Rejected β€” changes product behavior. + +--- + +## 11. Open questions + +1. **Dedicated bucket vs. reuse** (Β§4.1 decision point). Spec assumes dedicated, per sibling-feature precedent. Confirm. +2. **Backfill of legacy inline shares** β€” spec says skip (they're already small). Confirm no desire to normalize storage. +3. **SPA toast copy** β€” is the `detail` passthrough enough, or do we want PR-3? +4. **Object encryption** β€” S3-managed (SSE-S3/AES256) matches every sibling bucket. A share body is the same data-class as the session it came from (behind the same auth), so no case for KMS/CMK here (`feedback_governance_via_identity_claims`). Confirm. diff --git a/frontend/ai.client/package-lock.json b/frontend/ai.client/package-lock.json index e9c3db137..5d2980d48 100644 --- a/frontend/ai.client/package-lock.json +++ b/frontend/ai.client/package-lock.json @@ -1,12 +1,12 @@ { "name": "ai.client", - "version": "1.5.0", + "version": "1.6.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "ai.client", - "version": "1.5.0", + "version": "1.6.0", "dependencies": { "@angular/cdk": "21.2.14", "@angular/common": "21.2.17", diff --git a/frontend/ai.client/package.json b/frontend/ai.client/package.json index 2351021cb..4f83a631b 100644 --- a/frontend/ai.client/package.json +++ b/frontend/ai.client/package.json @@ -1,6 +1,6 @@ { "name": "ai.client", - "version": "1.5.0", + "version": "1.6.0", "scripts": { "ng": "ng", "start": "ng serve", diff --git a/frontend/ai.client/src/app/admin/manage-models/model-form.page.html b/frontend/ai.client/src/app/admin/manage-models/model-form.page.html index 8e64898c2..5481bf437 100644 --- a/frontend/ai.client/src/app/admin/manage-models/model-form.page.html +++ b/frontend/ai.client/src/app/admin/manage-models/model-form.page.html @@ -296,10 +296,11 @@

Access contr

- Select which application roles can access this model. + Select which application roles can access this model. Saving grants the model + to each selected role, exactly as if you had added it on the role's own page.

@if (rolesResource.isLoading()) {
Loading roles…
@@ -309,10 +310,11 @@

Access contr

} @else {
- @for (role of availableAppRoles(); track role.roleId) { + @for (role of selectableAppRoles(); track role.roleId) {
diff --git a/frontend/ai.client/src/app/admin/manage-models/model-form.page.ts b/frontend/ai.client/src/app/admin/manage-models/model-form.page.ts index 730c49eec..86e83c353 100644 --- a/frontend/ai.client/src/app/admin/manage-models/model-form.page.ts +++ b/frontend/ai.client/src/app/admin/manage-models/model-form.page.ts @@ -282,6 +282,20 @@ export class ModelFormPage implements OnInit { readonly rolesResource = this.appRolesService.rolesResource; readonly availableAppRoles = computed(() => this.appRolesService.getEnabledRoles()); + /** + * Roles that already reach this model without a direct grant β€” they hold a + * wildcard ('*') grant or inherit it from a parent. The server derives these; + * they can't be toggled here, so they're shown separately from the picker + * rather than as unchecked boxes that would look like "no access". + */ + readonly inheritedAppRoles = signal([]); + readonly selectableAppRoles = computed(() => + this.availableAppRoles().filter(r => !this.inheritedAppRoles().includes(r.roleId)), + ); + readonly inheritedRoleDetails = computed(() => + this.availableAppRoles().filter(r => this.inheritedAppRoles().includes(r.roleId)), + ); + // Form state readonly isEditMode = signal(false); readonly modelId = signal(null); @@ -303,7 +317,10 @@ export class ModelFormPage implements OnInit { outputModalities: this.fb.control([], { nonNullable: true, validators: [Validators.required] }), maxInputTokens: this.fb.control(0, { nonNullable: true, validators: [Validators.required, Validators.min(1)] }), maxOutputTokens: this.fb.control(null, { validators: [Validators.min(1)] }), - allowedAppRoles: this.fb.control([], { nonNullable: true, validators: [Validators.required] }), + // No `required` validator: a model can legitimately have zero direct grants + // and still be reachable via a role's wildcard grant or inheritance, so + // forcing a selection would block saving those models. + allowedAppRoles: this.fb.control([], { nonNullable: true }), availableToRoles: this.fb.control([], { nonNullable: true }), enabled: this.fb.control(true, { nonNullable: true }), isDefault: this.fb.control(false, { nonNullable: true }), @@ -803,6 +820,10 @@ export class ModelFormPage implements OnInit { try { const model = await this.managedModelsService.getModel(id); + // Roles that reach this model via a wildcard grant or inheritance. Held + // outside the form: they're server-derived and not editable here. + this.inheritedAppRoles.set(model.inheritedAppRoles ?? []); + // Populate form with model data this.modelForm.patchValue({ modelId: model.modelId, diff --git a/frontend/ai.client/src/app/admin/manage-models/models/managed-model.model.ts b/frontend/ai.client/src/app/admin/manage-models/models/managed-model.model.ts index ad461f603..f9dc66a0f 100644 --- a/frontend/ai.client/src/app/admin/manage-models/models/managed-model.model.ts +++ b/frontend/ai.client/src/app/admin/manage-models/models/managed-model.model.ts @@ -82,8 +82,17 @@ export interface ManagedModel { maxOutputTokens: number | null; /** Lifecycle status of the model (e.g., 'ACTIVE', 'LEGACY') */ modelLifecycle?: string | null; - /** AppRole IDs that have access to this model (preferred over availableToRoles) */ + /** + * AppRole IDs that grant this model DIRECTLY. Derived server-side from the + * role records (the source of truth); saving the form writes it back through + * to each role's grantedModels. + */ allowedAppRoles: string[]; + /** + * AppRole IDs that grant this model indirectly β€” via a wildcard ('*') grant or + * inheritance from a parent role. Read-only: change these by editing the role. + */ + inheritedAppRoles?: string[]; /** @deprecated Legacy JWT role names - use allowedAppRoles instead */ availableToRoles: string[]; /** Whether the model is enabled for use */ diff --git a/frontend/ai.client/src/app/assistants/assistant-form/services/preview-chat.service.ts b/frontend/ai.client/src/app/assistants/assistant-form/services/preview-chat.service.ts index b19259c82..b7d68447c 100644 --- a/frontend/ai.client/src/app/assistants/assistant-form/services/preview-chat.service.ts +++ b/frontend/ai.client/src/app/assistants/assistant-form/services/preview-chat.service.ts @@ -261,6 +261,11 @@ export class PreviewChatService { }, body: JSON.stringify(requestBody), signal: this.abortController.signal, + // Keep the stream alive when the tab is backgrounded. The library + // default aborts + reopens the POST on `visibilitychange`, which would + // re-issue the same turn on tab return. See the detailed note in + // chat-http.service.ts. + openWhenHidden: true, onmessage: (msg: EventSourceMessage) => { this.handleStreamEvent(msg, callbacks); }, diff --git a/frontend/ai.client/src/app/assistants/services/web-source.service.spec.ts b/frontend/ai.client/src/app/assistants/services/web-source.service.spec.ts new file mode 100644 index 000000000..3816332df --- /dev/null +++ b/frontend/ai.client/src/app/assistants/services/web-source.service.spec.ts @@ -0,0 +1,95 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { TestBed } from '@angular/core/testing'; +import { HttpTestingController, provideHttpClientTesting } from '@angular/common/http/testing'; +import { provideHttpClient } from '@angular/common/http'; +import { signal } from '@angular/core'; +import { WebSourceService, WebSourceError } from './web-source.service'; +import { ConfigService } from '../../services/config.service'; +import { CrawlJob } from '../models/web-source.model'; + +const BASE = 'http://localhost:8000/assistants/assistant1/web-sources'; + +function stubCrawl(overrides: Partial = {}): CrawlJob { + return { + crawlId: 'CRAWL-abc123', + assistantId: 'assistant1', + rootUrl: 'https://example.com/docs/', + status: 'complete', + settings: { maxDepth: 1, maxPages: 25 }, + discoveredCount: 12, + fetchedCount: 10, + failedCount: 2, + startedAt: '2026-07-14T00:00:00Z', + startedByUserId: 'user-1', + ...overrides, + } as CrawlJob; +} + +describe('WebSourceService', () => { + let service: WebSourceService; + let httpMock: HttpTestingController; + + beforeEach(() => { + TestBed.resetTestingModule(); + TestBed.configureTestingModule({ + providers: [ + provideHttpClient(), + provideHttpClientTesting(), + WebSourceService, + { provide: ConfigService, useValue: { appApiUrl: signal('http://localhost:8000') } }, + ], + }); + service = TestBed.inject(WebSourceService); + httpMock = TestBed.inject(HttpTestingController); + }); + + afterEach(() => { + httpMock.match(() => true); + TestBed.resetTestingModule(); + }); + + it('should list crawls', async () => { + const crawls = [stubCrawl()]; + + const promise = service.listCrawls('assistant1'); + + await vi.waitFor(() => { + const req = httpMock.expectOne(`${BASE}/crawls`); + expect(req.request.method).toBe('GET'); + req.flush({ crawls }); + }); + + expect(await promise).toEqual(crawls); + }); + + it('should delete a crawl', async () => { + const promise = service.deleteCrawl('assistant1', 'CRAWL-abc123'); + + await vi.waitFor(() => { + const req = httpMock.expectOne(`${BASE}/crawls/CRAWL-abc123`); + expect(req.request.method).toBe('DELETE'); + req.flush(null, { status: 204, statusText: 'No Content' }); + }); + + await expect(promise).resolves.toBeUndefined(); + }); + + it('should surface the server message when a running crawl refuses deletion', async () => { + const promise = service.deleteCrawl('assistant1', 'CRAWL-abc123'); + const caught = promise.catch((err: unknown) => err); + + await vi.waitFor(() => { + const req = httpMock.expectOne(`${BASE}/crawls/CRAWL-abc123`); + req.flush( + { detail: 'This crawl is still running. Wait for it to finish, then remove it.' }, + { status: 409, statusText: 'Conflict' }, + ); + }); + + const error = (await caught) as WebSourceError; + expect(error).toBeInstanceOf(WebSourceError); + expect(error.status).toBe(409); + expect(error.code).toBe('HTTP_409'); + expect(error.message).toContain('still running'); + }); +}); diff --git a/frontend/ai.client/src/app/assistants/services/web-source.service.ts b/frontend/ai.client/src/app/assistants/services/web-source.service.ts index 8ce07c0dd..509b4f26a 100644 --- a/frontend/ai.client/src/app/assistants/services/web-source.service.ts +++ b/frontend/ai.client/src/app/assistants/services/web-source.service.ts @@ -106,6 +106,24 @@ export class WebSourceService { } } + /** + * Remove a web source β€” its crawl record, every page it ingested, and any + * sync policy covering it. Rejects with `HTTP_409` while the crawl is still + * running (its pages are still being written). + */ + async deleteCrawl(assistantId: string, crawlId: string): Promise { + try { + await firstValueFrom( + this.http.delete( + `${this.baseUrl()}/assistants/${encodeURIComponent(assistantId)}/web-sources/crawls/${encodeURIComponent(crawlId)}`, + this.requestOptions(), + ), + ); + } catch (err) { + throw this.toError(err, 'Failed to remove web source'); + } + } + private toError(err: unknown, fallback: string): WebSourceError { if (err instanceof HttpErrorResponse) { const detail = diff --git a/frontend/ai.client/src/app/knowledge-base/knowledge-base-section.component.html b/frontend/ai.client/src/app/knowledge-base/knowledge-base-section.component.html index da4f0bdc4..0bb0b68f3 100644 --- a/frontend/ai.client/src/app/knowledge-base/knowledge-base-section.component.html +++ b/frontend/ai.client/src/app/knowledge-base/knowledge-base-section.component.html @@ -224,6 +224,19 @@

Web sour /> } + + } diff --git a/frontend/ai.client/src/app/knowledge-base/knowledge-base-section.component.ts b/frontend/ai.client/src/app/knowledge-base/knowledge-base-section.component.ts index 378491e70..cbea746cc 100644 --- a/frontend/ai.client/src/app/knowledge-base/knowledge-base-section.component.ts +++ b/frontend/ai.client/src/app/knowledge-base/knowledge-base-section.component.ts @@ -43,6 +43,10 @@ import { SyncPolicyControlComponent, SyncIntervalSelection, } from '../assistants/components/sync-policy-control.component'; +import { + ConfirmationDialogComponent, + ConfirmationDialogData, +} from '../components/confirmation-dialog'; import { UserConnectorsService } from '../settings/connectors/services/user-connectors.service'; import { OAuthConsentService } from '../services/oauth-consent/oauth-consent.service'; import { ToastService } from '../services/toast/toast.service'; @@ -775,6 +779,85 @@ export class KnowledgeBaseSectionComponent implements OnDestroy { } } + /** + * Remove a web source: the crawl record, every page it added to the + * knowledge base, and any sync policy covering it. Confirmed first β€” unlike + * a single document this can take dozens of pages with it, so the dialog + * names the count. + */ + async removeWebSource(crawl: CrawlJob): Promise { + const recordId = this.id(); + if (!recordId) { + return; + } + + const pages = crawl.fetchedCount; + const pageCount = `${pages} page${pages === 1 ? '' : 's'}`; + const dialogRef = this.dialog.open( + ConfirmationDialogComponent, + { + data: { + title: 'Remove web source', + message: + `${crawl.rootUrl} and the ${pageCount} it added will be removed from ` + + `this knowledge base, along with any sync schedule on it. ` + + `This cannot be undone.`, + confirmText: 'Remove', + destructive: true, + }, + }, + ); + if ((await firstValueFrom(dialogRef.closed)) !== true) { + return; + } + + // Optimistic, like deleteDocument: drop the source and its pages up front + // so the click lands immediately, and restore both lists if the call fails + // (a still-running crawl is refused with a 409). + const previousCrawls = this.webCrawls(); + const previousDocuments = this.uploadedDocuments(); + const pageDocuments = previousDocuments.filter((doc) => this.isPageOf(doc, crawl)); + + this.webCrawls.update((crawls) => crawls.filter((c) => c.crawlId !== crawl.crawlId)); + this.uploadedDocuments.update((docs) => docs.filter((doc) => !this.isPageOf(doc, crawl))); + // Stop polling any page still mid-processing, so its spinner can't + // reappear on the next tick before the GET starts 404ing. + this.pollingDocuments.update((set) => { + const next = new Set(set); + for (const doc of pageDocuments) { + next.delete(doc.documentId); + } + return next; + }); + + try { + await this.webSourceService.deleteCrawl(recordId, crawl.crawlId); + // The backend cascades the sync policy with its source β€” mirror that + // locally so the control disappears with the row. + const covering = this.syncPolicyFor(crawl.crawlId); + if (covering) { + this.removePolicy(covering.policyId); + } + this.toast.success('Web source removed.'); + } catch (error) { + this.webCrawls.set(previousCrawls); + this.uploadedDocuments.set(previousDocuments); + const message = error instanceof Error ? error.message : 'Failed to remove web source.'; + this.toast.error(message); + } + } + + /** + * Whether a document is one of the pages a crawl produced. Mirrors the + * backend's rule (a `web` document whose source URL sits under the crawl + * root) so the optimistic removal matches what the server actually deletes. + */ + private isPageOf(doc: Document, crawl: CrawlJob): boolean { + return ( + doc.sourceConnectorId === 'web' && !!doc.sourceFileId?.startsWith(crawl.rootUrl) + ); + } + // ── KB sync policy actions ────────────────────────────────────────────── /** True while the reconnect consent popup is open (guards the abort effect). */ diff --git a/frontend/ai.client/src/app/session/services/chat/chat-http.service.spec.ts b/frontend/ai.client/src/app/session/services/chat/chat-http.service.spec.ts index 0ec766f24..73d516207 100644 --- a/frontend/ai.client/src/app/session/services/chat/chat-http.service.spec.ts +++ b/frontend/ai.client/src/app/session/services/chat/chat-http.service.spec.ts @@ -33,7 +33,7 @@ describe('ChatHttpService', () => { { provide: StreamParserService, useValue: { getCurrentStreamId: vi.fn().mockReturnValue('stream-1'), parseEventSourceMessage: vi.fn() } }, { provide: ChatStateService, useValue: { abortRequest: vi.fn(), setChatLoading: vi.fn(), setLastTurnInterrupted: vi.fn(), seedSessionAggregates: vi.fn(), createAbortController: vi.fn().mockReturnValue(new AbortController()) } }, { provide: MessageMapService, useValue: { endStreaming: vi.fn() } }, - { provide: ErrorService, useValue: { handleHttpError: vi.fn() } }, + { provide: ErrorService, useValue: { handleHttpError: vi.fn(), addError: vi.fn() } }, ], }); service = TestBed.inject(ChatHttpService); @@ -144,4 +144,49 @@ describe('ChatHttpService', () => { vi.useRealTimers(); vi.restoreAllMocks(); }); + + it('surfaces a soft "Already responding" notice (not a hard error) on a 409 single-flight rejection', async () => { + // The inference-api single-flight guard rejects a duplicate turn while the + // prior one is still streaming server-side; the BFF relays it as 409. + const body = JSON.stringify({ + detail: 'A response is already streaming for this conversation. Wait for it to finish before sending another message.', + }); + vi.spyOn(globalThis, 'fetch').mockResolvedValue( + new Response(body, { status: 409, headers: { 'content-type': 'application/json' } }), + ); + const errorSvc = TestBed.inject(ErrorService) as any; + + await expect( + service.sendChatRequest({ session_id: 's1', message: 'hi' }), + ).rejects.toMatchObject({ name: 'AlreadyStreamingError' }); + + // Gentle, dismissible notice carrying the server's explanation β€” NOT the + // "Chat Request Failed" / network-error paths. + expect(errorSvc.addError).toHaveBeenCalledTimes(1); + const [title, message] = errorSvc.addError.mock.calls[0]; + expect(title).toBe('Already responding'); + expect(message).toContain('already streaming'); + // Loading is cleared so the user can retry once the prior turn finishes. + expect(chatStateService.setChatLoading).toHaveBeenCalledWith('s1', false); + vi.restoreAllMocks(); + }); + + it('unwraps a double-encoded 409 body from the BFF proxy', async () => { + // app-api relays inference-api's body verbatim inside its own `detail`, + // so the payload can be `{"detail":"{\"detail\":\"…\"}"}`. + const inner = JSON.stringify({ detail: 'This conversation is busy generating a response.' }); + const body = JSON.stringify({ detail: inner }); + vi.spyOn(globalThis, 'fetch').mockResolvedValue( + new Response(body, { status: 409, headers: { 'content-type': 'application/json' } }), + ); + const errorSvc = TestBed.inject(ErrorService) as any; + + await expect( + service.sendChatRequest({ session_id: 's1', message: 'hi' }), + ).rejects.toMatchObject({ name: 'AlreadyStreamingError' }); + + const [, message] = errorSvc.addError.mock.calls[0]; + expect(message).toBe('This conversation is busy generating a response.'); + vi.restoreAllMocks(); + }); }); diff --git a/frontend/ai.client/src/app/session/services/chat/chat-http.service.ts b/frontend/ai.client/src/app/session/services/chat/chat-http.service.ts index 02649f6ab..d49f43919 100644 --- a/frontend/ai.client/src/app/session/services/chat/chat-http.service.ts +++ b/frontend/ai.client/src/app/session/services/chat/chat-http.service.ts @@ -33,6 +33,45 @@ class UnauthorizedError extends Error { this.name = 'UnauthorizedError'; } } +/** + * Thrown by the SSE `onopen` when the BFF returned 409 β€” the inference-api + * single-flight guard rejected this turn because another turn for the same + * session is still streaming server-side (a second tab/device, a transport + * retry, or a Stop whose server turn hasn't ended, since a client abort does + * not propagate through the AgentCore Runtime). Not a failure: `onerror` + * surfaces it as a gentle notice rather than an error toast. + */ +class AlreadyStreamingError extends Error { + constructor(message: string) { + super(message); + this.name = 'AlreadyStreamingError'; + } +} + +/** + * Best-effort human message for a 409 from `/chat/stream`. The app-api proxy + * relays the inference-api body verbatim inside its own `detail`, so the + * payload can be double-encoded (`{"detail":"{\"detail\":\"…\"}"}`) β€” unwrap one + * nested layer, and fall back to a fixed message if the body is unreadable. + */ +async function parseConflictMessage(response: Response): Promise { + const fallback = + 'This conversation is still generating a response. Wait for it to finish before sending another message.'; + try { + const data = await response.json(); + let detail: unknown = data?.detail ?? data?.error?.message ?? data?.message; + if (typeof detail === 'string' && detail.trim().startsWith('{')) { + try { + detail = (JSON.parse(detail) as { detail?: unknown })?.detail ?? detail; + } catch { + // Not nested JSON after all β€” keep the string as-is. + } + } + return typeof detail === 'string' && detail.trim() ? detail : fallback; + } catch { + return fallback; + } +} interface GenerateTitleRequest { session_id: string; @@ -125,6 +164,23 @@ export class ChatHttpService { }, body: JSON.stringify(requestObject), signal: abortController.signal, + // Keep the stream open when the tab is backgrounded. With the library + // default (`openWhenHidden: false`) fetch-event-source aborts the + // connection on `visibilitychange` to hidden and REOPENS it β€” issuing + // a brand-new `POST /invocations` for the SAME turn β€” when the tab + // becomes visible again. That reopen happens inside the library, reusing + // this request and bypassing our per-session double-submit + streamId + // supersession guards, so nothing here catches it. Because a client + // abort does NOT propagate through the AgentCore Runtime data plane, the + // original backend agent keeps running; the reopened one runs the same + // turn concurrently, and both persist tool-use/tool-result events to the + // same AgentCore Memory session β€” corrupting history (duplicate / + // interleaved toolResult turns) and bricking the conversation with a + // Bedrock "toolResult blocks exceed toolUse blocks" ValidationException. + // Keeping the single stream alive across tab switches is also correct for + // long agentic turns. See the restore-time repair in + // TurnBasedSessionManager for the server-side safety net. + openWhenHidden: true, async onopen(response) { if (response.ok && response.headers.get('content-type')?.includes('text/event-stream')) { return; // everything's good @@ -151,6 +207,12 @@ export class ChatHttpService { } throw new FatalError(errorMessage); + } else if (response.status === 409) { + // Single-flight guard: another turn for this session is still + // streaming server-side. This is a benign rejection, not a + // failure β€” surface a gentle notice (handled in `onerror`) and + // don't tear down as fatal/retriable. + throw new AlreadyStreamingError(await parseConflictMessage(response)); } else if (response.status >= 400 && response.status < 500 && response.status !== 429) { // Client-side errors are usually non-retriable let errorMessage = `Request failed with status ${response.status}`; @@ -221,6 +283,14 @@ export class ChatHttpService { throw err; } + // 409 single-flight rejection is expected, not an error: the prior + // response is still generating. Show a soft, dismissible notice with + // the server's explanation rather than a "Chat Request Failed" toast. + if (err instanceof AlreadyStreamingError) { + this.errorService.addError('Already responding', err.message, undefined, undefined); + throw err; + } + // Display error message to user using ErrorService if (err instanceof FatalError) { this.errorService.addError('Chat Request Failed', err.message, undefined, undefined); diff --git a/infrastructure/lib/constructs/app-api/app-api-environment.ts b/infrastructure/lib/constructs/app-api/app-api-environment.ts index b1884c8cf..066ff0c19 100644 --- a/infrastructure/lib/constructs/app-api/app-api-environment.ts +++ b/infrastructure/lib/constructs/app-api/app-api-environment.ts @@ -91,6 +91,7 @@ export interface AppApiSsmParams { ragDocumentsBucketArn: string; sharedConversationsTableName: string; sharedConversationsTableArn: string; + sharedConversationsBucketName: string; memoryId: string; // Memory Spaces memorySpacesTableName: string; @@ -206,6 +207,7 @@ export function resolveAppApiParams( ragDocumentsBucketArn: refs.ragDocumentsBucket.bucketArn, sharedConversationsTableName: refs.sharedConversationsTable.tableName, sharedConversationsTableArn: refs.sharedConversationsTable.tableArn, + sharedConversationsBucketName: refs.sharedConversationsBucket.bucketName, memoryId: overrides.memoryId, // Memory Spaces memorySpacesTableName: refs.memorySpacesTable.tableName, @@ -265,6 +267,7 @@ export function buildAppApiEnvironment( COGNITO_DOMAIN_URL: params.cognitoDomainUrl, COGNITO_REGION: config.awsRegion, SHARED_CONVERSATIONS_TABLE_NAME: params.sharedConversationsTableName, + SHARED_CONVERSATIONS_BUCKET_NAME: params.sharedConversationsBucketName, BFF_SESSIONS_TABLE_NAME: params.bffSessionsTableName, BFF_COOKIE_SIGNING_KEY_ARN: params.bffCookieSigningKeyArn, BFF_COOKIE_DATA_KEY_SECRET_ARN: params.bffCookieDataKeySecretArn, diff --git a/infrastructure/lib/constructs/app-api/app-api-iam-grants.ts b/infrastructure/lib/constructs/app-api/app-api-iam-grants.ts index bf2f39dca..b29f59a7d 100644 --- a/infrastructure/lib/constructs/app-api/app-api-iam-grants.ts +++ b/infrastructure/lib/constructs/app-api/app-api-iam-grants.ts @@ -138,6 +138,23 @@ export function grantAppApiPermissions(props: AppApiIamGrantsProps): void { }), ); + // ── Shared-conversations snapshot bucket ── + // The share BODY (messages + metadata) is offloaded to S3 because a long + // conversation exceeds DynamoDB's 400 KB item limit; only a pointer stays + // in the shared-conversations table. app-api is the sole reader/writer: + // create writes, view/export read, revoke/session-cleanup delete. Mirrors + // MemorySpacesBucketReadWrite. See + // docs/specs/share-large-conversations-s3-offload.md. + const sharedConversationsBucketArn = props.refs.sharedConversationsBucket.bucketArn; + taskRole.addToPrincipalPolicy( + new iam.PolicyStatement({ + sid: 'SharedConversationsBucketReadWrite', + effect: iam.Effect.ALLOW, + actions: ['s3:GetObject', 's3:PutObject', 's3:DeleteObject', 's3:ListBucket'], + resources: [sharedConversationsBucketArn, `${sharedConversationsBucketArn}/*`], + }), + ); + // ── Core tables (OIDC, Users, Roles, API Keys, OAuth) ── const coreTables = [ { sid: 'OidcStateAccess', arn: props.refs.oidcStateTable.tableArn }, @@ -156,6 +173,7 @@ export function grantAppApiPermissions(props: AppApiIamGrantsProps): void { { sid: 'BffSessionsAccess', arn: props.refs.bffSessionsTable.tableArn }, { sid: 'VoiceTicketReplayAccess', arn: props.refs.voiceTicketReplayTable.tableArn }, { sid: 'UserFilesTableAccess', arn: props.refs.fileUploadTable.tableArn }, + { sid: 'SharedConversationsAccess', arn: props.refs.sharedConversationsTable.tableArn }, ]; for (const { sid, arn } of coreTables) { diff --git a/infrastructure/lib/constructs/data/shared-conversations-construct.ts b/infrastructure/lib/constructs/data/shared-conversations-construct.ts index 48b75b06f..a5925813f 100644 --- a/infrastructure/lib/constructs/data/shared-conversations-construct.ts +++ b/infrastructure/lib/constructs/data/shared-conversations-construct.ts @@ -1,8 +1,15 @@ +import * as cdk from 'aws-cdk-lib'; import * as dynamodb from 'aws-cdk-lib/aws-dynamodb'; +import * as s3 from 'aws-cdk-lib/aws-s3'; import * as ssm from 'aws-cdk-lib/aws-ssm'; import { Construct } from 'constructs'; -import { AppConfig, getResourceName, getRemovalPolicy } from '../../config'; +import { + AppConfig, + getAutoDeleteObjects, + getRemovalPolicy, + getResourceName, +} from '../../config'; export interface SharedConversationsConstructProps { config: AppConfig; @@ -18,9 +25,18 @@ export interface SharedConversationsConstructProps { * PK: share_id * GSI: SessionShareIndex β€” lookup by original session_id * GSI: OwnerShareIndex β€” list shares by owner, sorted by created_at + * + * The snapshot BODY (messages + metadata) is offloaded to the sibling S3 + * bucket, because a long conversation exceeds DynamoDB's 400 KB item limit + * (PutItem ValidationException). The DynamoDB item keeps only the small + * control fields plus an S3 pointer (`body_ref`); the bytes live in S3. + * Mirrors the Memory Spaces / Artifacts / Skills S3-offload pattern. + * Private, server-side-only (app-api reads/writes; never loaded cross-origin). + * See docs/specs/share-large-conversations-s3-offload.md. */ export class SharedConversationsConstruct extends Construct { public readonly table: dynamodb.Table; + public readonly bucket: s3.Bucket; constructor( scope: Construct, @@ -31,6 +47,25 @@ export class SharedConversationsConstruct extends Construct { const { config } = props; + // Snapshot-body bucket β€” private, bytes-only, fetched server-side by + // app-api. No expiration lifecycle rule: a share's object lives exactly + // as long as its DynamoDB row and is deleted explicitly on revoke / + // session cleanup. Age-based reaping would break live shares. + this.bucket = new s3.Bucket(this, 'SharedConversationsBucket', { + bucketName: getResourceName(config, 'shared-conversations'), + encryption: s3.BucketEncryption.S3_MANAGED, + blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL, + enforceSSL: true, + lifecycleRules: [ + { + id: 'abort-stale-multipart', + abortIncompleteMultipartUploadAfter: cdk.Duration.days(7), + }, + ], + removalPolicy: getRemovalPolicy(config), + autoDeleteObjects: getAutoDeleteObjects(config), + }); + this.table = new dynamodb.Table(this, 'SharedConversationsTable', { tableName: getResourceName(config, 'shared-conversations'), partitionKey: { @@ -69,5 +104,12 @@ export class SharedConversationsConstruct extends Construct { tier: ssm.ParameterTier.STANDARD, }); + new ssm.StringParameter(this, 'SharedConversationsBucketNameParameter', { + parameterName: `/${config.projectPrefix}/shares/shared-conversations-bucket-name`, + stringValue: this.bucket.bucketName, + description: 'Shared conversations snapshot-body S3 bucket name', + tier: ssm.ParameterTier.STANDARD, + }); + } } diff --git a/infrastructure/lib/constructs/platform-compute-refs.ts b/infrastructure/lib/constructs/platform-compute-refs.ts index c3248afdc..80d91c37b 100644 --- a/infrastructure/lib/constructs/platform-compute-refs.ts +++ b/infrastructure/lib/constructs/platform-compute-refs.ts @@ -79,6 +79,7 @@ export interface PlatformComputeRefs { userMenuLinksTable: dynamodb.ITable; systemPromptsTable: dynamodb.ITable; sharedConversationsTable: dynamodb.ITable; + sharedConversationsBucket: s3.IBucket; fileUploadBucket: s3.IBucket; fileUploadTable: dynamodb.ITable; diff --git a/infrastructure/lib/platform-stack.ts b/infrastructure/lib/platform-stack.ts index c5953b75e..48bfecd49 100644 --- a/infrastructure/lib/platform-stack.ts +++ b/infrastructure/lib/platform-stack.ts @@ -162,6 +162,7 @@ export class PlatformStack extends cdk.Stack { public readonly userMenuLinksTable: dynamodb.ITable; public readonly systemPromptsTable: dynamodb.ITable; public readonly sharedConversationsTable: dynamodb.ITable; + public readonly sharedConversationsBucket: s3.IBucket; public readonly fileUploadBucket: s3.IBucket; public readonly fileUploadTable: dynamodb.ITable; @@ -369,6 +370,7 @@ export class PlatformStack extends cdk.Stack { { config }, ); this.sharedConversationsTable = sharedConversations.table; + this.sharedConversationsBucket = sharedConversations.bucket; // ============================================================ // RAG data @@ -689,6 +691,7 @@ export class PlatformStack extends cdk.Stack { userMenuLinksTable: this.userMenuLinksTable, systemPromptsTable: this.systemPromptsTable, sharedConversationsTable: this.sharedConversationsTable, + sharedConversationsBucket: this.sharedConversationsBucket, fileUploadBucket: this.fileUploadBucket, fileUploadTable: this.fileUploadTable, ragDocumentsBucket: this.ragDocumentsBucket, diff --git a/infrastructure/package-lock.json b/infrastructure/package-lock.json index 1afe8880e..629ee6744 100644 --- a/infrastructure/package-lock.json +++ b/infrastructure/package-lock.json @@ -1,12 +1,12 @@ { "name": "infrastructure", - "version": "1.5.0", + "version": "1.6.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "infrastructure", - "version": "1.5.0", + "version": "1.6.0", "dependencies": { "aws-cdk-lib": "2.260.0", "constructs": "10.6.0" diff --git a/infrastructure/package.json b/infrastructure/package.json index a47cb6df4..d930087a1 100644 --- a/infrastructure/package.json +++ b/infrastructure/package.json @@ -1,6 +1,6 @@ { "name": "infrastructure", - "version": "1.5.0", + "version": "1.6.0", "bin": { "infrastructure": "bin/infrastructure.js" }, diff --git a/infrastructure/test/platform-stack.test.ts b/infrastructure/test/platform-stack.test.ts index 2c31e0917..90d98bd4d 100644 --- a/infrastructure/test/platform-stack.test.ts +++ b/infrastructure/test/platform-stack.test.ts @@ -117,8 +117,9 @@ describe('PlatformStack', () => { it('creates all data buckets', () => { // file-uploads, SPA static, mcp-sandbox, rag-documents, fine-tuning-data, // artifacts-content, skill-resources (admin-managed Skills reference files), - // memory-spaces (Memory Spaces feature content bucket) - template.resourceCountIs('AWS::S3::Bucket', 8); + // memory-spaces (Memory Spaces feature content bucket), + // shared-conversations (share snapshot-body offload) + template.resourceCountIs('AWS::S3::Bucket', 9); }); }); diff --git a/infrastructure/test/security-policy.test.ts b/infrastructure/test/security-policy.test.ts index d5e430f62..624438f3b 100644 --- a/infrastructure/test/security-policy.test.ts +++ b/infrastructure/test/security-policy.test.ts @@ -160,6 +160,60 @@ describe('Security policy hardening', () => { }); }); + // ────────────────────────────────────────────────────────── + // 2b. App-api DynamoDB table grants + // ────────────────────────────────────────────────────────── + + describe('App-api shared-conversations table grant', () => { + // Regression guard: the shared-conversations table was threaded + // into the app-api container as an env var but never granted on + // the task role, so every /conversations/{id}/share PutItem and + // /conversations/{id}/shares Query returned AccessDeniedException + // (surfaced to users as a 500 "Failed to create share"). The grant + // lives in app-api-iam-grants.ts under Sid 'SharedConversationsAccess'. + it("app-api role can PutItem/Query/GetItem on the shared-conversations table (incl. its GSIs)", () => { + // Iterate both AWS::IAM::Policy AND AWS::IAM::ManagedPolicy because + // CDK auto-splits oversized inline policies into managed overflow + // policies attached to the same role. + const candidates: PolicyStatement[] = []; + for (const [, r] of Object.entries(template.findResources('AWS::IAM::Policy'))) { + const stmts = ((r.Properties as { PolicyDocument?: { Statement?: PolicyStatement[] } })?.PolicyDocument?.Statement) ?? []; + for (const s of stmts) candidates.push(s); + } + for (const [, r] of Object.entries(template.findResources('AWS::IAM::ManagedPolicy'))) { + const stmts = ((r.Properties as { PolicyDocument?: { Statement?: PolicyStatement[] } })?.PolicyDocument?.Statement) ?? []; + for (const s of stmts) candidates.push(s); + } + + const matches = candidates.filter((s) => s.Sid === 'SharedConversationsAccess'); + + if (matches.length === 0) { + throw new Error( + "Could not locate the shared-conversations DynamoDB grant. " + + "Looked for Sid 'SharedConversationsAccess' in AWS::IAM::Policy + AWS::IAM::ManagedPolicy. " + + "Without it, creating/listing conversation shares fails with AccessDeniedException. " + + "If the Sid was renamed, update this test.", + ); + } + + for (const s of matches) { + const actions = asArray(s.Action); + // The share service does PutItem (create), Query on SessionShareIndex + // (list/delete-for-session), and GetItem (retrieve a single share). + expect(actions).toContain('dynamodb:PutItem'); + expect(actions).toContain('dynamodb:Query'); + expect(actions).toContain('dynamodb:GetItem'); + // Never a wildcard resource. + const resources = asArray(s.Resource); + expect(resources).not.toContain('*'); + // Must cover the table's GSIs (SessionShareIndex) β€” the list/revoke + // paths Query that index, which requires an index/* resource entry. + const resourceStrs = resources.map((r) => JSON.stringify(r)); + expect(resourceStrs.some((r) => r.includes('index/'))).toBe(true); + } + }); + }); + // ────────────────────────────────────────────────────────── // 3. S3 hardening (encryption + public-access-block) // ──────────────────────────────────────────────────────────