feat: improved compatability layer and responses api handing

tests/api/v1/test_routes_extended.py

@@ -181,6 +181,11 @@ class TestListModelsExtra:
 
 @contextmanager
 def _patch_v1_db(conn):
+    from application.storage.db.repositories.agents import AgentsRepository
+
+    if AgentsRepository(conn).find_by_key("x") is None:
+        AgentsRepository(conn).create("u-test", "Test Agent", "published", key="x")
+
     @contextmanager
     def _yield():
         yield conn
@@ -286,6 +291,217 @@ class TestChatCompletionsHappyPath:
         assert resp.status_code == 400
         assert resp.get_json()["error"]["type"] == "invalid_request"
 
+    def test_invalid_api_key_returns_401_before_processing(self, pg_conn):
+        app = _build_app()
+        with _patch_v1_db(pg_conn):
+            with app.test_client() as c:
+                resp = c.post(
+                    "/v1/chat/completions",
+                    headers={"Authorization": "Bearer invalid"},
+                    json={"messages": [{"role": "user", "content": "Hi"}]},
+                )
+        assert resp.status_code == 401
+
+    def test_model_field_is_accepted_and_ignored(self, pg_conn):
+        app = _build_app()
+        fake_processor = MagicMock()
+        fake_processor.decoded_token = {"sub": "u-test"}
+        fake_processor.conversation_id = None
+        fake_processor.agent_config = {}
+        fake_processor.agent_id = None
+        fake_processor.model_id = "m"
+        fake_processor.build_agent.return_value = MagicMock()
+        fake_helper = MagicMock()
+        fake_helper.check_usage.return_value = None
+        fake_helper.process_response_stream.return_value = {
+            "error": None,
+            "conversation_id": "conv-1",
+            "answer": "ok",
+            "sources": [],
+            "tool_calls": [],
+            "thought": "",
+            "extra": {},
+        }
+        with _patch_v1_db(pg_conn), patch(
+            "application.api.v1.routes.StreamProcessor",
+            return_value=fake_processor,
+        ), patch(
+            "application.api.v1.routes._V1AnswerHelper",
+            return_value=fake_helper,
+        ):
+            with app.test_client() as c:
+                resp = c.post(
+                    "/v1/chat/completions",
+                    headers={"Authorization": "Bearer x"},
+                    json={
+                        "model": "another-agent",
+                        "messages": [{"role": "user", "content": "Hi"}],
+                    },
+                )
+        assert resp.status_code == 200
+
+    def test_null_n_is_treated_as_default(self, pg_conn):
+        app = _build_app()
+        with _patch_v1_db(pg_conn), patch(
+            "application.api.v1.routes.translate_request",
+            side_effect=ValueError("reached translator"),
+        ):
+            with app.test_client() as c:
+                resp = c.post(
+                    "/v1/chat/completions",
+                    headers={"Authorization": "Bearer x"},
+                    json={
+                        "n": None,
+                        "messages": [{"role": "user", "content": "Hi"}],
+                    },
+                )
+        assert resp.status_code == 400
+        assert resp.get_json()["error"]["message"] == "Failed to process request"
+
+    def test_non_object_stream_options_returns_400(self, pg_conn):
+        app = _build_app()
+        with _patch_v1_db(pg_conn):
+            with app.test_client() as c:
+                resp = c.post(
+                    "/v1/chat/completions",
+                    headers={"Authorization": "Bearer x"},
+                    json={
+                        "stream": True,
+                        "stream_options": True,
+                        "messages": [{"role": "user", "content": "Hi"}],
+                    },
+                )
+        assert resp.status_code == 400
+        assert resp.get_json()["error"]["type"] == "invalid_request_error"
+
+    def test_conversation_id_must_belong_to_authenticated_agent(self, pg_conn):
+        from application.storage.db.repositories.agents import AgentsRepository
+        from application.storage.db.repositories.conversations import (
+            ConversationsRepository,
+        )
+
+        agents = AgentsRepository(pg_conn)
+        first = agents.create("same-owner", "First", "published", key="agent-a")
+        agents.create("same-owner", "Second", "published", key="agent-b")
+        conversation = ConversationsRepository(pg_conn).create(
+            "same-owner", "private", agent_id=str(first["id"])
+        )
+        app = _build_app()
+        with _patch_v1_db(pg_conn):
+            with app.test_client() as c:
+                resp = c.post(
+                    "/v1/chat/completions",
+                    headers={"Authorization": "Bearer agent-b"},
+                    json={
+                        "model": "ignored",
+                        "conversation_id": str(conversation["id"]),
+                        "messages": [{"role": "user", "content": "secret?"}],
+                    },
+                )
+        assert resp.status_code == 400
+        assert resp.get_json()["error"]["code"] == "conversation_not_found"
+
+    def test_stale_session_mapping_is_deleted_and_replaced(self, pg_conn):
+        app = _build_app()
+        fake_processor = MagicMock()
+        fake_processor.decoded_token = {"sub": "u-test"}
+        fake_processor.conversation_id = None
+        fake_processor.agent_config = {}
+        fake_processor.agent_id = None
+        fake_processor.model_id = "m"
+        fake_processor.build_agent.return_value = MagicMock()
+        fake_helper = MagicMock()
+        fake_helper.check_usage.return_value = None
+        fake_helper.process_response_stream.return_value = {
+            "error": None,
+            "conversation_id": "new-conversation",
+            "answer": "ok",
+            "sources": [],
+            "tool_calls": [],
+            "thought": "",
+            "extra": {},
+        }
+        with _patch_v1_db(pg_conn), patch(
+            "application.api.v1.routes.load_conversation",
+            return_value="deleted-conversation",
+        ), patch(
+            "application.api.v1.routes._conversation_belongs_to_agent",
+            return_value=False,
+        ), patch(
+            "application.api.v1.routes.delete_conversation",
+        ) as delete_mapping, patch(
+            "application.api.v1.routes.StreamProcessor",
+            return_value=fake_processor,
+        ), patch(
+            "application.api.v1.routes._V1AnswerHelper",
+            return_value=fake_helper,
+        ):
+            with app.test_client() as c:
+                resp = c.post(
+                    "/v1/chat/completions",
+                    headers={
+                        "Authorization": "Bearer x",
+                        "X-DocsGPT-Session-ID": "session",
+                    },
+                    json={"messages": [{"role": "user", "content": "Hi"}]},
+                )
+        assert resp.status_code == 200
+        delete_mapping.assert_called_once()
+
+    def test_duplicate_resume_returns_conflict(self, pg_conn):
+        from application.api.answer.services.continuation_service import (
+            ResumeInProgressError,
+        )
+
+        app = _build_app()
+        fake_processor = MagicMock()
+        fake_processor.decoded_token = {"sub": "u-test"}
+        with _patch_v1_db(pg_conn), patch(
+            "application.api.v1.routes.translate_request",
+            return_value={
+                "conversation_id": "conv-1",
+                "tool_actions": [{"call_id": "call-1", "result": "ok"}],
+                "messages": [],
+            },
+        ), patch(
+            "application.api.v1.routes._conversation_belongs_to_agent",
+            return_value=True,
+        ), patch(
+            "application.api.v1.routes.StreamProcessor",
+            return_value=fake_processor,
+        ), patch(
+            "application.api.v1.routes.ContinuationService.claim_state",
+            side_effect=ResumeInProgressError("Resume already in progress"),
+        ):
+            with app.test_client() as c:
+                resp = c.post(
+                    "/v1/chat/completions",
+                    headers={"Authorization": "Bearer x"},
+                    json={"messages": [{"role": "tool", "content": "ok"}]},
+                )
+        assert resp.status_code == 409
+        assert resp.get_json()["error"]["code"] == "resume_in_progress"
+
+    def test_unsupported_n_and_logprobs_are_explicit(self, pg_conn):
+        app = _build_app()
+        with _patch_v1_db(pg_conn):
+            with app.test_client() as c:
+                n_response = c.post(
+                    "/v1/chat/completions",
+                    headers={"Authorization": "Bearer x"},
+                    json={"n": 2, "messages": [{"role": "user", "content": "Hi"}]},
+                )
+                logprobs_response = c.post(
+                    "/v1/chat/completions",
+                    headers={"Authorization": "Bearer x"},
+                    json={
+                        "logprobs": True,
+                        "messages": [{"role": "user", "content": "Hi"}],
+                    },
+                )
+        assert n_response.status_code == 400
+        assert logprobs_response.status_code == 400
+
     def test_translate_error_returns_400(self, pg_conn):
         app = _build_app()
         with _patch_v1_db(pg_conn), patch(
@@ -542,7 +758,7 @@ class TestChatCompletionsHappyPath:
         translated_chunks: list = []
 
         def _fake_translate_stream_event(
-            event_data, completion_id, model_name, strip_reasoning_leak=False
+            event_data, completion_id, model_name, strip_reasoning_leak=False, state=None
         ):
             translated_chunks.append(event_data)
             return ['data: x\n\n']
@@ -583,5 +799,3 @@ class TestChatCompletionsHappyPath:
         assert "message_id" not in types_translated
         assert "id" in types_translated
         assert "answer" in types_translated
-
-

tests/api/v1/test_session_store.py

@@ -0,0 +1,74 @@
+"""Tests for privacy-preserving v1 client session correlation."""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from application.api.v1.session_store import (
+    delete_conversation,
+    identify_session,
+    load_conversation,
+    save_conversation,
+)
+
+
+@pytest.mark.unit
+def test_session_key_is_prompt_and_agent_scoped():
+    headers = {"X-Session-ID": "raw-secret-session"}
+    main = {"messages": [{"role": "system", "content": "coding"}]}
+    title = {"messages": [{"role": "system", "content": "title"}]}
+    one = identify_session(headers, main, "agent-1")
+    assert one != identify_session(headers, title, "agent-1")
+    assert one != identify_session(headers, main, "agent-2")
+    assert "raw-secret-session" not in one.key
+
+
+@pytest.mark.unit
+def test_explicit_docsgpt_session_precedes_generic_header():
+    data = {
+        "docsgpt": {"session_id": "explicit"},
+        "messages": [{"role": "user", "content": "hi"}],
+    }
+    explicit = identify_session({"X-Session-ID": "generic"}, data, "agent")
+    expected = identify_session(
+        {"X-DocsGPT-Session-ID": "explicit"},
+        {"messages": data["messages"]},
+        "agent",
+    )
+    assert explicit == expected
+
+
+@pytest.mark.unit
+def test_openai_user_field_is_not_a_chat_session():
+    data = {
+        "user": "stable-end-user-id",
+        "messages": [{"role": "user", "content": "hi"}],
+    }
+    assert identify_session({}, data, "agent") is None
+
+
+@pytest.mark.unit
+def test_save_and_load_use_bounded_redis_ttl(monkeypatch):
+    redis = MagicMock()
+    redis.get.return_value = b"conv-1"
+    monkeypatch.setattr(
+        "application.api.v1.session_store.get_redis_instance", lambda: redis
+    )
+    monkeypatch.setattr(
+        "application.api.v1.session_store.settings.V1_SESSION_TTL_SECONDS", 123
+    )
+    session = identify_session({"X-Session-ID": "s"}, {"messages": []}, "a")
+    save_conversation(session, "conv-1")
+    redis.set.assert_called_once_with(session.key, "conv-1", ex=123)
+    assert load_conversation(session) == "conv-1"
+
+
+@pytest.mark.unit
+def test_delete_conversation_invalidates_stale_mapping(monkeypatch):
+    redis = MagicMock()
+    monkeypatch.setattr(
+        "application.api.v1.session_store.get_redis_instance", lambda: redis
+    )
+    session = identify_session({"X-Session-ID": "s"}, {"messages": []}, "a")
+    delete_conversation(session)
+    redis.delete.assert_called_once_with(session.key)

tests/api/v1/test_v1_continuation_ownership.py

@@ -1,25 +1,9 @@
-"""Foreign-``conversation_id`` ownership gate on the ``/v1`` continuation path.
+"""Owner-and-agent ``conversation_id`` gates on the ``/v1`` continuation path.
 
-A ``/v1`` tool continuation persists its final answer into a client-supplied
-``conversation_id``. Ownership IS enforced deep in the stack:
-
-    route → ``build_continuation_from_messages`` → ``build_agent("")``
-          → ``initialize`` → ``_load_conversation_history``
-          → ``get_conversation(conversation_id, owner)``  ← owner-scoped read
-          → ``None`` for a foreign conversation
-          → ``ValueError("Conversation not found or unauthorized")``
-          → route's ``except ValueError`` → HTTP 400
-
-…but the existing E2E suite mocks ``build_agent``, so this security gate is
-never actually exercised. This test seeds a conversation owned by user A and an
-agent owned by user B, then POSTs a continuation with B's api_key targeting A's
-conversation. The ownership read runs against the **real** DB (not mocked): we
-only short-circuit ``StreamProcessor.initialize`` to the
-``_load_conversation_history`` step so the test doesn't have to stand up the
-full agent/model/source/retriever machinery to reach the gate.
-
-Asserted: the route returns 400 AND no assistant turn is appended to A's
-conversation (the foreign write is rejected before persistence).
+The route validates conversation metadata before initializing the stream
+processor. These tests use the real database to verify that a foreign-owner or
+same-owner/different-agent conversation cannot receive a continuation and that
+the rejected request never appends an assistant turn.
 """
 
 from __future__ import annotations
@@ -58,10 +42,12 @@ CLIENT_TOOL = {
 }
 
 
-def _seed_conversation(conn, user_id: str, name: str) -> str:
+def _seed_conversation(
+    conn, user_id: str, name: str, agent_id: str | None = None
+) -> str:
     """Create a conversation owned by ``user_id`` with one finalized turn."""
     repo = ConversationsRepository(conn)
-    conv = repo.create(user_id, name)
+    conv = repo.create(user_id, name, agent_id=agent_id)
     conv_id = str(conv["id"])
     repo.reserve_message(
         conv_id,
@@ -72,8 +58,10 @@ def _seed_conversation(conn, user_id: str, name: str) -> str:
     return conv_id
 
 
-def _seed_agent(conn, user_id: str, key: str) -> None:
-    AgentsRepository(conn).create(user_id, "Weather Agent", "published", key=key)
+def _seed_agent(conn, user_id: str, key: str) -> dict:
+    return AgentsRepository(conn).create(
+        user_id, "Weather Agent", "published", key=key
+    )
 
 
 def _build_app() -> Flask:
@@ -141,9 +129,11 @@ class TestForeignConversationOwnership:
             _seed_user(conn, user_a)
             _seed_user(conn, user_b)
             # Agent owned by B → the v1 route resolves decoded_token={"sub": B}.
-            _seed_agent(conn, user_b, api_key)
+            agent = _seed_agent(conn, user_b, api_key)
             # Conversation owned by A, with one existing complete turn.
-            conv_id = _seed_conversation(conn, user_a, "A's private chat")
+            conv_id = _seed_conversation(
+                conn, user_a, "A's private chat", agent_id=str(agent["id"])
+            )
 
         with pg_engine.connect() as conn:
             rows_before = _row_count(conn, conv_id)
@@ -173,7 +163,7 @@ class TestForeignConversationOwnership:
         # The ownership gate raised ValueError → route returns 400.
         assert resp.status_code == 400, resp.get_data(as_text=True)
         body = resp.get_json()
-        assert body["error"]["type"] == "invalid_request"
+        assert body["error"]["type"] == "invalid_request_error"
 
         # No assistant turn was appended to A's conversation — the foreign
         # write was rejected before any persistence.
@@ -191,9 +181,11 @@ class TestForeignConversationOwnership:
 
         with pg_engine.begin() as conn:
             _seed_user(conn, owner)
-            _seed_agent(conn, owner, api_key)
+            agent = _seed_agent(conn, owner, api_key)
             # Conversation owned by the agent owner (the caller).
-            conv_id = _seed_conversation(conn, owner, "owner's chat")
+            conv_id = _seed_conversation(
+                conn, owner, "owner's chat", agent_id=str(agent["id"])
+            )
 
         app = _build_app()
 

tests/api/v1/test_v1_tool_pause_finalization.py

@@ -586,23 +586,14 @@ class TestNativePauseUnchanged:
 
 
 # ---------------------------------------------------------------------------
-# Guarantee 2 — coherent Option B routing: a v1 continuation carrying a
-# conversation_id rebuilds STATELESSLY via build_continuation_from_messages
-# (never resume_from_tool_actions), and threads no reserved_message_id.
+# Guarantee 2 — legacy continuations without pending server state still fall
+# back to a stateless transcript rebuild.
 # ---------------------------------------------------------------------------
 
 
 @pytest.mark.unit
 class TestV1ContinuationRoutesStatelessly:
-    """Coherent Option B: a v1 tool-result continuation is fully stateless on
-    the resume side. Because the pause finalized the prior row ``complete`` and
-    wrote NO ``pending_tool_state``, the route must rebuild from the re-POSTed
-    history via ``build_continuation_from_messages`` — **even when a
-    ``conversation_id`` is carried** — and must NOT call
-    ``resume_from_tool_actions`` (whose ``load_state`` would 400). The
-    ``_continuation`` dict therefore carries no live ``reserved_message_id`` /
-    ``request_id`` (id-reuse is inert under this design).
-    """
+    """A carried conversation id with no pending state remains compatible."""
 
     def _build_app(self) -> Flask:
         app = Flask(__name__)
@@ -667,6 +658,15 @@ class TestV1ContinuationRoutesStatelessly:
             yield MagicMock()
 
         with patch(
+            "application.api.v1.routes._lookup_agent",
+            return_value={"id": "agent-1", "name": "Agent", "user_id": "owner"},
+        ), patch(
+            "application.api.v1.routes.ContinuationService.claim_state",
+            return_value=None,
+        ), patch(
+            "application.api.v1.routes._conversation_belongs_to_agent",
+            return_value=True,
+        ), patch(
             "application.api.v1.routes.translate_request",
             side_effect=_fake_translate,
         ), patch(
@@ -708,13 +708,12 @@ class TestV1ContinuationRoutesStatelessly:
 
         continuation = captured.get("_continuation")
         assert continuation is not None
-        # id-reuse is gone from the v1 continuation dict (inert under Option B).
-        assert "reserved_message_id" not in continuation
-        assert "request_id" not in continuation
-        # The v1 path still runs in stateless-finalize mode.
-        assert captured.get("finalize_tool_pause_as_complete") is True
-        # The incoming conversation_id is the persistence target.
+        assert continuation["reserved_message_id"] is None
+        assert continuation["request_id"] is None
+        assert captured.get("finalize_tool_pause_as_complete") is False
+        # Missing durable state runs without appending a blank sibling turn.
         assert captured.get("conversation_id") == "conv-1"
+        assert captured.get("should_persist") is False
 
 
 # ---------------------------------------------------------------------------
@@ -737,6 +736,9 @@ class _FakeLLMGen:
     def gen(self, *args, **kwargs) -> str:
         return "Title"
 
+    def gen_stream(self, *args, **kwargs):
+        yield "It is 72F and sunny in SF."
+
 
 class _FakeClientToolExecutor:
     """Tool executor stub exposing ``get_tools`` for the real
@@ -912,6 +914,9 @@ class TestV1ToolRoundTripEndToEnd:
             "application.api.answer.services.stream_processor.StreamProcessor"
             ".build_agent",
             _fake_build_agent,
+        ), patch(
+            "application.api.v1.routes._conversation_belongs_to_agent",
+            return_value=True,
         ):
             with app.test_client() as c:
                 resp1 = _post_chat(
@@ -919,7 +924,10 @@ class TestV1ToolRoundTripEndToEnd:
                     {
                         "messages": [{"role": "user", "content": "weather in SF?"}],
                         "tools": [self.CLIENT_TOOL],
-                        "docsgpt": {"save_conversation": True},
+                        "docsgpt": {
+                            "save_conversation": True,
+                            "session_id": "tool-session",
+                        },
                     },
                     api_key,
                 )
@@ -936,19 +944,21 @@ class TestV1ToolRoundTripEndToEnd:
         conv_id = body1.get("docsgpt", {}).get("conversation_id")
         assert conv_id
 
-        # Reserved row finalized ``complete`` (with tool_calls); no
-        # ``pending_tool_state`` and no non-terminal row.
+        # The pause remains resumable and owns the original WAL placeholder.
         with pg_engine.connect() as conn:
             statuses = _row_statuses(conn, conv_id)
-            assert statuses == ["complete"]
-            assert _pending_tool_state_count(conn, conv_id) == 0
-            assert self._count_non_terminal(conn, user_id) == 0
+            assert statuses[0] in ("pending", "streaming")
+            assert _pending_tool_state_count(conn, conv_id) == 1
+            assert self._count_non_terminal(conn, user_id) == 1
 
         # ---- POST #2: tool result + conversation_id -> agent answers ----
         with _wire_v1_route_db(pg_engine, monkeypatch), patch(
             "application.api.answer.services.stream_processor.StreamProcessor"
             ".build_agent",
             _fake_build_agent,
+        ), patch(
+            "application.api.v1.routes._conversation_belongs_to_agent",
+            return_value=True,
         ):
             with app.test_client() as c:
                 resp2 = _post_chat(
@@ -977,7 +987,10 @@ class TestV1ToolRoundTripEndToEnd:
                             },
                         ],
                         "conversation_id": conv_id,
-                        "docsgpt": {"save_conversation": True},
+                        "docsgpt": {
+                            "save_conversation": True,
+                            "session_id": "tool-session",
+                        },
                     },
                     api_key,
                 )
@@ -996,12 +1009,12 @@ class TestV1ToolRoundTripEndToEnd:
             from sqlalchemy import text
 
             # The answer is a NEW terminal turn appended to the same
-            # conversation: the original tool-call turn + the answer turn.
+            # conversation: the original WAL placeholder is finalized.
             statuses = _row_statuses(conn, conv_id)
-            assert statuses == ["complete", "complete"]
+            assert statuses == ["complete"]
             # Nothing left non-terminal anywhere for this user / conversation.
             assert self._count_non_terminal(conn, user_id) == 0
-            # The appended turn carries the assistant answer.
+            # The finalized turn carries the assistant answer.
             answer_rows = conn.execute(
                 text(
                     "SELECT response FROM conversation_messages "
@@ -1017,4 +1030,3 @@ class TestV1ToolRoundTripEndToEnd:
                 {"u": user_id},
             ).scalar()
             assert conv_count == 1
-