notebook-ai-extension/opencode_bridge/tests/test_routes.py

import json

import pytest
import tornado.httpclient


class FakeOpenCodeClient:
    """Drop-in replacement for OpenCodeClient with recording + canned responses."""

    def __init__(self):
        self.calls = []
        self.session_id = "fake-session-123"
        self.health_response = {"healthy": True, "version": "0.1.0"}
        self.providers_response = {"providers": [{"id": "anthropic", "models": [{"id": "claude"}]}]}
        # Canned session-messages list response (list_session_messages).
        # Default: one user + one assistant message, mixed text parts.
        self.messages_response = [
            {
                "info": {"role": "user", "id": "m1"},
                "parts": [{"type": "text", "text": "fix the bug"}],
            },
            {
                "info": {"role": "assistant", "id": "m2"},
                "parts": [{"type": "text", "text": "```python\nx = 1\n```"}],
            },
        ]

    @property
    def endpoint(self):
        return "http://fake-opencode"

    async def health(self):
        self.calls.append(("health",))
        return self.health_response

    async def list_providers(self):
        self.calls.append(("list_providers",))
        return self.providers_response

    async def create_session(self, title):
        self.calls.append(("create_session", title))
        return {"id": self.session_id, "title": title}

    async def send_message_async(
        self, session_id, parts, provider_id=None, model_id=None, system=None
    ):
        self.calls.append(
            ("send_message_async", session_id, parts, system, provider_id, model_id)
        )
        return None

    async def delete_session(self, session_id):
        self.calls.append(("delete_session", session_id))
        return True

    async def list_session_messages(self, session_id):
        self.calls.append(("list_session_messages", session_id))
        return self.messages_response

    async def reply_permission(self, session_id, permission_id, response):
        self.calls.append(("reply_permission", session_id, permission_id, response))
        return True

    async def reply_question(self, session_id, question_id, answer):
        self.calls.append(("reply_question", session_id, question_id, answer))
        return True

    async def stream_global_events(self):
        """Async generator — tests can monkey-patch this to control the stream."""
        if False:  # pragma: no cover - never executed, just makes this an asyncgen
            yield {}


class FakeSessionManager:
    """Drop-in replacement for SessionManager with recording."""

    def __init__(self, session_id: str = "fake-session-123") -> None:
        self._session_id = session_id
        self.calls: list = []

    async def get_or_create(self, notebook_path: str) -> str:
        self.calls.append(("get_or_create", notebook_path))
        return self._session_id

    def peek(self, notebook_path: str):
        # Mirror SessionManager.peek: return the session id without
        # creating one. None means "no session yet" (used by the no-session
        # test to short-circuit the history route).
        return self._session_id or None

    async def release(self, notebook_path: str) -> bool:
        self.calls.append(("release", notebook_path))
        return True

    def list_sessions(self) -> list:
        return [
            {"notebookPath": path, "sessionId": sid}
            for path, sid in sorted(self._sessions.items())
        ]

    def invalidate(self, notebook_path: str) -> bool:
        self.calls.append(("invalidate", notebook_path))
        return True


async def test_hello(jp_fetch):
    # When
    response = await jp_fetch("opencode-bridge", "hello")

    # Then
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload == {
            "data": (
                "Hello, world!"
                " This is the '/opencode-bridge/hello' endpoint."
                " Try visiting me in your browser!"
            ),
        }


async def test_health_handler(monkeypatch, jp_fetch):
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    response = await jp_fetch("opencode-bridge", "health")
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload["ok"] is True
    assert payload["version"] == "0.1.0"
    assert ("health",) in fake.calls


async def test_providers_handler(monkeypatch, jp_fetch):
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    response = await jp_fetch("opencode-bridge", "providers")
    assert response.code == 200
    payload = json.loads(response.body)
    assert "providers" in payload
    assert payload["providers"][0]["id"] == "anthropic"
    assert ("list_providers",) in fake.calls


async def test_edit_handler(monkeypatch, jp_fetch):
    """POST /opencode-bridge/edit fires prompt_async and returns sessionId only.

    The async path returns immediately with `{ok, sessionId, notebookPath}` —
    the actual LLM response is consumed via the /events SSE endpoint, not
    embedded in this response.
    """
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    fake_sm = FakeSessionManager(session_id="fake-session-123")
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager", lambda h: fake_sm
    )

    body = json.dumps({
        "prompt": "Add type hints",
        "context": {
            "notebookPath": "test.ipynb",
            "cellId": "cell-1",
            "language": "python",
            "cellIndex": 0,
            "totalCells": 1,
            "source": "def foo(): return 42\n",
            "previousCode": None,
            "error": None,
        },
    })
    response = await jp_fetch(
        "opencode-bridge", "edit",
        method="POST",
        body=body,
    )
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload == {
        "ok": True,
        "sessionId": "fake-session-123",
        "notebookPath": "test.ipynb",
    }
    # Async path: no `markdown` field. The LLM reply is consumed via SSE.
    assert "markdown" not in payload

    # Session manager was used (not direct create/delete on client).
    call_names = [c[0] for c in fake.calls]
    assert "create_session" not in call_names
    assert "delete_session" not in call_names
    assert "send_message_async" in call_names
    assert "send_message_sync" not in call_names  # legacy sync path is gone

    # send_message_async received the session ID from the manager.
    send_call = [c for c in fake.calls if c[0] == "send_message_async"][0]
    assert send_call[1] == "fake-session-123"
    # system prompt was passed.
    assert "你是一个代码助手" in send_call[3]

    # SessionManager.get_or_create was called with the notebook path.
    sm_call_names = [c[0] for c in fake_sm.calls]
    assert sm_call_names == ["get_or_create"]
    assert fake_sm.calls[0][1] == "test.ipynb"


async def test_edit_handler_includes_model_when_provider_and_model_given(
    monkeypatch, jp_fetch
):
    """providerId/modelId are forwarded into the async send body when both set."""
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager",
        lambda h: FakeSessionManager(),
    )

    body = json.dumps({
        "prompt": "x",
        "context": {
            "notebookPath": "n.ipynb",
            "cellId": "c",
            "language": "python",
            "cellIndex": 0,
            "totalCells": 1,
            "source": "",
            "previousCode": None,
            "error": None,
        },
        "providerId": "anthropic",
        "modelId": "claude-sonnet-4-20250514",
    })
    response = await jp_fetch("opencode-bridge", "edit", method="POST", body=body)
    assert response.code == 200

    send_call = [c for c in fake.calls if c[0] == "send_message_async"][0]
    # (send_message_async, sid, parts, system, provider_id, model_id)
    assert send_call[4] == "anthropic"
    assert send_call[5] == "claude-sonnet-4-20250514"


async def test_session_list_handler(monkeypatch, jp_fetch):
    fake_sm = FakeSessionManager()
    fake_sm._sessions = {
        "foo.ipynb": "sid-1",
        "bar.ipynb": "sid-2",
    }  # direct injection
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager", lambda h: fake_sm
    )

    response = await jp_fetch("opencode-bridge", "sessions")
    assert response.code == 200
    payload = json.loads(response.body)
    assert "sessions" in payload
    paths = {s["notebookPath"] for s in payload["sessions"]}
    assert paths == {"foo.ipynb", "bar.ipynb"}


async def test_session_messages_handler_returns_empty_when_no_session(
    monkeypatch, jp_fetch
) -> None:
    # No session registered for this notebook -> handler returns
    # {"messages": []} WITHOUT calling OpenCodeClient (peek short-circuits).
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    fake_sm = FakeSessionManager(session_id="")
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager", lambda h: fake_sm
    )

    response = await jp_fetch(
        "opencode-bridge", "session-messages",
        method="GET",
        params={"notebook": "fresh.ipynb"},
    )
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload == {"messages": []}
    # No OpenCode call was made.
    assert all(c[0] != "list_session_messages" for c in fake.calls)


async def test_session_messages_handler_projects_opencode_messages(
    monkeypatch, jp_fetch
) -> None:
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    fake_sm = FakeSessionManager(session_id="fake-session-123")
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager", lambda h: fake_sm
    )

    response = await jp_fetch(
        "opencode-bridge", "session-messages",
        method="GET",
        params={"notebook": "test.ipynb"},
    )
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload == {
        "messages": [
            {"role": "user", "content": "fix the bug"},
            {"role": "assistant", "content": "```python\nx = 1\n```"},
        ]
    }
    # The OpenCode client was called with the session id from the manager.
    assert ("list_session_messages", "fake-session-123") in fake.calls


async def test_session_messages_handler_requires_notebook(
    monkeypatch, jp_fetch
) -> None:
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)
    # jp_fetch raises HTTPClientError on 4xx; assert the handler 400s
    # (the body would contain "missing 'notebook' query parameter").
    with pytest.raises(tornado.httpclient.HTTPClientError) as exc_info:
        await jp_fetch("opencode-bridge", "session-messages", method="GET")
    assert exc_info.value.code == 400


async def test_session_release_handler(monkeypatch, jp_fetch):
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    fake_sm = FakeSessionManager()
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager", lambda h: fake_sm
    )

    response = await jp_fetch(
        "opencode-bridge", "session",
        method="DELETE",
        params={"notebook": "foo.ipynb"},
    )
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload["ok"] is True
    assert payload["notebookPath"] == "foo.ipynb"
    assert payload["deleted"] is True
    assert ("release", "foo.ipynb") in fake_sm.calls


async def test_permission_reply_forwards_response(monkeypatch, jp_fetch):
    """POST /opencode-bridge/permissions/:permId?session=:sid forwards to client."""
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    response = await jp_fetch(
        "opencode-bridge", "permissions", "perm-42",
        method="POST",
        params={"session": "ses-1"},
        body=json.dumps({"response": "once"}),
    )
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload == {
        "ok": True,
        "permissionId": "perm-42",
        "response": "once",
    }
    assert ("reply_permission", "ses-1", "perm-42", "once") in fake.calls


async def test_permission_reply_rejects_invalid_response(monkeypatch, jp_fetch):
    """response must be one of once/always/reject — 400 otherwise."""
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    with pytest.raises(tornado.httpclient.HTTPClientError) as exc_info:
        await jp_fetch(
            "opencode-bridge", "permissions", "perm-1",
            method="POST",
            params={"session": "ses-1"},
            body=json.dumps({"response": "yes"}),
        )
    assert exc_info.value.code == 400
    assert "yes" not in fake.calls and fake.calls == []


async def test_permission_reply_requires_session(monkeypatch, jp_fetch):
    """Missing ?session=... is a 400."""
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)
    with pytest.raises(tornado.httpclient.HTTPClientError) as exc_info:
        await jp_fetch(
            "opencode-bridge", "permissions", "perm-1",
            method="POST",
            body=json.dumps({"response": "once"}),
        )
    assert exc_info.value.code == 400


async def test_question_reply_forwards_answer(monkeypatch, jp_fetch):
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    response = await jp_fetch(
        "opencode-bridge", "questions", "q-7", "reply",
        method="POST",
        params={"session": "ses-2"},
        body=json.dumps({"answer": "use Python 3.12"}),
    )
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload == {"ok": True, "questionId": "q-7"}
    assert ("reply_question", "ses-2", "q-7", "use Python 3.12") in fake.calls


async def test_question_reply_rejects_empty_answer(monkeypatch, jp_fetch):
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)
    with pytest.raises(tornado.httpclient.HTTPClientError) as exc_info:
        await jp_fetch(
            "opencode-bridge", "questions", "q-1", "reply",
            method="POST",
            params={"session": "ses-1"},
            body=json.dumps({"answer": "   "}),
        )
    assert exc_info.value.code == 400
    assert fake.calls == []


async def test_global_event_handler_forwards_all_events_unfiltered(
    monkeypatch, jp_fetch
):
    """Server-side SSE proxy must NOT filter by session — it forwards
    everything so the client can decide. (Previously a too-eager server
    filter silently dropped events the client would have accepted.)

    We mock stream_global_events to push three events with different
    sessionID shapes (one matching, one not, one missing). The proxy
    should pass all three through as data: lines.
    """
    fake = FakeOpenCodeClient()

    async def fake_stream():
        yield {"type": "session.next.text.delta", "properties": {"sessionID": "ses-1", "delta": "A"}}
        yield {"type": "session.next.text.delta", "properties": {"sessionID": "ses-OTHER", "delta": "B"}}
        yield {"type": "server.connected", "properties": {}}

    fake.stream_global_events = fake_stream
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    # Use ?session=ses-1 to confirm the param is accepted but doesn't filter.
    response = await jp_fetch(
        "opencode-bridge", "events",
        params={"session": "ses-1"},
    )
    assert response.code == 200
    body = response.body.decode("utf-8")
    # All three events should appear in the body, regardless of sessionID.
    assert '"delta": "A"' in body
    assert '"delta": "B"' in body
    assert '"server.connected"' in body
    # The body should use the SSE format: each event as a `data:` line.
    assert body.count("data: ") == 3