feat(cli): Agent Mode bootstrap + claim flow (Python)
New behavior on `mem0 init`:
- With no `--email`/`--api-key` AND a positive agent signal (`--agent`,
global `--json`/`--agent`, or one of the recognized agent env vars
CLAUDECODE / CURSOR_AGENT / CODEX_CLI / CLINE / CONTINUE / AIDER /
GOOSE / WINDSURF), bootstrap an unattended Agent Mode account via
POST /api/v1/auth/agent_mode/. No email, no OTP, no dashboard.
- With `--email <addr>` AND an existing config that has agent_mode=true,
run the claim device-flow against the existing key instead of minting
a fresh one. The raw API key never leaves the device; backend confirms
claim via the existing CLILoginRequest poll path. Config flips
agent_mode=false and stamps claimed_at on success.
- Bare `mem0 init` with no signal + no TTY still errors out — auto-bootstrap
requires a positive agent signal to avoid surprising pipe-using humans.
New flags:
--agent Force unattended Agent Mode bootstrap.
--source Channel attribution string for signup_source PostHog property.
Config schema extensions on PlatformConfig:
agent_mode, created_via, claimed_at, default_user_id.
Telemetry (M1-M6 from the growth doc):
- cli.init: mode (agent|email|api_key|existing_key), agent_caller,
signup_source, claimed_agent_mode (bool when --email claims an
existing agent-mode config).
- All cli.* events: agent_mode reflects config.platform.agent_mode (the
bootstrap flag), not the output-format flag — per the growth-doc spec.
This commit is contained in:
Mgeeeek
@@ -0,0 +1,177 @@
|
||||
"""Agent Mode commands — bootstrap (unattended signup) and claim (human upgrade)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import secrets
|
||||
import time
|
||||
import webbrowser
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import typer
|
||||
from rich.console import Console
|
||||
|
||||
from mem0_cli.branding import (
|
||||
DIM_COLOR,
|
||||
print_error,
|
||||
print_info,
|
||||
print_success,
|
||||
)
|
||||
from mem0_cli.config import Mem0Config, save_config
|
||||
|
||||
console = Console()
|
||||
err_console = Console(stderr=True)
|
||||
|
||||
# Claim polling: 2-second interval, 10-minute timeout matches the backend's
|
||||
# CLILoginRequest expires_at (15 minutes — we give up before the token does).
|
||||
_POLL_INTERVAL_SECS = 2
|
||||
_POLL_TIMEOUT_SECS = 600
|
||||
|
||||
_SOURCE_HEADERS = {
|
||||
"X-Mem0-Source": "cli",
|
||||
"X-Mem0-Client-Language": "python",
|
||||
}
|
||||
|
||||
|
||||
def bootstrap_via_backend(
|
||||
config: Mem0Config,
|
||||
*,
|
||||
source: str | None = None,
|
||||
) -> None:
|
||||
"""POST /api/v1/auth/agent_mode/ and mutate config in place.
|
||||
|
||||
Returns nothing — the caller saves the config and prints follow-up messages.
|
||||
Raises typer.Exit(1) on failure.
|
||||
"""
|
||||
base_url = (config.platform.base_url or "https://api.mem0.ai").rstrip("/")
|
||||
body: dict[str, Any] = {}
|
||||
if source:
|
||||
body["source"] = source
|
||||
|
||||
try:
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
resp = client.post(
|
||||
f"{base_url}/api/v1/auth/agent_mode/",
|
||||
headers={**_SOURCE_HEADERS, "Content-Type": "application/json"},
|
||||
json=body,
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
print_error(err_console, f"Network error contacting Mem0: {exc}")
|
||||
raise typer.Exit(1) from exc
|
||||
|
||||
if resp.status_code == 429:
|
||||
print_error(err_console, "Rate-limited. Try again in a few minutes.")
|
||||
raise typer.Exit(1)
|
||||
if resp.status_code == 503:
|
||||
print_error(err_console, "Agent Mode is temporarily disabled. Try again later.")
|
||||
raise typer.Exit(1)
|
||||
if resp.status_code != 200:
|
||||
try:
|
||||
detail = resp.json().get("error", resp.text)
|
||||
except Exception:
|
||||
detail = resp.text
|
||||
print_error(err_console, f"Bootstrap failed: {detail}")
|
||||
raise typer.Exit(1)
|
||||
|
||||
envelope = resp.json()
|
||||
config.platform.api_key = envelope["api_key"]
|
||||
config.platform.base_url = base_url
|
||||
config.platform.agent_mode = True
|
||||
config.platform.created_via = "agent_mode"
|
||||
config.platform.claimed_at = ""
|
||||
config.platform.default_user_id = envelope["default_user_id"]
|
||||
# Adopt the slug-derived user_id as the default scope for memory ops.
|
||||
config.defaults.user_id = envelope["default_user_id"]
|
||||
save_config(config)
|
||||
|
||||
print_success(console, f"Agent Mode active. Default user_id: {envelope['default_user_id']}")
|
||||
console.print(f" [{DIM_COLOR}]To claim this account later: {envelope.get('claim_command', 'mem0 init --email <your-email>')}[/]")
|
||||
|
||||
|
||||
def claim_via_device_flow(config: Mem0Config, *, email: str) -> None:
|
||||
"""Run the claim flow against an existing agent-mode config.
|
||||
|
||||
Reuses the existing CLI device flow (initiate_cli_login → frontend OTP →
|
||||
associate_cli_token → get_api_key_from_cli_token poll). The raw API key
|
||||
never leaves the device — backend confirms claim, CLI updates only
|
||||
`platform.agent_mode` and `platform.claimed_at`.
|
||||
"""
|
||||
base_url = (config.platform.base_url or "https://api.mem0.ai").rstrip("/")
|
||||
if not config.platform.api_key or not config.platform.agent_mode:
|
||||
print_error(
|
||||
err_console,
|
||||
"This command requires an active Agent Mode config. Run `mem0 init` first.",
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
cli_token = secrets.token_urlsafe(32)
|
||||
raw_key = config.platform.api_key
|
||||
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
try:
|
||||
init_resp = client.post(
|
||||
f"{base_url}/api/v1/accounts/cli_login/",
|
||||
json={"token": cli_token, "claim_for_apikey": raw_key},
|
||||
headers=_SOURCE_HEADERS,
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
print_error(err_console, f"Could not initiate claim: {exc}")
|
||||
raise typer.Exit(1) from exc
|
||||
|
||||
if init_resp.status_code != 200:
|
||||
try:
|
||||
detail = init_resp.json().get("error", init_resp.text)
|
||||
except Exception:
|
||||
detail = init_resp.text
|
||||
print_error(err_console, f"Could not initiate claim: {detail}")
|
||||
raise typer.Exit(1)
|
||||
|
||||
login_url = init_resp.json().get("login_url", "")
|
||||
print_info(console, "Open in your browser to claim:")
|
||||
console.print(f" [{DIM_COLOR}]{login_url}[/]")
|
||||
try:
|
||||
webbrowser.open(login_url)
|
||||
except Exception:
|
||||
pass # Printing the URL is sufficient
|
||||
|
||||
# Poll for completion
|
||||
deadline = time.monotonic() + _POLL_TIMEOUT_SECS
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(_POLL_INTERVAL_SECS)
|
||||
try:
|
||||
poll = client.post(
|
||||
f"{base_url}/api/v1/accounts/get_api_key_from_cli_token/",
|
||||
json={"token": cli_token},
|
||||
headers=_SOURCE_HEADERS,
|
||||
)
|
||||
except httpx.HTTPError:
|
||||
continue # transient — keep polling
|
||||
|
||||
if poll.status_code != 200:
|
||||
# 400 "Token expired" / "Invalid token" → bail
|
||||
try:
|
||||
err = poll.json().get("error", "")
|
||||
except Exception:
|
||||
err = poll.text
|
||||
if "expired" in err.lower():
|
||||
print_error(err_console, "Claim link expired. Run `mem0 init --email <addr>` again.")
|
||||
raise typer.Exit(1)
|
||||
continue
|
||||
|
||||
body = poll.json()
|
||||
if body.get("claimed"):
|
||||
config.platform.agent_mode = False
|
||||
config.platform.claimed_at = body.get("claimed_at") or _utcnow_iso()
|
||||
config.platform.user_email = email
|
||||
config.platform.created_via = "email"
|
||||
save_config(config)
|
||||
print_success(console, f"Agent claimed to {email}. Your API key is unchanged.")
|
||||
return
|
||||
|
||||
print_error(err_console, "Claim timed out. Run `mem0 init --email <addr>` again.")
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
def _utcnow_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
@@ -182,21 +182,63 @@ def run_init(
|
||||
email: str | None = None,
|
||||
code: str | None = None,
|
||||
force: bool = False,
|
||||
source: str | None = None,
|
||||
agent: bool = False,
|
||||
) -> None:
|
||||
"""Interactive setup wizard for mem0 CLI.
|
||||
|
||||
When both *api_key* and *user_id* are supplied, all prompts are skipped
|
||||
(non-interactive mode). When running in a non-TTY without the required
|
||||
flags, an error message is printed.
|
||||
|
||||
Agent Mode dispatch (no email/api-key flags):
|
||||
- If existing config has an active API key → reuse (existing_key path).
|
||||
- Else if any positive agent signal (--agent, --json global, agent env
|
||||
var, or `agent` flag) → POST /api/v1/auth/agent_mode/ and write config.
|
||||
- Else fall through to the interactive wizard.
|
||||
|
||||
Claim dispatch:
|
||||
- If `--email` is set AND existing config has `agent_mode=true`, run the
|
||||
claim device-flow against the existing key instead of minting a new
|
||||
email-based key.
|
||||
"""
|
||||
from mem0_cli.agent_detect import detect_agent_caller
|
||||
from mem0_cli.commands.agent_mode_cmd import bootstrap_via_backend, claim_via_device_flow
|
||||
from mem0_cli.state import is_agent_mode as _global_agent_mode
|
||||
from mem0_cli.telemetry import capture_event
|
||||
|
||||
def _fire_init(mode: str, *, claimed: bool = False) -> None:
|
||||
"""Fire cli.init telemetry with M1-M6 properties."""
|
||||
props: dict = {"command": "init", "mode": mode}
|
||||
agent_caller = detect_agent_caller()
|
||||
if agent_caller:
|
||||
props["agent_caller"] = agent_caller
|
||||
if source:
|
||||
props["signup_source"] = source
|
||||
if claimed:
|
||||
props["claimed_agent_mode"] = True
|
||||
capture_event("cli.init", props)
|
||||
|
||||
config = Mem0Config()
|
||||
|
||||
base_url = os.environ.get("MEM0_BASE_URL", config.platform.base_url or DEFAULT_BASE_URL)
|
||||
config.platform.base_url = base_url
|
||||
|
||||
if code and not email:
|
||||
print_error(err_console, "--code requires --email.")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# ── Email + existing agent-mode config → claim flow ─────────────────
|
||||
if email and CONFIG_FILE.exists():
|
||||
existing = load_config()
|
||||
if existing.platform.agent_mode and existing.platform.api_key:
|
||||
email = email.strip().lower()
|
||||
_validate_email(email)
|
||||
print_info(console, f"Claiming Agent Mode account to {email}...")
|
||||
claim_via_device_flow(existing, email=email)
|
||||
_fire_init("email", claimed=True)
|
||||
return
|
||||
|
||||
# Warn if an existing config with an API key would be overwritten
|
||||
if not force and CONFIG_FILE.exists():
|
||||
existing = load_config()
|
||||
@@ -257,6 +299,16 @@ def run_init(
|
||||
console.print()
|
||||
return
|
||||
|
||||
# ── Agent Mode auto-bootstrap (no email, no api_key flag) ─────────
|
||||
# Positive agent signal required: explicit --agent flag (local or global)
|
||||
# OR a recognized agent env var. Pure "no TTY" alone is NOT enough — pipe
|
||||
# users would get surprised by a silent shadow signup.
|
||||
agent_ctx = agent or _global_agent_mode() or (detect_agent_caller() is not None)
|
||||
if not api_key and not email and agent_ctx:
|
||||
bootstrap_via_backend(config, source=source)
|
||||
_fire_init("agent")
|
||||
return
|
||||
|
||||
# ── API key flow (existing) ───────────────────────────────────────
|
||||
|
||||
# Non-TTY: resolve defaults so partial flags work in pipelines / CI
|
||||
@@ -265,7 +317,7 @@ def run_init(
|
||||
print_error(
|
||||
err_console,
|
||||
"Non-interactive terminal detected and --api-key is required.",
|
||||
hint="Run: mem0 init --api-key <key> [--user-id <id>]",
|
||||
hint="Run: mem0 init --api-key <key>, --email <addr>, or --agent for unattended Agent Mode bootstrap.",
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
user_id = user_id or os.environ.get("USER") or os.environ.get("USERNAME") or "mem0-cli"
|
||||
|
||||
Reference in New Issue
Block a user