From f716e5d37ef0250b9b7243a9a3d828f5bad019f0 Mon Sep 17 00:00:00 2001
From: Bill <bill@bballou.com>
Date: Wed, 3 Dec 2025 15:07:06 -0500
Subject: [PATCH] fix: implement token-based authentication at server startup

- Server now authenticates from GRIST_MCP_TOKEN env var or token parameter
- Removed unused code (_set_agent, nonlocal check)
- Added AuthError handling in main.py
- Updated test to pass token explicitly
---
 src/grist_mcp/main.py   |  7 ++++++-
 src/grist_mcp/server.py | 35 ++++++++++++++++++++---------------
 tests/test_server.py    |  2 +-
 3 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/src/grist_mcp/main.py b/src/grist_mcp/main.py
index a7584f1..1355f95 100644
--- a/src/grist_mcp/main.py
+++ b/src/grist_mcp/main.py
@@ -7,6 +7,7 @@ import sys
 from mcp.server.stdio import stdio_server
 
 from grist_mcp.server import create_server
+from grist_mcp.auth import AuthError
 
 
 async def main():
@@ -16,7 +17,11 @@ async def main():
         print(f"Error: Config file not found at {config_path}", file=sys.stderr)
         sys.exit(1)
 
-    server = create_server(config_path)
+    try:
+        server = create_server(config_path)
+    except AuthError as e:
+        print(f"Authentication error: {e}", file=sys.stderr)
+        sys.exit(1)
 
     async with stdio_server() as (read_stream, write_stream):
         await server.run(read_stream, write_stream, server.create_initialization_options())
diff --git a/src/grist_mcp/server.py b/src/grist_mcp/server.py
index 263f707..c9dc9dd 100644
--- a/src/grist_mcp/server.py
+++ b/src/grist_mcp/server.py
@@ -1,7 +1,9 @@
 """MCP server setup and tool registration."""
 
+import json
+import os
+
 from mcp.server import Server
-from mcp.server.stdio import stdio_server
 from mcp.types import Tool, TextContent
 
 from grist_mcp.config import load_config
@@ -21,14 +23,27 @@ 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(config_path: str) -> Server:
-    """Create and configure the MCP server."""
+def create_server(config_path: str, token: str | None = None) -> Server:
+    """Create and configure the MCP server.
+
+    Args:
+        config_path: Path to the configuration YAML file.
+        token: Agent token for authentication. If not provided, reads from
+               GRIST_MCP_TOKEN environment variable.
+
+    Raises:
+        AuthError: If token is invalid or not provided.
+    """
     config = load_config(config_path)
     auth = Authenticator(config)
     server = Server("grist-mcp")
 
-    # Current agent context (set during authentication)
-    _current_agent: Agent | None = None
+    # Authenticate agent from token (required for all tool calls)
+    auth_token = token or os.environ.get("GRIST_MCP_TOKEN")
+    if not auth_token:
+        raise AuthError("No token provided. Set GRIST_MCP_TOKEN environment variable.")
+
+    _current_agent: Agent = auth.authenticate(auth_token)
 
     @server.list_tools()
     async def list_tools() -> list[Tool]:
@@ -203,11 +218,6 @@ def create_server(config_path: str) -> Server:
 
     @server.call_tool()
     async def call_tool(name: str, arguments: dict) -> list[TextContent]:
-        nonlocal _current_agent
-
-        if _current_agent is None:
-            return [TextContent(type="text", text="Error: Not authenticated")]
-
         try:
             if name == "list_documents":
                 result = await _list_documents(_current_agent)
@@ -269,7 +279,6 @@ def create_server(config_path: str) -> Server:
             else:
                 return [TextContent(type="text", text=f"Unknown tool: {name}")]
 
-            import json
             return [TextContent(type="text", text=json.dumps(result))]
 
         except AuthError as e:
@@ -277,8 +286,4 @@ def create_server(config_path: str) -> Server:
         except Exception as e:
             return [TextContent(type="text", text=f"Error: {e}")]
 
-    # Store auth for external access
-    server._auth = auth
-    server._set_agent = lambda agent: setattr(server, '_current_agent', agent) or setattr(type(server), '_current_agent', agent)
-
     return server
diff --git a/tests/test_server.py b/tests/test_server.py
index e6e4145..8c883b5 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -21,7 +21,7 @@ tokens:
         permissions: [read, write, schema]
 """)
 
-    server = create_server(str(config_file))
+    server = create_server(str(config_file), token="test-token")
 
     # Server should have tools registered
     assert server is not None