Notify replaced plugin processes at upgrade grace start

Send rpc.superseded with the newly active hash before activation returns.
Add an opt-in SDK hook and drain helper that closes SDK listeners and idle
connections without cancelling active handlers or background work.
Document the additive notice and preserve existing grace, stop and restart rules.

Cover immediate notices, ignored notices, early service restarts, SDK draining
and duplicate-hook diagnostics.

Card: SAU-352

sdk/macros/src/lib.rs

@@ -153,6 +153,7 @@ pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream {
         let mut stop = quote!();
         let mut has_stop = false;
         let mut has_retired = false;
+        let mut has_superseded = false;
         for member in &mut item.items {
             let ImplItem::Fn(f) = member else { continue };
             let retired = f.attrs.iter().any(|a| {
@@ -161,6 +162,12 @@ pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream {
                     .last()
                     .is_some_and(|s| s.ident == "on_retired")
             });
+            let superseded = f.attrs.iter().any(|a| {
+                a.path()
+                    .segments
+                    .last()
+                    .is_some_and(|s| s.ident == "on_superseded")
+            });
             let stopping = f.attrs.iter().any(|a| {
                 a.path()
                     .segments
@@ -179,11 +186,16 @@ pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream {
                     "only one on_retired method is allowed",
                 ));
             }
+            if superseded && std::mem::replace(&mut has_superseded, true) {
+                return Err(syn::Error::new(
+                    f.span(),
+                    "only one on_superseded method is allowed",
+                ));
+            }
             f.attrs.retain(|a| {
-                !a.path()
-                    .segments
-                    .last()
-                    .is_some_and(|s| s.ident == "on_retired" || s.ident == "on_stop")
+                !a.path().segments.last().is_some_and(|s| {
+                    s.ident == "on_retired" || s.ident == "on_stop" || s.ident == "on_superseded"
+                })
             });
             documented(&f.attrs, f.span())?;
             if f.sig.inputs.is_empty()
@@ -224,7 +236,9 @@ pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream {
                 stop = quote!(fn stop(&self)->::std::pin::Pin<Box<dyn ::std::future::Future<Output=()> + '_>> {Box::pin(async move {let _=self.#name() #await_;})});
                 continue;
             }
-            let wire = if retired {
+            let wire = if superseded {
+                "rpc.superseded".to_string()
+            } else if retired {
                 "rpc.retired".to_string()
             } else {
                 name.unraw().to_string()
@@ -253,7 +267,7 @@ pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream {
             };
             let doc = docs(&f.attrs);
             arms.push(quote!(#wire => { fn params_type<T: ::sausage_plugin::Params>() {} params_type::<#p>(); let result=self.#name(#(#args),*) #await_ ?; Ok(::sausage_plugin::serde_json::to_value(result)?) }));
-            if !retired {
+            if !retired && !superseded {
                 entries.push(quote!(::sausage_plugin::method::<#p,#output>(&mut doc,#wire,#doc);));
             }
         }
@@ -328,6 +342,12 @@ pub fn import(input: TokenStream) -> TokenStream {
 pub fn on_retired(_: TokenStream, input: TokenStream) -> TokenStream {
     input
 }
+/// Mark the method that handles `rpc.superseded` at the start of upgrade grace.
+/// Calls still work here. Use `sausage_plugin::drain` to finish active work and exit.
+#[proc_macro_attribute]
+pub fn on_superseded(_: TokenStream, input: TokenStream) -> TokenStream {
+    input
+}
 /// Mark the hook that runs at stdin EOF without a live call transport.
 #[proc_macro_attribute]
 pub fn on_stop(_: TokenStream, input: TokenStream) -> TokenStream {

sdk/src/endpoints.rs

@@ -65,7 +65,7 @@ impl<T: serde::de::DeserializeOwned + 'static> IntoFuture for EndpointCall<T> {
     }
 }
 
-/// Publish a declared interface and keep its listener alive until stdin EOF.
+/// Publish a declared interface and keep its listener alive until stdin EOF or [`drain`].
 #[cfg(target_os = "wasi")]
 pub async fn listen<T: Interface>(path: impl AsRef<Path>, value: Rc<T>) -> Result<()> {
     use wstd::iter::AsyncIterator;
@@ -80,7 +80,7 @@ pub async fn listen<T: Interface>(path: impl AsRef<Path>, value: Rc<T>) -> Resul
         while let Some(Ok(stream)) = incoming.next().await {
             let value = value.clone();
             let name = name.clone();
-            crate::transport::endpoint_task(async move {
+            background(async move {
                 let _ = serve_connection(stream, value, name).await;
             });
         }
@@ -211,7 +211,11 @@ async fn serve_connection<T: Interface>(
     use futures::{FutureExt, StreamExt};
     let (mut reader, mut writer) = stream.split();
     let mut lines = crate::lines::Lines::default();
-    while let Some(request) = lines.rpc(&mut reader).await? {
+    while let Some(request) = crate::transport::until_drained(lines.rpc(&mut reader))
+        .await
+        .transpose()?
+        .flatten()
+    {
         let method = request["method"].as_str().unwrap_or_default();
         let (tx, mut rx) = futures::channel::mpsc::unbounded();
         let ctx = Ctx {

sdk/src/lib.rs

@@ -358,11 +358,16 @@
 //! `core.revert` restores code, not state. Preserve compatibility or migrate the state.
 //! `#[rpc(grace_seconds = N)]` declares the grace period for old processes. The default is 60 seconds.
 //! New calls use the new version immediately. Old processes can still use old imports during their grace period.
+//! [`macro@on_superseded`] handles the notice at grace start, with normal calls still available.
+//! It is opt-in: without a hook the process continues as before until core stops it.
+//! Call [`drain`] in this hook to close SDK endpoint listeners and idle connections,
+//! let active handlers and background tasks finish, and exit without waiting for grace.
+//! Custom listeners and long-lived background tasks must be ended by the plugin itself.
 //! [`macro@on_retired`] handles `rpc.retired` after the last process of an old version ends.
 //! This hook is best effort. If the host stops first, the hook does not run.
 //! New processes must handle unmigrated state. Check a migration marker and migrate if it is absent.
 //! See section 3.7 of `/proc/core/PROTOCOL.md`.
-//! Declare at most one stop hook and one retirement hook per interface.
+//! Declare at most one hook of each kind per interface.
 //!
 //! ```no_run
 //! use sausage_plugin::{rpc, schema, state_dir, write_atomic, Result};
@@ -374,8 +379,21 @@
 //!     /// The SHA-256 of the previous WASM file.
 //!     previous: String,
 //! }
+//! /// The version now receiving new calls.
+//! #[schema]
+//! struct SupersededParams {
+//!     /// The SHA-256 of the newly active WASM file.
+//!     current: String,
+//! }
 //! #[rpc]
 //! impl Migrator {
+//!     /// Finish current work and exit early during upgrade grace.
+//!     #[sausage_plugin::on_superseded]
+//!     fn superseded(&self, p: SupersededParams) -> Result<()> {
+//!         eprintln!("replaced by {}", p.current);
+//!         sausage_plugin::drain();
+//!         Ok(())
+//!     }
 //!     /// Record a retired version. Normal methods must also handle missing markers.
 //!     #[sausage_plugin::on_retired]
 //!     fn retired(&self, p: RetiredParams) -> Result<()> {
@@ -401,6 +419,6 @@ pub use helpers::*;
 pub use imports::*;
 #[doc(hidden)]
 pub use inventory;
-pub use sausage_plugin_macros::{import, on_retired, on_stop, rpc, schema, test};
+pub use sausage_plugin_macros::{import, on_retired, on_stop, on_superseded, rpc, schema, test};
 pub use transport::*;
 pub use {schemars, serde, serde_json};

sdk/src/transport.rs

@@ -96,6 +96,7 @@ struct State {
     active: bool,
     endpoints: HashMap<std::any::TypeId, String>,
     stopping: bool,
+    draining: bool,
     next: u64,
     pending: HashMap<u64, Pending>,
     #[cfg(target_os = "wasi")]
@@ -125,17 +126,48 @@ pub(crate) fn endpoint_name<T: 'static>() -> Result<String> {
 }
 #[cfg(target_os = "wasi")]
 pub(crate) fn endpoint_task(future: impl Future<Output = ()> + 'static) {
+    background(async move {
+        let _ = until_drained(future).await;
+    });
+}
+/// Cancel an accept/read wait when draining, but never an active handler.
+#[cfg(target_os = "wasi")]
+pub(crate) async fn until_drained<F: Future>(future: F) -> Option<F::Output> {
     let (handle, registration) = futures::future::AbortHandle::new_pair();
     let id = STATE.with(|s| {
         let mut s = s.borrow_mut();
+        if s.draining {
+            return None;
+        }
         s.next += 1;
         let id = s.next;
         s.listeners.insert(id, handle);
-        id
-    });
-    background(async move {
-        let _ = futures::future::Abortable::new(future, registration).await;
-        STATE.with(|s| s.borrow_mut().listeners.remove(&id));
+        Some(id)
+    })?;
+    let result = futures::future::Abortable::new(future, registration)
+        .await
+        .ok();
+    STATE.with(|s| s.borrow_mut().listeners.remove(&id));
+    result
+}
+/// Stop accepting SDK endpoint work, then exit when handlers and background tasks finish.
+///
+/// Calls and notifications still work. This does not run the EOF stop hook or cancel
+/// background tasks: services with their own listeners/tasks must arrange for those
+/// to finish too. Repeated calls are harmless. Core's upgrade grace remains the bound.
+pub fn drain() {
+    STATE.with(|s| {
+        let mut s = s.borrow_mut();
+        s.draining = true;
+        #[cfg(target_os = "wasi")]
+        {
+            for (_, listener) in s.listeners.drain() {
+                listener.abort();
+            }
+            if let Some(w) = s.wake.take() {
+                w.wake();
+            }
+        }
     });
 }
 #[cfg(target_os = "wasi")]

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

@@ -16,6 +16,12 @@ struct Retired {
     /// The previous hash.
     previous: String,
 }
+/// The replacement version.
+#[schema]
+struct Superseded {
+    /// The current hash.
+    current: String,
+}
 /// A large payload.
 #[schema]
 struct Echo {
@@ -111,6 +117,15 @@ impl Probe {
     fn retired(&self, p: Retired) -> Result<String> {
         Ok(p.previous)
     }
+    /// Drain on replacement while calls are still available.
+    #[sausage_plugin::on_superseded]
+    async fn superseded(&self, p: Superseded) -> Result<()> {
+        call("other.superseded", json!({"current":p.current})).await?;
+        sausage_plugin::drain();
+        sausage_plugin::drain();
+        Ctx::default().notify("draining", json!({}))?;
+        Ok(())
+    }
     /// Wait until EOF, then run the stop hook.
     async fn wait(&self) -> Result<String> {
         call("other.wait", json!({})).await?;
@@ -193,6 +208,11 @@ impl Control {
         ctx.notify("endpoint-progress", json!({"value":7}))?;
         Ok(7)
     }
+    /// Finish an in-flight call after draining begins.
+    async fn wait(&self) -> Result<String> {
+        call("other.inflight", json!({})).await?;
+        Ok("finished".into())
+    }
     /// End the endpoint process after the response.
     fn stop(&self, ctx: &Ctx) -> Result<()> {
         ctx.exit_after_response();

sdk/tests/ui/superseded.rs

@@ -0,0 +1,11 @@
+struct Supersedes;
+#[sausage_plugin::rpc]
+impl Supersedes {
+    /// First superseded hook.
+    #[sausage_plugin::on_superseded]
+    fn a(&self) -> sausage_plugin::Result<()> { Ok(()) }
+    /// Second superseded hook.
+    #[sausage_plugin::on_superseded]
+    fn b(&self) -> sausage_plugin::Result<()> { Ok(()) }
+}
+fn main() {}

sdk/tests/ui/superseded.stderr

@@ -0,0 +1,5 @@
+error: only one on_superseded method is allowed
+ --> tests/ui/superseded.rs:7:5
+  |
+7 |     /// Second superseded hook.
+  |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^

sdk/tests/wasi.rs

@@ -201,6 +201,88 @@ fn protocol() {
     assert_eq!(outgoing["method"], "other.background");
     p.answer(&outgoing, json!({}));
     assert!(p.child.wait().unwrap().success());
+    // Replacement is a notification, not EOF. Drain keeps both endpoint handlers
+    // and background calls alive, while closing idle connections and the listener.
+    assert!(
+        !d["result"]["methods"]
+            .as_array()
+            .unwrap()
+            .iter()
+            .any(|m| m["name"] == "rpc.superseded" || m["name"] == "superseded")
+    );
+    std::fs::remove_file(dir.join("stopped")).unwrap();
+    let mut provider = start(&wasm, &dir);
+    provider.request("serve", json!({}));
+    assert_eq!(provider.read()["result"], Value::Null);
+    provider.request("start", json!({}));
+    let first = provider.read();
+    let second = provider.read();
+    let background = if first["method"] == "other.background" {
+        first
+    } else {
+        second
+    };
+    assert_eq!(background["method"], "other.background");
+    let address = std::fs::read_to_string(dir.join("endpoint"))
+        .unwrap()
+        .split_whitespace()
+        .next()
+        .unwrap()
+        .trim_start_matches("tcp:")
+        .to_owned();
+    let mut idle = std::net::TcpStream::connect(&address).unwrap();
+    idle.set_read_timeout(Some(std::time::Duration::from_secs(4)))
+        .unwrap();
+    writeln!(
+        idle,
+        "{}",
+        json!({"jsonrpc":"2.0","id":0,"method":"rpc.endpoint","params":{}})
+    )
+    .unwrap();
+    let mut idle = BufReader::new(idle);
+    let mut line = String::new();
+    idle.read_line(&mut line).unwrap();
+    assert_eq!(
+        sausage_plugin::serde_json::from_str::<Value>(&line).unwrap()["id"],
+        0
+    );
+    let mut active = std::net::TcpStream::connect(&address).unwrap();
+    active
+        .set_read_timeout(Some(std::time::Duration::from_secs(4)))
+        .unwrap();
+    writeln!(
+        active,
+        "{}",
+        json!({"jsonrpc":"2.0","id":1,"method":"wait","params":{}})
+    )
+    .unwrap();
+    let inflight = provider.read();
+    assert_eq!(inflight["method"], "other.inflight");
+    provider
+        .send(json!({"jsonrpc":"2.0","method":"rpc.superseded","params":{"current":"new-hash"}}));
+    let hook = provider.read();
+    assert_eq!(hook["method"], "other.superseded");
+    assert_eq!(hook["params"]["current"], "new-hash");
+    provider.answer(&hook, json!({}));
+    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());
+    assert!(provider.child.try_wait().unwrap().is_none());
+    provider.answer(&inflight, json!({}));
+    line.clear();
+    BufReader::new(active).read_line(&mut line).unwrap();
+    assert_eq!(
+        sausage_plugin::serde_json::from_str::<Value>(&line).unwrap()["result"],
+        "finished"
+    );
+    assert!(provider.child.try_wait().unwrap().is_none());
+    provider.answer(&background, json!({}));
+    provider.finishes();
+    assert!(
+        !dir.join("stopped").exists(),
+        "drain is not the EOF stop hook"
+    );
     let provider_doc = &d["result"];
     let imports = target.join("imports");
     std::fs::create_dir_all(&imports).unwrap();