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.
This commit is contained in:
@@ -6,6 +6,9 @@ from mcp.server import Server
|
|||||||
from mcp.types import Tool, TextContent
|
from mcp.types import Tool, TextContent
|
||||||
|
|
||||||
from grist_mcp.auth import Authenticator, Agent, AuthError
|
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.discovery import list_documents as _list_documents
|
||||||
from grist_mcp.tools.read import list_tables as _list_tables
|
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
|
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.
|
"""Create and configure the MCP server for an authenticated agent.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
auth: Authenticator instance for permission checks.
|
auth: Authenticator instance for permission checks.
|
||||||
agent: The authenticated agent for this server instance.
|
agent: The authenticated agent for this server instance.
|
||||||
|
token_manager: Optional session token manager for HTTP proxy access.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Configured MCP Server instance.
|
Configured MCP Server instance.
|
||||||
@@ -203,6 +207,34 @@ def create_server(auth: Authenticator, agent: Agent) -> Server:
|
|||||||
"required": ["document", "table", "column_id"],
|
"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()
|
@server.call_tool()
|
||||||
@@ -265,6 +297,17 @@ def create_server(auth: Authenticator, agent: Agent) -> Server:
|
|||||||
_current_agent, auth, arguments["document"], arguments["table"],
|
_current_agent, auth, arguments["document"], arguments["table"],
|
||||||
arguments["column_id"],
|
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:
|
else:
|
||||||
return [TextContent(type="text", text=f"Unknown tool: {name}")]
|
return [TextContent(type="text", text=f"Unknown tool: {name}")]
|
||||||
|
|
||||||
|
|||||||
@@ -53,5 +53,48 @@ tokens:
|
|||||||
assert "modify_column" in tool_names
|
assert "modify_column" in tool_names
|
||||||
assert "delete_column" in tool_names
|
assert "delete_column" in tool_names
|
||||||
|
|
||||||
# Should have all 12 tools
|
# Session tools (always registered)
|
||||||
assert len(result.root.tools) == 12
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user