Add SnapshotableContainer protocol and Docker snapshot implementation#94
Add SnapshotableContainer protocol and Docker snapshot implementation#94joshgreaves wants to merge 1 commit intomainfrom
Conversation
|
Warning Rate limit exceeded
⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (5)
📝 WalkthroughWalkthroughAdds a SnapshotableContainer protocol and exposes it; implements snapshot/from_snapshot/delete_snapshot/delete_snapshot_sync in DockerContainer; adds unit tests for snapshot behavior; and deduplicates preset registration to avoid duplicate presets. Changes
Sequence Diagram(s)sequenceDiagram
participant T as Test / Client
participant C as DockerContainer
participant D as DockerClient
participant I as ImageStore
T->>C: snapshot()
C->>D: docker commit/create image (tag)
D-->>C: image_id (snapshot_id)
C-->>T: return snapshot_id
Note over T,C: restore from snapshot
T->>C: from_snapshot(snapshot_id)
C->>D: reference image(snapshot_id)
D-->>C: image metadata
C-->>T: new DockerContainer instance
T->>C: delete_snapshot(snapshot_id)
C->>D: remove image(snapshot_id)
D-->>I: delete image (may be not-found)
D-->>C: success / ImageNotFound handled
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
Comment |
![baz-reviewer[bot]](https://avatars.githubusercontent.com/in/933528?s=60&v=4){kind=small}
| async def snapshot(self) -> str: | ||
| """Commit current container state as a Docker image.""" | ||
| if self._container is None: | ||
| raise RuntimeError("Container not started, snapshot is not possible.") |
There was a problem hiding this comment.
snapshot() raises raw RuntimeError when _container is None, violating the repo's error-handling guideline (CLAUDE.md) and preventing callers from distinguishing container lifecycle failures. Can we introduce or reuse a container-specific exception (e.g. ContainerNotStartedError) that derives from the repo exception hierarchy?
Finding type: AI Coding Guidelines | Severity: 🟢 Low
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents:
In src/ares/containers/docker.py around lines 191-194, the snapshot() method currently
raises a raw RuntimeError when self._container is None. Replace this with a
repo-specific container lifecycle exception (for example ContainerNotStartedError) that
derives from the project's error hierarchy described in CLAUDE.md. Add an import for
that exception at the top of the file (or create it in the appropriate exceptions module
if it doesn't exist), change the raise site to raise ContainerNotStartedError with a
clear message, and update any other lifecycle methods in this file that currently raise
RuntimeError to raise the same container-specific exception for consistency.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/ares/containers/docker_snapshot_test.py (1)
49-60: Exercise the restore path, not just field assignment.This only checks constructor state. It still passes if
from_snapshot()stops using the snapshot ID onstart(), which is the contract consumers care about. Make this async and assert the mocked Docker client is started with the snapshot image.🧪 Tighten the restore assertion
-def test_from_snapshot_creates_container(): +@pytest.mark.asyncio +async def test_from_snapshot_creates_container(mock_docker_client): """Test that from_snapshot() creates a DockerContainer with the snapshot as image.""" container = docker.DockerContainer.from_snapshot( "sha256:abc123", name="restored", default_workdir="/workspace", @@ assert isinstance(container, docker.DockerContainer) assert container.image == "sha256:abc123" assert container.name == "restored" assert container.default_workdir == "/workspace" + + await container.start() + call_kwargs = mock_docker_client.containers.run.call_args + assert call_kwargs[1]["image"] == "sha256:abc123"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/ares/containers/docker_snapshot_test.py` around lines 49 - 60, The test only verifies field assignment; update test_from_snapshot_creates_container to exercise the restore path by making it async, creating/mocking the Docker client used by docker.DockerContainer (or the factory it uses), calling await container.start() and asserting the mocked client's start/run method was invoked with image "sha256:abc123" (or that the container was created/started using that image); specifically target DockerContainer.from_snapshot and DockerContainer.start to ensure the snapshot ID is actually used at start time and assert the mock received the expected image argument.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/ares/containers/docker.py`:
- Around line 196-217: delete_snapshot and delete_snapshot_sync remove any image
ID given, so validate the ARES marker label ("ares-go-explore" == "true") before
calling images.remove to avoid deleting unrelated images; in delete_snapshot use
asyncio.to_thread to first fetch the image metadata (e.g.,
self._client.images.get or inspect) and check image.attrs['Config']['Labels']
(or image.labels) contains "ares-go-explore":"true" and only then call
images.remove (preserve the existing ImageNotFound handling and
_LOGGER.info/_LOGGER.debug behavior), and in delete_snapshot_sync do the same
synchronous inspect/get+label check before self._client.images.remove, using
contextlib.suppress(docker.errors.ImageNotFound) to keep idempotence; if the
label is absent or not matching, log a warning/debug and skip removal.
---
Nitpick comments:
In `@src/ares/containers/docker_snapshot_test.py`:
- Around line 49-60: The test only verifies field assignment; update
test_from_snapshot_creates_container to exercise the restore path by making it
async, creating/mocking the Docker client used by docker.DockerContainer (or the
factory it uses), calling await container.start() and asserting the mocked
client's start/run method was invoked with image "sha256:abc123" (or that the
container was created/started using that image); specifically target
DockerContainer.from_snapshot and DockerContainer.start to ensure the snapshot
ID is actually used at start time and assert the mock received the expected
image argument.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: d3645b6d-39f5-4aca-a10b-1e8858cd0f97
📒 Files selected for processing (5)
src/ares/containers/__init__.pysrc/ares/containers/containers.pysrc/ares/containers/docker.pysrc/ares/containers/docker_snapshot_test.pysrc/ares/presets.py
| tag = f"ares-snapshot-{uuid.uuid4().hex[:12]}" | ||
| image = await asyncio.to_thread( | ||
| self._container.commit, | ||
| repository="ares-go-explore", | ||
| tag=tag, | ||
| conf={"Labels": {"ares-go-explore": "true"}}, | ||
| ) | ||
| _LOGGER.info("Snapshot created: %s (tag: %s)", image.id, tag) | ||
| return image.id | ||
|
|
||
| async def delete_snapshot(self, snapshot_id: str) -> None: | ||
| """Delete a Docker image created by snapshot().""" | ||
| try: | ||
| await asyncio.to_thread(self._client.images.remove, snapshot_id) | ||
| _LOGGER.info("Snapshot deleted: %s", snapshot_id) | ||
| except docker.errors.ImageNotFound: | ||
| _LOGGER.debug("Snapshot %s already deleted.", snapshot_id) | ||
|
|
||
| def delete_snapshot_sync(self, snapshot_id: str) -> None: | ||
| """Synchronous version for atexit cleanup.""" | ||
| with contextlib.suppress(docker.errors.ImageNotFound): | ||
| self._client.images.remove(snapshot_id) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n src/ares/containers/docker.py | sed -n '190,225p'Repository: withmartian/ares
Length of output: 1691
🏁 Script executed:
# Get more context around the snapshot-related methods
cat -n src/ares/containers/docker.py | sed -n '170,230p'Repository: withmartian/ares
Length of output: 2987
🏁 Script executed:
# Search for the snapshot method implementation
rg -n "def snapshot" src/ares/containers/docker.py -A 15Repository: withmartian/ares
Length of output: 765
🏁 Script executed:
# Check if label validation exists anywhere
rg -n "ares-go-explore" src/ares/containers/docker.pyRepository: withmartian/ares
Length of output: 167
🏁 Script executed:
# Check the full class structure for context
rg -n "class DockerContainer" src/ares/containers/docker.py -A 50Repository: withmartian/ares
Length of output: 2559
Use the snapshot label as a delete guard.
snapshot() marks ARES-created images with ares-go-explore=true, but both delete paths ignore that marker and will remove any image ID they are given. That makes a bad or stale snapshot_id capable of deleting an unrelated local image. Validate the label before images.remove() and keep the current not-found idempotence on top.
🛡️ Suggested guard
+ def _assert_snapshot_image(self, snapshot_id: str) -> None:
+ image = self._client.images.get(snapshot_id)
+ labels = image.attrs.get("Config", {}).get("Labels") or {}
+ if labels.get("ares-go-explore") != "true":
+ raise ValueError(f"{snapshot_id} is not an ARES snapshot")
+
async def delete_snapshot(self, snapshot_id: str) -> None:
"""Delete a Docker image created by snapshot()."""
try:
+ await asyncio.to_thread(self._assert_snapshot_image, snapshot_id)
await asyncio.to_thread(self._client.images.remove, snapshot_id)
_LOGGER.info("Snapshot deleted: %s", snapshot_id)
except docker.errors.ImageNotFound:
_LOGGER.debug("Snapshot %s already deleted.", snapshot_id)
def delete_snapshot_sync(self, snapshot_id: str) -> None:
"""Synchronous version for atexit cleanup."""
with contextlib.suppress(docker.errors.ImageNotFound):
+ self._assert_snapshot_image(snapshot_id)
self._client.images.remove(snapshot_id)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| tag = f"ares-snapshot-{uuid.uuid4().hex[:12]}" | |
| image = await asyncio.to_thread( | |
| self._container.commit, | |
| repository="ares-go-explore", | |
| tag=tag, | |
| conf={"Labels": {"ares-go-explore": "true"}}, | |
| ) | |
| _LOGGER.info("Snapshot created: %s (tag: %s)", image.id, tag) | |
| return image.id | |
| async def delete_snapshot(self, snapshot_id: str) -> None: | |
| """Delete a Docker image created by snapshot().""" | |
| try: | |
| await asyncio.to_thread(self._client.images.remove, snapshot_id) | |
| _LOGGER.info("Snapshot deleted: %s", snapshot_id) | |
| except docker.errors.ImageNotFound: | |
| _LOGGER.debug("Snapshot %s already deleted.", snapshot_id) | |
| def delete_snapshot_sync(self, snapshot_id: str) -> None: | |
| """Synchronous version for atexit cleanup.""" | |
| with contextlib.suppress(docker.errors.ImageNotFound): | |
| self._client.images.remove(snapshot_id) | |
| tag = f"ares-snapshot-{uuid.uuid4().hex[:12]}" | |
| image = await asyncio.to_thread( | |
| self._container.commit, | |
| repository="ares-go-explore", | |
| tag=tag, | |
| conf={"Labels": {"ares-go-explore": "true"}}, | |
| ) | |
| _LOGGER.info("Snapshot created: %s (tag: %s)", image.id, tag) | |
| return image.id | |
| def _assert_snapshot_image(self, snapshot_id: str) -> None: | |
| image = self._client.images.get(snapshot_id) | |
| labels = image.attrs.get("Config", {}).get("Labels") or {} | |
| if labels.get("ares-go-explore") != "true": | |
| raise ValueError(f"{snapshot_id} is not an ARES snapshot") | |
| async def delete_snapshot(self, snapshot_id: str) -> None: | |
| """Delete a Docker image created by snapshot().""" | |
| try: | |
| await asyncio.to_thread(self._assert_snapshot_image, snapshot_id) | |
| await asyncio.to_thread(self._client.images.remove, snapshot_id) | |
| _LOGGER.info("Snapshot deleted: %s", snapshot_id) | |
| except docker.errors.ImageNotFound: | |
| _LOGGER.debug("Snapshot %s already deleted.", snapshot_id) | |
| def delete_snapshot_sync(self, snapshot_id: str) -> None: | |
| """Synchronous version for atexit cleanup.""" | |
| with contextlib.suppress(docker.errors.ImageNotFound): | |
| self._assert_snapshot_image(snapshot_id) | |
| self._client.images.remove(snapshot_id) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ares/containers/docker.py` around lines 196 - 217, delete_snapshot and
delete_snapshot_sync remove any image ID given, so validate the ARES marker
label ("ares-go-explore" == "true") before calling images.remove to avoid
deleting unrelated images; in delete_snapshot use asyncio.to_thread to first
fetch the image metadata (e.g., self._client.images.get or inspect) and check
image.attrs['Config']['Labels'] (or image.labels) contains
"ares-go-explore":"true" and only then call images.remove (preserve the existing
ImageNotFound handling and _LOGGER.info/_LOGGER.debug behavior), and in
delete_snapshot_sync do the same synchronous inspect/get+label check before
self._client.images.remove, using
contextlib.suppress(docker.errors.ImageNotFound) to keep idempotence; if the
label is absent or not matching, log a warning/debug and skip removal.
|
@coderabbitai full review |
✅ Actions performedFull review triggered. |
|
@coderabbitaidev full review |
|
✅ Actions performedFull review triggered. |
|
@coderabbitaidev full review |
|
✅ Actions performedFull review triggered. |
There was a problem hiding this comment.
♻️ Duplicate comments (1)
src/ares/containers/docker.py (1)
206-217:⚠️ Potential issue | 🟠 MajorStill missing an ownership guard before image deletion.
snapshot()labels ARES-created images withares-go-explore=true, but both delete paths still remove any image ID they're given. A stale or wrongsnapshot_idcan still delete an unrelated local image unless you inspect the label first and skip or reject non-ARES images.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/ares/containers/docker.py` around lines 206 - 217, Both delete_snapshot and delete_snapshot_sync must verify the image is owned by ARES before removing it: call the Docker API to inspect/get the image (e.g., via self._client.images.get(snapshot_id) or .get and check image.attrs.get("Config", {}).get("Labels") or image.labels) and only proceed with self._client.images.remove if the label ares-go-explore == "true"; if the image is missing, preserve the existing ImageNotFound handling, and if the label is absent or false, log a warning/debug and skip deletion (do this check in both delete_snapshot and delete_snapshot_sync).
🧹 Nitpick comments (2)
src/ares/containers/docker.py (1)
219-229: Instantiate viaclsin this classmethod.Hardcoding
DockerContainer(...)meansSomeDockerSubclass.from_snapshot(...)would return the base type. Usingcls(...)preserves subclass constructors and overrides.♻️ Minimal change
- return DockerContainer(image=snapshot_id, name=name, resources=resources, default_workdir=default_workdir) + return cls(image=snapshot_id, name=name, resources=resources, default_workdir=default_workdir)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/ares/containers/docker.py` around lines 219 - 229, The classmethod from_snapshot currently instantiates the concrete DockerContainer class directly, which breaks subclassing; change the instantiation to use cls(...) instead of DockerContainer(...) so that SomeDockerSubclass.from_snapshot(...) returns an instance of the actual subclass; update the return expression in from_snapshot to return cls(image=snapshot_id, name=name, resources=resources, default_workdir=default_workdir).src/ares/presets.py (1)
128-138: Surface preset-name collisions instead of silently dropping them.Because
preset_nameis derived from_make_harbor_dataset_id(), two different Harbor specs can still collapse onto the same public name. Right now the later one disappears with no trace, which makes a missing preset hard to diagnose. Consider tracking the first(name, version)and logging when another spec hits the samepreset_name.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/ares/presets.py` around lines 128 - 138, The loop over code_env.list_harbor_datasets() silently drops duplicates when preset_name (built via _make_harbor_dataset_id and code_agent_id) collides because of the seen set; change this to map preset_name -> first ds_spec identifier (e.g., (ds_spec.name, ds_spec.version)) instead of a plain seen set, and when encountering a preset_name already present, emit a clear log or warning (including the existing mapped (name,version) and the new ds_spec) rather than silently continue; update references in this block (seen -> a dict, ds_spec, _make_harbor_dataset_id, preset_name) and keep the original dedup behavior but surface collisions for debugging.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@src/ares/containers/docker.py`:
- Around line 206-217: Both delete_snapshot and delete_snapshot_sync must verify
the image is owned by ARES before removing it: call the Docker API to
inspect/get the image (e.g., via self._client.images.get(snapshot_id) or .get
and check image.attrs.get("Config", {}).get("Labels") or image.labels) and only
proceed with self._client.images.remove if the label ares-go-explore == "true";
if the image is missing, preserve the existing ImageNotFound handling, and if
the label is absent or false, log a warning/debug and skip deletion (do this
check in both delete_snapshot and delete_snapshot_sync).
---
Nitpick comments:
In `@src/ares/containers/docker.py`:
- Around line 219-229: The classmethod from_snapshot currently instantiates the
concrete DockerContainer class directly, which breaks subclassing; change the
instantiation to use cls(...) instead of DockerContainer(...) so that
SomeDockerSubclass.from_snapshot(...) returns an instance of the actual
subclass; update the return expression in from_snapshot to return
cls(image=snapshot_id, name=name, resources=resources,
default_workdir=default_workdir).
In `@src/ares/presets.py`:
- Around line 128-138: The loop over code_env.list_harbor_datasets() silently
drops duplicates when preset_name (built via _make_harbor_dataset_id and
code_agent_id) collides because of the seen set; change this to map preset_name
-> first ds_spec identifier (e.g., (ds_spec.name, ds_spec.version)) instead of a
plain seen set, and when encountering a preset_name already present, emit a
clear log or warning (including the existing mapped (name,version) and the new
ds_spec) rather than silently continue; update references in this block (seen ->
a dict, ds_spec, _make_harbor_dataset_id, preset_name) and keep the original
dedup behavior but surface collisions for debugging.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 52c9edd3-5569-4e90-b8bd-2c25edfcf20e
📒 Files selected for processing (5)
src/ares/containers/__init__.pysrc/ares/containers/containers.pysrc/ares/containers/docker.pysrc/ares/containers/docker_snapshot_test.pysrc/ares/presets.py
2 participants
User description
Summary
SnapshotableContainerprotocol extendingContainerwithsnapshot(),from_snapshot(),delete_snapshot(), anddelete_snapshot_sync()methodsDockerContainerusingdocker committo capture filesystem state as tagged imagesSnapshotableContainerfromares.containerspresets.pyContext
Part of Go-Explore support (Wave 1). This PR is independent of the other two Wave 1 PRs and can be reviewed/merged in parallel:
go-explore/agent-state-serialization)go-explore/checkpoint-protocol)Daytona support is not included -- their runtime sandbox snapshotting API is not yet available (daytonaio/daytona#2519). The
SnapshotableContainerprotocol is designed so Daytona support slots in without changes to consumers.Tests
6 new tests covering snapshot creation, error handling,
from_snapshot, deletion, and idempotent deletion.Generated description
Below is a concise technical summary of the changes proposed in this PR:
Introduce the
SnapshotableContainerprotocol inares.containers.containersand updateDockerContainerto capture filesystem snapshots via Docker images so Go-Explore workflows can restore prior environment states. Fix duplicate Harbor preset registration inpresets.pyto ensure each built-in preset is registered only once.SnapshotableContainerand DockerContainer snapshot lifecycle so Go-Explore flows can capture and restore filesystem state.Modified files (4)
Latest Contributors(2)
Modified files (1)
Latest Contributors(2)
Summary by CodeRabbit
New Features
Bug Fixes
Tests