From b57faab2aea4bea31e0b40fac9e4fe0553599ddd Mon Sep 17 00:00:00 2001
From: Rustin170506 <tech@rustin.me>
Date: Mon, 11 Nov 2024 22:48:23 +0800
Subject: [PATCH] feat: Improve error message when state streaming API is
 unimplemented

Signed-off-by: Rustin170506 <tech@rustin.me>
---
 tokio-console/src/conn.rs | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/tokio-console/src/conn.rs b/tokio-console/src/conn.rs
index c27893ffe..10a1417d4 100644
--- a/tokio-console/src/conn.rs
+++ b/tokio-console/src/conn.rs
@@ -122,7 +122,18 @@ impl Connection {
                 let update_stream =
                     Box::new(client.watch_updates(update_request).await?.into_inner());
                 let state_request = tonic::Request::new(StateRequest {});
-                let state_stream = Box::new(client.watch_state(state_request).await?.into_inner());
+                let state_stream = match client.watch_state(state_request).await {
+                    Ok(stream) => Box::new(stream.into_inner()),
+                    Err(e) => {
+                        if e.code() == tonic::Code::Unimplemented {
+                            tracing::error!(
+                                "The server at {} does not support state streaming. Please update the console-subscriber to v0.5.0 or later version.",
+                                self.target
+                            );
+                        }
+                        return Err(e.into());
+                    }
+                };
                 Ok::<State, Box<dyn Error + Send + Sync>>(State::Connected {
                     client,
                     update_stream,