Grok CLI MCP Server

"""FastMCP server for grok-cli-mcp.""" from __future__ import annotations import json import sys from typing import Optional from mcp.server.fastmcp.server import Context, FastMCP from .types import GrokMessage from .utils import _collect_assistant_text, _run_grok # Server server = FastMCP( name="grok-mcp", instructions=( "This MCP server exposes the Grok CLI via tools for general queries, chat-style prompts, " "and code-oriented tasks. Use grok_query for generic prompts, grok_chat for role-based " "messages, and grok_code when asking for code with optional language/context. " "You can optionally pass a Grok model name. Requires GROK_API_KEY." ), debug=False, log_level="INFO", ) @server.tool( name="grok_query", title="Grok Query", description=( "Send a single prompt to Grok via CLI headless mode. Returns the assistant's text. " "Use raw_output=true to get raw CLI output and parsed messages." ), ) async def grok_query( prompt: str, model: Optional[str] = None, raw_output: bool = False, timeout_s: float = 120.0, ctx: Optional[Context] = None, ) -> str | dict: """ Send a single prompt to Grok. Args: prompt: The user prompt. model: Optional Grok model name (passed with -m if provided). raw_output: If true, returns {text, messages, raw, model}. timeout_s: Process timeout in seconds. ctx: FastMCP context. Returns: Assistant's text response, or dict with full details if raw_output=True. """ result = await _run_grok(prompt, model=model, timeout_s=timeout_s, ctx=ctx) # Collate assistant text assistant_text = _collect_assistant_text(result.messages) if result.messages else (result.raw or "") if raw_output: return { "text": assistant_text, "messages": [m.model_dump() for m in result.messages], "raw": result.raw, "model": result.model, } return assistant_text @server.tool( name="grok_chat", title="Grok Chat", description=( "Send a list of role/content messages to Grok by flattening into a single prompt. " "Useful for multi-turn context when the CLI only supports a single '-p' prompt." ), ) async def grok_chat( messages: list[GrokMessage], model: Optional[str] = None, raw_output: bool = False, timeout_s: float = 120.0, ctx: Optional[Context] = None, ) -> str | dict: """ Send multi-turn conversation to Grok. Args: messages: Array of {role, content}; content may be string or structured. model: Optional Grok model name. raw_output: If true, returns structured output. timeout_s: Process timeout in seconds. ctx: FastMCP context. Returns: Assistant's text response, or dict with full details if raw_output=True. """ # Flatten messages to a single prompt with role tags # This preserves some conversational context even if CLI accepts only single prompt. prompt_lines: list[str] = [] for m in messages: content_str: str if isinstance(m.content, str): content_str = m.content else: try: content_str = json.dumps(m.content, ensure_ascii=False) except Exception: content_str = str(m.content) prompt_lines.append(f"{m.role.capitalize()}: {content_str}") prompt = "\n".join(prompt_lines) result = await _run_grok(prompt, model=model, timeout_s=timeout_s, ctx=ctx) assistant_text = _collect_assistant_text(result.messages) if result.messages else (result.raw or "") if raw_output: return { "text": assistant_text, "messages": [m.model_dump() for m in result.messages], "raw": result.raw, "model": result.model, } return assistant_text @server.tool( name="grok_code", title="Grok Code Task", description=( "Ask Grok for code or code-related guidance. You can provide a language hint and context " "(e.g., file snippets or requirements). Returns assistant text by default." ), ) async def grok_code( task: str, language: Optional[str] = None, context: Optional[str] = None, model: Optional[str] = None, raw_output: bool = False, timeout_s: float = 180.0, ctx: Optional[Context] = None, ) -> str | dict: """ Ask Grok for code generation or guidance. Args: task: Description of what code/help you need. language: Optional language hint (e.g., 'python', 'typescript'). context: Optional context (repo constraints, file snippets, tests, etc.). model: Optional Grok model name. raw_output: If true, returns structured output. timeout_s: Process timeout in seconds. ctx: FastMCP context. Returns: Assistant's text response, or dict with full details if raw_output=True. """ sys_instructions = [ "You are an expert software engineer.", "Respond with clear, correct, directly usable code and concise explanations.", "Prefer minimal dependencies and explain tradeoffs when relevant.", ] if language: sys_instructions.append(f"Primary language: {language}") if context: sys_instructions.append("Context:\n" + context.strip()) prompt = "\n\n".join( [ "\n".join(sys_instructions), "Task:", task.strip(), ] ) result = await _run_grok(prompt, model=model, timeout_s=timeout_s, ctx=ctx) assistant_text = _collect_assistant_text(result.messages) if result.messages else (result.raw or "") if raw_output: return { "text": assistant_text, "messages": [m.model_dump() for m in result.messages], "raw": result.raw, "model": result.model, } return assistant_text def main() -> None: """Run the MCP server with stdio transport.""" try: server.run("stdio") except KeyboardInterrupt: sys.exit(130)