turnstone/tests/test_server_attachments_on_create.py

"""Tests for the multipart variant of POST /v1/api/workstreams/new.

Exercises:
- The pure helpers ``validate_and_save_uploaded_files`` (stages to the
  per-node buffer) and ``resolve_staged_attachments`` (peeks them back).
- The full create endpoint via TestClient with a FakeSession factory so
  the initial-message dispatch thread runs end-to-end without an LLM.
"""

from __future__ import annotations

import json
import queue
import threading
import time

import pytest
from starlette.testclient import TestClient

# Magic-byte-valid 1x1 PNG (matches the fixture in test_server_attachments_endpoints.py)
PNG_1x1 = (
    b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
    b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\rIDATx\x9cc\xfc\xcf"
    b"\xc0\xc0\xc0\x00\x00\x00\x05\x00\x01\xa5\xf6E@\x00\x00\x00\x00IEND\xaeB`\x82"
)

_TEST_JWT_SECRET = "test-jwt-secret-minimum-32-chars!"


def _make_jwt(user_id: str) -> str:
    from turnstone.core.auth import JWT_AUD_SERVER, create_jwt

    return create_jwt(
        user_id=user_id,
        scopes=frozenset({"read", "write"}),
        source="test",
        secret=_TEST_JWT_SECRET,
        audience=JWT_AUD_SERVER,
        # ``workstreams.create`` is now a real gate on POST /workstreams/new
        # — see PR adding 057_role_permission_overrides.  Embed the perm so
        # the multipart-create flow under test stays exercising the create
        # path and not the new 403.
        permissions=frozenset({"workstreams.create"}),
    )


def _auth(user: str) -> dict[str, str]:
    return {"Authorization": f"Bearer {_make_jwt(user)}"}


# ---------------------------------------------------------------------------
# Pure helper tests
# ---------------------------------------------------------------------------


class TestValidateAndSaveUploadedFiles:
    def test_stages_image_and_text_to_buffer(self, tmp_path):
        import hashlib

        from turnstone.core.attachment_buffer import get_attachment_buffer
        from turnstone.core.attachments import (
            validate_and_save_uploaded_files as _validate_and_save_uploaded_files,
        )

        buf = get_attachment_buffer()
        buf.clear()
        try:
            files = [
                ("hi.png", "image/png", PNG_1x1),
                ("notes.md", "text/markdown", b"# Hello\n"),
            ]
            ids, err = _validate_and_save_uploaded_files(files, "ws-X", "userA")
            assert err is None
            assert len(ids) == 2
            # Ids are content hashes (content-addressed staging).
            assert ids[0] == hashlib.sha256(PNG_1x1).hexdigest()
            staged = buf.list_for(ws_id="ws-X", user_id="userA")
            assert len(staged) == 2
            assert {s.kind for s in staged} == {"image", "text"}
        finally:
            buf.clear()

    def test_rejects_oversized_image(self, tmp_path):
        from turnstone.core.attachments import IMAGE_SIZE_CAP
        from turnstone.core.attachments import (
            validate_and_save_uploaded_files as _validate_and_save_uploaded_files,
        )
        from turnstone.core.storage import init_storage, reset_storage

        reset_storage()
        init_storage("sqlite", path=str(tmp_path / "t.db"), run_migrations=False)
        try:
            # Magic-byte-valid PNG header padded past the cap.
            oversized = PNG_1x1 + b"\x00" * (IMAGE_SIZE_CAP + 1)
            files = [("big.png", "image/png", oversized)]
            ids, err = _validate_and_save_uploaded_files(files, "ws-X", "userA")
            assert err is not None
            assert err.status_code == 413
            assert ids == []
        finally:
            reset_storage()

    def test_rejects_unsupported_text(self, tmp_path):
        from turnstone.core.attachments import (
            validate_and_save_uploaded_files as _validate_and_save_uploaded_files,
        )
        from turnstone.core.storage import init_storage, reset_storage

        reset_storage()
        init_storage("sqlite", path=str(tmp_path / "t.db"), run_migrations=False)
        try:
            # No image magic, MIME isn't text/*, extension not allowlisted
            files = [("evil.bin", "application/octet-stream", b"\x00\x01\x02")]
            ids, err = _validate_and_save_uploaded_files(files, "ws-X", "userA")
            assert err is not None
            assert err.status_code == 400
            assert ids == []
        finally:
            reset_storage()


# (The per-user pending-upload cap was removed with the content-addressing
# cutover; ``validate_and_save_uploaded_files`` no longer 409s — the buffer's
# own size/TTL ceilings bound a flood.)


class TestResolveStagedAttachments:
    def _stage(self, ws_id, user_id, filename, mime, kind, content):
        from turnstone.core.attachment_buffer import get_attachment_buffer

        return get_attachment_buffer().stage(
            ws_id=ws_id,
            user_id=user_id,
            filename=filename,
            mime_type=mime,
            kind=kind,
            content=content,
        )

    def test_resolves_staged_to_attachments(self):
        from turnstone.core.attachment_buffer import get_attachment_buffer
        from turnstone.core.attachments import Attachment
        from turnstone.core.attachments import (
            resolve_staged_attachments as _resolve_staged_attachments,
        )

        buf = get_attachment_buffer()
        buf.clear()
        try:
            a1 = self._stage("ws-X", "userA", "a.txt", "text/plain", "text", b"hello")
            a2 = self._stage("ws-X", "userA", "b.png", "image/png", "image", PNG_1x1)
            resolved, taken, dropped = _resolve_staged_attachments(
                [a1.attachment_id, a2.attachment_id], "ws-X", "userA"
            )
            assert taken == [a1.attachment_id, a2.attachment_id]
            assert dropped == []
            assert len(resolved) == 2
            assert all(isinstance(a, Attachment) for a in resolved)
            assert [a.kind for a in resolved] == ["text", "image"]
        finally:
            buf.clear()

    def test_unknown_id_is_dropped(self):
        from turnstone.core.attachment_buffer import get_attachment_buffer
        from turnstone.core.attachments import (
            resolve_staged_attachments as _resolve_staged_attachments,
        )

        buf = get_attachment_buffer()
        buf.clear()
        try:
            a1 = self._stage("ws-X", "userA", "a.txt", "text/plain", "text", b"hello")
            resolved, taken, dropped = _resolve_staged_attachments(
                [a1.attachment_id, "never-staged"], "ws-X", "userA"
            )
            assert taken == [a1.attachment_id]
            assert dropped == ["never-staged"]
            assert len(resolved) == 1
        finally:
            buf.clear()

    def test_cross_user_id_not_resolved(self):
        # A staged id scoped to another user (same ws) must not resolve.
        from turnstone.core.attachment_buffer import get_attachment_buffer
        from turnstone.core.attachments import (
            resolve_staged_attachments as _resolve_staged_attachments,
        )

        buf = get_attachment_buffer()
        buf.clear()
        try:
            other = self._stage("ws-X", "userB", "a.txt", "text/plain", "text", b"hello")
            resolved, taken, dropped = _resolve_staged_attachments(
                [other.attachment_id], "ws-X", "userA"
            )
            assert resolved == []
            assert dropped == [other.attachment_id]
        finally:
            buf.clear()


# ---------------------------------------------------------------------------
# End-to-end create endpoint tests (multipart variant)
# ---------------------------------------------------------------------------


class _FakeSession:
    """Minimal stand-in: records send() invocations from the dispatch thread.

    Knows its own ``ws_id`` and ``user_id`` so it can faithfully simulate the
    real ChatSession's attachment-consume step against storage.  Tests then
    assert that pending attachments are gone after dispatch.
    """

    def __init__(self, ws_id: str = "", user_id: str = ""):
        self.ws_id = ws_id
        self.user_id = user_id
        self.model = "test-model"
        self.model_alias = "test-model"
        self.messages = []
        self.sends: list[tuple[str, list, str | None]] = []
        self._lock = threading.Lock()
        self._cancel_event = threading.Event()
        self.notify_targets = ""
        self._notify_on_complete = "[]"

    def send(self, text, attachments=None, send_id=None):
        with self._lock:
            self.sends.append((text, list(attachments or []), send_id))
            # Simulate the real ChatSession commit: write each attachment
            # content-addressed, record the ref-list on a (synthetic) message
            # row, and drain the staged handles from the buffer — so callers
            # can assert the lifecycle landed.
            if attachments and self.ws_id and self.user_id:
                from turnstone.core.attachment_buffer import get_attachment_buffer
                from turnstone.core.memory import (
                    save_attachment,
                    save_message,
                    set_message_attachments,
                )

                mid = save_message(self.ws_id, "user", text)
                buf = get_attachment_buffer()
                ref_ids = []
                for a in attachments:
                    save_attachment(
                        a.attachment_id,
                        a.filename,
                        a.mime_type,
                        len(a.content),
                        a.kind,
                        a.content,
                    )
                    ref_ids.append(a.attachment_id)
                    buf.discard(a.attachment_id, ws_id=self.ws_id, user_id=self.user_id)
                set_message_attachments(self.ws_id, mid, ref_ids)

    # Methods the create handler may call but we don't care about
    def set_watch_runner(self, *_a, **_kw):
        pass

    def queue_message(self, *_a, **_kw):
        return ("", "normal", "msg-x")

    def request_title_refresh(self, *_a, **_kw):
        pass

    def resume(self, *_a, **_kw):
        return False


class _FakeUI:
    def __init__(self, ws_id="", user_id=""):
        self.ws_id = ws_id
        self._user_id = user_id
        self.auto_approve = False
        self.auto_approve_tools: set[str] = set()
        self.events: list[dict] = []
        self._enqueued: list[dict] = []

    def _enqueue(self, ev):
        self._enqueued.append(ev)

    def on_stream_end(self):
        pass

    def on_state_change(self, state):
        self.events.append({"type": "state_change", "state": state})

    def on_error(self, msg):
        self.events.append({"type": "error", "message": msg})


@pytest.fixture
def app_client(tmp_path, monkeypatch):
    """End-to-end app with a fake session factory + SessionManager."""
    from turnstone.core.adapters.interactive_adapter import InteractiveAdapter
    from turnstone.core.metrics import MetricsCollector
    from turnstone.core.session_manager import SessionManager
    from turnstone.core.storage import get_storage, init_storage, reset_storage
    from turnstone.server import WebUI, create_app

    reset_storage()
    init_storage("sqlite", path=str(tmp_path / "t.db"), run_migrations=False)

    metrics = MetricsCollector()
    metrics.model = "test-model"
    monkeypatch.setattr("turnstone.server._metrics", metrics)
    # Replace WebUI with our fake so the create handler's isinstance check passes.
    monkeypatch.setattr("turnstone.server.WebUI", _FakeUI)

    fake_sessions: list[_FakeSession] = []

    def _factory(ui, _model, ws_id, **_kw):
        # The user_id rides on the WebUI factory closure; pull it off the
        # ui instance so the FakeSession's consume step uses the right scope.
        user_id = getattr(ui, "_user_id", "")
        s = _FakeSession(ws_id=ws_id, user_id=user_id)
        fake_sessions.append(s)
        return s

    gq: queue.Queue[dict] = queue.Queue(maxsize=1000)
    WebUI._global_queue = gq
    adapter = InteractiveAdapter(
        global_queue=gq,
        ui_factory=lambda ws: _FakeUI(
            ws_id=ws.id,
            user_id=ws.user_id,
        ),
        session_factory=_factory,
    )
    mgr = SessionManager(
        adapter, storage=get_storage(), max_active=10, node_id="node-test", event_emitter=adapter
    )

    app = create_app(
        workstreams=mgr,
        global_queue=gq,
        global_listeners=[],
        global_listeners_lock=threading.Lock(),
        skip_permissions=False,
        jwt_secret=_TEST_JWT_SECRET,
    )
    # Pending uploads live in the process-global per-node buffer; clear it so
    # staged uploads can't leak across tests.
    from turnstone.core.attachment_buffer import get_attachment_buffer

    get_attachment_buffer().clear()

    client = TestClient(app, raise_server_exceptions=False)
    try:
        yield client, fake_sessions, gq
    finally:
        client.close()
        get_attachment_buffer().clear()
        reset_storage()


class TestCreateMultipart:
    def test_create_with_image_and_initial_message(self, app_client):
        import hashlib

        from turnstone.core.attachment_buffer import get_attachment_buffer
        from turnstone.core.memory import attachment_referenced_in_ws, get_attachment

        client, sessions, _gq = app_client
        meta = {"name": "demo", "initial_message": "describe this image"}
        resp = client.post(
            "/v1/api/workstreams/new",
            data={"meta": json.dumps(meta)},
            files=[("file", ("tiny.png", PNG_1x1, "image/png"))],
            headers=_auth("userA"),
        )
        assert resp.status_code == 200, resp.text
        data = resp.json()
        ws_id = data["ws_id"]
        assert ws_id
        assert len(data["attachment_ids"]) == 1
        aid = data["attachment_ids"][0]
        assert aid == hashlib.sha256(PNG_1x1).hexdigest()  # content-addressed id

        # Wait briefly for the dispatch thread
        deadline = time.time() + 2.0
        while time.time() < deadline and not sessions:
            time.sleep(0.02)
        deadline = time.time() + 2.0
        while time.time() < deadline and not sessions[0].sends:
            time.sleep(0.02)
        assert sessions
        assert sessions[0].sends, "session.send was not invoked"
        text, atts, send_id = sessions[0].sends[0]
        assert text == "describe this image"
        assert send_id  # tracking token threaded through
        assert len(atts) == 1
        assert atts[0].kind == "image"

        # Lifecycle: the FakeSession commit wrote the blob content-addressed +
        # recorded the ref-list + drained the buffer.  Poll for the worker.
        deadline = time.time() + 2.0
        while time.time() < deadline:
            if get_attachment(aid) is not None and attachment_referenced_in_ws(aid, ws_id):
                break
            time.sleep(0.02)
        row = get_attachment(aid)
        assert row is not None and row["refcount"] >= 1
        assert attachment_referenced_in_ws(aid, ws_id) is True
        # Drained from the buffer post-commit.
        assert get_attachment_buffer().get(aid, ws_id=ws_id, user_id="userA") is None

    def test_create_drains_staged_synchronously(self, app_client, monkeypatch):
        """A create-time attachment dispatched on the first turn must be drained
        by the create handler itself, not only by the async dispatch worker —
        else the freshly-opened pane's rehydrate races the worker's write-time
        drain and paints the image as a still-pending composer chip.

        Neuter the worker's drain (stub ``send``) so only the synchronous
        post-install drain can clear the buffer, then assert it's empty right
        after the response with NO polling."""
        from turnstone.core.attachment_buffer import get_attachment_buffer

        client, _sessions, _gq = app_client
        monkeypatch.setattr(_FakeSession, "send", lambda self, *a, **k: None)
        meta = {"name": "demo", "initial_message": "describe this image"}
        resp = client.post(
            "/v1/api/workstreams/new",
            data={"meta": json.dumps(meta)},
            files=[("file", ("tiny.png", PNG_1x1, "image/png"))],
            headers=_auth("userA"),
        )
        assert resp.status_code == 200, resp.text
        ws_id = resp.json()["ws_id"]
        aid = resp.json()["attachment_ids"][0]
        # No poll: the create handler drained it before returning, so the new
        # pane's rehydrate can't observe it as still-staged.
        assert get_attachment_buffer().get(aid, ws_id=ws_id, user_id="userA") is None

    def test_create_with_attachments_no_initial_message_keeps_staged(self, app_client):
        import hashlib

        from turnstone.core.attachment_buffer import get_attachment_buffer

        client, _, _gq = app_client
        meta = {"name": "stash"}
        resp = client.post(
            "/v1/api/workstreams/new",
            data={"meta": json.dumps(meta)},
            files=[("file", ("notes.md", b"# hello\n", "text/markdown"))],
            headers=_auth("userA"),
        )
        assert resp.status_code == 200, resp.text
        data = resp.json()
        ws_id = data["ws_id"]
        assert data["attachment_ids"] == [hashlib.sha256(b"# hello\n").hexdigest()]
        # No initial_message → no send → the upload stays staged in the buffer.
        staged = get_attachment_buffer().list_for(ws_id=ws_id, user_id="userA")
        assert len(staged) == 1
        assert staged[0].filename == "notes.md"

    def test_create_rejects_oversized_image_and_rolls_back(self, app_client):
        from turnstone.core.attachments import IMAGE_SIZE_CAP

        client, _, gq = app_client
        oversized = PNG_1x1 + b"\x00" * (IMAGE_SIZE_CAP + 1)
        meta = {"name": "fails"}
        resp = client.post(
            "/v1/api/workstreams/new",
            data={"meta": json.dumps(meta)},
            files=[("file", ("big.png", oversized, "image/png"))],
            headers=_auth("userA"),
        )
        assert resp.status_code == 413
        # Regression: ws_created must NOT have been emitted for a request
        # that's about to be rejected.  Otherwise SSE consumers see a
        # phantom workstream flash on dashboards.
        events: list[dict] = []
        while not gq.empty():
            events.append(gq.get_nowait())
        kinds = {e.get("type") for e in events}
        assert "ws_created" not in kinds, f"phantom ws_created emitted for failed create: {events}"

    def test_create_missing_meta_returns_400(self, app_client):
        client, _, _gq = app_client
        resp = client.post(
            "/v1/api/workstreams/new",
            files=[("file", ("notes.md", b"hello", "text/markdown"))],
            headers=_auth("userA"),
        )
        assert resp.status_code == 400

    def test_create_invalid_meta_json_returns_400(self, app_client):
        client, _, _gq = app_client
        resp = client.post(
            "/v1/api/workstreams/new",
            data={"meta": "{not json}"},
            files=[],
            headers=_auth("userA"),
        )
        assert resp.status_code == 400

    def test_attachments_with_resume_ws_returns_400(self, app_client):
        from turnstone.core.memory import register_workstream

        client, _, _gq = app_client
        register_workstream("ws-resume-target", name="resume target")
        meta = {"name": "fork", "resume_ws": "ws-resume-target"}
        resp = client.post(
            "/v1/api/workstreams/new",
            data={"meta": json.dumps(meta)},
            files=[("file", ("notes.md", b"hello", "text/markdown"))],
            headers=_auth("userA"),
        )
        assert resp.status_code == 400


class TestCreateJsonStillWorks:
    """The JSON path must remain byte-for-byte identical (back-compat)."""

    def test_create_json_no_attachments(self, app_client):
        client, _, _gq = app_client
        resp = client.post(
            "/v1/api/workstreams/new",
            json={"name": "json-only"},
            headers=_auth("userA"),
        )
        assert resp.status_code == 200, resp.text
        data = resp.json()
        assert data["ws_id"]
        # New optional field, but always emitted (empty list when absent)
        assert data["attachment_ids"] == []