server.py

import asyncio
import json
import os
import subprocess
import time
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, AsyncIterator

import boto3
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse


app = FastAPI()

DEFAULT_MODEL = "amazon-bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
DEFAULT_BEDROCK_MODEL = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
OPENCODE_URL = "http://127.0.0.1:4097"
OPENCODE_DIRECTORY = "/app"

opencode_proc: subprocess.Popen[str] | None = None
runtime_sessions: dict[str, str] = {}


@app.on_event("startup")
async def startup() -> None:
    await asyncio.to_thread(ensure_opencode_server)


@app.get("/ping")
def ping() -> dict[str, object]:
    return {"status": "Healthy", "time_of_last_update": int(time.time())}


@app.post("/invocations")
async def invocations(request: Request) -> StreamingResponse:
    body = await request.json()
    prompt = str(body.get("prompt") or "").strip()
    if not prompt:
        return JSONResponse({"error": "prompt is required"}, status_code=400)

    return StreamingResponse(stream_response(body, request), media_type="text/event-stream")


async def stream_response(body: dict[str, Any], request: Request) -> AsyncIterator[str]:
    try:
        result = await asyncio.to_thread(invoke_opencode, body, request)
    except Exception as error:
        yield sse({"type": "opencode_error", "message": str(error)})
        return

    for part in result.get("parts") or []:
        if not isinstance(part, dict):
            continue
        yield sse({"type": normalize_part_type(str(part.get("type") or "part")), "part": part})

    info = result.get("info")
    if isinstance(info, dict):
        yield sse({"type": "message_finish", "info": info})


def invoke_opencode(body: dict[str, Any], request: Request) -> dict[str, Any]:
    ensure_opencode_server()
    session_id = get_or_create_opencode_session(body, request)
    model = parse_model()
    payload = {
        "agent": os.environ.get("OPENCODE_AGENT", "build"),
        "model": {
            "providerID": model["providerID"],
            "modelID": model["modelID"],
        },
        "system": str(body.get("systemPrompt") or body.get("system") or "").strip() or None,
        "parts": [{"type": "text", "text": str(body["prompt"])}],
    }
    payload = {key: value for key, value in payload.items() if value is not None}
    return post_json(f"/session/{session_id}/message?directory={urllib.parse.quote(OPENCODE_DIRECTORY)}", payload)


def get_or_create_opencode_session(body: dict[str, Any], request: Request) -> str:
    runtime_session_id = runtime_session_key(body, request)
    if runtime_session_id in runtime_sessions:
        return runtime_sessions[runtime_session_id]

    model = parse_model()
    session = post_json(
        f"/session?directory={urllib.parse.quote(OPENCODE_DIRECTORY)}",
        {
            "agent": os.environ.get("OPENCODE_AGENT", "build"),
            "model": {
                "providerID": model["providerID"],
                "id": model["modelID"],
            },
            "title": f"AgentCore {runtime_session_id}",
        },
    )
    session_id = session["id"]
    runtime_sessions[runtime_session_id] = session_id
    return session_id


def runtime_session_key(body: dict[str, Any], request: Request) -> str:
    metadata = body.get("metadata")
    if isinstance(metadata, dict):
        for key in ("runtimeSessionId", "controlPlaneSessionId"):
            value = metadata.get(key)
            if isinstance(value, str) and value:
                return value
    for header in (
        "x-amz-bedrock-agentcore-runtime-session-id",
        "x-amzn-bedrock-agentcore-runtime-session-id",
    ):
        value = request.headers.get(header)
        if value:
            return value
    return "default"


def ensure_opencode_server() -> None:
    global opencode_proc
    if opencode_proc and opencode_proc.poll() is None and opencode_healthcheck():
        return

    env = {
        **os.environ,
        "AWS_REGION": region(),
        "AWS_DEFAULT_REGION": region(),
    }
    apply_boto_credentials(env)
    opencode_proc = subprocess.Popen(
        [
            "opencode",
            "serve",
            "--pure",
            "--hostname",
            "127.0.0.1",
            "--port",
            "4097",
        ],
        env=env,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.STDOUT,
        text=True,
    )
    deadline = time.time() + 15
    while time.time() < deadline:
        if opencode_proc.poll() is not None:
            raise RuntimeError(f"opencode server exited with status {opencode_proc.returncode}")
        if opencode_healthcheck():
            return
        time.sleep(0.2)
    raise RuntimeError("opencode server did not become healthy")


def opencode_healthcheck() -> bool:
    try:
        request = urllib.request.Request(f"{OPENCODE_URL}/session/status", method="GET")
        with urllib.request.urlopen(request, timeout=2) as response:
            return response.status < 500
    except Exception:
        return False


def post_json(path: str, payload: dict[str, Any]) -> dict[str, Any]:
    data = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(
        f"{OPENCODE_URL}{path}",
        data=data,
        headers={"content-type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=300) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as error:
        details = error.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"opencode server returned {error.code}: {details}") from error


def parse_model() -> dict[str, str]:
    configured = os.environ.get("OPENCODE_MODEL", DEFAULT_MODEL)
    if "/" in configured:
        provider_id, model_id = configured.split("/", 1)
    else:
        provider_id = "amazon-bedrock"
        model_id = os.environ.get("BEDROCK_MODEL_ID", configured or DEFAULT_BEDROCK_MODEL)
    return {"providerID": provider_id, "modelID": model_id}


def apply_boto_credentials(env: dict[str, str]) -> None:
    credentials = boto3.Session(region_name=region()).get_credentials()
    if credentials is None:
        return
    frozen = credentials.get_frozen_credentials()
    if frozen.access_key and frozen.secret_key:
        env["AWS_ACCESS_KEY_ID"] = frozen.access_key
        env["AWS_SECRET_ACCESS_KEY"] = frozen.secret_key
    if frozen.token:
        env["AWS_SESSION_TOKEN"] = frozen.token


def region() -> str:
    return os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "us-west-2"


def normalize_part_type(part_type: str) -> str:
    return part_type.replace("-", "_")


def sse(payload: dict[str, object]) -> str:
    return f"data: {json.dumps(payload, separators=(',', ':'))}\n\n"