Clean up dropped SDK calls and make drain assertions order-independent

PROTOCOL.md

@@ -274,10 +274,14 @@ IO; only overlapping opens and rename/removal wait for one another.
 Watch the narrowest useful path: delivery is not rate-limited, so watching a
 busy parent can cause a rescan for every write.
 
-The watch ends with its caller; there is no cancellation API or polling/retry
-loop. SDK callers should poll the watch future alongside their handler's work
-and drop it on return, not keep it in a background task that prevents an idle
-process from exiting. Dropping the future does not cancel the host watch.
+The host-side watch ends only when its caller process exits, not on a change
+notification; there is no cancellation API or polling/retry loop. SDK callers
+should poll the watch future alongside their handler's work and drop it on
+return, not keep it in a background task that prevents an idle process from
+exiting. Dropping the future removes the SDK's pending entry and callback, so
+late notifications and responses are ignored; it does not cancel the host
+watch. A long-lived service keeps one watch per purpose for its lifetime,
+rather than repeatedly opening and dropping watches.
 
 ## 5. Endpoints
 

sdk/src/calls.rs

@@ -1,5 +1,19 @@
 //! Outgoing core calls and notification callbacks.
 use super::*;
+/// Removes a dropped call's local routing state, without cancelling host work.
+struct PendingCall(u64);
+impl Drop for PendingCall {
+    fn drop(&mut self) {
+        let pending = STATE.with(|s| s.borrow_mut().pending.remove(&self.0));
+        if let Some(mut pending) = pending {
+            // A lock may still be granted: release it when its late response arrives.
+            if let Some(lease) = pending.unlock.take() {
+                STATE.with(|s| s.borrow_mut().abandoned_locks.insert(self.0, lease));
+            }
+            // Drop user-owned callback captures outside the STATE borrow.
+        }
+    }
+}
 /// An outgoing request builder.
 #[must_use = "await the call to send it"]
 pub struct Call {
@@ -26,12 +40,16 @@ impl Call {
         self
     }
     /// Receive progress notifications before the result arrives.
+    /// Dropping the call future removes its local pending entry and callback;
+    /// later notifications and responses are ignored. It does not cancel host work.
     ///
     /// For `core.watch`, poll the call future alongside work in the handler, then
     /// drop it when the handler returns. Do not put it in `background()` just to
     /// keep the watch alive: background tasks keep the process alive, while the
     /// watch only ends when the process ends. Dropping the future does not cancel
-    /// the host watch, but it lets an otherwise idle process exit normally.
+    /// the host watch, but it lets an otherwise idle process exit normally. A
+    /// long-lived service should keep one watch per purpose for its lifetime,
+    /// rather than repeatedly opening and dropping watches on the host.
     ///
     /// Scan only after `armed`, and rescan on `changed`. Events may coalesce, but
     /// are not rate-limited: watch the narrowest useful path to avoid rescanning
@@ -71,6 +89,7 @@ impl IntoFuture for Call {
                 );
                 Ok(id)
             })?;
+            let _pending = PendingCall(id);
             let (method, params) = if self.detach {
                 (
                     "core.call".to_string(),
@@ -80,7 +99,8 @@ impl IntoFuture for Call {
                 (self.method, self.params)
             };
             if let Err(e) = send(json!({"jsonrpc":"2.0","id":id,"method":method,"params":params})) {
-                STATE.with(|s| s.borrow_mut().pending.remove(&id));
+                let pending = STATE.with(|s| s.borrow_mut().pending.remove(&id));
+                drop(pending);
                 return Err(e);
             }
             rx.await
@@ -88,3 +108,154 @@ impl IntoFuture for Call {
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{STATE, State, block_on, call, clear_pending, lock, receive};
+    use crate::Result;
+    use futures::{FutureExt, future::LocalBoxFuture};
+    use serde_json::{Value, json};
+    use std::{cell::RefCell, future::IntoFuture, rc::Rc, task::Context};
+
+    fn activate() {
+        STATE.with(|s| {
+            *s.borrow_mut() = State {
+                active: true,
+                ..State::default()
+            };
+        });
+    }
+    fn start(
+        mut call: LocalBoxFuture<'static, Result<Value>>,
+    ) -> (u64, LocalBoxFuture<'static, Result<Value>>) {
+        assert!(
+            call.as_mut()
+                .poll(&mut Context::from_waker(std::task::Waker::noop()))
+                .is_pending()
+        );
+        (STATE.with(|s| s.borrow().next), call)
+    }
+    fn notification(id: u64, name: &str) {
+        receive(json!({"method":"core.notify","params":{"call":id,"name":name,"params":{}}}));
+    }
+    fn response(id: u64, result: Value) {
+        receive(json!({"id":id,"result":result}));
+    }
+
+    #[test]
+    fn dropped_call_removes_pending_and_callback() {
+        activate();
+        let events = Rc::new(RefCell::new(Vec::new()));
+        let captured = events.clone();
+        let (id, future) = start(
+            call("core.watch", json!({"path":"/mailbox"}))
+                .on_notify(move |name, _| captured.borrow_mut().push(name.to_owned()))
+                .into_future(),
+        );
+        notification(id, "armed");
+        assert_eq!(&*events.borrow(), &["armed"]);
+        assert_eq!(Rc::strong_count(&events), 2);
+        drop(future);
+        STATE.with(|s| {
+            assert!(s.borrow().pending.is_empty());
+            assert!(s.borrow().abandoned_locks.is_empty());
+        });
+        assert_eq!(Rc::strong_count(&events), 1);
+        notification(id, "changed");
+        response(id, Value::Null);
+        receive(json!({"id":id,"error":{"code":-32000,"message":"late error"}}));
+        assert_eq!(&*events.borrow(), &["armed"]);
+    }
+
+    #[test]
+    fn kept_watch_delivers_until_result() {
+        activate();
+        let events = Rc::new(RefCell::new(Vec::new()));
+        let captured = events.clone();
+        let (id, future) = start(
+            call("core.watch", json!({"path":"/mailbox"}))
+                .on_notify(move |name, _| captured.borrow_mut().push(name.to_owned()))
+                .into_future(),
+        );
+        notification(id, "armed");
+        notification(id, "changed");
+        notification(id, "changed");
+        response(id, json!({"done":true}));
+        assert_eq!(block_on(future).unwrap(), json!({"done":true}));
+        notification(id, "late");
+        assert_eq!(&*events.borrow(), &["armed", "changed", "changed"]);
+        STATE.with(|s| assert!(s.borrow().pending.is_empty()));
+        assert_eq!(Rc::strong_count(&events), 1);
+    }
+
+    #[test]
+    fn callback_can_drop_its_own_call_without_being_restored() {
+        activate();
+        let holder = Rc::new(RefCell::new(None));
+        let captured = holder.clone();
+        let (id, future) = start(
+            call("core.watch", json!({"path":"/mailbox"}))
+                .on_notify(move |_, _| {
+                    // Panics on a second invocation; the first must remove the entry.
+                    drop(captured.borrow_mut().take().expect("live call"));
+                })
+                .into_future(),
+        );
+        *holder.borrow_mut() = Some(future);
+        notification(id, "armed");
+        STATE.with(|s| assert!(s.borrow().pending.is_empty()));
+        notification(id, "changed");
+        response(id, Value::Null);
+        assert_eq!(Rc::strong_count(&holder), 1);
+    }
+
+    #[test]
+    fn eof_cleanup_can_drop_callback_owned_futures() {
+        activate();
+        let holder = Rc::new(RefCell::new(None));
+        let weak = Rc::downgrade(&holder);
+        let captured = holder.clone();
+        let (_, future) = start(
+            call("core.watch", json!({"path":"/mailbox"}))
+                .on_notify(move |_, _| {
+                    let _ = &captured;
+                })
+                .into_future(),
+        );
+        *holder.borrow_mut() = Some(future);
+        drop(holder);
+        STATE.with(|s| s.borrow_mut().stopping = true);
+        clear_pending();
+        assert!(weak.upgrade().is_none());
+        STATE.with(|s| assert!(s.borrow().pending.is_empty()));
+    }
+
+    #[test]
+    fn unpolled_call_never_registers() {
+        activate();
+        drop(call("core.watch", json!({"path":"/mailbox"})).into_future());
+        STATE.with(|s| {
+            assert!(s.borrow().pending.is_empty());
+            assert_eq!(s.borrow().next, 0);
+        });
+    }
+
+    #[test]
+    fn dropped_lock_keeps_only_late_grant_cleanup() {
+        activate();
+        let (id, future) = start(
+            lock("/late")
+                .into_future()
+                .map(|r| r.map(|_| Value::Null))
+                .boxed_local(),
+        );
+        drop(future);
+        STATE.with(|s| {
+            assert!(s.borrow().pending.is_empty());
+            assert_eq!(s.borrow().abandoned_locks.len(), 1);
+        });
+        notification(id, "ignored");
+        response(id, json!({"held":true}));
+        STATE.with(|s| assert!(s.borrow().abandoned_locks.is_empty()));
+    }
+}

sdk/src/transport.rs

@@ -99,6 +99,8 @@ struct State {
     draining: bool,
     next: u64,
     pending: HashMap<u64, Pending>,
+    // Only leases, never callbacks or response senders, survive dropped lock acquisitions.
+    abandoned_locks: HashMap<u64, std::rc::Rc<Lease>>,
     #[cfg(target_os = "wasi")]
     tasks: usize,
     #[cfg(target_os = "wasi")]
@@ -109,6 +111,18 @@ struct State {
     listeners: HashMap<u64, futures::future::AbortHandle>,
 }
 thread_local! { static STATE: RefCell<State> = RefCell::new(State::default()); }
+#[cfg(any(target_os = "wasi", test))]
+fn clear_pending() {
+    let pending = STATE.with(|s| {
+        let mut s = s.borrow_mut();
+        (
+            std::mem::take(&mut s.pending),
+            std::mem::take(&mut s.abandoned_locks),
+        )
+    });
+    // Callback captures may own other call futures whose drop guards use STATE.
+    drop(pending);
+}
 pub(crate) fn endpoint_name<T: 'static>() -> Result<String> {
     STATE
         .with(|s| {
@@ -224,16 +238,16 @@ pub fn test_run<F: Future>(f: F) -> F::Output {
             STATE.with(|s| {
                 let mut s = s.borrow_mut();
                 s.stopping = true;
-                s.pending.clear();
             });
+            clear_pending();
         });
         let result = f.await;
         drop(reader);
         STATE.with(|s| {
             let mut s = s.borrow_mut();
             s.active = false;
-            s.pending.clear();
         });
+        clear_pending();
         result
     })
 }
@@ -305,10 +319,19 @@ pub fn receive(value: Value) {
                 .set(result.as_ref().is_ok_and(|v| v["held"] == true));
         }
         let _ = p.tx.send(result);
+    } else {
+        let lease = value["id"]
+            .as_u64()
+            .and_then(|id| STATE.with(|s| s.borrow_mut().abandoned_locks.remove(&id)));
+        if let Some(lease) = lease {
+            lease
+                .held
+                .set(value.get("error").is_none() && value["result"]["held"] == true);
         } else {
             eprintln!("note: response with unknown ID {}", value["id"]);
         }
     }
+}
 
 /// Serve the main interface until its work ends or stdin reaches EOF.
 #[cfg(target_os = "wasi")]
@@ -356,11 +379,11 @@ async fn serve<T: Interface>(value: T) {
                 STATE.with(|s| {
                     let mut s = s.borrow_mut();
                     s.stopping = true;
-                    s.pending.clear();
                     for (_, listener) in s.listeners.drain() {
                         listener.abort();
                     }
                 });
+                clear_pending();
                 value.stop().await;
                 break;
             }
@@ -411,8 +434,8 @@ async fn serve<T: Interface>(value: T) {
     STATE.with(|s| {
         let mut s = s.borrow_mut();
         s.active = false;
-        s.pending.clear();
     });
+    clear_pending();
     let _ = std::io::stdout().flush();
     std::process::exit(0);
 }

sdk/tests/plugins/probe/src/main.rs

@@ -98,6 +98,36 @@ impl Probe {
         .await
     }
 
+    /// Drop one watch while another still delivers, in a process that remains alive.
+    async fn drop_watch(&self) -> Result<Vec<String>> {
+        use std::{cell::RefCell, future::IntoFuture, rc::Rc, task::Poll};
+        let events = Rc::new(RefCell::new(Vec::new()));
+        let captured = events.clone();
+        let mut dropped = call("core.watch", json!({"path":"/dropped"}))
+            .on_notify(move |_, _| captured.borrow_mut().push("dropped".to_owned()))
+            .into_future();
+        std::future::poll_fn(|cx| {
+            assert!(dropped.as_mut().poll(cx).is_pending());
+            Poll::Ready(())
+        })
+        .await;
+        drop(dropped);
+        let captured = events.clone();
+        let mut kept = call("core.watch", json!({"path":"/kept"}))
+            .on_notify(move |name, _| captured.borrow_mut().push(name.to_owned()))
+            .into_future();
+        std::future::poll_fn(|cx| {
+            assert!(kept.as_mut().poll(cx).is_pending());
+            Poll::Ready(())
+        })
+        .await;
+        // The harness sends late traffic and kept-watch events before answering.
+        call("other.barrier", json!({})).await?;
+        drop(kept);
+        let result = events.borrow().clone();
+        Ok(result)
+    }
+
     /// Publish a control endpoint until EOF.
     async fn serve(&self) -> Result<()> {
         sausage_plugin::listen("/endpoint", std::rc::Rc::new(Control)).await

sdk/tests/wasi.rs

@@ -117,6 +117,25 @@ fn protocol() {
     p.send(json!({"jsonrpc":"2.0","method":"core.notify","params":{"call":watch["id"],"name":"armed","params":{}}}));
     assert_eq!(p.read()["result"], Value::Null);
     p.finishes();
+    // Keep the process alive after dropping a call: late traffic is ignored,
+    // while the kept watch continues to deliver every notification.
+    let mut p = start(&wasm, &dir);
+    p.request("drop_watch", json!({}));
+    let dropped = p.read();
+    let kept = p.read();
+    let barrier = p.read();
+    assert_eq!(dropped["params"]["path"], "/dropped");
+    assert_eq!(kept["params"]["path"], "/kept");
+    assert_eq!(barrier["method"], "other.barrier");
+    for watch in [&dropped, &kept] {
+        for name in ["armed", "changed", "changed"] {
+            p.send(json!({"jsonrpc":"2.0","method":"core.notify","params":{"call":watch["id"],"name":name,"params":{}}}));
+        }
+    }
+    p.answer(&dropped, Value::Null);
+    p.answer(&barrier, json!({}));
+    assert_eq!(p.read()["result"], json!(["armed", "changed", "changed"]));
+    p.finishes();
     let mut p = start(&wasm, &dir);
     p.request("rpc.discover", json!({}));
     let d = p.read();
@@ -276,7 +295,21 @@ fn protocol() {
     assert_eq!(provider.read()["params"]["name"], "draining");
     line.clear();
     assert_eq!(idle.read_line(&mut line).unwrap(), 0);
-    assert!(std::net::TcpStream::connect(&address).is_err());
+    // drain() wakes independent accept/read cancellations. Idle EOF does not
+    // order listener destruction: TCP may still complete a handshake, but that
+    // connection must close without serving work. Do not retry the connect.
+    if let Ok(mut late) = std::net::TcpStream::connect(&address) {
+        eprintln!("TCP handshake raced listener cancellation; waiting for close");
+        // Reuse the four-second socket bound used by idle/active above.
+        late.set_read_timeout(Some(std::time::Duration::from_secs(4)))
+            .unwrap();
+        let closed = std::io::Read::read(&mut late, &mut [0]);
+        assert!(
+            matches!(closed, Ok(0))
+                || matches!(&closed, Err(e) if e.kind() == std::io::ErrorKind::ConnectionReset),
+            "late connection survived drain: {closed:?}"
+        );
+    }
     assert!(provider.child.try_wait().unwrap().is_none());
     provider.answer(&inflight, json!({}));
     line.clear();
@@ -288,6 +321,7 @@ fn protocol() {
     assert!(provider.child.try_wait().unwrap().is_none());
     provider.answer(&background, json!({}));
     provider.finishes();
+    assert!(std::net::TcpStream::connect(&address).is_err());
     assert!(
         !dir.join("stopped").exists(),
         "drain is not the EOF stop hook"