diff --git a/.github/workflows/standalone-product.yml b/.github/workflows/standalone-product.yml index f797057..bd27017 100644 --- a/.github/workflows/standalone-product.yml +++ b/.github/workflows/standalone-product.yml @@ -54,6 +54,32 @@ jobs: run: python -m pip install --disable-pip-version-check fastapi==0.150.0 pydantic==2.14.1 - name: Verify session continuation stays on its own agent run: python -m unittest discover -s tests/backend -p test_session_scope.py -v + session-storage: + runs-on: ubuntu-latest + timeout-minutes: 5 + services: + postgres: + image: postgres:16.6-alpine + env: + POSTGRES_USER: session_storage_admin + POSTGRES_PASSWORD: synthetic-ci-only + POSTGRES_DB: session_storage + ports: + - 127.0.0.1:15441:5432 + options: >- + --health-cmd "pg_isready -U session_storage_admin -d session_storage" + --health-interval 2s --health-timeout 2s --health-retries 20 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - name: Install only the pinned API-test, model and PostgreSQL driver dependencies + run: python -m pip install --disable-pip-version-check fastapi==0.150.0 pydantic==2.14.1 sqlalchemy==2.0.36 sqlmodel==0.0.22 asyncpg==0.32.0 + - name: Session isolation on the product's queries against real PostgreSQL + env: + SESSION_STORAGE_TEST_URL: postgresql+asyncpg://session_storage_admin:synthetic-ci-only@127.0.0.1:15441/session_storage + run: python -m unittest discover -s tests/backend -p test_session_storage.py -v evaluation-reliability: runs-on: ubuntu-latest timeout-minutes: 5 diff --git a/backend/api/v1/routers/chat.py b/backend/api/v1/routers/chat.py index dc74259..c7a318c 100644 --- a/backend/api/v1/routers/chat.py +++ b/backend/api/v1/routers/chat.py @@ -350,12 +350,20 @@ async def export_session_csv( build_export_s3_key, generate_export_presigned_url, check_s3_object_exists, CSV_EXPORT_URL_EXPIRY, ) - chat_log = await get_chat_log_for_export(session_id, chat_id) + # This path names no agent, so an agent key is scoped by an `agent_id` query parameter the + # caller chooses: only that agent's turns can be found with it. + key_agent = auth.get("agent_id") if auth.get("auth_type") == "agent_api_key" else None + chat_log = await get_chat_log_for_export(session_id, chat_id, agent_id=key_agent) if not chat_log: detail = f"Chat turn {chat_id} not found" if chat_id else "No SQL result found for this session" raise HTTPException(status_code=404, detail=detail) if not chat_log.sql_query: raise HTTPException(status_code=404, detail="This chat turn has no SQL result to export") + # The CSV is the turn's rows: an agent that hides result rows does not hand them out here + # either. Checked against the agent that produced the turn. + _, _, show_db_result, _ = await get_agent_with_visibility(chat_log.agent_id) + if not show_db_result: + raise HTTPException(status_code=403, detail="Result rows are hidden for this agent") s3_key = build_export_s3_key(chat_log.agent_id, session_id, chat_log.chat_id) if not await check_s3_object_exists(s3_key): raise HTTPException(status_code=404, detail="Export file not available. The query may have been made before export was enabled.") diff --git a/backend/api/v1/routers/sessions.py b/backend/api/v1/routers/sessions.py index f0d91c3..2276449 100644 --- a/backend/api/v1/routers/sessions.py +++ b/backend/api/v1/routers/sessions.py @@ -28,6 +28,7 @@ UpdateSessionRequest, ) from agent.utils.authorization import require_agent_permission +from api.v1.utils.chat_helper import get_agent_with_visibility router = APIRouter(prefix="/v1") @@ -144,8 +145,10 @@ async def get_session( detail=f"Session {session_id} not found for agent {agent_id}" ) - # Get conversation data - conversation_data = await get_session_data(session_id) + # Get conversation data, under the agent's answer visibility settings + _, show_sql_query, show_db_result, _ = await get_agent_with_visibility(agent_id) + conversation_data = await get_session_data( + session_id, agent_id=agent_id, show_sql_query=show_sql_query, show_db_result=show_db_result) return { "session_id": session_details['session_id'], diff --git a/backend/api/v1/utils/chat_helper.py b/backend/api/v1/utils/chat_helper.py index 9e15f9e..a4201d5 100644 --- a/backend/api/v1/utils/chat_helper.py +++ b/backend/api/v1/utils/chat_helper.py @@ -17,14 +17,21 @@ from sqlmodel import select -async def get_past_conversation_safe(session_id): - """Get past conversation for a session""" +async def get_past_conversation_safe(session_id, agent_id): + """Get past conversation for a session — this agent's turns only. + + A session id can be shared by more than one agent (a chat service may send its own thread + id to every agent it calls). Reading by session id alone replayed one agent's questions and + answers into another agent's model context. + """ try: start_time = time.time() messages = [] async for session in get_async_session(): res = await session.execute( - select(ChatLogs).where(ChatLogs.session_id == session_id).order_by(ChatLogs.created_at) + select(ChatLogs).where( + ChatLogs.session_id == session_id, ChatLogs.agent_id == agent_id, + ).order_by(ChatLogs.created_at) ) rows = res.scalars().all() for chat_log in rows: @@ -45,7 +52,7 @@ async def get_past_messages_safe(session_id, user_id, agent_id, user_query): session_info = await agent_log_new_session(user_id, agent_id, user_query) return {"past_messages": [], "session_id": session_info['session_id']} else: - past_messages = await get_past_conversation_safe(session_id) + past_messages = await get_past_conversation_safe(session_id, agent_id) return {"past_messages": past_messages, "session_id": session_id} except Exception as e: logger.error(f"Error in get_past_messages_safe: {str(e)}") @@ -53,16 +60,19 @@ async def get_past_messages_safe(session_id, user_id, agent_id, user_query): async def ensure_session_belongs_to_agent(session_id: Optional[str], agent_id: str) -> None: - """Refuse to continue a session that was not created on this agent. + """Refuse a session_id that this API created for a different agent. - History is read by session_id alone (get_past_conversation_safe), so an unchecked id would - replay another agent's questions and answers into this agent's model context and append the - new turn to that foreign session. A new conversation (no id) is always allowed. + Runs before any history read or model call, so continuing another agent's conversation is + refused outright rather than only stripped of its history (which + `get_past_conversation_safe` does in any case). A session id this API never created (a chat + service's own thread id) continues with this agent's turns only. A new conversation (no id) + is always allowed. """ if not session_id or session_id in ["", "string"]: return - from session.db_utils.session_db_utils import get_session_details - if not await get_session_details(session_id, agent_id): + from session.db_utils.session_db_utils import get_session_agent_ids + owners = await get_session_agent_ids(session_id) + if owners and agent_id not in owners: raise HTTPException(status_code=404, detail=f"Session {session_id} not found for agent {agent_id}") @@ -328,14 +338,16 @@ async def get_agent_runtime_db(agent_id: str) -> Tuple[dict, float]: return db_response.get("db_details"), elapsed -async def get_chat_log_for_export(session_id: str, chat_id: Optional[str]): - """Fetch the ChatLogs row for a CSV export request.""" +async def get_chat_log_for_export(session_id: str, chat_id: Optional[str], agent_id: Optional[str] = None): + """Fetch the ChatLogs row for a CSV export request; only `agent_id`'s turns when given.""" + scope = [ChatLogs.agent_id == agent_id] if agent_id else [] async for db_session in get_async_session(): if chat_id: res = await db_session.execute( select(ChatLogs).where( ChatLogs.chat_id == chat_id, ChatLogs.session_id == session_id, + *scope, ).limit(1) ) else: @@ -343,26 +355,23 @@ async def get_chat_log_for_export(session_id: str, chat_id: Optional[str]): select(ChatLogs).where( ChatLogs.session_id == session_id, ChatLogs.sql_query.isnot(None), + *scope, ).order_by(ChatLogs.created_at.desc()).limit(1) ) return res.scalar_one_or_none() async def get_session_messages_data(session_id: str, agent_id: str) -> dict: - """Verify agent exists then return session conversation data.""" - async for session in get_async_session(): - result = await session.execute( - select(ChatbotDetails).where(ChatbotDetails.agent_id == agent_id) - ) - agent = result.scalar_one_or_none() - if not agent: - raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") + """Verify agent exists then return session conversation data, under the agent's answer + visibility settings (the same ones live chat applies).""" + _, show_sql_query, show_db_result, _ = await get_agent_with_visibility(agent_id) from session.db_utils.session_db_utils import get_session_data, get_session_details session_detail = await get_session_details(session_id, agent_id) if not session_detail: return {"message": "No session data found"} - session_data = await get_session_data(session_id) + session_data = await get_session_data( + session_id, agent_id=agent_id, show_sql_query=show_sql_query, show_db_result=show_db_result) if session_data: return { "session_id": session_detail['session_id'], diff --git a/backend/api/v1/utils/feedback_helper.py b/backend/api/v1/utils/feedback_helper.py index a1e4757..a47a05d 100644 --- a/backend/api/v1/utils/feedback_helper.py +++ b/backend/api/v1/utils/feedback_helper.py @@ -49,7 +49,8 @@ async def process_feedback(agent_id: str, body) -> dict: session_id=body.session_id, chat_id=body.chat_id, feedback=normalized_label, - corrected_sql=sql_to_store + corrected_sql=sql_to_store, + agent_id=agent_id, ) if normalized_label == "rejected" and body.corrected_query: diff --git a/backend/session/db_utils/session_db_utils.py b/backend/session/db_utils/session_db_utils.py index 9e1577d..aff8df6 100644 --- a/backend/session/db_utils/session_db_utils.py +++ b/backend/session/db_utils/session_db_utils.py @@ -59,6 +59,14 @@ async def get_past_conversation(session_id): logger.error(f"Error getting past conversation: {e}") raise +async def get_session_agent_ids(session_id: str) -> set: + """The agents a session created by this API belongs to; empty when it never created one.""" + async with get_async_session() as session: + res = await session.execute( + select(ChatSessions.chatbot_id, ChatSessions.agent_id).where(ChatSessions.session_id == session_id) + ) + return {agent for row in res.all() for agent in row if agent} + async def get_session_details(session_id: str, agent_id: str): """ Return dict with session_id, created_time, session_title or None. @@ -382,13 +390,24 @@ def _chart_for_replay(row) -> Optional[dict]: return None -async def get_session_data(session_id: str): +async def get_session_data(session_id: str, *, agent_id: str, show_sql_query: bool, show_db_result: bool): + """A saved conversation — this agent's turns — under the agent's answer visibility settings. + + All three are required (no default) so every caller decides. Only `agent_id`'s turns are + served: a session id can carry another agent's turns too (a chat service may send its own + thread id to every agent it calls). Live chat leaves the SQL out when `show_sql_query` is + off and the rows out when `show_db_result` is off, and a replay must not hand back what the + live answer withheld. Pass the flags from `get_agent_with_visibility`, the same resolution + live chat uses. + """ try: logger.info(f"Getting session data for session_id: {session_id}") session_data = [] async with get_async_session() as session: res = await session.execute( - select(ChatLogs).where(ChatLogs.session_id == session_id).order_by(ChatLogs.created_at) + select(ChatLogs).where( + ChatLogs.session_id == session_id, ChatLogs.agent_id == agent_id, + ).order_by(ChatLogs.created_at) ) rows = res.all() logger.info(f"Found {len(rows) if rows else 0} messages for session {session_id}") @@ -417,12 +436,16 @@ async def get_session_data(session_id: str): } # Build debug info (no dedicated debug_info column in DB) if any([row_obj.user_query, row_obj.refined_query, row_obj.sql_query, row_obj.db_output]): - assistant_message["debug_info"] = { + debug_info = { "query": row_obj.user_query, "refine_query": row_obj.refined_query, - "sql_query": row_obj.sql_query, - "db_output": row_obj.db_output } + # Same omission as the live answer: the key is left out, not nulled. + if show_sql_query: + debug_info["sql_query"] = row_obj.sql_query + if show_db_result: + debug_info["db_output"] = row_obj.db_output + assistant_message["debug_info"] = debug_info else: assistant_message["debug_info"] = None # The ONLY change to this payload: the chart is sent instead of a @@ -430,10 +453,11 @@ async def get_session_data(session_id: str): assistant_message["graph_output"] = _chart_for_replay(row_obj) # Additive: the column-labelled result block stored for this turn # ({columns, rows, row_count, truncated}); `db_output` above keeps bare rows - # without column names. None for turns stored before the block existed. + # without column names. None for turns stored before the block existed, and + # when the agent hides result rows. stored_debug = getattr(row_obj, "debug_info", None) assistant_message["result"] = ( - stored_debug.get("db_result") if isinstance(stored_debug, dict) else None + stored_debug.get("db_result") if show_db_result and isinstance(stored_debug, dict) else None ) session_data.append(assistant_message) logger.info(f"Returning {len(session_data)} total messages for session {session_id}") @@ -453,10 +477,13 @@ async def get_session_chat_counter_from_db(session_id: str): logger.error(f"Error getting session chat counter: {e}") return None -async def add_user_feedback(session_id: str, chat_id: str, feedback: str, corrected_sql: str = None): +async def add_user_feedback(session_id: str, chat_id: str, feedback: str, corrected_sql: str = None, *, agent_id: str): + """Record feedback on one of `agent_id`'s turns. A turn of another agent is not found, so a + caller allowed to use one agent cannot rewrite another agent's stored verdict or SQL.""" try: async with get_async_session() as session: - res = await session.execute(select(ChatLogs).where(ChatLogs.session_id == session_id, ChatLogs.chat_id == chat_id)) + res = await session.execute(select(ChatLogs).where( + ChatLogs.session_id == session_id, ChatLogs.chat_id == chat_id, ChatLogs.agent_id == agent_id)) row = res.first() if not row: return diff --git a/docs/api.md b/docs/api.md index 9745289..a9b7243 100644 --- a/docs/api.md +++ b/docs/api.md @@ -24,6 +24,19 @@ uses the signed HTTP-only session cookie; no hosted Auth service is required. | `GET /v1/agents/{id}/sessions` | Saved conversations on the agent | | `GET /v1/agents/{id}/sessions/{session_id}` | A saved conversation, each answer with its SQL, `result` rows and chart | +An agent only ever sees its own turns of a conversation: chat history and replays +are read for that agent alone. A `session_id` this API created for another agent +returns HTTP404 before any history is read or a model is called; a `session_id` +the API did not create (for example a chat service's own thread id) continues +with this agent's turns only. + +`show_sql_query` and `show_db_result` apply to saved conversations exactly as to +live answers: a replay leaves out the SQL and the rows (`debug_info.db_output` +and `result`) the agent hides, and `GET /v1/sessions/{session_id}/export` returns +HTTP403 for an agent that hides result rows; with an agent key it finds only that +agent's turns. Both settings are read when the +conversation is served, so turning one back on shows the stored SQL or rows again. + Use the generated OpenAPI contract for complete request/response fields and streaming options. `tests/evaluation/run.py` is a small working client example: it loads configuration privately, creates the sample agent/source/table, issues diff --git a/docs/mcp.md b/docs/mcp.md index 9574f95..9f55ead 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -127,14 +127,16 @@ and `chart_output`. Send a follow-up to - The agent's source credentials should be a dedicated read-only database account. The adapter cannot widen what that account or the agent's table visibility allows; it only forwards questions. -- A conversation can be continued only on the agent that created it: a - `session_id` from another agent returns HTTP 404 before any history is read - or any model is called. Deployments running an API version older than this - change do not perform that check. -- Saved conversations replayed through the API include SQL and rows regardless - of the agent's visibility settings. The adapter applies the agent's - `show_sql_query` and `show_db_result` to them, matching live answers, but a - caller using the API directly with the same key can read them. +- An agent only ever sees its own turns of a conversation: history and replays + are read for that agent alone, and a `session_id` created on another agent + returns HTTP 404 before any history is read or any model is called. + Deployments running an API version older than this change do not perform + that check. +- Saved conversations follow the agent's `show_sql_query` and `show_db_result` + in the API itself, exactly as live answers do: a replay leaves out the SQL and + rows the agent hides, and CSV export returns HTTP 403 when rows are hidden. + Deployments running an API version older than this change return them; the + adapter applies the same settings either way. - An agent key can list and read every conversation on its own agent, not only the ones it created. @@ -142,8 +144,8 @@ and `chart_output`. Send a follow-up to `tests/mcp` checks the adapter through real MCP client sessions, in memory and over stdio, against a synthetic API with the product's credential rules. CI runs -it with `tests/backend/test_session_scope.py`, which covers the session check -and saved-result replay in the API itself. +it with `tests/backend/test_session_scope.py`, which covers the session check, +saved-result replay and saved-conversation visibility in the API itself. The integration was also run end to end on a disposable installation of this branch (Compose stack, default model profile, the synthetic `sample_sales` diff --git a/docs/source-preservation.json b/docs/source-preservation.json index 3aaf70a..50500b4 100644 --- a/docs/source-preservation.json +++ b/docs/source-preservation.json @@ -1100,9 +1100,9 @@ }, "backend/api/v1/routers/chat.py": { "baseline_sha256": "2649938a3a23278b9a3ad952853436987b27e2918787af6b813b272059a47d78", - "current_sha256": "cf916be23d558e1b8da82355f21e2cff7ae87f04ee4cfa3a496242c24c192642", - "issue": "DA-011 / DA-MCP-001", - "reason": "Central versioned model configuration and explicit nonsecret configuration error response; existing prompts and SQL/security policy semantics preserved. Runtime acceptance pending. DA-MCP-001: session-agent scope check before history/model." + "current_sha256": "c505fdf62dd8af040a340cb513642016656026606469ac50b067c88642582dd2", + "issue": "DA-011 / DA-MCP-001, DA-MCP-002", + "reason": "Central versioned model configuration and explicit nonsecret configuration error response; existing prompts and SQL/security policy semantics preserved. Runtime acceptance pending. DA-MCP-001: session-agent scope check before history/model. DA-MCP-002: Chat history and saved replay read only the requesting agent's turns (a session created on another agent is still refused with 404) and an agent key exports only its own agent's turns. Saved conversations apply the agent's answer settings at the API: replay (session detail and messages) leaves out the SQL and rows that show_sql_query / show_db_result hide, including the column-labelled result block, and CSV export returns 403 for an agent that hides rows. Prompts, models, SQL generation and security guards unchanged." }, "backend/api/v1/utils/tables_helper.py": { "baseline_sha256": "7d02daefc0bd22fe516fdfbfd26d034c3b55c05655fcb1fd4bd9a42029d1d09b", @@ -1124,15 +1124,27 @@ }, "backend/api/v1/utils/chat_helper.py": { "baseline_sha256": "a19c1b056d58c9cf570b77eb5d6ff7e6e845a01648056636ba2b8f6f20d66385", - "current_sha256": "d545db45dc86dc5521a561abb31ea63652f100142d206184e61f59d7823fedd6", - "issue": "DA-MCP-001", - "reason": "Session continuation is scoped to the agent that created the session (cross-agent history replay closed), and saved-conversation replay additively serves the stored column-labelled result block; required by the MCP adapter. No prompt, model, SQL generation or execution change." + "current_sha256": "ea1c447f054c6b3ff25d5f47835ad3a3373be821bee84c75b87825cc95a6f4dd", + "issue": "DA-MCP-001, DA-MCP-002", + "reason": "Session continuation is scoped to the agent that created the session (cross-agent history replay closed), and saved-conversation replay additively serves the stored column-labelled result block; required by the MCP adapter. No prompt, model, SQL generation or execution change. DA-MCP-002: Chat history and saved replay read only the requesting agent's turns (a session created on another agent is still refused with 404) and an agent key exports only its own agent's turns. Saved conversations apply the agent's answer settings at the API: replay (session detail and messages) leaves out the SQL and rows that show_sql_query / show_db_result hide, including the column-labelled result block, and CSV export returns 403 for an agent that hides rows. Prompts, models, SQL generation and security guards unchanged." }, "backend/session/db_utils/session_db_utils.py": { "baseline_sha256": "18e03d5966a451f0b1bc6d5777e2254effc1c1ad248f033d0067c7d7faf64d11", - "current_sha256": "08b3aa026c754695bf795f23191cab5f1445eebe5306cab3ced501b9b0e89499", - "issue": "DA-MCP-001", - "reason": "Session continuation is scoped to the agent that created the session (cross-agent history replay closed), and saved-conversation replay additively serves the stored column-labelled result block; required by the MCP adapter. No prompt, model, SQL generation or execution change." + "current_sha256": "f8432c65a862153d033a5003b3afbed4a0a1848c813b3716a9322f5b72625be4", + "issue": "DA-MCP-001, DA-MCP-002", + "reason": "Session continuation is scoped to the agent that created the session (cross-agent history replay closed), and saved-conversation replay additively serves the stored column-labelled result block; required by the MCP adapter. No prompt, model, SQL generation or execution change. DA-MCP-002: Chat history and saved replay read only the requesting agent's turns (a session created on another agent is still refused with 404) and an agent key exports only its own agent's turns. Saved conversations apply the agent's answer settings at the API: replay (session detail and messages) leaves out the SQL and rows that show_sql_query / show_db_result hide, including the column-labelled result block, and CSV export returns 403 for an agent that hides rows. Prompts, models, SQL generation and security guards unchanged." + }, + "backend/api/v1/routers/sessions.py": { + "baseline_sha256": "77f78284e95e36bec167d39cbe89ed8f8969ce958425882905360f1650f7e638", + "current_sha256": "dc345228a30bd8d3fd539c7e2b90c1e7c7b0688306c46f446b6a9346bec638cf", + "issue": "DA-MCP-002", + "reason": "Chat history and saved replay read only the requesting agent's turns (a session created on another agent is still refused with 404) and an agent key exports only its own agent's turns. Saved conversations apply the agent's answer settings at the API: replay (session detail and messages) leaves out the SQL and rows that show_sql_query / show_db_result hide, including the column-labelled result block, and CSV export returns 403 for an agent that hides rows. Prompts, models, SQL generation and security guards unchanged." + }, + "backend/api/v1/utils/feedback_helper.py": { + "baseline_sha256": "307272276f3ff5242d2277db6e0ee5e62ee33861951f22c45b20459ded10f349", + "current_sha256": "ef4bdadb07b2cc06a95543ec9e3d86ab00ca17a0ad5d15f9974599278639712c", + "issue": "DA-MCP-002", + "reason": "Feedback is recorded only on a turn of the agent in the path, so a caller allowed to use one agent cannot rewrite another agent's stored verdict or SQL." } } } diff --git a/tests/backend/test_session_scope.py b/tests/backend/test_session_scope.py index de1140d..e6b0bb1 100644 --- a/tests/backend/test_session_scope.py +++ b/tests/backend/test_session_scope.py @@ -1,132 +1,305 @@ -"""A supplied session_id continues a conversation only on the agent that created it. - -Chat history is read by session_id alone, so without this check a caller holding a key for -agent A could pass a session created on agent B and have B's questions and answers replayed -into A's model context (and appended to B's session). Synthetic boundaries; no external calls. -""" -import asyncio -import importlib.util -import json -import sys -import types -import unittest -import uuid -from unittest.mock import AsyncMock, patch - -from fastapi import BackgroundTasks, Depends, HTTPException, Path, Query -from fastapi.responses import JSONResponse -from starlette.requests import Request -from access_policy_harness import BACKEND, load_function - -sys.path.insert(0, str(BACKEND)) -from agent.utils.model_profiles import ModelConfigurationError - - -def session_store(owner_agent): - """get_session_details as the product implements it: match on session AND agent.""" - async def get_session_details(session_id, agent_id): - if session_id == 'session-of-b' and agent_id == owner_agent: - return {'session_id': session_id, 'created_time': None, 'session_title': 't'} - return None - return types.SimpleNamespace(get_session_details=AsyncMock(side_effect=get_session_details)) - - -class SessionScopeTests(unittest.IsolatedAsyncioTestCase): - def setUp(self): - self.guard = load_function('api/v1/utils/chat_helper.py', 'ensure_session_belongs_to_agent', - HTTPException=HTTPException) - spec = importlib.util.spec_from_file_location('synthetic_chat_schemas', BACKEND/'api/v1/schemas/chat.py') - self.schemas = importlib.util.module_from_spec(spec) - spec.loader.exec_module(self.schemas) - self.history = AsyncMock(return_value={'past_messages': [], 'session_id': 'session-of-b'}) - self.flow = AsyncMock(return_value={'response': 'ok', 'is_answered': True}) - self.endpoint = load_function('api/v1/routers/chat.py', 'chat_with_agent', - Path=Path, Query=Query, Depends=Depends, BackgroundTasks=BackgroundTasks, - ChatRequest=self.schemas.ChatRequest, Request=Request, - get_api_key_or_session=lambda: {}, require_agent_permission=lambda *args: lambda: {}, - HTTPException=HTTPException, ModelConfigurationError=ModelConfigurationError, - JSONResponse=JSONResponse, uuid=uuid, asyncio=asyncio, - get_agent_with_visibility=AsyncMock(return_value=(object(), True, True, True)), - apply_agent_answer_settings=lambda *args: {}, get_owner_email_by_agent_id=AsyncMock(return_value='owner'), - enforce_query_cap=AsyncMock(), ensure_session_belongs_to_agent=self.guard, - get_past_messages_safe=self.history, - get_agent_runtime_db=AsyncMock(return_value=({'db_id': 'db'}, 0)), - run_agent_flow=self.flow, _generate_follow_up_safe=AsyncMock(return_value=None), - maybe_upload_csv_export=AsyncMock(return_value=None), persist_chat_logs=AsyncMock(), - persist_health_event=AsyncMock(), - build_chat_response=lambda result, body, session_id, *rest: {'session_id': session_id, - 'response': result['response']}) - - async def ask(self, agent_id, session_id): - body = self.schemas.ChatRequest(query='how many orders?', session_id=session_id) - with patch.dict(sys.modules, { - 'session': types.ModuleType('session'), 'session.db_utils': types.ModuleType('session.db_utils'), - 'session.db_utils.session_db_utils': session_store('agent-b'), - 'agent.utils.common_utils': types.SimpleNamespace(serialize_db_value=lambda value: value)}): - return await self.endpoint(agent_id=agent_id, stream=False, body=body, - auth={'auth_type': 'agent_api_key', 'agent_id': agent_id}, user={}) - - async def test_foreign_session_is_refused_before_history_or_model(self): - with self.assertRaises(HTTPException) as caught: - await self.ask('agent-a', 'session-of-b') - self.assertEqual(caught.exception.status_code, 404) - self.history.assert_not_awaited() - self.flow.assert_not_awaited() - - async def test_unknown_session_is_refused(self): - with self.assertRaises(HTTPException) as caught: - await self.ask('agent-b', 'never-created') - self.assertEqual(caught.exception.status_code, 404) - self.flow.assert_not_awaited() - - async def test_own_session_continues(self): - response = await self.ask('agent-b', 'session-of-b') - self.assertEqual(response.status_code, 200) - self.assertEqual(json.loads(response.body)['session_id'], 'session-of-b') - self.history.assert_awaited_once() - self.flow.assert_awaited_once() - - async def test_new_conversation_needs_no_session(self): - for missing in (None, '', 'string'): - with self.subTest(session_id=missing): - response = await self.ask('agent-a', missing) - self.assertEqual(response.status_code, 200) - - -class SavedResultReplayTests(unittest.IsolatedAsyncioTestCase): - """GET .../sessions/{id} additively serves the column-labelled block stored per turn.""" - - async def replay(self, debug_info): - from contextlib import asynccontextmanager - from datetime import datetime - from unittest.mock import MagicMock - row = types.SimpleNamespace(chat_id='c1', user_query='totals?', gpt_response='North 250', refined_query=None, - sql_query='SELECT 1', db_output='[["North", 250]]', feedback=None, - created_at=datetime(2026, 10, 23), debug_info=debug_info, - chart_output=None, graph_output=None) - - @asynccontextmanager - async def get_async_session(): - yield types.SimpleNamespace(execute=AsyncMock(return_value=types.SimpleNamespace(all=lambda: [(row,)]))) - - chart = load_function('session/db_utils/session_db_utils.py', '_chart_for_replay') - replay = load_function('session/db_utils/session_db_utils.py', 'get_session_data', - get_async_session=get_async_session, select=MagicMock(), ChatLogs=MagicMock(), - Row=tuple, _chart_for_replay=chart) - return await replay('s1') - - async def test_stored_result_block_is_served_alongside_unchanged_fields(self): - block = {'columns': [{'name': 'region', 'type': 'text'}], 'rows': [{'region': 'North'}], - 'row_count': 1, 'truncated': False} - user, assistant = await self.replay({'db_result': block, 'sql_query': 'SELECT 1'}) - self.assertEqual(assistant['result'], block) - self.assertEqual(assistant['debug_info']['db_output'], '[["North", 250]]') - self.assertNotIn('result', user) - - async def test_turns_without_a_stored_block_replay_none(self): - _, assistant = await self.replay(None) - self.assertIsNone(assistant['result']) - - -if __name__ == '__main__': - unittest.main(verbosity=2) +"""Saved conversations stay on their agent and under their agent's answer visibility. + +1. A session this API created for agent B cannot be continued on agent A: refused before any + history read or model call. A session id the API never created (the platform + chat service sends its own thread id) continues, with only this agent's turns as history — + test_session_storage.py checks that scoping on real SQL. +2. A replayed conversation applies the agent's `show_sql_query` / `show_db_result` exactly as + the live answer does, on every route that serves saved turns (session detail, session + messages, CSV export) — not only in clients that choose to hide the fields. + +Synthetic agents, sessions and rows; no database, network or model calls. +""" +import ast +import asyncio +import importlib.util +import json +import sys +import types +import unittest +import uuid +from contextlib import asynccontextmanager +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +from fastapi import BackgroundTasks, Depends, HTTPException, Path, Query +from fastapi.responses import JSONResponse +from starlette.requests import Request +from access_policy_harness import BACKEND, load_function + +sys.path.insert(0, str(BACKEND)) +from agent.utils.model_profiles import ModelConfigurationError + +SQL = 'SELECT region, SUM(amount) FROM orders GROUP BY region' +ROWS = '[["North", 250.0], ["row-only-region", 200.0]]' # a value that appears only in the rows + + +def session_store(owner_agent): + """get_session_agent_ids: the agents a session this API created belongs to.""" + async def get_session_agent_ids(session_id): + return {owner_agent} if session_id == 'session-of-b' else set() + return AsyncMock(side_effect=get_session_agent_ids) + + +def session_module(**attrs): + return {'session': types.ModuleType('session'), 'session.db_utils': types.ModuleType('session.db_utils'), + 'session.db_utils.session_db_utils': types.SimpleNamespace(**attrs)} + + +class SessionScopeTests(unittest.IsolatedAsyncioTestCase): + def setUp(self): + self.guard = load_function('api/v1/utils/chat_helper.py', 'ensure_session_belongs_to_agent', + HTTPException=HTTPException) + spec = importlib.util.spec_from_file_location('synthetic_chat_schemas', BACKEND/'api/v1/schemas/chat.py') + self.schemas = importlib.util.module_from_spec(spec) + spec.loader.exec_module(self.schemas) + self.history = AsyncMock(return_value={'past_messages': [], 'session_id': 'session-of-b'}) + self.flow = AsyncMock(return_value={'response': 'ok', 'is_answered': True}) + self.endpoint = load_function('api/v1/routers/chat.py', 'chat_with_agent', + Path=Path, Query=Query, Depends=Depends, BackgroundTasks=BackgroundTasks, + ChatRequest=self.schemas.ChatRequest, Request=Request, + get_api_key_or_session=lambda: {}, require_agent_permission=lambda *args: lambda: {}, + HTTPException=HTTPException, ModelConfigurationError=ModelConfigurationError, + JSONResponse=JSONResponse, uuid=uuid, asyncio=asyncio, + get_agent_with_visibility=AsyncMock(return_value=(object(), True, True, True)), + apply_agent_answer_settings=lambda *args: {}, get_owner_email_by_agent_id=AsyncMock(return_value='owner'), + enforce_query_cap=AsyncMock(), + ensure_session_belongs_to_agent=self.guard, get_past_messages_safe=self.history, + get_agent_runtime_db=AsyncMock(return_value=({'db_id': 'db'}, 0)), + run_agent_flow=self.flow, _generate_follow_up_safe=AsyncMock(return_value=None), + maybe_upload_csv_export=AsyncMock(return_value=None), persist_chat_logs=AsyncMock(), + persist_health_event=AsyncMock(), + build_chat_response=lambda result, body, session_id, *rest: {'session_id': session_id, + 'response': result['response']}) + + async def ask(self, agent_id, session_id): + body = self.schemas.ChatRequest(query='how many orders?', session_id=session_id) + with patch.dict(sys.modules, { + **session_module(get_session_agent_ids=session_store('agent-b')), + 'agent.utils.common_utils': types.SimpleNamespace(serialize_db_value=lambda value: value)}): + return await self.endpoint(agent_id=agent_id, stream=False, body=body, + auth={'auth_type': 'agent_api_key', 'agent_id': agent_id}, user={}) + + async def test_foreign_session_is_refused_before_history_or_model(self): + # The query-cap check before it only reads the owner's count; nothing is charged. + with self.assertRaises(HTTPException) as caught: + await self.ask('agent-a', 'session-of-b') + self.assertEqual(caught.exception.status_code, 404) + self.history.assert_not_awaited() + self.flow.assert_not_awaited() + + async def test_session_id_the_api_never_created_continues_on_this_agent(self): + # The platform chat service's own thread id: not refused; history is asked for this + # agent's turns only (the read itself is scoped; see test_session_storage.py). + response = await self.ask('agent-a', 'chat-service-thread') + self.assertEqual(response.status_code, 200) + session_id, _, agent_id, _ = self.history.await_args.args + self.assertEqual((session_id, agent_id), ('chat-service-thread', 'agent-a')) + self.flow.assert_awaited_once() + + async def test_own_session_continues(self): + response = await self.ask('agent-b', 'session-of-b') + self.assertEqual(response.status_code, 200) + self.assertEqual(json.loads(response.body)['session_id'], 'session-of-b') + self.history.assert_awaited_once() + self.flow.assert_awaited_once() + + async def test_new_conversation_needs_no_session(self): + for missing in (None, '', 'string'): + with self.subTest(session_id=missing): + response = await self.ask('agent-a', missing) + self.assertEqual(response.status_code, 200) + + +RESULT = {'columns': [{'name': 'region', 'type': 'text'}, {'name': 'total', 'type': 'numeric'}], + 'rows': [{'region': 'North', 'total': 250.0}, {'region': 'row-only-region', 'total': 200.0}], + 'row_count': 2, 'truncated': False} + + +def stored_row(debug_info=None): + return types.SimpleNamespace(chat_id='c1', user_query='totals by region?', gpt_response='North 250, South 200', + refined_query='totals by region', sql_query=SQL, db_output=ROWS, feedback=None, + created_at=datetime(2026, 10, 23), chart_output=None, + debug_info={'sql_query': SQL, 'db_result': RESULT} if debug_info is None else debug_info, + graph_output={'graph_suggestion': 'bar', 'graph_data': [{'label': 'North', 'value': 250}], + 'metadata': {}}) + + +def replay_function(rows): + @asynccontextmanager + async def get_async_session(): + yield types.SimpleNamespace(execute=AsyncMock(return_value=types.SimpleNamespace(all=lambda: [(r,) for r in rows]))) + chart = load_function('session/db_utils/session_db_utils.py', '_chart_for_replay') + return load_function('session/db_utils/session_db_utils.py', 'get_session_data', + get_async_session=get_async_session, select=MagicMock(), ChatLogs=MagicMock(), + Row=tuple, _chart_for_replay=chart) + + +class ReplayVisibilityTests(unittest.IsolatedAsyncioTestCase): + """get_session_data serves only what the live answer would have shown.""" + + async def replay(self, show_sql_query, show_db_result): + return await replay_function([stored_row()])( + 's1', agent_id='agent-visible', show_sql_query=show_sql_query, show_db_result=show_db_result) + + async def test_each_setting_controls_its_own_field(self): + for show_sql, show_rows in ((True, True), (False, True), (True, False), (False, False)): + with self.subTest(show_sql_query=show_sql, show_db_result=show_rows): + user, assistant = await self.replay(show_sql, show_rows) + debug = assistant['debug_info'] + self.assertEqual(debug['query'], 'totals by region?') + self.assertEqual(debug['refine_query'], 'totals by region') + self.assertEqual(debug.get('sql_query'), SQL if show_sql else None) + self.assertEqual('sql_query' in debug, show_sql) + self.assertEqual(debug.get('db_output'), ROWS if show_rows else None) + self.assertEqual('db_output' in debug, show_rows) + # The column-labelled block stored for the turn follows the rows setting. + self.assertEqual(assistant['result'], RESULT if show_rows else None) + self.assertNotIn('result', user) + # Anything else on the turn is untouched by the settings. + self.assertEqual(assistant['message'], 'North 250, South 200') + self.assertEqual(assistant['graph_output']['chart_suggestion'], 'bar') + self.assertEqual(user['message'], 'totals by region?') + if not show_sql: + self.assertNotIn('SELECT', json.dumps([user, assistant])) + if not show_rows: + self.assertNotIn('row-only-region', json.dumps([user, assistant])) + + async def test_turns_without_a_stored_block_replay_none(self): + _, assistant = await replay_function([stored_row(debug_info={})])( + 's1', agent_id='agent-visible', show_sql_query=True, show_db_result=True) + self.assertIsNone(assistant['result']) + self.assertEqual(assistant['debug_info']['db_output'], ROWS) + + async def test_callers_must_state_the_settings(self): + with self.assertRaises(TypeError): + await replay_function([stored_row()])('s1') + + def test_every_product_caller_passes_both_settings(self): + """A gate on the call sites: no route can serve saved turns without the agent and its settings.""" + callers = [] + for path in BACKEND.rglob('*.py'): + if 'tests' in path.parts: + continue + tree = ast.parse(path.read_bytes()) + for node in ast.walk(tree): + if isinstance(node, ast.Call): + name = getattr(node.func, 'id', None) or getattr(node.func, 'attr', None) + if name == 'get_session_data': + callers.append((path.relative_to(BACKEND).as_posix(), {k.arg for k in node.keywords})) + self.assertTrue(callers) + for where, keywords in callers: + with self.subTest(caller=where): + self.assertLessEqual({'agent_id', 'show_sql_query', 'show_db_result'}, keywords) + + +def visibility(flags_by_agent): + async def get_agent_with_visibility(agent_id): + if agent_id not in flags_by_agent: + raise HTTPException(status_code=404, detail=f'Agent {agent_id} not found') + show_sql, show_rows = flags_by_agent[agent_id] + return object(), show_sql, show_rows, show_sql or show_rows + return AsyncMock(side_effect=get_agent_with_visibility) + + +AGENTS = {'agent-visible': (True, True), 'agent-hidden': (False, False)} + + +def owned_by(agent): + async def details(session_id, agent_id): + return {'session_id': session_id, 'created_time': None, 'session_title': 't'} if agent_id == agent else None + return details + + +class SavedConversationRouteTests(unittest.IsolatedAsyncioTestCase): + """Allowed and denied behaviour on each route that serves saved turns.""" + + async def assert_route(self, call, agent): + payload = await call(agent) + text = json.dumps(payload, default=str) + if AGENTS[agent][0]: + self.assertIn(SQL, text) + self.assertIn('row-only-region', text) + else: + self.assertNotIn('SELECT', text) + self.assertNotIn('row-only-region', text) + self.assertIn('North 250, South 200', text) # the answer itself is always served + + async def test_session_detail_route(self): + async def call(agent): + route = load_function('api/v1/routers/sessions.py', 'get_session', + Path=Path, Request=Request, Depends=Depends, HTTPException=HTTPException, + get_api_key_or_session=lambda: {}, require_agent_permission=lambda *a: lambda: {}, + get_session_details=AsyncMock(side_effect=owned_by(agent)), + get_session_data=replay_function([stored_row()]), + get_agent_with_visibility=visibility(AGENTS), datetime=datetime, + traceback=MagicMock()) + return await route(agent_id=agent, session_id='s1', auth={'user_id': None}, user={}) + for agent in AGENTS: + with self.subTest(agent=agent): + await self.assert_route(call, agent) + + async def test_session_messages_route(self): + async def call(agent): + helper = load_function('api/v1/utils/chat_helper.py', 'get_session_messages_data', + get_agent_with_visibility=visibility(AGENTS)) + with patch.dict(sys.modules, session_module( + get_session_details=AsyncMock(side_effect=owned_by(agent)), + get_session_data=replay_function([stored_row()]))): + return await helper('s1', agent) + for agent in AGENTS: + with self.subTest(agent=agent): + await self.assert_route(call, agent) + + async def test_unknown_agent_is_refused(self): + helper = load_function('api/v1/utils/chat_helper.py', 'get_session_messages_data', + get_agent_with_visibility=visibility(AGENTS)) + with self.assertRaises(HTTPException) as caught: + await helper('s1', 'agent-missing') + self.assertEqual(caught.exception.status_code, 404) + + +class CsvExportVisibilityTests(unittest.IsolatedAsyncioTestCase): + """The CSV export is the turn's rows: refused for an agent that hides them, and an agent + key reaches only its own agent's turns.""" + + async def export(self, agent, auth=None): + presign = self.presign = AsyncMock(return_value='https://synthetic.example.test/export.csv') + s3 = types.SimpleNamespace(build_export_s3_key=lambda *parts: '/'.join(parts), + generate_export_presigned_url=presign, + check_s3_object_exists=AsyncMock(return_value=True), CSV_EXPORT_URL_EXPIRY=900) + turn = types.SimpleNamespace(agent_id=agent, chat_id='c1', sql_query=SQL) + + async def lookup(session_id, chat_id, agent_id=None): # the product query's agent filter + return turn if agent_id in (None, turn.agent_id) else None + route = load_function('api/v1/routers/chat.py', 'export_session_csv', + Path=Path, Query=Query, Depends=Depends, HTTPException=HTTPException, + JSONResponse=JSONResponse, get_api_key_or_session=lambda: {}, + get_chat_log_for_export=AsyncMock(side_effect=lookup), + get_agent_with_visibility=visibility(AGENTS)) + parents = {n: types.ModuleType(n) for n in ('agent', 'agent.utils')} + with patch.dict(sys.modules, {**parents, 'agent.utils.s3_export_utils': s3}): + return await route(session_id="s1", chat_id="c1", auth=auth or {'auth_type': 'master_api_key'}) + + async def test_agent_key_reaches_only_its_own_agents_turns(self): + own = await self.export('agent-visible', {'auth_type': 'agent_api_key', 'agent_id': 'agent-visible'}) + self.assertEqual(own.status_code, 200) + with self.assertRaises(HTTPException) as caught: + await self.export('agent-visible', {'auth_type': 'agent_api_key', 'agent_id': 'agent-other'}) + self.assertEqual(caught.exception.status_code, 404) + self.presign.assert_not_awaited() + + async def test_visible_agent_gets_a_download_link(self): + response = await self.export("agent-visible") + self.assertEqual(response.status_code, 200) + self.assertEqual(json.loads(response.body)['export_url'], 'https://synthetic.example.test/export.csv') + + async def test_hidden_agent_is_refused_before_a_link_is_signed(self): + with self.assertRaises(HTTPException) as caught: + await self.export('agent-hidden') + self.assertEqual(caught.exception.status_code, 403) + self.presign.assert_not_awaited() + + +if __name__ == '__main__': + unittest.main(verbosity=2) diff --git a/tests/backend/test_session_storage.py b/tests/backend/test_session_storage.py new file mode 100644 index 0000000..3582259 --- /dev/null +++ b/tests/backend/test_session_storage.py @@ -0,0 +1,151 @@ +"""Session isolation on actual SQL: the product's queries against real PostgreSQL tables. + +Runs against the disposable database named by SESSION_STORAGE_TEST_URL (an asyncpg URL; the +session-storage CI job starts one) and is skipped without it. The tables are the product's own +models (chat_sessions, nl2sql_logs); the functions are the product's source bodies, bound to +that database. Synthetic agents, sessions and rows only. + +Covers the case the mocked tests cannot: a session id shared by two agents — the platform chat +service sends its own thread id to every agent it calls — must neither be refused nor replay +the other agent's turns. +""" +import os +import types +import unittest +import uuid +from contextlib import asynccontextmanager +from datetime import datetime, timedelta +from unittest.mock import patch +import sys + +from fastapi import HTTPException +from access_policy_harness import BACKEND, load_function + +sys.path.insert(0, str(BACKEND)) +TEST_URL = os.getenv("SESSION_STORAGE_TEST_URL") + + +@unittest.skipUnless(TEST_URL, "SESSION_STORAGE_TEST_URL not set") +class SessionStorageTests(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + from sqlalchemy.engine.row import Row + from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + from sqlmodel import select + from models.sql_models import ChatLogs, ChatSessions + + from sqlalchemy import text + + # A schema of its own: the shared CI database may already hold older shapes of these + # tables (created by other suites), which create_all would silently keep. + run = uuid.uuid4().hex[:8] + self.schema = f"session_storage_{run}" + boot = create_async_engine(TEST_URL) + async with boot.begin() as conn: + await conn.execute(text(f'CREATE SCHEMA "{self.schema}"')) + await boot.dispose() + self.engine = create_async_engine(TEST_URL, connect_args={"server_settings": {"search_path": self.schema}}) + async with self.engine.begin() as conn: + await conn.run_sync(lambda c: ChatSessions.metadata.create_all( + c, tables=[ChatSessions.__table__, ChatLogs.__table__])) + factory = async_sessionmaker(self.engine, class_=AsyncSession, expire_on_commit=False) + + @asynccontextmanager + async def session_cm(): + async with factory() as s: + yield s + + async def session_gen(): + async with factory() as s: + yield s + + self.a, self.b = f"NL2SQL_A_{run}", f"NL2SQL_B_{run}" + self.created_on_b = f"created-on-b-{run}" # a session this API created for agent B + self.thread = f"chat-service-thread-{run}" # an id only the chat service knows + t0 = datetime(2026, 10, 23, 12, 0) + async with factory() as s: + s.add(ChatSessions(session_id=self.created_on_b, chatbot_id=self.b, agent_id=self.b, + user_id="b-user", session_title="b", created_time=t0)) + for i, (session_id, agent, question) in enumerate([ + (self.created_on_b, self.b, "B private question"), + (self.thread, self.b, "B turn in shared thread"), + (self.thread, self.a, "A turn in shared thread")]): + s.add(ChatLogs(session_id=session_id, agent_id=agent, chat_id=f"{run}-{i}", user_id="u", + user_query=question, gpt_response=f"answer to {question}", + sql_query=f"SELECT '{agent}'", db_output=f'[["{agent}-row"]]', + created_at=t0 + timedelta(minutes=i))) + await s.commit() + + common = dict(select=select, ChatLogs=ChatLogs, ChatSessions=ChatSessions, Row=Row) + self.owners = load_function("session/db_utils/session_db_utils.py", "get_session_agent_ids", + get_async_session=session_cm, **common) + self.details = load_function("session/db_utils/session_db_utils.py", "get_session_details", + get_async_session=session_cm, **common) + chart = load_function("session/db_utils/session_db_utils.py", "_chart_for_replay") + self.replay = load_function("session/db_utils/session_db_utils.py", "get_session_data", + get_async_session=session_cm, _chart_for_replay=chart, **common) + self.past = load_function("api/v1/utils/chat_helper.py", "get_past_conversation_safe", + get_async_session=session_gen, **common) + self.guard = load_function("api/v1/utils/chat_helper.py", "ensure_session_belongs_to_agent", + HTTPException=HTTPException) + self.export_row = load_function("api/v1/utils/chat_helper.py", "get_chat_log_for_export", + get_async_session=session_gen, **common) + self.feedback = load_function("session/db_utils/session_db_utils.py", "add_user_feedback", + get_async_session=session_cm, **common) + + async def asyncTearDown(self): + from sqlalchemy import text + async with self.engine.begin() as conn: + await conn.execute(text(f'DROP SCHEMA "{self.schema}" CASCADE')) + await self.engine.dispose() + + async def guarded(self, session_id, agent_id): + stub = {"session": types.ModuleType("session"), "session.db_utils": types.ModuleType("session.db_utils"), + "session.db_utils.session_db_utils": types.SimpleNamespace(get_session_agent_ids=self.owners)} + with patch.dict(sys.modules, stub): + await self.guard(session_id, agent_id) + + async def test_another_agents_session_is_refused(self): + with self.assertRaises(HTTPException) as caught: + await self.guarded(self.created_on_b, self.a) + self.assertEqual(caught.exception.status_code, 404) + self.assertIsNone(await self.details(self.created_on_b, self.a)) + + async def test_own_new_and_chat_service_sessions_continue(self): + for session_id, agent in ((self.created_on_b, self.b), (self.thread, self.a), (self.thread, self.b), + (f"never-seen-{uuid.uuid4()}", self.a), (None, self.a)): + with self.subTest(session=session_id, agent=agent): + await self.guarded(session_id, agent) + + async def test_history_is_this_agents_turns_only(self): + a_in_thread = await self.past(self.thread, self.a) + self.assertEqual([m["content"] for m in a_in_thread], + ["A turn in shared thread", "answer to A turn in shared thread"]) + self.assertEqual(await self.past(self.created_on_b, self.a), []) + self.assertEqual([m["content"] for m in await self.past(self.created_on_b, self.b)][0], "B private question") + + async def test_replay_is_this_agents_turns_only(self): + turns = await self.replay(self.thread, agent_id=self.a, show_sql_query=True, show_db_result=True) + self.assertEqual([t["message"] for t in turns if t["role"] == "user"], ["A turn in shared thread"]) + self.assertNotIn(self.b, repr(turns)) + hidden = await self.replay(self.thread, agent_id=self.a, show_sql_query=False, show_db_result=False) + self.assertNotIn("SELECT", repr(hidden)) + self.assertNotIn(f"{self.a}-row", repr(hidden)) + + async def test_export_lookup_is_scoped_to_the_key_agent(self): + latest_for_a = await self.export_row(self.thread, None, agent_id=self.a) + self.assertEqual(latest_for_a.agent_id, self.a) + b_turn = (await self.export_row(self.thread, None, agent_id=self.b)).chat_id + self.assertIsNone(await self.export_row(self.thread, b_turn, agent_id=self.a)) + self.assertEqual((await self.export_row(self.thread, b_turn)).agent_id, self.b) # master key: unscoped + + async def test_feedback_reaches_only_this_agents_turns(self): + b_turn = (await self.export_row(self.created_on_b, None, agent_id=self.b)).chat_id + await self.feedback(self.created_on_b, b_turn, "rejected", "SELECT 'rewritten by A'", agent_id=self.a) + unchanged = await self.export_row(self.created_on_b, b_turn) + self.assertEqual((unchanged.feedback, unchanged.sql_query), (None, f"SELECT '{self.b}'")) + await self.feedback(self.created_on_b, b_turn, "accepted", None, agent_id=self.b) + self.assertEqual((await self.export_row(self.created_on_b, b_turn)).feedback, "accepted") + + +if __name__ == "__main__": + unittest.main(verbosity=2)