fix: Update confluence-collab proxy with proper async lifecycle (CF-1812)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

confluence-collab/src/confluence_collab/proxy.py

@@ -3,6 +3,8 @@
 Spawns mcp-atlassian as a subprocess (stdio), proxies all its tools, and
 registers the confluence_section_* tools from this package. Claude Code
 sees a single MCP server with all tools under one prefix.
+
+Usage: python -m confluence_collab.proxy
 """
 
 from __future__ import annotations
@@ -13,81 +15,165 @@ import logging
 import os
 import sys
 
-from mcp.server.fastmcp import FastMCP
+from mcp import types
+from mcp.server.lowlevel import Server
+from mcp.server.stdio import stdio_server
 from mcp.client.session import ClientSession
 from mcp.client.stdio import StdioServerParameters, stdio_client
 
-from confluence_collab.server import (
+from confluence_collab.client import Auth
-    confluence_section_list,
+from confluence_collab.editor import (
-    confluence_section_get,
+    section_list,
-    confluence_section_update,
+    section_get,
-    confluence_section_append,
+    section_update,
-    confluence_section_delete,
+    section_append,
+    section_delete,
 )
 
 logger = logging.getLogger("confluence-collab-proxy")
 
-mcp = FastMCP("atlassian-with-sections")
+# Section tool definitions (added alongside proxied mcp-atlassian tools)
+SECTION_TOOLS = [
+    types.Tool(
+        name="confluence_section_list",
+        description="List all sections (headings) on a Confluence page. Returns JSON array of {heading, level}.",
+        inputSchema={
+            "type": "object",
+            "properties": {
+                "page_id": {"type": "string", "description": "Confluence page ID"},
+            },
+            "required": ["page_id"],
+        },
+    ),
+    types.Tool(
+        name="confluence_section_get",
+        description="Get the HTML content of a specific section by heading text. Uses fuzzy matching (case-insensitive, partial match).",
+        inputSchema={
+            "type": "object",
+            "properties": {
+                "page_id": {"type": "string", "description": "Confluence page ID"},
+                "heading": {"type": "string", "description": "Section heading text to find"},
+            },
+            "required": ["page_id", "heading"],
+        },
+    ),
+    types.Tool(
+        name="confluence_section_update",
+        description="Update the content of a section identified by heading. Replaces only the target section, preserving the rest of the page. Handles version conflicts with automatic retry (exponential backoff).",
+        inputSchema={
+            "type": "object",
+            "properties": {
+                "page_id": {"type": "string", "description": "Confluence page ID"},
+                "heading": {"type": "string", "description": "Section heading text to find (fuzzy matched)"},
+                "body": {"type": "string", "description": "New HTML content for the section"},
+            },
+            "required": ["page_id", "heading", "body"],
+        },
+    ),
+    types.Tool(
+        name="confluence_section_append",
+        description="Append HTML content to the end of a section. Adds content after existing section content, before the next heading.",
+        inputSchema={
+            "type": "object",
+            "properties": {
+                "page_id": {"type": "string", "description": "Confluence page ID"},
+                "heading": {"type": "string", "description": "Section heading text to find (fuzzy matched)"},
+                "body": {"type": "string", "description": "HTML content to append"},
+            },
+            "required": ["page_id", "heading", "body"],
+        },
+    ),
+    types.Tool(
+        name="confluence_section_delete",
+        description="Delete an entire section (heading + content) from a page. Removes the heading tag and all content up to the next same-or-higher-level heading.",
+        inputSchema={
+            "type": "object",
+            "properties": {
+                "page_id": {"type": "string", "description": "Confluence page ID"},
+                "heading": {"type": "string", "description": "Section heading text to find (fuzzy matched)"},
+            },
+            "required": ["page_id", "heading"],
+        },
+    ),
+]
 
-# Register section tools directly (they're already decorated with @mcp.tool in server.py,
+SECTION_TOOL_NAMES = {t.name for t in SECTION_TOOLS}
-# but we need to re-register them on this new FastMCP instance)
-mcp.tool()(confluence_section_list)
-mcp.tool()(confluence_section_get)
-mcp.tool()(confluence_section_update)
-mcp.tool()(confluence_section_append)
-mcp.tool()(confluence_section_delete)
 
 
-async def _proxy_upstream_tools() -> None:
+def _get_auth() -> Auth:
-    """Connect to mcp-atlassian subprocess and proxy its tools."""
+    return Auth.from_env()
+
+
+async def _handle_section_tool(name: str, arguments: dict) -> list[types.TextContent]:
+    """Handle a section tool call and return MCP content."""
+    auth = _get_auth()
+    page_id = arguments["page_id"]
+
+    if name == "confluence_section_list":
+        sections = await section_list(page_id, auth)
+        text = json.dumps([{"heading": s.heading, "level": s.level} for s in sections], indent=2)
+    elif name == "confluence_section_get":
+        content = await section_get(page_id, arguments["heading"], auth)
+        text = content if content is not None else f"Section '{arguments['heading']}' not found on page {page_id}"
+    elif name == "confluence_section_update":
+        result = await section_update(page_id, arguments["heading"], arguments["body"], auth)
+        text = json.dumps({"status": "ok" if result.ok else "error", "message": result.message, "version": result.version, "retries": result.retries})
+    elif name == "confluence_section_append":
+        result = await section_append(page_id, arguments["heading"], arguments["body"], auth)
+        text = json.dumps({"status": "ok" if result.ok else "error", "message": result.message, "version": result.version})
+    elif name == "confluence_section_delete":
+        result = await section_delete(page_id, arguments["heading"], auth)
+        text = json.dumps({"status": "ok" if result.ok else "error", "message": result.message, "version": result.version})
+    else:
+        text = f"Unknown section tool: {name}"
+
+    return [types.TextContent(type="text", text=text)]
+
+
+async def run_proxy():
+    """Run the composite MCP server with upstream proxy + section tools."""
+    server = Server("atlassian")
+
+    # Connect to upstream mcp-atlassian
     cmd = "uvx"
     args = ["--python", "3.13", "mcp-atlassian"]
-
     server_params = StdioServerParameters(command=cmd, args=args, env=dict(os.environ))
 
-    async with stdio_client(server_params) as (read, write):
+    async with stdio_client(server_params) as (upstream_read, upstream_write):
-        async with ClientSession(read, write) as session:
+        async with ClientSession(upstream_read, upstream_write) as upstream:
-            await session.initialize()
+            await upstream.initialize()
 
-            # List upstream tools
+            # Discover upstream tools
-            tools_result = await session.list_tools()
+            upstream_tools_result = await upstream.list_tools()
-            logger.info("Proxying %d upstream tools from mcp-atlassian", len(tools_result.tools))
+            upstream_tools = upstream_tools_result.tools
+            logger.info("Proxying %d upstream tools + %d section tools", len(upstream_tools), len(SECTION_TOOLS))
 
-            # Register each upstream tool as a proxy on our server
+            # Combined tool list
-            for tool in tools_result.tools:
+            all_tools = list(upstream_tools) + SECTION_TOOLS
-                _register_proxy_tool(tool, session)
 
-            # Keep running until interrupted
+            @server.list_tools()
-            await asyncio.Event().wait()
+            async def handle_list_tools() -> list[types.Tool]:
+                return all_tools
 
+            @server.call_tool()
+            async def handle_call_tool(name: str, arguments: dict | None) -> list[types.TextContent]:
+                arguments = arguments or {}
 
-def _register_proxy_tool(tool, session: ClientSession) -> None:
+                # Section tools handled locally
-    """Register a proxied tool from the upstream MCP server."""
+                if name in SECTION_TOOL_NAMES:
+                    return await _handle_section_tool(name, arguments)
 
-    async def proxy_handler(**kwargs):
+                # Proxy to upstream mcp-atlassian
-        result = await session.call_tool(tool.name, kwargs)
+                result = await upstream.call_tool(name, arguments)
-        # Return concatenated text content
+                return result.content
-        texts = []
-        for content in result.content:
-            if hasattr(content, "text"):
-                texts.append(content.text)
-        return "\n".join(texts) if texts else ""
 
-    proxy_handler.__name__ = tool.name
+            # Run our server on stdio
-    proxy_handler.__doc__ = tool.description or ""
+            async with stdio_server() as (read, write):
-
+                await server.run(read, write, server.create_initialization_options())
-    # Build parameter annotations from tool schema
-    mcp.tool(name=tool.name, description=tool.description or "")(proxy_handler)
 
 
 def main():
-    """Run the composite MCP server.
+    asyncio.run(run_proxy())
-
-    Note: The proxy approach requires running the upstream mcp-atlassian
-    as a subprocess. For simpler deployment, use the standalone server.py
-    which only provides section tools, and keep mcp-atlassian separate.
-    """
-    mcp.run(transport="stdio")
 
 
 if __name__ == "__main__":