notebook-ai-extension/opencode_bridge/routes.py

import json

import logging

from jupyter_server.base.handlers import APIHandler
from jupyter_server.utils import url_path_join
import tornado

from .config import resolve_config
from .opencode_client import OpenCodeClient, OpenCodeError
from .session_manager import SessionManager


log = logging.getLogger("opencode_bridge.routes")


# Unified system prompt — the LLM (OpenCode) is asked to return its reply
# as MARKDOWN (code wrapped in ```language fences; a brief explanation is
# fine). The frontend renders this with `marked` directly — no fence
# stripping on the server, so the response keeps the structure that makes
# markdown rendering meaningful (code blocks, headings, etc.).
UNIFIED_SYSTEM_PROMPT = (
    "你是一个代码助手。基于提供的代码上下文（以及可选的 traceback），"
    "按照用户的指令修改代码。"
    "可附简短说明。"
)


def make_client(handler: APIHandler) -> OpenCodeClient:
    """Factory for OpenCodeClient. Tests monkey-patch this."""
    cfg = resolve_config(handler.settings.get("opencode_bridge", {}))
    return OpenCodeClient(cfg)


def get_session_manager(handler: APIHandler) -> SessionManager:
    """Return the SessionManager singleton for this web app, creating on first use.

    Stored in handler.settings["opencode_bridge_session_manager"] so it survives
    across requests but is per-server-instance. Tests monkey-patch this.
    """
    sm = handler.settings.get("opencode_bridge_session_manager")
    if sm is None:
        def client_factory() -> OpenCodeClient:
            cfg = resolve_config(handler.settings.get("opencode_bridge", {}))
            return OpenCodeClient(cfg)
        sm = SessionManager(client_factory)
        handler.settings["opencode_bridge_session_manager"] = sm
    return sm


def _build_request_body(prompt: str, context: dict) -> dict:
    """Build full request body for OpenCode POST /session/:id/prompt_async.

    Returns dict with 'parts' (list) and 'system' (str) keys. No mode concept:
    the LLM interprets the user's natural-language instruction, with the code
    context and the optional traceback in front of it.
    """
    parts: list[dict] = []

    if context.get("previousCode"):
        parts.append({
            "type": "text",
            "text": "<previous_cell>\n%s\n</previous_cell>\n" % context["previousCode"],
        })

    error = context.get("error")
    if error:
        parts.append({
            "type": "text",
            "text": (
                "<traceback>\n%s: %s\n" % (error["ename"], error["evalue"])
                + "\n".join(error.get("traceback", []))
                + "\n</traceback>\n"
            ),
        })

    parts.append({
        "type": "text",
        "text": "<cell language='%s'>\n%s\n</cell>\n" % (
            context.get("language", "python"),
            context["source"],
        ),
    })

    if prompt:
        parts.append({"type": "text", "text": "<instruction>\n%s\n</instruction>" % prompt})

    return {"parts": parts, "system": UNIFIED_SYSTEM_PROMPT}


def _strip_code_fence(s: str) -> str:
    """Strip ```language ... ``` fences from LLM output.

    Kept for any future "apply-to-cell" path that needs clean source.
    Not used by the current markdown-rendering display flow.
    """
    s = s.strip()
    if s.startswith("```"):
        lines = s.split("\n")
        if lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].startswith("```"):
            lines = lines[:-1]
        return "\n".join(lines).strip()
    return s


def _extract_session_id(event: dict) -> str | None:
    """Best-effort session ID extraction from an OpenCode event.

    OpenCode's event shape is not strictly documented; this mirrors the
    heuristics in demo.html ``handleGlobalEvent``:
      properties.sessionID / properties.sessionId /
      payload.properties.sessionID / syncEvent.aggregateID
    """
    props = event.get("properties") or {}
    sid = props.get("sessionID") or props.get("sessionId")
    if sid:
        return sid
    payload_props = (event.get("payload") or {}).get("properties") or {}
    sid = payload_props.get("sessionID") or payload_props.get("sessionId")
    if sid:
        return sid
    sync = event.get("syncEvent") or (event.get("payload") or {}).get("syncEvent") or {}
    return sync.get("aggregateID")


class HelloRouteHandler(APIHandler):
    # The following decorator should be present on all verb methods (head, get, post,
    # patch, put, delete, options) to ensure only authorized user can request the
    # Jupyter server
    @tornado.web.authenticated
    def get(self):
        self.finish(json.dumps({
            "data": (
                "Hello, world!"
                " This is the '/opencode-bridge/hello' endpoint."
                " Try visiting me in your browser!"
            ),
        }))


class HealthHandler(APIHandler):
    @tornado.web.authenticated
    async def get(self):
        try:
            client = make_client(self)
            data = await client.health()
            self.finish(json.dumps({
                "ok": data.get("healthy", False),
                "version": data.get("version"),
                "endpoint": client.endpoint,
            }))
        except Exception as e:
            log.exception("health check failed")
            self.set_status(503)
            self.finish(json.dumps({
                "ok": False,
                "error": str(e),
                "endpoint": make_client(self).endpoint,
            }))


class ProvidersHandler(APIHandler):
    @tornado.web.authenticated
    async def get(self):
        try:
            data = await make_client(self).list_providers()
            self.finish(json.dumps(data))
        except Exception as e:
            log.exception("providers list failed")
            self.set_status(502)
            self.finish(json.dumps({"error": str(e)}))


class EditHandler(APIHandler):
    """Async edit: fire prompt_async, return sessionId immediately.

    The LLM's response is consumed via the separate ``/opencode-bridge/events``
    SSE endpoint. Frontend subscribes to SSE after this call returns and
    processes events incrementally.
    """

    @tornado.web.authenticated
    async def post(self):
        try:
            body = json.loads(self.request.body)
            prompt = body.get("prompt", "")
            context = body["context"]
            provider_id = body.get("providerId") or None
            model_id = body.get("modelId") or None
            notebook_path = context.get("notebookPath", "")
        except (KeyError, json.JSONDecodeError) as e:
            self.set_status(400)
            self.finish(json.dumps({"error": "bad request: %s" % e}))
            return

        if not notebook_path:
            self.set_status(400)
            self.finish(json.dumps({"error": "missing context.notebookPath"}))
            return

        client = make_client(self)
        sm = get_session_manager(self)
        try:
            sid = await sm.get_or_create(notebook_path)
            request_body = _build_request_body(prompt, context)
            await client.send_message_async(
                sid,
                request_body["parts"],
                provider_id=provider_id,
                model_id=model_id,
                system=request_body["system"],
            )
            self.finish(json.dumps({
                "ok": True,
                "sessionId": sid,
                "notebookPath": notebook_path,
            }))
        except OpenCodeError as e:
            # If session is invalid on OpenCode side, invalidate cache.
            if "404" in str(e) or "not found" in str(e).lower():
                sm.invalidate(notebook_path)
                log.warning("invalidated dead session for %s", notebook_path)
            log.exception("edit (async) failed")
            self.set_status(502)
            self.finish(json.dumps({"ok": False, "error": str(e)}))
        except Exception as e:
            log.exception("edit (async) failed")
            self.set_status(502)
            self.finish(json.dumps({"ok": False, "error": str(e)}))


class GlobalEventHandler(APIHandler):
    """SSE proxy for OpenCode Serve's ``GET /global/event``.

    Streams every event from the upstream SSE feed to the client as
    ``text/event-stream`` so EventSource / fetch+reader on the frontend
    can consume them directly. The frontend filters by session ID —
    we deliberately do NOT filter on the server because OpenCode's
    event-shape has varied across versions and a too-eager server
    filter can silently drop events the client would have accepted.

    The connection is held open until the client disconnects, the upstream
    closes, or an unrecoverable error occurs.
    """

    @tornado.web.authenticated
    async def get(self):
        # The ?session= query param is accepted for API compatibility
        # but the server does NOT use it to filter — see docstring.
        _ = self.get_query_argument("session", None)

        self.set_header("Content-Type", "text/event-stream")
        self.set_header("Cache-Control", "no-cache")
        # Disable proxy buffering (e.g. nginx) so chunks reach the client
        # immediately.
        self.set_header("X-Accel-Buffering", "no")
        # Prevent tornado from closing the connection on its own keep-alive
        # timer — we hold it open as long as the upstream is alive.
        self.set_header("Connection", "keep-alive")

        client = make_client(self)
        try:
            async for event in client.stream_global_events():
                payload = json.dumps(event, ensure_ascii=False)
                self.write("data: %s\n\n" % payload)
                # flush() pushes the chunk to the client; without it tornado
                # would buffer until the response is "finished".
                await self.flush()
        except tornado.iostream.StreamClosedError:
            # Client went away — normal disconnect, not an error.
            log.info("SSE client disconnected")
        except Exception as e:
            log.exception("SSE proxy error")
            try:
                self.write("data: %s\n\n" % json.dumps({
                    "type": "error",
                    "message": str(e),
                }))
                await self.flush()
            except Exception:
                # Connection already gone; nothing to send.
                pass


class PermissionReplyHandler(APIHandler):
    """Respond to a `permission.asked` event for a given session.

    POST ``/opencode-bridge/permissions/:permId?session=:sid`` with body
    ``{"response": "once"|"always"|"reject"}``. Forwards to OpenCode
    Serve's ``POST /session/:sid/permissions/:permId``.
    """

    VALID_RESPONSES = ("once", "always", "reject")

    @tornado.web.authenticated
    async def post(self, perm_id: str):
        session_id = self.get_query_argument("session", "")
        if not session_id:
            self.set_status(400)
            self.finish(json.dumps({"error": "missing 'session' query parameter"}))
            return
        try:
            body = json.loads(self.request.body) if self.request.body else {}
        except json.JSONDecodeError as e:
            self.set_status(400)
            self.finish(json.dumps({"error": "bad request: %s" % e}))
            return
        response = body.get("response")
        if response not in self.VALID_RESPONSES:
            self.set_status(400)
            self.finish(json.dumps({
                "error": "response must be one of %s" % list(self.VALID_RESPONSES)
            }))
            return
        client = make_client(self)
        try:
            ok = await client.reply_permission(session_id, perm_id, response)
            self.finish(json.dumps({
                "ok": ok,
                "permissionId": perm_id,
                "response": response,
            }))
        except OpenCodeError as e:
            log.exception("permission reply failed")
            self.set_status(502)
            self.finish(json.dumps({"ok": False, "error": str(e)}))
        except Exception as e:
            log.exception("permission reply failed")
            self.set_status(502)
            self.finish(json.dumps({"ok": False, "error": str(e)}))


class QuestionReplyHandler(APIHandler):
    """Reply to a `question.asked` event with a freeform answer.

    POST ``/opencode-bridge/questions/:qId/reply?session=:sid`` with body
    ``{"answer": "..."}``. Forwards to OpenCode Serve's
    ``POST /session/:sid/question/:qId/reply``.
    """

    @tornado.web.authenticated
    async def post(self, question_id: str):
        session_id = self.get_query_argument("session", "")
        if not session_id:
            self.set_status(400)
            self.finish(json.dumps({"error": "missing 'session' query parameter"}))
            return
        try:
            body = json.loads(self.request.body) if self.request.body else {}
        except json.JSONDecodeError as e:
            self.set_status(400)
            self.finish(json.dumps({"error": "bad request: %s" % e}))
            return
        answer = body.get("answer")
        if not isinstance(answer, str) or not answer.strip():
            self.set_status(400)
            self.finish(json.dumps({"error": "missing 'answer'"}))
            return
        client = make_client(self)
        try:
            ok = await client.reply_question(session_id, question_id, answer)
            self.finish(json.dumps({
                "ok": ok,
                "questionId": question_id,
            }))
        except OpenCodeError as e:
            log.exception("question reply failed")
            self.set_status(502)
            self.finish(json.dumps({"ok": False, "error": str(e)}))
        except Exception as e:
            log.exception("question reply failed")
            self.set_status(502)
            self.finish(json.dumps({"ok": False, "error": str(e)}))


class SessionListHandler(APIHandler):
    """List all active notebook -> session mappings. Debug endpoint."""

    @tornado.web.authenticated
    def get(self):
        sm = get_session_manager(self)
        self.finish(json.dumps({"sessions": sm.list_sessions()}))


class SessionReleaseHandler(APIHandler):
    """Release the OpenCode session for a specific notebook.

    Query param: notebook=<notebook path, URL-encoded>
    """

    @tornado.web.authenticated
    async def delete(self):
        notebook_path = self.get_query_argument("notebook", "")
        if not notebook_path:
            self.set_status(400)
            self.finish(json.dumps({"error": "missing 'notebook' query parameter"}))
            return
        sm = get_session_manager(self)
        deleted = await sm.release(notebook_path)
        self.finish(json.dumps({
            "ok": True,
            "notebookPath": notebook_path,
            "deleted": deleted,
        }))


class SessionMessagesHandler(APIHandler):
    """List the current session's messages for a notebook (scrollable history).

    Query param: notebook=<notebook path, URL-encoded>
    Returns: { messages: [{ role: "user"|"assistant", content: string }] }
    If no session exists for the notebook yet, returns { messages: [] }
    (does NOT create a session just to report emptiness).
    """

    @tornado.web.authenticated
    async def get(self):
        notebook_path = self.get_query_argument("notebook", "")
        if not notebook_path:
            self.set_status(400)
            self.finish(json.dumps({"error": "missing 'notebook' query parameter"}))
            return
        sm = get_session_manager(self)
        sid = sm.peek(notebook_path)
        if sid is None:
            self.finish(json.dumps({"messages": []}))
            return
        try:
            client = make_client(self)
            raw = await client.list_session_messages(sid)
        except OpenCodeError as e:
            if "404" in str(e) or "not found" in str(e).lower():
                sm.invalidate(notebook_path)
                log.warning("invalidated dead session for %s", notebook_path)
            log.exception("list session messages failed")
            self.set_status(502)
            self.finish(json.dumps({"ok": False, "error": str(e)}))
            return
        except Exception as e:
            log.exception("list session messages failed")
            self.set_status(502)
            self.finish(json.dumps({"ok": False, "error": str(e)}))
            return

        # Project OpenCode's {info, parts}[] into a frontend-friendly
        # {role, content}[] by joining the text parts.
        messages = []
        for m in raw or []:
            info = m.get("info") or {}
            role = info.get("role") or "assistant"
            parts = m.get("parts") or []
            content = "\n".join(
                p.get("text", "") for p in parts if p.get("type") == "text"
            ).strip()
            messages.append({"role": role, "content": content})
        self.finish(json.dumps({"messages": messages}))


def setup_route_handlers(web_app):
    host_pattern = ".*$"
    base_url = web_app.settings["base_url"]

    handlers = [
        (url_path_join(base_url, "opencode-bridge", "hello"), HelloRouteHandler),
        (url_path_join(base_url, "opencode-bridge", "health"), HealthHandler),
        (url_path_join(base_url, "opencode-bridge", "providers"), ProvidersHandler),
        (url_path_join(base_url, "opencode-bridge", "edit"), EditHandler),
        (url_path_join(base_url, "opencode-bridge", "events"), GlobalEventHandler),
        (url_path_join(base_url, "opencode-bridge", "sessions"), SessionListHandler),
        (url_path_join(base_url, "opencode-bridge", "session"), SessionReleaseHandler),
        (url_path_join(base_url, "opencode-bridge", "session-messages"), SessionMessagesHandler),
        # /permissions/<perm_id>  and  /questions/<q_id>/reply
        # use a URLSpec with a regex so the path param is forwarded to
        # the handler as a positional argument.
        tornado.web.URLSpec(
            url_path_join(base_url, "opencode-bridge", "permissions", r"([A-Za-z0-9_\-]+)"),
            PermissionReplyHandler,
        ),
        tornado.web.URLSpec(
            url_path_join(base_url, "opencode-bridge", "questions", r"([A-Za-z0-9_\-]+)", "reply"),
            QuestionReplyHandler,
        ),
    ]

    web_app.add_handlers(host_pattern, handlers)