fix(client): preserve operation observation ownership

crates/wip-client/src/lib.rs

@@ -398,6 +398,7 @@ struct EntryRecord {
 struct ObjectRecord {
     object: Object,
     revision: RequestId,
+    call_owners: usize,
     touched: u64,
 }
 
@@ -455,7 +456,10 @@ impl Session {
         })
     }
 
-    fn mark_object_stale(&mut self, key: u64) {
+    fn mark_object_stale(&mut self, key: u64, revision: RequestId) {
+        if let Some(record) = self.objects.get_mut(&key) {
+            record.revision = record.revision.max(revision);
+        }
         for entry in self.entries.values_mut() {
             if entry.object_key == Some(key) {
                 entry.state = ObservationState::Stale;
@@ -463,6 +467,29 @@ impl Session {
         }
     }
 
+    fn pin_call_object(&mut self, key: u64) {
+        let record = self
+            .objects
+            .get_mut(&key)
+            .expect("call context object remains observed");
+        record.call_owners = record.call_owners.saturating_add(1);
+    }
+
+    fn release_call_object(&mut self, key: u64) {
+        let Some(record) = self.objects.get_mut(&key) else {
+            return;
+        };
+        record.call_owners = record.call_owners.saturating_sub(1);
+        let remove = record.call_owners == 0
+            && !self
+                .entries
+                .values()
+                .any(|entry| entry.object_key == Some(key));
+        if remove {
+            remove_object(self, key);
+        }
+    }
+
     fn push_diagnostic(&mut self, value: Diagnostic, maximum: usize) {
         push_bounded(&mut self.diagnostics, value, maximum);
     }
@@ -863,6 +890,8 @@ impl Client {
         )
         .map_err(|error| ClientError::new(ClientErrorKind::InvalidRequest, error.to_string()))?;
         let id = self.allocate_request_id();
+        self.session_mut(&context.session)?
+            .pin_call_object(context.object_key);
         self.pending.insert(
             id,
             Pending::Call {
@@ -908,7 +937,13 @@ impl Client {
             Pending::Interface { session, request } => {
                 self.complete_interface(id, &session, request, response)
             }
-            Pending::Call { context, .. } => self.complete_call(id, *context, response),
+            Pending::Call { context, .. } => {
+                let session = context.session.clone();
+                let object_key = context.object_key;
+                let result = self.complete_call(id, *context, response);
+                self.session_mut(&session)?.release_call_object(object_key);
+                result
+            }
         }
     }
 
@@ -924,16 +959,24 @@ impl Client {
             Pending::Call {
                 context,
                 dispatched: true,
-            } => self.finish_unknown_call(
+            } => {
+                let session = context.session.clone();
+                let object_key = context.object_key;
+                let result = self.finish_unknown_call(
                     id,
                     *context,
                     OutcomeUnknownReason::Disconnect,
                     ClientError::new(ClientErrorKind::Transport, detail),
-            ),
+                );
+                self.session_mut(&session)?.release_call_object(object_key);
+                result
+            }
             Pending::Call {
                 context,
                 dispatched: false,
             } => {
+                let session = context.session.clone();
+                let object_key = context.object_key;
                 let record = CallRecord {
                     request_id: id,
                     context: *context,
@@ -943,8 +986,9 @@ impl Client {
                     )),
                 };
                 let maximum = self.limits.history;
-                self.session_mut(record.context.session())?
+                let state = self.session_mut(&session)?;
-                    .push_call(record.clone(), maximum);
+                state.push_call(record.clone(), maximum);
+                state.release_call_object(object_key);
                 Ok(Completion::Call(Box::new(record)))
             }
             other => {
@@ -998,12 +1042,16 @@ impl Client {
             DispatchedTransportFailure::Timeout => OutcomeUnknownReason::Timeout,
             DispatchedTransportFailure::Disconnect => OutcomeUnknownReason::Disconnect,
         };
-        self.finish_unknown_call(
+        let session = context.session.clone();
+        let object_key = context.object_key;
+        let result = self.finish_unknown_call(
             id,
             *context,
             reason,
             ClientError::new(ClientErrorKind::Transport, detail),
-        )
+        );
+        self.session_mut(&session)?.release_call_object(object_key);
+        result
     }
 
     fn complete_fetch(
@@ -1213,18 +1261,11 @@ impl Client {
                 let session_id = context.session.clone();
                 let maximum = self.limits.history;
                 let state = self.session_mut(&session_id)?;
-                if state
-                    .entries
-                    .get(&context.request.target.path)
-                    .and_then(|entry| entry.object_key)
-                    == Some(context.object_key)
-                {
                 if let Some(object) = state.objects.get_mut(&context.object_key) {
                     object.object.validator = value.validator.clone();
-                        object.revision = id;
+                    object.revision = object.revision.max(id);
-                    }
-                    state.mark_object_stale(context.object_key);
                 }
+                state.mark_object_stale(context.object_key, id);
                 let record = CallRecord {
                     request_id: id,
                     context,
@@ -1242,7 +1283,7 @@ impl Client {
                 if error.code == ProtocolErrorCode::OperationOutcomeUnknown {
                     let maximum = self.limits.history;
                     let state = self.session_mut(&session_id)?;
-                    state.mark_object_stale(context.object_key);
+                    state.mark_object_stale(context.object_key, id);
                     let record = CallRecord {
                         request_id: id,
                         context,
@@ -1260,7 +1301,7 @@ impl Client {
                     ProtocolErrorCode::ValidatorRequired
                     | ProtocolErrorCode::ValidatorMismatch
                     | ProtocolErrorCode::InterfaceMismatch => {
-                        state.mark_object_stale(context.object_key);
+                        state.mark_object_stale(context.object_key, id);
                     }
                     ProtocolErrorCode::InterfaceValidatorRequired
                     | ProtocolErrorCode::InterfaceValidatorMismatch => {
@@ -1298,7 +1339,7 @@ impl Client {
         let session_id = context.session.clone();
         let maximum = self.limits.history;
         let state = self.session_mut(&session_id)?;
-        state.mark_object_stale(context.object_key);
+        state.mark_object_stale(context.object_key, id);
         state.push_diagnostic(
             Diagnostic {
                 request_id: id,
@@ -1596,6 +1637,12 @@ fn ensure_entry_slot(state: &mut Session, path: &str, maximum: usize) -> Result<
         .iter()
         .find(|(path, entry)| {
             !matches!(entry.state, ObservationState::Loading(_))
+                && !entry.object_key.is_some_and(|key| {
+                    state
+                        .objects
+                        .get(&key)
+                        .is_some_and(|record| record.call_owners > 0)
+                })
                 && !state
                     .trees
                     .get(*path)
@@ -1605,7 +1652,7 @@ fn ensure_entry_slot(state: &mut Session, path: &str, maximum: usize) -> Result<
         .ok_or_else(|| {
             ClientError::new(
                 ClientErrorKind::Capacity,
-                "all entry observations are loading",
+                "all entry observations are owned by in-flight work",
             )
         })?;
     remove_entry(state, &candidate);
@@ -1668,9 +1715,16 @@ fn insert_object(
     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);
+    let existing_path_entry = state.entries.get(path).cloned();
+    let existing_path_key = existing_path_entry
+        .as_ref()
+        .and_then(|entry| entry.object_key);
     let replaceable_key = existing_path_key.filter(|key| {
-        !state
+        state
+            .objects
+            .get(key)
+            .is_none_or(|record| record.call_owners == 0)
+            && !state
                 .entries
                 .iter()
                 .any(|(entry_path, entry)| entry_path != path && entry.object_key == Some(*key))
@@ -1678,7 +1732,6 @@ fn insert_object(
     let key = if let Some(reference) = &object.r#ref {
         if let Some(key) = state.refs.get(reference).copied() {
             if let Some(existing) = state.objects.get(&key)
-                && object_ownership_revision(state, key).is_none_or(|owner| owner <= revision)
                 && existing.object.validator.is_some()
                 && existing.object.validator == object.validator
                 && existing.object != object
@@ -1717,8 +1770,13 @@ fn insert_object(
     }
     let touched = state.tick();
     let ownership_revision = object_ownership_revision(state, key).unwrap_or(revision);
+    let suppressed = ownership_revision > revision;
+    let call_owners = state
+        .objects
+        .get(&key)
+        .map_or(0, |record| record.call_owners);
     match state.objects.get_mut(&key) {
-        Some(existing) if ownership_revision > revision => {
+        Some(existing) if suppressed => {
             existing.revision = ownership_revision;
             existing.touched = touched;
         }
@@ -1728,19 +1786,36 @@ fn insert_object(
                 ObjectRecord {
                     object,
                     revision,
+                    call_owners,
                     touched,
                 },
             );
         }
     }
-    state.entries.insert(
+    let entry = if suppressed {
-        path.to_owned(),
+        match existing_path_entry {
+            Some(mut entry) if entry.object_key == Some(key) => {
+                entry.revision = ownership_revision;
+                if matches!(entry.state, ObservationState::Loading(current) if current == revision)
+                {
+                    entry.state = ObservationState::Fresh;
+                }
+                entry
+            }
+            _ => EntryRecord {
+                object_key: Some(key),
+                revision: ownership_revision,
+                state: ObservationState::Stale,
+            },
+        }
+    } else {
         EntryRecord {
             object_key: Some(key),
             revision,
             state: ObservationState::Fresh,
-        },
+        }
-    );
+    };
+    state.entries.insert(path.to_owned(), entry);
     Ok(key)
 }
 
@@ -1754,8 +1829,9 @@ fn make_object_slot(
             state
                 .objects
                 .iter()
-                .filter(|(key, _)| {
+                .filter(|(key, record)| {
-                    !state.entries.iter().any(|(path, entry)| {
+                    record.call_owners == 0
+                        && !state.entries.iter().any(|(path, entry)| {
                             entry.object_key == Some(**key)
                                 && (matches!(entry.state, ObservationState::Loading(_))
                                     || state.trees.get(path).is_some_and(|tree| {
@@ -1769,7 +1845,7 @@ fn make_object_slot(
         let candidate = candidate.ok_or_else(|| {
             ClientError::new(
                 ClientErrorKind::Capacity,
-                "all object cache entries are owned by loading observations",
+                "all object cache entries are owned by in-flight work",
             )
         })?;
         remove_object(state, candidate);
@@ -1780,6 +1856,13 @@ fn make_object_slot(
 }
 
 fn remove_object(state: &mut Session, key: u64) {
+    if state
+        .objects
+        .get(&key)
+        .is_some_and(|record| record.call_owners > 0)
+    {
+        return;
+    }
     if let Some(record) = state.objects.remove(&key)
         && let Some(reference) = record.object.r#ref
         && state.refs.get(&reference) == Some(&key)
@@ -1834,6 +1917,12 @@ fn reserve_tree_metadata(
             .find(|(path, entry)| {
                 !paths.contains(path)
                     && !matches!(entry.state, ObservationState::Loading(_))
+                    && !entry.object_key.is_some_and(|key| {
+                        state
+                            .objects
+                            .get(&key)
+                            .is_some_and(|record| record.call_owners > 0)
+                    })
                     && !state
                         .trees
                         .get(*path)
@@ -1843,7 +1932,7 @@ fn reserve_tree_metadata(
             .ok_or_else(|| {
                 ClientError::new(
                     ClientErrorKind::Capacity,
-                    "tree response cannot evict loading entry observations",
+                    "tree response cannot evict entries owned by in-flight work",
                 )
             })?;
         remove_entry(state, &candidate);

crates/wip-client/tests/runtime.rs

@@ -288,6 +288,45 @@ fn normalized_aliases_preserve_newer_request_ownership() {
         client.object(&session, "/a/shared").unwrap().validator,
         Some(b"newest".to_vec())
     );
+
+    let older_tree = client.prepare_tree(&session, "/a", 1, 4).unwrap();
+    let newer_alias = client.prepare_fetch(&session, "/b/shared").unwrap();
+    client
+        .complete(
+            newer_alias.id,
+            fetch_response(
+                "/b/shared",
+                &object("shared", Some("shared-ref"), b"newest"),
+            ),
+        )
+        .unwrap();
+    let mut contradictory = object("shared", Some("shared-ref"), b"newest");
+    contradictory.description = Some("contradictory representation".into());
+    let error = client
+        .complete(
+            older_tree.id,
+            encode_fetch_tree_response(
+                &request,
+                &FetchTreeResponse {
+                    root: ObjectTreeNode {
+                        object: object("a", None, b"root"),
+                        children: vec![ObjectTreeNode {
+                            object: contradictory,
+                            children: vec![],
+                        }],
+                    },
+                    truncated: false,
+                },
+                wire_limits(),
+            )
+            .unwrap(),
+        )
+        .unwrap_err();
+    assert_eq!(error.kind, ClientErrorKind::InvalidResponse);
+    assert_eq!(
+        client.object(&session, "/b/shared").unwrap().validator,
+        Some(b"newest".to_vec())
+    );
 }
 
 #[test]
@@ -445,6 +484,184 @@ fn call_context_fixes_membership_validators_descriptor_and_result_validation() {
     );
 }
 
+#[test]
+fn older_retrieval_cannot_restore_freshness_after_operation_completion() {
+    let mut client = client(8, 8, 8, 8);
+    let session = client
+        .open_session("https://example.test/wip", SecurityContext::new("subject"))
+        .unwrap();
+    observe_object(
+        &mut client,
+        &session,
+        "/item",
+        object("item", Some("shared-ref"), b"object-v1"),
+    );
+    observe_interface(
+        &mut client,
+        &session,
+        FetchInterfaceResponse {
+            interface: "math".into(),
+            descriptor: descriptor(TypeExpr::Integer),
+            validator: Some(b"interface-v1".to_vec()),
+        },
+    );
+
+    let older_alias = client.prepare_fetch(&session, "/alias/item").unwrap();
+    let call = client
+        .prepare_call(&session, "/item", "math", "calculate", args(7))
+        .unwrap();
+    let request = client
+        .pending_call_context(call.id)
+        .unwrap()
+        .request()
+        .clone();
+    let fixed_descriptor = client
+        .pending_call_context(call.id)
+        .unwrap()
+        .descriptor()
+        .clone();
+    client.mark_dispatched(call.id).unwrap();
+    client
+        .complete(
+            call.id,
+            encode_call_operation_response(
+                &request,
+                &fixed_descriptor,
+                &CallOperationResponse {
+                    result: Value::Integer(14),
+                    validator: Some(b"object-v2".to_vec()),
+                },
+                wire_limits(),
+            )
+            .unwrap(),
+        )
+        .unwrap();
+    assert_eq!(
+        client.object(&session, "/item").unwrap().state,
+        ObservationState::Stale
+    );
+
+    client
+        .complete(
+            older_alias.id,
+            fetch_response(
+                "/alias/item",
+                &object("item", Some("shared-ref"), b"object-v1"),
+            ),
+        )
+        .unwrap();
+    let alias = client.object(&session, "/alias/item").unwrap();
+    assert_eq!(alias.state, ObservationState::Stale);
+    assert_eq!(alias.validator, Some(b"object-v2".to_vec()));
+    assert_eq!(
+        client
+            .prepare_call(&session, "/alias/item", "math", "calculate", args(1))
+            .unwrap_err()
+            .kind,
+        ClientErrorKind::MissingObservation
+    );
+
+    observe_object(
+        &mut client,
+        &session,
+        "/item",
+        object("item", Some("shared-ref"), b"object-v2"),
+    );
+    let older_alias = client.prepare_fetch(&session, "/second/item").unwrap();
+    let call = client
+        .prepare_call(&session, "/item", "math", "calculate", args(8))
+        .unwrap();
+    client.mark_dispatched(call.id).unwrap();
+    let unknown = ProtocolError {
+        code: ProtocolErrorCode::OperationOutcomeUnknown,
+        message: "commit uncertain".into(),
+    };
+    client
+        .complete(
+            call.id,
+            encode_protocol_error_response(
+                ProtocolInteraction::CallOperation,
+                &unknown,
+                wire_limits(),
+            )
+            .unwrap(),
+        )
+        .unwrap();
+    client
+        .complete(
+            older_alias.id,
+            fetch_response(
+                "/second/item",
+                &object("item", Some("shared-ref"), b"object-v2"),
+            ),
+        )
+        .unwrap();
+    assert_eq!(
+        client.object(&session, "/second/item").unwrap().state,
+        ObservationState::Stale
+    );
+}
+
+#[test]
+fn in_flight_calls_pin_target_identity_until_staleness_is_applied() {
+    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,
+        "/item",
+        object("item", Some("shared-ref"), b"object-v1"),
+    );
+    observe_interface(
+        &mut client,
+        &session,
+        FetchInterfaceResponse {
+            interface: "math".into(),
+            descriptor: descriptor(TypeExpr::Integer),
+            validator: Some(b"interface-v1".to_vec()),
+        },
+    );
+    let call = client
+        .prepare_call(&session, "/item", "math", "calculate", args(7))
+        .unwrap();
+    client.mark_dispatched(call.id).unwrap();
+
+    assert_eq!(
+        client.prepare_fetch(&session, "/other").unwrap_err().kind,
+        ClientErrorKind::Capacity
+    );
+    let unknown = ProtocolError {
+        code: ProtocolErrorCode::OperationOutcomeUnknown,
+        message: "commit uncertain".into(),
+    };
+    client
+        .complete(
+            call.id,
+            encode_protocol_error_response(
+                ProtocolInteraction::CallOperation,
+                &unknown,
+                wire_limits(),
+            )
+            .unwrap(),
+        )
+        .unwrap();
+    assert_eq!(
+        client.object(&session, "/item").unwrap().state,
+        ObservationState::Stale
+    );
+
+    observe_object(
+        &mut client,
+        &session,
+        "/other",
+        object("other", None, b"other"),
+    );
+    assert!(client.object(&session, "/item").is_none());
+    assert!(client.object(&session, "/other").is_some());
+}
+
 #[test]
 fn validator_mismatch_stales_only_its_observation_and_never_retries() {
     let mut client = client(8, 8, 8, 8);