fix(client): retain request ownership revisions

crates/wip-client/src/lib.rs

@@ -219,6 +219,8 @@ pub struct TreeObservation {
     pub children: Vec<String>,
     /// Whether the latest expansion rooted here omitted indexable entries.
     pub truncated: bool,
+    /// Request revision that most recently owned this tree subject.
+    pub revision: RequestId,
     /// Current expansion lifecycle.
     pub state: ObservationState,
 }
@@ -388,6 +390,7 @@ pub enum Completion {
 #[derive(Debug, Clone)]
 struct EntryRecord {
     object_key: Option<u64>,
+    revision: RequestId,
     state: ObservationState,
 }
 
@@ -632,8 +635,10 @@ impl Client {
             .entry(request.path.clone())
             .or_insert(EntryRecord {
                 object_key: None,
+                revision: id,
                 state: ObservationState::Loading(id),
             });
+        entry.revision = id;
         entry.state = ObservationState::Loading(id);
         self.pending.insert(
             id,
@@ -678,8 +683,10 @@ impl Client {
                 path: request.path.clone(),
                 children: Vec::new(),
                 truncated: false,
+                revision: id,
                 state: ObservationState::Loading(id),
             });
+        tree.revision = id;
         tree.state = ObservationState::Loading(id);
         self.pending.insert(
             id,
@@ -1001,7 +1008,9 @@ impl Client {
             Ok(DecodedResponse::Success(object)) => {
                 let maximum = self.limits.objects_per_session;
                 let mut candidate = self.session(session)?.clone();
-                if let Err(error) = insert_object(&mut candidate, &request.path, object, maximum) {
+                if let Err(error) =
+                    insert_object(&mut candidate, &request.path, object, maximum, id)
+                {
                     self.session_mut(session)?
                         .entries
                         .get_mut(&request.path)
@@ -1549,7 +1558,13 @@ fn ensure_entry_slot(state: &mut Session, path: &str, maximum: usize) -> Result<
     let candidate = state
         .entries
         .iter()
-        .find(|(_, entry)| !matches!(entry.state, ObservationState::Loading(_)))
+        .find(|(path, entry)| {
+            !matches!(entry.state, ObservationState::Loading(_))
+                && !state
+                    .trees
+                    .get(*path)
+                    .is_some_and(|tree| matches!(tree.state, ObservationState::Loading(_)))
+        })
         .map(|(path, _)| path.clone())
         .ok_or_else(|| {
             ClientError::new(
@@ -1604,6 +1619,7 @@ fn insert_object(
     path: &str,
     object: Object,
     maximum: usize,
+    revision: RequestId,
 ) -> Result<u64, ClientError> {
     ensure_entry_slot(state, path, maximum)?;
     let existing_path_key = state.entries.get(path).and_then(|entry| entry.object_key);
@@ -1658,6 +1674,7 @@ fn insert_object(
         path.to_owned(),
         EntryRecord {
             object_key: Some(key),
+            revision,
             state: ObservationState::Fresh,
         },
     );
@@ -1675,9 +1692,12 @@ fn make_object_slot(
                 .objects
                 .iter()
                 .filter(|(key, _)| {
-                    !state.entries.values().any(|entry| {
+                    !state.entries.iter().any(|(path, entry)| {
                         entry.object_key == Some(**key)
-                            && matches!(entry.state, ObservationState::Loading(_))
+                            && (matches!(entry.state, ObservationState::Loading(_))
+                                || state.trees.get(path).is_some_and(|tree| {
+                                    matches!(tree.state, ObservationState::Loading(_))
+                                }))
                     })
                 })
                 .min_by_key(|(_, record)| record.touched)
@@ -1712,12 +1732,24 @@ fn remove_object(state: &mut Session, key: u64) {
     state
         .entries
         .retain(|_, entry| entry.object_key != Some(key));
-    for path in &removed_paths {
+    let removable_paths: Vec<&String> = removed_paths
-        state.trees.remove(path);
+        .iter()
+        .filter(|path| {
+            !state
+                .trees
+                .get(*path)
+                .is_some_and(|tree| matches!(tree.state, ObservationState::Loading(_)))
+        })
+        .collect();
+    for path in &removable_paths {
+        state.trees.remove(*path);
     }
     for tree in state.trees.values_mut() {
-        tree.children
+        tree.children.retain(|path| {
-            .retain(|path| !removed_paths.iter().any(|removed| removed == path));
+            !removable_paths
+                .iter()
+                .any(|removed| removed.as_str() == path)
+        });
     }
 }
 
@@ -1737,7 +1769,12 @@ fn reserve_tree_metadata(
             .entries
             .iter()
             .find(|(path, entry)| {
-                !paths.contains(path) && !matches!(entry.state, ObservationState::Loading(_))
+                !paths.contains(path)
+                    && !matches!(entry.state, ObservationState::Loading(_))
+                    && !state
+                        .trees
+                        .get(*path)
+                        .is_some_and(|tree| matches!(tree.state, ObservationState::Loading(_)))
             })
             .map(|(path, _)| path.clone())
             .ok_or_else(|| {
@@ -1773,19 +1810,15 @@ fn reserve_tree_metadata(
     Ok(())
 }
 
-fn mark_subtree_stale(state: &mut Session, root: &str) {
+fn mark_subtree_stale(state: &mut Session, root: &str, request_id: RequestId) {
     let prefix = format!("{root}/");
     for (path, entry) in &mut state.entries {
-        if (path == root || path.starts_with(&prefix))
+        if (path == root || path.starts_with(&prefix)) && entry.revision <= request_id {
-            && !matches!(entry.state, ObservationState::Loading(_))
-        {
             entry.state = ObservationState::Stale;
         }
     }
     for (path, tree) in &mut state.trees {
-        if (path == root || path.starts_with(&prefix))
+        if (path == root || path.starts_with(&prefix)) && tree.revision <= request_id {
-            && !matches!(tree.state, ObservationState::Loading(_))
-        {
             tree.state = ObservationState::Stale;
         }
     }
@@ -1812,7 +1845,11 @@ fn apply_tree(
     if !response.truncated {
         let mut omitted = Vec::new();
         for (path, _, children) in &nodes {
-            if let Some(previous) = state.trees.get(path) {
+            if let Some(previous) = state
+                .trees
+                .get(path)
+                .filter(|previous| previous.revision <= request_id)
+            {
                 omitted.extend(
                     previous
                         .children
@@ -1823,24 +1860,24 @@ fn apply_tree(
             }
         }
         for path in omitted {
-            mark_subtree_stale(state, &path);
+            mark_subtree_stale(state, &path, request_id);
         }
     }
 
     for (path, object, _) in &nodes {
-        let newer_owner = matches!(
+        let newer_owner = state
-            state.entries.get(path).map(|entry| &entry.state),
+            .entries
-            Some(ObservationState::Loading(owner)) if *owner > request_id
+            .get(path)
-        );
+            .is_some_and(|entry| entry.revision > request_id);
         if !newer_owner {
-            insert_object(state, path, object.clone(), maximum)?;
+            insert_object(state, path, object.clone(), maximum, request_id)?;
         }
     }
     for (path, _, children) in nodes {
-        let newer_owner = matches!(
+        let newer_owner = state
-            state.trees.get(&path).map(|tree| &tree.state),
+            .trees
-            Some(ObservationState::Loading(owner)) if *owner > request_id
+            .get(&path)
-        );
+            .is_some_and(|tree| tree.revision > request_id);
         if newer_owner {
             continue;
         }
@@ -1851,6 +1888,7 @@ fn apply_tree(
                 path,
                 children,
                 truncated: is_root && response.truncated,
+                revision: request_id,
                 state: ObservationState::Fresh,
             },
         );

crates/wip-client/tests/runtime.rs

@@ -535,6 +535,54 @@ fn overlapping_tree_work_preserves_newer_subject_request_ownership() {
         Some(b"newer-fetch".to_vec())
     );
 
+    let older_tree = client.prepare_tree(&session, "/", 1, 4).unwrap();
+    let completed_newer_fetch = client.prepare_fetch(&session, "/child").unwrap();
+    client
+        .complete(
+            completed_newer_fetch.id,
+            fetch_response("/child", &object("child", None, b"completed-newer")),
+        )
+        .unwrap();
+    client
+        .complete(
+            older_tree.id,
+            encode_fetch_tree_response(&tree_request, &tree_value, wire_limits()).unwrap(),
+        )
+        .unwrap();
+    assert_eq!(
+        client.object(&session, "/child").unwrap().validator,
+        Some(b"completed-newer".to_vec())
+    );
+
+    let older_omitting_tree = client.prepare_tree(&session, "/", 1, 4).unwrap();
+    let newer_fetch = client.prepare_fetch(&session, "/child").unwrap();
+    client
+        .complete(
+            newer_fetch.id,
+            fetch_response("/child", &object("child", None, b"newer-than-omission")),
+        )
+        .unwrap();
+    client
+        .complete(
+            older_omitting_tree.id,
+            encode_fetch_tree_response(
+                &tree_request,
+                &FetchTreeResponse {
+                    root: ObjectTreeNode {
+                        object: object("", None, b"older-root"),
+                        children: vec![],
+                    },
+                    truncated: false,
+                },
+                wire_limits(),
+            )
+            .unwrap(),
+        )
+        .unwrap();
+    let child = client.object(&session, "/child").unwrap();
+    assert_eq!(child.state, ObservationState::Fresh);
+    assert_eq!(child.validator, Some(b"newer-than-omission".to_vec()));
+
     let old_tree = client.prepare_tree(&session, "/", 1, 4).unwrap();
     let newer_child_tree = client.prepare_tree(&session, "/child", 0, 1).unwrap();
     client
@@ -569,6 +617,40 @@ fn overlapping_tree_work_preserves_newer_subject_request_ownership() {
         client.tree(&session, "/child").unwrap().state,
         ObservationState::Fresh
     );
+
+    let older_parent = client.prepare_tree(&session, "/", 1, 4).unwrap();
+    let newer_child = client.prepare_tree(&session, "/child", 0, 1).unwrap();
+    client
+        .complete(
+            newer_child.id,
+            encode_fetch_tree_response(
+                &child_request,
+                &FetchTreeResponse {
+                    root: ObjectTreeNode {
+                        object: object("child", None, b"completed-newer-tree"),
+                        children: vec![],
+                    },
+                    truncated: false,
+                },
+                wire_limits(),
+            )
+            .unwrap(),
+        )
+        .unwrap();
+    client
+        .complete(
+            older_parent.id,
+            encode_fetch_tree_response(&tree_request, &tree_value, wire_limits()).unwrap(),
+        )
+        .unwrap();
+    assert_eq!(
+        client.object(&session, "/child").unwrap().validator,
+        Some(b"completed-newer-tree".to_vec())
+    );
+    assert_eq!(
+        client.tree(&session, "/child").unwrap().revision,
+        newer_child.id
+    );
 }
 
 #[test]
@@ -701,6 +783,49 @@ fn validator_free_interface_observation_remains_immutable() {
     assert_eq!(client.interface(&session, "math").unwrap().validator, None);
 }
 
+#[test]
+fn entry_eviction_never_discards_a_loading_tree_subject() {
+    let mut client = client(1, 8, 8, 8);
+    let session = client
+        .open_session("https://example.test/wip", SecurityContext::new("subject"))
+        .unwrap();
+    observe_object(&mut client, &session, "/", object("", None, b"root"));
+    let tree = client.prepare_tree(&session, "/", 0, 1).unwrap();
+    assert_eq!(
+        client.prepare_fetch(&session, "/other").unwrap_err().kind,
+        ClientErrorKind::Capacity
+    );
+    assert!(matches!(
+        client.tree(&session, "/").unwrap().state,
+        ObservationState::Loading(id) if id == tree.id
+    ));
+    client
+        .complete(
+            tree.id,
+            encode_fetch_tree_response(
+                &FetchTreeRequest {
+                    path: "/".into(),
+                    depth: 0,
+                    limit: 1,
+                },
+                &FetchTreeResponse {
+                    root: ObjectTreeNode {
+                        object: object("", None, b"tree"),
+                        children: vec![],
+                    },
+                    truncated: false,
+                },
+                wire_limits(),
+            )
+            .unwrap(),
+        )
+        .unwrap();
+    assert_eq!(
+        client.tree(&session, "/").unwrap().state,
+        ObservationState::Fresh
+    );
+}
+
 #[test]
 fn failed_entry_and_tree_observation_metadata_is_bounded() {
     let mut client = client(2, 8, 2, 8);