valkey-cli auto-exit from subscribed mode

src/valkey-cli.c

@@ -218,6 +218,8 @@ static struct config {
     int shutdown;
     int monitor_mode;
     int pubsub_mode;
+    int pubsub_unsharded_count; /* channels and patterns */
+    int pubsub_sharded_count;   /* shard channels */
     int blocking_state_aborted; /* used to abort monitor_mode and pubsub_mode. */
     int latency_mode;
     int latency_dist_mode;
@@ -2224,6 +2226,7 @@ static int cliReadReply(int output_raw_strings) {
         fflush(stdout);
         sdsfree(out);
     }
+
     return REDIS_OK;
 }
 
@@ -2402,6 +2405,25 @@ static int cliSendCommand(int argc, char **argv, long repeat) {
                             config.pubsub_mode = 1;
                             cliRefreshPrompt();
                         }
+
+                        /* Handle pubsub mode */
+                        if (config.last_reply->elements >= 3) {
+                            char *cmd = config.last_reply->element[0]->str;
+                            int count = config.last_reply->element[2]->integer;
+
+                            if (strcmp(cmd, "subscribe") == 0 || strcmp(cmd, "unsubscribe") == 0 ||
+                                strcmp(cmd, "psubscribe") == 0 || strcmp(cmd, "punsubscribe") == 0) {
+                                config.pubsub_unsharded_count = count;
+                            } else if (strcmp(cmd, "ssubscribe") == 0 || strcmp(cmd, "sunsubscribe") == 0) {
+                                config.pubsub_sharded_count = count;
+                            }
+
+                            if (config.pubsub_unsharded_count == 0 && config.pubsub_sharded_count == 0) {
+                                config.pubsub_mode = 0;
+                                cliRefreshPrompt();
+                            }
+                        }
+
                         if (--num_expected_pubsub_push > 0) {
                             continue; /* We need more of these. */
                         }
@@ -9493,6 +9515,8 @@ int main(int argc, char **argv) {
     config.shutdown = 0;
     config.monitor_mode = 0;
     config.pubsub_mode = 0;
+    config.pubsub_unsharded_count = 0;
+    config.pubsub_sharded_count = 0;
     config.blocking_state_aborted = 0;
     config.latency_mode = 0;
     config.latency_dist_mode = 0;

tests/integration/valkey-cli.tcl

@@ -608,6 +608,47 @@ if {!$::tls} { ;# fake_redis_node doesn't support TLS
         assert_equal "a\n1\nb\n2\nc\n3" [exec {*}$cmdline ZRANGE new_zset 0 -1 WITHSCORES]
     }
 
+    test {valkey-cli pubsub mode with multiple subscription types} {
+        set fd [open_cli]
+
+        # Subscribe to a regular channel
+        write_cli $fd "SUBSCRIBE channel1"
+        assert_match "*subscribe*channel1*" [read_cli $fd]
+
+        # Subscribe to a pattern
+        write_cli $fd "PSUBSCRIBE pattern*"
+        assert_match "*psubscribe*pattern**" [read_cli $fd]
+
+        # Subscribe to a shard channel
+        write_cli $fd "SSUBSCRIBE schannel1"
+        assert_match "*ssubscribe*schannel1*" [read_cli $fd]
+
+        # Unsubscribe from regular channel
+        write_cli $fd "UNSUBSCRIBE channel1"
+        assert_match "*unsubscribe*channel1*" [read_cli $fd]
+
+        # Verify still in pubsub mode
+        catch {run_command $fd "SET key value"} err
+        assert_match "*ERR*only*SUBSCRIBE*UNSUBSCRIBE*allowed*" $err
+
+        # Unsubscribe from pattern
+        write_cli $fd "PUNSUBSCRIBE pattern*"
+        assert_match "*punsubscribe*pattern**" [read_cli $fd]
+
+         # Verify still in pubsub mode
+        catch {run_command $fd "GET key"} err
+        assert_match "*ERR*only*SUBSCRIBE*UNSUBSCRIBE*allowed*" $err
+
+        # Unsubscribe from shard channel
+        write_cli $fd "SUNSUBSCRIBE schannel1"
+        assert_match "*sunsubscribe*schannel1*" [read_cli $fd]
+
+        # Verify that we've exited pubsub mode
+        assert_equal "PONG" [run_command $fd "PING"]
+
+        close_cli $fd
+    }
+
     test "Valid Connection Scheme: redis://" {
         set cmdline [valkeycliuri "redis://" [srv host] [srv port]]
         assert_equal {PONG} [exec {*}$cmdline PING]