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")


MODE_SYSTEM_PROMPTS = {
    "optimize": (
        "你是一个代码优化专家。请基于用户提供的代码上下文，"
        "返回只包含优化后代码的回复，不要任何解释或 markdown 围栏。"
    ),
    "fix": (
        "你是一个 Python 排错专家。用户给出了一段产生错误的代码和 traceback，"
        "请返回只包含修复后代码的回复，不要任何解释或 markdown 围栏。"
    ),
    "edit": (
        "你是一个代码编辑助手。基于用户的指令修改给定代码，"
        "返回只包含修改后完整代码的回复，不要任何解释或 markdown 围栏。"
    ),
}


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(mode: str, prompt: str, context: dict) -> dict:
    """Build full request body for OpenCode POST /session/:id/message.

    Returns dict with 'parts' (list) and 'system' (str) keys.
    """
    system = MODE_SYSTEM_PROMPTS[mode]
    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 mode == "edit" and prompt:
        parts.append({"type": "text", "text": "<instruction>\n%s\n</instruction>" % prompt})

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


def _strip_code_fence(s: str) -> str:
    """Strip ```language ... ``` fences from LLM output."""
    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


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):
    @tornado.web.authenticated
    async def post(self):
        try:
            body = json.loads(self.request.body)
            mode = body["mode"]
            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(mode, prompt, context)
            result = await client.send_message_sync(
                sid,
                request_body["parts"],
                provider_id=provider_id,
                model_id=model_id,
                system=request_body["system"],
            )

            text_parts = [
                p.get("text", "")
                for p in result.get("parts", [])
                if p.get("type") == "text"
            ]
            final_source = _strip_code_fence("\n".join(text_parts).strip())

            self.finish(json.dumps({
                "ok": True,
                "mode": mode,
                "finalSource": final_source,
                "sessionId": sid,
                "notebookPath": notebook_path,
            }))
        except OpenCodeError as e:
            # If session is invalid on OpenCode side, invalidate cache.
            # 404 / 410 / "session not found" -> next get_or_create will recreate.
            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 failed")
            self.set_status(502)
            self.finish(json.dumps({"ok": False, "error": str(e)}))
        except Exception as e:
            log.exception("edit failed")
            self.set_status(502)
            self.finish(json.dumps({"ok": False, "error": str(e)}))
        # NO finally delete — session is reused per notebook.


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,
        }))


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", "sessions"), SessionListHandler),
        (url_path_join(base_url, "opencode-bridge", "session"), SessionReleaseHandler),
    ]

    web_app.add_handlers(host_pattern, handlers)