"""Gemini Chat MCP - Simple MCP server for chatting with Gemini and custom Gems."""
import asyncio
import logging
import os
from pathlib import Path
from typing import Any, Optional
import httpx
from dotenv import load_dotenv
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent, CallToolResult
from gemini_webapi import GeminiClient, ChatSession, set_log_level
from gemini_webapi.constants import Model
try:
from .cookie_sync import CookieSyncServer, get_cookie_store
except ImportError:
from cookie_sync import CookieSyncServer, get_cookie_store
# Setup logging to file
LOG_DIR = Path(__file__).parent.parent.parent / "logs"
LOG_DIR.mkdir(exist_ok=True)
LOG_FILE = LOG_DIR / "gemini-chat-mcp.log"
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
handlers=[
logging.FileHandler(LOG_FILE, encoding="utf-8"),
logging.StreamHandler()
]
)
logger = logging.getLogger("gemini-chat-mcp")
logger.info(f"Logging to {LOG_FILE}")
set_log_level("WARNING")
load_dotenv()
# Constants
MODEL_MAP = {
"gemini-3-pro": Model.G_3_0_PRO,
"gemini-2.5-flash": Model.G_2_5_FLASH,
}
DEFAULT_MODEL = Model.G_3_0_PRO
# Global state
_client: Optional[GeminiClient] = None
_need_reinit = False
_callback_registered = False
# Conversation session management
# Key format: "gem_id:model" or "default:model" for no gem
# Value: ChatSession object
_chat_sessions: dict[str, ChatSession] = {}
# Persistent storage for conversation metadata
_sessions_file: Optional[Path] = None
def _get_sessions_file() -> Path:
"""Get path to sessions persistence file."""
global _sessions_file
if _sessions_file is None:
_sessions_file = get_project_root() / ".gemini_sessions.json"
return _sessions_file
def _save_sessions_to_file() -> None:
"""Save all session metadata to file for persistence."""
sessions_data = {}
for key, session in _chat_sessions.items():
if session.cid: # Only save if we have a valid cid
sessions_data[key] = {
"cid": session.cid,
"rid": session.rid,
"rcid": session.rcid,
"gem": session.gem if isinstance(session.gem, str) else None,
}
try:
with open(_get_sessions_file(), "w", encoding="utf-8") as f:
import json
json.dump(sessions_data, f, indent=2)
logger.debug(f"Saved {len(sessions_data)} sessions to file")
except Exception as e:
logger.warning(f"Failed to save sessions: {e}")
def _load_sessions_from_file(client: GeminiClient) -> None:
"""Load session metadata from file and recreate ChatSession objects."""
global _chat_sessions
sessions_file = _get_sessions_file()
if not sessions_file.exists():
return
try:
with open(sessions_file, "r", encoding="utf-8") as f:
import json
sessions_data = json.load(f)
for key, data in sessions_data.items():
# Recreate ChatSession with full metadata
metadata = [data.get("cid"), data.get("rid"), data.get("rcid")]
# Parse model from key
parts = key.split(":")
model_str = parts[1] if len(parts) > 1 else "gemini-3-pro"
session = client.start_chat(
metadata=metadata,
gem=data.get("gem"),
model=get_model(model_str)
)
_chat_sessions[key] = session
logger.info(f"Loaded {len(sessions_data)} sessions from file")
except Exception as e:
logger.warning(f"Failed to load sessions: {e}")
server = Server("gemini-chat-mcp")
def _get_session_key(gem_id: Optional[str], model: str) -> str:
"""Generate a unique key for session storage."""
gem_key = gem_id or "default"
return f"{gem_key}:{model}"
def on_cookie_change(name: str, value: str) -> None:
"""Handle cookie updates from Chrome extension - mark client for reinit."""
global _need_reinit, _chat_sessions
logger.info(f"Cookie changed: {name}, will reinitialize client on next request")
_need_reinit = True
# Clear all sessions on cookie change since they'll be invalid
_chat_sessions.clear()
def register_cookie_callback() -> None:
"""Register callback to handle cookie changes."""
global _callback_registered
if not _callback_registered:
cookie_store = get_cookie_store()
cookie_store.on_change(on_cookie_change)
_callback_registered = True
logger.debug("Cookie change callback registered")
def get_project_root() -> Path:
return Path(__file__).parent.parent.parent
def update_env_file(key: str, value: str) -> None:
"""Update a key in the .env file."""
env_path = get_project_root() / ".env"
lines = []
key_found = False
if env_path.exists():
with open(env_path, "r", encoding="utf-8") as f:
lines = f.readlines()
for i, line in enumerate(lines):
if line.startswith(f"{key}="):
lines[i] = f"{key}={value}\n"
key_found = True
break
if not key_found:
lines.append(f"{key}={value}\n")
with open(env_path, "w", encoding="utf-8") as f:
f.writelines(lines)
COOKIE_SYNC_URL = "http://127.0.0.1:8765/cookies"
async def fetch_cookies_from_sync_server() -> Optional[dict[str, str]]:
"""Fetch cookies from existing cookie sync server (e.g., nano-banana-mcp)."""
try:
async with httpx.AsyncClient(timeout=2.0) as client:
resp = await client.get(COOKIE_SYNC_URL)
if resp.status_code == 200:
data = resp.json()
if data.get("success") and data.get("cookies"):
cookies = data["cookies"]
psid = cookies.get("__Secure-1PSID")
psidts = cookies.get("__Secure-1PSIDTS")
if psid and psidts:
logger.info("Fetched cookies from existing sync server")
return {"psid": psid, "psidts": psidts}
except Exception as e:
logger.debug(f"Could not fetch from sync server: {e}")
return None
async def get_cookies() -> Optional[dict[str, str]]:
"""Get cookies from cookie store, sync server, or environment."""
cookie_store = get_cookie_store()
# Try local cookie store first (from extension direct to this server)
psid = cookie_store.get_psid()
psidts = cookie_store.get_psidts()
if psid and psidts:
logger.debug("Using cookies from local cookie store")
return {"psid": psid, "psidts": psidts}
# Try fetching from existing sync server (e.g., nano-banana-mcp)
cookies = await fetch_cookies_from_sync_server()
if cookies:
return cookies
# Fall back to environment
psid = os.getenv("GEMINI_SECURE_1PSID")
psidts = os.getenv("GEMINI_SECURE_1PSIDTS")
if psid and psidts:
logger.debug("Using cookies from environment")
return {"psid": psid, "psidts": psidts}
return None
async def get_client() -> GeminiClient:
"""Get or create Gemini client."""
global _client, _need_reinit
if _client and not _need_reinit:
return _client
# Close old client if reinitializing
if _client and _need_reinit:
try:
await _client.close()
logger.info("Closed old Gemini client for reinitialization")
except Exception as e:
logger.warning(f"Error closing old client: {e}")
_client = None
cookies = await get_cookies()
if not cookies:
raise ValueError("No cookies. Set GEMINI_SECURE_1PSID/PSIDTS in .env or use Chrome extension.")
_client = GeminiClient(cookies["psid"], cookies["psidts"])
await _client.init(timeout=30, auto_close=False, auto_refresh=True)
_need_reinit = False
# Load saved sessions after client init
_load_sessions_from_file(_client)
logger.info("Gemini client initialized")
return _client
def get_model(model_str: Optional[str]) -> Model:
"""Map model string to Model enum."""
if not model_str:
return DEFAULT_MODEL
return MODEL_MAP.get(model_str.lower(), DEFAULT_MODEL)
def get_model_str(model_str: Optional[str]) -> str:
"""Get normalized model string."""
if not model_str:
return "gemini-3-pro"
return model_str.lower() if model_str.lower() in MODEL_MAP else "gemini-3-pro"
def parse_conversation_url(url: str) -> Optional[dict[str, str]]:
"""Parse Gemini conversation URL to extract gem_id and conversation_id.
URL format: https://gemini.google.com/gem/{gem_id}/{conversation_id}
or: https://gemini.google.com/app/{conversation_id}
"""
if not url:
return None
try:
# Remove trailing slashes and extract path
url = url.rstrip("/")
if "/gem/" in url:
# Format: /gem/{gem_id}/{conversation_id}
parts = url.split("/gem/")[1].split("/")
if len(parts) >= 2:
return {"gem_id": parts[0], "conversation_id": parts[1]}
elif len(parts) == 1:
return {"gem_id": parts[0]}
elif "/app/" in url:
# Format: /app/{conversation_id}
parts = url.split("/app/")[1].split("/")
if parts:
return {"conversation_id": parts[0]}
except Exception:
pass
return None
# Tools definition
TOOLS = [
Tool(
name="chat",
description="""Chat with Gemini - can generate text and images.
By default, continues the existing conversation (multi-turn).
Optionally chat with a custom Gem by providing gem_id.
To resume a specific conversation, provide conversation_id or full URL (the gem_id will be auto-extracted from URL).""",
inputSchema={
"type": "object",
"properties": {
"message": {"type": "string", "description": "Message to send to Gemini"},
"conversation_id": {
"type": "string",
"description": "Optional: Conversation ID to resume. Can be just the ID or full Gemini URL (https://gemini.google.com/gem/xxx/yyy)"
},
"gem_id": {
"type": "string",
"description": "Optional: Custom Gem ID to chat with (e.g., '10wwoIn-UvdTVzg4rFH3szi4wk8oTo8Cf'). Not needed if conversation_id is a full URL."
},
"model": {
"type": "string",
"description": "Optional: Model to use. Options: 'gemini-3-pro' (default), 'gemini-2.5-flash'.",
"enum": ["gemini-3-pro", "gemini-2.5-flash"],
"default": "gemini-3-pro"
}
},
"required": ["message"]
}
),
Tool(
name="new_conversation",
description="Start a new conversation with Gemini, discarding the current conversation history. Use this when you want to start fresh.",
inputSchema={
"type": "object",
"properties": {
"message": {"type": "string", "description": "First message for the new conversation"},
"gem_id": {
"type": "string",
"description": "Optional: Custom Gem ID to chat with"
},
"model": {
"type": "string",
"description": "Optional: Model to use. Options: 'gemini-3-pro' (default), 'gemini-2.5-flash'.",
"enum": ["gemini-3-pro", "gemini-2.5-flash"],
"default": "gemini-3-pro"
}
},
"required": ["message"]
}
),
Tool(
name="get_conversation_info",
description="Get information about the current conversation session, including conversation ID and URL.",
inputSchema={
"type": "object",
"properties": {
"gem_id": {
"type": "string",
"description": "Optional: Gem ID to check. If not provided, checks the default conversation."
},
"model": {
"type": "string",
"description": "Optional: Model to check. Default is gemini-3-pro.",
"enum": ["gemini-3-pro", "gemini-2.5-flash"],
"default": "gemini-3-pro"
}
},
"required": []
}
),
Tool(
name="set_cookie",
description="Update __Secure-1PSIDTS cookie value. Use this when cookie expires. Just paste the new PSIDTS value.",
inputSchema={
"type": "object",
"properties": {
"psidts": {"type": "string", "description": "The new __Secure-1PSIDTS cookie value (starts with 'sidts-')"},
"psid": {"type": "string", "description": "Optional: The __Secure-1PSID cookie value (starts with 'g.a'). Only needed if PSID also expired."}
},
"required": ["psidts"]
}
),
]
@server.list_tools()
async def list_tools() -> list[Tool]:
return TOOLS
async def handle_chat(arguments: dict[str, Any], force_new: bool = False) -> CallToolResult:
"""Handle chat tool with conversation persistence."""
global _chat_sessions
try:
client = await get_client()
gem_id = arguments.get("gem_id")
model_str = get_model_str(arguments.get("model"))
model = get_model(arguments.get("model"))
message = arguments["message"]
conversation_id = arguments.get("conversation_id")
# Parse conversation_id if it's a URL
if conversation_id and conversation_id.startswith("http"):
parsed = parse_conversation_url(conversation_id)
logger.info(f"Parsed URL: {parsed}")
if parsed:
conversation_id = parsed.get("conversation_id")
# If gem_id not explicitly provided but found in URL, use it
if not gem_id and parsed.get("gem_id"):
gem_id = parsed["gem_id"]
logger.info(f"Using gem_id={gem_id}, conversation_id={conversation_id}, model={model_str}")
session_key = _get_session_key(gem_id, model_str)
chat_session: Optional[ChatSession] = None
logger.info(f"Current sessions in memory: {list(_chat_sessions.keys())}")
# Determine which session to use
if force_new:
# Force new conversation - remove existing session
if session_key in _chat_sessions:
del _chat_sessions[session_key]
logger.info(f"Starting new conversation (force_new=True)")
elif conversation_id:
# Resume specific conversation by ID
# Note: We only have cid from URL. rid and rcid will be None.
# The conversation will continue but won't have full history context.
chat_session = client.start_chat(
cid=conversation_id,
gem=gem_id,
model=model
)
_chat_sessions[session_key] = chat_session
logger.info(f"Resuming conversation with cid={conversation_id}, metadata={chat_session.metadata}")
elif session_key in _chat_sessions:
# Reuse existing session
chat_session = _chat_sessions[session_key]
logger.info(f"Continuing existing conversation: {chat_session.cid}")
# Send message
if chat_session:
logger.info(f"Sending via existing session, cid={chat_session.cid}")
response = await chat_session.send_message(message)
else:
# First message - create new session
chat_session = client.start_chat(gem=gem_id, model=model)
response = await chat_session.send_message(message)
_chat_sessions[session_key] = chat_session
logger.info(f"Created new conversation: {chat_session.cid}")
# Save session metadata after each message for persistence
_save_sessions_to_file()
logger.debug(f"Session metadata after send: cid={chat_session.cid}, rid={chat_session.rid}, rcid={chat_session.rcid}")
text = response.text or ""
# Build conversation info
cid = chat_session.cid if chat_session else None
if cid:
if gem_id:
conv_url = f"https://gemini.google.com/gem/{gem_id}/{cid}"
else:
conv_url = f"https://gemini.google.com/app/{cid}"
text = f"[Conversation: {conv_url}]\n\n{text}"
# Save images if any
if response.images:
output_dir = get_project_root() / "output"
output_dir.mkdir(exist_ok=True)
paths = []
for i, img in enumerate(response.images):
name = f"chat_{i}.png"
path = output_dir / name
c = 1
while path.exists():
name = f"chat_{i}_{c}.png"
path = output_dir / name
c += 1
await img.save(path=str(output_dir), filename=name, verbose=False)
paths.append(str(path))
text += "\n\nGenerated images:\n" + "\n".join(f"• {p}" for p in paths)
return CallToolResult(content=[TextContent(type="text", text=text or "No response")])
except Exception as e:
import traceback
error_detail = traceback.format_exc()
logger.error(f"Chat failed: {e}\n{error_detail}")
return CallToolResult(content=[TextContent(type="text", text=f"Chat failed: {e}")], isError=True)
async def handle_new_conversation(arguments: dict[str, Any]) -> CallToolResult:
"""Handle new_conversation tool - starts a fresh conversation."""
return await handle_chat(arguments, force_new=True)
async def handle_get_conversation_info(arguments: dict[str, Any]) -> CallToolResult:
"""Handle get_conversation_info tool."""
gem_id = arguments.get("gem_id")
model_str = get_model_str(arguments.get("model"))
session_key = _get_session_key(gem_id, model_str)
if session_key not in _chat_sessions:
return CallToolResult(content=[TextContent(
type="text",
text=f"No active conversation for gem_id='{gem_id or 'default'}', model='{model_str}'"
)])
chat_session = _chat_sessions[session_key]
cid = chat_session.cid
rid = chat_session.rid
rcid = chat_session.rcid
if gem_id:
conv_url = f"https://gemini.google.com/gem/{gem_id}/{cid}"
else:
conv_url = f"https://gemini.google.com/app/{cid}"
info = f"""**Active Conversation**
- URL: {conv_url}
- Conversation ID (cid): {cid}
- Reply ID (rid): {rid}
- Reply Candidate ID (rcid): {rcid}
- Gem ID: {gem_id or 'None (default)'}
- Model: {model_str}
To resume this conversation later, use:
- conversation_id: "{cid}"
- or URL: "{conv_url}"
"""
return CallToolResult(content=[TextContent(type="text", text=info)])
async def handle_set_cookie(arguments: dict[str, Any]) -> CallToolResult:
"""Handle set_cookie tool."""
global _need_reinit, _chat_sessions
psidts = arguments["psidts"]
psid = arguments.get("psid")
if not psidts.startswith("sidts-"):
return CallToolResult(
content=[TextContent(type="text", text="Invalid PSIDTS. Should start with 'sidts-'")],
isError=True
)
if psid and not psid.startswith("g.a"):
return CallToolResult(
content=[TextContent(type="text", text="Invalid PSID. Should start with 'g.a'")],
isError=True
)
try:
# Update cookie store
cookie_store = get_cookie_store()
cookie_store.set("__Secure-1PSIDTS", psidts)
if psid:
cookie_store.set("__Secure-1PSID", psid)
# Also update env file
update_env_file("GEMINI_SECURE_1PSIDTS", psidts)
os.environ["GEMINI_SECURE_1PSIDTS"] = psidts
if psid:
update_env_file("GEMINI_SECURE_1PSID", psid)
os.environ["GEMINI_SECURE_1PSID"] = psid
_need_reinit = True
# Clear sessions on cookie change
_chat_sessions.clear()
return CallToolResult(content=[TextContent(type="text", text="Cookie updated! Next request will use new cookie.")])
except Exception as e:
return CallToolResult(content=[TextContent(type="text", text=f"Failed: {e}")], isError=True)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult:
"""Route tool calls."""
handlers = {
"chat": handle_chat,
"new_conversation": handle_new_conversation,
"get_conversation_info": handle_get_conversation_info,
"set_cookie": handle_set_cookie,
}
if handler := handlers.get(name):
return await handler(arguments)
return CallToolResult(content=[TextContent(type="text", text=f"Unknown tool: {name}")], isError=True)
async def run_server():
# Register cookie change callback
register_cookie_callback()
# Start cookie sync server
cookie_sync = CookieSyncServer(host="127.0.0.1", port=8765)
async with cookie_sync:
logger.info("Cookie sync server running at http://127.0.0.1:8765")
async with stdio_server() as (read_stream, write_stream):
await server.run(read_stream, write_stream, server.create_initialization_options())
def main():
asyncio.run(run_server())
if __name__ == "__main__":
main()
