From ba88ba01f32d9c017d00f641926f7ddaa8f69083 Mon Sep 17 00:00:00 2001
From: Bill <bill@bballou.com>
Date: Fri, 2 Jan 2026 13:51:47 -0500
Subject: [PATCH] feat(server): register session token tools

Add get_proxy_documentation and request_session_token tools to the MCP
server. The create_server function now accepts an optional token_manager
parameter (SessionTokenManager | None) to maintain backward compatibility.

When token_manager is None, request_session_token returns an error
message instead of creating tokens.
---
 src/grist_mcp/server.py   | 45 ++++++++++++++++++++++++++++++++++++-
 tests/unit/test_server.py | 47 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 89 insertions(+), 3 deletions(-)

diff --git a/src/grist_mcp/server.py b/src/grist_mcp/server.py
index 6081568..3c3e400 100644
--- a/src/grist_mcp/server.py
+++ b/src/grist_mcp/server.py
@@ -6,6 +6,9 @@ from mcp.server import Server
 from mcp.types import Tool, TextContent
 
 from grist_mcp.auth import Authenticator, Agent, AuthError
+from grist_mcp.session import SessionTokenManager
+from grist_mcp.tools.session import get_proxy_documentation as _get_proxy_documentation
+from grist_mcp.tools.session import request_session_token as _request_session_token
 
 from grist_mcp.tools.discovery import list_documents as _list_documents
 from grist_mcp.tools.read import list_tables as _list_tables
@@ -21,12 +24,13 @@ from grist_mcp.tools.schema import modify_column as _modify_column
 from grist_mcp.tools.schema import delete_column as _delete_column
 
 
-def create_server(auth: Authenticator, agent: Agent) -> Server:
+def create_server(auth: Authenticator, agent: Agent, token_manager: SessionTokenManager | None = None) -> Server:
     """Create and configure the MCP server for an authenticated agent.
 
     Args:
         auth: Authenticator instance for permission checks.
         agent: The authenticated agent for this server instance.
+        token_manager: Optional session token manager for HTTP proxy access.
 
     Returns:
         Configured MCP Server instance.
@@ -203,6 +207,34 @@ def create_server(auth: Authenticator, agent: Agent) -> Server:
                     "required": ["document", "table", "column_id"],
                 },
             ),
+            Tool(
+                name="get_proxy_documentation",
+                description="Get complete documentation for the HTTP proxy API",
+                inputSchema={"type": "object", "properties": {}, "required": []},
+            ),
+            Tool(
+                name="request_session_token",
+                description="Request a short-lived token for direct HTTP API access. Use this to delegate bulk data operations to scripts.",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "document": {
+                            "type": "string",
+                            "description": "Document name to grant access to",
+                        },
+                        "permissions": {
+                            "type": "array",
+                            "items": {"type": "string", "enum": ["read", "write", "schema"]},
+                            "description": "Permission levels to grant",
+                        },
+                        "ttl_seconds": {
+                            "type": "integer",
+                            "description": "Token lifetime in seconds (max 3600, default 300)",
+                        },
+                    },
+                    "required": ["document", "permissions"],
+                },
+            ),
         ]
 
     @server.call_tool()
@@ -265,6 +297,17 @@ def create_server(auth: Authenticator, agent: Agent) -> Server:
                     _current_agent, auth, arguments["document"], arguments["table"],
                     arguments["column_id"],
                 )
+            elif name == "get_proxy_documentation":
+                result = await _get_proxy_documentation()
+            elif name == "request_session_token":
+                if token_manager is None:
+                    return [TextContent(type="text", text="Session tokens not enabled")]
+                result = await _request_session_token(
+                    _current_agent, auth, token_manager,
+                    arguments["document"],
+                    arguments["permissions"],
+                    ttl_seconds=arguments.get("ttl_seconds", 300),
+                )
             else:
                 return [TextContent(type="text", text=f"Unknown tool: {name}")]
 
diff --git a/tests/unit/test_server.py b/tests/unit/test_server.py
index 218e801..ede381f 100644
--- a/tests/unit/test_server.py
+++ b/tests/unit/test_server.py
@@ -53,5 +53,48 @@ tokens:
     assert "modify_column" in tool_names
     assert "delete_column" in tool_names
 
-    # Should have all 12 tools
-    assert len(result.root.tools) == 12
+    # Session tools (always registered)
+    assert "get_proxy_documentation" in tool_names
+    assert "request_session_token" in tool_names
+
+    # Should have all 14 tools
+    assert len(result.root.tools) == 14
+
+
+@pytest.mark.asyncio
+async def test_create_server_registers_session_tools(tmp_path):
+    from grist_mcp.session import SessionTokenManager
+
+    config_file = tmp_path / "config.yaml"
+    config_file.write_text("""
+documents:
+  test-doc:
+    url: https://grist.example.com
+    doc_id: abc123
+    api_key: test-key
+
+tokens:
+  - token: valid-token
+    name: test-agent
+    scope:
+      - document: test-doc
+        permissions: [read, write, schema]
+""")
+
+    config = load_config(str(config_file))
+    auth = Authenticator(config)
+    agent = auth.authenticate("valid-token")
+    token_manager = SessionTokenManager()
+    server = create_server(auth, agent, token_manager)
+
+    # Get the list_tools handler and call it
+    handler = server.request_handlers.get(ListToolsRequest)
+    assert handler is not None
+
+    req = ListToolsRequest(method="tools/list")
+    result = await handler(req)
+
+    tool_names = [t.name for t in result.root.tools]
+
+    assert "get_proxy_documentation" in tool_names
+    assert "request_session_token" in tool_names