From 9e7459a0e255bd5e7e5a2a84161d43bddeda2f59 Mon Sep 17 00:00:00 2001 From: Drake Aiman Date: Wed, 4 Nov 2026 23:06:05 +0000 Subject: [PATCH 1/3] feat(state): keep attached sessions in their own store file A session attached to a runtime the CLI did not create must never be unassigned by an older CLI that cannot tell it apart. StateStore now keeps attached records in a second owner-only locked file derived from the sessions path (sessions.attached.json by default, .attached.json for --config), reads and lists both files and removes from both, so existing callers work unchanged. The file a record is read from decides its attached flag, so refreshing a record writes it back to the file it came from and never moves it to the other one. Records in sessions.json gain an "attached": false key, which older versions ignore because they do not forbid extra fields. --- AGENTS.md | 2 +- CHANGELOG.md | 4 + README.md | 4 +- docs/01_session_management.md | 2 + src/colab_cli/cli.py | 6 +- src/colab_cli/state.py | 56 +++++++++----- tests/test_state.py | 142 ++++++++++++++++++++++++++++++++++ 7 files changed, 194 insertions(+), 22 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index e4fbc70a..dce4c035 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -11,7 +11,7 @@ - `colab.research.google.com` (session backend / `tun/m/...`): accepts the `userinfo.email` scope. Handles assign, unassign, CCU info, and the contents API. - **Runtime**: `ColabRuntime` wraps `jupyter-kernel-client` for execution. - **State**: - - `StateStore` persists session metadata in `~/.config/colab-cli/sessions.json`. + - `StateStore` persists session metadata in `~/.config/colab-cli/sessions.json`; attached sessions (a runtime the CLI did not create) and their tokens go to the derived `.attached.json` beside it, and the file a record is read from decides its `attached` flag. - Persistent settings are in `~/.config/colab-cli/settings.json`. - **History**: `HistoryLogger` records structured events in `~/.config/colab-cli/history/*.jsonl`. diff --git a/CHANGELOG.md b/CHANGELOG.md index 9031725d..511311e8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,10 @@ below corresponds to a tag of the same name. ### Added - **ssh:** `colab ssh [-s NAME] -- CMD` runs one command on an existing runtime and exits with its status. Output is byte-exact when piped (a PTY only when stdin and stdout are terminals), and it never creates a runtime. Its handshake is tagged `ssh-mode-cmd`. +- **sessions:** Sessions attached to a runtime the CLI did not create are kept + in their own owner-only file beside the sessions file + (`sessions.attached.json`, or `.attached.json` with + `--config .json`), so older versions never see them. ### Fixed diff --git a/README.md b/README.md index d0324153..234a285e 100644 --- a/README.md +++ b/README.md @@ -108,7 +108,7 @@ Run `colab --help` to view specific options, defaults, and detailed he ### Global Options * `--auth {oauth2,adc}` — Authentication strategy for the Colab API (default: `adc`). * `-c, --client-oauth-config PATH` — Path to public OAuth client credentials configuration (default: `~/.colab-cli-oauth-config.json`). -* `--config PATH` — Path to local session metadata storage (default: `~/.config/colab-cli/sessions.json`). +* `--config PATH` — Path to local session metadata storage (default: `~/.config/colab-cli/sessions.json`). Sessions attached to a runtime the CLI did not create, and their tokens, are kept beside it in `.attached.json`. * `--logtostderr` — Direct debug logging output to stderr. --- @@ -146,7 +146,7 @@ colab stop -s analysis * **Machine shape:** Use `--high-mem` with `colab new`, `colab run`, or `colab ssh` (when auto-creating a runtime) to request a high-RAM machine shape. Requires Colab Pro or Pro+ entitlement for supported accelerators (CPU, T4, A100, etc.). L4 and TPU runtimes ignore this flag because they only offer one shape. Machine shape is shown in `colab sessions` and `colab status`. * **TTY Requirements:** The interactive commands `repl` and `console` require a local TTY. When running inside automated scripts or pipelines, make sure to pipe stdin (e.g., `echo "print(1)" | colab repl`) to trigger non-interactive execution modes. * **Transparent Code Execution:** When calling `colab exec -f file.py`, the CLI reads the file locally and transmits its content to the remote kernel. You do not need to manually upload files before execution. -* **Storage & State Paths:** Session tokens and metadata are stored at `~/.config/colab-cli/sessions.json`. Global CLI settings are located at `~/.config/colab-cli/settings.json`. These can be customized or isolated via the global `--config` flag. +* **Storage & State Paths:** Session tokens and metadata are stored at `~/.config/colab-cli/sessions.json`; sessions attached to a runtime the CLI did not create, and their tokens, are kept beside it in `sessions.attached.json` (`.attached.json` for `--config .json`). Global CLI settings are located at `~/.config/colab-cli/settings.json`. These can be customized or isolated via the global `--config` flag. ### Ephemeral Accelerator Jobs diff --git a/docs/01_session_management.md b/docs/01_session_management.md index b916a9d3..b8387fdd 100644 --- a/docs/01_session_management.md +++ b/docs/01_session_management.md @@ -1,5 +1,6 @@ --- log: +2026-11-04: `StateStore` keeps attached sessions in their own owner-only file beside the sessions file. 2026-11-04: `--gpu` and `--tpu` now reject unknown values instead of silently selecting `A100` (GPU) or `v6e1` (TPU). `resolve_runtime_options()` uses the new `GPU_ACCELERATORS` / `TPU_ACCELERATORS` maps and raises a `typer.BadParameter` that names the flag and the supported list; `--help` is generated from the same maps so it cannot drift. Validation runs before any assignment request, in both `colab new` and `colab run`. 2026-11-02: Added `_check_linux_keyring_locked` in `auth.py` to check if the active Linux `SecretService` keyring collection is locked via `secretstorage` before reading or writing OAuth tokens, logging and printing an immediate warning with the collection label, timeout, and `--auth=adc` hint while still attempting the keyring operation. 2026-10-29: The debug log (`colab.log` and `--logtostderr`) masks access tokens, runtime proxy tokens, OAuth codes, client secrets, auth headers and cookies, including in percent-encoded URLs; `jupyter_kernel_client` logs at INFO so kernel stdin replies are not recorded. @@ -140,6 +141,7 @@ The runtime proxy token (`runtimeProxyInfo.token`) authenticates every request t ### State Persistence & File Locking - **Stores**: `StateStore` (`sessions.json`) and `SettingsStore` (`settings.json`) both derive from `_LockedFileStore`, which guards concurrent access across independent `colab` invocations. +- **Attached sessions**: `StateStore` keeps records with `attached=True` (a runtime the CLI did not create) in a second locked, owner-only file derived from the sessions path (`sessions.attached.json`, or `.attached.json` for `--config .json`), so versions that read only `sessions.json` never see them. `add` routes on `attached`, `get` and `list` read both files, and `remove` removes from both. The file a record is read from decides its `attached` flag, so a refresh writes each record back to its own file. - **Cross-platform locking**: Locking uses the [`filelock`](https://pypi.org/project/filelock/) library rather than `fcntl.flock`. `fcntl` is POSIX-only and is unavailable on Windows, so the original implementation crashed on import there. `filelock` provides the same advisory cross-process locking on Linux, macOS, and Windows. - **Shared vs. exclusive**: Each store owns a `filelock.ReadWriteLock` bound to a sidecar file (`.lock`). Reads acquire `read_lock()` (shared — multiple concurrent readers allowed) and writes acquire `write_lock()` (exclusive). This preserves the `LOCK_SH`/`LOCK_EX` distinction of the previous `fcntl` implementation. - **`is_singleton=False`**: The `ReadWriteLock` is created with `is_singleton=False`. With `filelock`'s default (`True`), two `ReadWriteLock` objects for the same path *within a single process* are deduplicated into one reentrant lock; its reentrancy guard then raises `RuntimeError` when two threads each construct their own `StateStore` and contend for the write lock. Disabling the singleton registry makes each store's lock independent so they serialize via the underlying file lock instead. diff --git a/src/colab_cli/cli.py b/src/colab_cli/cli.py index fa56f58f..d272d175 100644 --- a/src/colab_cli/cli.py +++ b/src/colab_cli/cli.py @@ -59,7 +59,11 @@ def callback( Optional[str], typer.Option( "--config", - help="Path to session state file (~/.config/colab-cli/sessions.json)", + help=( + "Path to session state file" + " (~/.config/colab-cli/sessions.json); attached sessions and" + " their tokens are kept beside it in .attached.json" + ), ), ] = None, logtostderr: Annotated[ diff --git a/src/colab_cli/state.py b/src/colab_cli/state.py index 2538577e..4f9f0caa 100644 --- a/src/colab_cli/state.py +++ b/src/colab_cli/state.py @@ -35,6 +35,7 @@ class SessionState(BaseModel): session_id: Optional[str] = None last_execution: Optional[Tuple[str, Optional[str], str]] = None running: Optional[str] = None + attached: bool = False class Settings(BaseModel): @@ -128,19 +129,30 @@ def save(self, settings: Settings): class StateStore(_LockedFileStore): + """Session records: CLI-created ones in `path`, attached ones beside it. + + Attached records live in their own file so that older CLI versions, which + read only `path`, never see them and so never unassign their runtimes. + """ + def __init__(self, path: Optional[str] = None): if not path: path = os.path.expanduser("~/.config/colab-cli/sessions.json") super().__init__(path) + self._attached = _LockedFileStore( + os.path.splitext(path)[0] + ".attached.json", + ) + make_private(self._attached.lock_path) - def _load_raw(self, f) -> Dict[str, SessionState]: + def _load_raw(self, f, attached: bool) -> Dict[str, SessionState]: try: f.seek(0) content = f.read() if not content or content.isspace(): return {} data = json.loads(content) - return {k: SessionState(**v) for k, v in data.items()} + origin = {"attached": attached} + return {k: SessionState(**{**v, **origin}) for k, v in data.items()} except Exception: return {} @@ -150,28 +162,36 @@ def _save_raw(self, f, sessions: Dict[str, SessionState]): ) self._write_data(f, content) + def _read(self, store: _LockedFileStore) -> Dict[str, SessionState]: + with store._lock_shared() as f: + if f is None: + return {} + return self._load_raw(f, store is self._attached) + + def _remove_from(self, store: _LockedFileStore, name: str) -> None: + if not os.path.exists(store.path): + return + with store._lock_exclusive() as f: + sessions = self._load_raw(f, store is self._attached) + if name in sessions: + del sessions[name] + self._save_raw(f, sessions) + def add(self, state: SessionState): - with self._lock_exclusive() as f: - sessions = self._load_raw(f) + store = self._attached if state.attached else self + with store._lock_exclusive() as f: + sessions = self._load_raw(f, state.attached) sessions[state.name] = state self._save_raw(f, sessions) def get(self, name: str) -> Optional[SessionState]: - with self._lock_shared() as f: - if f is None: - return None - sessions = self._load_raw(f) - return sessions.get(name) + own, attached_records = self._read(self), self._read(self._attached) + return own.get(name) or attached_records.get(name) def remove(self, name: str): - with self._lock_exclusive() as f: - sessions = self._load_raw(f) - if name in sessions: - del sessions[name] - self._save_raw(f, sessions) + self._remove_from(self, name) + self._remove_from(self._attached, name) def list(self) -> Dict[str, SessionState]: - with self._lock_shared() as f: - if f is None: - return {} - return self._load_raw(f) + own, attached_records = self._read(self), self._read(self._attached) + return {**attached_records, **own} diff --git a/tests/test_state.py b/tests/test_state.py index 6408426a..8692df0c 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import multiprocessing import os import pytest +import stat import tempfile import threading from datetime import datetime, timezone @@ -344,3 +346,143 @@ def test_state_store_multiprocess_concurrency(temp_config): assert p1.exitcode == 0 assert p2.exitcode == 0 assert len(StateStore(temp_config).list()) == 80 + + +OWNER_ONLY = 0o600 + + +def _session_state(name, attached=False): + return SessionState( + name=name, token="t", url="u", endpoint=f"{name}-ep", attached=attached + ) + + +def _record(name, **fields): + return dict(name=name, token="t", url="u", endpoint=f"{name}-ep", **fields) + + +@pytest.mark.parametrize( + "config, sessions_file, attached_file", + [ + ( + None, + ".config/colab-cli/sessions.json", + ".config/colab-cli/sessions.attached.json", + ), + ("custom.json", "custom.json", "custom.attached.json"), + ], + ids=["default", "config-file"], +) +def test_attached_record_goes_to_its_own_file( + tmp_path, monkeypatch, config, sessions_file, attached_file +): + monkeypatch.setenv("HOME", str(tmp_path)) + (tmp_path / ".config/colab-cli").mkdir(parents=True) + store = StateStore(str(tmp_path / config) if config else None) + + store.add(_session_state("att", attached=True)) + + assert not (tmp_path / sessions_file).exists() + assert (tmp_path / attached_file).exists() + assert store.get("att").attached is True + + +def test_store_reads_and_lists_both_files(tmp_path): + path = tmp_path / "sessions.json" + store = StateStore(str(path)) + store.add(_session_state("cli")) + store.add(_session_state("att", attached=True)) + + assert '"cli"' in path.read_text() + assert '"att"' not in path.read_text() + attached_text = (tmp_path / "sessions.attached.json").read_text() + assert '"att"' in attached_text + assert '"cli"' not in attached_text + reloaded = StateStore(str(path)) + assert reloaded.get("cli").attached is False + assert reloaded.get("att").attached is True + assert sorted(reloaded.list()) == ["att", "cli"] + + +def test_remove_deletes_the_record_from_either_file(tmp_path): + path = tmp_path / "sessions.json" + store = StateStore(str(path)) + store.add(_session_state("cli")) + store.add(_session_state("att", attached=True)) + reloaded = StateStore(str(path)) + + reloaded.remove("att") + reloaded.remove("cli") + + assert StateStore(str(path)).list() == {} + + +def test_attached_add_leaves_sessions_file_byte_identical(tmp_path): + path = tmp_path / "sessions.json" + store = StateStore(str(path)) + store.add(_session_state("cli")) + before = path.read_bytes() + + store.add(_session_state("att", attached=True)) + + assert path.read_bytes() == before + + +def test_attached_remove_leaves_sessions_file_byte_identical(tmp_path): + path = tmp_path / "sessions.json" + store = StateStore(str(path)) + store.add(_session_state("cli")) + store.add(_session_state("att", attached=True)) + before = path.read_bytes() + + store.remove("att") + + assert path.read_bytes() == before + + +@pytest.mark.parametrize( + "stored_flag", [{}, {"attached": False}], ids=["no-flag", "flag-false"] +) +def test_refreshed_attached_record_stays_attached(tmp_path, stored_flag): + path = tmp_path / "sessions.json" + store = StateStore(str(path)) + store.add(_session_state("cli")) + before = path.read_bytes() + attached_file = tmp_path / "sessions.attached.json" + attached_file.write_text(json.dumps({"att": _record("att", **stored_flag)})) + + store.add(store.get("att")) + + assert path.read_bytes() == before + assert sorted(store.list()) == ["att", "cli"] + assert store.get("att").attached is True + + +def test_refreshed_cli_record_stays_out_of_attached_file(tmp_path): + path = tmp_path / "sessions.json" + path.write_text(json.dumps({"cli": _record("cli", attached=True)})) + store = StateStore(str(path)) + + store.add(store.get("cli")) + + assert not (tmp_path / "sessions.attached.json").exists() + assert store.get("cli").attached is False + + +@pytest.mark.parametrize("existing", [False, True], ids=["fresh", "was-0644"]) +def test_attached_file_and_lock_are_private(tmp_path, existing): + attached = tmp_path / "sessions.attached.json" + lock = tmp_path / "sessions.attached.json.lock" + old_umask = os.umask(0o022) + try: + if existing: + for p in (attached, lock): + p.touch() + os.chmod(p, 0o644) + store = StateStore(str(tmp_path / "sessions.json")) + store.add(_session_state("att", attached=True)) + finally: + os.umask(old_umask) + + assert stat.S_IMODE(attached.stat().st_mode) == OWNER_ONLY + assert stat.S_IMODE(lock.stat().st_mode) == OWNER_ONLY From 48d7bd323b0eda261c4d0bc9467b7860cefc5e17 Mon Sep 17 00:00:00 2001 From: Drake Aiman Date: Thu, 5 Nov 2026 00:25:44 +0000 Subject: [PATCH 2/3] feat(state): refuse a session name present in both store files A name can end up in both the sessions file and the attached file, for example when `colab new -s NAME` (in this or an older version) creates a session under a name that is already attached. StateStore now raises DuplicateSessionError for such a name from get and list, so commands that read the store stop before any Colab request, and the CLI prints the error with the recovery command: colab stop -s NAME forgets the attached record without unassigning and leaves the session the CLI created. --- AGENTS.md | 2 +- CHANGELOG.md | 5 +-- README.md | 2 +- docs/01_session_management.md | 8 ++-- skills/colab-operator/SKILL.md | 4 +- src/colab_cli/cli.py | 14 ++++++- src/colab_cli/commands/session.py | 12 +++++- src/colab_cli/state.py | 37 +++++++++++++++++- tests/test_cli.py | 63 +++++++++++++++++++++++++++++++ tests/test_state.py | 49 +++++++++++++++++++++++- 10 files changed, 178 insertions(+), 18 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index dce4c035..992db2e2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -47,7 +47,7 @@ The package version is derived from the git tag via `hatch-vcs` (see `pyproject. 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 ..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. + - **Major** (`X.0.0`): any breaking change to the CLI surface, flag semantics, on-disk state schema (`sessions.json` / `.attached.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. diff --git a/CHANGELOG.md b/CHANGELOG.md index 511311e8..9b54ba4e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,10 +21,7 @@ below corresponds to a tag of the same name. ### Added - **ssh:** `colab ssh [-s NAME] -- CMD` runs one command on an existing runtime and exits with its status. Output is byte-exact when piped (a PTY only when stdin and stdout are terminals), and it never creates a runtime. Its handshake is tagged `ssh-mode-cmd`. -- **sessions:** Sessions attached to a runtime the CLI did not create are kept - in their own owner-only file beside the sessions file - (`sessions.attached.json`, or `.attached.json` with - `--config .json`), so older versions never see them. +- **sessions:** Sessions attached to a runtime the CLI did not create are kept in their own owner-only file beside the sessions file (`sessions.attached.json`, or `.attached.json` with `--config .json`), so older versions never see them. If one name is in both files, commands that read the store stop with an error before any Colab request, and `colab stop -s NAME` forgets the attached record without stopping its runtime. ### Fixed diff --git a/README.md b/README.md index 234a285e..883904e9 100644 --- a/README.md +++ b/README.md @@ -108,7 +108,7 @@ Run `colab --help` to view specific options, defaults, and detailed he ### Global Options * `--auth {oauth2,adc}` — Authentication strategy for the Colab API (default: `adc`). * `-c, --client-oauth-config PATH` — Path to public OAuth client credentials configuration (default: `~/.colab-cli-oauth-config.json`). -* `--config PATH` — Path to local session metadata storage (default: `~/.config/colab-cli/sessions.json`). Sessions attached to a runtime the CLI did not create, and their tokens, are kept beside it in `.attached.json`. +* `--config PATH` — Path to local session metadata storage (default: `~/.config/colab-cli/sessions.json`). Sessions attached to a runtime the CLI did not create, and their tokens, are kept beside it in `.attached.json`, where `` is the path without its extension (`sessions.attached.json` by default). * `--logtostderr` — Direct debug logging output to stderr. --- diff --git a/docs/01_session_management.md b/docs/01_session_management.md index b8387fdd..665227df 100644 --- a/docs/01_session_management.md +++ b/docs/01_session_management.md @@ -1,6 +1,6 @@ --- log: -2026-11-04: `StateStore` keeps attached sessions in their own owner-only file beside the sessions file. +2026-11-04: `StateStore` keeps attached sessions in their own owner-only file beside the sessions file; a name in both files fails with an error that names `colab stop -s NAME`, which forgets the attached record without unassigning. 2026-11-04: `--gpu` and `--tpu` now reject unknown values instead of silently selecting `A100` (GPU) or `v6e1` (TPU). `resolve_runtime_options()` uses the new `GPU_ACCELERATORS` / `TPU_ACCELERATORS` maps and raises a `typer.BadParameter` that names the flag and the supported list; `--help` is generated from the same maps so it cannot drift. Validation runs before any assignment request, in both `colab new` and `colab run`. 2026-11-02: Added `_check_linux_keyring_locked` in `auth.py` to check if the active Linux `SecretService` keyring collection is locked via `secretstorage` before reading or writing OAuth tokens, logging and printing an immediate warning with the collection label, timeout, and `--auth=adc` hint while still attempting the keyring operation. 2026-10-29: The debug log (`colab.log` and `--logtostderr`) masks access tokens, runtime proxy tokens, OAuth codes, client secrets, auth headers and cookies, including in percent-encoded URLs; `jupyter_kernel_client` logs at INFO so kernel stdin replies are not recorded. @@ -134,14 +134,14 @@ The runtime proxy token (`runtimeProxyInfo.token`) authenticates every request t - **Authentication**: Uses `google-auth-oauthlib` to perform a local server OAuth flow. - **Global Flags**: - `-c`, `--client-oauth-config`: Path to the client secrets JSON file (default: `~/.colab-cli-oauth-config.json`). - - `--config`: Path to the session state JSON file (default: `~/.config/colab-cli/sessions.json`). + - `--config`: Path to the session state JSON file (default: `~/.config/colab-cli/sessions.json`). Attached sessions and their tokens are kept beside it in `.attached.json`, where `` is the path without its extension. - **Token Storage**: Credentials are persisted to the OS native keyring via the `keyring` library (`service_name="google-colab-cli"`, `username="oauth_token"`) after the initial flow. If the keyring is inaccessible or locked, `_run_remote_flow` is invoked for the current run and a warning is printed to stderr. A legacy `~/.config/colab-cli/token.json` is migrated into the keyring on first use and deleted after a successful keyring write, or immediately if the keyring already holds a token. - Use `requests` for robust HTTP interactions and `pydantic` for schema validation. - Handle authentication headers (likely `Authorization: Bearer ` or cookies). ### State Persistence & File Locking - **Stores**: `StateStore` (`sessions.json`) and `SettingsStore` (`settings.json`) both derive from `_LockedFileStore`, which guards concurrent access across independent `colab` invocations. -- **Attached sessions**: `StateStore` keeps records with `attached=True` (a runtime the CLI did not create) in a second locked, owner-only file derived from the sessions path (`sessions.attached.json`, or `.attached.json` for `--config .json`), so versions that read only `sessions.json` never see them. `add` routes on `attached`, `get` and `list` read both files, and `remove` removes from both. The file a record is read from decides its `attached` flag, so a refresh writes each record back to its own file. +- **Attached sessions**: `StateStore` keeps records with `attached=True` (a runtime the CLI did not create) in a second locked, owner-only file derived from the sessions path (`sessions.attached.json`, or `.attached.json` for `--config .json`), so versions that read only `sessions.json` never see them. `add` routes on `attached`, `get` and `list` read both files, and `remove` removes from both. The file a record is read from decides its `attached` flag, so a refresh writes each record back to its own file. A name present in both raises `DuplicateSessionError` from `get` and `list`, so commands that read the store stop before any Colab request (commands that assign first, such as `colab new -s`, are not covered yet); `colab stop -s NAME` recovers by forgetting the attached record without unassigning. - **Cross-platform locking**: Locking uses the [`filelock`](https://pypi.org/project/filelock/) library rather than `fcntl.flock`. `fcntl` is POSIX-only and is unavailable on Windows, so the original implementation crashed on import there. `filelock` provides the same advisory cross-process locking on Linux, macOS, and Windows. - **Shared vs. exclusive**: Each store owns a `filelock.ReadWriteLock` bound to a sidecar file (`.lock`). Reads acquire `read_lock()` (shared — multiple concurrent readers allowed) and writes acquire `write_lock()` (exclusive). This preserves the `LOCK_SH`/`LOCK_EX` distinction of the previous `fcntl` implementation. - **`is_singleton=False`**: The `ReadWriteLock` is created with `is_singleton=False`. With `filelock`'s default (`True`), two `ReadWriteLock` objects for the same path *within a single process* are deduplicated into one reentrant lock; its reentrancy guard then raises `RuntimeError` when two threads each construct their own `StateStore` and contend for the write lock. Disabling the singleton registry makes each store's lock independent so they serialize via the underlying file lock instead. @@ -157,7 +157,7 @@ TDD is mandatory for all session management features. ### 2. State Store Validation - **Test Case**: Verify `StateStore` correctly handles file locking and multiple concurrent reads/writes. -- **Test Case**: Verify `--config` override correctly directs all operations to the specified file path. +- **Test Case**: Verify `--config` override correctly directs all operations to the specified file path and its `.attached.json`. - **Test Case (cross-platform locking)**: Verify the store locks via `filelock.ReadWriteLock` on the `.lock` sidecar and does not import the POSIX-only `fcntl`. - **Test Case (shared/exclusive semantics)**: Verify reads go through `read_lock()` and writes through `write_lock()`. - **Test Case (cross-process exclusion)**: Hold the write lock from a separate process and confirm the store's in-process write blocks until release. diff --git a/skills/colab-operator/SKILL.md b/skills/colab-operator/SKILL.md index c8892d70..8d6a089c 100644 --- a/skills/colab-operator/SKILL.md +++ b/skills/colab-operator/SKILL.md @@ -79,8 +79,8 @@ by running `uv tool install google-colab-cli` or `pip install google-colab-cli`. ## Safety - **Always `colab stop -s ` when done** — idle VMs burn compute units. `colab run` (without `--keep`) self-cleans even if the script errors. -- Local state lives in `~/.config/colab-cli/sessions.json` (settings in `settings.json`, history in `history/*.jsonl`). Don't edit by hand. -- **Isolate parallel/agent runs** with the global `--config ` flag to point session state at a scratch file (e.g. `colab --config /tmp/agent.json new -s job`). +- Local state lives in `~/.config/colab-cli/sessions.json`, plus `sessions.attached.json` beside it for sessions attached to a runtime the CLI did not create, with their tokens (settings in `settings.json`, history in `history/*.jsonl`). Don't edit by hand. +- **Isolate parallel/agent runs** with the global `--config ` flag to point session state at a scratch file (e.g. `colab --config /tmp/agent.json new -s job`). Attached sessions then go to `.attached.json` (here `/tmp/agent.attached.json`), which holds tokens: delete it together with the scratch file. ## Recovery - "Session not found" / 404 / 401 on exec: the backend pruned the VM. `colab exec`/`repl` detect this and clean up local state automatically — run `colab sessions` and re-create with `colab new`. diff --git a/src/colab_cli/cli.py b/src/colab_cli/cli.py index d272d175..a6a3c9de 100644 --- a/src/colab_cli/cli.py +++ b/src/colab_cli/cli.py @@ -13,7 +13,7 @@ # limitations under the License. import os -from typing import Optional +from typing import Any, Optional import click import typer @@ -24,6 +24,7 @@ from colab_cli.auth import AuthProvider from colab_cli.common import state, setup_logging from colab_cli.commands import session, execution, files, automation, run, ssh, utility, usage +from colab_cli.state import DuplicateSessionError class AlphabeticalGroup(TyperGroup): @@ -37,6 +38,14 @@ class AlphabeticalGroup(TyperGroup): def list_commands(self, ctx: click.Context) -> list[str]: return sorted(super().list_commands(ctx)) + def invoke(self, ctx: click.Context) -> Any: + """Reports a `DuplicateSessionError` as a CLI error, not a traceback.""" + try: + return super().invoke(ctx) + except DuplicateSessionError as e: + typer.echo(f"[colab] Error: {e}", err=True) + raise typer.Exit(1) + app = typer.Typer( help="Colab CLI", @@ -62,7 +71,8 @@ def callback( help=( "Path to session state file" " (~/.config/colab-cli/sessions.json); attached sessions and" - " their tokens are kept beside it in .attached.json" + " their tokens are kept beside it in .attached.json," + " where is the path without its extension" ), ), ] = None, diff --git a/src/colab_cli/commands/session.py b/src/colab_cli/commands/session.py index aaf79f56..102ca8cd 100644 --- a/src/colab_cli/commands/session.py +++ b/src/colab_cli/commands/session.py @@ -29,7 +29,7 @@ shape_display_label, ) from colab_cli.utils import get_status_code -from colab_cli.state import SessionState +from colab_cli.state import DuplicateSessionError, SessionState from colab_cli.runtime import ColabRuntime @@ -360,7 +360,15 @@ def stop( from colab_cli.common import state name = state.resolve_session(session) - s = state.get_session(name, ignore_missing_session=True) + try: + s = state.get_session(name, ignore_missing_session=True) + except DuplicateSessionError: + state.store.remove_attached(name) + typer.echo( + f"[colab] Forgot the attached record for '{name}';" + " its runtime keeps running." + ) + return if not s: typer.echo(f"[colab] Session '{name}' not found.") return diff --git a/src/colab_cli/state.py b/src/colab_cli/state.py index 4f9f0caa..d665029c 100644 --- a/src/colab_cli/state.py +++ b/src/colab_cli/state.py @@ -128,6 +128,17 @@ def save(self, settings: Settings): self._write_data(f, settings.model_dump_json(indent=2)) +class DuplicateSessionError(Exception): + """A session name is in both the sessions file and the attached file.""" + + def __init__(self, name: str): + super().__init__( + f"Session '{name}' is recorded both as a CLI session and as an " + f"attached session. Run 'colab stop -s {name}' to forget the " + "attached record." + ) + + class StateStore(_LockedFileStore): """Session records: CLI-created ones in `path`, attached ones beside it. @@ -178,6 +189,11 @@ def _remove_from(self, store: _LockedFileStore, name: str) -> None: self._save_raw(f, sessions) def add(self, state: SessionState): + """Writes `state` to the file its `attached` flag names. + + A new record defaults to a CLI session, so writers of attached sessions + set the flag. + """ store = self._attached if state.attached else self with store._lock_exclusive() as f: sessions = self._load_raw(f, state.attached) @@ -185,13 +201,32 @@ def add(self, state: SessionState): self._save_raw(f, sessions) def get(self, name: str) -> Optional[SessionState]: + """Returns the record named `name` from either file, or None. + + Raises: + DuplicateSessionError: If `name` is in both files. + """ own, attached_records = self._read(self), self._read(self._attached) + if name in own and name in attached_records: + raise DuplicateSessionError(name) return own.get(name) or attached_records.get(name) def remove(self, name: str): self._remove_from(self, name) self._remove_from(self._attached, name) + def remove_attached(self, name: str) -> None: + """Forgets only the attached record for `name`.""" + self._remove_from(self._attached, name) + def list(self) -> Dict[str, SessionState]: + """Returns the records of both files by name. + + Raises: + DuplicateSessionError: If a name is in both files. + """ own, attached_records = self._read(self), self._read(self._attached) - return {**attached_records, **own} + duplicates = own.keys() & attached_records.keys() + if duplicates: + raise DuplicateSessionError(min(duplicates)) + return {**own, **attached_records} diff --git a/tests/test_cli.py b/tests/test_cli.py index 013be89d..e7c791ef 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -22,12 +22,17 @@ from colab_cli.cli import app from colab_cli.client import ( Assignment, + Client, ColabRequestError, PostAssignmentResponse, + Prod, RuntimeProxyInfo, TooManyAssignmentsError, ) +from colab_cli.commands import session as session_cmd +from colab_cli.common import State from colab_cli.consumption import ConsumptionUserInfo +from colab_cli.state import SessionState, StateStore runner = CliRunner() @@ -343,6 +348,64 @@ def test_cli_stop(mock_client, mock_store, mock_common_state): mock_store.remove.assert_called_with("s1") +def _dup_session(attached): + return SessionState( + name="dup", token="t", url="u", endpoint="dup-ep", attached=attached + ) + + +class _RecordingSession: + """Stands in for the authorized HTTP session and records every request.""" + + def __init__(self): + self.requests = [] + + def request(self, method, url, **kwargs): + self.requests.append((method, url)) + raise RuntimeError(f"unexpected request: {method} {url}") + + +@pytest.fixture +def dup_state(tmp_path, monkeypatch, mocker): + monkeypatch.setenv("HOME", str(tmp_path)) + mocker.patch("colab_cli.cli.auto_update.run_background_check") + st = State() + st._store = StateStore(str(tmp_path / "sessions.json")) + st._client = Client(Prod(), _RecordingSession()) + st._history = MagicMock() + st.store.add(_dup_session(attached=False)) + st.store.add(_dup_session(attached=True)) + mocker.patch("colab_cli.common.state", st) + return st + + +@pytest.mark.parametrize( + "args", + [["status"], ["sessions"], ["stop"], ["restart-kernel", "-s", "dup"]], + ids=["status", "sessions", "stop-without-name", "named-command"], +) +def test_duplicate_name_fails_before_any_request(dup_state, args): + result = runner.invoke(app, args) + + assert result.exit_code == 1 + assert "colab stop -s dup" in result.output + assert dup_state.client.session.requests == [] + + +def test_stop_recovers_duplicate_without_unassigning(dup_state, tmp_path): + sessions_before = (tmp_path / "sessions.json").read_bytes() + + result = runner.invoke(app, ["stop", "-s", "dup"]) + + assert result.exit_code == 0, result.output + assert dup_state.client.session.requests == [] + session_cmd.ColabRuntime.assert_not_called() + remaining = StateStore(str(tmp_path / "sessions.json")).list() + assert list(remaining) == ["dup"] + assert remaining["dup"].attached is False + assert (tmp_path / "sessions.json").read_bytes() == sessions_before + + def test_cli_sessions_prune(mock_common_state): mock_assignment = MagicMock() mock_session_state1 = MagicMock() diff --git a/tests/test_state.py b/tests/test_state.py index 8692df0c..4db26d66 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -23,7 +23,13 @@ import filelock -from colab_cli.state import StateStore, SessionState, SettingsStore, Settings +from colab_cli.state import ( + DuplicateSessionError, + StateStore, + SessionState, + SettingsStore, + Settings, +) @pytest.fixture @@ -440,6 +446,20 @@ def test_attached_remove_leaves_sessions_file_byte_identical(tmp_path): assert path.read_bytes() == before +def test_remove_attached_keeps_the_cli_record(tmp_path): + path = tmp_path / "sessions.json" + store = StateStore(str(path)) + store.add(_session_state("dup")) + store.add(_session_state("dup", attached=True)) + before = path.read_bytes() + + store.remove_attached("dup") + + assert '"dup"' not in (tmp_path / "sessions.attached.json").read_text() + assert store.get("dup").attached is False + assert path.read_bytes() == before + + @pytest.mark.parametrize( "stored_flag", [{}, {"attached": False}], ids=["no-flag", "flag-false"] ) @@ -486,3 +506,30 @@ def test_attached_file_and_lock_are_private(tmp_path, existing): assert stat.S_IMODE(attached.stat().st_mode) == OWNER_ONLY assert stat.S_IMODE(lock.stat().st_mode) == OWNER_ONLY + + +@pytest.mark.parametrize( + "read", + [lambda store: store.get("dup"), lambda store: store.list()], + ids=["get", "list"], +) +def test_name_in_both_files_raises_with_recovery_command(tmp_path, read): + store = StateStore(str(tmp_path / "sessions.json")) + store.add(_session_state("dup")) + store.add(_session_state("dup", attached=True)) + + with pytest.raises(DuplicateSessionError) as exc: + read(store) + + assert "colab stop -s dup" in str(exc.value) + + +def test_get_of_another_name_ignores_the_duplicate(tmp_path): + store = StateStore(str(tmp_path / "sessions.json")) + store.add(_session_state("dup")) + store.add(_session_state("dup", attached=True)) + store.add(_session_state("other")) + + other = store.get("other") + + assert other.endpoint == "other-ep" From c53e7479d76109ff2fc80774f34908222f4f7e7c Mon Sep 17 00:00:00 2001 From: Drake Aiman Date: Thu, 5 Nov 2026 00:25:44 +0000 Subject: [PATCH 3/3] feat(state): refuse an attached store file as the sessions file Reading one store's attached file as another store's sessions file, for example with `--config ~/.config/colab-cli/sessions.attached.json`, would stamp its records as CLI sessions and write them back that way; the reverse, a derived attached path that links to another store's sessions file, would stamp CLI records as attached. StateStore now refuses a sessions path ending in .attached.json (in any case, or a symlink to one) and a path whose derived attached path resolves to a file not named *.attached.json, before touching any file. The CLI runs the same check on --config at startup, because the store opens lazily and `colab new` or `colab run` would otherwise assign a runtime before the refusal and record it nowhere. Hard links and renamed copies are not detected in either direction: an attached file linked or copied under a sessions name, or a derived attached path hard-linked to another store's sessions file. DuplicateSessionError and the new AttachedStorePathError share StoreError, which the CLI reports as an error message instead of a traceback. --- CHANGELOG.md | 2 +- README.md | 2 +- docs/01_session_management.md | 8 ++--- src/colab_cli/cli.py | 10 +++--- src/colab_cli/state.py | 52 ++++++++++++++++++++++++++---- tests/test_cli.py | 54 +++++++++++++++++++++++++++++--- tests/test_state.py | 59 +++++++++++++++++++++++++++++++++++ 7 files changed, 167 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9b54ba4e..8ab51506 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,7 +21,7 @@ below corresponds to a tag of the same name. ### Added - **ssh:** `colab ssh [-s NAME] -- CMD` runs one command on an existing runtime and exits with its status. Output is byte-exact when piped (a PTY only when stdin and stdout are terminals), and it never creates a runtime. Its handshake is tagged `ssh-mode-cmd`. -- **sessions:** Sessions attached to a runtime the CLI did not create are kept in their own owner-only file beside the sessions file (`sessions.attached.json`, or `.attached.json` with `--config .json`), so older versions never see them. If one name is in both files, commands that read the store stop with an error before any Colab request, and `colab stop -s NAME` forgets the attached record without stopping its runtime. +- **sessions:** Sessions attached to a runtime the CLI did not create are kept in their own owner-only file beside the sessions file (`sessions.attached.json`, or `.attached.json` with `--config .json`), so older versions never see them; `--config` refuses a path ending in `.attached.json`. If one name is in both files, commands that read the store stop with an error before any Colab request, and `colab stop -s NAME` forgets the attached record without stopping its runtime. ### Fixed diff --git a/README.md b/README.md index 883904e9..4bcf9a1b 100644 --- a/README.md +++ b/README.md @@ -108,7 +108,7 @@ Run `colab --help` to view specific options, defaults, and detailed he ### Global Options * `--auth {oauth2,adc}` — Authentication strategy for the Colab API (default: `adc`). * `-c, --client-oauth-config PATH` — Path to public OAuth client credentials configuration (default: `~/.colab-cli-oauth-config.json`). -* `--config PATH` — Path to local session metadata storage (default: `~/.config/colab-cli/sessions.json`). Sessions attached to a runtime the CLI did not create, and their tokens, are kept beside it in `.attached.json`, where `` is the path without its extension (`sessions.attached.json` by default). +* `--config PATH` — Path to local session metadata storage (default: `~/.config/colab-cli/sessions.json`). Sessions attached to a runtime the CLI did not create, and their tokens, are kept beside it in `.attached.json`, where `` is the path without its extension (`sessions.attached.json` by default); a path ending in `.attached.json` is refused. * `--logtostderr` — Direct debug logging output to stderr. --- diff --git a/docs/01_session_management.md b/docs/01_session_management.md index 665227df..b2b20b43 100644 --- a/docs/01_session_management.md +++ b/docs/01_session_management.md @@ -1,6 +1,6 @@ --- log: -2026-11-04: `StateStore` keeps attached sessions in their own owner-only file beside the sessions file; a name in both files fails with an error that names `colab stop -s NAME`, which forgets the attached record without unassigning. +2026-11-04: `StateStore` keeps attached sessions in their own owner-only file beside the sessions file; a name in both files fails with an error that names `colab stop -s NAME`, which forgets the attached record without unassigning. A sessions path that names an attached file, or whose attached path links to another kind of file, is refused (hard links and copies are not detected). 2026-11-04: `--gpu` and `--tpu` now reject unknown values instead of silently selecting `A100` (GPU) or `v6e1` (TPU). `resolve_runtime_options()` uses the new `GPU_ACCELERATORS` / `TPU_ACCELERATORS` maps and raises a `typer.BadParameter` that names the flag and the supported list; `--help` is generated from the same maps so it cannot drift. Validation runs before any assignment request, in both `colab new` and `colab run`. 2026-11-02: Added `_check_linux_keyring_locked` in `auth.py` to check if the active Linux `SecretService` keyring collection is locked via `secretstorage` before reading or writing OAuth tokens, logging and printing an immediate warning with the collection label, timeout, and `--auth=adc` hint while still attempting the keyring operation. 2026-10-29: The debug log (`colab.log` and `--logtostderr`) masks access tokens, runtime proxy tokens, OAuth codes, client secrets, auth headers and cookies, including in percent-encoded URLs; `jupyter_kernel_client` logs at INFO so kernel stdin replies are not recorded. @@ -134,14 +134,14 @@ The runtime proxy token (`runtimeProxyInfo.token`) authenticates every request t - **Authentication**: Uses `google-auth-oauthlib` to perform a local server OAuth flow. - **Global Flags**: - `-c`, `--client-oauth-config`: Path to the client secrets JSON file (default: `~/.colab-cli-oauth-config.json`). - - `--config`: Path to the session state JSON file (default: `~/.config/colab-cli/sessions.json`). Attached sessions and their tokens are kept beside it in `.attached.json`, where `` is the path without its extension. + - `--config`: Path to the session state JSON file (default: `~/.config/colab-cli/sessions.json`). Attached sessions and their tokens are kept beside it in `.attached.json`, where `` is the path without its extension; a path ending in `.attached.json` is refused. - **Token Storage**: Credentials are persisted to the OS native keyring via the `keyring` library (`service_name="google-colab-cli"`, `username="oauth_token"`) after the initial flow. If the keyring is inaccessible or locked, `_run_remote_flow` is invoked for the current run and a warning is printed to stderr. A legacy `~/.config/colab-cli/token.json` is migrated into the keyring on first use and deleted after a successful keyring write, or immediately if the keyring already holds a token. - Use `requests` for robust HTTP interactions and `pydantic` for schema validation. - Handle authentication headers (likely `Authorization: Bearer ` or cookies). ### State Persistence & File Locking - **Stores**: `StateStore` (`sessions.json`) and `SettingsStore` (`settings.json`) both derive from `_LockedFileStore`, which guards concurrent access across independent `colab` invocations. -- **Attached sessions**: `StateStore` keeps records with `attached=True` (a runtime the CLI did not create) in a second locked, owner-only file derived from the sessions path (`sessions.attached.json`, or `.attached.json` for `--config .json`), so versions that read only `sessions.json` never see them. `add` routes on `attached`, `get` and `list` read both files, and `remove` removes from both. The file a record is read from decides its `attached` flag, so a refresh writes each record back to its own file. A name present in both raises `DuplicateSessionError` from `get` and `list`, so commands that read the store stop before any Colab request (commands that assign first, such as `colab new -s`, are not covered yet); `colab stop -s NAME` recovers by forgetting the attached record without unassigning. +- **Attached sessions**: `StateStore` keeps records with `attached=True` (a runtime the CLI did not create) in a second locked, owner-only file derived from the sessions path (`sessions.attached.json`, or `.attached.json` for `--config .json`), so versions that read only `sessions.json` never see them. `add` routes on `attached`, `get` and `list` read both files, and `remove` removes from both. The file a record is read from decides its `attached` flag, so a refresh writes each record back to its own file. A name present in both raises `DuplicateSessionError` from `get` and `list`, so commands that read the store stop before any Colab request (commands that assign first, such as `colab new -s`, are not covered yet); `colab stop -s NAME` recovers by forgetting the attached record without unassigning. A sessions path ending in `.attached.json` (in any case, or a symlink to one), or one whose derived attached path is a symlink to a file not named `*.attached.json`, raises `AttachedStorePathError`. The CLI applies this check to `--config` at startup, before any command runs. Hard links and renamed copies are not detected in either direction: do not hard-link or copy an attached file or a sessions file under another store's name. - **Cross-platform locking**: Locking uses the [`filelock`](https://pypi.org/project/filelock/) library rather than `fcntl.flock`. `fcntl` is POSIX-only and is unavailable on Windows, so the original implementation crashed on import there. `filelock` provides the same advisory cross-process locking on Linux, macOS, and Windows. - **Shared vs. exclusive**: Each store owns a `filelock.ReadWriteLock` bound to a sidecar file (`.lock`). Reads acquire `read_lock()` (shared — multiple concurrent readers allowed) and writes acquire `write_lock()` (exclusive). This preserves the `LOCK_SH`/`LOCK_EX` distinction of the previous `fcntl` implementation. - **`is_singleton=False`**: The `ReadWriteLock` is created with `is_singleton=False`. With `filelock`'s default (`True`), two `ReadWriteLock` objects for the same path *within a single process* are deduplicated into one reentrant lock; its reentrancy guard then raises `RuntimeError` when two threads each construct their own `StateStore` and contend for the write lock. Disabling the singleton registry makes each store's lock independent so they serialize via the underlying file lock instead. @@ -157,7 +157,7 @@ TDD is mandatory for all session management features. ### 2. State Store Validation - **Test Case**: Verify `StateStore` correctly handles file locking and multiple concurrent reads/writes. -- **Test Case**: Verify `--config` override correctly directs all operations to the specified file path and its `.attached.json`. +- **Test Case**: Verify `--config` override correctly directs all operations to the specified file path and its `.attached.json`, and refuses a path ending in `.attached.json`. - **Test Case (cross-platform locking)**: Verify the store locks via `filelock.ReadWriteLock` on the `.lock` sidecar and does not import the POSIX-only `fcntl`. - **Test Case (shared/exclusive semantics)**: Verify reads go through `read_lock()` and writes through `write_lock()`. - **Test Case (cross-process exclusion)**: Hold the write lock from a separate process and confirm the store's in-process write blocks until release. diff --git a/src/colab_cli/cli.py b/src/colab_cli/cli.py index a6a3c9de..af40677b 100644 --- a/src/colab_cli/cli.py +++ b/src/colab_cli/cli.py @@ -24,7 +24,7 @@ from colab_cli.auth import AuthProvider from colab_cli.common import state, setup_logging from colab_cli.commands import session, execution, files, automation, run, ssh, utility, usage -from colab_cli.state import DuplicateSessionError +from colab_cli.state import StoreError, resolve_store_paths class AlphabeticalGroup(TyperGroup): @@ -39,10 +39,10 @@ def list_commands(self, ctx: click.Context) -> list[str]: return sorted(super().list_commands(ctx)) def invoke(self, ctx: click.Context) -> Any: - """Reports a `DuplicateSessionError` as a CLI error, not a traceback.""" + """Reports a `StoreError` as a CLI error, not a traceback.""" try: return super().invoke(ctx) - except DuplicateSessionError as e: + except StoreError as e: typer.echo(f"[colab] Error: {e}", err=True) raise typer.Exit(1) @@ -72,7 +72,8 @@ def callback( "Path to session state file" " (~/.config/colab-cli/sessions.json); attached sessions and" " their tokens are kept beside it in .attached.json," - " where is the path without its extension" + " where is the path without its extension;" + " *.attached.json paths are refused" ), ), ] = None, @@ -98,6 +99,7 @@ def callback( state.config_path = config state.logtostderr = logtostderr state.auth_provider = auth + resolve_store_paths(config) setup_logging(logtostderr) # Daily fetch + cached banner on every invocation. diff --git a/src/colab_cli/state.py b/src/colab_cli/state.py index d665029c..bf59e5fb 100644 --- a/src/colab_cli/state.py +++ b/src/colab_cli/state.py @@ -128,7 +128,50 @@ def save(self, settings: Settings): self._write_data(f, settings.model_dump_json(indent=2)) -class DuplicateSessionError(Exception): +_ATTACHED_SUFFIX = ".attached.json" + + +class StoreError(Exception): + """A session store problem that the CLI reports as an error message.""" + + +class AttachedStorePathError(StoreError): + """The sessions path, or its attached path, aliases the other kind.""" + + def __init__(self, path: str, reason: str = ""): + if not reason: + reason = "it is an attached sessions file; pass a sessions file" + super().__init__(f"Cannot use '{path}' as the sessions file: {reason}.") + + +def resolve_store_paths(path: Optional[str] = None) -> Tuple[str, str]: + """Returns the sessions file and its attached file for `path`. + + Hard links and copies of either file are not detected. + + Args: + path: The sessions file; None or empty means + ~/.config/colab-cli/sessions.json. + + Raises: + AttachedStorePathError: If one store's attached file would be read as + another store's sessions file, or the reverse. + """ + if not path: + path = os.path.expanduser("~/.config/colab-cli/sessions.json") + for name in (path, os.path.realpath(path)): + if name.lower().endswith(_ATTACHED_SUFFIX): + raise AttachedStorePathError(path) + attached_path = os.path.splitext(path)[0] + _ATTACHED_SUFFIX + resolved = os.path.realpath(attached_path).lower() + if not resolved.endswith(_ATTACHED_SUFFIX): + raise AttachedStorePathError( + path, f"'{attached_path}' links to a file of another kind" + ) + return path, attached_path + + +class DuplicateSessionError(StoreError): """A session name is in both the sessions file and the attached file.""" def __init__(self, name: str): @@ -147,12 +190,9 @@ class StateStore(_LockedFileStore): """ def __init__(self, path: Optional[str] = None): - if not path: - path = os.path.expanduser("~/.config/colab-cli/sessions.json") + path, attached_path = resolve_store_paths(path) super().__init__(path) - self._attached = _LockedFileStore( - os.path.splitext(path)[0] + ".attached.json", - ) + self._attached = _LockedFileStore(attached_path) make_private(self._attached.lock_path) def _load_raw(self, f, attached: bool) -> Dict[str, SessionState]: diff --git a/tests/test_cli.py b/tests/test_cli.py index e7c791ef..82dbac78 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import time from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch @@ -366,19 +367,25 @@ def request(self, method, url, **kwargs): @pytest.fixture -def dup_state(tmp_path, monkeypatch, mocker): +def real_state(tmp_path, monkeypatch, mocker): monkeypatch.setenv("HOME", str(tmp_path)) mocker.patch("colab_cli.cli.auto_update.run_background_check") st = State() - st._store = StateStore(str(tmp_path / "sessions.json")) st._client = Client(Prod(), _RecordingSession()) st._history = MagicMock() - st.store.add(_dup_session(attached=False)) - st.store.add(_dup_session(attached=True)) mocker.patch("colab_cli.common.state", st) + mocker.patch("colab_cli.cli.state", st) return st +@pytest.fixture +def dup_state(real_state, tmp_path): + real_state._store = StateStore(str(tmp_path / "sessions.json")) + real_state.store.add(_dup_session(attached=False)) + real_state.store.add(_dup_session(attached=True)) + return real_state + + @pytest.mark.parametrize( "args", [["status"], ["sessions"], ["stop"], ["restart-kernel", "-s", "dup"]], @@ -406,6 +413,45 @@ def test_stop_recovers_duplicate_without_unassigning(dup_state, tmp_path): assert (tmp_path / "sessions.json").read_bytes() == sessions_before +@pytest.mark.parametrize( + "args", + [["status"], ["new", "-s", "job"], ["run", "job.py"]], + ids=["status", "new", "run"], +) +def test_config_naming_an_attached_file_is_refused( + real_state, tmp_path, monkeypatch, args +): + monkeypatch.chdir(tmp_path) + (tmp_path / "job.py").write_text("print(1)\n") + attached = tmp_path / "sessions.attached.json" + attached.write_text(json.dumps({"dup": _dup_session(True).model_dump()})) + before = attached.read_bytes() + + result = runner.invoke(app, ["--config", str(attached), *args]) + + assert result.exit_code == 1 + assert isinstance(result.exception, SystemExit) + assert "[colab] Error: Cannot use" in result.output + assert real_state.client.session.requests == [] + assert attached.read_bytes() == before + + +def test_config_with_linked_attached_path_is_refused(real_state, tmp_path): + sessions = tmp_path / "sessions.json" + sessions.write_text(json.dumps({"dup": _dup_session(False).model_dump()})) + before = sessions.read_bytes() + (tmp_path / "agent.attached.json").symlink_to(sessions) + config = str(tmp_path / "agent.json") + + result = runner.invoke(app, ["--config", config, "sessions"]) + + assert result.exit_code == 1 + assert isinstance(result.exception, SystemExit) + assert "links to a file of another kind" in result.output + assert real_state.client.session.requests == [] + assert sessions.read_bytes() == before + + def test_cli_sessions_prune(mock_common_state): mock_assignment = MagicMock() mock_session_state1 = MagicMock() diff --git a/tests/test_state.py b/tests/test_state.py index 4db26d66..3467db4a 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -24,6 +24,7 @@ import filelock from colab_cli.state import ( + AttachedStorePathError, DuplicateSessionError, StateStore, SessionState, @@ -533,3 +534,61 @@ def test_get_of_another_name_ignores_the_duplicate(tmp_path): other = store.get("other") assert other.endpoint == "other-ep" + + +@pytest.mark.parametrize( + "name", + ["sessions.attached.json", "Sessions.ATTACHED.json", "link.json"], + ids=["attached-file", "other-case", "symlink"], +) +def test_store_refuses_an_attached_file_as_sessions_file(tmp_path, name): + attached = tmp_path / "sessions.attached.json" + attached.write_text(json.dumps({"att": _record("att", attached=True)})) + before = attached.read_bytes() + path = tmp_path / name + if name == "link.json": + path.symlink_to(attached) + + with pytest.raises(AttachedStorePathError): + StateStore(str(path)) + + assert attached.read_bytes() == before + + +def test_store_refuses_an_attached_path_linked_to_another_file(tmp_path): + sessions = tmp_path / "sessions.json" + sessions.write_text(json.dumps({"cli": _record("cli")})) + before = sessions.read_bytes() + (tmp_path / "agent.attached.json").symlink_to(sessions) + + with pytest.raises(AttachedStorePathError): + StateStore(str(tmp_path / "agent.json")) + + assert sessions.read_bytes() == before + + +@pytest.mark.parametrize( + "config, link, target", + [ + ("sessions.json", "sessions.json", "real-sessions.json"), + ("agent.json", "agent.attached.json", "other.attached.json"), + ], + ids=["sessions-file", "attached-file"], +) +def test_store_accepts_links_to_files_of_the_same_kind(tmp_path, config, link, target): + (tmp_path / target).write_text("{}") + (tmp_path / link).symlink_to(tmp_path / target) + store = StateStore(str(tmp_path / config)) + + store.add(_session_state("att", attached=True)) + + assert store.get("att").attached is True + + +def test_store_refuses_a_dangling_attached_link(tmp_path): + (tmp_path / "a.attached.json").symlink_to(tmp_path / "nowhere.json") + + with pytest.raises(AttachedStorePathError): + StateStore(str(tmp_path / "a.json")) + + assert [p.name for p in tmp_path.iterdir()] == ["a.attached.json"]