diff --git a/CHANGELOG.md b/CHANGELOG.md index e9c90e9c..3f722f3f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -312,6 +312,17 @@ Entries before v0.5.0 were written retroactively as summaries. ### Security +- **Agent and MCP runtime sessions are bound to their owner.** New session ids created by + the Web UI start with a prefix derived from the signed-in user (`shared/session_owner.py`, + mirrored in `web-ui/src/lib/sessionId.ts`), and the agent and the MCP server admit such an + id only for that user. Each runtime process serves one user. Existing decks keep their + chat: a chat id from before this change is admitted when AgentCore Memory holds that + user's history for it, and a deck without history continues under the user's own id. The + agent forwards an owner-prefixed id to the MCP runtime (also for composer groups, + `-gN`). A refused call gets an error event with no details; a chat tab opened before + the update may be asked to reload the page once (new error code + `session_reload_required`). Deploy the Web UI first, then SdpmRuntime, then SdpmAgent. + Local mode is unchanged. - **The AgentCore runtimes allow only the operations the app uses.** The agent and MCP runtimes, their `DEFAULT` endpoints and the named endpoints `sdpm_agent_endpoint` and `sdpm_endpoint` now carry a resource-based policy (`AWS::BedrockAgentCore::ResourcePolicy`, diff --git a/agent/Dockerfile b/agent/Dockerfile index ccec2eea..b01f0273 100644 --- a/agent/Dockerfile +++ b/agent/Dockerfile @@ -27,7 +27,8 @@ RUN uv export --frozen --directory /tmp/workspace --package sdpm-agent --no-dev RUN useradd -m -u 1000 bedrock_agentcore USER bedrock_agentcore -COPY agent/basic_agent.py agent/cost_logger.py agent/errors.py agent/factory.py agent/composition.py agent/mcp_clients.py agent/message_hooks.py agent/model_profiles.py agent/resilience.py agent/session.py agent/stop_guard.py agent/streaming.py agent/turn_fence.py ./ +COPY agent/basic_agent.py agent/cost_logger.py agent/errors.py agent/factory.py agent/composition.py agent/mcp_clients.py agent/message_hooks.py agent/model_profiles.py agent/resilience.py agent/session.py agent/session_guard.py agent/stop_guard.py agent/streaming.py agent/turn_fence.py ./ +COPY shared/__init__.py shared/session_owner.py shared/ COPY agent/prompts/ prompts/ COPY agent/modes/ modes/ COPY agent/tools/ tools/ diff --git a/agent/README.md b/agent/README.md index e73ef090..c2934f3e 100644 --- a/agent/README.md +++ b/agent/README.md @@ -44,7 +44,7 @@ Three connection patterns are supported. Each returns an `MCPClient` instance ad ### Pattern 1: AgentCore Runtime + JWT Bearer -For MCP servers deployed on Amazon Bedrock AgentCore Runtime. The caller's JWT is forwarded as-is for authentication and user_id propagation, and the caller's session id is forwarded as `Mcp-Session-Id` so that AgentCore routes the requests of one session to the same microVM. +For MCP servers deployed on Amazon Bedrock AgentCore Runtime. The caller's JWT is forwarded as-is for authentication and user_id propagation, and the caller's session id is forwarded as `Mcp-Session-Id`, with the owner prefix of the token's user, so that AgentCore routes the requests of one session to the same microVM. ```python def _mcp_agentcore_runtime(jwt_token: str, session_id: str = "") -> MCPClient: diff --git a/agent/README_ja.md b/agent/README_ja.md index fcaf6deb..1a33d3db 100644 --- a/agent/README_ja.md +++ b/agent/README_ja.md @@ -32,7 +32,7 @@ User → AgentCore Runtime → Agent ### Pattern 1: AgentCore Runtime + JWT Bearer -AgentCore Runtime 上の MCP サーバーに JWT を転送して接続する。あわせて呼び出し元のセッション ID を `Mcp-Session-Id` として転送し、同一セッションのリクエストが同じ microVM に routing されるようにする。 +AgentCore Runtime 上の MCP サーバーに JWT を転送して接続する。あわせて呼び出し元のセッション ID を、トークンの利用者の持ち主プレフィックス付きで `Mcp-Session-Id` として転送し、同一セッションのリクエストが同じ microVM に routing されるようにする。 ```python def _mcp_agentcore_runtime(jwt_token: str, session_id: str = "") -> MCPClient: diff --git a/agent/basic_agent.py b/agent/basic_agent.py index 12a170a9..22e2368f 100644 --- a/agent/basic_agent.py +++ b/agent/basic_agent.py @@ -19,6 +19,7 @@ from bedrock_agentcore.runtime import BedrockAgentCoreApp +import session_guard import stop_guard from errors import BAD_INPUT, AUTH, error_event from factory import create_agent @@ -80,6 +81,13 @@ async def agent_stream(payload, context): yield error_event("Could not extract user_id from JWT sub claim", code=AUTH) return + # Only the session's owner gets in (session_guard.py); before the cancel registry, + # so a refused call never touches a running turn. + session_id, refusal = await session_guard.admit(session_id, getattr(context, "session_id", None), user_id) + if refusal is not None: + yield refusal + return + # Cancel previous request for the same session prev_cancel = _cancel_events.get(session_id) if prev_cancel is not None: diff --git a/agent/errors.py b/agent/errors.py index 2b2afd6b..a9e3e087 100644 --- a/agent/errors.py +++ b/agent/errors.py @@ -17,6 +17,8 @@ BAD_INPUT = "bad_input" MAX_OUTPUT = "max_output" INTERNAL = "internal" +# Sent only with an explicit code: the page holds a session id it has to replace. +SESSION_RELOAD_REQUIRED = "session_reload_required" _PATTERNS: list[tuple[tuple[str, ...], str]] = [ (("Too much media", "too long"), CONVERSATION_TOO_LONG), diff --git a/agent/mcp_clients.py b/agent/mcp_clients.py index 4ffd9dbd..a8445888 100644 --- a/agent/mcp_clients.py +++ b/agent/mcp_clients.py @@ -13,6 +13,8 @@ from mcp.client.streamable_http import streamablehttp_client from strands.tools.mcp import MCPClient +from shared.session_owner import owned_mcp_session_id, sub_from_jwt + logger = logging.getLogger("sdpm.agent") # MCP server definitions: (display_name, required) @@ -61,7 +63,10 @@ def mcp_agentcore_runtime( # after is back to 0.18s. No re-initialize fallback is needed. This holds # for the stateless server we run (`stateless_http=True` in # servers/remote/server.py); stateful MCP is documented to 404 instead. - headers["Mcp-Session-Id"] = session_id + # + # The id always carries the owner prefix of the token's user: an id of an + # earlier format gets it prepended (shared/session_owner.py). + headers["Mcp-Session-Id"] = owned_mcp_session_id(session_id, sub_from_jwt(jwt_token)) kwargs: dict = {} if tool_filters: diff --git a/agent/session_guard.py b/agent/session_guard.py new file mode 100644 index 00000000..b5fa64ac --- /dev/null +++ b/agent/session_guard.py @@ -0,0 +1,88 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Admit an invocation only for the owner of its session. + +The AgentCore Runtime authorizer has validated the JWT before the request reaches +this container; ``sub`` comes from that token. A session id is admitted when: + +- it carries the owner prefix of ``sub`` (``shared.session_owner``), or +- it has an earlier format and AgentCore Memory already holds this user's events + for it (sessions created before ids carried the prefix); + +and this process serves no other user. Refusals carry a code and no details. +""" + +import asyncio +import logging +import os + +from errors import AUTH, BAD_INPUT, SERVICE_UNAVAILABLE, SESSION_RELOAD_REQUIRED, error_event +from shared.session_owner import ProcessOwner, belongs_to, is_owner_prefixed + +logger = logging.getLogger("sdpm.agent") + +# The one user this process serves (set by the first admitted invocation). +_process_owner = ProcessOwner() +# (sub, session id) pairs of earlier-format ids confirmed through Memory. +_owned_legacy: set[tuple[str, str]] = set() +_client = None + + +def _memory_client(): + global _client + if _client is None: + import boto3 + from botocore.config import Config + + region = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION", "us-east-1")) + _client = boto3.client( + "bedrock-agentcore", + region_name=region, + config=Config(connect_timeout=2, read_timeout=5, retries={"max_attempts": 2}), + ) + return _client + + +def _has_own_events(memory_id: str, sub: str, session_id: str) -> bool: + response = _memory_client().list_events( + memoryId=memory_id, actorId=sub, sessionId=session_id, maxResults=1, includePayloads=False, + ) + return bool(response.get("events")) + + +def _refuse(reason: str, session_id: str, event: dict) -> tuple[None, dict]: + logger.warning("Session refused (%s) for session %s", reason, (session_id or "")[:12]) + return None, event + + +async def admit( + payload_session_id: str, header_session_id: str | None, sub: str, +) -> tuple[str | None, dict | None]: + """Return ``(session_id, None)`` when admitted, else ``(None, error_event)``. + + The platform routes by the session header, so that id is the one checked; + without the header (local runs) the payload id is used. + """ + if header_session_id and header_session_id != payload_session_id: + return _refuse("id mismatch", header_session_id, error_event("Session id mismatch", code=BAD_INPUT)) + session_id = header_session_id or payload_session_id + + if is_owner_prefixed(session_id): + if not belongs_to(session_id, sub): + return _refuse("owner", session_id, error_event("Session not available", code=AUTH)) + elif (sub, session_id) not in _owned_legacy: + memory_id = os.environ.get("MEMORY_ID", "") + if not memory_id or memory_id == "PLACEHOLDER": + return _refuse("legacy", session_id, error_event("Reload the page", code=SESSION_RELOAD_REQUIRED)) + try: + owned = await asyncio.to_thread(_has_own_events, memory_id, sub, session_id) + except Exception: + logger.exception("Session check failed for session %s", session_id[:12]) + return None, error_event("Session check unavailable", code=SERVICE_UNAVAILABLE) + if not owned: + return _refuse("legacy", session_id, error_event("Reload the page", code=SESSION_RELOAD_REQUIRED)) + _owned_legacy.add((sub, session_id)) + + if not _process_owner.claim(sub): + return _refuse("process", session_id, error_event("Session not available", code=AUTH)) + return session_id, None diff --git a/docs/en/security.md b/docs/en/security.md index 67d55499..96bd0b26 100644 --- a/docs/en/security.md +++ b/docs/en/security.md @@ -29,7 +29,9 @@ These settings come with every deployment of the CDK stacks in `infra/`: - **AgentCore Runtime**: the agent and MCP runtimes validate a JWT on every call (`customJwtAuthorizer`). The user ID comes from the token's `sub` claim. A resource-based policy on each runtime and its endpoints allows only `InvokeAgentRuntime` and - `StopRuntimeSession`, the operations the app uses. + `StopRuntimeSession`, the operations the app uses. The agent and the MCP server admit a + session only for its owner: the session id carries a prefix derived from the token's + `sub`, and each runtime process serves one user (`shared/session_owner.py`). - **Per-deck authorization**: a user is the deck's owner, a viewer (the deck is public), or has no access. The rules live in one place, `shared/authz.py`, used by both the API and the MCP server. See [Architecture — Authorization](architecture.md#authorization-role-based-access-control). diff --git a/servers/remote/server.py b/servers/remote/server.py index adb3c5c6..86434dee 100644 --- a/servers/remote/server.py +++ b/servers/remote/server.py @@ -33,6 +33,7 @@ from shared.authz import authorize # noqa: E402 from shared import user_library # noqa: E402 +from shared.session_owner import ProcessOwner, belongs_to, is_owner_prefixed, sub_from_jwt # noqa: E402 from storage import write_fence # noqa: E402 from storage.aws import AwsStorage # noqa: E402 from tools import assets, reference, preview, generate # noqa: E402 @@ -279,6 +280,44 @@ def _client_quiet() -> None: _start_soffice_warm_up_once() +# --- Session ownership --- +# +# A session is admitted only for its owner (shared/session_owner.py). An +# owner-prefixed id must carry the prefix of the token's user; ids of other +# formats are admitted under the process binding alone. Every deck operation +# still goes through _check_deck_access. + +_process_owner = ProcessOwner() +_SESSION_HEADERS = ("mcp-session-id", "x-amzn-bedrock-agentcore-runtime-session-id") + + +def _session_refusal(headers: dict) -> str | None: + """Return the refused session id (or ``"-"``) when the request may not proceed, else None.""" + auth = headers.get("authorization") + if auth is None: + return None # platform health checks; every data tool requires the JWT itself + token = auth[7:].strip() if auth.lower().startswith("bearer ") else "" + sub = sub_from_jwt(token) + if not sub: + return "-" + for name in _SESSION_HEADERS: + session_id = headers.get(name) + if session_id and is_owner_prefixed(session_id) and not belongs_to(session_id, sub): + return session_id + if not _process_owner.claim(sub): + return headers.get("mcp-session-id") or "-" + return None + + +async def _send_forbidden(send) -> None: # type: ignore + await send({ + "type": "http.response.start", + "status": 403, + "headers": [(b"content-type", b"application/json")], + }) + await send({"type": "http.response.body", "body": b'{"error": "forbidden"}'}) + + class _CaptureHeadersMiddleware: """Raw ASGI middleware to capture HTTP headers into a ContextVar. @@ -298,6 +337,11 @@ async def __call__(self, scope, receive, send): # type: ignore """Capture headers from HTTP requests into ContextVar.""" if scope["type"] == "http": headers = {k.decode(): v.decode() for k, v in scope.get("headers", [])} + refused = _session_refusal(headers) + if refused is not None: + logger.warning("Session refused for session %s", refused[:12]) + await _send_forbidden(send) + return client = scope.get("method") == "POST" and "mcp-session-id" in headers if client: _client_request_started() diff --git a/shared/__init__.py b/shared/__init__.py index 9a5a9697..acf76481 100644 --- a/shared/__init__.py +++ b/shared/__init__.py @@ -8,6 +8,7 @@ - artifact_keys: key layout of deck artifacts (previews, compose, HTML) and the latest-epoch selector - review: change request contract of the Web UI comment tool - user_library: the user's style / template library (S3 layout, pins, template lists) +- session_owner: runtime session ids bound to their owner (also used by agent/) Upload file conversion (PDF/DOCX/XLSX/PPTX → Markdown/JSON) lives in the core, ``sdpm.tools.attachment.ingest``. diff --git a/shared/session_owner.py b/shared/session_owner.py new file mode 100644 index 00000000..88cc02b0 --- /dev/null +++ b/shared/session_owner.py @@ -0,0 +1,91 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Session ownership — bind runtime sessions to the user who owns them. + +A session id created by the Web UI starts with ``owner_prefix(sub)``, the first 16 +hex characters of SHA-256 over the signed-in user's ``sub``. The agent and the +MCP server admit an owner-prefixed id only for that user, and each runtime +process serves a single user (``ProcessOwner``). + +The JWT has already been validated by the AgentCore Runtime authorizer before +the request reaches the container, so ``sub_from_jwt`` only decodes the payload. + +The Web UI mirror is ``web-ui/src/lib/sessionId.ts``; both test suites pin the +same test vector. +""" + +import base64 +import hashlib +import hmac +import json +import re +import threading + +PREFIX_LEN = 16 + +# A prefixed id: 16 lowercase hex characters, then "-". No id of the earlier +# formats matches: a UUID has its first "-" at index 8, and a padded deck id +# has no "-" at all. +_PREFIXED = re.compile(r"^[0-9a-f]{16}-") + + +def owner_prefix(sub: str) -> str: + """Return the session id prefix of ``sub`` (16 lowercase hex characters).""" + return hashlib.sha256(sub.encode("utf-8")).hexdigest()[:PREFIX_LEN] + + +def is_owner_prefixed(session_id: str) -> bool: + """True when ``session_id`` has the owner-prefixed shape.""" + return bool(_PREFIXED.match(session_id or "")) + + +def belongs_to(session_id: str, sub: str) -> bool: + """True when ``session_id`` is owner-prefixed and its prefix is ``sub``'s. + + Suffixes after the prefix (composer groups ``-gN``) do not matter. + """ + if not sub or not is_owner_prefixed(session_id): + return False + return hmac.compare_digest(session_id[:PREFIX_LEN], owner_prefix(sub)) + + +def owned_mcp_session_id(session_id: str, sub: str) -> str: + """Return the ``Mcp-Session-Id`` the agent sends for ``session_id``. + + An id that already belongs to ``sub`` is sent unchanged; any other id gets + ``sub``'s prefix in front, so the MCP session always carries its owner. + Without a ``sub`` the id is returned unchanged. + """ + if not sub or not session_id or belongs_to(session_id, sub): + return session_id + return f"{owner_prefix(sub)}-{session_id}" + + +def sub_from_jwt(token: str) -> str: + """Return the ``sub`` claim of a JWT (payload decode only), or ``""``.""" + try: + payload = token.split(".")[1] + payload += "=" * (-len(payload) % 4) + claims = json.loads(base64.urlsafe_b64decode(payload)) + sub = claims.get("sub", "") if isinstance(claims, dict) else "" + except (IndexError, ValueError, AttributeError): + return "" + return sub if isinstance(sub, str) else "" + + +class ProcessOwner: + """The one user a runtime process serves: the first ``sub`` that claims it.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._sub: str | None = None + + def claim(self, sub: str) -> bool: + """Bind the process to ``sub`` if unbound; True when ``sub`` is the owner.""" + if not sub: + return False + with self._lock: + if self._sub is None: + self._sub = sub + return True + return hmac.compare_digest(self._sub.encode("utf-8"), sub.encode("utf-8")) diff --git a/tests/agent_stop_server.py b/tests/agent_stop_server.py index c8db45c6..b7be7653 100644 --- a/tests/agent_stop_server.py +++ b/tests/agent_stop_server.py @@ -11,7 +11,7 @@ from pathlib import Path ROOT = Path(__file__).resolve().parents[1] -sys.path[:0] = [str(ROOT / "agent"), str(ROOT / "tests")] +sys.path[:0] = [str(ROOT / "agent"), str(ROOT / "tests"), str(ROOT)] os.environ["MEMORY_ID"] = "memory-test" os.environ.setdefault("AWS_REGION", "us-east-1") diff --git a/tests/test_agent_errors.py b/tests/test_agent_errors.py index cb1e30e7..a5de7db5 100644 --- a/tests/test_agent_errors.py +++ b/tests/test_agent_errors.py @@ -26,6 +26,7 @@ MAX_OUTPUT = _errors.MAX_OUTPUT MODEL_NOT_READY = _errors.MODEL_NOT_READY SERVICE_UNAVAILABLE = _errors.SERVICE_UNAVAILABLE +SESSION_RELOAD_REQUIRED = _errors.SESSION_RELOAD_REQUIRED THROTTLED = _errors.THROTTLED TIMEOUT = _errors.TIMEOUT _PATTERNS = _errors._PATTERNS @@ -66,6 +67,10 @@ def test_explicit_code_wins(self): ev = error_event("No JWT token", code=AUTH) assert ev["code"] == AUTH + def test_session_reload_required_is_sent_with_its_code(self): + ev = error_event("Reload the page", code=SESSION_RELOAD_REQUIRED) + assert ev == {"status": "error", "code": "session_reload_required", "error": "Reload the page"} + class TestTypeScriptParity: """Guard against drift between agent/errors.py and lib/agentErrors.ts.""" @@ -75,6 +80,7 @@ def test_all_python_codes_exist_in_ts(self): py_codes = { CONVERSATION_TOO_LONG, THROTTLED, TIMEOUT, MODEL_NOT_READY, SERVICE_UNAVAILABLE, AUTH, BAD_INPUT, MAX_OUTPUT, INTERNAL, + SESSION_RELOAD_REQUIRED, } for code in py_codes: assert f'"{code}"' in ts, f"code {code!r} missing from agentErrors.ts" @@ -88,6 +94,7 @@ def test_ts_has_no_extra_codes(self): py_codes = { CONVERSATION_TOO_LONG, THROTTLED, TIMEOUT, MODEL_NOT_READY, SERVICE_UNAVAILABLE, AUTH, BAD_INPUT, MAX_OUTPUT, INTERNAL, + SESSION_RELOAD_REQUIRED, } assert ts_codes == py_codes, ( f"code sets differ — TS only: {ts_codes - py_codes}, Python only: {py_codes - ts_codes}" diff --git a/tests/test_agent_session_guard.py b/tests/test_agent_session_guard.py new file mode 100644 index 00000000..5cd29812 --- /dev/null +++ b/tests/test_agent_session_guard.py @@ -0,0 +1,192 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""The agent entry admits an invocation only for the owner of its session. + +Drives basic_agent.agent_stream in process with the agent build and the stream faked, +and AgentCore Memory replaced by a recording fake. +""" + +import asyncio +import base64 +import json + +import pytest + +from shared.session_owner import ProcessOwner +from test_agent_tool_exposure import _basic_agent + +P1 = "c6c289e49e9c05b2" # owner_prefix("user-1") # pragma: allowlist secret +P2 = "d92b69cfb82cecab" # owner_prefix("user-2") # pragma: allowlist secret +OWN = f"{P1}-00000000-0000-4000-8000-000000000000" +OTHER = f"{P2}-00000000-0000-4000-8000-000000000000" +LEGACY = "0f8c6a4e-2b1d-4c3e-9a7f-5d6e7f8a9b0c" + + +class _Context: + def __init__(self, sub: str, session_id: str | None = None): + claims = base64.urlsafe_b64encode(json.dumps({"sub": sub}).encode()).decode().rstrip("=") + self.request_headers = {"authorization": f"Bearer h.{claims}.s"} + self.session_id = session_id + + +class _Memory: + def __init__(self, events=(), error: Exception | None = None): + self.events = list(events) + self.error = error + self.calls: list[dict] = [] + + def list_events(self, **kwargs): + self.calls.append(kwargs) + if self.error: + raise self.error + return {"events": self.events} + + +@pytest.fixture() +def agent(monkeypatch): + basic_agent = _basic_agent() + guard = basic_agent.session_guard + monkeypatch.setattr(guard, "_process_owner", ProcessOwner()) + monkeypatch.setattr(guard, "_owned_legacy", set()) + monkeypatch.setenv("MEMORY_ID", "memory-test") + built: list[dict] = [] + + def fake_create_agent(**kwargs): + built.append(kwargs) + return object(), [] + + async def fake_stream(agent, user_query, session_id, cancel, stop): + yield {"done": True} + + monkeypatch.setattr(basic_agent, "create_agent", fake_create_agent) + monkeypatch.setattr(basic_agent, "stream_agent", fake_stream) + basic_agent.built = built + yield basic_agent + del basic_agent.built + + +def _memory(monkeypatch, agent, **kwargs) -> _Memory: + memory = _Memory(**kwargs) + monkeypatch.setattr(agent.session_guard, "_memory_client", lambda: memory) + return memory + + +def _run(agent, session_id: str, sub: str = "user-1", header: str | None = None) -> list[dict]: + async def collect(): + payload = {"prompt": "hello", "runtimeSessionId": session_id} + return [e async for e in agent.agent_stream(payload, _Context(sub, header))] + + return asyncio.run(collect()) + + +def _code(events): + assert len(events) == 1 and events[0]["status"] == "error", events + return events[0]["code"] + + +def test_own_prefixed_id_reaches_the_agent(agent, monkeypatch): + memory = _memory(monkeypatch, agent) + + events = _run(agent, OWN, header=OWN) + + assert events[-1] == {"done": True} + assert agent.built[0]["session_id"] == OWN + assert memory.calls == [] + + +def test_other_users_prefixed_id_is_refused_before_the_cancel_registry(agent): + running = asyncio.Event() + agent._cancel_events[OTHER] = running + try: + events = _run(agent, OTHER, header=OTHER) + finally: + del agent._cancel_events[OTHER] + + assert events == [{"status": "error", "code": "auth", "error": "Session not available"}] + assert agent.built == [] + assert not running.is_set() + + +def test_header_and_payload_mismatch_is_bad_input(agent): + assert _code(_run(agent, OWN, header=f"{P1}-ffffffff-0000-4000-8000-000000000000")) == "bad_input" + assert agent.built == [] + + +def test_header_id_is_checked_and_used(agent): + assert _code(_run(agent, OTHER, header=OTHER)) == "auth" + + _run(agent, OWN, header=OWN) + + assert agent.built[0]["session_id"] == OWN + + +def test_legacy_id_with_own_events_is_admitted_and_cached(agent, monkeypatch): + memory = _memory(monkeypatch, agent, events=[{"eventId": "e1"}]) + + assert _run(agent, LEGACY)[-1] == {"done": True} + assert _run(agent, LEGACY)[-1] == {"done": True} + + assert memory.calls == [{ + "memoryId": "memory-test", "actorId": "user-1", "sessionId": LEGACY, + "maxResults": 1, "includePayloads": False, + }] + assert [b["session_id"] for b in agent.built] == [LEGACY, LEGACY] + + +def test_legacy_id_without_own_events_asks_for_a_reload(agent, monkeypatch): + memory = _memory(monkeypatch, agent) + + events = _run(agent, LEGACY) + + assert events == [{"status": "error", "code": "session_reload_required", "error": "Reload the page"}] + assert agent.built == [] + # A negative answer is not cached. + _run(agent, LEGACY) + assert len(memory.calls) == 2 + + +@pytest.mark.parametrize("memory_id", ["", "PLACEHOLDER"]) +def test_legacy_id_without_memory_asks_for_a_reload(agent, monkeypatch, memory_id): + memory = _memory(monkeypatch, agent, events=[{"eventId": "e1"}]) + monkeypatch.setenv("MEMORY_ID", memory_id) + + assert _code(_run(agent, LEGACY)) == "session_reload_required" + assert memory.calls == [] + + +def test_memory_error_fails_closed(agent, monkeypatch): + _memory(monkeypatch, agent, error=RuntimeError("boom")) + + events = _run(agent, LEGACY) + + assert _code(events) == "service_unavailable" + assert "boom" not in events[0]["error"] + assert agent.built == [] + + +def test_deck_id_of_an_owner_is_admitted_without_memory(agent, monkeypatch): + memory = _memory(monkeypatch, agent) + deck_sid = P1 + "-" + "a1b2c3d4".ljust(36, "0") + assert len(deck_sid) == 53 + + assert _run(agent, deck_sid)[-1] == {"done": True} + assert memory.calls == [] + + +def test_deck_id_of_an_owner_is_refused_for_another_user(agent): + deck_sid = P1 + "-" + "a1b2c3d4".ljust(36, "0") + + assert _code(_run(agent, deck_sid, sub="user-2")) == "auth" + + +def test_process_serves_only_its_first_user(agent): + assert _run(agent, OWN)[-1] == {"done": True} + + assert _code(_run(agent, OTHER, sub="user-2")) == "auth" + assert len(agent.built) == 1 + + +def test_refused_call_does_not_bind_the_process(agent): + assert _code(_run(agent, OTHER)) == "auth" + + assert _run(agent, OTHER, sub="user-2")[-1] == {"done": True} diff --git a/tests/test_agent_stop_guard.py b/tests/test_agent_stop_guard.py index bf288172..103a599e 100644 --- a/tests/test_agent_stop_guard.py +++ b/tests/test_agent_stop_guard.py @@ -28,6 +28,8 @@ ROOT = Path(__file__).resolve().parents[1] AGENT_DIR = ROOT / "agent" +# A session id of user-1 (owner prefix of "user-1", see shared/session_owner.py). +OWNED = "c6c289e49e9c05b2-00000000-0000-4000-8000-000000000000" from agent_stop_fakes import FakeMCP, Recorder, ScriptedModel # noqa: E402 @@ -84,7 +86,7 @@ def client() -> None: # keeps the connection open, like the Web UI's stream try: with httpx.stream( "POST", f"http://127.0.0.1:{port}/invocations", timeout=30, - json={"prompt": "Make the deck.", "runtimeSessionId": "s" * 33}, + json={"prompt": "Make the deck.", "runtimeSessionId": OWNED}, headers={"Authorization": f"Bearer h.{claims}.s"}, ) as response: for _ in response.iter_lines(): diff --git a/tests/test_agent_tool_exposure.py b/tests/test_agent_tool_exposure.py index 7517498c..c6df04e8 100644 --- a/tests/test_agent_tool_exposure.py +++ b/tests/test_agent_tool_exposure.py @@ -21,6 +21,8 @@ from mcp.shared.memory import create_client_server_memory_streams AGENT_DIR = Path(__file__).resolve().parents[1] / "agent" +# A session id of user-1 (owner prefix of "user-1", see shared/session_owner.py). +OWNED = "c6c289e49e9c05b2-00000000-0000-4000-8000-000000000000" def _factory(): @@ -176,7 +178,7 @@ async def fake_stream(agent, user_query, session_id, cancel, cancel_signal): monkeypatch.setattr(basic_agent, "create_agent", fake_create_agent) monkeypatch.setattr(basic_agent, "stream_agent", fake_stream) - events = _run(basic_agent, {"prompt": "hello", "runtimeSessionId": "s" * 33}) + events = _run(basic_agent, {"prompt": "hello", "runtimeSessionId": OWNED}) assert seen["prompt"] == "hello" assert events[-1] == {"done": True} diff --git a/tests/test_agent_turn_fence.py b/tests/test_agent_turn_fence.py index 615aa718..d84e4129 100644 --- a/tests/test_agent_turn_fence.py +++ b/tests/test_agent_turn_fence.py @@ -10,6 +10,9 @@ import threading +# A session id of user-1 (owner prefix of "user-1", see shared/session_owner.py). +OWNED = "c6c289e49e9c05b2-00000000-0000-4000-8000-000000000000" + # Imported lazily: test_mcp_parallel_connect puts agent/ on sys.path when it is imported, # and agent/tools would then shadow servers/remote's tools package for the test modules @@ -219,7 +222,7 @@ def fake_create_agent(**kwargs): monkeypatch.setattr(basic_agent, "TurnFence", _FenceSpy) monkeypatch.setattr(basic_agent, "create_agent", fake_create_agent) monkeypatch.setattr(basic_agent, "stream_agent", stream) - _run(basic_agent, {"prompt": "hello", "runtimeSessionId": "s" * 33}) + _run(basic_agent, {"prompt": "hello", "runtimeSessionId": OWNED}) (fence,) = _FenceSpy.instances assert seen["fence"] is fence return fence diff --git a/tests/test_mcp_session_affinity.py b/tests/test_mcp_session_affinity.py index bf479dc2..65701a14 100644 --- a/tests/test_mcp_session_affinity.py +++ b/tests/test_mcp_session_affinity.py @@ -122,3 +122,32 @@ def test_terminate_on_close_stays_false(monkeypatch): mcp_clients.mcp_agentcore_runtime(jwt_token="tok", session_id="s-3") assert seen["transport_kwargs"]["terminate_on_close"] is False + + +# --- owner prefix ------------------------------------------------------------- + +P1 = "c6c289e49e9c05b2" # owner_prefix("user-1") # pragma: allowlist secret +OWN = f"{P1}-00000000-0000-4000-8000-000000000000" + + +def _jwt_of(sub: str) -> str: + import base64 + import json + + body = base64.urlsafe_b64encode(json.dumps({"sub": sub}).encode()).decode().rstrip("=") + return f"h.{body}.s" + + +@pytest.mark.parametrize("session_id, sent", [ + ("sess-1234567890", f"{P1}-sess-1234567890"), + (OWN, OWN), + (f"{OWN}-g2", f"{OWN}-g2"), + ("sess-1234567890-g2", f"{P1}-sess-1234567890-g2"), +]) +def test_mcp_session_id_carries_the_owner_prefix(monkeypatch, session_id, sent): + mcp_clients = _require_mcp_clients() + seen = _capture(monkeypatch, mcp_clients) + + mcp_clients.mcp_agentcore_runtime(jwt_token=_jwt_of("user-1"), session_id=session_id) + + assert seen["transport_kwargs"]["headers"]["Mcp-Session-Id"] == sent diff --git a/tests/test_remote_session_gate.py b/tests/test_remote_session_gate.py new file mode 100644 index 00000000..b131ece0 --- /dev/null +++ b/tests/test_remote_session_gate.py @@ -0,0 +1,128 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""The remote MCP server admits a session only for its owner. + +An owner-prefixed session id must carry the prefix of the token's user (also for a +composer group's ``-gN``), and each server process serves one user. Requests go +through the app as server.py runs it (FastMCP + the header middleware). +""" + +import asyncio + +import pytest + +import test_remote_warmup_preload as harness + +server = harness.server # the fixture: servers/remote/server.py imported fresh per test +_bearer = harness._bearer +_post = harness._post + +P1 = "c6c289e49e9c05b2" # owner_prefix("user-1") # pragma: allowlist secret +P2 = "d92b69cfb82cecab" # owner_prefix("user-2") # pragma: allowlist secret +OWN = f"{P1}-00000000-0000-4000-8000-000000000000" +OTHER = f"{P2}-00000000-0000-4000-8000-000000000000" +LIST = {"jsonrpc": "2.0", "id": 1, "method": "tools/list"} +PING = {"jsonrpc": "2.0", "id": 1, "method": "ping"} + + +def _headers(session_id: str | None, sub: str | None = "user-1", **extra) -> dict: + headers = dict(extra) + if sub is not None: + headers["authorization"] = _bearer(sub) + if session_id is not None: + headers["mcp-session-id"] = session_id + return headers + + +def _list(module, session_id: str | None, sub: str | None = "user-1", **extra): + return _post(module, LIST, headers=_headers(session_id, sub, **extra)) + + +def _statuses(module, *requests: tuple[dict, dict]) -> list[int]: + """POST each (body, headers) in turn through one app run (FastMCP's manager runs once).""" + import httpx + + app = module.mcp.streamable_http_app() + app.add_middleware(module._CaptureHeadersMiddleware) + + async def run(): + async with module.mcp.session_manager.run(): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + out = [] + for body, headers in requests: + r = await client.post( + "/mcp", json=body, + headers={"accept": "application/json, text/event-stream", **headers}, + ) + out.append(r.status_code) + return out + + return asyncio.run(run()) + + +@pytest.mark.parametrize("session_id, status", [ + (OWN, 200), + (OTHER, 403), + (f"{OWN}-g3", 200), + (f"{OTHER}-g3", 403), +]) +def test_owner_prefixed_ids(server, session_id, status): + module, *_ = server + + r = _list(module, session_id) + + assert r.status_code == status + if status == 403: + assert r.json() == {"error": "forbidden"} + + +def test_runtime_session_header_of_another_user_is_refused(server): + module, *_ = server + + r = _list(module, OWN, **{"x-amzn-bedrock-agentcore-runtime-session-id": OTHER}) + + assert r.status_code == 403 + + +def test_legacy_id_binds_the_process_to_its_first_user(server): + module, *_ = server + + assert _statuses( + module, + (LIST, _headers("sess-1")), + (LIST, _headers("sess-1", sub="user-2")), + (LIST, _headers("sess-1")), + ) == [200, 403, 200] + + +def test_other_users_own_prefixed_id_is_refused_on_a_bound_process(server): + module, *_ = server + + assert _statuses(module, (LIST, _headers(OWN)), (LIST, _headers(OTHER, sub="user-2"))) == [200, 403] + + +def test_ping_without_authorization_passes_and_does_not_bind(server): + module, *_ = server + + assert _statuses(module, (PING, {}), (LIST, _headers(OTHER, sub="user-2"))) == [200, 200] + + +def test_authorization_without_sub_is_refused(server): + module, *_ = server + + r = _post(module, LIST, headers={"authorization": "Bearer not-a-jwt", "mcp-session-id": "sess-1"}) + + assert r.status_code == 403 + + +def test_refused_request_does_not_start_the_warm_up(server, monkeypatch): + module, _, _, warm_ups = server + monkeypatch.setattr(module, "_WARMUP_QUIET_S", 0.05) + + r = _post(module, LIST, headers={"authorization": _bearer(), "mcp-session-id": OTHER}, linger=0.2) + + assert r.status_code == 403 + assert warm_ups == [] + assert module._client_requests_in_flight == 0 + assert module._warmup_timer is None diff --git a/tests/test_session_owner.py b/tests/test_session_owner.py new file mode 100644 index 00000000..10de94e5 --- /dev/null +++ b/tests/test_session_owner.py @@ -0,0 +1,136 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Session ids carry their owner's prefix, and a runtime process serves one user.""" + +import base64 +import json +import threading + +import pytest + +from shared import session_owner +from shared.session_owner import ( + ProcessOwner, + belongs_to, + is_owner_prefixed, + owned_mcp_session_id, + owner_prefix, + sub_from_jwt, +) + +# web-ui/src/lib/sessionId.test.ts pins the same vectors. +P1 = "c6c289e49e9c05b2" # owner_prefix("user-1") # pragma: allowlist secret +P2 = "d92b69cfb82cecab" # owner_prefix("user-2") # pragma: allowlist secret +UUID = "0f8c6a4e-2b1d-4c3e-9a7f-5d6e7f8a9b0c" +PADDED_DECK = "a1b2c3d4".ljust(36, "0") + + +def _jwt(claims) -> str: + body = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip("=") + return f"h.{body}.s" + + +def test_owner_prefix_vectors(): + assert owner_prefix("user-1") == P1 + assert owner_prefix("user-2") == P2 + + +@pytest.mark.parametrize("sid, expected", [ + (f"{P1}-{UUID}", True), + (f"{P1}-{PADDED_DECK}", True), + (f"{P1}-{UUID}-g12", True), + (UUID, False), + (PADDED_DECK, False), + ("s" * 33, False), + (f"{UUID}-g12", False), + (f"{P1.upper()}-{UUID}", False), + ("", False), +]) +def test_is_owner_prefixed(sid, expected): + assert is_owner_prefixed(sid) is expected + + +def test_prefixed_ids_fit_the_runtime_session_id_bounds(): + for sid in (f"{P1}-{UUID}", f"{P1}-{PADDED_DECK}"): + assert len(sid) == 53 + assert 33 <= len(sid) <= 256 + + +def test_belongs_to(): + assert belongs_to(f"{P1}-{UUID}", "user-1") + assert belongs_to(f"{P1}-{UUID}-g3", "user-1") + assert not belongs_to(f"{P1}-{UUID}", "user-2") + assert not belongs_to(f"{P1}-{UUID}", "") + assert not belongs_to(UUID, "user-1") + + +def test_same_deck_gets_different_ids_per_user(): + assert f"{owner_prefix('user-1')}-{PADDED_DECK}" != f"{owner_prefix('user-2')}-{PADDED_DECK}" + assert not belongs_to(f"{P2}-{PADDED_DECK}", "user-1") + + +def test_belongs_to_compares_in_constant_time(monkeypatch): + calls: list[tuple] = [] + real = session_owner.hmac.compare_digest + + def recorder(a, b): + calls.append((a, b)) + return real(a, b) + + monkeypatch.setattr(session_owner.hmac, "compare_digest", recorder) + + assert belongs_to(f"{P1}-{UUID}", "user-1") + assert not belongs_to(f"{P2}-{UUID}", "user-1") + assert calls == [(P1, P1), (P2, P1)] + + +def test_owned_mcp_session_id(): + assert owned_mcp_session_id("sess-1234567890", "user-1") == f"{P1}-sess-1234567890" + assert owned_mcp_session_id(f"{P1}-{UUID}", "user-1") == f"{P1}-{UUID}" + assert owned_mcp_session_id(f"{P1}-{UUID}-g3", "user-1") == f"{P1}-{UUID}-g3" + assert owned_mcp_session_id(f"{P2}-{UUID}", "user-1") == f"{P1}-{P2}-{UUID}" + assert owned_mcp_session_id("sess-1", "") == "sess-1" + assert owned_mcp_session_id("", "user-1") == "" + + +def test_process_owner_first_claim_wins(): + owner = ProcessOwner() + assert owner.claim("user-1") + assert owner.claim("user-1") + assert not owner.claim("user-2") + + +def test_process_owner_empty_sub_never_claims(): + owner = ProcessOwner() + assert not owner.claim("") + assert owner.claim("user-2") + + +def test_process_owner_concurrent_claims_bind_one_sub(): + owner = ProcessOwner() + barrier = threading.Barrier(8) + results: list[bool] = [] + lock = threading.Lock() + + def claim(i): + barrier.wait() + ok = owner.claim(f"user-{i}") + with lock: + results.append(ok) + + threads = [threading.Thread(target=claim, args=(i,)) for i in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert results.count(True) == 1 + + +def test_sub_from_jwt(): + assert sub_from_jwt(_jwt({"sub": "user-1"})) == "user-1" + assert sub_from_jwt(_jwt({"sub": 42})) == "" + assert sub_from_jwt(_jwt(["sub"])) == "" + assert sub_from_jwt("garbage") == "" + assert sub_from_jwt("a.!!!.b") == "" + assert sub_from_jwt("") == "" diff --git a/web-ui/messages/en.json b/web-ui/messages/en.json index 820dc2e8..a1b436b7 100644 --- a/web-ui/messages/en.json +++ b/web-ui/messages/en.json @@ -272,7 +272,8 @@ "fetchWebImagesDescription": "Include images from websites in presentations", "sessionBusy": "This chat is still answering in another tab or panel. Wait for it to finish, then send again.", "sessionLocked": "This chat is in use by another kiro-cli process, such as a terminal session. Close it there, then send again.", - "errorStopGeneration": "Failed to stop generation" + "errorStopGeneration": "Failed to stop generation", + "errorSessionReloadRequired": "This page is out of date. Reload the page, then send again." }, "snippet": { "title": "Text Snippet", diff --git a/web-ui/messages/ja.json b/web-ui/messages/ja.json index c8b39449..9ee7fe82 100644 --- a/web-ui/messages/ja.json +++ b/web-ui/messages/ja.json @@ -272,7 +272,8 @@ "fetchWebImagesDescription": "Web サイトの画像をプレゼンテーションに取り込みます", "sessionBusy": "このチャットは別のタブまたはパネルで応答中です。応答が終わってから、もう一度送信してください。", "sessionLocked": "このチャットは別の kiro-cli プロセス（ターミナルなど）で使用中です。そちらを終了してから、もう一度送信してください。", - "errorStopGeneration": "生成を停止できませんでした" + "errorStopGeneration": "生成を停止できませんでした", + "errorSessionReloadRequired": "このページは古くなっています。ページを再読み込みしてから、もう一度送信してください。" }, "snippet": { "title": "テキストスニペット", diff --git a/web-ui/package-lock.json b/web-ui/package-lock.json index 9f51bdcb..96d510a6 100644 --- a/web-ui/package-lock.json +++ b/web-ui/package-lock.json @@ -9,6 +9,7 @@ "version": "0.3.0", "license": "MIT-0", "dependencies": { + "@aws-crypto/sha256-js": "^5.2.0", "@dnd-kit/helpers": "0.5.1", "@dnd-kit/react": "0.5.1", "@radix-ui/react-alert-dialog": "^1.2.1", diff --git a/web-ui/package.json b/web-ui/package.json index d1ce242b..61e32f24 100644 --- a/web-ui/package.json +++ b/web-ui/package.json @@ -17,6 +17,7 @@ "clean": "rm -rf build/ node_modules/ .next/" }, "dependencies": { + "@aws-crypto/sha256-js": "^5.2.0", "@dnd-kit/helpers": "0.5.1", "@dnd-kit/react": "0.5.1", "@radix-ui/react-alert-dialog": "^1.2.1", diff --git a/web-ui/src/components/chat/ChatPanel.session.test.tsx b/web-ui/src/components/chat/ChatPanel.session.test.tsx new file mode 100644 index 00000000..f0eb3880 --- /dev/null +++ b/web-ui/src/components/chat/ChatPanel.session.test.tsx @@ -0,0 +1,91 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: MIT-0 +/** + * ChatPanel session ids (Cloud mode): a new chat and a deck without a linked chat get ids + * with the owner prefix of the signed-in user; a linked chat of an earlier format is kept + * while it has history and replaced by the user's own deck id when it has none. + */ + +import { createRef } from "react" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" +import { act, cleanup, waitFor } from "@testing-library/react" +import { renderWithIntl } from "@/test/renderWithIntl" +import { deckSessionId, ownerPrefix } from "@/lib/sessionId" + +const { invokeAgentCore, getChatHistory } = vi.hoisted(() => ({ + invokeAgentCore: vi.fn(), + getChatHistory: vi.fn(), +})) + +vi.mock("@/hooks/useAuth", () => { + const auth = { isAuthenticated: true, user: { id_token: "tok", access_token: "at", profile: { sub: "u1" } } } + return { useAuth: () => auth } +}) +vi.mock("@/hooks/UseMobile", () => ({ useIsMobile: () => false })) +vi.mock("sonner", () => ({ toast: Object.assign(vi.fn(), { dismiss: vi.fn(), info: vi.fn(), error: vi.fn() }) })) +vi.mock("@/lib/runtimeConfig", () => ({ + loadRuntimeConfig: vi.fn(async () => ({ agentRuntimeArn: "arn:test", awsRegion: "us-east-1" })), +})) +vi.mock("@/services/agentCoreService", async (importOriginal) => ({ + ...(await importOriginal()), + setAgentConfig: vi.fn(), + stopRuntimeSession: vi.fn(), + invokeAgentCore, +})) +vi.mock("@/services/deckService", () => ({ + getChatHistory, + patchDeck: vi.fn(), + createReview: vi.fn(), + batchGetSlidePreviewUrls: vi.fn(async () => new Map()), + fetchStyles: vi.fn(async () => []), + fetchTemplates: vi.fn(async () => []), +})) +vi.mock("@/services/kiroSessionsService", () => ({ forkKiroSession: vi.fn(), listKiroSessions: vi.fn(async () => ({ recent: [] })) })) + +import { ChatPanel, type ChatPanelHandle } from "./ChatPanel" + +const LEGACY = "0f8c6a4e-2b1d-4c3e-9a7f-5d6e7f8a9b0c" +const HISTORY = [ + { role: "user", content: [{ text: "hello" }] }, + { role: "assistant", content: [{ text: "hi" }] }, +] + +/** Render, wait for the history load(s), send one message and return the session id it used. */ +async function sentSessionId(props: { deckId: string; chatSessionId?: string }, historyCalls = 1): Promise { + const ref = createRef() + renderWithIntl() + await waitFor(() => expect(getChatHistory).toHaveBeenCalledTimes(historyCalls)) + await act(async () => { await ref.current!.sendMessage("go") }) + await waitFor(() => expect(invokeAgentCore).toHaveBeenCalled()) + return invokeAgentCore.mock.calls[0][1] as string +} + +describe("ChatPanel session ids (Cloud)", () => { + beforeEach(() => { + vi.clearAllMocks() + Element.prototype.scrollIntoView = vi.fn() + getChatHistory.mockResolvedValue({ messages: [], truncated: false }) + invokeAgentCore.mockResolvedValue("") + }) + afterEach(() => cleanup()) + + it("gives a new chat an id with the user's prefix", async () => { + const id = await sentSessionId({ deckId: "new" }) + expect(id.startsWith(`${ownerPrefix("u1")}-`)).toBe(true) + expect(id).toHaveLength(53) + }) + + it("gives a deck without a linked chat the user's deck id", async () => { + expect(await sentSessionId({ deckId: "d1" })).toBe(deckSessionId("u1", "d1")) + }) + + it("keeps a linked chat of an earlier format that has history", async () => { + getChatHistory.mockResolvedValue({ messages: HISTORY, truncated: false }) + expect(await sentSessionId({ deckId: "d1", chatSessionId: LEGACY })).toBe(LEGACY) + }) + + it("replaces a linked chat of an earlier format that has no history", async () => { + expect(await sentSessionId({ deckId: "d1", chatSessionId: LEGACY }, 2)).toBe(deckSessionId("u1", "d1")) + expect(getChatHistory.mock.calls.map((c) => c[0])).toEqual([LEGACY, deckSessionId("u1", "d1")]) + }) +}) diff --git a/web-ui/src/components/chat/ChatPanel.tsx b/web-ui/src/components/chat/ChatPanel.tsx index 39a2582f..906e002f 100644 --- a/web-ui/src/components/chat/ChatPanel.tsx +++ b/web-ui/src/components/chat/ChatPanel.tsx @@ -15,6 +15,7 @@ import { IS_LOCAL } from "@/lib/mode" import { loadRuntimeConfig } from "@/lib/runtimeConfig" import { buildAttachedMarkers } from "@/lib/attachmentMarker" import { generateSessionId, setAgentConfig } from "@/services/agentCoreService" +import { deckSessionId, isOwnerPrefixed } from "@/lib/sessionId" import { getChatHistory, patchDeck } from "@/services/deckService" import * as localService from "@/services/localService" import type { UploadedFile } from "@/services/uploadService" @@ -75,12 +76,18 @@ export const ChatPanel = forwardRef(function Ch const tSessionPicker = useTranslations("sessionPicker") const tReview = useTranslations("review") const locale = useLocale() + const auth = useAuth() // --- Session --- + // Cloud ids carry the owner prefix of the signed-in user (lib/sessionId.ts). + const sub = auth.user?.profile?.sub const [sessionId, setSessionId] = useState(() => { if (chatSessionId) return chatSessionId - if (deckId === "new") return generateSessionId() - return deckId.padEnd(36, "0") + if (deckId === "new") return generateSessionId(sub) + return IS_LOCAL || !sub ? deckId.padEnd(36, "0") : deckSessionId(sub, deckId) }) + // A linked chat id of an earlier format whose history is empty is replaced for + // this page load (see loadHistory); the prop-swap effect must not bring it back. + const replacedLegacyIdRef = useRef(null) const sessionIdRef = useRef(sessionId) sessionIdRef.current = sessionId // Mid-stream guard: when a deck is created while the agent is still @@ -92,7 +99,7 @@ export const ChatPanel = forwardRef(function Ch const pendingChatSessionIdRef = useRef(null) const isLoadingRef = useRef(false) useEffect(() => { - if (chatSessionId && chatSessionId !== sessionId) { + if (chatSessionId && chatSessionId !== sessionId && chatSessionId !== replacedLegacyIdRef.current) { if (isLoadingRef.current) { pendingChatSessionIdRef.current = chatSessionId } else { @@ -146,7 +153,6 @@ export const ChatPanel = forwardRef(function Ch const currentDeckId = useRef(deckId) currentDeckId.current = deckId - const auth = useAuth() const isMobile = useIsMobile() const slashPickerItems = useSlashPickerItems(auth.user?.id_token) // Mirror of the chat input; the HTML pill derives its state from the text alone. @@ -260,7 +266,7 @@ export const ChatPanel = forwardRef(function Ch return } console.error("Failed to continue from kiro session", err) - const freshSessionId = generateSessionId() + const freshSessionId = generateSessionId(sub) sessionIdRef.current = freshSessionId setSessionId(freshSessionId) setContinuedOrigin(undefined) @@ -320,6 +326,10 @@ export const ChatPanel = forwardRef(function Ch } const parsed = parseCloudHistory(history) if (parsed.length > 0) stream.setMessages(parsed) + } else if (!IS_LOCAL && sub && deckId && deckId !== "new" && !isOwnerPrefixed(sessionId)) { + // An earlier-format id with no history: continue under this user's own id. + replacedLegacyIdRef.current = sessionId + setSessionId(deckSessionId(sub, deckId)) } } finally { setHistoryLoading(false) diff --git a/web-ui/src/components/chat/StyleChatPanel.tsx b/web-ui/src/components/chat/StyleChatPanel.tsx index 0fa55c27..b04927ce 100644 --- a/web-ui/src/components/chat/StyleChatPanel.tsx +++ b/web-ui/src/components/chat/StyleChatPanel.tsx @@ -37,7 +37,7 @@ interface StyleChatPanelProps { export function StyleChatPanel({ styleId, onStyleWritten, onStyleSaved }: StyleChatPanelProps) { const t = useTranslations("styleChat") const auth = useAuth() - const sessionId = useRef(generateSessionId()).current + const sessionId = useRef(generateSessionId(auth.user?.profile?.sub)).current const messagesEndRef = useRef(null) const scrollContainerRef = useRef(null) const shouldAutoScroll = useRef(true) diff --git a/web-ui/src/hooks/useChatStream.ts b/web-ui/src/hooks/useChatStream.ts index b31c8ac3..90188c61 100644 --- a/web-ui/src/hooks/useChatStream.ts +++ b/web-ui/src/hooks/useChatStream.ts @@ -289,6 +289,7 @@ export function useChatStream({ sessionId, mode, deckId, onToolEvent, onSendComp mode, deckId, options?.continuedFrom, + (code: string) => code === "session_reload_required" ? t("errorSessionReloadRequired") : undefined, ) } catch (err) { if (err instanceof DOMException && err.name === "AbortError") { diff --git a/web-ui/src/lib/agentErrors.test.ts b/web-ui/src/lib/agentErrors.test.ts index 3d0b26f7..7221260d 100644 --- a/web-ui/src/lib/agentErrors.test.ts +++ b/web-ui/src/lib/agentErrors.test.ts @@ -39,3 +39,20 @@ describe("agentErrorMessage", () => { expect(agentErrorMessage("timed out waiting", "future_code")).toMatch(/took too long/) }) }) + +describe("agentErrorMessage with session_reload_required", () => { + it("maps the code to the English reload message", () => { + expect(agentErrorMessage("Reload the page", "session_reload_required")).toBe( + "This page is out of date. Reload the page, then send again.", + ) + }) + + it("uses the translator when it returns a string", () => { + const translate = (code: string) => (code === "session_reload_required" ? "translated" : undefined) + expect(agentErrorMessage("Reload the page", "session_reload_required", translate)).toBe("translated") + }) + + it("falls back to English when the translator returns undefined", () => { + expect(agentErrorMessage("raw", "throttled", () => undefined)).toMatch(/temporarily busy/) + }) +}) diff --git a/web-ui/src/lib/agentErrors.ts b/web-ui/src/lib/agentErrors.ts index 86116808..33ce4aa9 100644 --- a/web-ui/src/lib/agentErrors.ts +++ b/web-ui/src/lib/agentErrors.ts @@ -17,6 +17,7 @@ export type AgentErrorCode = | "bad_input" | "max_output" | "internal" + | "session_reload_required" const MESSAGES: Record = { conversation_too_long: "This conversation is too long for the model to process. Please start a new chat to continue.", @@ -28,6 +29,7 @@ const MESSAGES: Record = { bad_input: "The request could not be processed. Please adjust your input and try again.", max_output: "The response hit the model's output limit. Ask for a smaller step and try again.", internal: "Sorry, something went wrong. Please try again.", + session_reload_required: "This page is out of date. Reload the page, then send again.", } const PATTERNS: [string[], AgentErrorCode][] = [ @@ -54,11 +56,16 @@ export function classifyAgentError(message: string): AgentErrorCode { * * @param message - Raw error text from the event * @param code - Stable code from the event, when present + * @param translate - Optional translator for a resolved code; English is used when it returns undefined */ -export function agentErrorMessage(message: string, code?: string): string { +export function agentErrorMessage( + message: string, + code?: string, + translate?: (code: AgentErrorCode) => string | undefined, +): string { const resolved = (code && code in MESSAGES ? code : classifyAgentError(message)) as AgentErrorCode // Unclassified internal errors keep the raw message — it's more actionable // than a generic apology when nothing matched. if (resolved === "internal" && !code && message) return message - return MESSAGES[resolved] + return translate?.(resolved) ?? MESSAGES[resolved] } diff --git a/web-ui/src/lib/sessionId.test.ts b/web-ui/src/lib/sessionId.test.ts new file mode 100644 index 00000000..8e61552a --- /dev/null +++ b/web-ui/src/lib/sessionId.test.ts @@ -0,0 +1,63 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: MIT-0 +import { describe, expect, it } from "vitest" +import { deckSessionId, isOwnerPrefixed, newSessionId, ownerPrefix } from "./sessionId" + +// runtimeSessionId accepts 33 to 256 characters. +const MIN = 33 +const MAX = 256 + +describe("ownerPrefix", () => { + it("matches the Python vector (tests/test_session_owner.py)", () => { + expect(ownerPrefix("user-1")).toBe("c6c289e49e9c05b2") // pragma: allowlist secret + expect(ownerPrefix("user-2")).toBe("d92b69cfb82cecab") // pragma: allowlist secret + }) +}) + +describe("newSessionId", () => { + it("is -, 53 characters", () => { + const id = newSessionId("user-1") + expect(id).toMatch(/^[0-9a-f]{16}-[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/) + expect(id.startsWith("c6c289e49e9c05b2-")).toBe(true) + expect(id).toHaveLength(53) + expect(id.length).toBeGreaterThanOrEqual(MIN) + expect(id.length).toBeLessThanOrEqual(MAX) + }) + + it("is new every time", () => { + expect(newSessionId("user-1")).not.toBe(newSessionId("user-1")) + }) +}) + +describe("deckSessionId", () => { + it("is -, 53 characters, the same on every call", () => { + const id = deckSessionId("user-1", "a1b2c3d4") + expect(id).toBe("c6c289e49e9c05b2-" + "a1b2c3d4".padEnd(36, "0")) + expect(id).toHaveLength(53) + expect(id.length).toBeGreaterThanOrEqual(MIN) + expect(id.length).toBeLessThanOrEqual(MAX) + expect(deckSessionId("user-1", "a1b2c3d4")).toBe(id) + }) + + it("differs per user for the same deck, each with its own prefix", () => { + const a = deckSessionId("u1", "d1") + const b = deckSessionId("u2", "d1") + expect(a).not.toBe(b) + expect(a.slice(0, 16)).toBe(ownerPrefix("u1")) + expect(b.slice(0, 16)).toBe(ownerPrefix("u2")) + }) +}) + +describe("isOwnerPrefixed", () => { + it.each([ + [newSessionId("user-1"), true], + [deckSessionId("user-1", "d1"), true], + [`${newSessionId("user-1")}-g2`, true], + [crypto.randomUUID(), false], + ["a1b2c3d4".padEnd(36, "0"), false], + ["s".repeat(36), false], + ["C6C289E49E9C05B2-x", false], + ])("%s → %s", (id, expected) => { + expect(isOwnerPrefixed(id)).toBe(expected) + }) +}) diff --git a/web-ui/src/lib/sessionId.ts b/web-ui/src/lib/sessionId.ts new file mode 100644 index 00000000..2c01b7e7 --- /dev/null +++ b/web-ui/src/lib/sessionId.ts @@ -0,0 +1,47 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: MIT-0 +/** + * Runtime session ids bound to their owner — mirrors shared/session_owner.py. + * + * A Cloud session id starts with the owner prefix of the signed-in user: the first + * 16 hex characters of SHA-256 over the token's `sub`. The agent and the MCP server + * admit an owner-prefixed id only for that user. Both test suites pin the same vector. + * + * SHA-256 runs synchronously (`@aws-crypto/sha256-js`) because ids are created in + * `useState` / `useRef` initialisers; Web Crypto is async only. + */ + +import { Sha256 } from "@aws-crypto/sha256-js" + +const PREFIXED = /^[0-9a-f]{16}-/ + +/** Session id prefix of `sub` (16 lowercase hex characters). */ +export function ownerPrefix(sub: string): string { + const hash = new Sha256() + hash.update(new TextEncoder().encode(sub)) + return Array.from(hash.digestSync(), (b) => b.toString(16).padStart(2, "0")).join("").slice(0, 16) +} + +/** A new session id for `sub`: `-` (53 characters). */ +export function newSessionId(sub: string): string { + return `${ownerPrefix(sub)}-${crypto.randomUUID()}` +} + +/** + * The session id of a deck that has no linked chat yet: `-` + * (53 characters, the same on every page load for this user and deck). + * + * Ids never collide across users: the prefix comes from the caller's own `sub`, so + * another user's id for the same deck has a different prefix, and the agent and the + * MCP server refuse an id whose prefix is not the caller's. Two users share a 16-hex + * prefix with probability 2^-64 per pair, and each runtime process still serves only + * one user. + */ +export function deckSessionId(sub: string, deckId: string): string { + return `${ownerPrefix(sub)}-${deckId.padEnd(36, "0")}` +} + +/** True when `id` has the owner-prefixed shape (ids of earlier formats do not). */ +export function isOwnerPrefixed(id: string): boolean { + return PREFIXED.test(id) +} diff --git a/web-ui/src/services/agentCoreService.js b/web-ui/src/services/agentCoreService.js index 35c39319..17b28019 100644 --- a/web-ui/src/services/agentCoreService.js +++ b/web-ui/src/services/agentCoreService.js @@ -9,6 +9,7 @@ import * as parser from './strandsParser.js'; import { getAllowedModels } from "@/lib/allowedModels"; import { IS_LOCAL } from "@/lib/mode"; +import { newSessionId } from "@/lib/sessionId"; // Read the user's selected chat model from localStorage, validated against the allowed list. function readSelectedChatModelId() { @@ -56,7 +57,7 @@ export const setAgentConfig = (runtimeArn, region = "us-east-1") => { /** * Invokes the AgentCore runtime with streaming support */ -export const invokeAgentCore = async (query, sessionId, onStreamUpdate, accessToken, userId, onToolUse, signal, mode, deckId, continuedFrom) => { +export const invokeAgentCore = async (query, sessionId, onStreamUpdate, accessToken, userId, onToolUse, signal, mode, deckId, continuedFrom, translateError) => { // Local mode: proxy through Next.js API Route → kiro-cli acp if (IS_LOCAL) { return invokeLocalAgent(query, sessionId, onStreamUpdate, onToolUse, signal, mode, deckId, continuedFrom); @@ -144,7 +145,7 @@ export const invokeAgentCore = async (query, sessionId, onStreamUpdate, accessTo for (const line of lines) { if (line.trim()) { - completion = parser.parseStreamingChunk(line, completion, onStreamUpdate, onToolUse); + completion = parser.parseStreamingChunk(line, completion, onStreamUpdate, onToolUse, translateError); } } } @@ -190,10 +191,10 @@ export const stopRuntimeSession = async (sessionId, accessToken) => { } /** - * Generate a new session ID + * Generate a new session ID. Cloud ids carry the owner prefix of `sub` (lib/sessionId.ts). */ -export const generateSessionId = () => { - return generateId() +export const generateSessionId = (sub) => { + return IS_LOCAL || !sub ? generateId() : newSessionId(sub) } /** diff --git a/web-ui/src/services/agentCoreService.test.ts b/web-ui/src/services/agentCoreService.test.ts new file mode 100644 index 00000000..291b22ee --- /dev/null +++ b/web-ui/src/services/agentCoreService.test.ts @@ -0,0 +1,27 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: MIT-0 +import { afterEach, describe, expect, it, vi } from "vitest" + +const mode = vi.hoisted(() => ({ local: false })) +vi.mock("@/lib/mode", () => ({ + get IS_LOCAL() { return mode.local }, +})) + +import { generateSessionId } from "./agentCoreService" + +const UUID = /^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/ + +describe("generateSessionId", () => { + afterEach(() => { mode.local = false }) + + it("carries the owner prefix in Cloud mode", () => { + const id = generateSessionId("user-1") + expect(id.startsWith("c6c289e49e9c05b2-")).toBe(true) + expect(id.slice(17)).toMatch(UUID) + }) + + it("stays a plain UUID in Local mode", () => { + mode.local = true + expect(generateSessionId("user-1")).toMatch(UUID) + }) +}) diff --git a/web-ui/src/services/strandsParser.js b/web-ui/src/services/strandsParser.js index d24f0920..1d2d32db 100644 --- a/web-ui/src/services/strandsParser.js +++ b/web-ui/src/services/strandsParser.js @@ -37,9 +37,10 @@ export const resetParserState = () => { * @param {string} currentCompletion - The accumulated completion text * @param {Function} updateCallback - Callback to update the UI with new text * @param {Function} [toolCallback] - Callback for tool use events: (toolName, toolUseData) + * @param {Function} [translateError] - Optional translator for error codes: (code) => string | undefined * @returns {string} Updated completion text */ -export const parseStreamingChunk = (line, currentCompletion, updateCallback, toolCallback) => { +export const parseStreamingChunk = (line, currentCompletion, updateCallback, toolCallback, translateError) => { if (!line || !line.trim()) return currentCompletion; if (!line.startsWith('data: ')) return currentCompletion; @@ -54,7 +55,7 @@ export const parseStreamingChunk = (line, currentCompletion, updateCallback, too // Agent-level error (e.g. Bedrock ValidationException, image/token limit exceeded) if (json.status === 'error' && json.error) { - const newCompletion = currentCompletion + `⚠️ ${agentErrorMessage(json.error, json.code)}`; + const newCompletion = currentCompletion + `⚠️ ${agentErrorMessage(json.error, json.code, translateError)}`; updateCallback(newCompletion); return newCompletion; } diff --git a/web-ui/src/services/strandsParser.test.ts b/web-ui/src/services/strandsParser.test.ts new file mode 100644 index 00000000..2aaad9bb --- /dev/null +++ b/web-ui/src/services/strandsParser.test.ts @@ -0,0 +1,21 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: MIT-0 +import { describe, it, expect, vi } from "vitest" +import { parseStreamingChunk } from "./strandsParser.js" + +describe("parseStreamingChunk error events", () => { + const line = `data: ${JSON.stringify({ status: "error", code: "session_reload_required", error: "Reload the page" })}` + + it("appends the translated message when a translator is given", () => { + const update = vi.fn() + const translate = (code: string) => (code === "session_reload_required" ? "translated" : undefined) + const out = parseStreamingChunk(line, "", update, undefined, translate) + expect(out).toBe("⚠️ translated") + expect(update).toHaveBeenCalledWith("⚠️ translated") + }) + + it("appends the English message without a translator", () => { + const out = parseStreamingChunk(line, "", vi.fn()) + expect(out).toBe("⚠️ This page is out of date. Reload the page, then send again.") + }) +})