diff --git a/.github/workflows/python-test.yml b/.github/workflows/python-test.yml
index 48979cc..8925f33 100644
--- a/.github/workflows/python-test.yml
+++ b/.github/workflows/python-test.yml
@@ -32,10 +32,16 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Setup Rust
         uses: actions-rust-lang/setup-rust-toolchain@v1
+        with:
+          cache: false
       - uses: Swatinem/rust-cache@v2
         with:
           workspaces: python
           cache-targets: false
+      - name: Install dependencies
+        run: |
+          sudo apt update
+          sudo apt install -y protobuf-compiler
       - name: Install uv
         uses: astral-sh/setup-uv@v3
       - name: Create virtual environment and install dependencies
diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml
index a065dc1..42042b2 100644
--- a/.github/workflows/rust-test.yml
+++ b/.github/workflows/rust-test.yml
@@ -39,6 +39,10 @@ jobs:
         with:
           workspaces: rust/lance-context
           cache-targets: false
+      - name: Install dependencies
+        run: |
+          sudo apt update
+          sudo apt install -y protobuf-compiler
       - name: Build tests
         run: cargo test --manifest-path rust/lance-context/Cargo.toml --no-run
       - name: Run unit tests
diff --git a/python/python/lance_context/__init__.py b/python/python/lance_context/__init__.py
index 1d2eaf2..27c69c5 100644
--- a/python/python/lance_context/__init__.py
+++ b/python/python/lance_context/__init__.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from ._internal import Context, version as _version
+from ._internal import Context  # pyright: ignore[reportMissingImports]
+from ._internal import version as _version  # pyright: ignore[reportMissingImports]
 
 __all__ = ["Context", "__version__"]
 
diff --git a/python/src/lib.rs b/python/src/lib.rs
index d96da56..fb002c7 100644
--- a/python/src/lib.rs
+++ b/python/src/lib.rs
@@ -47,7 +47,7 @@ impl Context {
     }
 
     #[pyo3(signature = (label = None))]
-    fn snapshot(&self, label: Option<&str>) -> String {
+    fn snapshot(&mut self, label: Option<&str>) -> String {
         self.inner.snapshot(label)
     }
 
diff --git a/rust/lance-context/src/lib.rs b/rust/lance-context/src/lib.rs
index af97519..3ce93cb 100644
--- a/rust/lance-context/src/lib.rs
+++ b/rust/lance-context/src/lib.rs
@@ -1,10 +1,14 @@
 //! Core types for the lance-context storage layer.
 
+use std::collections::HashMap;
+
 #[derive(Debug, Clone)]
 pub struct Context {
     uri: String,
     branch: String,
-    entries: u64,
+    next_id: u64,
+    entries: Vec<ContextEntry>,
+    snapshots: HashMap<String, Snapshot>,
 }
 
 impl Context {
@@ -12,7 +16,9 @@ impl Context {
         Self {
             uri: uri.into(),
             branch: "main".to_string(),
-            entries: 0,
+            next_id: 1,
+            entries: Vec::new(),
+            snapshots: HashMap::new(),
         }
     }
 
@@ -25,27 +31,66 @@ impl Context {
     }
 
     pub fn entries(&self) -> u64 {
-        self.entries
+        self.entries.len() as u64
     }
 
-    pub fn add(&mut self, _role: &str, _content: &str, _data_type: Option<&str>) {
-        self.entries += 1;
+    pub fn add(&mut self, role: &str, content: &str, data_type: Option<&str>) -> u64 {
+        let entry_id = self.next_id;
+        self.next_id += 1;
+        self.entries.push(ContextEntry {
+            id: entry_id,
+            role: role.to_string(),
+            data_type: data_type.map(|value| value.to_string()),
+            content: content.to_string(),
+        });
+        entry_id
     }
 
-    pub fn snapshot(&self, label: Option<&str>) -> String {
-        match label {
+    pub fn snapshot(&mut self, label: Option<&str>) -> String {
+        let id = match label {
             Some(label) if !label.is_empty() => label.to_string(),
-            _ => format!("snapshot-{}", self.entries),
-        }
+            _ => format!("snapshot-{}", self.entries.len()),
+        };
+        let snapshot = Snapshot {
+            id: id.clone(),
+            label: label.map(|value| value.to_string()),
+            entry_count: self.entries.len() as u64,
+            branch: self.branch.clone(),
+        };
+        self.snapshots.insert(id.clone(), snapshot);
+        id
     }
 
     pub fn fork(&self, branch_name: impl Into<String>) -> Self {
         Self {
             uri: self.uri.clone(),
             branch: branch_name.into(),
-            entries: self.entries,
+            next_id: self.next_id,
+            entries: self.entries.clone(),
+            snapshots: self.snapshots.clone(),
+        }
+    }
+
+    pub fn checkout(&mut self, snapshot_id: &str) {
+        if let Some(snapshot) = self.snapshots.get(snapshot_id) {
+            self.entries.truncate(snapshot.entry_count as usize);
+            self.next_id = self.entries.len() as u64 + 1;
         }
     }
+}
 
-    pub fn checkout(&mut self, _snapshot_id: &str) {}
+#[derive(Debug, Clone)]
+pub struct ContextEntry {
+    pub id: u64,
+    pub role: String,
+    pub data_type: Option<String>,
+    pub content: String,
+}
+
+#[derive(Debug, Clone)]
+pub struct Snapshot {
+    pub id: String,
+    pub label: Option<String>,
+    pub entry_count: u64,
+    pub branch: String,
 }