diff --git a/infrastructure/template.yaml b/infrastructure/template.yaml index fa3a191..6dd2b1d 100644 --- a/infrastructure/template.yaml +++ b/infrastructure/template.yaml @@ -126,6 +126,31 @@ Resources: - Key: Project Value: nat + # SessionStateTable - server-authoritative confirmation + undo state + # Holds per-session destructive-action confirmations and undo history so that + # neither can be forged via the request body. Rows expire via the TTL. + SessionStateTable: + Type: AWS::DynamoDB::Table + DeletionPolicy: Retain + UpdateReplacePolicy: Retain + Properties: + TableName: !Sub 'nat-session-state-${Environment}' + BillingMode: PAY_PER_REQUEST + AttributeDefinitions: + - AttributeName: session_id + AttributeType: S + KeySchema: + - AttributeName: session_id + KeyType: HASH + TimeToLiveSpecification: + AttributeName: expires_at + Enabled: true + Tags: + - Key: Environment + Value: !Ref Environment + - Key: Project + Value: nat + # OAuthStateTable - single-use CSRF nonces for the NationBuilder OAuth flow. # Each record is written at OAuth init and deleted on callback (single-use); # the TTL on expires_at evicts any that are never consumed. @@ -188,6 +213,8 @@ Resources: - !Sub '${TenantsTable.Arn}/index/*' - !GetAtt UsersTable.Arn - !Sub '${UsersTable.Arn}/index/*' + - !GetAtt SessionStateTable.Arn + - !Sub '${SessionStateTable.Arn}/index/*' - !GetAtt OAuthStateTable.Arn - PolicyName: SecretsManagerAccess PolicyDocument: @@ -874,6 +901,7 @@ Resources: NATIONS_TABLE: !Ref NationsTable USERS_TABLE: !Ref UsersTable TENANTS_TABLE: !Ref TenantsTable + SESSION_STATE_TABLE: !Ref SessionStateTable ANTHROPIC_API_KEY_SECRET: !Sub 'nat/anthropic-api-key-${Environment}' SESSION_JWT_SECRET_NAME: !Sub 'nat/session-jwt-secret-${Environment}' CLAUDE_MODEL: 'claude-haiku-4-5-20251001' @@ -987,6 +1015,18 @@ Outputs: Export: Name: !Sub '${AWS::StackName}-UsersTableArn' + SessionStateTableName: + Description: Name of the Session State DynamoDB table + Value: !Ref SessionStateTable + Export: + Name: !Sub '${AWS::StackName}-SessionStateTableName' + + SessionStateTableArn: + Description: ARN of the Session State DynamoDB table + Value: !GetAtt SessionStateTable.Arn + Export: + Name: !Sub '${AWS::StackName}-SessionStateTableArn' + LambdaExecutionRoleArn: Description: ARN of the Lambda execution role Value: !GetAtt LambdaExecutionRole.Arn diff --git a/src/lambdas/nat_agent_streaming/handler.py b/src/lambdas/nat_agent_streaming/handler.py index 1cb226a..19e5ad1 100644 --- a/src/lambdas/nat_agent_streaming/handler.py +++ b/src/lambdas/nat_agent_streaming/handler.py @@ -33,12 +33,30 @@ SessionTokenError, authenticate_request, ) + from src.lambdas.shared.session_state import ( + append_undo_entry, + compute_tool_id, + consume_confirmation, + filter_authorized_confirmations, + get_undo_stack, + make_session_id, + record_pending_confirmation, + ) except ModuleNotFoundError: # pragma: no cover - exercised only in Lambda from shared.session_token import ( # type: ignore[no-redef] SessionContext, SessionTokenError, authenticate_request, ) + from shared.session_state import ( # type: ignore[no-redef] + append_undo_entry, + compute_tool_id, + consume_confirmation, + filter_authorized_confirmations, + get_undo_stack, + make_session_id, + record_pending_confirmation, + ) logger = logging.getLogger() logger.setLevel(logging.INFO) @@ -306,6 +324,43 @@ def _get_undo_instruction( return "This action cannot be undone" +def _build_undo_entry( + tool_name: str, tool_input: dict[str, Any] +) -> dict[str, Any] | None: + """ + Build a server-side undo entry from an executed tool call, using only data + the server can observe (the tool name and input). + + Actions whose reversal depends on the tool *result* (e.g. the id of a newly + created record, or a tagging id) cannot be reconstructed server-side without + result visibility, so they are not recorded as undoable. Failing to offer an + undo is safe; offering a forged one is not. + + Returns ``None`` when the action is not server-reconstructable. + """ + if tool_name == "add_to_list": + person_id = tool_input.get("person_id") + list_id = tool_input.get("list_id") + if person_id and list_id: + return { + "description": f"Add person {person_id} to list {list_id}", + "toolName": tool_name, + "undoType": "remove_from_list", + "undoData": {"person_id": person_id, "list_id": list_id}, + } + elif tool_name == "remove_from_list": + person_id = tool_input.get("person_id") + list_id = tool_input.get("list_id") + if person_id and list_id: + return { + "description": f"Remove person {person_id} from list {list_id}", + "toolName": tool_name, + "undoType": "add_to_list", + "undoData": {"person_id": person_id, "list_id": list_id}, + } + return None + + async def stream_agent_response( query: str, nb_slug: str, @@ -313,7 +368,7 @@ async def stream_agent_response( model: str, context: dict[str, Any] | None = None, confirmed_tools: set[str] | None = None, - undo_stack: list[dict[str, Any]] | None = None + session_id: str | None = None ) -> AsyncGenerator[str, None]: """ Stream responses from the Nat agent as SSE events. @@ -324,8 +379,11 @@ async def stream_agent_response( nb_token: NationBuilder API token model: Claude model to use context: Optional page context from extension - confirmed_tools: Set of tool_id values that have been confirmed by user - undo_stack: List of recent undoable actions from this session + confirmed_tools: Set of tool_id values authorised server-side for this + turn (already validated against the server's pending-confirmation + record; never trusted directly from the client) + session_id: Opaque server-derived session id used to record/consume + confirmations and to load/append the authoritative undo history Yields: SSE formatted strings @@ -362,8 +420,12 @@ async def stream_agent_response( if context_parts: context_sections.append(f"[Context: {', '.join(context_parts)}]") - # Add undo stack context if the user might be asking to undo - # Check for common undo phrases + # Add undo stack context if the user might be asking to undo. + # The undo history is loaded from the authoritative server-side store, NOT + # from the request body: a client-supplied stack is forgeable and could + # steer an "undo" into deleting an arbitrary record. + undo_stack = get_undo_stack(session_id) if session_id else [] + query_lower = query.lower() undo_phrases = ["undo", "revert", "reverse", "take back", "cancel that", "nevermind"] is_undo_request = any(phrase in query_lower for phrase in undo_phrases) @@ -410,10 +472,16 @@ async def stream_agent_response( "text": block.text }) elif isinstance(block, ToolUseBlock): - # Check if this is a destructive tool requiring confirmation - tool_id = f"{block.name}_{hash(json.dumps(block.input, sort_keys=True))}" + # Stable id for this (tool, input) pair so the prompt + # and the follow-up confirmation request agree on it. + tool_id = compute_tool_id(block.name, block.input) if block.name in DESTRUCTIVE_TOOLS and tool_id not in confirmed: + # Record server-side that we legitimately prompted + # for this action; only a recorded tool_id can be + # honoured as confirmed on resubmission. + if session_id: + record_pending_confirmation(session_id, tool_id) # Send confirmation_required event and pause yield format_sse_event(SSE_EVENT_CONFIRMATION_REQUIRED, { "tool_id": tool_id, @@ -424,6 +492,22 @@ async def stream_agent_response( # Return early - client must re-submit with confirmation return + # A confirmed destructive action is single-use: consume + # the record so it cannot be replayed for another turn. + if ( + session_id + and block.name in DESTRUCTIVE_TOOLS + and tool_id in confirmed + ): + consume_confirmation(session_id, tool_id) + + # Record any server-reconstructable undoable action to + # the authoritative server-side undo stack. + if session_id: + undo_entry = _build_undo_entry(block.name, block.input) + if undo_entry is not None: + append_undo_entry(session_id, undo_entry) + # Notify client about tool invocation tool_info = { "name": block.name, @@ -439,10 +523,13 @@ async def stream_agent_response( "is_error": message.is_error, }) - # Send final done event with complete response + # Send final done event with complete response. The undo stack returned + # here is the authoritative server-side history (the client renders it + # but never supplies it back). yield format_sse_event(SSE_EVENT_DONE, { "response": "".join(full_response), "tool_calls": tool_calls, + "undo_stack": get_undo_stack(session_id) if session_id else [], }) except Exception as e: @@ -491,9 +578,10 @@ async def process_streaming_request(body: dict[str, Any]) -> AsyncGenerator[str, user_id = body.get("user_id") nation_slug = body.get("nation_slug") # NEW: Nation identifier page_context = body.get("context", {}) + # Client-supplied confirmations are *claims*, not authority: each must be + # backed by a server-side record that we actually prompted for it. The + # client's undo_stack is ignored entirely (maintained server-side instead). confirmed_tools_list: list[str] = body.get("confirmed_tools", []) - confirmed_tools = set(confirmed_tools_list) if confirmed_tools_list else set() - undo_stack: list[dict[str, Any]] = body.get("undo_stack", []) if not query: yield format_sse_event(SSE_EVENT_ERROR, { @@ -560,6 +648,15 @@ async def process_streaming_request(body: dict[str, Any]) -> AsyncGenerator[str, nb_token, verified_slug = tokens + # Server-authoritative confirmation gate: keep only the tool_ids the server + # itself previously emitted a confirmation_required for, in this session. A + # forged confirmed_tools entry is dropped here and will trigger a real + # confirmation round-trip rather than executing a destructive tool. + session_id = make_session_id(user_id, nation_slug) + confirmed_tools = filter_authorized_confirmations( + session_id, confirmed_tools_list + ) + # Stream the agent response and track whether it succeeded stream_succeeded = False async for event in stream_agent_response( @@ -569,7 +666,7 @@ async def process_streaming_request(body: dict[str, Any]) -> AsyncGenerator[str, model=CLAUDE_MODEL, context=page_context, confirmed_tools=confirmed_tools, - undo_stack=undo_stack, + session_id=session_id, ): yield event # Check if this is a successful done event (no error in response) diff --git a/src/lambdas/shared/session_state.py b/src/lambdas/shared/session_state.py new file mode 100644 index 0000000..0ad75ca --- /dev/null +++ b/src/lambdas/shared/session_state.py @@ -0,0 +1,215 @@ +""" +Server-side session state: destructive-action confirmations and undo history. + +Confirmation state and undo history are *authoritative server-side*. Both were +previously trusted from the streaming request body (``confirmed_tools`` and +``undo_stack``), which are forgeable: a malicious client could pre-confirm a +destructive tool it was never prompted for (bypassing the confirmation gate), +or spoof undo instructions so an "undo" deletes the wrong record. + +This module persists that state in DynamoDB keyed by an opaque ``session_id`` +derived from the verified ``(user_id, nation_slug)`` identity claims (never from +client-supplied fields). State expires automatically via a DynamoDB TTL. + +Security posture +---------------- +Confirmation checks **fail closed**: if the store is unavailable, an action is +treated as *not* confirmed, so the worst case is an extra confirmation +round-trip — never a silent destructive execution. Undo helpers **fail open**: +undo is a convenience, so a store error simply yields no undo context. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +import time +from typing import Any + +import boto3 +from botocore.exceptions import ClientError + +logger = logging.getLogger() +logger.setLevel(logging.INFO) + +# DynamoDB table holding per-session confirmation + undo state. +SESSION_STATE_TABLE = os.environ.get("SESSION_STATE_TABLE", "nat-session-state-dev") + +# How long session state lives before the DynamoDB TTL reaps it (seconds). +SESSION_STATE_TTL_SECONDS = int( + os.environ.get("SESSION_STATE_TTL_SECONDS", "3600") +) + +# Cap the server-side undo stack so the item stays small. +MAX_UNDO_ENTRIES = 20 + + +def get_dynamodb_resource() -> Any: + """Get DynamoDB resource (allows mocking in tests).""" + return boto3.resource("dynamodb") + + +def get_current_timestamp() -> int: + """Get current Unix timestamp (allows mocking in tests).""" + return int(time.time()) + + +def make_session_id(user_id: str, nation_slug: str) -> str: + """ + Build the opaque session id for confirmation/undo state. + + Derived solely from the verified identity claims so it cannot be forged or + aimed at another user's session. The nation slug is length-prefixed so the + encoding is unambiguous even if a field were ever to contain the delimiter + (no two distinct (user, nation) pairs can collide on the same session id). + """ + return f"v1:{len(nation_slug)}:{nation_slug}#{user_id}" + + +def compute_tool_id(tool_name: str, tool_input: dict[str, Any]) -> str: + """ + Compute a stable, deterministic id for a (tool, input) pair. + + Uses SHA-256 over the canonicalised input rather than the builtin ``hash()``, + which is salted per process (``PYTHONHASHSEED``) and would therefore differ + between the request that prompts for confirmation and the follow-up request + that confirms it — breaking the round-trip and the server-side record match. + """ + canonical = json.dumps(tool_input, sort_keys=True, ensure_ascii=False) + digest = hashlib.sha256(f"{tool_name}:{canonical}".encode("utf-8")).hexdigest() + return f"{tool_name}_{digest[:16]}" + + +def _table() -> Any: + return get_dynamodb_resource().Table(SESSION_STATE_TABLE) + + +def record_pending_confirmation(session_id: str, tool_id: str) -> None: + """ + Record that the server legitimately prompted ``session_id`` to confirm + ``tool_id``. + + Only a tool_id recorded here may later be honoured as confirmed. Best-effort: + on failure the confirmation simply will not be authorised on resubmission + (the user is re-prompted), so the gate fails closed. + """ + expires_at = get_current_timestamp() + SESSION_STATE_TTL_SECONDS + try: + _table().update_item( + Key={"session_id": session_id}, + UpdateExpression="ADD pending_tool_ids :tid SET expires_at = :ttl", + ExpressionAttributeValues={ + ":tid": {tool_id}, + ":ttl": expires_at, + }, + ) + except ClientError as e: + logger.error( + f"Failed to record pending confirmation for session {session_id}: {e}" + ) + + +def filter_authorized_confirmations( + session_id: str, tool_ids: list[str] | set[str] +) -> set[str]: + """ + Return the subset of client-supplied ``tool_ids`` that the server actually + prompted for in this session. + + This is the authoritative check that closes the bypass: a forged tool_id the + server never issued a ``confirmation_required`` for is dropped here, so it + cannot execute a destructive tool without a real confirmation round-trip. + Fails closed (returns empty) on any error. + """ + requested = {t for t in tool_ids if t} + if not requested: + return set() + + try: + response = _table().get_item( + Key={"session_id": session_id}, + ProjectionExpression="pending_tool_ids", + ) + except ClientError as e: + logger.error( + f"Failed to read pending confirmations for session {session_id}: {e}" + ) + return set() + + item = response.get("Item") or {} + pending = item.get("pending_tool_ids") or set() + # DynamoDB returns string sets as Python sets; be defensive about lists too. + pending_set = set(pending) + return requested & pending_set + + +def consume_confirmation(session_id: str, tool_id: str) -> None: + """ + Remove a confirmation once it has been honoured so it cannot be replayed for + a different agent turn. Best-effort. + """ + try: + _table().update_item( + Key={"session_id": session_id}, + UpdateExpression="DELETE pending_tool_ids :tid", + ExpressionAttributeValues={":tid": {tool_id}}, + ) + except ClientError as e: + logger.error( + f"Failed to consume confirmation {tool_id} for session {session_id}: {e}" + ) + + +def append_undo_entry(session_id: str, entry: dict[str, Any]) -> None: + """ + Append a server-observed undoable action to the session's undo stack. + + The stack is stored as a JSON string (preserving exact id types) and capped + to the most recent ``MAX_UNDO_ENTRIES``. Best-effort / fails open. + """ + try: + stack = get_undo_stack(session_id) + stack.append(entry) + stack = stack[-MAX_UNDO_ENTRIES:] + expires_at = get_current_timestamp() + SESSION_STATE_TTL_SECONDS + _table().update_item( + Key={"session_id": session_id}, + UpdateExpression="SET undo_stack_json = :stack, expires_at = :ttl", + ExpressionAttributeValues={ + ":stack": json.dumps(stack, ensure_ascii=False), + ":ttl": expires_at, + }, + ) + except (ClientError, ValueError, TypeError) as e: + logger.error(f"Failed to append undo entry for session {session_id}: {e}") + + +def get_undo_stack(session_id: str) -> list[dict[str, Any]]: + """ + Return the server-maintained undo stack for a session (never the client's). + + Fails open: returns an empty list on any error so undo simply offers no + history rather than blocking the request. + """ + try: + response = _table().get_item( + Key={"session_id": session_id}, + ProjectionExpression="undo_stack_json", + ) + except ClientError as e: + logger.error(f"Failed to read undo stack for session {session_id}: {e}") + return [] + + item = response.get("Item") or {} + raw = item.get("undo_stack_json") + if not raw: + return [] + try: + parsed = json.loads(raw) + except (ValueError, TypeError): + return [] + if not isinstance(parsed, list): + return [] + return [e for e in parsed if isinstance(e, dict)] diff --git a/tests/test_nat_agent_streaming.py b/tests/test_nat_agent_streaming.py index 07d9948..77acf3b 100644 --- a/tests/test_nat_agent_streaming.py +++ b/tests/test_nat_agent_streaming.py @@ -19,13 +19,18 @@ handler, authenticated_body, process_streaming_request, + stream_agent_response, + _build_undo_entry, _get_undo_instruction, SSE_EVENT_TEXT, SSE_EVENT_TOOL_USE, SSE_EVENT_TOOL_RESULT, SSE_EVENT_ERROR, SSE_EVENT_DONE, + SSE_EVENT_CONFIRMATION_REQUIRED, ) +from src.lambdas.shared import session_state +from src.lambdas.shared.session_state import compute_tool_id, make_session_id from src.lambdas.shared.session_token import mint_session_token @@ -804,3 +809,361 @@ def test_authenticated_body_overrides_client_identity(self) -> None: assert result["nation_slug"] == "real-nation" # The original body dict is not mutated. assert body["nation_slug"] == "victim-nation" + + +# ============================================================================= +# Server-side confirmation / undo (issue #12) +# ============================================================================= + + +class _SessionStateFakeTable: + """In-memory DynamoDB Table stand-in for the session-state store.""" + + def __init__(self) -> None: + self.items: dict[str, dict[str, Any]] = {} + + def update_item( + self, + Key: dict[str, Any], + UpdateExpression: str, + ExpressionAttributeValues: dict[str, Any], + ) -> None: + item = self.items.setdefault(Key["session_id"], {}) + if "ADD pending_tool_ids :tid" in UpdateExpression: + cur = set(item.get("pending_tool_ids", set())) + cur |= set(ExpressionAttributeValues[":tid"]) + item["pending_tool_ids"] = cur + if "DELETE pending_tool_ids :tid" in UpdateExpression: + cur = set(item.get("pending_tool_ids", set())) + cur -= set(ExpressionAttributeValues[":tid"]) + if cur: + item["pending_tool_ids"] = cur + else: + item.pop("pending_tool_ids", None) + if "SET undo_stack_json = :stack" in UpdateExpression: + item["undo_stack_json"] = ExpressionAttributeValues[":stack"] + if ":ttl" in ExpressionAttributeValues: + item["expires_at"] = ExpressionAttributeValues[":ttl"] + + def get_item( + self, Key: dict[str, Any], ProjectionExpression: str | None = None + ) -> dict[str, Any]: + item = self.items.get(Key["session_id"]) + return {"Item": dict(item)} if item is not None else {} + + +class _FakeTextBlock: + def __init__(self, text: str) -> None: + self.text = text + + +class _FakeToolUseBlock: + def __init__(self, name: str, tool_input: dict[str, Any]) -> None: + self.name = name + self.input = tool_input + + +class _FakeAssistantMessage: + def __init__(self, content: list[Any]) -> None: + self.content = content + + +class _FakeResultMessage: + def __init__(self, result: str, is_error: bool = False) -> None: + self.result = result + self.is_error = is_error + + +def _make_fake_client(messages: list[Any]) -> Any: + """Build a ClaudeSDKClient replacement yielding scripted messages.""" + + class _FakeClient: + def __init__(self, *args: Any, **kwargs: Any) -> None: + pass + + async def __aenter__(self) -> "_FakeClient": + return self + + async def __aexit__(self, *exc: Any) -> bool: + return False + + async def query(self, prompt: str) -> None: + self.prompt = prompt + + async def receive_response(self) -> Any: + for message in messages: + yield message + + return _FakeClient + + +def _patch_streaming_sdk(messages: list[Any], table: Any) -> Any: + """Patch the agent SDK symbols, options builder, and session-state store.""" + resource = MagicMock() + resource.Table.return_value = table + patches = [ + patch("claude_agent_sdk.ClaudeSDKClient", _make_fake_client(messages)), + patch("claude_agent_sdk.AssistantMessage", _FakeAssistantMessage), + patch("claude_agent_sdk.TextBlock", _FakeTextBlock), + patch("claude_agent_sdk.ToolUseBlock", _FakeToolUseBlock), + patch("claude_agent_sdk.ResultMessage", _FakeResultMessage), + patch("nat.agent.create_nat_options", return_value=MagicMock()), + patch.object( + session_state, "get_dynamodb_resource", return_value=resource + ), + ] + + class _Combined: + def __enter__(self) -> None: + for p in patches: + p.start() + + def __exit__(self, *exc: Any) -> bool: + for p in patches: + p.stop() + return False + + return _Combined() + + +def _collect_stream(**kwargs: Any) -> list[str]: + async def _run() -> list[str]: + events = [] + async for event in stream_agent_response(**kwargs): + events.append(event) + return events + + return run_async(_run()) + + +class TestServerSideConfirmation: + """The confirmation gate must be authoritative server-side (issue #12).""" + + DELETE_INPUT = {"id": "contact-5"} + SESSION = "testnation#user-test-12345" + + def _delete_message(self) -> _FakeAssistantMessage: + return _FakeAssistantMessage( + [_FakeToolUseBlock("delete_contact", self.DELETE_INPUT)] + ) + + def test_unconfirmed_destructive_tool_prompts_and_records(self) -> None: + """First touch of a destructive tool prompts AND records server-side.""" + table = _SessionStateFakeTable() + with _patch_streaming_sdk([self._delete_message()], table): + events = _collect_stream( + query="delete contact 5", + nb_slug=TEST_NB_SLUG, + nb_token=TEST_NB_TOKEN, + model="claude-haiku-4-5-20251001", + confirmed_tools=set(), + session_id=self.SESSION, + ) + + # The destructive tool was NOT executed; a confirmation was requested. + assert any("event: confirmation_required" in e for e in events) + assert not any("event: tool_use" in e for e in events) + # And the server recorded that it legitimately prompted for this tool_id. + tool_id = compute_tool_id("delete_contact", self.DELETE_INPUT) + assert tool_id in table.items[self.SESSION]["pending_tool_ids"] + + def test_forged_confirmation_does_not_execute(self) -> None: + """A forged confirmed_tools value (no server record) cannot execute. + + This is the core bypass from issue #12: a client that pre-confirms a + tool_id it was never prompted for must still be forced through a real + confirmation round-trip. + """ + table = _SessionStateFakeTable() + forged_tool_id = compute_tool_id("delete_contact", self.DELETE_INPUT) + # The handler validates client claims via filter_authorized_confirmations, + # which returns empty because nothing was recorded -> confirmed is empty. + with _patch_streaming_sdk([self._delete_message()], table): + authorized = session_state.filter_authorized_confirmations( + self.SESSION, [forged_tool_id] + ) + assert authorized == set() + events = _collect_stream( + query="delete contact 5", + nb_slug=TEST_NB_SLUG, + nb_token=TEST_NB_TOKEN, + model="claude-haiku-4-5-20251001", + confirmed_tools=authorized, + session_id=self.SESSION, + ) + + assert any("event: confirmation_required" in e for e in events) + assert not any("event: tool_use" in e for e in events) + + def test_legitimate_confirm_then_execute(self) -> None: + """Recorded confirmation -> the destructive tool executes on resubmit.""" + table = _SessionStateFakeTable() + + # Round 1: prompt + record. + with _patch_streaming_sdk([self._delete_message()], table): + _collect_stream( + query="delete contact 5", + nb_slug=TEST_NB_SLUG, + nb_token=TEST_NB_TOKEN, + model="claude-haiku-4-5-20251001", + confirmed_tools=set(), + session_id=self.SESSION, + ) + + tool_id = compute_tool_id("delete_contact", self.DELETE_INPUT) + + # Round 2: client resubmits; server validates the claim against its record. + with _patch_streaming_sdk([self._delete_message()], table): + authorized = session_state.filter_authorized_confirmations( + self.SESSION, [tool_id] + ) + assert authorized == {tool_id} + events = _collect_stream( + query="delete contact 5", + nb_slug=TEST_NB_SLUG, + nb_token=TEST_NB_TOKEN, + model="claude-haiku-4-5-20251001", + confirmed_tools=authorized, + session_id=self.SESSION, + ) + + # Now the tool executes (tool_use emitted, no new confirmation prompt). + assert any("event: tool_use" in e for e in events) + assert not any("event: confirmation_required" in e for e in events) + # The confirmation was consumed (single-use) and cannot be replayed. + assert "pending_tool_ids" not in table.items[self.SESSION] + + def test_non_destructive_tool_executes_without_confirmation(self) -> None: + table = _SessionStateFakeTable() + msg = _FakeAssistantMessage( + [_FakeToolUseBlock("list_signups", {"filter": {"email": "a@b.c"}})] + ) + with _patch_streaming_sdk([msg], table): + events = _collect_stream( + query="find people", + nb_slug=TEST_NB_SLUG, + nb_token=TEST_NB_TOKEN, + model="claude-haiku-4-5-20251001", + confirmed_tools=set(), + session_id=self.SESSION, + ) + assert any("event: tool_use" in e for e in events) + assert not any("event: confirmation_required" in e for e in events) + + +class TestServerSideUndo: + """Undo history is server-maintained and never trusted from the client.""" + + SESSION = "testnation#user-test-12345" + + def test_client_supplied_undo_stack_is_ignored(self) -> None: + """A forged undo_stack in the body must not reach the agent prompt.""" + table = _SessionStateFakeTable() + captured: dict[str, Any] = {} + + async def fake_stream(**kwargs: Any) -> Any: + captured.update(kwargs) + yield format_sse_event(SSE_EVENT_DONE, {"response": "ok", "tool_calls": []}) + + with patch( + "src.lambdas.nat_agent_streaming.handler.stream_agent_response", + side_effect=lambda **kw: fake_stream(**kw), + ), patch( + "src.lambdas.nat_agent_streaming.handler.verify_nation_subscription", + return_value={"valid": True}, + ), patch( + "src.lambdas.nat_agent_streaming.handler.check_rate_limit", + return_value=None, + ), patch( + "src.lambdas.nat_agent_streaming.handler.check_and_reset_billing_cycle_nation", + return_value=False, + ), patch( + "src.lambdas.nat_agent_streaming.handler.track_query_usage_nation", + return_value=1, + ), patch( + "src.lambdas.nat_agent_streaming.handler.get_nb_tokens_by_nation", + return_value=(TEST_NB_TOKEN, TEST_NB_SLUG), + ): + body = { + "query": "undo that", + "user_id": TEST_USER_ID, + "nation_slug": TEST_NATION_SLUG, + "undo_stack": [ + { + "description": "attacker controlled", + "toolName": "create_signup", + "undoType": "delete_created", + "undoData": {"signup_id": "victim-999"}, + } + ], + } + run_async(_drain(process_streaming_request(body))) + + # stream_agent_response is invoked with a session_id, NOT the client stack. + assert "undo_stack" not in captured + assert captured["session_id"] == make_session_id( + TEST_USER_ID, TEST_NATION_SLUG + ) + + def test_confirmed_tools_filtered_through_server(self) -> None: + """process_streaming_request only forwards server-authorized tool_ids.""" + table = _SessionStateFakeTable() + captured: dict[str, Any] = {} + + async def fake_stream(**kwargs: Any) -> Any: + captured.update(kwargs) + yield format_sse_event(SSE_EVENT_DONE, {"response": "ok", "tool_calls": []}) + + resource = MagicMock() + resource.Table.return_value = table + + with patch( + "src.lambdas.nat_agent_streaming.handler.stream_agent_response", + side_effect=lambda **kw: fake_stream(**kw), + ), patch( + "src.lambdas.nat_agent_streaming.handler.verify_nation_subscription", + return_value={"valid": True}, + ), patch( + "src.lambdas.nat_agent_streaming.handler.check_rate_limit", + return_value=None, + ), patch( + "src.lambdas.nat_agent_streaming.handler.check_and_reset_billing_cycle_nation", + return_value=False, + ), patch( + "src.lambdas.nat_agent_streaming.handler.track_query_usage_nation", + return_value=1, + ), patch( + "src.lambdas.nat_agent_streaming.handler.get_nb_tokens_by_nation", + return_value=(TEST_NB_TOKEN, TEST_NB_SLUG), + ), patch.object( + session_state, "get_dynamodb_resource", return_value=resource + ): + forged = compute_tool_id("delete_contact", {"id": "1"}) + body = { + "query": "do it", + "user_id": TEST_USER_ID, + "nation_slug": TEST_NATION_SLUG, + "confirmed_tools": [forged], + } + run_async(_drain(process_streaming_request(body))) + + # The forged tool_id is dropped: nothing was recorded server-side. + assert captured["confirmed_tools"] == set() + + def test_build_undo_entry_add_to_list(self) -> None: + entry = _build_undo_entry("add_to_list", {"person_id": "1", "list_id": "2"}) + assert entry is not None + assert entry["undoType"] == "remove_from_list" + assert entry["undoData"] == {"person_id": "1", "list_id": "2"} + + def test_build_undo_entry_unreconstructable_returns_none(self) -> None: + # create_signup's undo needs the created id from the tool *result*, which + # the server cannot observe -> not recorded (safe, not forgeable). + assert _build_undo_entry("create_signup", {"first_name": "A"}) is None + + +async def _drain(agen: Any) -> list[str]: + events = [] + async for event in agen: + events.append(event) + return events diff --git a/tests/test_session_state.py b/tests/test_session_state.py new file mode 100644 index 0000000..a8ea54d --- /dev/null +++ b/tests/test_session_state.py @@ -0,0 +1,241 @@ +""" +Unit tests for the server-side session state store (confirmations + undo). + +These cover the authoritative-confirmation invariant that closes the bypass in +issue #12: a tool_id is only honoured as confirmed if the *server* recorded that +it prompted for it; client claims alone carry no authority. +""" + +from __future__ import annotations + +import json +from typing import Any +from unittest.mock import MagicMock, patch + +from botocore.exceptions import ClientError + +from src.lambdas.shared import session_state +from src.lambdas.shared.session_state import ( + append_undo_entry, + compute_tool_id, + consume_confirmation, + filter_authorized_confirmations, + get_undo_stack, + make_session_id, + record_pending_confirmation, +) + + +class FakeTable: + """Minimal in-memory stand-in for a boto3 DynamoDB Table resource. + + Supports the narrow slice of update_item/get_item semantics this module + uses: ADD/DELETE on a string-set attribute and SET on scalar attributes. + """ + + def __init__(self) -> None: + self.items: dict[str, dict[str, Any]] = {} + + def _key(self, key: dict[str, Any]) -> str: + return key["session_id"] + + def update_item( + self, + Key: dict[str, Any], + UpdateExpression: str, + ExpressionAttributeValues: dict[str, Any], + ) -> None: + item = self.items.setdefault(self._key(Key), {}) + expr = UpdateExpression + if "ADD pending_tool_ids :tid" in expr: + current = set(item.get("pending_tool_ids", set())) + current |= set(ExpressionAttributeValues[":tid"]) + item["pending_tool_ids"] = current + if "DELETE pending_tool_ids :tid" in expr: + current = set(item.get("pending_tool_ids", set())) + current -= set(ExpressionAttributeValues[":tid"]) + if current: + item["pending_tool_ids"] = current + else: + item.pop("pending_tool_ids", None) + if "SET undo_stack_json = :stack" in expr: + item["undo_stack_json"] = ExpressionAttributeValues[":stack"] + if ":ttl" in ExpressionAttributeValues: + item["expires_at"] = ExpressionAttributeValues[":ttl"] + + def get_item( + self, Key: dict[str, Any], ProjectionExpression: str | None = None + ) -> dict[str, Any]: + item = self.items.get(self._key(Key)) + if item is None: + return {} + return {"Item": dict(item)} + + +def _patch_table(table: Any) -> Any: + resource = MagicMock() + resource.Table.return_value = table + return patch.object( + session_state, "get_dynamodb_resource", return_value=resource + ) + + +class TestComputeToolId: + def test_is_deterministic(self) -> None: + a = compute_tool_id("delete_contact", {"id": "5"}) + b = compute_tool_id("delete_contact", {"id": "5"}) + assert a == b + + def test_key_order_independent(self) -> None: + a = compute_tool_id("update_signup", {"a": 1, "b": 2}) + b = compute_tool_id("update_signup", {"b": 2, "a": 1}) + assert a == b + + def test_differs_by_input(self) -> None: + a = compute_tool_id("delete_contact", {"id": "5"}) + b = compute_tool_id("delete_contact", {"id": "6"}) + assert a != b + + def test_differs_by_tool(self) -> None: + a = compute_tool_id("delete_contact", {"id": "5"}) + b = compute_tool_id("delete_signup", {"id": "5"}) + assert a != b + + def test_prefixed_with_tool_name(self) -> None: + tid = compute_tool_id("delete_event", {"id": "9"}) + assert tid.startswith("delete_event_") + + +class TestSessionId: + def test_derived_from_identity(self) -> None: + assert make_session_id("user-1", "nation-a") == "v1:8:nation-a#user-1" + + def test_distinct_per_user(self) -> None: + assert make_session_id("u1", "n") != make_session_id("u2", "n") + + def test_distinct_per_nation(self) -> None: + assert make_session_id("u", "n1") != make_session_id("u", "n2") + + def test_no_collision_with_delimiter_in_fields(self) -> None: + # Length-prefixing makes the encoding unambiguous even with '#'/':'. + assert make_session_id("b#c", "a") != make_session_id("c", "a#b") + + +class TestConfirmations: + def test_recorded_confirmation_is_authorized(self) -> None: + table = FakeTable() + with _patch_table(table): + tid = compute_tool_id("delete_contact", {"id": "5"}) + record_pending_confirmation("sess", tid) + assert filter_authorized_confirmations("sess", [tid]) == {tid} + + def test_unprompted_tool_id_is_rejected(self) -> None: + """The bypass: a forged tool_id the server never prompted for is dropped.""" + table = FakeTable() + with _patch_table(table): + forged = compute_tool_id("delete_contact", {"id": "999"}) + assert filter_authorized_confirmations("sess", [forged]) == set() + + def test_only_matching_subset_authorized(self) -> None: + table = FakeTable() + with _patch_table(table): + real = compute_tool_id("delete_contact", {"id": "5"}) + forged = compute_tool_id("delete_signup", {"person_id": "7"}) + record_pending_confirmation("sess", real) + result = filter_authorized_confirmations("sess", [real, forged]) + assert result == {real} + + def test_confirmation_is_session_scoped(self) -> None: + table = FakeTable() + with _patch_table(table): + tid = compute_tool_id("delete_contact", {"id": "5"}) + record_pending_confirmation("session-a", tid) + # A different session cannot use session-a's confirmation. + assert filter_authorized_confirmations("session-b", [tid]) == set() + + def test_consume_removes_authorization(self) -> None: + table = FakeTable() + with _patch_table(table): + tid = compute_tool_id("delete_contact", {"id": "5"}) + record_pending_confirmation("sess", tid) + consume_confirmation("sess", tid) + assert filter_authorized_confirmations("sess", [tid]) == set() + + def test_empty_request_returns_empty(self) -> None: + table = FakeTable() + with _patch_table(table): + assert filter_authorized_confirmations("sess", []) == set() + + def test_record_sets_ttl(self) -> None: + table = FakeTable() + with _patch_table(table), patch.object( + session_state, "get_current_timestamp", return_value=1000 + ): + tid = compute_tool_id("delete_contact", {"id": "5"}) + record_pending_confirmation("sess", tid) + assert table.items["sess"]["expires_at"] == ( + 1000 + session_state.SESSION_STATE_TTL_SECONDS + ) + + def test_filter_fails_closed_on_dynamo_error(self) -> None: + """A store outage must not authorize an action (fail closed).""" + resource = MagicMock() + bad_table = MagicMock() + bad_table.get_item.side_effect = ClientError( + {"Error": {"Code": "InternalServerError", "Message": "boom"}}, + "GetItem", + ) + resource.Table.return_value = bad_table + with patch.object( + session_state, "get_dynamodb_resource", return_value=resource + ): + assert filter_authorized_confirmations("sess", ["anything"]) == set() + + +class TestUndoStack: + def test_append_and_read_roundtrip(self) -> None: + table = FakeTable() + with _patch_table(table): + entry = { + "description": "Add person 1 to list 2", + "toolName": "add_to_list", + "undoType": "remove_from_list", + "undoData": {"person_id": "1", "list_id": "2"}, + } + append_undo_entry("sess", entry) + stack = get_undo_stack("sess") + assert stack == [entry] + + def test_get_empty_when_missing(self) -> None: + table = FakeTable() + with _patch_table(table): + assert get_undo_stack("sess") == [] + + def test_append_caps_history(self) -> None: + table = FakeTable() + with _patch_table(table): + for i in range(session_state.MAX_UNDO_ENTRIES + 5): + append_undo_entry("sess", {"i": i}) + stack = get_undo_stack("sess") + assert len(stack) == session_state.MAX_UNDO_ENTRIES + # Oldest entries dropped; newest retained. + assert stack[-1] == {"i": session_state.MAX_UNDO_ENTRIES + 4} + + def test_get_handles_malformed_json(self) -> None: + table = FakeTable() + table.items["sess"] = {"undo_stack_json": "not json{"} + with _patch_table(table): + assert get_undo_stack("sess") == [] + + def test_get_fails_open_on_dynamo_error(self) -> None: + resource = MagicMock() + bad_table = MagicMock() + bad_table.get_item.side_effect = ClientError( + {"Error": {"Code": "InternalServerError", "Message": "boom"}}, + "GetItem", + ) + resource.Table.return_value = bad_table + with patch.object( + session_state, "get_dynamodb_resource", return_value=resource + ): + assert get_undo_stack("sess") == []