Skip to content

AGENTS.md

Latest commit

 

History

History
97 lines (86 loc) · 24.8 KB

AGENTS.md

File metadata and controls

97 lines (86 loc) · 24.8 KB
Raw

Colab CLI: Agent Guidelines

Architecture Overview

  • CLI: Modular Typer based entry point in cli.py with subcommands in commands/.
  • Common: common.py centralizes shared State (lazy-loading) and session resolution.
  • Client: ColabClient handles API interactions (assignment, unassignment).
  • Auth: auth.py exposes a single get_credentials(config_path, provider) facade that dispatches on the AuthProvider enum. Two providers are supported, selected via the global --auth=oauth2|adc flag (default oauth2):
    • oauth2: public google-auth-oauthlib InstalledAppFlow, token cached at ~/.config/colab-cli/token.json. Reads the client OAuth config from -c/--client-oauth-config (default ~/.colab-cli-oauth-config.json), falling back to the bundled src/colab_cli/oauth_config.json resource (re-added in PR #41 / 9f44fe2, 2026-05-29 — the earlier "removed in 20eb88e" note was stale/incorrect; the file exists and auth.py:_get_google_auth_credentials loads it via importlib.resources). As of 2026-06-11 the flow is a remote copy-paste flow, not a localhost server: _run_remote_flow sets redirect_uri=https://sdk.cloud.google.com/applicationdefaultauthcode.html + token_usage=remote, prints the URL, and reads the pasted code via input(). NEVER revert to OOB (urn:ietf:wg:oauth:2.0:oob) — Google blocked it in 2022 ("OOB flow has been blocked"); the sdk.cloud.google.com redirect is registered only to the bundled cloud-SDK client (764086051850-...), so any other client id gets redirect_uri_mismatch. Server-side acceptance/rejection of these variants is verifiable GET-only by building the authorization URL and inspecting whether Google reaches sign-in vs. an OAuth error page (no resources allocated).
    • adc: Google Application Default Credentials via google.auth.default(). The CLI passes scopes=PUBLIC_SCOPES (which includes colaboratory) and re-applies via creds.with_scopes() for credential types that support it. User credentials minted by gcloud auth application-default login ignore the scopes= kwarg AND raise NotImplementedError on with_scopes: ADC users must explicitly re-authenticate with gcloud auth application-default login --scopes=openid,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/userinfo.email,https://www.googleapis.com/auth/colaboratory. userinfo.email is required by the session backend at colab.research.google.com (assign/unassign/sessions/keep-alive return 401 without it); colaboratory is retained for forward compatibility and other Colab features (keep-alive no longer uses colab.pa.googleapis.com — see the Keep-alive note below); openid and cloud-platform are mandated by gcloud itself, which rejects scope lists that omit cloud-platform with Invalid value for [--scopes]. Service-account / GCE / GKE / impersonated creds get the right scopes transparently via with_scopes.
  • Backend Hosts: Two distinct backends with different requirements:
    • colab.research.google.com (session backend / tun/m/...): accepts the userinfo.email scope. Handles assign, unassign, the contents API, and keep-alive (see below).
    • Keep-alive (2026-06-15, issue #14): keep-alive is a Tunnel Frontend (TFE) HTTP ping — GET https://colab.research.google.com/tun/m/<endpoint>/keep-alive/ with header X-Colab-Tunnel: Google, authenticated by the user's own Gaia bearer token (same host/credential as assign). TFE records LastActiveTime before forwarding, refreshing the idle timer. The VM usually doesn't answer on this path, so the request commonly read-times-out even on success — client.keep_alive_assignment therefore catches requests.exceptions.ReadTimeout and treats it as success, while genuine HTTP errors (e.g. 404 for a deleted assignment) propagate. This mirrors the official colab-vscode extension's sendKeepAlive (src/colab/client.ts). DO NOT revert to the colab.pa.googleapis.com RuntimeService/KeepAliveAssignment RPC: that RPC requires the caller to be a serviceusage consumer of Colab's internal project 1014160490159, which no ordinary user account is, so it returned HTTP 403 USER_PROJECT_DENIED for every external user (issue #14) and silently idle-pruned their sessions within minutes. The browser only made that RPC work by riding the user's google.com cookie through an internal cookie-proxy (colab.clients6.google.com), which a headless bearer-token CLI cannot use. Verified live 2026-06-15 with a third-party account (the RPC 403'd; the tunnel ping succeeded and kept the VM alive).
  • Runtime: ColabRuntime wraps jupyter-kernel-client for execution.
  • State:
    • StateStore persists session metadata in ~/.config/colab-cli/sessions.json.
    • Persistent settings are in ~/.config/colab-cli/settings.json.
  • History: HistoryLogger records structured events in ~/.config/colab-cli/history/*.jsonl.

Core Mandates

  • Minimalism: Favor standard library where possible (e.g., urllib) while utilizing Typer for CLI ergonomics.

  • Piping: Always consider piped input (stdin) vs. interactive TTY.

  • Trace Alignment: When implementing new endpoints, validate against captured browser traces (HAR files).

  • TDD (Test-Driven Development): Always implement tests first. Verify they fail before implementing the solution to make them pass. Every design must include a testing strategy and specific test cases.

  • Jupyter Protocol Deviations: Google Colab uses custom extensions to the Jupyter protocol. Examples include colab_request messages over the iopub channel and input_reply wrapping colab_reply payloads on the stdin channel. These require monkey-patching or specialized handlers within jupyter-kernel-client (e.g., wsclient.kernel_socket.on_message interceptors).

  • Integration Testing: Unit tests and mocks are not enough. Before declaring any feature complete, you MUST perform a real-world, end-to-end integration test against a live Colab environment using the CLI. Never rely solely on mocked unit tests to verify a feature's correctness.

    • Integration tests are located in integration/ (e.g., integration/repro_plot_redirection/test.sh).
    • To run an integration test, use: uv run bash integration/repro_<name>/test.sh.
    • uv run ensures the colab command (entry point) is available in the shell environment.

  • Continuous Improvement: Whenever the user provides feedback, workflow advice, or corrections, immediately encode that advice into this AGENTS.md file. The goal is to learn from review and never repeat the same errors.

Tools & Workflow

  • Workflow:
    1. Draft: Plan and start the task. Create a new git branch before working on new features or changes.
    2. Refine: Implement changes and verify with tests and linting. Run tests using uv run pytest tests/ and resolve any lint errors using uv run ruff check . --fix.
    3. Finalize: Ensure everything is complete and correct. Whenever features are added or behaviors change, you MUST re-review the corresponding design document in docs/ and update it to reflect the new state. You should also add a brief log entry to the frontmatter of the updated design document with the current date summarizing the change. Finally, commit the finished changes to the git branch for review.

Subcommand Workflows

  • Session Management: new, sessions, status, stop.
  • Execution: repl, exec, console.
  • Files: ls, rm, upload, download, edit.
  • Automation: auth, drivemount, install, log, pay, version, update.

Release Tagging Workflow

The package version is derived from the git tag via hatch-vcs (see pyproject.toml), so a release is just (1) a CHANGELOG.md entry and (2) a vX.Y.Z tag on the merged commit. Follow this workflow ONLY when the user explicitly requests a release (e.g. "cut a release", "tag v0.7.0", "release the keep-alive fix"). NEVER propose or perform a release proactively.

  1. Verify a clean, current main: git fetch origin && git checkout main && git pull --ff-only origin main. Confirm git status is clean. If main has diverged or there are unmerged feature branches the user expects in the release, stop and ask.
  2. Identify unreleased commits: Find the last release tag with git describe --tags --abbrev=0, then list commits since then with git log <last-tag>..HEAD --oneline. These are what the new changelog section must cover. Cross-reference each commit's PR number (the (#NN) suffix in the merge commit subject) — every entry in the changelog should cite at least one PR.
  3. Infer the next version (semver): Categorize each commit using the Keep-a-Changelog buckets already present in CHANGELOG.md (Added, Changed, Fixed, Removed, Deprecated, Security). Pick the bump:
    • Major (X.0.0): any breaking change to the CLI surface, flag semantics, on-disk state schema (sessions.json / settings.json), or persisted token format.
    • Minor (0.X.0): at least one Added entry (new subcommand, new flag, new capability) with no breaking changes.
    • Patch (0.0.X): only Fixed, Changed, Removed (non-breaking cleanup), Deprecated, or Security entries. Confirm the inferred version with the user before proceeding — never tag without that confirmation, even when the bump seems obvious.
  4. Draft and commit the CHANGELOG entry on a branch: Create a release branch (e.g. release-v0.7.0), then prepend a new ## [X.Y.Z] - YYYY-MM-DD section above the previous release. Group entries by category in the order already established in CHANGELOG.md (Changed, Added, Fixed, Removed). Each bullet should be one short paragraph naming the user-visible behavior change (not the implementation detail) and ending with the PR reference (#NN). Append a new link reference at the bottom: [X.Y.Z]: https://github.com/googlecolab/google-colab-cli/compare/v<PREV>...vX.Y.Z. Commit with message docs: add CHANGELOG.md for vX.Y.Z release, push the branch, and open a PR with gh pr create.
  5. Wait for the changelog PR to merge: Do NOT tag the pre-merge commit on your branch. The tag must land on the squash-merge commit that appears on main (this is what hatch-vcs will see and what users will git checkout). After the user reports the PR is merged (or you observe it via gh pr view <num> --json state,mergeCommit), git fetch origin && git checkout main && git pull --ff-only origin main.
  6. Tag the merged changelog commit and push the tag: Verify git log -1 --oneline is the merged changelog commit (subject matches docs: add CHANGELOG.md for vX.Y.Z release (#NN)). Create an annotated tag whose message is the changelog section body: git tag -a vX.Y.Z -m "vX.Y.Z" -m "<changelog section body>". Push with git push origin vX.Y.Z. Do NOT create a GitHub Release entry via gh release create — the tag alone is sufficient because hatch-vcs reads it directly, and the canonical release notes already live in CHANGELOG.md.
  7. Verify: Run git describe --tags --exact-match HEAD (should print vX.Y.Z) and uv run colab version (the version string should match). Confirm with gh api repos/googlecolab/google-colab-cli/tags --jq '.[0].name' that the tag is visible on the remote.

Implementation Principles

  1. Direct Execution: Code for auth, drivemount, etc., should be injected and executed on the VM kernel.
  2. Contents API: Use the Jupyter Contents API for file management as seen in the browser traces.
  3. Transparent Storage: Local state must be overridable via flags.
  4. No netrc: Avoid netrc for token persistence in this project.
  5. Mocking Interactivity: When testing commands that branch on stdin.isatty(), use the is_stdin_tty helper in execution.py and mock it via mocker.patch("colab_cli.commands.execution.is_stdin_tty", return_value=...). This ensures tests don't hang in CI/agent environments.
  6. State Isolation: Always patch the colab_cli.common.state singleton in tests to control session persistence and client behavior. Refer to tests/conftest.py for the standard global fixture.
  7. Fire-and-Forget Architecture: The Colab CLI is a "fire-and-forget" tool. Avoid using background threads for long-running tasks within the main command flows. For persistent needs such as keep-alive, utilize detached background daemon processes (with PID tracking in the session state).
  8. Verify the Local Install: A globally-installed colab may exist on PATH (e.g. at ~/.local/bin/colab) and can shadow the project's editable install when uv run is invoked from outside the repo. ALWAYS run shell commands with the repo as the working directory (e.g. via the workdir parameter, never cd && cmd) so uv run colab ... resolves to .venv/bin/colab. Confirm with which colab and uv run which colab if a CLI test produces unexpected results (e.g. flag-not-recognized errors for flags you just added). Shebang invocations always resolve via $PATH, so a script like #!/usr/bin/env -S colab run ... will pick up the stale global tool even when the editable install is current — when testing shebang-based behavior after a code change, always run uv tool install --reinstall --force --from . colab first, then verify with colab version (the version string includes the git short SHA). Encoded 2026-05-12 after the SystemExit-suppression fix appeared not to work in examples/hello_colab.py because the shebang resolved to a uv-tool install pinned to the prior commit.
  9. Isolate the Regression First: When a user reports an error in code you just touched, do NOT assume your change caused it. First, reproduce the failure on main (or the branch point) to determine whether the bug is pre-existing. Only after confirming the regression is yours should you start debugging the new code. Encoded after spending a turn debugging "ADC broke colab new" only to discover colab new --gpu A100 was already failing on main due to an A100-quota-vs-default issue unrelated to ADC.
  10. Live Probes Allocate Real Resources: Probing the Colab API to debug an issue creates real, billable assignments — every successful POST /tun/m/assign reserves a VM. Prefer GET-only (read) probes whenever possible. For any state-mutating call, (a) record every endpoint you create as you go, and (b) clean up via client.unassign(endpoint) (or colab stop) before declaring the investigation done. Then verify with colab sessions that nothing was orphaned.
  11. Push Freshness: The remote may have advanced during a session (other contributors land commits while you work). ALWAYS git fetch <remote> immediately before pushing or merging. If git log main..<remote>/main is non-empty, reset local main to the remote, rebase feature branches onto it, retest, then push. NEVER force-push main to recover from divergence.
  12. Amend Safety: Before git commit --amend, explicitly verify all three preconditions: (a) the user requested amend OR a pre-commit hook auto-modified files for an otherwise-successful commit, (b) HEAD was created by you in this conversation (git log -1 --format='%an %ae'), (c) the commit has not been pushed to a remote. If any precondition fails, create a new commit instead. NEVER amend a failed/rejected commit — fix the issue and create a new one.
  13. Run Integration Tests Yourself: Re-read AGENTS.md "Agent Execution Limitations" before claiming you can't run a test. The CANNOT-run list is exclusively interactive commands (colab auth, colab drivemount, colab repl, colab console). Tests built on colab new / colab stop / colab log are non-interactive and you MUST run them yourself before declaring a fix complete. Encoded after running through a full implement-and-document cycle for a fix that the integration test would have falsified in 30 seconds — the cure was uv run bash integration/repro_keep_alive/test.sh.
  14. Heed Research Caveats: When a research subagent surfaces a caveat ("the proto allows it but the policy may reject"), treat it as a TODO to verify, not as a footnote. The validation pattern: shell out to the actual service with the proposed inputs and confirm the response matches expectations BEFORE writing the code that depends on it. For policy-gated paths, run a one-shot probe before committing to a design.
  15. Two distinct auths — never confuse them: This codebase has two unrelated authentication concerns: (a) CLI-to-Colab-control-plane: how Client authenticates HTTP requests to colab.research.google.com and colab.pa.googleapis.com. Driven by the global --auth={oauth2,adc} flag and auth.py:get_credentials. The OAuth2 path is bootstrapped automatically by google_auth_oauthlib's InstalledAppFlow on first invocation when ~/.config/colab-cli/token.json doesn't exist; no separate command is needed — any --auth=oauth2 invocation (e.g. colab --auth=oauth2 sessions) triggers the browser consent flow. (b) VM-side credentials: how the colab auth subcommand injects user GCP credentials into the running Colab kernel so notebook code (e.g. gcloud, BigQuery client) can make authenticated calls from inside the VM. This uses the USE_AUTH_EPHEM='0' gcloud-fallback path executed on the kernel via ColabRuntime. NEVER tell a user to "run colab auth" as a prerequisite for fixing CLI-side auth issues — they're orthogonal layers and the suggestion misleads.
  16. Detached daemons inherit nothing useful: When spawning a detached child process via subprocess.Popen (e.g. spawn_keep_alive), the child does NOT inherit the parent's parsed Typer flags. It re-parses argv from scratch, so any global flag the parent saw via --auth=adc or --config /tmp/foo.json MUST be re-emitted as part of the child's command line. Forgetting this causes silent fallback to Typer defaults: in 2026-04-30 this manifested as the keep-alive daemon (a) using OAUTH2 instead of the parent's ADC, and (b) reading from ~/.config/colab-cli/sessions.json instead of the parent's --config path. Always propagate every relevant global flag in cmd = [sys.executable, "-m", ...] BEFORE the subcommand name (Typer requires global flags before the subcommand).
  17. Persist-before-spawn for daemons that read shared state: When a detached child reads from a state file the parent owns (e.g. state.store.get(session_name)), the parent MUST persist BEFORE spawning. Otherwise the child can race ahead of the parent's add() call and observe an empty store. Symptom in 2026-04-30: keep-alive daemon exits immediately with keep_alive_stopped reason=session_not_found iters=1 duration=0.0s. Fix: call state.store.add(s) once before spawn_keep_alive, and again after (to capture the PID).
  18. [SUPERSEDED 2026-06-15 — keep-alive no longer calls colab.pa.googleapis.com; see issue #14 / the Keep-alive note above] colab.pa.googleapis.com rejects ADC user creds without X-Goog-User-Project: When calling colab.pa.googleapis.com from ADC user credentials minted by gcloud auth application-default login, the bearer token carries the user's gcloud quota project. The colab API key sent alongside it belongs to a different project (Colab, 1014160490159). The backend enforces project-match between the two and returns HTTP 400 CONSUMER_INVALID ("The API Key and the authentication credential are from different projects."). The old fix was to send X-Goog-User-Project: 1014160490159 to pin the consumer project — but that in turn required serviceusage.serviceUsageConsumer on project 1014160490159, which ordinary users lack, yielding HTTP 403 USER_PROJECT_DENIED (the issue #14 root cause). Both failure modes are now moot because keep-alive uses the TFE tunnel ping on colab.research.google.com instead. Retained here as institutional knowledge: if any future code path must call colab.pa.googleapis.com with a bearer token, it will face this same project-entitlement wall for non-internal accounts.
  19. Pydantic validation requires a schema: _issue_request accepts an optional schema= and historically called TypeAdapter(schema).validate_python(...) unconditionally. When schema=None (a caller that doesn't care about the response body — e.g. fire-and-forget RPCs like KeepAliveAssignment that return []), this raises pydantic.ValidationError: Input should be None. Always guard with if schema is None: return after the empty-body short-circuit.
  20. Suggest the branch-diff review command after committing: The user reviews changes with git diff main..<branch-name> (full cumulative diff against main, not just the latest commit). After landing one or more commits on a feature branch, ALWAYS suggest the exact command — e.g. "Review with git diff main..sort-help-commands" — instead of git show <sha> (which only shows a single commit and misses context when a branch has multiple commits). Encoded 2026-05-05 after suggesting git show 9d9c7da for a branch the user wanted to review holistically.
  21. Verify research-tool claims with primary sources: Research tools and AI assistants can be confidently wrong, especially about edge cases or features outside their training corpus. When such a tool says "X is not used / not parsed / doesn't exist", treat it as a hypothesis to verify, not a fact. Always cross-check against the primary source (the actual code or config) — and when a tool names files, check whether the indirection chain it describes actually exists. The cost of believing the tool when it's wrong is shipping a non-functional feature; the cost of double-checking is small. Encoded 2026-05-05 after a colab url first-cut shipped the wrong URL format because of unverified output.
  22. Clean up orphaned assignments before finishing live tests: After running a live integration test, colab sessions may show server-side assignments that the local colab stop couldn't see (e.g. assignments leaked from earlier in the session, or from crashed prior runs). Always run colab sessions as the final cleanup step, and for any [?]-marked orphan, run python -c "from colab_cli.common import state; state.client.unassign('<endpoint>')" (the CLI doesn't expose a direct unassign-by-endpoint command). Re-verify with colab sessions returning "No active sessions". Encoded 2026-05-05 after the first colab url live test left an orphan from a prior conversation turn that would have idled-out and billed compute units.
  23. Forward unknown args through Typer with context_settings: Typer/Click consumes any token starting with -/-- as a flag of the parent command unless told otherwise. For a subcommand like colab run script.py --some-script-flag (where --some-script-flag belongs to the user's script, not to colab), declare it with app.command(name="run", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) and accept the positional with Annotated[Optional[List[str]], typer.Argument(...)] = None. Also use repr() (not f-string interpolation) when embedding those forwarded strings into kernel-side Python source — repr() produces a safe round-trippable literal regardless of the user's shell-passed quotes, backslashes, or non-ASCII bytes. Encoded 2026-05-12 while implementing colab run.

Agent Execution Limitations (What I Can vs Cannot Run)

As an AI agent operating via non-interactive shell tools (run_shell_command), there are strict limits on what I can test autonomously without human intervention:

  • I CAN Run:
    • Automated tests (pytest), linting (ruff), and headless execution scripts.
    • Subcommands that don't pause for user input (e.g., colab new, colab status, colab stop, colab ls, colab install, colab exec <file.py>).
    • Piped colab repl and colab console: as of 2026-05-07 both commands support piped stdin and exit cleanly on EOF (echo 'cmd' | colab console -s s returns in ~1.2s). The unpiped, interactive variants still cannot be run autonomously.
    • Specially crafted mock scripts that simulate timeouts or API calls.
  • I CANNOT Run (Requires User Assistance):
    • colab auth: This command relies on the traditional Gcloud fallback input_request (via USE_AUTH_EPHEM='0'), which prompts the user via Python's input() to click a URL, sign in, and paste back an authorization code. My shell tool will hang indefinitely on this input().
    • colab drivemount: This command prompts the user via sys.stdin.readline() (specifically querying /dev/tty to ensure input is captured) to press Enter after granting OAuth consent in the browser. My shell tool will timeout/hang waiting for Enter.
    • Interactive (TTY) colab repl / colab console: When stdin is a real terminal these commands drop into interactive raw-TTY modes that require real-time keystroke streaming. My shell tools cannot support this. (Piped stdin is fine — see above.)

Whenever working on interactive commands, I must build the core logic, write mock tests, and explicitly ask the user to run the live test in their terminal to verify success.