GitHubHelp

LucasSilvaFerreira/colab-agent-bridge

★ 0Forks 0PythonGitHub ↗Compare

README

🚀 Colab Agent Bridge (colab-agent-bridge)

License: Apache 2.0 Python 3.11+ MCP Standard

A seamless bridge connecting local AI coding agents (Google Antigravity, Gemini CLI, Claude Code, Cursor, Windsurf) and local terminal scripts to persistent Google Colab cloud sessions via the Model Context Protocol (googlecolab/colab-mcp).

🌟 Why Colab Agent Bridge?

Challenge with Raw Colab MCP How Colab Agent Bridge Solves It
New Tabs Opened Repeatedly: One-shot tools trigger open <url> on every execution. Zero New Tabs: Runs a persistent background daemon that reuses your open Colab tab.
Lost In-Memory State: One-off scripts disconnect and lose session continuity. Shared Kernel Memory: Variables, loaded models, and imports remain in memory across calls.
Silent Disconnections: Agents fail if the browser tab hasn't finished connecting. Automatic Tab Verification: Automatically detects connection status (--check) and guides you.
Complex Setup: Manual socket configuration and proxy management. 1-Click Agent Skill & CLI: Native integration with Antigravity and a simple colab-exec CLI.

🏗️ Architecture

┌─────────────────────────────────────────────────────────────┐
│ Google Colab Web Notebook (Browser Tab - Opened ONCE)       │
└──────────────────────────┬──────────────────────────────────┘
                           │ Persistent WebSocket (ws://localhost:56595)
┌──────────────────────────▼──────────────────────────────────┐
│ Colab Bridge Daemon (colab-daemon)                          │
│ Maintains persistent MCP transport & shared kernel session  │
└──────────────────────────▲──────────────────────────────────┘
                           │ Fast Local Control API (http://localhost:56594)
┌──────────────────────────┴──────────────────────────────────┐
│ AI Coding Agents (Antigravity, Claude, Gemini) / CLI        │
│ Executes code on-demand in ~2s with zero browser disruption │
└─────────────────────────────────────────────────────────────┘

⚡ Quick Start

1. Installation

Clone and install dependencies via uv (recommended) or pip:

git clone https://github.com/LucasSilvaFerreira/colab-agent-bridge.git
cd colab-agent-bridge
./install.sh

Or install in editable mode:

uv pip install -e .
# or
pip install -e .

2. Verify Connection

Check if your Colab tab is currently connected:

colab-exec --check

If no tab is open, the tool will provide the connection URL and can open it once in your browser.

3. Execute Code in Colab

# Inline snippet
colab-exec -c "import torch; print('CUDA Available:', torch.cuda.is_available())"

# Execute a Python script
colab-exec -f examples/02_gpu_benchmark.py

# Inspect all cells in the active notebook
colab-exec --cells

⚡ Real Benchmark (CPU vs. NVIDIA L4 GPU)

Tested on Google Colab with examples/02_gpu_benchmark.py ($4000 \times 4000$ matrix multiplication):

  • Standard CPU: 2.008s
  • NVIDIA L4 GPU (24 GB VRAM): 0.201s (~10x speedup)

🤖 Using as an AI Agent Skill (Antigravity & Gemini CLI)

This repository includes a ready-to-use skill in .agents/skills/colab-mcp/:

Installing the Skill

  • For Current Workspace: Automatically discovered in .agents/skills/colab-mcp/.
  • For All Workspaces (Global): 
    mkdir -p ~/.gemini/config/skills/colab-mcp
cp -r .agents/skills/colab-mcp/* ~/.gemini/config/skills/colab-mcp/

Sample Agent Prompts

Once installed, your agent will automatically activate the skill whenever you ask:

  • "Run this PyTorch training benchmark in Google Colab"
  • "Check available GPU VRAM in Colab"
  • "Add a cell to my open Colab notebook and display the results"

🐍 Programmatic Python API

You can also use colab_bridge in your own Python automation scripts:

from colab_bridge import execute, check_status

# 1. Verify connection
connected, status = check_status()
print(f"Colab Connected: {connected}")

# 2. Execute code
success, output, result = execute("""
import os, platform
print(f"Running on: {platform.platform()}")
print(f"Storage: {os.statvfs('/').f_bavail * 4096 / (1024**3):.2f} GB free")
""")

print(output)

🛠️ Native Colab MCP Tools Exposed

Under the hood, Google Colab provides 7 native MCP tools managed by the bridge:

Tool Parameters Description
add_code_cell cellIndex, language, code Inserts a new code cell and returns newCellId.
add_text_cell cellIndex, content Inserts a Markdown/LaTeX cell.
run_code_cell cellId Runs the cell in the cloud kernel and captures stdout/stderr.
get_cells cellIndexStart, cellIndexEnd, includeOutputs Retrieves cell content and execution streams.
update_cell cellId, content Updates existing cell code or text.
delete_cell cellId Deletes a cell by ID.
move_cell cellId, cellIndex Re-orders cells in the notebook.

📁 Repository Structure

colab-agent-bridge/
├── .agents/skills/colab-mcp/          # AI Agent Skill (Antigravity, Gemini CLI)
│   ├── SKILL.md                       # Skill definition & execution runbook
│   ├── references/
│   │   └── colab_tools.md             # Complete Colab MCP tools reference
│   └── scripts/
│       ├── colab_daemon.py            # WebSocket + HTTP control daemon
│       ├── colab_exec.py              # CLI client
│       └── run_colab_code.py          # Standalone runner
├── src/colab_bridge/                  # Core Python package
│   ├── __init__.py                    # Public API exports
│   ├── daemon.py                      # Daemon engine
│   ├── client.py                      # Python client functions
│   └── cli.py                         # CLI entry point (colab-exec)
├── examples/                          # Ready-to-run examples
│   ├── 01_hello_colab.py              # System & directory check
│   ├── 02_gpu_benchmark.py            # PyTorch CUDA tensor benchmark
│   └── 03_shared_state.py             # Shared kernel state verification
├── install.sh                         # 1-step installer
├── pyproject.toml                     # Package configuration
└── README.md                          # Documentation

📄 License

Apache License 2.0. See LICENSE for details.

Contributors

LucasSilvaFerreira

Issues