okb 1.1.0a0__py3-none-any.whl → 1.1.2__py3-none-any.whl

okb/http_server.py

@@ -1,9 +1,15 @@
 """HTTP transport server for MCP with token authentication.
 
-This module provides an HTTP server that serves the LKB MCP server with
-token-based authentication. Tokens can be passed via Authorization header
-or query parameter. A single HTTP server can serve multiple databases,
-with the token determining which database to use.
+This module provides an HTTP server that serves the OKB MCP server with
+token-based authentication using Streamable HTTP transport. Tokens can be
+passed via Authorization header or query parameter. A single HTTP server
+can serve multiple databases, with the token determining which database to use.
+
+Transport: Streamable HTTP (RFC 9728 compliant)
+- POST /mcp → send JSON-RPC messages, get SSE response
+- GET /mcp → optional standalone SSE for server notifications
+- DELETE /mcp → terminate session
+- Session ID in Mcp-Session-Id header
 """
 
 from __future__ import annotations
@@ -12,12 +18,11 @@ import sys
 from typing import Any
 
 from mcp.server import Server
-from mcp.server.sse import SseServerTransport
+from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
 from mcp.types import CallToolResult, TextContent, Tool
-from starlette.applications import Starlette
+from starlette.middleware.cors import CORSMiddleware
 from starlette.requests import Request
-from starlette.responses import JSONResponse, Response
-from starlette.routing import Mount, Route
+from starlette.responses import JSONResponse
 
 from .config import config
 from .local_embedder import warmup
@@ -37,9 +42,16 @@ READ_ONLY_TOOLS = frozenset(
         "get_document",
         "list_sources",
         "list_projects",
+        "list_documents_by_project",
+        "get_project_stats",
         "recent_documents",
         "get_actionable_items",
         "get_database_info",
+        "list_sync_sources",
+        "list_pending_entities",
+        "list_pending_merges",
+        "get_topic_clusters",
+        "get_entity_relationships",
     }
 )
 
@@ -51,6 +63,17 @@ WRITE_TOOLS = frozenset(
         "add_todo",
         "trigger_sync",
         "trigger_rescan",
+        "enrich_document",
+        "approve_entity",
+        "reject_entity",
+        "analyze_knowledge_base",
+        "find_entity_duplicates",
+        "merge_entities",
+        "approve_merge",
+        "reject_merge",
+        "run_consolidation",
+        "rename_project",
+        "set_document_project",
     }
 )
 
@@ -66,14 +89,14 @@ def extract_token(request: Request) -> str | None:
 
 
 class HTTPMCPServer:
-    """HTTP server for MCP with token authentication."""
+    """HTTP server for MCP with token authentication using Streamable HTTP transport."""
 
     def __init__(self):
         self.knowledge_bases: dict[str, KnowledgeBase] = {}
         self.server = Server("knowledge-base")
-        # Single shared transport instance for all connections
-        self.transport = SseServerTransport("/messages/")
-        # Map session_id (hex string) -> token_info
+        # Session manager handles all transport complexity
+        self.session_manager = StreamableHTTPSessionManager(app=self.server)
+        # Map mcp-session-id -> token_info
         self.session_tokens: dict[str, TokenInfo] = {}
         self._setup_handlers()
 
@@ -208,6 +231,83 @@ class HTTPMCPServer:
                     content=[TextContent(type="text", text=f"## Projects\n\n{project_list}")]
                 )
 
+            elif name == "list_documents_by_project":
+                project = arguments["project"]
+                limit = arguments.get("limit", 100)
+                docs = kb.list_documents_by_project(project, limit)
+                if not docs:
+                    return CallToolResult(
+                        content=[
+                            TextContent(
+                                type="text", text=f"No documents found for project '{project}'."
+                            )
+                        ]
+                    )
+                output = [f"## Documents in '{project}' ({len(docs)} documents)\n"]
+                for d in docs:
+                    output.append(f"- **{d['title'] or d['source_path']}** ({d['source_type']})")
+                    output.append(f"  - `{d['source_path']}`")
+                return CallToolResult(content=[TextContent(type="text", text="\n".join(output))])
+
+            elif name == "get_project_stats":
+                stats = kb.get_project_stats()
+                if not stats:
+                    return CallToolResult(
+                        content=[TextContent(type="text", text="No projects found.")]
+                    )
+                output = ["## Project Statistics\n"]
+                for p in stats:
+                    output.append(f"- **{p['project']}**: {p['document_count']} documents")
+                return CallToolResult(content=[TextContent(type="text", text="\n".join(output))])
+
+            elif name == "rename_project":
+                old_name = arguments["old_name"]
+                new_name = arguments["new_name"]
+                if old_name == new_name:
+                    return CallToolResult(
+                        content=[TextContent(type="text", text="Old and new names are the same.")]
+                    )
+                count = kb.rename_project(old_name, new_name)
+                if count == 0:
+                    return CallToolResult(
+                        content=[
+                            TextContent(
+                                type="text", text=f"No documents found with project '{old_name}'."
+                            )
+                        ]
+                    )
+                return CallToolResult(
+                    content=[
+                        TextContent(
+                            type="text",
+                            text=f"Renamed project '{old_name}' to '{new_name}' "
+                            f"({count} documents updated).",
+                        )
+                    ]
+                )
+
+            elif name == "set_document_project":
+                source_path = arguments["source_path"]
+                project = arguments.get("project")
+                success = kb.set_document_project(source_path, project)
+                if not success:
+                    return CallToolResult(
+                        content=[
+                            TextContent(type="text", text=f"Document not found: {source_path}")
+                        ]
+                    )
+                if project:
+                    return CallToolResult(
+                        content=[
+                            TextContent(
+                                type="text", text=f"Set project to '{project}' for {source_path}"
+                            )
+                        ]
+                    )
+                return CallToolResult(
+                    content=[TextContent(type="text", text=f"Cleared project for {source_path}")]
+                )
+
             elif name == "recent_documents":
                 from .mcp_server import format_relative_time, get_document_date
 
@@ -265,13 +365,13 @@ class HTTPMCPServer:
                 deleted = kb.delete_knowledge(arguments["source_path"])
                 if deleted:
                     return CallToolResult(
-                        content=[TextContent(type="text", text="Knowledge entry deleted.")]
+                        content=[TextContent(type="text", text="Document deleted.")]
                     )
                 return CallToolResult(
                     content=[
                         TextContent(
                             type="text",
-                            text="Could not delete. Entry not found or not a Claude-saved entry.",
+                            text="Could not delete. Document not found.",
                         )
                     ]
                 )
@@ -394,6 +494,134 @@ class HTTPMCPServer:
                 )
                 return CallToolResult(content=[TextContent(type="text", text=result)])
 
+            elif name == "list_sync_sources":
+                from .mcp_server import _list_sync_sources
+
+                token_info = getattr(self.server, "_current_token_info", None)
+                db_name = token_info.database if token_info else config.get_database().name
+                result = _list_sync_sources(kb.db_url, db_name)
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "enrich_document":
+                from .mcp_server import _enrich_document
+
+                result = _enrich_document(
+                    kb.db_url,
+                    source_path=arguments["source_path"],
+                    extract_todos=arguments.get("extract_todos", True),
+                    extract_entities=arguments.get("extract_entities", True),
+                    auto_create_entities=arguments.get("auto_create_entities", False),
+                )
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "list_pending_entities":
+                from .mcp_server import _list_pending_entities
+
+                result = _list_pending_entities(
+                    kb.db_url,
+                    entity_type=arguments.get("entity_type"),
+                    limit=arguments.get("limit", 20),
+                )
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "approve_entity":
+                from .mcp_server import _approve_entity
+
+                result = _approve_entity(kb.db_url, arguments["pending_id"])
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "reject_entity":
+                from .mcp_server import _reject_entity
+
+                result = _reject_entity(kb.db_url, arguments["pending_id"])
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "analyze_knowledge_base":
+                from .mcp_server import _analyze_knowledge_base
+
+                result = _analyze_knowledge_base(
+                    kb.db_url,
+                    project=arguments.get("project"),
+                    sample_size=arguments.get("sample_size", 15),
+                    auto_update=arguments.get("auto_update", True),
+                )
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            # Entity consolidation tools
+            elif name == "find_entity_duplicates":
+                from .mcp_server import _find_entity_duplicates
+
+                result = _find_entity_duplicates(
+                    kb.db_url,
+                    similarity_threshold=arguments.get("similarity_threshold", 0.85),
+                    limit=arguments.get("limit", 50),
+                )
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "merge_entities":
+                from .mcp_server import _merge_entities
+
+                result = _merge_entities(
+                    kb.db_url,
+                    canonical_path=arguments["canonical_path"],
+                    duplicate_path=arguments["duplicate_path"],
+                )
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "list_pending_merges":
+                from .mcp_server import _list_pending_merges
+
+                result = _list_pending_merges(
+                    kb.db_url,
+                    limit=arguments.get("limit", 50),
+                )
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "approve_merge":
+                from .mcp_server import _approve_merge
+
+                result = _approve_merge(kb.db_url, arguments["merge_id"])
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "reject_merge":
+                from .mcp_server import _reject_merge
+
+                result = _reject_merge(kb.db_url, arguments["merge_id"])
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "get_topic_clusters":
+                from .mcp_server import _get_topic_clusters
+
+                result = _get_topic_clusters(
+                    kb.db_url,
+                    limit=arguments.get("limit", 20),
+                )
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "get_entity_relationships":
+                from .mcp_server import _get_entity_relationships
+
+                result = _get_entity_relationships(
+                    kb.db_url,
+                    entity_name=arguments.get("entity_name"),
+                    relationship_type=arguments.get("relationship_type"),
+                    limit=arguments.get("limit", 50),
+                )
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
+            elif name == "run_consolidation":
+                from .mcp_server import _run_consolidation
+
+                result = _run_consolidation(
+                    kb.db_url,
+                    detect_duplicates=arguments.get("detect_duplicates", True),
+                    detect_cross_doc=arguments.get("detect_cross_doc", True),
+                    build_clusters=arguments.get("build_clusters", True),
+                    extract_relationships=arguments.get("extract_relationships", True),
+                    dry_run=arguments.get("dry_run", False),
+                )
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+
             else:
                 return CallToolResult(
                     content=[TextContent(type="text", text=f"Unknown tool: {name}")]
@@ -402,95 +630,131 @@ class HTTPMCPServer:
         except Exception as e:
             return CallToolResult(content=[TextContent(type="text", text=f"Error: {e!s}")])
 
-    def create_app(self) -> Starlette:
+    def create_app(self):
         """Create the Starlette application."""
         verifier = OKBTokenVerifier(self._get_db_url)
+        session_header_name = "mcp-session-id"
 
-        async def handle_sse(request: Request) -> Response:
-            """Handle SSE connections for MCP."""
-            # Verify token
-            token = extract_token(request)
-            if not token:
-                return JSONResponse(
-                    {"error": "Missing token. Use Authorization header or ?token= parameter"},
-                    status_code=401,
-                )
-
-            token_info = verifier.verify(token)
-            if not token_info:
-                return JSONResponse(
-                    {"error": "Invalid or expired token"},
-                    status_code=401,
-                )
-
-            # Track existing sessions before connecting
-            existing_sessions = set(self.transport._read_stream_writers.keys())
-
-            async with self.transport.connect_sse(
-                request.scope, request.receive, request._send
-            ) as (read_stream, write_stream):
-                # Find the new session ID by comparing before/after
-                current_sessions = set(self.transport._read_stream_writers.keys())
-                new_sessions = current_sessions - existing_sessions
-                if not new_sessions:
-                    return JSONResponse(
-                        {"error": "Failed to establish session"},
-                        status_code=500,
-                    )
-                session_id = new_sessions.pop()
-                session_id_hex = session_id.hex
+        def create_mcp_handler():
+            """Create an ASGI handler for MCP with auth."""
 
-                # Store token mapping for this session
-                self.session_tokens[session_id_hex] = token_info
-                self.server._current_token_info = token_info
+            async def handle_mcp(scope, receive, send):
+                """Handle all MCP requests (GET, POST, DELETE) with auth."""
+                request = Request(scope, receive)
 
-                try:
-                    await self.server.run(
-                        read_stream, write_stream, self.server.create_initialization_options()
+                # Extract and verify token
+                token = extract_token(request)
+                if not token:
+                    response = JSONResponse(
+                        {"error": "Missing token. Use Authorization header or ?token= param"},
+                        status_code=401,
                     )
-                finally:
-                    # Clean up session on disconnect
-                    self.session_tokens.pop(session_id_hex, None)
-
-            return Response()
-
-        async def handle_messages(scope, receive, send):
-            """Handle POST messages for MCP (raw ASGI handler)."""
-            request = Request(scope, receive)
-
-            # Look up session from query params
-            session_id = request.query_params.get("session_id")
-            if not session_id:
-                response = JSONResponse({"error": "Missing session_id"}, status_code=400)
-                await response(scope, receive, send)
-                return
-
-            token_info = self.session_tokens.get(session_id)
-            if not token_info:
-                response = JSONResponse({"error": "Invalid or expired session"}, status_code=401)
-                await response(scope, receive, send)
-                return
+                    await response(scope, receive, send)
+                    return
+
+                token_info = verifier.verify(token)
+                if not token_info:
+                    response = JSONResponse(
+                        {"error": "Invalid or expired token"},
+                        status_code=401,
+                    )
+                    await response(scope, receive, send)
+                    return
+
+                # Check if this is an existing session
+                session_id = request.headers.get(session_header_name)
+                if session_id:
+                    # Verify token matches existing session (compare by hash and db, not object)
+                    existing_token = self.session_tokens.get(session_id)
+                    if existing_token:
+                        # Token must match the one used to create the session
+                        if (
+                            existing_token.token_hash != token_info.token_hash
+                            or existing_token.database != token_info.database
+                        ):
+                            response = JSONResponse(
+                                {"error": "Token mismatch for session"},
+                                status_code=401,
+                            )
+                            await response(scope, receive, send)
+                            return
 
-            # Set current token info for tool calls
-            self.server._current_token_info = token_info
+                # Set current token info for tool calls
+                self.server._current_token_info = token_info
 
-            await self.transport.handle_post_message(scope, receive, send)
+                # Wrap send to capture the session ID from response headers
+                captured_session_id = None
 
-        async def health(request: Request) -> JSONResponse:
-            """Health check endpoint."""
-            return JSONResponse({"status": "ok"})
+                async def send_wrapper(message):
+                    nonlocal captured_session_id
+                    if message["type"] == "http.response.start":
+                        headers = message.get("headers", [])
+                        for name, value in headers:
+                            header_name = (
+                                name.lower() if isinstance(name, bytes) else name.lower().encode()
+                            )
+                            if header_name == session_header_name.encode():
+                                captured_session_id = (
+                                    value.decode() if isinstance(value, bytes) else value
+                                )
+                                # Store immediately since SSE keeps connection open
+                                if captured_session_id not in self.session_tokens:
+                                    self.session_tokens[captured_session_id] = token_info
+                                break
+                    await send(message)
+
+                # Delegate to session manager
+                await self.session_manager.handle_request(scope, receive, send_wrapper)
+
+            return handle_mcp
+
+        # Create the MCP handler ASGI app
+        mcp_handler = create_mcp_handler()
+
+        # Custom ASGI app that routes /mcp and /sse to MCP handler
+        async def router(scope, receive, send):
+            if scope["type"] == "http":
+                path = scope["path"].rstrip("/")  # Handle trailing slash
+                if path in ("/mcp", "/sse"):
+                    await mcp_handler(scope, receive, send)
+                    return
+                elif path == "/health" or scope["path"] == "/health":
+                    response = JSONResponse({"status": "ok"})
+                    await response(scope, receive, send)
+                    return
+            # 404 for unknown paths
+            response = JSONResponse({"error": "Not found"}, status_code=404)
+            await response(scope, receive, send)
+
+        # Wrap with lifespan handling
+        async def app_with_lifespan(scope, receive, send):
+            if scope["type"] == "lifespan":
+                async with self.session_manager.run():
+                    # Handle lifespan protocol
+                    while True:
+                        message = await receive()
+                        if message["type"] == "lifespan.startup":
+                            await send({"type": "lifespan.startup.complete"})
+                        elif message["type"] == "lifespan.shutdown":
+                            await send({"type": "lifespan.shutdown.complete"})
+                            return
+            else:
+                await router(scope, receive, send)
 
-        routes = [
-            Route("/health", health, methods=["GET"]),
-            Route("/sse", handle_sse, methods=["GET"]),
-            Mount("/messages", app=handle_messages),
-        ]
+        # Add CORS for browser clients - wrap the raw ASGI app
+        app = CORSMiddleware(
+            app_with_lifespan,
+            allow_origins=["*"],
+            allow_methods=["GET", "POST", "DELETE"],
+            allow_headers=["Authorization", "Content-Type", session_header_name],
+            expose_headers=[session_header_name],
+        )
 
-        return Starlette(routes=routes)
+        return app
 
 
 def run_http_server(host: str = "127.0.0.1", port: int = 8080):
-    """Run the HTTP MCP server."""
+    """Run the HTTP MCP server with Streamable HTTP transport."""
     import uvicorn
 
     print("Warming up embedding model...", file=sys.stderr)
@@ -501,8 +765,9 @@ def run_http_server(host: str = "127.0.0.1", port: int = 8080):
     app = http_server.create_app()
 
     print(f"Starting HTTP MCP server on http://{host}:{port}", file=sys.stderr)
-    print("  SSE endpoint: /sse", file=sys.stderr)
-    print("  Messages endpoint: /messages/", file=sys.stderr)
+    print("  MCP endpoint: /mcp (GET, POST, DELETE)", file=sys.stderr)
+    print("  MCP endpoint: /sse (alias for /mcp)", file=sys.stderr)
     print("  Health endpoint: /health", file=sys.stderr)
+    print("  Transport: Streamable HTTP", file=sys.stderr)
 
     uvicorn.run(app, host=host, port=port, log_level="info")