zk-data-agent/agent_platform/gateway/app.py

import asyncio
from contextlib import asynccontextmanager
from typing import Annotated
from urllib.parse import quote

import uvicorn
from fastapi import Depends, FastAPI, Header, HTTPException, Query
from fastapi.responses import Response, StreamingResponse

from agent_platform.auth import UserIdentity, decode_openwebui_identity, verify_service_bearer
from agent_platform.config import Settings, get_settings
from agent_platform.gateway.provider import ExecutionProvider, create_execution_provider
from agent_platform.gateway.schemas import (
    ApplyPatchRequest,
    ExecRequest,
    GitDiffRequest,
    GitRequest,
    ListFilesRequest,
    ProcessRequest,
    ReadFileRequest,
    SearchFilesRequest,
    StartProcessRequest,
    ToolResult,
    WorkspaceListing,
    WorkspaceStatus,
    WriteFileRequest,
)


def create_app(
    settings: Settings | None = None,
    provider: ExecutionProvider | None = None,
) -> FastAPI:
    settings = settings or get_settings()

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        settings.validate_gateway()
        app.state.provider = provider or create_execution_provider(settings)
        yield

    app = FastAPI(
        title="K1412 Workspace Gateway",
        description="Authenticated tools for one isolated workspace per user.",
        version="0.1.0",
        lifespan=lifespan,
    )
    mutation_locks: dict[str, asyncio.Lock] = {}
    mutation_locks_guard = asyncio.Lock()

    def require_identity(
        authorization: Annotated[
            str | None,
            Header(alias="Authorization", include_in_schema=False),
        ] = None,
        user_jwt: Annotated[
            str | None,
            Header(alias="X-OpenWebUI-User-Jwt", include_in_schema=False),
        ] = None,
    ) -> UserIdentity:
        verify_service_bearer(authorization, settings.internal_gateway_key)
        return decode_openwebui_identity(user_jwt, settings.openwebui_forward_jwt_secret)

    Identity = Annotated[UserIdentity, Depends(require_identity)]

    def executor() -> ExecutionProvider:
        return app.state.provider

    def translate_value_error(exc: ValueError) -> HTTPException:
        return HTTPException(status_code=400, detail=str(exc))

    async def mutation_lock(user_id: str) -> asyncio.Lock:
        async with mutation_locks_guard:
            return mutation_locks.setdefault(user_id, asyncio.Lock())

    def attachment_header(filename: str) -> str:
        fallback = "".join(character for character in filename if character.isascii() and character.isalnum())
        fallback = fallback[:80] or "download"
        return f"attachment; filename={fallback}; filename*=UTF-8''{quote(filename)}"

    @app.get("/health", include_in_schema=False)
    async def health() -> dict:
        return {"status": "ok", "provider": settings.execution_provider}

    @app.post("/v1/tools/workspace_status", response_model=WorkspaceStatus, operation_id="workspace_status")
    async def workspace_status(identity: Identity) -> WorkspaceStatus:
        return await executor().status(identity.user_id)

    @app.post("/v1/tools/list_files", response_model=ToolResult, operation_id="list_files")
    async def list_files(body: ListFilesRequest, identity: Identity) -> ToolResult:
        try:
            return await executor().list_files(identity.user_id, body.path, body.max_depth, body.limit)
        except ValueError as exc:
            raise translate_value_error(exc) from exc

    @app.get("/v1/files", response_model=WorkspaceListing, operation_id="browse_workspace_files")
    async def browse_workspace_files(
        identity: Identity,
        path: Annotated[str, Query(max_length=4096)] = ".",
    ) -> WorkspaceListing:
        try:
            return await executor().browse_files(identity.user_id, path)
        except ValueError as exc:
            raise translate_value_error(exc) from exc
        except FileNotFoundError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc

    @app.get("/v1/files/download", operation_id="download_workspace_file")
    async def download_workspace_file(
        identity: Identity,
        path: Annotated[str, Query(max_length=4096)],
    ) -> Response:
        try:
            download = await executor().download_file(identity.user_id, path)
        except ValueError as exc:
            raise translate_value_error(exc) from exc
        except (FileNotFoundError, IsADirectoryError) as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc
        return Response(
            content=download.content or b"",
            media_type=download.media_type,
            headers={"Content-Disposition": attachment_header(download.filename)},
        )

    @app.get("/v1/files/archive", operation_id="archive_workspace_files")
    async def archive_workspace_files(
        identity: Identity,
        path: Annotated[str, Query(max_length=4096)] = ".",
    ) -> StreamingResponse:
        try:
            download = await executor().archive_files(identity.user_id, path)
        except ValueError as exc:
            raise translate_value_error(exc) from exc
        except FileNotFoundError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc
        return StreamingResponse(
            download.chunks or iter(()),
            media_type=download.media_type,
            headers={"Content-Disposition": attachment_header(download.filename)},
        )

    @app.post("/v1/tools/read_file", response_model=ToolResult, operation_id="read_file")
    async def read_file(body: ReadFileRequest, identity: Identity) -> ToolResult:
        try:
            return await executor().read_file(identity.user_id, body.path, body.start_line, body.max_lines)
        except ValueError as exc:
            raise translate_value_error(exc) from exc

    @app.post("/v1/tools/write_file", response_model=ToolResult, operation_id="write_file")
    async def write_file(body: WriteFileRequest, identity: Identity) -> ToolResult:
        try:
            lock = await mutation_lock(identity.user_id)
            async with lock:
                return await executor().write_file(identity.user_id, body.path, body.content)
        except ValueError as exc:
            raise translate_value_error(exc) from exc

    @app.post("/v1/tools/search_files", response_model=ToolResult, operation_id="search_files")
    async def search_files(body: SearchFilesRequest, identity: Identity) -> ToolResult:
        try:
            return await executor().search_files(
                identity.user_id,
                body.query,
                body.path,
                body.glob,
                body.limit,
            )
        except ValueError as exc:
            raise translate_value_error(exc) from exc

    @app.post("/v1/tools/exec", response_model=ToolResult, operation_id="exec")
    async def exec_command(body: ExecRequest, identity: Identity) -> ToolResult:
        try:
            lock = await mutation_lock(identity.user_id)
            async with lock:
                return await executor().exec(
                    identity.user_id,
                    body.command,
                    body.cwd,
                    min(body.timeout_seconds, settings.tool_timeout_seconds),
                )
        except ValueError as exc:
            raise translate_value_error(exc) from exc

    @app.post("/v1/tools/apply_patch", response_model=ToolResult, operation_id="apply_patch")
    async def apply_patch(body: ApplyPatchRequest, identity: Identity) -> ToolResult:
        try:
            lock = await mutation_lock(identity.user_id)
            async with lock:
                return await executor().apply_patch(identity.user_id, body.patch, body.cwd)
        except ValueError as exc:
            raise translate_value_error(exc) from exc

    @app.post("/v1/tools/git_status", response_model=ToolResult, operation_id="git_status")
    async def git_status(body: GitRequest, identity: Identity) -> ToolResult:
        try:
            return await executor().exec(identity.user_id, "git status --short --branch", body.cwd, 30)
        except ValueError as exc:
            raise translate_value_error(exc) from exc

    @app.post("/v1/tools/git_diff", response_model=ToolResult, operation_id="git_diff")
    async def git_diff(body: GitDiffRequest, identity: Identity) -> ToolResult:
        command = "git diff --cached" if body.staged else "git diff"
        try:
            return await executor().exec(identity.user_id, command, body.cwd, 30)
        except ValueError as exc:
            raise translate_value_error(exc) from exc

    @app.post("/v1/tools/start_process", response_model=ToolResult, operation_id="start_process")
    async def start_process(body: StartProcessRequest, identity: Identity) -> ToolResult:
        try:
            lock = await mutation_lock(identity.user_id)
            async with lock:
                return await executor().start_process(identity.user_id, body.command, body.cwd)
        except ValueError as exc:
            raise translate_value_error(exc) from exc

    @app.post("/v1/tools/poll_process", response_model=ToolResult, operation_id="poll_process")
    async def poll_process(body: ProcessRequest, identity: Identity) -> ToolResult:
        return await executor().poll_process(identity.user_id, body.process_id)

    @app.post("/v1/tools/cancel_process", response_model=ToolResult, operation_id="cancel_process")
    async def cancel_process(body: ProcessRequest, identity: Identity) -> ToolResult:
        lock = await mutation_lock(identity.user_id)
        async with lock:
            return await executor().cancel_process(identity.user_id, body.process_id)

    return app


app = create_app()


def run() -> None:
    uvicorn.run("agent_platform.gateway.app:app", host="0.0.0.0", port=8001)  # noqa: S104


if __name__ == "__main__":
    run()