Custom Claude templates and MCP

from pathlib import Path

from gravixlayer import GravixLayer, TemplateBuilder

client = GravixLayer()

dockerfile = Path(
    "examples/templates/dockerfiles/agent-claude.Dockerfile"
).read_text()

builder = (
    TemplateBuilder("agent-claude", "Claude Code on Gravix Layer base")
    .dockerfile(dockerfile)
    .vcpu(2)
    .memory(2048)
    .disk(8192)
)

status = client.templates.build_and_wait(builder, timeout_secs=1800)
print(status.template_id, status.status)

# Use either the template name or template_id when creating runtimes:
# client.runtime.create(template=status.template_id, ...)

export GRAVIXLAYER_CLAUDE_TEMPLATE="agent-claude"
# or the UUID printed as status.template_id

from gravixlayer import GravixLayer, TemplateBuilder

client = GravixLayer()

builder = (
    TemplateBuilder("my-claude", "Claude Code + pytest + project CLAUDE.md")
    .from_image("your-registry.example.com/agent-claude:latest")
    .vcpu(2)
    .memory(2048)
    .disk(8192)
    .pip_install("pytest", "ruff", "httpx")
    .copy_file(
        "# Always prefer small PRs and add tests for new endpoints.\n",
        "/workspace/CLAUDE.md",
    )
)

status = client.templates.build_and_wait(builder, timeout_secs=1800)
print(status.template_id, status.status)

runtime.file.create_directory("/workspace/.claude/skills/review")
runtime.file.write(
    "/workspace/.claude/skills/review/SKILL.md",
    """---
name: review
description: Prefer security and test coverage in PR reviews
---
When reviewing, always check authz, input validation, and missing tests.
""",
)

import os
import uuid

from gravixlayer import GravixLayer

TEMPLATE = os.environ["GRAVIXLAYER_CLAUDE_TEMPLATE"]
ANTHROPIC_API_KEY = os.environ["ANTHROPIC_API_KEY"]
MCP_URL = os.environ["MCP_URL"]  # e.g. https://mcp.example.com/mcp
MCP_TOKEN = os.environ["MCP_TOKEN"]
MCP_HOST = os.environ["MCP_HOST"]  # e.g. mcp.example.com

client = GravixLayer()

provider = client.identity.providers.create(
    name=f"anthropic-mcp-{uuid.uuid4().hex[:8]}",
    provider_type="api_key",
    secrets=[
        {"key": "ANTHROPIC_API_KEY", "value": ANTHROPIC_API_KEY},
        {"key": "MCP_TOKEN", "value": MCP_TOKEN},
    ],
)

policy = client.network_policies.create(
    name=f"claude-mcp-{uuid.uuid4().hex[:8]}",
    egress_mode="allowlist",
    rules=[
        {"destination": "api.anthropic.com", "port": 443, "protocol": "tcp"},
        {"destination": MCP_HOST, "port": 443, "protocol": "tcp"},
    ],
)

runtime = client.runtime.create(
    template=TEMPLATE,
    providers=[provider.id],
    network_policy_ids=[policy.id],
    env_vars={
        "ANTHROPIC_API_KEY": ANTHROPIC_API_KEY,
        "MCP_TOKEN": MCP_TOKEN,
    },
    timeout=1800,
)

add = runtime.run_cmd(
    command="claude",
    args=[
        "mcp",
        "add",
        "--transport",
        "http",
        "team-mcp",
        MCP_URL,
        "--header",
        f"Authorization: Bearer {MCP_TOKEN}",
    ],
    timeout=120,
)
print(add.stdout, add.stderr)
assert add.success, add.stderr

result = runtime.run_cmd(
    command="claude",
    args=[
        "--dangerously-skip-permissions",
        "-p",
        "Use the team-mcp tools to list available tools and summarize them in three bullets.",
        "--output-format",
        "text",
    ],
    timeout=900,
)
print(result.stdout)
print("exit_code:", result.exit_code)

runtime.kill()
client.identity.providers.delete(provider.id)
client.network_policies.delete(policy.id)

# Same flow via CLI (PROVIDER_ID / POLICY_ID already include MCP host egress)
RT=$(gravixlayer runtime create \
  --template "$GRAVIXLAYER_CLAUDE_TEMPLATE" \
  --provider "$ANTHROPIC_PROVIDER_ID" \
  --network-policy "$CLAUDE_NETWORK_POLICY_ID" \
  --env "ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY" \
  --env "MCP_TOKEN=$MCP_TOKEN" \
  --timeout 1800 \
  --wait \
  --output json | jq -r '.runtime_id')

gravixlayer runtime exec "$RT" --timeout 120 -- \
  claude mcp add --transport http team-mcp "$MCP_URL" \
  --header "Authorization: Bearer $MCP_TOKEN"

gravixlayer runtime exec "$RT" --timeout 900 -- \
  claude --dangerously-skip-permissions \
  -p 'Use the team-mcp tools to list available tools.' \
  --output-format text

gravixlayer runtime kill "$RT" -y