Repository: numaproj/numaflow

Issue:
Performance benchmarking and publish report for ServingPipeline


Hints from issue comments:

Solution patch:
diff --git a/config/advanced-install/minimal-crds.yaml b/config/advanced-install/minimal-crds.yaml
index edf1df35..a1c2a60e 100644
--- a/config/advanced-install/minimal-crds.yaml
+++ b/config/advanced-install/minimal-crds.yaml
@@ -189,6 +189,55 @@ spec:
 ---
 apiVersion: apiextensions.k8s.io/v1
 kind: CustomResourceDefinition
+metadata:
+  name: servingpipelines.numaflow.numaproj.io
+spec:
+  group: numaflow.numaproj.io
+  names:
+    kind: ServingPipeline
+    listKind: ServingPipelineList
+    plural: servingpipelines
+    shortNames:
+    - spl
+    singular: servingpipeline
+  scope: Namespaced
+  versions:
+  - additionalPrinterColumns:
+    - jsonPath: .status.phase
+      name: Phase
+      type: string
+    - jsonPath: .metadata.creationTimestamp
+      name: Age
+      type: date
+    - jsonPath: .status.message
+      name: Message
+      type: string
+    name: v1alpha1
+    schema:
+      openAPIV3Schema:
+        properties:
+          apiVersion:
+            type: string
+          kind:
+            type: string
+          metadata:
+            type: object
+          spec:
+            type: object
+            x-kubernetes-preserve-unknown-fields: true
+          status:
+            type: object
+            x-kubernetes-preserve-unknown-fields: true
+        required:
+        - spec
+        type: object
+    served: true
+    storage: true
+    subresources:
+      status: {}
+---
+apiVersion: apiextensions.k8s.io/v1
+kind: CustomResourceDefinition
 metadata:
   name: vertices.numaflow.numaproj.io
 spec:
diff --git a/config/base/crds/minimal/kustomization.yaml b/config/base/crds/minimal/kustomization.yaml
index 6a183196..0c6ed431 100644
--- a/config/base/crds/minimal/kustomization.yaml
+++ b/config/base/crds/minimal/kustomization.yaml
@@ -6,3 +6,4 @@ resources:
   - numaflow.numaproj.io_pipelines.yaml
   - numaflow.numaproj.io_vertices.yaml
   - numaflow.numaproj.io_monovertices.yaml
+  - numaflow.numaproj.io_servingpipelines.yaml
diff --git a/config/extensions/webhook/numaflow-webhook-deployment.yaml b/config/extensions/webhook/numaflow-webhook-deployment.yaml
index 11fa9d08..9bb80f20 100644
--- a/config/extensions/webhook/numaflow-webhook-deployment.yaml
+++ b/config/extensions/webhook/numaflow-webhook-deployment.yaml
@@ -17,22 +17,22 @@ spec:
         app.kubernetes.io/component: numaflow-webhook
     spec:
       containers:
-      - name: webhook
-        image: quay.io/numaproj/numaflow:latest
-        imagePullPolicy: Always
-        args:
-        - webhook-service
-        env:
-        - name: NAMESPACE
-          valueFrom:
-            fieldRef:
-              fieldPath: metadata.namespace
-        - name: PORT
-          value: "443"              
-        - name: DEPLOYMENT_NAME
-          value: numaflow-webhook
-        - name: SERVICE_NAME
-          value: numaflow-webhook
-        - name: CLUSTER_ROLE_NAME
-          value: numaflow-webhook
+        - name: webhook
+          image: quay.io/numaproj/numaflow:latest
+          imagePullPolicy: Always
+          args:
+            - webhook-service
+          env:
+            - name: NAMESPACE
+              valueFrom:
+                fieldRef:
+                  fieldPath: metadata.namespace
+            - name: PORT
+              value: "443"
+            - name: DEPLOYMENT_NAME
+              value: numaflow-webhook
+            - name: SERVICE_NAME
+              value: numaflow-webhook
+            - name: CLUSTER_ROLE_NAME
+              value: numaflow-webhook
       serviceAccountName: numaflow-webhook-sa
diff --git a/config/extensions/webhook/rbac/numaflow-webhook-cluster-role.yaml b/config/extensions/webhook/rbac/numaflow-webhook-cluster-role.yaml
index b744196a..a213d3ec 100644
--- a/config/extensions/webhook/rbac/numaflow-webhook-cluster-role.yaml
+++ b/config/extensions/webhook/rbac/numaflow-webhook-cluster-role.yaml
@@ -55,6 +55,7 @@ rules:
     resources:
       - interstepbufferservices
       - pipelines
+      - servingpipelines
   - apiGroups:
       - rbac.authorization.k8s.io
     resources:
diff --git a/config/validating-webhook-install.yaml b/config/validating-webhook-install.yaml
index d0186578..3fc6375a 100644
--- a/config/validating-webhook-install.yaml
+++ b/config/validating-webhook-install.yaml
@@ -57,6 +57,7 @@ rules:
   resources:
   - interstepbufferservices
   - pipelines
+  - servingpipelines
   verbs:
   - get
   - list
diff --git a/pkg/apis/numaflow/v1alpha1/const.go b/pkg/apis/numaflow/v1alpha1/const.go
index 6118d621..d3261d90 100644
--- a/pkg/apis/numaflow/v1alpha1/const.go
+++ b/pkg/apis/numaflow/v1alpha1/const.go
@@ -157,7 +157,9 @@ const (
 	EnvK8sServerVersion                 = "K8S_SERVER_VERSION"
 	EnvNumaflowRuntime                  = "NUMAFLOW_RUNTIME"
 	EnvServingSettings                  = "NUMAFLOW_SERVING_SOURCE_SETTINGS"
-	EnvServingStore                     = "NUMAFLOW_SERVING_KV_STORE"
+	EnvServingCallbackStore             = "NUMAFLOW_SERVING_CALLBACK_STORE"
+	EnvServingResponseStore             = "NUMAFLOW_SERVING_RESPONSE_STORE"
+	EnvServingStatusStore               = "NUMAFLOW_SERVING_STATUS_STORE"
 
 	NumaflowRustBinary          = "/bin/numaflow-rs"
 	PathVarRun                  = "/var/run/numaflow"
diff --git a/pkg/apis/numaflow/v1alpha1/serving_pipeline_types.go b/pkg/apis/numaflow/v1alpha1/serving_pipeline_types.go
index 1fd631d6..54d659dc 100644
--- a/pkg/apis/numaflow/v1alpha1/serving_pipeline_types.go
+++ b/pkg/apis/numaflow/v1alpha1/serving_pipeline_types.go
@@ -184,22 +184,16 @@ func (sp ServingPipeline) GetServingDeploymentObj(req GetServingPipelineResource
 		return nil, fmt.Errorf("failed to marshal serving source settings: %w", err)
 	}
 	encodedServingSourceSettings := base64.StdEncoding.EncodeToString(servingSourceSettings)
-
 	envVars := []corev1.EnvVar{
 		{Name: EnvNamespace, ValueFrom: &corev1.EnvVarSource{FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.namespace"}}},
 		{Name: EnvPod, ValueFrom: &corev1.EnvVarSource{FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.name"}}},
 		{Name: EnvServingMinPipelineSpec, Value: encodedPipelineSpec},
-		{Name: "NUMAFLOW_SERVING_SOURCE_SETTINGS", Value: encodedServingSourceSettings},
-		{Name: "NUMAFLOW_SERVING_KV_STORE", Value: fmt.Sprintf("%s_SERVING_KV_STORE", sp.GetServingStoreName())},
+		{Name: EnvServingSettings, Value: encodedServingSourceSettings},
+		{Name: EnvServingCallbackStore, Value: fmt.Sprintf("%s_SERVING_CALLBACK_STORE", sp.GetServingStoreName())},
+		{Name: EnvServingResponseStore, Value: fmt.Sprintf("%s_SERVING_RESPONSE_STORE", sp.GetServingStoreName())},
+		{Name: EnvServingStatusStore, Value: fmt.Sprintf("%s_SERVING_STATUS_STORE", sp.GetServingStoreName())},
 		{Name: EnvServingPort, Value: strconv.Itoa(ServingServicePort)},
-		{ // TODO: do we still need it?
-			Name: EnvServingHostIP,
-			ValueFrom: &corev1.EnvVarSource{
-				FieldRef: &corev1.ObjectFieldSelector{
-					FieldPath: "status.podIP",
-				},
-			},
-		},
+		{Name: EnvReplica, ValueFrom: &corev1.EnvVarSource{FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.annotations['" + KeyReplica + "']"}}},
 	}
 	envVars = append(envVars, req.Env...)
 
@@ -325,7 +319,9 @@ func (sp ServingPipeline) GetPipelineObj(req GetServingPipelineResourceReq) Pipe
 			plSpec.Vertices[i].ContainerTemplate.Env,
 			corev1.EnvVar{Name: EnvCallbackEnabled, Value: "true"},
 			corev1.EnvVar{Name: EnvServingSettings, Value: encodedServingSourceSettings},
-			corev1.EnvVar{Name: EnvServingStore, Value: fmt.Sprintf("%s_SERVING_KV_STORE", sp.GetServingStoreName())},
+			corev1.EnvVar{Name: EnvServingCallbackStore, Value: fmt.Sprintf("%s_SERVING_CALLBACK_STORE", sp.GetServingStoreName())},
+			corev1.EnvVar{Name: EnvServingResponseStore, Value: fmt.Sprintf("%s_SERVING_RESPONSE_STORE", sp.GetServingStoreName())},
+			corev1.EnvVar{Name: EnvServingStatusStore, Value: fmt.Sprintf("%s_SERVING_STATUS_STORE", sp.GetServingStoreName())},
 			corev1.EnvVar{Name: EnvNumaflowRuntime, Value: "rust"},
 		)
 		if plSpec.Vertices[i].Scale.Min == nil {
diff --git a/pkg/apis/proto/reduce/wal/aligned/segment.proto b/pkg/apis/proto/reduce/wal/aligned/segment.proto
deleted file mode 100644
index 9dc4c27f..00000000
--- a/pkg/apis/proto/reduce/wal/aligned/segment.proto
+++ /dev/null
@@ -1,14 +0,0 @@
-syntax = "proto3";
-import "google/protobuf/timestamp.proto";
-
-package aligned_wal;
-
-// SegmentHeader for the Aligned WAL.
-message SegmentHeader {
-  // start_time of the window
-  google.protobuf.Timestamp start_time = 1;
-  // end_time = start_time + window
-  google.protobuf.Timestamp end_time = 2;
-  // this is currently set to 0 (in the future we can split large windows to multiple slots)
-  string slot = 3;
-}
diff --git a/pkg/apis/proto/reduce/wal/wal.proto b/pkg/apis/proto/reduce/wal/wal.proto
new file mode 100644
index 00000000..7b930f8e
--- /dev/null
+++ b/pkg/apis/proto/reduce/wal/wal.proto
@@ -0,0 +1,20 @@
+syntax = "proto3";
+import "google/protobuf/timestamp.proto";
+
+// stores all the proto definitions that are required for the WAL
+package wal;
+
+// GCEvent is the event that is persisted in the WAL when a window is garbage collected
+message GCEvent {
+  // start_time of the window
+  google.protobuf.Timestamp start_time = 1;
+  // end time of the window
+  google.protobuf.Timestamp end_time = 2;
+  // keys of the window, it will be empty for aligned windows
+  repeated string keys = 3;
+}
+
+message Footer {
+  // the latest event time in the current Segment
+  google.protobuf.Timestamp latest_event_time = 1;
+}
\ No newline at end of file
diff --git a/pkg/isbsvc/jetstream_service.go b/pkg/isbsvc/jetstream_service.go
index 9da7d22e..a9fa10af 100644
--- a/pkg/isbsvc/jetstream_service.go
+++ b/pkg/isbsvc/jetstream_service.go
@@ -89,23 +89,59 @@ func (jss *jetStreamSvc) CreateBuffersAndBuckets(ctx context.Context, buffers, b
 	}
 
 	if servingSourceStore != "" {
-		kvName := JetStreamServingSourceStoreKVName(servingSourceStore)
-		if _, err := jss.js.KeyValue(kvName); err != nil {
+		cbStoreName := JetStreamServingCallbackStoreName(servingSourceStore)
+		if _, err := jss.js.KeyValue(cbStoreName); err != nil {
 			if !errors.Is(err, nats.ErrBucketNotFound) && !errors.Is(err, nats.ErrStreamNotFound) {
-				return fmt.Errorf("failed to query information of KV %q, %w", kvName, err)
+				return fmt.Errorf("failed to query information of KV %q, %w", cbStoreName, err)
 			}
 			if _, err := jss.js.CreateKeyValue(&nats.KeyValueConfig{
-				Bucket:       kvName,
+				Bucket:       cbStoreName,
+				MaxValueSize: 0,
+				History:      1,
+				TTL:          30 * time.Minute,
+				MaxBytes:     0,
+				Storage:      nats.FileStorage,
+				Replicas:     v.GetInt("stream.replicas"),
+			}); err != nil {
+				return fmt.Errorf("failed to create serving source KV %q, %w", cbStoreName, err)
+			}
+			log.Infow("Succeeded to create a serving source Callback KV", zap.String("kvName", cbStoreName))
+		}
+		rsStoreName := JetStreamServingResponseStoreName(servingSourceStore)
+		if _, err := jss.js.KeyValue(rsStoreName); err != nil {
+			if !errors.Is(err, nats.ErrBucketNotFound) && !errors.Is(err, nats.ErrStreamNotFound) {
+				return fmt.Errorf("failed to query information of KV %q, %w", rsStoreName, err)
+			}
+			if _, err := jss.js.CreateKeyValue(&nats.KeyValueConfig{
+				Bucket:       rsStoreName,
 				MaxValueSize: 0,
-				History:      64,
-				TTL:          time.Hour * 24 * 1, // 1 day
+				History:      1,
+				TTL:          2 * time.Hour,
 				MaxBytes:     0,
 				Storage:      nats.FileStorage,
 				Replicas:     v.GetInt("stream.replicas"),
 			}); err != nil {
-				return fmt.Errorf("failed to create serving source KV %q, %w", kvName, err)
+				return fmt.Errorf("failed to create serving source KV %q, %w", rsStoreName, err)
 			}
-			log.Infow("Succeeded to create a serving source KV", zap.String("kvName", kvName))
+			log.Infow("Succeeded to create a serving source Response KV", zap.String("kvName", rsStoreName))
+		}
+		statusStoreName := JetStreamServingStatusStoreName(servingSourceStore)
+		if _, err := jss.js.KeyValue(statusStoreName); err != nil {
+			if !errors.Is(err, nats.ErrBucketNotFound) && !errors.Is(err, nats.ErrStreamNotFound) {
+				return fmt.Errorf("failed to query information of KV %q, %w", statusStoreName, err)
+			}
+			if _, err := jss.js.CreateKeyValue(&nats.KeyValueConfig{
+				Bucket:       statusStoreName,
+				MaxValueSize: 0,
+				History:      1,
+				TTL:          2 * time.Hour,
+				MaxBytes:     0,
+				Storage:      nats.FileStorage,
+				Replicas:     v.GetInt("stream.replicas"),
+			}); err != nil {
+				return fmt.Errorf("failed to create serving source KV %q, %w", statusStoreName, err)
+			}
+			log.Infow("Succeeded to create a serving source Status KV", zap.String("kvName", statusStoreName))
 		}
 	}
 
@@ -246,11 +282,20 @@ func (jss *jetStreamSvc) DeleteBuffersAndBuckets(ctx context.Context, buffers, b
 	}
 
 	if servingSourceStore != "" {
-		servingSourceStoreKVName := JetStreamServingSourceStoreKVName(servingSourceStore)
-		if err := jss.js.DeleteKeyValue(servingSourceStoreKVName); err != nil && !errors.Is(err, nats.ErrBucketNotFound) && !errors.Is(err, nats.ErrStreamNotFound) {
-			return fmt.Errorf("failed to serving source store %q, %w", servingSourceStoreKVName, err)
+		servingSourceCallbackStoreKVName := JetStreamServingCallbackStoreName(servingSourceStore)
+		if err := jss.js.DeleteKeyValue(servingSourceCallbackStoreKVName); err != nil && !errors.Is(err, nats.ErrBucketNotFound) && !errors.Is(err, nats.ErrStreamNotFound) {
+			return fmt.Errorf("failed to delete serving source callback store %q, %w", servingSourceCallbackStoreKVName, err)
+		}
+		log.Infow("Succeeded to delete a serving source store", zap.String("kvName", servingSourceCallbackStoreKVName))
+		servingSourceResponseStoreKVName := JetStreamServingResponseStoreName(servingSourceStore)
+		if err := jss.js.DeleteKeyValue(servingSourceResponseStoreKVName); err != nil && !errors.Is(err, nats.ErrBucketNotFound) && !errors.Is(err, nats.ErrStreamNotFound) {
+			return fmt.Errorf("failed to delete serving source response store %q, %w", servingSourceResponseStoreKVName, err)
 		}
-		log.Infow("Succeeded to delete a serving source store", zap.String("kvName", servingSourceStoreKVName))
+		servingSourceStatusStoreKVName := JetStreamServingStatusStoreName(servingSourceStore)
+		if err := jss.js.DeleteKeyValue(servingSourceStatusStoreKVName); err != nil && !errors.Is(err, nats.ErrBucketNotFound) && !errors.Is(err, nats.ErrStreamNotFound) {
+			return fmt.Errorf("failed to delete serving source status store %q, %w", servingSourceStatusStoreKVName, err)
+		}
+		log.Infow("Succeeded to delete a serving source response store", zap.String("kvName", servingSourceResponseStoreKVName))
 	}
 	return nil
 }
@@ -284,9 +329,17 @@ func (jss *jetStreamSvc) ValidateBuffersAndBuckets(ctx context.Context, buffers,
 		}
 	}
 	if servingSourceStore != "" {
-		servingSourceStoreKVName := JetStreamServingSourceStoreKVName(servingSourceStore)
-		if _, err := jss.js.KeyValue(servingSourceStoreKVName); err != nil {
-			return fmt.Errorf("failed to query serving source store KV %q, %w", servingSourceStoreKVName, err)
+		servingSourceCallbackStoreKVName := JetStreamServingCallbackStoreName(servingSourceStore)
+		if _, err := jss.js.KeyValue(servingSourceCallbackStoreKVName); err != nil {
+			return fmt.Errorf("failed to query serving source callback store KV %q, %w", servingSourceCallbackStoreKVName, err)
+		}
+		servingSourceResponseStoreKVName := JetStreamServingResponseStoreName(servingSourceStore)
+		if _, err := jss.js.KeyValue(servingSourceResponseStoreKVName); err != nil {
+			return fmt.Errorf("failed to query serving source response store KV %q, %w", servingSourceResponseStoreKVName, err)
+		}
+		servingSourceStatusStoreKVName := JetStreamServingStatusStoreName(servingSourceStore)
+		if _, err := jss.js.KeyValue(servingSourceStatusStoreKVName); err != nil {
+			return fmt.Errorf("failed to query serving source status store KV %q, %w", servingSourceStatusStoreKVName, err)
 		}
 	}
 	return nil
@@ -343,6 +396,14 @@ func JetStreamSideInputsStoreKVName(sideInputStoreName string) string {
 	return fmt.Sprintf("%s_SIDE_INPUTS", sideInputStoreName)
 }
 
-func JetStreamServingSourceStoreKVName(servingSourceStoreName string) string {
-	return fmt.Sprintf("%s_SERVING_KV_STORE", servingSourceStoreName)
+func JetStreamServingCallbackStoreName(servingSourceStoreName string) string {
+	return fmt.Sprintf("%s_SERVING_CALLBACK_STORE", servingSourceStoreName)
+}
+
+func JetStreamServingResponseStoreName(servingSourceStoreName string) string {
+	return fmt.Sprintf("%s_SERVING_RESPONSE_STORE", servingSourceStoreName)
+}
+
+func JetStreamServingStatusStoreName(servingSourceStoreName string) string {
+	return fmt.Sprintf("%s_SERVING_STATUS_STORE", servingSourceStoreName)
 }
diff --git a/rust/Cargo.lock b/rust/Cargo.lock
index 633a1a08..02bdfc37 100644
--- a/rust/Cargo.lock
+++ b/rust/Cargo.lock
@@ -78,9 +78,9 @@ dependencies = [
 
 [[package]]
 name = "async-nats"
-version = "0.39.0"
+version = "0.40.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a798aab0c0203b31d67d501e5ed1f3ac6c36a329899ce47fc93c3bea53f3ae89"
+checksum = "e23419d455dc57d3ae60a2f4278cf561fc74fe866e548e14d2b0ad3e1b8ca0b2"
 dependencies = [
  "base64 0.22.1",
  "bytes",
@@ -3503,6 +3503,7 @@ dependencies = [
  "thiserror 2.0.12",
  "tokio",
  "tokio-stream",
+ "tokio-util",
  "tonic",
  "tower",
  "tower-http",
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index b17a9499..24a8362c 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -66,7 +66,7 @@ numaflow-sqs = { path = "extns/numaflow-sqs" }
 tokio = { version = "1.44.2", features = ["macros", "rt-multi-thread"] }
 bytes = "1.9.0"
 tracing = "0.1.41"
-async-nats = "0.39.0"
+async-nats = "0.40.0"
 thiserror = "2.0.3"
 axum = { version = "0.8.3", default-features = false, features = [
     "http1",
diff --git a/rust/README.md b/rust/README.md
index c9d0b685..0552d428 100644
--- a/rust/README.md
+++ b/rust/README.md
@@ -1,13 +1,4 @@
-# Numaserve
+# Numaflow Data Plane
 
-Numaserve is the serving/HTTP endpoint for [Numaflow](https://numaflow.numaproj.io/) which is a distributed scalable general-purpose
-async processing platform.
-
-## Problem
-
- [Numaflow](https://numaflow.numaproj.io/) being a DCG (Directed Compute Graph), there is no direct way to interact via request/response 
- protocol (in other words, the Source of Numaflow is decoupled with the Sink). There are a couple of use-cases where request/response semantics
- are required. E.g.,
- 
- * Response with the status of completion (ability to say whether the processing is complete)
- * Response with output of a complex computation over the DCG directly to the callee (say, something like [kserve](https://kserve.github.io/website/latest/)).
+Numaflow Data Plane is being rewritten in Rust for better performance, memory safety, and more importantly the
+maintainability.
diff --git a/rust/numaflow-core/src/config.rs b/rust/numaflow-core/src/config.rs
index 8594343c..a5520b1e 100644
--- a/rust/numaflow-core/src/config.rs
+++ b/rust/numaflow-core/src/config.rs
@@ -13,7 +13,8 @@ const ENV_VERTEX_OBJ: &str = "NUMAFLOW_VERTEX_OBJECT";
 const ENV_CALLBACK_ENABLED: &str = "NUMAFLOW_CALLBACK_ENABLED";
 const ENV_CALLBACK_CONCURRENCY: &str = "NUMAFLOW_CALLBACK_CONCURRENCY";
 const ENV_NUMAFLOW_SERVING_SOURCE_SETTINGS: &str = "NUMAFLOW_SERVING_SOURCE_SETTINGS";
-const ENV_NUMAFLOW_SERVING_KV_STORE: &str = "NUMAFLOW_SERVING_KV_STORE";
+const ENV_NUMAFLOW_SERVING_CALLBACK_STORE: &str = "NUMAFLOW_SERVING_CALLBACK_STORE";
+const ENV_NUMAFLOW_SERVING_RESPONSE_STORE: &str = "NUMAFLOW_SERVING_RESPONSE_STORE";
 const DEFAULT_CALLBACK_CONCURRENCY: usize = 100;
 
 /// Building blocks (Source, Sink, Transformer, FallBack, Metrics, etc.) to build a Pipeline or a
diff --git a/rust/numaflow-core/src/config/monovertex.rs b/rust/numaflow-core/src/config/monovertex.rs
index aa66fc11..17ef138b 100644
--- a/rust/numaflow-core/src/config/monovertex.rs
+++ b/rust/numaflow-core/src/config/monovertex.rs
@@ -162,7 +162,7 @@ impl MonovertexConfig {
                 })?;
             callback_config = Some(ServingCallbackConfig {
                 callback_store: Box::leak(Box::new(format!(
-                    "{}-{}_SERVING_KV_STORE",
+                    "{}-{}_SERVING_CALLBACK_STORE",
                     get_namespace(),
                     get_pipeline_name(),
                 ))),
diff --git a/rust/numaflow-core/src/config/pipeline.rs b/rust/numaflow-core/src/config/pipeline.rs
index ff2da552..002c6395 100644
--- a/rust/numaflow-core/src/config/pipeline.rs
+++ b/rust/numaflow-core/src/config/pipeline.rs
@@ -9,7 +9,10 @@ use serde::Deserialize;
 use serde_json::from_slice;
 use tracing::info;
 
-use super::{DEFAULT_CALLBACK_CONCURRENCY, ENV_CALLBACK_CONCURRENCY, ENV_CALLBACK_ENABLED};
+use super::{
+    DEFAULT_CALLBACK_CONCURRENCY, ENV_CALLBACK_CONCURRENCY, ENV_CALLBACK_ENABLED,
+    ENV_NUMAFLOW_SERVING_RESPONSE_STORE,
+};
 use crate::config::components::metrics::MetricsConfig;
 use crate::config::components::sink::SinkConfig;
 use crate::config::components::sink::SinkType;
@@ -22,7 +25,7 @@ use crate::config::pipeline::map::MapVtxConfig;
 use crate::config::pipeline::watermark::WatermarkConfig;
 use crate::config::pipeline::watermark::{BucketConfig, EdgeWatermarkConfig};
 use crate::config::pipeline::watermark::{IdleConfig, SourceWatermarkConfig};
-use crate::config::ENV_NUMAFLOW_SERVING_KV_STORE;
+use crate::config::ENV_NUMAFLOW_SERVING_CALLBACK_STORE;
 use crate::config::ENV_NUMAFLOW_SERVING_SOURCE_SETTINGS;
 use crate::error::Error;
 use crate::Result;
@@ -205,7 +208,8 @@ pub(crate) struct UserDefinedStoreConfig {
 
 #[derive(Debug, Deserialize, Clone, PartialEq)]
 pub(crate) struct NatsStoreConfig {
-    pub(crate) name: String,
+    /// Response (or result) Store Name.
+    pub(crate) rs_store_name: String,
 }
 
 impl Default for UserDefinedStoreConfig {
@@ -260,7 +264,8 @@ impl PipelineConfig {
                     ENV_CALLBACK_ENABLED,
                     ENV_CALLBACK_CONCURRENCY,
                     ENV_NUMAFLOW_SERVING_SOURCE_SETTINGS,
-                    ENV_NUMAFLOW_SERVING_KV_STORE,
+                    ENV_NUMAFLOW_SERVING_CALLBACK_STORE,
+                    ENV_NUMAFLOW_SERVING_RESPONSE_STORE,
                 ]
                 .contains(&key.as_str())
             })
@@ -356,7 +361,7 @@ impl PipelineConfig {
                         )
                     })?;
                 let serving_spec: numaflow_models::models::ServingSpec =
-                    serde_json::from_slice(serving_settings_decoded.as_slice()).map_err(|e| {
+                    from_slice(serving_settings_decoded.as_slice()).map_err(|e| {
                         Error::Config(
                             format!(
                                 "Failed to base64 decode value of environment variable '{ENV_NUMAFLOW_SERVING_SOURCE_SETTINGS}'. value='{serving_settings}'. Err={e:?}"
@@ -365,8 +370,10 @@ impl PipelineConfig {
                     })?;
 
                 if serving_spec.store.is_none() {
-                    let kv_store = get_var(ENV_NUMAFLOW_SERVING_KV_STORE)?;
-                    Some(ServingStoreType::Nats(NatsStoreConfig { name: kv_store }))
+                    let rs_kv_store = get_var(ENV_NUMAFLOW_SERVING_RESPONSE_STORE)?;
+                    Some(ServingStoreType::Nats(NatsStoreConfig {
+                        rs_store_name: rs_kv_store,
+                    }))
                 } else {
                     Some(ServingStoreType::UserDefined(
                         UserDefinedStoreConfig::default(),
@@ -510,8 +517,8 @@ impl PipelineConfig {
                     ))
                 })?;
 
-            let kv_store = env::var("NUMAFLOW_SERVING_KV_STORE").map_err(|_| {
-                    Error::Config("Serving store is default, but environment variable NUMAFLOW_SERVING_KV_STORE is not set".into())
+            let kv_store = env::var(ENV_NUMAFLOW_SERVING_CALLBACK_STORE).map_err(|_| {
+                    Error::Config("Serving store is default, but environment variable NUMAFLOW_SERVING_CALLBACK_STORE is not set".into())
                 })?;
             callback_config = Some(ServingCallbackConfig {
                 callback_store: Box::leak(kv_store.into_boxed_str()),
@@ -714,7 +721,8 @@ mod tests {
         let env_vars = [
             ("NUMAFLOW_ISBSVC_JETSTREAM_URL", "localhost:4222"),
             ("NUMAFLOW_SERVING_SOURCE_SETTINGS", "eyJhdXRoIjpudWxsLCJzZXJ2aWNlIjp0cnVlLCJtc2dJREhlYWRlcktleSI6IlgtTnVtYWZsb3ctSWQifQ=="),
-            ("NUMAFLOW_SERVING_KV_STORE", "test-kv-store"),
+            ("NUMAFLOW_SERVING_CALLBACK_STORE", "test-kv-store"),
+            ("NUMAFLOW_SERVING_RESPONSE_STORE", "test-kv-store"),
         ];
         let pipeline_config = PipelineConfig::load(pipeline_cfg_base64, env_vars).unwrap();
 
@@ -746,7 +754,7 @@ mod tests {
                 },
                 fb_sink_config: None,
                 serving_store_config: Some(ServingStoreType::Nats(NatsStoreConfig {
-                    name: "test-kv-store".into(),
+                    rs_store_name: "test-kv-store".into(),
                 })),
             }),
             metrics_config: MetricsConfig {
diff --git a/rust/numaflow-core/src/lib.rs b/rust/numaflow-core/src/lib.rs
index 5b55acfb..bea27f57 100644
--- a/rust/numaflow-core/src/lib.rs
+++ b/rust/numaflow-core/src/lib.rs
@@ -1,5 +1,6 @@
 use crate::config::{config, CustomResourceType};
 use bytes::Bytes;
+use numaflow_monitor::runtime;
 use tokio::signal;
 use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
@@ -66,7 +67,11 @@ mod serving_store;
 /// [Watermark]: https://numaflow.numaproj.io/core-concepts/watermarks/
 mod watermark;
 
-use numaflow_monitor::runtime;
+/// [Reduce] is a function which "collects" a group of items and then perform some "reduction" operation
+/// on all of them, thus reducing them to a single value.
+///
+/// [Reduce]:https://numaflow.numaproj.io/user-guide/user-defined-functions/reduce/reduce/
+mod reduce;
 
 pub async fn run() -> Result<()> {
     let cln_token = CancellationToken::new();
diff --git a/rust/numaflow-core/src/message.rs b/rust/numaflow-core/src/message.rs
index 6a827619..9ff29264 100644
--- a/rust/numaflow-core/src/message.rs
+++ b/rust/numaflow-core/src/message.rs
@@ -242,6 +242,15 @@ impl fmt::Display for MessageID {
     }
 }
 
+impl TryFrom<Message> for Bytes {
+    type Error = Error;
+
+    fn try_from(value: Message) -> Result<Self, Self::Error> {
+        let b: BytesMut = value.try_into()?;
+        Ok(b.freeze())
+    }
+}
+
 impl TryFrom<Message> for BytesMut {
     type Error = Error;
 
diff --git a/rust/numaflow-core/src/monovertex.rs b/rust/numaflow-core/src/monovertex.rs
index bf5e6cc6..b6c22403 100644
--- a/rust/numaflow-core/src/monovertex.rs
+++ b/rust/numaflow-core/src/monovertex.rs
@@ -4,7 +4,7 @@ use tracing::info;
 use crate::config::is_mono_vertex;
 use crate::config::monovertex::MonovertexConfig;
 use crate::error::{self};
-use crate::metrics::{LagReader, PendingReaderTasks, PendingReaderTasks_};
+use crate::metrics::{LagReader, PendingReaderTasks_};
 use crate::shared::create_components;
 use crate::sink::SinkWriter;
 use crate::source::Source;
diff --git a/rust/numaflow-core/src/pipeline.rs b/rust/numaflow-core/src/pipeline.rs
index 04abe6cd..7ac561df 100644
--- a/rust/numaflow-core/src/pipeline.rs
+++ b/rust/numaflow-core/src/pipeline.rs
@@ -13,9 +13,7 @@ use crate::config::pipeline::isb::Stream;
 use crate::config::pipeline::map::MapVtxConfig;
 use crate::config::pipeline::watermark::WatermarkConfig;
 use crate::config::pipeline::{PipelineConfig, ServingStoreType, SinkVtxConfig, SourceVtxConfig};
-use crate::metrics::{
-    ComponentHealthChecks, LagReader, PendingReaderTasks, PendingReaderTasks_, PipelineComponents,
-};
+use crate::metrics::{ComponentHealthChecks, LagReader, PendingReaderTasks, PipelineComponents};
 use crate::pipeline::forwarder::source_forwarder;
 use crate::pipeline::isb::jetstream::reader::JetStreamReader;
 use crate::pipeline::isb::jetstream::writer::JetstreamWriter;
diff --git a/rust/numaflow-core/src/reduce.rs b/rust/numaflow-core/src/reduce.rs
new file mode 100644
index 00000000..25ff1369
--- /dev/null
+++ b/rust/numaflow-core/src/reduce.rs
@@ -0,0 +1,2 @@
+/// Write Head Log for storing reduce data until the reduction is complete.
+pub(crate) mod wal;
diff --git a/rust/numaflow-core/src/reduce/wal.rs b/rust/numaflow-core/src/reduce/wal.rs
new file mode 100644
index 00000000..7dd6b4c3
--- /dev/null
+++ b/rust/numaflow-core/src/reduce/wal.rs
@@ -0,0 +1,8 @@
+//! Write Ahead Log for both Aligned and Unaligned Reduce Operation. The WAL is to persist the data
+//! we read from the ISB and store it until the processing is complete.
+
+/// A WAL Segment.
+pub(crate) mod segment;
+
+/// All the errors WAL could face.
+pub(crate) mod error;
diff --git a/rust/numaflow-core/src/reduce/wal/error.rs b/rust/numaflow-core/src/reduce/wal/error.rs
new file mode 100644
index 00000000..2f9d6fe9
--- /dev/null
+++ b/rust/numaflow-core/src/reduce/wal/error.rs
@@ -0,0 +1,48 @@
+use std::fmt::{Display, Formatter};
+use std::io;
+
+/// The result of any operation on the WAL.
+pub(crate) type WalResult<T> = Result<T, Error>;
+
+#[derive(Debug)]
+pub(crate) enum Error {
+    /// Error while writing a particular WAL entry.
+    WriteEntry {
+        /// The unique message ID. This can be used for NACK if needed.
+        id: String,
+        /// Reason of the error.
+        err: String,
+    },
+    /// IO Errors.
+    Io(io::Error),
+    /// Other errors.
+    Other(String),
+}
+
+impl Display for Error {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Error::WriteEntry { id, err } => {
+                write!(f, "Error::WriteEntry {id} {err}")
+            }
+            Error::Io(err) => {
+                write!(f, "Error::Io {err}")
+            }
+            Error::Other(err) => {
+                write!(f, "Error::Other {err}")
+            }
+        }
+    }
+}
+
+impl From<io::Error> for Error {
+    fn from(value: io::Error) -> Self {
+        Error::Io(value)
+    }
+}
+
+impl From<String> for Error {
+    fn from(value: String) -> Self {
+        Error::Other(value)
+    }
+}
diff --git a/rust/numaflow-core/src/reduce/wal/segment.rs b/rust/numaflow-core/src/reduce/wal/segment.rs
new file mode 100644
index 00000000..6f8f587a
--- /dev/null
+++ b/rust/numaflow-core/src/reduce/wal/segment.rs
@@ -0,0 +1,428 @@
+//! There are three types of Segments that goes into the WAL.
+//!
+//! ### Data
+//! The Data WAL contains the raw data [crate::message::Message] that is streamed from the ISB.
+//!
+//! ### GC Events
+//! The GC Events contract is that any GC event that is written means all the messages < the window
+//! end time can be deleted from the segment WAL(if keys are present only the messages with same
+//! keys will be deleted).
+//! E.g., of GC Events,
+//! Fixed window of 60s duration GC Events - (60, 120), (120, 180), (180, 240)
+//! Sliding window of 60s length and slide 10s GC Events - (60, 70), (70, 80), (80, 90)
+//! Session window with 10s timeout GC Events - (60, 100, `[key1, key2]`), (120, 1000, `[key3]`)
+//!
+//! ### Compaction
+//! While compacting, a new data segment is created called Compaction Segments. These contain the
+//! data elements that cannot be deleted because the reduction is not complete yet.
+//! Subsequent compactions should also consider already compacted files and order is very important
+//! during compaction we should make sure the order of the messages does not change.
+
+use crate::shared::grpc::{prost_timestamp_from_utc, utc_from_timestamp};
+use chrono::{DateTime, Utc};
+use numaflow_pb::objects::wal::GcEvent;
+
+/// Reading from WAL for Replays.
+pub(crate) mod replay;
+
+/// Writing/Append to the WAL.
+pub(crate) mod append;
+
+/// A compactor to truncate processed WAL segments. It can be both Data WAL and GC WAL.
+pub(crate) mod compactor;
+
+/// WAL is made of three types, the Data, GC, and Compaction WAL.
+#[derive(Debug, Clone)]
+pub(crate) enum WalType {
+    /// Data WAL contains the raw data.
+    Data,
+    /// GC WAL contains the completed-processing events.
+    Gc,
+    /// Compaction WAL is a Data WAL, but it has the compacted data after purging based on GC Events.
+    Compact,
+}
+
+impl WalType {
+    /// Some WALs have footers and some does not. It is mostly for optimizations.
+    pub(crate) fn has_footer(&self) -> bool {
+        // TODO: Set footer to true for Segment and Compaction for optimizations if needed.
+        match self {
+            WalType::Data => false,
+            WalType::Gc => false,
+            WalType::Compact => false,
+        }
+    }
+
+    /// Prefix of the WAL Segment as stored in the disk.
+    pub(crate) fn segment_prefix(&self) -> &'static str {
+        match self {
+            WalType::Data => "data",
+            WalType::Gc => "gc",
+            WalType::Compact => "compaction",
+        }
+    }
+
+    /// Suffix of the WAL Segment as stored in the disk. Not all WAL Segments have prefix,
+    /// it is used to filter our work-in-progress WAL Segments that are derived from other WALs.
+    pub(crate) fn segment_suffix(&self) -> &'static str {
+        match self {
+            WalType::Data => "",
+            WalType::Gc => "",
+            WalType::Compact => "",
+        }
+    }
+}
+
+impl std::fmt::Display for WalType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            WalType::Data => write!(f, "Data"),
+            WalType::Gc => write!(f, "GC"),
+            WalType::Compact => write!(f, "Compact"),
+        }
+    }
+}
+
+/// An entry in the GC WAL about the GC action.
+#[derive(Debug)]
+pub(crate) struct GcEventEntry {
+    /// Start time of the Window
+    start_time: DateTime<Utc>,
+    /// The end time of the Window. Anything after this timestamp has not been processed.
+    end_time: DateTime<Utc>,
+    /// Keys are used only for Unaligned window since sessions are defined not purely on [crate::watermark]
+    /// but also on Keys.
+    keys: Option<Vec<String>>,
+}
+
+impl From<GcEvent> for GcEventEntry {
+    fn from(value: GcEvent) -> Self {
+        Self {
+            start_time: utc_from_timestamp(value.start_time),
+            end_time: utc_from_timestamp(value.end_time),
+            keys: value.keys.into(),
+        }
+    }
+}
+
+impl From<GcEventEntry> for GcEvent {
+    fn from(value: GcEventEntry) -> Self {
+        Self {
+            start_time: Some(prost_timestamp_from_utc(value.start_time)),
+            end_time: Some(prost_timestamp_from_utc(value.end_time)),
+            keys: value.keys.unwrap_or_default(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::message::{Message, MessageID};
+    use crate::reduce::wal::segment::append::{AppendOnlyWal, SegmentWriteMessage};
+    use crate::reduce::wal::segment::compactor::{Compactor, WindowKind};
+    use crate::reduce::wal::segment::replay::{ReplayWal, SegmentEntry};
+    use crate::reduce::wal::segment::WalType;
+    use bytes::Bytes;
+    use chrono::{TimeZone, Utc};
+    use std::sync::Arc;
+    use tokio_stream::wrappers::ReceiverStream;
+    use tokio_stream::StreamExt;
+
+    #[tokio::test]
+    async fn test_gc_wal_and_aligned_compaction() {
+        let test_path = tempfile::tempdir().unwrap().into_path();
+
+        // Create GC WAL
+        let gc_wal = AppendOnlyWal::new(
+            WalType::Gc,
+            test_path.clone(),
+            1,    // 1MB
+            1000, // 1s flush interval
+            500,  // channel buffer
+        )
+        .await
+        .unwrap();
+
+        // Create and write GC events
+        let gc_start = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 5).unwrap();
+        let gc_end = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 10).unwrap();
+
+        let gc_event = GcEvent {
+            start_time: Some(prost_timestamp_from_utc(gc_start)),
+            end_time: Some(prost_timestamp_from_utc(gc_end)),
+            keys: vec![],
+        };
+
+        let (tx, rx) = tokio::sync::mpsc::channel(10);
+        let (_offset_stream, handle) = gc_wal
+            .streaming_write(ReceiverStream::new(rx))
+            .await
+            .unwrap();
+
+        tx.send(SegmentWriteMessage::WriteData {
+            id: Some("gc".to_string()),
+            data: bytes::Bytes::from(prost::Message::encode_to_vec(&gc_event)),
+        })
+        .await
+        .unwrap();
+
+        drop(tx);
+        handle.await.unwrap().unwrap();
+
+        // Create segment WAL
+        let segment_wal = AppendOnlyWal::new(
+            WalType::Data,
+            test_path.clone(),
+            1,    // 20MB
+            1000, // 1s flush interval
+            500,  // channel buffer
+        )
+        .await
+        .unwrap();
+
+        // Write 100 segment entries
+        let (tx, rx) = tokio::sync::mpsc::channel(100);
+        let (_offset_stream, handle) = segment_wal
+            .streaming_write(ReceiverStream::new(rx))
+            .await
+            .unwrap();
+
+        let start_time = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 0).unwrap();
+        let time_increment = chrono::Duration::seconds(1);
+
+        for i in 1..=100 {
+            let mut message = Message::default();
+            message.event_time = start_time + (time_increment * i);
+            message.keys = Arc::from(vec!["test-key".to_string()]);
+            message.value = bytes::Bytes::from(vec![1, 2, 3]);
+            message.id = MessageID {
+                vertex_name: "test-vertex".to_string().into(),
+                offset: i.to_string().into(),
+                index: 0,
+            };
+
+            let proto_message: Bytes = message.try_into().unwrap();
+            tx.send(SegmentWriteMessage::WriteData {
+                id: Some(format!("msg-{}", i)),
+                data: proto_message,
+            })
+            .await
+            .unwrap();
+        }
+
+        drop(tx);
+        handle.await.unwrap().unwrap();
+
+        // Create and run compactor
+        let compactor = Compactor::new(
+            test_path.clone(),
+            WindowKind::Aligned,
+            1,    // 20MB
+            1000, // 1s flush interval
+            500,  // channel buffer
+        )
+        .await
+        .unwrap();
+
+        compactor.compact().await.unwrap();
+
+        // Verify compacted data
+        let compaction_wal = ReplayWal::new(WalType::Compact, test_path);
+        let (mut rx, handle) = compaction_wal.streaming_read().unwrap();
+
+        let mut remaining_message_count = 0;
+        while let Some(entry) = rx.next().await {
+            if let SegmentEntry::DataEntry { data, .. } = entry {
+                let msg: numaflow_pb::objects::isb::Message = prost::Message::decode(data).unwrap();
+                if let Some(header) = msg.header {
+                    if let Some(message_info) = header.message_info {
+                        let event_time = utc_from_timestamp(message_info.event_time);
+                        assert!(
+                            event_time > gc_end,
+                            "Found message with event_time <= gc_end"
+                        );
+                    }
+                }
+                remaining_message_count += 1
+            }
+        }
+        handle.await.unwrap().unwrap();
+
+        // we send 100 messages out of which 10 messages have event times less than the window end
+        // time so the remaining message cound should be 90
+        assert_eq!(
+            remaining_message_count, 90,
+            "Expected 90 messages to remain after compaction"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_gc_wal_and_unaligned_compaction() {
+        let test_path = tempfile::tempdir().unwrap().into_path();
+
+        // Create GC WAL
+        let gc_wal = AppendOnlyWal::new(
+            WalType::Gc,
+            test_path.clone(),
+            1,    // 1MB
+            1000, // 1s flush interval
+            500,  // channel buffer
+        )
+        .await
+        .unwrap();
+
+        // Create and write GC events with different keys
+        let gc_start = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 0).unwrap();
+        let gc_end_key1 = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 10).unwrap();
+        let gc_end_key2 = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 15).unwrap();
+
+        let gc_event1 = GcEvent {
+            start_time: Some(prost_timestamp_from_utc(gc_start)),
+            end_time: Some(prost_timestamp_from_utc(gc_end_key1)),
+            keys: vec!["key1".to_string()],
+        };
+
+        let gc_event2 = GcEvent {
+            start_time: Some(prost_timestamp_from_utc(gc_start)),
+            end_time: Some(prost_timestamp_from_utc(gc_end_key2)),
+            keys: vec!["key2".to_string()],
+        };
+
+        let (tx, rx) = tokio::sync::mpsc::channel(10);
+        let (_offset_stream, handle) = gc_wal
+            .streaming_write(ReceiverStream::new(rx))
+            .await
+            .unwrap();
+
+        // Send both GC events
+        tx.send(SegmentWriteMessage::WriteData {
+            id: Some("gc1".to_string()),
+            data: bytes::Bytes::from(prost::Message::encode_to_vec(&gc_event1)),
+        })
+        .await
+        .unwrap();
+
+        tx.send(SegmentWriteMessage::WriteData {
+            id: Some("gc2".to_string()),
+            data: bytes::Bytes::from(prost::Message::encode_to_vec(&gc_event2)),
+        })
+        .await
+        .unwrap();
+
+        drop(tx);
+        handle.await.unwrap().unwrap();
+
+        // Create segment WAL
+        let segment_wal = AppendOnlyWal::new(
+            WalType::Data,
+            test_path.clone(),
+            1,    // 20MB
+            1000, // 1s flush interval
+            500,  // channel buffer
+        )
+        .await
+        .unwrap();
+
+        // Write segment entries with different keys
+        let (tx, rx) = tokio::sync::mpsc::channel(100);
+        let (_offset_stream, handle) = segment_wal
+            .streaming_write(ReceiverStream::new(rx))
+            .await
+            .unwrap();
+
+        let start_time = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 0).unwrap();
+        let time_increment = chrono::Duration::seconds(1);
+
+        // Create 50 messages with key1 and 50 messages with key2
+        for i in 1..=100 {
+            let mut message = Message::default();
+            message.event_time = start_time + (time_increment * i);
+
+            // Alternate between key1 and key2
+            let key = if i % 2 == 0 { "key1" } else { "key2" };
+            message.keys = Arc::from(vec![key.to_string()]);
+
+            message.value = bytes::Bytes::from(vec![1, 2, 3]);
+            message.id = MessageID {
+                vertex_name: "test-vertex".to_string().into(),
+                offset: i.to_string().into(),
+                index: 0,
+            };
+
+            let proto_message: Bytes = message.try_into().unwrap();
+            tx.send(SegmentWriteMessage::WriteData {
+                id: Some(format!("msg-{}", i)),
+                data: proto_message,
+            })
+            .await
+            .unwrap();
+        }
+
+        drop(tx);
+        handle.await.unwrap().unwrap();
+
+        // Create and run compactor with unaligned window kind
+        let compactor = Compactor::new(
+            test_path.clone(),
+            WindowKind::Unaligned,
+            1,    // 20MB
+            1000, // 1s flush interval
+            500,  // channel buffer
+        )
+        .await
+        .unwrap();
+
+        compactor.compact().await.unwrap();
+
+        // Verify compacted data
+        let compaction_wal = ReplayWal::new(WalType::Compact, test_path);
+        let (mut rx, handle) = compaction_wal.streaming_read().unwrap();
+
+        let mut remaining_message_count = 0;
+        let mut key1_count = 0;
+        let mut key2_count = 0;
+
+        while let Some(entry) = rx.next().await {
+            if let SegmentEntry::DataEntry { data, .. } = entry {
+                let msg: numaflow_pb::objects::isb::Message = prost::Message::decode(data).unwrap();
+                if let Some(header) = msg.header {
+                    // Check event time based on key
+                    if let (Some(message_info), keys) = (header.message_info.as_ref(), header.keys)
+                    {
+                        let event_time = utc_from_timestamp(message_info.event_time);
+
+                        if keys.contains(&"key1".to_string()) {
+                            key1_count += 1;
+                            assert!(
+                                event_time > gc_end_key1,
+                                "Found key1 message with event_time <= gc_end_key1"
+                            );
+                        } else if keys.contains(&"key2".to_string()) {
+                            key2_count += 1;
+                            assert!(
+                                event_time > gc_end_key2,
+                                "Found key2 message with event_time <= gc_end_key2"
+                            );
+                        }
+                    }
+                }
+                remaining_message_count += 1
+            }
+        }
+        handle.await.unwrap().unwrap();
+
+        // For key1, we should have removed messages with event time <= 10s (10 messages)
+        // For key2, we should have removed messages with event time <= 15s (15 messages)
+        // Total messages: 100 - (5 key1 messages + 8 key2 messages) = 87
+        // Note: Since we alternate keys and start from 1, key1 is at odd indices and key2 at even indices
+        assert_eq!(
+            remaining_message_count, 87,
+            "Expected 87 messages to remain after compaction"
+        );
+
+        // We expect 45 key1 messages (50 - 5) and 40 key2 messages (50 - 8)
+        assert_eq!(key1_count, 45, "Expected 45 key1 messages to remain");
+        assert_eq!(key2_count, 42, "Expected 42 key2 messages to remain");
+    }
+}
diff --git a/rust/numaflow-core/src/reduce/wal/segment/append.rs b/rust/numaflow-core/src/reduce/wal/segment/append.rs
new file mode 100644
index 00000000..0d4fdbad
--- /dev/null
+++ b/rust/numaflow-core/src/reduce/wal/segment/append.rs
@@ -0,0 +1,678 @@
+use crate::reduce::wal::error::WalResult;
+use crate::reduce::wal::segment::WalType;
+use bytes::Bytes;
+use chrono::{DateTime, Utc};
+use std::path::Path;
+use std::{io, path::PathBuf};
+use tokio::io::BufWriter;
+use tokio::task::JoinHandle;
+use tokio::{
+    fs::{File, OpenOptions},
+    io::AsyncWriteExt,
+    sync::mpsc::{self, Sender},
+    time::{interval, Duration},
+};
+use tokio_stream::wrappers::ReceiverStream;
+use tokio_stream::StreamExt;
+use tracing::{debug, info};
+
+/// Duration after which the WAL Segment is considered stale.
+const ROTATE_IF_STALE_DURATION: chrono::Duration = chrono::Duration::seconds(30);
+
+/// The Command that has to be operated on the Segment.
+pub(crate) enum SegmentWriteMessage {
+    /// Writes the given payload to the WAL.
+    WriteData {
+        /// Unique ID of the payload. Useful to detect write failures.
+        id: Option<String>,
+        /// Data to be written on do the WAL.
+        data: Bytes,
+    },
+    /// Rotates the file on demand.
+    Rotate {
+        /// When "true" Rotate only if the size is >= the size
+        on_size: bool,
+    },
+}
+
+/// Append only Segment writer that manages and rotates the WAL Segment.
+struct SegmentWriteActor {
+    /// Kind of WAL
+    wal_type: WalType,
+    /// Path where the WALs are persisted.
+    base_path: PathBuf,
+    /// Name of the current WAL Segment.
+    current_file_name: String,
+    /// Time when the segment was created.
+    create_time: DateTime<Utc>,
+    /// The current segment writer exposed via a buffer.
+    current_file_buf: BufWriter<File>,
+    /// Size of the current segment.
+    current_size: u64,
+    /// The file index since the last restart.
+    file_index: usize,
+    /// The interval at which the files have to be flushed.
+    flush_interval: Duration,
+    /// Maximum file size per segment.
+    max_file_size: u64,
+    /// The result of [SegmentWriteMessage] operation.
+    result_tx: Sender<String>,
+}
+
+impl Drop for SegmentWriteActor {
+    fn drop(&mut self) {
+        // TODO: do drop on shutdown on error path
+    }
+}
+
+impl SegmentWriteActor {
+    /// Creates a new SegmentWriteActor.
+    fn new(
+        wal_type: WalType,
+        base_path: PathBuf,
+        current_file_name: String,
+        create_time: DateTime<Utc>,
+        current_file: BufWriter<File>,
+        max_file_size: u64,
+        flush_interval: Duration,
+        result_tx: Sender<String>,
+    ) -> Self {
+        Self {
+            wal_type,
+            base_path,
+            current_file_name,
+            create_time,
+            current_file_buf: current_file,
+            current_size: 0,
+            file_index: 0,
+            max_file_size,
+            flush_interval,
+            result_tx,
+        }
+    }
+
+    /// Starts processing the [SegmentWriteMessage] operation.
+    async fn start_processing(
+        mut self,
+        in_rx: ReceiverStream<SegmentWriteMessage>,
+    ) -> WalResult<()> {
+        let mut flush_timer = interval(self.flush_interval);
+
+        let in_rx = in_rx.timeout(Duration::from_secs(15));
+        tokio::pin!(in_rx);
+
+        loop {
+            tokio::select! {
+                maybe_msg = in_rx.next() => {
+                    let Some(msg) = maybe_msg else {
+                        break;
+                    };
+                    match msg {
+                        Ok(msg) => {
+                            self.handle_message(msg).await?;
+                        }
+                        Err(_) => {
+                            debug!("calling flush and possible rotate on timeout");
+                            self.flush_and_rotate().await?;
+                        }
+                    }
+                }
+                _ = flush_timer.tick() => {
+                    self.flush().await?;
+                }
+            }
+        }
+
+        info!(?self.wal_type, "Stopping, doing a final flush and rotate!");
+        self.rotate_file(false).await
+    }
+
+    /// Processes each [SegmentWriteMessage] operation. This should only return critical errors, and
+    /// we should exit upon errors.
+    async fn handle_message(&mut self, msg: SegmentWriteMessage) -> WalResult<()> {
+        match msg {
+            SegmentWriteMessage::WriteData { id, data } => {
+                self.write_data(data).await?;
+                // we need to respond only if ID is provided
+                if let Some(id) = id {
+                    self.result_tx.send(id).await.unwrap();
+                }
+            }
+            SegmentWriteMessage::Rotate { on_size } => {
+                // Rotate if forced (`on_size` is false) OR if size threshold is met
+                if !on_size || self.current_size >= self.max_file_size {
+                    self.rotate_file(true).await?;
+                } else {
+                    debug!(
+                        current_size = self.current_size,
+                        max_size = self.max_file_size,
+                        "Skipping rotation: size threshold not met and not forced."
+                    );
+                }
+            }
+        }
+        Ok(())
+    }
+
+    /// Writes the data to the Segment.
+    /// The writes are in this format `<u64(data_len)><[u8;data_len]>`.
+    /// ### CANCEL SAFETY:
+    /// This is not Cancel Safe. We can make it cancel safe by removing the buffering and instead of
+    /// `write_all`, we should use `write`.
+    async fn write_data(&mut self, data: Bytes) -> WalResult<()> {
+        if self.current_size > 0 && self.current_size + data.len() as u64 >= self.max_file_size {
+            self.rotate_file(true).await?;
+        }
+
+        let data_len = data.len() as u64;
+
+        // TODO: this will have data loss if disk is full, we need to fix this late to be CANCEL SAFE.
+        self.current_file_buf.write_u64_le(data_len).await?;
+        self.current_file_buf.write_all(&data).await?;
+
+        self.current_size += data_len;
+
+        Ok(())
+    }
+
+    /// Flush the buffer to the underlying Segment file.
+    async fn flush(&mut self) -> WalResult<()> {
+        self.current_file_buf.flush().await?;
+        Ok(())
+    }
+
+    /// Open a segment for Appending. The file will be truncated if it exists (it shouldn't!).
+    async fn open_segment(
+        wal_type: &WalType,
+        base_path: &Path,
+        idx: usize,
+    ) -> WalResult<(String, BufWriter<File>)> {
+        let timestamp_nanos = Utc::now()
+            .timestamp_nanos_opt()
+            .unwrap_or(Utc::now().timestamp_micros());
+
+        // the name is important. sorting is based on the timestamp and then on file-index.
+        let filename = format!(
+            "{}_{}_{}.wal{}",
+            wal_type.segment_prefix(),
+            idx,
+            timestamp_nanos,
+            wal_type.segment_suffix()
+        );
+
+        let new_path = base_path.join(filename.clone());
+
+        debug!(path = %new_path.display(), "Opening new WAL segment file");
+
+        let file = OpenOptions::new()
+            .create(true)
+            .write(true)
+            .append(false)
+            .truncate(true)
+            .open(&new_path)
+            .await?;
+
+        Ok((new_path.display().to_string(), BufWriter::new(file)))
+    }
+
+    /// Flush and possibly rotate since there are no new writes to the file.
+    async fn flush_and_rotate(&mut self) -> WalResult<()> {
+        // if the file has not been rotated in these many seconds, let's rotate
+        if Utc::now().signed_duration_since(self.create_time) > ROTATE_IF_STALE_DURATION {
+            debug!(duration = ?ROTATE_IF_STALE_DURATION, "Rotating stale file, no entries for a while");
+            self.rotate_file(true).await
+        } else {
+            self.flush().await
+        }
+    }
+
+    /// Rotates the file and opens a new one if `open_new` is `true`. It renames the file on rotation
+    /// and resets the internal fields.
+    async fn rotate_file(&mut self, open_new: bool) -> WalResult<()> {
+        if self.current_size == 0 {
+            return Ok(());
+        }
+
+        info!(
+            current_size = ?self.current_size,
+            file_name = ?self.current_file_name,
+            "Rotating WAL segment file"
+        );
+        self.flush().await?;
+
+        // rename the current file before we start a new one. Remove the suffix before freezing.
+        let to_file_name = if self.wal_type.segment_suffix().is_empty() {
+            format!("{}.frozen", self.current_file_name)
+        } else {
+            // trim the suffix if suffix exists
+            let to_file_name = self
+                .current_file_name
+                .trim_end_matches(&self.wal_type.segment_suffix());
+            format!("{}.frozen", to_file_name)
+        };
+        tokio::fs::rename(&self.current_file_name, &to_file_name).await?;
+        info!(?self.current_file_name, ?to_file_name, "rename successful");
+
+        if open_new {
+            // open new segment
+            self.file_index += 1;
+            let (file_name, buf_file) =
+                Self::open_segment(&self.wal_type, &self.base_path, self.file_index).await?;
+
+            self.current_file_name = file_name;
+            self.current_file_buf = buf_file;
+            self.create_time = Utc::now();
+            self.current_size = 0;
+        }
+
+        Ok(())
+    }
+}
+
+/// Append Only WAL.
+#[derive(Clone)]
+pub(crate) struct AppendOnlyWal {
+    wal_type: WalType,
+    base_path: PathBuf,
+    max_file_size_mb: u64,
+    flush_interval_ms: u64,
+    channel_buffer_size: usize,
+}
+
+impl AppendOnlyWal {
+    /// Creates an [AppendOnlyWal]
+    pub(crate) async fn new(
+        wal_type: WalType,
+        base_path: PathBuf,
+        max_file_size_mb: u64,
+        flush_interval_ms: u64,
+        channel_buffer_size: usize,
+    ) -> Result<Self, io::Error> {
+        tokio::fs::create_dir_all(&base_path).await?;
+        Ok(Self {
+            wal_type,
+            base_path,
+            max_file_size_mb,
+            flush_interval_ms,
+            channel_buffer_size,
+        })
+    }
+
+    /// Start the WAL for streaming write.
+    /// CANCEL SAFETY: This is not cancel safe. This is because our writes are buffered.
+    pub(crate) async fn streaming_write(
+        self,
+        stream: ReceiverStream<SegmentWriteMessage>,
+    ) -> WalResult<(ReceiverStream<String>, JoinHandle<WalResult<()>>)> {
+        let (file_name, file_buf) =
+            SegmentWriteActor::open_segment(&self.wal_type, &self.base_path, 0).await?;
+
+        let (result_tx, result_rx) = mpsc::channel::<String>(self.channel_buffer_size);
+
+        let max_file_size_bytes = self.max_file_size_mb * 1024 * 1024;
+        let flush_duration = Duration::from_millis(self.flush_interval_ms);
+
+        let mut actor = SegmentWriteActor::new(
+            self.wal_type,
+            self.base_path,
+            file_name,
+            Utc::now(),
+            file_buf,
+            max_file_size_bytes,
+            flush_duration,
+            result_tx,
+        );
+
+        actor.current_size = 0;
+        actor.file_index = 0;
+
+        let handle = tokio::spawn(async move {
+            actor.start_processing(stream).await?;
+            Ok(())
+        });
+
+        info!("FileWriterActor spawned and running.");
+
+        Ok((ReceiverStream::new(result_rx), handle))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::reduce::wal::segment::WalType;
+    use bytes::{Bytes, BytesMut};
+    use futures::stream::StreamExt;
+    use std::fs;
+    use tempfile::tempdir;
+
+    #[tokio::test]
+    async fn test_wal_write_receive_results_and_rotate() {
+        let temp_dir = tempdir().unwrap();
+        let base_path = temp_dir.path().to_path_buf();
+        let max_file_size_mb = 1;
+        let flush_interval_ms = 50;
+        let channel_buffer = 10;
+
+        let wal_writer = AppendOnlyWal::new(
+            WalType::Data,
+            base_path.clone(),
+            max_file_size_mb,
+            flush_interval_ms,
+            channel_buffer,
+        )
+        .await
+        .expect("WAL creation failed");
+
+        let (wal_tx, wal_rx) = mpsc::channel::<SegmentWriteMessage>(channel_buffer);
+        let (mut result_rx, _writer_handle) = wal_writer
+            .streaming_write(ReceiverStream::new(wal_rx))
+            .await
+            .expect("Failed to start WAL service");
+
+        let mut expected_ids = Vec::new();
+        let mut received_results = Vec::new();
+
+        let id1 = "msg-001".to_string();
+        let data1 = Bytes::from("some initial data");
+        expected_ids.push(id1.clone());
+        wal_tx
+            .send(SegmentWriteMessage::WriteData {
+                id: Some(id1),
+                data: data1.clone(),
+            })
+            .await
+            .unwrap();
+
+        let id2 = "msg-002".to_string();
+        let data2 = Bytes::from(" more data here");
+        expected_ids.push(id2.clone());
+        wal_tx
+            .send(SegmentWriteMessage::WriteData {
+                id: Some(id2),
+                data: data2.clone(),
+            })
+            .await
+            .unwrap();
+
+        let large_data_size = (0.6 * 1024.0 * 1024.0) as usize;
+        let large_data1 = Bytes::from(vec![b'A'; large_data_size]);
+        let id3 = "msg-large-1".to_string();
+        expected_ids.push(id3.clone());
+        wal_tx
+            .send(SegmentWriteMessage::WriteData {
+                id: Some(id3),
+                data: large_data1.clone(),
+            })
+            .await
+            .unwrap();
+
+        let large_data2 = Bytes::from(vec![b'B'; large_data_size]);
+        let id4 = "msg-large-2".to_string();
+        expected_ids.push(id4.clone());
+        wal_tx
+            .send(SegmentWriteMessage::WriteData {
+                id: Some(id4),
+                data: large_data2.clone(),
+            })
+            .await
+            .unwrap();
+
+        let id5 = "msg-005-rotated".to_string();
+        let data5 = Bytes::from("data after rotation");
+        expected_ids.push(id5.clone());
+        wal_tx
+            .send(SegmentWriteMessage::WriteData {
+                id: Some(id5),
+                data: data5.clone(),
+            })
+            .await
+            .unwrap();
+
+        let id6 = "msg-006-rotated".to_string();
+        let data6 = Bytes::from(" more data after rotation");
+        expected_ids.push(id6.clone());
+        wal_tx
+            .send(SegmentWriteMessage::WriteData {
+                id: Some(id6),
+                data: data6.clone(),
+            })
+            .await
+            .unwrap();
+
+        drop(wal_tx);
+
+        while let Some(result) = result_rx.next().await {
+            received_results.push(result);
+        }
+
+        assert_eq!(
+            received_results.len(),
+            expected_ids.len(),
+            "Should receive a result for every message sent"
+        );
+
+        let mut received_ids = Vec::new();
+        for result in received_results {
+            received_ids.push(result);
+        }
+
+        received_ids.sort();
+        expected_ids.sort();
+        assert_eq!(
+            received_ids, expected_ids,
+            "Mismatch between expected and received successful IDs"
+        );
+
+        let mut files: Vec<_> = fs::read_dir(&base_path)
+            .unwrap()
+            .map(|entry| entry.unwrap().path())
+            .filter(|path| {
+                path.extension()
+                    .map_or(false, |ext| ext == "wal" || ext == "frozen")
+            }) // Ensure we only check .wal files
+            .collect();
+        files.sort();
+
+        assert_eq!(
+            files.len(),
+            2,
+            "There should be at-most 2 WAL segment files (segment_0 and segment_1)"
+        );
+
+        let first_file_content = fs::read(&files[0]).unwrap();
+        let expected_first_file_content =
+            [data1.as_ref(), data2.as_ref(), large_data1.as_ref()].concat();
+        assert_eq!(
+            first_file_content.len(),
+            expected_first_file_content.len() + size_of::<[u8; 8]>() * 3,
+            "Content mismatch in first segment file"
+        );
+
+        let second_file_content = fs::read(&files[1]).unwrap();
+        let expected_second_file_content =
+            [large_data2.as_ref(), data5.as_ref(), data6.as_ref()].concat();
+        assert_eq!(
+            second_file_content.len(),
+            expected_second_file_content.len() + size_of::<[u8; 8]>() * 3,
+            "Content mismatch in second segment file"
+        );
+
+        // FIXME: why does this hang?
+        // writer_handle
+        //     .await
+        //     .expect("writer shouldn't fail")
+        //     .expect("WAL shouldn't raise error");
+    }
+
+    #[tokio::test]
+    async fn test_wal_periodic_flush() {
+        let temp_dir = tempdir().unwrap();
+        let base_path = temp_dir.path().to_path_buf();
+        let max_file_size_mb = 10;
+        let flush_interval_ms = 50;
+        let channel_buffer = 10;
+
+        let wal_writer = AppendOnlyWal::new(
+            WalType::Data,
+            base_path.clone(),
+            max_file_size_mb,
+            flush_interval_ms,
+            channel_buffer,
+        )
+        .await
+        .expect("failed to create wal");
+
+        let (wal_tx, wal_rx) = mpsc::channel::<SegmentWriteMessage>(channel_buffer);
+        let (mut result_rx, _writer_handle) = wal_writer
+            .streaming_write(ReceiverStream::new(wal_rx))
+            .await
+            .expect("Failed to start WAL service");
+
+        let id1 = Some("flush-test-1".to_string());
+        let data1 = Bytes::from("Data to be flushed");
+        wal_tx
+            .send(SegmentWriteMessage::WriteData {
+                id: id1.clone(),
+                data: data1.clone(),
+            })
+            .await
+            .unwrap();
+
+        tokio::time::sleep(Duration::from_millis(flush_interval_ms * 3)).await;
+
+        let files: Vec<_> = fs::read_dir(&base_path)
+            .unwrap()
+            .map(|entry| entry.unwrap().path())
+            .filter(|path| {
+                path.extension()
+                    .map_or(false, |ext| ext == "wal" || ext == "frozen")
+            })
+            .collect();
+
+        assert_eq!(files.len(), 1, "There should be 1 WAL file");
+        let file_content = fs::read(&files[0]).unwrap();
+        let mut expected = BytesMut::new();
+        expected.extend_from_slice(&data1.len().to_le_bytes());
+        expected.extend_from_slice(&data1);
+        assert_eq!(
+            file_content, expected,
+            "File content should match flushed data"
+        );
+
+        let result = result_rx.next().await.expect("Should receive a result");
+        assert_eq!(
+            result,
+            id1.unwrap(),
+            "Should have received the correct ID back"
+        );
+
+        // FIXME: why does this hang?
+        // writer_handle
+        //     .await
+        //     .expect("writer shouldn't fail")
+        //     .expect("WAL shouldn't raise error");
+    }
+
+    #[tokio::test]
+    async fn test_wal_rotate_command() {
+        let temp_dir = tempdir().unwrap();
+        let base_path = temp_dir.path().to_path_buf();
+        let max_file_size_mb = 10;
+        let flush_interval_ms = 50;
+        let channel_buffer = 10;
+
+        let wal_writer = AppendOnlyWal::new(
+            WalType::Data,
+            base_path.clone(),
+            max_file_size_mb,
+            flush_interval_ms,
+            channel_buffer,
+        )
+        .await
+        .expect("failed to create wal");
+
+        let (wal_tx, wal_rx) = mpsc::channel::<SegmentWriteMessage>(channel_buffer);
+        let (_result_rx, _writer_handle) = wal_writer
+            .streaming_write(ReceiverStream::new(wal_rx))
+            .await
+            .expect("Failed to start WAL service");
+
+        // Send some data to the WAL
+        let id1 = Some("msg-001".to_string());
+        let data1 = Bytes::from("data before rotation");
+        wal_tx
+            .send(SegmentWriteMessage::WriteData {
+                id: id1.clone(),
+                data: data1.clone(),
+            })
+            .await
+            .unwrap();
+
+        // Send the Rotate command
+        wal_tx
+            .send(SegmentWriteMessage::Rotate { on_size: false })
+            .await
+            .unwrap();
+
+        // Send more data after rotation
+        let id2 = Some("msg-002".to_string());
+        let data2 = Bytes::from("data after rotation");
+        wal_tx
+            .send(SegmentWriteMessage::WriteData {
+                id: id2.clone(),
+                data: data2.clone(),
+            })
+            .await
+            .unwrap();
+
+        drop(wal_tx);
+
+        // Allow some time for the rotation to complete
+        tokio::time::sleep(Duration::from_millis(100)).await;
+
+        // Verify the files
+        let mut files: Vec<_> = fs::read_dir(&base_path)
+            .unwrap()
+            .map(|entry| entry.unwrap().path())
+            .filter(|path| path.extension().map_or(false, |ext| ext == "frozen"))
+            .collect();
+        files.sort();
+
+        assert_eq!(files.len(), 2, "There should be 2 WAL segment files");
+
+        // Verify the rotated file is renamed correctly
+        assert!(files[0]
+            .file_name()
+            .unwrap()
+            .to_str()
+            .unwrap()
+            .ends_with(".frozen"));
+
+        // Verify the content of the rotated file
+        let rotated_file_content = fs::read(&files[0]).unwrap();
+        let mut expected_rotated_content = BytesMut::new();
+        expected_rotated_content.extend_from_slice(&data1.len().to_le_bytes());
+        expected_rotated_content.extend_from_slice(&data1);
+        assert_eq!(
+            rotated_file_content, expected_rotated_content,
+            "Rotated file content mismatch"
+        );
+
+        // Verify the content of the active file
+        let active_file_content = fs::read(&files[1]).unwrap();
+        let mut expected_active_content = BytesMut::new();
+        expected_active_content.extend_from_slice(&data2.len().to_le_bytes());
+        expected_active_content.extend_from_slice(&data2);
+        assert_eq!(
+            active_file_content, expected_active_content,
+            "Active file content mismatch"
+        );
+
+        // FIXME: why does this hang?
+        // writer_handle
+        //     .await
+        //     .expect("writer shouldn't fail")
+        //     .expect("WAL shouldn't raise error");
+    }
+}
diff --git a/rust/numaflow-core/src/reduce/wal/segment/compactor.rs b/rust/numaflow-core/src/reduce/wal/segment/compactor.rs
new file mode 100644
index 00000000..9d3070c3
--- /dev/null
+++ b/rust/numaflow-core/src/reduce/wal/segment/compactor.rs
@@ -0,0 +1,788 @@
+//! Compactor has business logic. It knows what kind of WALs have been created and will
+//! compact based on the type. WAL inherently is agnostic to data. The compactor will be given
+//! multiple WAL types (data, gc, etc.) and it decides how to purge (aka compact).
+//!
+//! Compaction is done by combing the [WalType::Gc] + [WalType::Data] + [WalType::Compact]
+//!
+//! ### Compaction Logic
+//! #### Aligned kind
+//! 1. Replay all the GC events and store the max end time
+//! 2. Replay all the data events and only retain the messages with event time > max end time
+//!
+//! #### Unaligned kind
+//! 1. Replay all the GC events and store the max end time for every key combination(map[key] = max end time)
+//! 2. Replay all the data events and only retain the messages with event time > max end time for that key
+//!
+
+use crate::reduce::wal::error::WalResult;
+use crate::reduce::wal::segment::append::{AppendOnlyWal, SegmentWriteMessage};
+use crate::reduce::wal::segment::replay::{ReplayWal, SegmentEntry};
+use crate::reduce::wal::segment::GcEventEntry;
+use crate::reduce::wal::segment::WalType;
+use crate::shared::grpc::utc_from_timestamp;
+use chrono::{DateTime, Utc};
+use numaflow_pb::objects::isb;
+use numaflow_pb::objects::wal::GcEvent;
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::time::{Duration, UNIX_EPOCH};
+use tokio::sync::mpsc;
+use tokio::task::JoinHandle;
+use tokio_stream::{wrappers::ReceiverStream, StreamExt};
+use tokio_util::sync::CancellationToken;
+use tracing::info;
+
+/// WALs can represent two Kinds of Windows and data is different for each Kind.
+pub(crate) enum WindowKind {
+    /// Aligned represents Fixed and Sliding Windows.
+    Aligned,
+    /// Unaligned represents Session Windows and Accumulators (almost like Global Windows).
+    Unaligned,
+}
+
+/// A Compactor that compacts based on the GC and Segment WAL files in the given path. It can
+/// compact both [WindowKind] of WALs.
+pub(crate) struct Compactor {
+    /// GC WAL
+    gc_wal: ReplayWal,
+    /// Segment WAL
+    segment_wal: ReplayWal,
+    /// Compaction WAL for Read Only
+    compaction_ro_wal: ReplayWal,
+    /// New Compaction WAL for writing compacted data
+    compaction_ao_wal: AppendOnlyWal,
+    /// Kind of Window, Aligned or Unaligned
+    kind: WindowKind,
+}
+
+const WAL_KEY_SEPERATOR: &'static str = ":";
+
+impl Compactor {
+    /// Creates a new Compactor.
+    pub(crate) async fn new(
+        path: PathBuf,
+        kind: WindowKind,
+        max_file_size_mb: u64,
+        flush_interval_ms: u64,
+        channel_buffer_size: usize,
+    ) -> WalResult<Self> {
+        let segment_wal = ReplayWal::new(WalType::Data, path.clone());
+        let gc_wal = ReplayWal::new(WalType::Gc, path.clone());
+        let compaction_ro_wal = ReplayWal::new(WalType::Compact, path.clone());
+        let compaction_ao_wal = AppendOnlyWal::new(
+            WalType::Compact,
+            path.clone(),
+            max_file_size_mb,
+            flush_interval_ms,
+            channel_buffer_size,
+        )
+        .await?;
+
+        Ok(Self {
+            gc_wal,
+            segment_wal,
+            compaction_ro_wal,
+            compaction_ao_wal,
+            kind,
+        })
+    }
+
+    /// Starts the compaction process every `interval_ms` milliseconds.
+    pub(crate) async fn start_compaction(
+        self,
+        duration: Duration,
+        cln_token: CancellationToken,
+    ) -> WalResult<JoinHandle<WalResult<()>>> {
+        let mut tick = tokio::time::interval(duration);
+        let result: JoinHandle<WalResult<()>> = tokio::spawn(async move {
+            loop {
+                tokio::select! {
+                _ = tick.tick() => {
+                        self.compact().await?;
+                     }
+                _ = cln_token.cancelled() => {
+                        break;
+                    }
+                }
+            }
+            Ok(())
+        });
+        Ok(result)
+    }
+
+    /// Compact first needs to get all the GC files and build a compaction map. This map will have
+    /// the oldest data before which all can be deleted.
+    pub(crate) async fn compact(&self) -> WalResult<()> {
+        match self.kind {
+            WindowKind::Aligned => self.compact_aligned().await?,
+            WindowKind::Unaligned => self.compact_unaligned().await?,
+        }
+        Ok(())
+    }
+
+    /// Compacts Aligned Segments.
+    /// ## Logic
+    /// - Get the oldest time by calling Build Aligned Compaction
+    /// - Get all the Segment files
+    /// - Open Append WAL for Compaction
+    /// - For each Segment file, deserialize data to [crate::message::Message]
+    /// - Compare the "event_time" of the Message with the oldest time
+    /// - If event_time is <= oldest_time, skip it, otherwise write it into the Compaction Append WAL.
+    /// - Send the Rotate message to the Compaction Append WAL after each Segment file has been processed.
+    /// - Delete the Segment File after the Rotate is complete.
+    async fn compact_aligned(&self) -> WalResult<()> {
+        // Get the oldest time and scanned GC files
+        let (oldest_time, gc_files) = self.build_aligned_compaction().await?;
+
+        let compact = AlignedCompaction(oldest_time);
+
+        // Compact the compaction_ro_wal
+        self.process_wal_stream(&self.compaction_ro_wal, &compact)
+            .await?;
+
+        // Compact the segment_wal
+        self.process_wal_stream(&self.segment_wal, &compact).await?;
+
+        // Delete the GC files
+        for gc_file in gc_files {
+            info!(gc_file = %gc_file.display(), "removing segment file");
+            tokio::fs::remove_file(gc_file).await?;
+        }
+
+        Ok(())
+    }
+
+    /// Compacts Unaligned Segments.
+    /// ## Logic
+    /// - Get the oldest time by calling Build Unaligned Compaction
+    /// - Get all the Segment files
+    /// - Open Append WAL for Compaction
+    /// - For each Segment file, deserialize data to [crate::message::Message]
+    /// - Compare the "event_time" of the Message with the compaction hashmap. The key is the joined keys
+    ///   are the key from the Message::Header.
+    /// - If event_time is <= oldest_time, skip it, otherwise write it into the Compaction Append WAL.
+    /// - Send the Rotate message to the Compaction Append WAL after each Segment file has been processed.
+    /// - Delete the Segment File after the Rotate is complete.
+    async fn compact_unaligned(&self) -> WalResult<()> {
+        // Get the oldest time map and scanned GC files
+        let (oldest_time_map, gc_files) = self.build_unaligned_compaction().await?;
+
+        let compact = UnalignedCompaction(oldest_time_map);
+
+        // Compact the compaction_ro_wal
+        self.process_wal_stream(&self.compaction_ro_wal, &compact)
+            .await?;
+
+        // Compact the segment_wal
+        self.process_wal_stream(&self.segment_wal, &compact).await?;
+
+        // Delete the GC files
+        for gc_file in gc_files {
+            info!(gc_file = %gc_file.display(), "removing segment file");
+            tokio::fs::remove_file(gc_file).await?;
+        }
+
+        Ok(())
+    }
+
+    /// Process the WAL stream (Readonly) and write the compacted data to the compaction WAL.
+    /// Every message is checked against the should_retain function and if it returns true, it is
+    /// written to the compaction WAL.
+    async fn process_wal_stream<T: ShouldRetain>(
+        &self,
+        wal: &ReplayWal,
+        should_retain: &T,
+    ) -> WalResult<()> {
+        // Get a streaming reader for the WAL
+        let (mut rx, handle) = wal.clone().streaming_read()?;
+
+        let (wal_tx, wal_rx) = mpsc::channel(100);
+        let (_result_rx, writer_handle) = self
+            .compaction_ao_wal
+            .clone()
+            .streaming_write(ReceiverStream::new(wal_rx))
+            .await
+            .map_err(|e| format!("Failed to start compaction WAL writer: {}", e))?;
+
+        // Process each WAL entry
+        while let Some(entry) = rx.next().await {
+            match entry {
+                SegmentEntry::DataEntry { data, .. } => {
+                    // Deserialize the message
+                    let msg: isb::Message = prost::Message::decode(data.clone())
+                        .map_err(|e| format!("Failed to decode message: {}", e))?;
+
+                    if should_retain.should_retain_message(&msg)? {
+                        // Send the message to the compaction WAL
+                        wal_tx
+                            .send(SegmentWriteMessage::WriteData { id: None, data })
+                            .await
+                            .map_err(|e| {
+                                format!("Failed to send message to compaction WAL: {}", e)
+                            })?;
+                    }
+                }
+                SegmentEntry::DataFooter { .. } => {
+                    unimplemented!()
+                }
+                SegmentEntry::CmdFileSwitch { filename } => {
+                    // Send rotate message after processing each file
+                    wal_tx
+                        .send(SegmentWriteMessage::Rotate { on_size: false })
+                        .await
+                        .map_err(|e| format!("Failed to send rotate command: {}", e))?;
+
+                    // Delete the processed segment file
+                    tokio::fs::remove_file(&filename).await.map_err(|e| {
+                        format!(
+                            "Failed to delete segment file {}: {}",
+                            filename.display(),
+                            e
+                        )
+                    })?;
+
+                    info!(filename = %filename.display(), "removing segment file");
+                }
+            }
+        }
+
+        // Wait for the WAL reader to complete
+        handle
+            .await
+            .map_err(|e| format!("WAL reader failed: {}", e))??;
+
+        // Drop the sender to close the channel
+        drop(wal_tx);
+
+        // Wait for the writer to complete
+        writer_handle
+            .await
+            .map_err(|e| format!("Compaction writer failed: {}", e))??;
+
+        Ok(())
+    }
+
+    /// Builds the oldest time below which all data has been processed. For Aligned WAL, all we need
+    /// to track is the oldest timestamp. We do not have to worry about the keys.
+    /// It returns the list of GC files scanned.
+    async fn build_aligned_compaction(&self) -> WalResult<(DateTime<Utc>, Vec<PathBuf>)> {
+        // The oldest time across all GC Segments.
+        let mut oldest_time = DateTime::from(UNIX_EPOCH);
+
+        // list of GC Segments scanned
+        let mut scanned_files = vec![];
+
+        // Get a streaming reader for the GC WAL.
+        let (mut rx, handle) = self.gc_wal.clone().streaming_read()?;
+
+        while let Some(entry) = rx.next().await {
+            match entry {
+                SegmentEntry::DataEntry { data, .. } => {
+                    // Decode the GC event.
+                    let gc: GcEvent = prost::Message::decode(data)
+                        .map_err(|e| format!("Failed to decode GC event: {e}"))?;
+                    let gc: GcEventEntry = gc.into();
+
+                    // Update the oldest time if the current GC event's end time is newer.
+                    oldest_time = oldest_time.max(gc.end_time);
+                }
+                SegmentEntry::DataFooter { .. } => {
+                    unimplemented!()
+                }
+                SegmentEntry::CmdFileSwitch { filename } => {
+                    // Add the scanned file to the list.
+                    scanned_files.push(filename);
+                }
+            }
+        }
+
+        // Wait for the WAL reader to complete.
+        handle
+            .await
+            .map_err(|e| format!("WAL reader failed: {e}"))??;
+
+        Ok((oldest_time, scanned_files))
+    }
+
+    /// Builds the oldest time below which all data has been processed. For Unaligned WAL, we need
+    /// to track the oldest timestamp for the given keys.
+    /// It returns the list of GC files scanned.
+    async fn build_unaligned_compaction(
+        &self,
+    ) -> WalResult<(HashMap<String, DateTime<Utc>>, Vec<PathBuf>)> {
+        // Map of keys to their oldest time
+        let mut oldest_time_map = HashMap::new();
+        // List of GC segment files scanned
+        let mut scanned_files = Vec::new();
+
+        // Get a streaming reader for the GC WAL
+        let (mut rx, handle) = self.gc_wal.clone().streaming_read()?;
+
+        while let Some(entry) = rx.next().await {
+            match entry {
+                SegmentEntry::DataEntry { data, .. } => {
+                    // Decode the GC event
+                    let gc: GcEvent = prost::Message::decode(data)
+                        .map_err(|e| format!("Failed to decode GC event: {e}"))?;
+                    let gc: GcEventEntry = gc.into();
+
+                    // Update the oldest time map for the given keys
+                    if let Some(keys) = gc.keys {
+                        let key = keys.join(WAL_KEY_SEPERATOR);
+                        oldest_time_map
+                            .entry(key)
+                            .and_modify(|dt| {
+                                if gc.end_time > *dt {
+                                    *dt = gc.end_time;
+                                }
+                            })
+                            .or_insert(gc.end_time);
+                    }
+                }
+                SegmentEntry::DataFooter { .. } => {
+                    unimplemented!()
+                }
+                SegmentEntry::CmdFileSwitch { filename } => {
+                    // Add the scanned file to the list
+                    scanned_files.push(filename);
+                }
+            }
+        }
+
+        // Wait for the WAL reader to complete
+        handle
+            .await
+            .map_err(|e| format!("WAL reader failed: {e}"))??;
+
+        Ok((oldest_time_map, scanned_files))
+    }
+}
+
+/// ShouldRetain trait defines a method to determine whether a message should be retained.
+trait ShouldRetain {
+    /// Determines whether a message should be retained based on its event time in the message
+    fn should_retain_message(&self, msg: &isb::Message) -> WalResult<bool>;
+}
+
+/// AlignedCompaction holds the oldest time below which all data has been processed. The event time
+/// can be the same for all Keys.
+struct AlignedCompaction(DateTime<Utc>);
+
+/// UnalignedCompaction holds the oldest time below which all data has been processed. The event time
+/// will be different for each Key.
+struct UnalignedCompaction(HashMap<String, DateTime<Utc>>);
+
+impl ShouldRetain for AlignedCompaction {
+    /// Determines whether a message should be retained based on its event time.
+    fn should_retain_message(&self, msg: &isb::Message) -> WalResult<bool> {
+        // Extract the event time from the message
+        let event_time = msg
+            .header
+            .as_ref()
+            .and_then(|header| header.message_info.as_ref())
+            .map(|info| utc_from_timestamp(info.event_time))
+            .expect("Failed to extract event time from message");
+
+        // Retain the message if its event time is greater than the oldest time
+        Ok(event_time > self.0)
+    }
+}
+
+impl ShouldRetain for UnalignedCompaction {
+    /// Determines whether an unaligned message should be retained based on its event time and keys.
+    fn should_retain_message(&self, msg: &isb::Message) -> WalResult<bool> {
+        // Extract the event time from the message
+        let event_time = msg
+            .header
+            .as_ref()
+            .and_then(|header| header.message_info.as_ref())
+            .map(|info| utc_from_timestamp(info.event_time))
+            .expect("Failed to extract event time from message");
+
+        // Extract the keys from the message
+        let keys = msg
+            .header
+            .as_ref()
+            .map(|header| header.keys.clone())
+            .unwrap_or_default();
+
+        // Join the keys with the separator
+        let key = keys.join(WAL_KEY_SEPERATOR);
+
+        // Check if the key exists in the map
+        if let Some(oldest_time) = self.0.get(&key) {
+            // Retain the message if its event time is greater than the oldest time for this key
+            return Ok(event_time > *oldest_time);
+        }
+
+        // If the key is not in the map, retain the message
+        Ok(true)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::message::{Message, MessageID};
+    use crate::shared::grpc::prost_timestamp_from_utc;
+    use bytes::Bytes;
+    use chrono::TimeZone;
+    use prost_types;
+    use std::fs;
+    use std::sync::Arc;
+    use tempfile::tempdir;
+    use tokio::sync::mpsc;
+    use tokio_stream::wrappers::ReceiverStream;
+
+    #[tokio::test]
+    async fn test_build_aligned_compaction() -> WalResult<()> {
+        // Create a temporary directory for the test
+        let temp_dir = tempdir().unwrap();
+        let path = temp_dir.path().to_path_buf();
+
+        // Create WAL directories
+        fs::create_dir_all(path.join(WalType::Gc.to_string())).unwrap();
+        fs::create_dir_all(path.join(WalType::Data.to_string())).unwrap();
+        fs::create_dir_all(path.join(WalType::Compact.to_string())).unwrap();
+
+        // Create a compactor with default test values
+        let compactor = Compactor::new(
+            path.clone(),
+            WindowKind::Aligned,
+            100,  // max_file_size_mb
+            1000, // flush_interval_ms
+            1000, // channel_buffer_size
+        )
+        .await?;
+
+        // Create some test GC events with different timestamps
+        let gc_events = vec![
+            GcEvent {
+                end_time: Some(prost_types::Timestamp {
+                    seconds: 1000,
+                    nanos: 0,
+                }),
+                ..Default::default()
+            },
+            GcEvent {
+                end_time: Some(prost_types::Timestamp {
+                    seconds: 2000,
+                    nanos: 0,
+                }),
+                ..Default::default()
+            },
+        ];
+
+        // Create a WAL writer for GC events
+        let gc_writer = AppendOnlyWal::new(
+            WalType::Gc,
+            path.clone(),
+            100,  // max_file_size_mb
+            1000, // flush_interval_ms
+            100,  // channel_buffer_size
+        )
+        .await?;
+
+        // Set up streaming write for the GC WAL
+        let (tx, rx) = mpsc::channel(100);
+        let (_result_rx, writer_handle) =
+            gc_writer.streaming_write(ReceiverStream::new(rx)).await?;
+
+        // Write GC events to the WAL
+        for event in gc_events {
+            let mut buf = Vec::new();
+            prost::Message::encode(&event, &mut buf)
+                .map_err(|e| format!("Failed to encode GC event: {e}"))?;
+            tx.send(SegmentWriteMessage::WriteData {
+                id: None,
+                data: Bytes::from(buf),
+            })
+            .await
+            .map_err(|e| format!("Failed to send data: {e}"))?;
+        }
+
+        // Drop the sender to close the channel
+        drop(tx);
+
+        // Wait for the writer to complete
+        writer_handle
+            .await
+            .map_err(|e| format!("Writer failed: {e}"))??;
+
+        // Run build_aligned_compaction
+        let (oldest_time, scanned_files) = compactor.build_aligned_compaction().await?;
+
+        // Verify the results
+        assert_eq!(oldest_time, Utc.timestamp_opt(2000, 0).unwrap());
+        assert!(!scanned_files.is_empty());
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_build_unaligned_compaction() -> WalResult<()> {
+        // Create a temporary directory for the test
+        let temp_dir = tempdir().unwrap();
+        let path = temp_dir.path().to_path_buf();
+
+        // Create WAL directories
+        fs::create_dir_all(path.join(WalType::Gc.to_string())).unwrap();
+        fs::create_dir_all(path.join(WalType::Data.to_string())).unwrap();
+        fs::create_dir_all(path.join(WalType::Compact.to_string())).unwrap();
+
+        // Create a compactor with default test values
+        let compactor = Compactor::new(
+            path.clone(),
+            WindowKind::Unaligned,
+            100,  // max_file_size_mb
+            1000, // flush_interval_ms
+            1000, // channel_buffer_size
+        )
+        .await?;
+
+        // Create some test GC events with different timestamps and keys
+        let gc_events = vec![
+            GcEvent {
+                end_time: Some(prost_types::Timestamp {
+                    seconds: 1000,
+                    nanos: 0,
+                }),
+                keys: vec!["key1".to_string(), "key2".to_string()],
+                ..Default::default()
+            },
+            GcEvent {
+                end_time: Some(prost_types::Timestamp {
+                    seconds: 2000,
+                    nanos: 0,
+                }),
+                keys: vec!["key1".to_string(), "key2".to_string()],
+                ..Default::default()
+            },
+            GcEvent {
+                end_time: Some(prost_types::Timestamp {
+                    seconds: 1500,
+                    nanos: 0,
+                }),
+                keys: vec!["key3".to_string(), "key4".to_string()],
+                ..Default::default()
+            },
+        ];
+
+        // Create a WAL writer for GC events
+        let gc_writer = AppendOnlyWal::new(
+            WalType::Gc,
+            path.clone(),
+            100,  // max_file_size_mb
+            1000, // flush_interval_ms
+            100,  // channel_buffer_size
+        )
+        .await?;
+
+        // Set up streaming write for the GC WAL
+        let (tx, rx) = mpsc::channel(100);
+        let (_result_rx, writer_handle) =
+            gc_writer.streaming_write(ReceiverStream::new(rx)).await?;
+
+        // Write GC events to the WAL
+        for event in gc_events {
+            let mut buf = Vec::new();
+            prost::Message::encode(&event, &mut buf)
+                .map_err(|e| format!("Failed to encode GC event: {e}"))?;
+            tx.send(SegmentWriteMessage::WriteData {
+                id: None,
+                data: Bytes::from(buf),
+            })
+            .await
+            .map_err(|e| format!("Failed to send data: {e}"))?;
+        }
+
+        // Drop the sender to close the channel
+        drop(tx);
+
+        // Wait for the writer to complete
+        writer_handle
+            .await
+            .map_err(|e| format!("Writer failed: {e}"))??;
+
+        // Run build_unaligned_compaction
+        let (oldest_time_map, scanned_files) = compactor.build_unaligned_compaction().await?;
+
+        // Verify the results
+        assert!(!scanned_files.is_empty());
+        assert_eq!(oldest_time_map.len(), 2); // We expect two different key combinations
+
+        // Check the timestamps for each key combination
+        let key1_2 = "key1:key2".to_string();
+        let key3_4 = "key3:key4".to_string();
+
+        assert_eq!(
+            oldest_time_map.get(&key1_2),
+            Some(&Utc.timestamp_opt(2000, 0).unwrap())
+        );
+        assert_eq!(
+            oldest_time_map.get(&key3_4),
+            Some(&Utc.timestamp_opt(1500, 0).unwrap())
+        );
+
+        Ok(())
+    }
+
+    // FIXME
+    #[tokio::test]
+    async fn test_gc_wal_and_compaction_with_multiple_files() {
+        let test_path = tempfile::tempdir().unwrap().into_path();
+
+        // Create GC WAL
+        let gc_wal = AppendOnlyWal::new(
+            WalType::Gc,
+            test_path.clone(),
+            1,    // 1MB
+            1000, // 1s flush interval
+            500,  // channel buffer
+        )
+        .await
+        .unwrap();
+
+        // Create and write 2 GC events
+        let gc_start_1 = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 5).unwrap();
+        let gc_end_1 = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 10).unwrap();
+        let gc_event_1 = GcEvent {
+            start_time: Some(prost_timestamp_from_utc(gc_start_1)),
+            end_time: Some(prost_timestamp_from_utc(gc_end_1)),
+            keys: vec![],
+        };
+
+        let gc_start_2 = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 15).unwrap();
+        let gc_end_2 = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 20).unwrap();
+        let gc_event_2 = GcEvent {
+            start_time: Some(prost_timestamp_from_utc(gc_start_2)),
+            end_time: Some(prost_timestamp_from_utc(gc_end_2)),
+            keys: vec![],
+        };
+
+        let (tx, rx) = tokio::sync::mpsc::channel(10);
+        let (_offset_stream, handle) = gc_wal
+            .streaming_write(ReceiverStream::new(rx))
+            .await
+            .unwrap();
+
+        tx.send(SegmentWriteMessage::WriteData {
+            id: Some("gc1".to_string()),
+            data: bytes::Bytes::from(prost::Message::encode_to_vec(&gc_event_1)),
+        })
+        .await
+        .unwrap();
+
+        tx.send(SegmentWriteMessage::WriteData {
+            id: Some("gc2".to_string()),
+            data: bytes::Bytes::from(prost::Message::encode_to_vec(&gc_event_2)),
+        })
+        .await
+        .unwrap();
+
+        drop(tx);
+        handle.await.unwrap().unwrap();
+
+        // Create segment WAL
+        let segment_wal = AppendOnlyWal::new(
+            WalType::Data,
+            test_path.clone(),
+            1,    // 1MB
+            1000, // 1s flush interval
+            500,  // channel buffer
+        )
+        .await
+        .unwrap();
+
+        // Write 1000 segment entries across 10 files
+        let (tx, rx) = mpsc::channel(100);
+        let (mut offset_stream, handle) = segment_wal
+            .streaming_write(ReceiverStream::new(rx))
+            .await
+            .unwrap();
+
+        let write_result_cnt = tokio::spawn(async move {
+            let mut counter = 0;
+            while let Some(msg) = offset_stream.next().await {
+                counter += 1;
+            }
+            return counter;
+        });
+
+        let start_time = Utc.with_ymd_and_hms(2025, 4, 1, 1, 0, 0).unwrap();
+        let time_increment = chrono::Duration::seconds(1);
+
+        for i in 1..=1000 {
+            let mut message = Message::default();
+            message.event_time = start_time + (time_increment * i);
+            message.keys = Arc::from(vec!["test-key".to_string()]);
+            message.value = bytes::Bytes::from(vec![1, 2, 3]);
+            message.id = MessageID {
+                vertex_name: "test-vertex".to_string().into(),
+                offset: i.to_string().into(),
+                index: 0,
+            };
+
+            let proto_message: Bytes = message.try_into().unwrap();
+            tx.send(SegmentWriteMessage::WriteData {
+                id: Some(format!("msg-{}", i)),
+                data: proto_message,
+            })
+            .await
+            .unwrap();
+
+            // Rotate every 100 messages to create 10 files
+            if i % 100 == 0 {
+                tx.send(SegmentWriteMessage::Rotate { on_size: false })
+                    .await
+                    .unwrap();
+            }
+        }
+        drop(tx);
+        handle.await.unwrap().unwrap();
+
+        assert_eq!(write_result_cnt.await.unwrap(), 1000);
+
+        // Create and run compactor
+        let compactor = Compactor::new(
+            test_path.clone(),
+            WindowKind::Aligned,
+            1,    // 1MB
+            1000, // 1s flush interval
+            500,  // channel buffer
+        )
+        .await
+        .unwrap();
+
+        compactor.compact().await.unwrap();
+
+        // Verify compacted data
+        let compaction_wal = ReplayWal::new(WalType::Compact, test_path);
+        let (mut rx, handle) = compaction_wal.streaming_read().unwrap();
+
+        let mut remaining_message_count = 0;
+        while let Some(entry) = rx.next().await {
+            if let SegmentEntry::DataEntry { data, .. } = entry {
+                let msg: numaflow_pb::objects::isb::Message = prost::Message::decode(data).unwrap();
+                if let Some(header) = msg.header {
+                    if let Some(message_info) = header.message_info {
+                        let event_time = utc_from_timestamp(message_info.event_time);
+                        assert!(
+                            event_time > gc_end_2,
+                            "Found message with event_time <= gc_end_2"
+                        );
+                    }
+                }
+                remaining_message_count += 1;
+            }
+        }
+        handle.await.unwrap().unwrap();
+
+        // Verify the number of remaining messages
+        // Messages with event_time <= gc_end_2 should be removed
+        assert_eq!(
+            remaining_message_count, 980,
+            "Expected 980 messages to remain after compaction"
+        );
+    }
+}
diff --git a/rust/numaflow-core/src/reduce/wal/segment/replay.rs b/rust/numaflow-core/src/reduce/wal/segment/replay.rs
new file mode 100644
index 00000000..33d76405
--- /dev/null
+++ b/rust/numaflow-core/src/reduce/wal/segment/replay.rs
@@ -0,0 +1,244 @@
+use crate::reduce::wal::error::WalResult;
+use crate::reduce::wal::segment::WalType;
+use bytes::Bytes;
+use std::cmp::Ordering;
+use std::fs;
+use std::io;
+use std::path::PathBuf;
+use tokio::fs::OpenOptions;
+use tokio::task::JoinHandle;
+use tokio::{
+    io::{AsyncReadExt, BufReader},
+    sync::mpsc::{self, Sender},
+    task,
+};
+use tokio_stream::wrappers::ReceiverStream;
+use tracing::{debug, info};
+
+/// Segment Entry as recorded in the WAL.
+#[derive(Debug)]
+pub(crate) enum SegmentEntry {
+    /// Data entry in the Segment
+    DataEntry { size: u64, data: Bytes },
+    /// The file has been switched
+    CmdFileSwitch { filename: PathBuf },
+    /// Data footer in the Segment.
+    /// TODO: This is for optimization which is yet to be implemented
+    #[allow(dead_code)]
+    DataFooter { size: u64, data: Bytes },
+}
+
+/// Replay the WAL in-order.
+#[derive(Debug, Clone)]
+pub(crate) struct ReplayWal {
+    wal_type: WalType,
+    base_path: PathBuf,
+}
+
+impl ReplayWal {
+    /// Creates a new Replayer for the WAL.
+    pub(crate) fn new(wal_type: WalType, base_path: PathBuf) -> Self {
+        Self {
+            wal_type,
+            base_path,
+        }
+    }
+
+    /// Reads the WAL files and streams it via the stream. Stream will be closed once all the
+    /// entries are read.
+    pub(crate) fn streaming_read(
+        self,
+    ) -> WalResult<(ReceiverStream<SegmentEntry>, JoinHandle<WalResult<()>>)> {
+        let mut files: Vec<PathBuf> = list_files(&self.wal_type, self.base_path.clone());
+        files = sort_filenames(files);
+
+        debug!(count = files.len(), "Found WAL segment files for replay");
+
+        let (tx, rx) = mpsc::channel::<SegmentEntry>(128);
+
+        let handle = task::spawn(async move {
+            info!("Starting WAL replay...");
+            for file_path in files {
+                info!(file = %file_path.display(), "Replaying");
+                Self::read_segment(&file_path, tx.clone()).await?;
+
+                tx.send(SegmentEntry::CmdFileSwitch {
+                    filename: file_path,
+                })
+                .await
+                .expect("rx dropped")
+            }
+            info!("Finished WAL replay task...");
+
+            Ok(())
+        });
+
+        Ok((ReceiverStream::new(rx), handle))
+    }
+
+    /// Read the segment file in-order and write to the channel.
+    async fn read_segment(path: &PathBuf, tx: Sender<SegmentEntry>) -> WalResult<()> {
+        let file = OpenOptions::new().read(true).open(path).await?;
+
+        let mut reader = BufReader::new(file);
+
+        // we read data_len first and then move our offset up till the len
+        // refresher: each entry in our file is <u64(data_len)><[u8;data_len]>
+        loop {
+            // read len first
+            let data_len_result = reader.read_u64_le().await;
+            let data_len = match data_len_result {
+                Ok(len) => len,
+                Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
+                Err(e) => return Err(format!("expected to read_u64 but couldn't {}", e).into()),
+            };
+
+            // make sure we have data for that len
+            let mut buffer = vec![0; data_len as usize];
+
+            // this is a critical error, we should be able to read data of len data_len
+            if let Err(e) = reader.read_exact(&mut buffer).await {
+                return Err(format!(
+                    "expected to read {}, but couldn't read_exact {}",
+                    data_len, e
+                )
+                .into());
+            }
+
+            // send each line
+            tx.send(SegmentEntry::DataEntry {
+                size: data_len,
+                data: Bytes::from(buffer),
+            })
+            .await
+            .expect("rx dropped while replaying");
+        }
+
+        Ok(())
+    }
+}
+
+/// Sort the filenames based on the file name. It is first sorted based on the timestamp and on
+/// conflict sorted on the file-index.
+fn sort_filenames(mut files: Vec<PathBuf>) -> Vec<PathBuf> {
+    files.sort_by(|a, b| {
+        let parse = |s: &str| {
+            let parts: Vec<&str> = s.split('_').collect();
+            let index = parts
+                .get(1)
+                .and_then(|s| s.parse::<u32>().ok())
+                .unwrap_or(0);
+            let ts_part = parts.get(2).unwrap_or(&"0");
+            let timestamp = ts_part
+                .split('.')
+                .next()
+                .unwrap_or("0")
+                .parse::<u64>()
+                .unwrap_or(0);
+            (timestamp, index)
+        };
+
+        let (ts_a, idx_a) = parse(a.to_str().expect("filename is valid"));
+        let (ts_b, idx_b) = parse(b.to_str().expect("filename is valid"));
+
+        // first sort on timestamp, if it matches, then sort on index
+        match ts_a.cmp(&ts_b) {
+            Ordering::Equal => idx_a.cmp(&idx_b),
+            v => v,
+        }
+    });
+
+    files
+}
+
+/// List all the files for the given [WalType].
+fn list_files(wal_type: &WalType, base_path: PathBuf) -> Vec<PathBuf> {
+    fs::read_dir(&base_path)
+        .unwrap_or_else(|_| panic!("directory {} to be present", base_path.display()))
+        .map(|entry| entry.expect("expect dirEntry to be good").path())
+        .filter(|path| path.is_file())
+        .filter(|path| {
+            path.file_name()
+                .expect("filename expected")
+                .to_str()
+                .expect("conversion should work")
+                .starts_with(wal_type.segment_prefix())
+        })
+        .filter(|path| path.extension().is_some_and(|ext| ext == "frozen"))
+        .collect::<Vec<_>>()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::reduce::wal::segment::WalType;
+    use std::fs::File;
+    use tempfile::tempdir;
+
+    #[test]
+    fn test_list_files() {
+        let temp_dir = tempdir().expect("Failed to create temp dir");
+        let base_path = temp_dir.path().to_path_buf();
+
+        // Create some test files
+        let _file1 = File::create(base_path.join("data_1.frozen")).expect("Failed to create file");
+        let _file2 = File::create(base_path.join("data_2.frozen")).expect("Failed to create file");
+        let _file3 = File::create(base_path.join("data_3.txt")).expect("Failed to create file");
+        let _file4 =
+            File::create(base_path.join("other_file.frozen")).expect("Failed to create file");
+
+        // Call the function
+        let result = list_files(&WalType::Data, base_path);
+
+        // Verify the result
+        let mut result_paths: Vec<String> = result
+            .iter()
+            .map(|path| path.file_name().unwrap().to_str().unwrap().to_string())
+            .collect();
+        result_paths.sort();
+
+        assert_eq!(result_paths, vec!["data_1.frozen", "data_2.frozen"]);
+    }
+
+    #[test]
+    fn test_sort_filenames() {
+        let temp_dir = tempdir().expect("Failed to create temp dir");
+        let base_path = temp_dir.path().to_path_buf();
+
+        // Create test files with the specified format
+        let _file1 = File::create(base_path.join("segment_000001_1000000001.wal"))
+            .expect("Failed to create file");
+        let _file4 = File::create(base_path.join("segment_000002_999999999.wal"))
+            .expect("Failed to create file");
+        let _file2 = File::create(base_path.join("segment_000001_1000000000.wal"))
+            .expect("Failed to create file");
+        let _file3 = File::create(base_path.join("segment_000001_999999999.wal"))
+            .expect("Failed to create file");
+
+        // Collect the file paths
+        let mut files: Vec<PathBuf> = fs::read_dir(&base_path)
+            .expect("Failed to read directory")
+            .map(|entry| entry.expect("Failed to read entry").path())
+            .collect();
+
+        // Sort the files using the function
+        files = sort_filenames(files);
+
+        // Verify the sorted order
+        let sorted_filenames: Vec<String> = files
+            .iter()
+            .map(|path| path.file_name().unwrap().to_str().unwrap().to_string())
+            .collect();
+
+        assert_eq!(
+            sorted_filenames,
+            vec![
+                "segment_000001_999999999.wal",
+                "segment_000002_999999999.wal",
+                "segment_000001_1000000000.wal",
+                "segment_000001_1000000001.wal",
+            ],
+            "File names are not sorted correctly"
+        );
+    }
+}
diff --git a/rust/numaflow-core/src/serving_store.rs b/rust/numaflow-core/src/serving_store.rs
index 1d2344e1..4e8206aa 100644
--- a/rust/numaflow-core/src/serving_store.rs
+++ b/rust/numaflow-core/src/serving_store.rs
@@ -1,3 +1,5 @@
+use bytes::Bytes;
+
 /// User defined serving store to store the serving responses.
 pub(crate) mod user_defined;
 
@@ -10,3 +12,17 @@ pub(crate) enum ServingStore {
     UserDefined(user_defined::UserDefinedStore),
     Nats(nats::NatsServingStore),
 }
+
+/// Entry in the serving store.
+#[derive(Clone, Debug)]
+pub(crate) struct StoreEntry {
+    /// Pod Hash is for filtering the stream by each request originating pod while listening for
+    /// "put" onto the KV store. This is used only on ISB KV store and this enables SSE (converse is
+    /// that SSE is not supported on user defined store).
+    pub(crate) pod_hash: String,
+    /// Unique ID Of the request to which the result belongs. There could be multiple results for
+    /// the same request.
+    pub(crate) id: String,
+    /// The result of the computation.
+    pub(crate) value: Bytes,
+}
diff --git a/rust/numaflow-core/src/serving_store/nats.rs b/rust/numaflow-core/src/serving_store/nats.rs
index 8c474fc4..f37db7df 100644
--- a/rust/numaflow-core/src/serving_store/nats.rs
+++ b/rust/numaflow-core/src/serving_store/nats.rs
@@ -1,9 +1,10 @@
+use crate::config::pipeline::NatsStoreConfig;
+use crate::serving_store::StoreEntry;
 use async_nats::jetstream::kv::Store;
 use async_nats::jetstream::Context;
-use bytes::Bytes;
 use chrono::Utc;
-
-use crate::config::pipeline::NatsStoreConfig;
+use tokio::task::JoinSet;
+use tracing::{info, trace};
 
 /// Nats serving store to store the serving responses.
 #[derive(Clone)]
@@ -18,29 +19,46 @@ impl NatsServingStore {
         nats_store_config: NatsStoreConfig,
     ) -> crate::Result<Self> {
         let store = js_context
-            .get_key_value(nats_store_config.name.as_str())
+            .get_key_value(nats_store_config.rs_store_name.as_str())
             .await
             .map_err(|e| crate::Error::Connection(format!("Failed to get kv store: {e:?}")))?;
         Ok(Self { store })
     }
 
-    /// Puts a datum into the serving store.
+    /// Puts multiple data items into the serving store concurrently.
     pub(crate) async fn put_datum(
         &mut self,
-        id: &str,
         origin: &str,
-        payload: Vec<u8>,
+        payloads: Vec<StoreEntry>,
     ) -> crate::Result<()> {
-        let id = format!(
-            "rs.{id}.{}.{}",
-            origin,
-            Utc::now().timestamp_nanos_opt().unwrap()
-        );
+        let mut jhset = JoinSet::new();
 
-        self.store
-            .put(id, Bytes::from(payload))
-            .await
-            .map_err(|e| crate::Error::Sink(format!("Failed to put datum: {e:?}")))?;
+        for payload in payloads {
+            let id = format!(
+                "rs.{}.{}.{}.{}",
+                payload.pod_hash,
+                payload.id,
+                origin,
+                Utc::now()
+                    .timestamp_nanos_opt()
+                    .unwrap_or(Utc::now().timestamp_micros())
+            );
+
+            trace!(?id, length = ?payload.value.len(), "Putting datum");
+
+            let store = self.store.clone();
+            jhset.spawn(async move {
+                store
+                    .put(id, payload.value)
+                    .await
+                    .map_err(|e| crate::Error::Sink(format!("Failed to put datum: {e:?}")))
+            });
+        }
+
+        while let Some(task) = jhset.join_next().await {
+            let result = task.map_err(|e| crate::Error::Sink(format!("Task failed: {e:?}")))?;
+            result?; // Propagate the first error, if any
+        }
         Ok(())
     }
 }
diff --git a/rust/numaflow-core/src/serving_store/user_defined.rs b/rust/numaflow-core/src/serving_store/user_defined.rs
index afd6b98a..e717ce51 100644
--- a/rust/numaflow-core/src/serving_store/user_defined.rs
+++ b/rust/numaflow-core/src/serving_store/user_defined.rs
@@ -1,8 +1,10 @@
 use numaflow_pb::clients::serving::serving_store_client::ServingStoreClient;
 use numaflow_pb::clients::serving::{Payload, PutRequest};
+use tokio::task::JoinSet;
 use tonic::transport::Channel;
 
 use crate::config::pipeline::UserDefinedStoreConfig;
+use crate::serving_store::StoreEntry;
 use crate::shared;
 
 /// User defined serving store to store the serving responses.
@@ -22,20 +24,34 @@ impl UserDefinedStore {
     /// Puts a datum into the serving store.
     pub(crate) async fn put_datum(
         &mut self,
-        id: &str,
         origin: &str,
-        payload: Vec<u8>,
+        payloads: Vec<StoreEntry>,
     ) -> crate::Result<()> {
-        let request = PutRequest {
-            id: id.to_string(),
-            payloads: vec![Payload {
-                origin: origin.to_string(),
-                value: payload.to_vec(),
-            }],
-        };
-        self.client.put(request).await.map_err(|e| {
-            crate::Error::Sink(format!("gRPC Put request failed on serving store: {e:?}"))
-        })?;
+        let mut jhset = JoinSet::new();
+
+        for payload in payloads {
+            let origin = origin.to_string();
+            let mut client = self.client.clone();
+
+            jhset.spawn(async move {
+                let request = PutRequest {
+                    id: payload.id,
+                    payloads: vec![Payload {
+                        origin: origin.clone(),
+                        value: payload.value.to_vec(),
+                    }],
+                };
+                client.put(request).await.map_err(|e| {
+                    crate::Error::Sink(format!("gRPC Put request failed on serving store: {e:?}"))
+                })
+            });
+        }
+
+        while let Some(task) = jhset.join_next().await {
+            let result = task.map_err(|e| crate::Error::Sink(format!("Task failed: {e:?}")))?;
+            result?; // Propagate the first error, if any
+        }
+
         Ok(())
     }
 
@@ -48,6 +64,7 @@ impl UserDefinedStore {
 mod tests {
     use super::*;
     use crate::config::pipeline::UserDefinedStoreConfig;
+    use bytes::Bytes;
     use numaflow::serving_store;
     use std::collections::HashMap;
     use std::sync::{Arc, Mutex};
@@ -104,8 +121,12 @@ mod tests {
         let mut store = UserDefinedStore::new(config).await.unwrap();
         let id = "test_id";
         let origin = "test_origin";
-        let payload = vec![1, 2, 3];
-        let result = store.put_datum(id, origin, payload).await;
+        let payloads = vec![StoreEntry {
+            pod_hash: "test_pod_hash".to_string(),
+            id: id.to_string(),
+            value: Bytes::from("test_value"),
+        }];
+        let result = store.put_datum(origin, payloads).await;
         assert!(result.is_ok());
 
         drop(store);
diff --git a/rust/numaflow-core/src/sink.rs b/rust/numaflow-core/src/sink.rs
index cc677f77..5af4c9a0 100644
--- a/rust/numaflow-core/src/sink.rs
+++ b/rust/numaflow-core/src/sink.rs
@@ -3,13 +3,13 @@
 //!
 //! [Sink]: https://numaflow.numaproj.io/user-guide/sinks/overview/
 
-use std::collections::HashMap;
-use std::time::Duration;
-
 use numaflow_pb::clients::serving::serving_store_client::ServingStoreClient;
 use numaflow_pb::clients::sink::sink_client::SinkClient;
 use numaflow_pb::clients::sink::sink_response;
 use numaflow_pb::clients::sink::Status::{Failure, Fallback, Serve, Success};
+use serving::{DEFAULT_ID_HEADER, DEFAULT_POD_HASH_KEY};
+use std::collections::HashMap;
+use std::time::Duration;
 use tokio::sync::mpsc::Receiver;
 use tokio::sync::{mpsc, oneshot};
 use tokio::task::JoinHandle;
@@ -30,7 +30,7 @@ use crate::metrics::{
     monovertex_metrics, mvtx_forward_metric_labels, pipeline_forward_metric_labels,
     pipeline_metrics,
 };
-use crate::serving_store::ServingStore;
+use crate::serving_store::{ServingStore, StoreEntry};
 use crate::tracker::TrackerHandle;
 use crate::Result;
 
@@ -51,8 +51,6 @@ mod serve;
 /// [User-Defined Sink]: https://numaflow.numaproj.io/user-guide/sinks/user-defined-sinks/
 mod user_defined;
 
-const NUMAFLOW_ID_HEADER: &str = "X-Numaflow-Id";
-
 /// Set of items to be implemented be a Numaflow Sink.
 ///
 /// [Sink]: https://numaflow.numaproj.io/user-guide/sinks/overview/
@@ -399,6 +397,7 @@ impl SinkWriter {
     ) -> Result<JoinHandle<Result<()>>> {
         let handle: JoinHandle<Result<()>> = tokio::spawn({
             async move {
+                info!(?self.batch_size, ?self.chunk_timeout, "Starting sink writer");
                 let chunk_stream =
                     messages_stream.chunks_timeout(self.batch_size, self.chunk_timeout);
 
@@ -857,23 +856,36 @@ impl SinkWriter {
             ));
         };
 
+        // convert Message to StoreEntry
+        let mut payloads = Vec::with_capacity(serving_msgs.len());
         for msg in serving_msgs {
-            let payload = msg.value.clone().into();
-            let id = msg.headers.get(NUMAFLOW_ID_HEADER).ok_or(Error::Sink(
-                "Missing numaflow id header. \
-                Numaflow Id header should be there in the serving message headers."
-                    .to_string(),
-            ))?;
+            let id = msg
+                .headers
+                .get(DEFAULT_ID_HEADER)
+                .ok_or(Error::Sink("Missing numaflow id header".to_string()))?
+                .clone();
+            let pod_hash = msg
+                .headers
+                .get(DEFAULT_POD_HASH_KEY)
+                .ok_or(Error::Sink("Missing pod replica header".to_string()))?
+                .clone();
+            payloads.push(StoreEntry {
+                pod_hash,
+                id,
+                value: msg.value.clone(),
+            });
+        }
 
-            match serving_store {
-                ServingStore::UserDefined(ud_store) => {
-                    ud_store.put_datum(id, get_vertex_name(), payload).await?;
-                }
-                ServingStore::Nats(nats_store) => {
-                    nats_store.put_datum(id, get_vertex_name(), payload).await?;
-                }
+        // push to corresponding store
+        match serving_store {
+            ServingStore::UserDefined(ud_store) => {
+                ud_store.put_datum(get_vertex_name(), payloads).await?;
+            }
+            ServingStore::Nats(nats_store) => {
+                nats_store.put_datum(get_vertex_name(), payloads).await?;
             }
         }
+
         Ok(())
     }
 
@@ -1266,7 +1278,7 @@ mod tests {
             NatsServingStore::new(
                 context.clone(),
                 NatsStoreConfig {
-                    name: serving_store.to_string(),
+                    rs_store_name: serving_store.to_string(),
                 },
             )
             .await
@@ -1310,7 +1322,8 @@ mod tests {
         let messages: Vec<Message> = (0..10)
             .map(|i| {
                 let mut headers = HashMap::new();
-                headers.insert(NUMAFLOW_ID_HEADER.to_string(), format!("id_{}", i));
+                headers.insert(DEFAULT_ID_HEADER.to_string(), format!("id_{}", i));
+                headers.insert(DEFAULT_POD_HASH_KEY.to_string(), "abcd".to_string());
                 Message {
                     typ: Default::default(),
                     keys: Arc::from(vec!["serve".to_string()]),
diff --git a/rust/numaflow-core/src/tracker.rs b/rust/numaflow-core/src/tracker.rs
index 5e8428c6..bc658644 100644
--- a/rust/numaflow-core/src/tracker.rs
+++ b/rust/numaflow-core/src/tracker.rs
@@ -16,7 +16,7 @@ use std::sync::Arc;
 
 use chrono::{DateTime, Utc};
 use serving::callback::CallbackHandler;
-use serving::DEFAULT_ID_HEADER;
+use serving::{DEFAULT_ID_HEADER, DEFAULT_POD_HASH_KEY};
 use tokio::sync::{mpsc, oneshot};
 use tracing::error;
 
@@ -87,6 +87,7 @@ struct Tracker {
 #[derive(Debug)]
 struct ServingCallbackInfo {
     id: String,
+    pod_hash: String,
     from_vertex: String,
     /// at the moment these are just tags.
     responses: Vec<Option<Vec<String>>>,
@@ -106,6 +107,16 @@ impl TryFrom<&Message> for ServingCallbackInfo {
             })?
             .to_owned();
 
+        let pod_hash = message
+            .headers
+            .get(DEFAULT_POD_HASH_KEY)
+            .ok_or_else(|| {
+                Error::Source(format!(
+                    "{DEFAULT_POD_HASH_KEY} is not found in message headers",
+                ))
+            })?
+            .to_owned();
+
         let from_vertex = message
             .metadata
             .as_ref()
@@ -115,6 +126,7 @@ impl TryFrom<&Message> for ServingCallbackInfo {
 
         Ok(ServingCallbackInfo {
             id: uuid,
+            pod_hash,
             from_vertex,
             responses: vec![None],
         })
@@ -329,6 +341,7 @@ impl Tracker {
         let result = callback_handler
             .callback(
                 callback_info.id,
+                callback_info.pod_hash,
                 callback_info.from_vertex,
                 callback_info.responses,
             )
@@ -516,7 +529,7 @@ mod tests {
         let callback_info: super::Result<ServingCallbackInfo> = TryFrom::try_from(&message);
         assert!(callback_info.is_err());
 
-        let headers = [(DEFAULT_ID_HEADER, "1234")]
+        let headers = [(DEFAULT_ID_HEADER, "1234"), (DEFAULT_POD_HASH_KEY, "abcd")]
             .into_iter()
             .map(|(k, v)| (k.to_string(), v.to_string()))
             .collect();
@@ -723,6 +736,7 @@ mod tests {
 
         let mut headers = HashMap::new();
         headers.insert(DEFAULT_ID_HEADER.to_string(), "1234".to_string());
+        headers.insert(DEFAULT_POD_HASH_KEY.to_string(), "abcd".to_string());
 
         let offset = Offset::String(StringOffset::new("offset1".to_string(), 0));
         let message = Message {
diff --git a/rust/numaflow-core/src/transformer.rs b/rust/numaflow-core/src/transformer.rs
index 84784497..bad0d058 100644
--- a/rust/numaflow-core/src/transformer.rs
+++ b/rust/numaflow-core/src/transformer.rs
@@ -208,6 +208,7 @@ impl Transformer {
             .iter()
             .filter(|message| message.dropped())
             .count();
+
         if dropped_messages_count > 0 {
             monovertex_metrics()
                 .transformer
diff --git a/rust/numaflow-pb/src/main.rs b/rust/numaflow-pb/src/main.rs
index 297465e0..ac4f9392 100644
--- a/rust/numaflow-pb/src/main.rs
+++ b/rust/numaflow-pb/src/main.rs
@@ -34,7 +34,7 @@ fn build_objects() {
         .compile_protos(
             &[
                 "proto/isb/message.proto",
-                "proto/reduce/wal/aligned/segment.proto",
+                "proto/reduce/wal/wal.proto",
                 "proto/watermark/watermark.proto",
             ],
             &["proto"],
diff --git a/rust/numaflow-pb/src/objects.rs b/rust/numaflow-pb/src/objects.rs
index 6ef7f862..7aecc40d 100644
--- a/rust/numaflow-pb/src/objects.rs
+++ b/rust/numaflow-pb/src/objects.rs
@@ -6,6 +6,6 @@ pub mod isb;
 #[rustfmt::skip]
 pub mod watermark;
 
-#[path = "objects/aligned_wal.rs"]
+#[path = "objects/wal.rs"]
 #[rustfmt::skip]
-pub mod aligned_wal;
+pub mod wal;
diff --git a/rust/numaflow-pb/src/objects/aligned_wal.rs b/rust/numaflow-pb/src/objects/aligned_wal.rs
deleted file mode 100644
index b2d016b6..00000000
--- a/rust/numaflow-pb/src/objects/aligned_wal.rs
+++ /dev/null
@@ -1,14 +0,0 @@
-// This file is @generated by prost-build.
-/// SegmentHeader for the Aligned WAL.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct SegmentHeader {
-    /// start_time of the window
-    #[prost(message, optional, tag = "1")]
-    pub start_time: ::core::option::Option<::prost_types::Timestamp>,
-    /// end_time = start_time + window
-    #[prost(message, optional, tag = "2")]
-    pub end_time: ::core::option::Option<::prost_types::Timestamp>,
-    /// this is currently set to 0 (in the future we can split large windows to multiple slots)
-    #[prost(string, tag = "3")]
-    pub slot: ::prost::alloc::string::String,
-}
diff --git a/rust/numaflow-pb/src/objects/wal.rs b/rust/numaflow-pb/src/objects/wal.rs
new file mode 100644
index 00000000..8c0a3474
--- /dev/null
+++ b/rust/numaflow-pb/src/objects/wal.rs
@@ -0,0 +1,20 @@
+// This file is @generated by prost-build.
+/// GCEvent is the event that is persisted in the WAL when a window is garbage collected
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct GcEvent {
+    /// start_time of the window
+    #[prost(message, optional, tag = "1")]
+    pub start_time: ::core::option::Option<::prost_types::Timestamp>,
+    /// end time of the window
+    #[prost(message, optional, tag = "2")]
+    pub end_time: ::core::option::Option<::prost_types::Timestamp>,
+    /// keys of the window, it will be empty for aligned windows
+    #[prost(string, repeated, tag = "3")]
+    pub keys: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
+}
+#[derive(Clone, Copy, PartialEq, ::prost::Message)]
+pub struct Footer {
+    /// the latest event time in the current Segment
+    #[prost(message, optional, tag = "1")]
+    pub latest_event_time: ::core::option::Option<::prost_types::Timestamp>,
+}
diff --git a/rust/serving/Cargo.toml b/rust/serving/Cargo.toml
index f489e1f9..db8d0d6c 100644
--- a/rust/serving/Cargo.toml
+++ b/rust/serving/Cargo.toml
@@ -36,12 +36,13 @@ prometheus-client = "0.23.1"
 thiserror = "2.0.12"
 reqwest = { workspace = true, features = ["rustls-tls", "json", "stream"] }
 http = "1.2.0"
-async-nats = "0.39.0"
+async-nats = "0.40.0"
 tokio-stream = "0.1.17"
 tonic = "0.13.0"
+tokio-util = "0.7.14"
 
 [dev-dependencies]
 reqwest = { workspace = true, features = ["json"] }
 rustls.workspace = true
 numaflow = { git = "https://github.com/numaproj/numaflow-rs.git", rev = "c6aa3a8d64a189357260f78131ef4692d387ba6c" }
-tempfile = "3.14.0"
+tempfile = "3.14.0"
\ No newline at end of file
diff --git a/rust/serving/README.md b/rust/serving/README.md
index e69de29b..5275c09c 100644
--- a/rust/serving/README.md
+++ b/rust/serving/README.md
@@ -0,0 +1,35 @@
+# Serving
+
+Serving is the HTTP endpoint for [Numaflow](https://numaflow.numaproj.io/) which is a distributed scalable
+general-purpose
+async processing platform. Unlike the normal HTTP endpoint, here we get the response from the DCG (Directed Compute
+Graph) directly.
+
+## High Level Data Movement
+
+```mermaid
+graph TD
+    subgraph Stores
+        CB[Callback Store<br>cb.*]
+        RS[Response Store<br>rs.*]
+        SS[Status Store<br>status.*]
+    end
+
+    subgraph Request Processing
+        REQ[Incoming Request] -->|Register| SS
+        SS -->|Track Status| PS[ProcessingStatus<br>- InProgress<br>- Completed<br>- Failed]
+        REQ --> RS
+        REQ -->|Fetch Response| RS
+    end
+
+    subgraph Orchestration
+        MG[Message Graph] -->|Watch Callbacks| CB
+        MG -->|Deregister| SS
+    end
+
+    subgraph Keys
+        K1[status.id] --> SS
+        K2[cb.pod-hash.id.vertex.timestamp] --> CB
+        K3[rs.id] --> RS
+    end
+```
\ No newline at end of file
diff --git a/rust/serving/src/app.rs b/rust/serving/src/app.rs
index 634159c3..0063659d 100644
--- a/rust/serving/src/app.rs
+++ b/rust/serving/src/app.rs
@@ -13,6 +13,7 @@ use hyper_util::client::legacy::connect::HttpConnector;
 use hyper_util::rt::TokioExecutor;
 use serde::Deserialize;
 use tokio::signal;
+use tokio_util::sync::CancellationToken;
 use tower::ServiceBuilder;
 use tower_http::classify::ServerErrorsFailureClass;
 use tower_http::timeout::TimeoutLayer;
@@ -21,7 +22,7 @@ use tracing::{info, info_span, Span};
 use uuid::Uuid;
 
 use self::{
-    direct_proxy::direct_proxy, jetstream_proxy::jetstream_proxy, message_path::get_message_path,
+    direct_proxy::direct_proxy, jetstream_proxy::serving_proxy, message_path::get_message_path,
 };
 use crate::app::store::cbstore::CallbackStore;
 use crate::app::store::datastore::DataStore;
@@ -46,6 +47,7 @@ pub(crate) mod tracker;
 pub(crate) async fn start_main_server<T, U>(
     app: AppState<T, U>,
     tls_config: RustlsConfig,
+    cancel_token: CancellationToken,
 ) -> crate::Result<()>
 where
     T: Clone + Send + Sync + DataStore + 'static,
@@ -60,6 +62,7 @@ where
     tokio::spawn(graceful_shutdown(
         handle.clone(),
         app.settings.drain_timeout_secs,
+        cancel_token.clone(),
     ));
 
     info!(?app_addr, "Starting application server");
@@ -72,6 +75,7 @@ where
         .await
         .map_err(|e| InitError(format!("Starting web server for metrics: {}", e)))?;
 
+    info!(?app_addr, "Server stopped");
     Ok(())
 }
 
@@ -164,9 +168,9 @@ where
     Ok(setup_app(app).await?.layer(layers))
 }
 
-// Gracefully shutdown the server on receiving SIGINT or SIGTERM
-// by sending a shutdown signal to the server using the handle.
-async fn graceful_shutdown(handle: Handle, duration_secs: u64) {
+/// Gracefully shutdown the server on receiving SIGINT or SIGTERM
+/// by sending a shutdown signal to the server using the handle.
+async fn graceful_shutdown(handle: Handle, duration_secs: u64, cancel_token: CancellationToken) {
     let ctrl_c = async {
         signal::ctrl_c()
             .await
@@ -185,10 +189,10 @@ async fn graceful_shutdown(handle: Handle, duration_secs: u64) {
         _ = terminate => {},
     }
 
-    info!("sending graceful shutdown signal");
-
+    info!("Got terminate signal, gracefully shutting down the server");
     // Signal the server to shut down using Handle.
     handle.graceful_shutdown(Some(Duration::from_secs(duration_secs)));
+    cancel_token.cancel();
 }
 
 const PUBLISH_ENDPOINTS: [&str; 3] = ["/v1/process/sync", "/v1/process/async", "/v1/process/fetch"];
@@ -293,7 +297,7 @@ async fn routes<
     app_state: AppState<T, U>,
 ) -> crate::Result<Router> {
     let state = app_state.orchestrator_state.clone();
-    let jetstream_proxy = jetstream_proxy(app_state.clone()).await?;
+    let jetstream_proxy = serving_proxy(app_state.clone()).await?;
 
     let message_path_handler = get_message_path(state);
     Ok(jetstream_proxy.merge(message_path_handler))
@@ -312,6 +316,7 @@ mod tests {
     use super::*;
     use crate::app::store::cbstore::memstore::InMemoryCallbackStore;
     use crate::app::store::datastore::inmemory::InMemoryDataStore;
+    use crate::app::store::status::StatusTracker;
     use crate::Settings;
 
     const PIPELINE_SPEC_ENCODED: &str = "eyJ2ZXJ0aWNlcyI6W3sibmFtZSI6ImluIiwic291cmNlIjp7InNlcnZpbmciOnsiYXV0aCI6bnVsbCwic2VydmljZSI6dHJ1ZSwibXNnSURIZWFkZXJLZXkiOiJYLU51bWFmbG93LUlkIiwic3RvcmUiOnsidXJsIjoicmVkaXM6Ly9yZWRpczo2Mzc5In19fSwiY29udGFpbmVyVGVtcGxhdGUiOnsicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIiLCJlbnYiOlt7Im5hbWUiOiJSVVNUX0xPRyIsInZhbHVlIjoiZGVidWcifV19LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6InBsYW5uZXIiLCJ1ZGYiOnsiY29udGFpbmVyIjp7ImltYWdlIjoiYXNjaWk6MC4xIiwiYXJncyI6WyJwbGFubmVyIl0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sImJ1aWx0aW4iOm51bGwsImdyb3VwQnkiOm51bGx9LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6InRpZ2VyIiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsidGlnZXIiXSwicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwiYnVpbHRpbiI6bnVsbCwiZ3JvdXBCeSI6bnVsbH0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiZG9nIiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsiZG9nIl0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sImJ1aWx0aW4iOm51bGwsImdyb3VwQnkiOm51bGx9LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6ImVsZXBoYW50IiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsiZWxlcGhhbnQiXSwicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwiYnVpbHRpbiI6bnVsbCwiZ3JvdXBCeSI6bnVsbH0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiYXNjaWlhcnQiLCJ1ZGYiOnsiY29udGFpbmVyIjp7ImltYWdlIjoiYXNjaWk6MC4xIiwiYXJncyI6WyJhc2NpaWFydCJdLCJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJidWlsdGluIjpudWxsLCJncm91cEJ5IjpudWxsfSwiY29udGFpbmVyVGVtcGxhdGUiOnsicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwic2NhbGUiOnsibWluIjoxfSwidXBkYXRlU3RyYXRlZ3kiOnsidHlwZSI6IlJvbGxpbmdVcGRhdGUiLCJyb2xsaW5nVXBkYXRlIjp7Im1heFVuYXZhaWxhYmxlIjoiMjUlIn19fSx7Im5hbWUiOiJzZXJ2ZS1zaW5rIiwic2luayI6eyJ1ZHNpbmsiOnsiY29udGFpbmVyIjp7ImltYWdlIjoic2VydmVzaW5rOjAuMSIsImVudiI6W3sibmFtZSI6Ik5VTUFGTE9XX0NBTExCQUNLX1VSTF9LRVkiLCJ2YWx1ZSI6IlgtTnVtYWZsb3ctQ2FsbGJhY2stVXJsIn0seyJuYW1lIjoiTlVNQUZMT1dfTVNHX0lEX0hFQURFUl9LRVkiLCJ2YWx1ZSI6IlgtTnVtYWZsb3ctSWQifV0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn19LCJyZXRyeVN0cmF0ZWd5Ijp7fX0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiZXJyb3Itc2luayIsInNpbmsiOnsidWRzaW5rIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6InNlcnZlc2luazowLjEiLCJlbnYiOlt7Im5hbWUiOiJOVU1BRkxPV19DQUxMQkFDS19VUkxfS0VZIiwidmFsdWUiOiJYLU51bWFmbG93LUNhbGxiYWNrLVVybCJ9LHsibmFtZSI6Ik5VTUFGTE9XX01TR19JRF9IRUFERVJfS0VZIiwidmFsdWUiOiJYLU51bWFmbG93LUlkIn1dLCJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9fSwicmV0cnlTdHJhdGVneSI6e319LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19XSwiZWRnZXMiOlt7ImZyb20iOiJpbiIsInRvIjoicGxhbm5lciIsImNvbmRpdGlvbnMiOm51bGx9LHsiZnJvbSI6InBsYW5uZXIiLCJ0byI6ImFzY2lpYXJ0IiwiY29uZGl0aW9ucyI6eyJ0YWdzIjp7Im9wZXJhdG9yIjoib3IiLCJ2YWx1ZXMiOlsiYXNjaWlhcnQiXX19fSx7ImZyb20iOiJwbGFubmVyIiwidG8iOiJ0aWdlciIsImNvbmRpdGlvbnMiOnsidGFncyI6eyJvcGVyYXRvciI6Im9yIiwidmFsdWVzIjpbInRpZ2VyIl19fX0seyJmcm9tIjoicGxhbm5lciIsInRvIjoiZG9nIiwiY29uZGl0aW9ucyI6eyJ0YWdzIjp7Im9wZXJhdG9yIjoib3IiLCJ2YWx1ZXMiOlsiZG9nIl19fX0seyJmcm9tIjoicGxhbm5lciIsInRvIjoiZWxlcGhhbnQiLCJjb25kaXRpb25zIjp7InRhZ3MiOnsib3BlcmF0b3IiOiJvciIsInZhbHVlcyI6WyJlbGVwaGFudCJdfX19LHsiZnJvbSI6InRpZ2VyIiwidG8iOiJzZXJ2ZS1zaW5rIiwiY29uZGl0aW9ucyI6bnVsbH0seyJmcm9tIjoiZG9nIiwidG8iOiJzZXJ2ZS1zaW5rIiwiY29uZGl0aW9ucyI6bnVsbH0seyJmcm9tIjoiZWxlcGhhbnQiLCJ0byI6InNlcnZlLXNpbmsiLCJjb25kaXRpb25zIjpudWxsfSx7ImZyb20iOiJhc2NpaWFydCIsInRvIjoic2VydmUtc2luayIsImNvbmRpdGlvbnMiOm51bGx9LHsiZnJvbSI6InBsYW5uZXIiLCJ0byI6ImVycm9yLXNpbmsiLCJjb25kaXRpb25zIjp7InRhZ3MiOnsib3BlcmF0b3IiOiJvciIsInZhbHVlcyI6WyJlcnJvciJdfX19XSwibGlmZWN5Y2xlIjp7fSwid2F0ZXJtYXJrIjp7fX0=";
@@ -326,9 +331,27 @@ mod tests {
 
         let datum_store = InMemoryDataStore::new(None);
         let callback_store = InMemoryCallbackStore::new(None);
+        let store_name = "test_health_check_endpoints";
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let _ = context.delete_key_value(store_name).await;
+
+        let _ = context
+            .create_key_value(jetstream::kv::Config {
+                bucket: store_name.to_string(),
+                history: 5,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
 
+        let status_tracker = StatusTracker::new(context.clone(), store_name, "0".to_string(), None)
+            .await
+            .unwrap();
         let msg_graph = MessageGraph::from_pipeline(&settings.pipeline_spec)?;
-        let callback_state = CallbackState::new(msg_graph, datum_store, callback_store).await?;
+        let callback_state =
+            CallbackState::new("0", msg_graph, datum_store, callback_store, status_tracker).await?;
 
         let nats_connection = async_nats::connect("localhost:4222")
             .await
@@ -338,6 +361,7 @@ mod tests {
             js_context,
             settings,
             orchestrator_state: callback_state,
+            cancellation_token: CancellationToken::new(),
         };
 
         let router = setup_app(app).await.unwrap();
@@ -366,10 +390,29 @@ mod tests {
 
         let datum_store = InMemoryDataStore::new(None);
         let callback_store = InMemoryCallbackStore::new(None);
+        let store_name = "test_auth_middleware";
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let _ = context.delete_key_value(store_name).await;
+
+        let _ = context
+            .create_key_value(jetstream::kv::Config {
+                bucket: store_name.to_string(),
+                history: 5,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let status_tracker = StatusTracker::new(context.clone(), store_name, "0".to_string(), None)
+            .await
+            .unwrap();
 
         let pipeline_spec = PIPELINE_SPEC_ENCODED.parse().unwrap();
         let msg_graph = MessageGraph::from_pipeline(&pipeline_spec)?;
-        let callback_state = CallbackState::new(msg_graph, datum_store, callback_store).await?;
+        let callback_state =
+            CallbackState::new("0", msg_graph, datum_store, callback_store, status_tracker).await?;
 
         let nats_connection = async_nats::connect("localhost:4222")
             .await
@@ -380,6 +423,7 @@ mod tests {
             js_context,
             settings: Arc::new(settings),
             orchestrator_state: callback_state,
+            cancellation_token: CancellationToken::new(),
         };
 
         let router = router_with_auth(app_state).await.unwrap();
diff --git a/rust/serving/src/app/direct_proxy.rs b/rust/serving/src/app/direct_proxy.rs
index aaa7dd32..c5e8cb97 100644
--- a/rust/serving/src/app/direct_proxy.rs
+++ b/rust/serving/src/app/direct_proxy.rs
@@ -1,3 +1,5 @@
+//! A simple direct reverse-proxy that forwards the request returns the response as-is from the container.
+
 use axum::{
     body::Body,
     extract::{Request, State},
@@ -14,8 +16,6 @@ use crate::app::response::ApiError;
 
 pub(crate) type Client = hyper_util::client::legacy::Client<HttpConnector, Body>;
 
-/// A simple direct reverse-proxy that forwards the request returns the response as-is from the container.
-
 // TODO
 // - [ ] fallbacks for better errors https://docs.rs/axum/latest/axum/struct.Router.html#fallbacks
 // - [ ] should wild captures `/*upstream` be captured using https://docs.rs/axum/latest/axum/extract/struct.Path.html?
diff --git a/rust/serving/src/app/jetstream_proxy.rs b/rust/serving/src/app/jetstream_proxy.rs
index 479a5104..e07e25a4 100644
--- a/rust/serving/src/app/jetstream_proxy.rs
+++ b/rust/serving/src/app/jetstream_proxy.rs
@@ -1,10 +1,13 @@
+//! Serving proxy receives the serving request and writes the request to the ISB so the pipeline can
+//! process the request.
+
 use std::{str::FromStr, sync::Arc};
 
 use super::{orchestrator, store::datastore::DataStore, AppState};
 use super::{FetchQueryParams, Tid};
 use crate::app::response::{ApiError, ServeResponse};
 use crate::app::store::cbstore::CallbackStore;
-use crate::app::store::datastore::Error as StoreError;
+use crate::config::{RequestType, DEFAULT_POD_HASH_KEY};
 use crate::metrics::serving_metrics;
 use crate::Error;
 use async_nats::jetstream::Context;
@@ -22,7 +25,8 @@ use axum::{
 use serde_json::json;
 use tokio::time::Instant;
 use tokio_stream::{Stream, StreamExt};
-use tracing::{error, Instrument};
+use tokio_util::sync::CancellationToken;
+use tracing::{debug, error, Instrument};
 
 const NUMAFLOW_RESP_ARRAY_LEN: &str = "Numaflow-Array-Len";
 const NUMAFLOW_RESP_ARRAY_IDX_LEN: &str = "Numaflow-Array-Index-Len";
@@ -31,10 +35,13 @@ struct ProxyState<T, U> {
     js_context: Context,
     stream: String,
     tid_header: String,
+    pod_hash: String,
     orchestrator: orchestrator::OrchestratorState<T, U>,
+    cancellation_token: CancellationToken,
 }
 
-pub(crate) async fn jetstream_proxy<
+/// Expose a router for the serving proxy and its endpoints.
+pub(crate) async fn serving_proxy<
     T: Clone + Send + Sync + DataStore + 'static,
     U: Clone + Send + Sync + CallbackStore + 'static,
 >(
@@ -44,7 +51,9 @@ pub(crate) async fn jetstream_proxy<
         js_context: state.js_context.clone(),
         stream: state.settings.js_message_stream.clone(),
         tid_header: state.settings.tid_header.clone(),
+        pod_hash: state.settings.pod_hash.clone(),
         orchestrator: state.orchestrator_state.clone(),
+        cancellation_token: state.cancellation_token.clone(),
     });
 
     let router = Router::new()
@@ -64,7 +73,12 @@ async fn fetch<
     Query(FetchQueryParams { id }): Query<FetchQueryParams>,
 ) -> Response {
     let datum_retrieve_start = Instant::now();
-    let pipeline_result = match proxy_state.orchestrator.clone().retrieve_saved(&id).await {
+    let pipeline_result = match proxy_state
+        .orchestrator
+        .clone()
+        .retrieve_saved(&id, RequestType::Async)
+        .await
+    {
         Ok(result) => {
             serving_metrics()
                 .datum_retrive_duration
@@ -138,37 +152,34 @@ where
         );
     }
     msg_headers.insert(proxy_state.tid_header.clone(), id.clone());
+    msg_headers.insert(DEFAULT_POD_HASH_KEY, proxy_state.pod_hash.clone());
 
     let mut orchestrator_state = proxy_state.orchestrator.clone();
-    let response_stream = orchestrator_state.stream_response(&id).await.unwrap();
-
-    let save_start = Instant::now();
-    proxy_state
-        .js_context
-        .publish_with_headers(proxy_state.stream.clone(), msg_headers, body)
+    let response_stream = orchestrator_state
+        .stream_response(&id, RequestType::Sse)
         .await
-        .inspect_err(|err| {
-            tracing::error!(
-                ?err,
-                id,
-                stream = proxy_state.stream,
-                "Saving message to Jetstream"
-            )
-        })
-        .map_err(|_| ApiError::BadGateway("Failed to write message to Jetstream".to_string()))?
-        .await
-        .inspect_err(|err| {
-            tracing::error!(
-                ?err,
-                id,
-                stream = proxy_state.stream,
-                "Waiting for message save confirmation from Jetstream"
-            )
-        })
-        .map_err(|_| ApiError::BadGateway("Failed to write message to Jetstream".to_string()))?;
-    serving_metrics()
-        .payload_save_duration
-        .observe(save_start.elapsed().as_micros() as f64);
+        .map_err(|e| ApiError::InternalServerError(e.to_string()))?;
+
+    // if the write to jetstream fails before returning we have to discard the state that we have
+    // stored in the tracker because it's not published to the pipeline otherwise when the users
+    // retry again we will show them its duplicate.
+    if let Err(err) = publish_to_jetstream(
+        proxy_state.js_context.clone(),
+        &proxy_state.stream,
+        &id,
+        msg_headers,
+        body,
+    )
+    .await
+    {
+        proxy_state
+            .orchestrator
+            .clone()
+            .discard(&id)
+            .await
+            .expect("Failed to discard");
+        return Err(err);
+    }
 
     let stream = response_stream
         .map(|response| Ok(Event::default().data(String::from_utf8_lossy(&response))));
@@ -176,6 +187,54 @@ where
     Ok(Sse::new(stream))
 }
 
+/// publishes the message to jetstream, if we are not able to publish we will return 429(too many
+/// requests) error so that the users can retry.
+async fn publish_to_jetstream(
+    js_context: Context,
+    stream: &str,
+    id: &str,
+    msg_headers: async_nats::HeaderMap,
+    body: Bytes,
+) -> Result<(), ApiError> {
+    let save_start = Instant::now();
+
+    // Publish the message to JetStream
+    let publish_result = js_context
+        .publish_with_headers(stream.to_string(), msg_headers, body)
+        .await
+        .inspect_err(|err| error!(?err, id, stream, "Saving message to JetStream"));
+
+    if let Err(err) = publish_result {
+        error!(?err, id, "Publish to JetStream failed, invoking discard");
+        return Err(ApiError::TooManyRequests(
+            "Failed to write message to JetStream".to_string(),
+        ));
+    }
+
+    // Wait for confirmation from JetStream
+    let confirmation_result = publish_result.unwrap().await.inspect_err(|err| {
+        error!(
+            ?err,
+            id, stream, "Waiting for message save confirmation from JetStream"
+        )
+    });
+
+    if let Err(err) = confirmation_result {
+        error!(
+            ?err,
+            id, "Message save confirmation failed, invoking discard"
+        );
+        return Err(ApiError::TooManyRequests(
+            "Failed to write message to JetStream".to_string(),
+        ));
+    }
+
+    serving_metrics()
+        .payload_save_duration
+        .observe(save_start.elapsed().as_micros() as f64);
+    Ok(())
+}
+
 async fn sync_publish<
     T: Send + Sync + Clone + DataStore + 'static,
     U: Send + Sync + Clone + CallbackStore + 'static,
@@ -193,19 +252,20 @@ async fn sync_publish<
         );
     }
     msg_headers.insert(proxy_state.tid_header.clone(), id.clone());
+    msg_headers.insert(DEFAULT_POD_HASH_KEY, proxy_state.pod_hash.clone());
 
     let processing_start = Instant::now();
     // Register the ID in the callback proxy state
     let notify = match proxy_state
         .orchestrator
         .clone()
-        .process_request(id.as_str())
+        .process_request(id.as_str(), RequestType::Sync)
         .await
     {
         Ok(result) => result,
         Err(e) => {
             error!(error = ?e, "Registering request in data store");
-            if let StoreError::DuplicateRequest(id) = e {
+            if let Error::Duplicate(id) = e {
                 return Err(ApiError::Conflict(format!(
                     "Specified request id {id} already exists in the store"
                 )));
@@ -216,39 +276,36 @@ async fn sync_publish<
         }
     };
 
-    let save_start = Instant::now();
-    proxy_state
-        .js_context
-        .publish_with_headers(proxy_state.stream.clone(), msg_headers, body)
-        .await
-        .inspect_err(|err| {
-            tracing::error!(
-                ?err,
-                id,
-                stream = proxy_state.stream,
-                "Saving message to Jetstream"
-            )
-        })
-        .map_err(|_| ApiError::BadGateway("Failed to write message to Jetstream".to_string()))?
-        .await
-        .inspect_err(|err| {
-            tracing::error!(
-                ?err,
-                id,
-                stream = proxy_state.stream,
-                "Waiting for message save confirmation from Jetstream"
-            )
-        })
-        .map_err(|_| ApiError::BadGateway("Failed to write message to Jetstream".to_string()))?;
-    serving_metrics()
-        .payload_save_duration
-        .observe(save_start.elapsed().as_micros() as f64);
+    // if the write to jetstream fails before returning we have to discard the state that we have
+    // stored in the tracker because it's not published to the pipeline otherwise when the users
+    // retry again we will show them its duplicate.
+    if let Err(err) = publish_to_jetstream(
+        proxy_state.js_context.clone(),
+        &proxy_state.stream,
+        &id,
+        msg_headers,
+        body,
+    )
+    .await
+    {
+        proxy_state
+            .orchestrator
+            .clone()
+            .discard(&id)
+            .await
+            .expect("Failed to discard");
+        return Err(err);
+    }
 
     let processing_result = match notify.await {
         Ok(processing_result) => {
             serving_metrics()
                 .processing_time
                 .observe(processing_start.elapsed().as_micros() as f64);
+            debug!(
+                "Time taken to process message={:?}",
+                processing_start.elapsed().as_millis()
+            );
             processing_result
         }
         Err(e) => {
@@ -260,7 +317,7 @@ async fn sync_publish<
     };
 
     if let Err(err) = processing_result {
-        tracing::error!(
+        error!(
             ?err,
             "The request processing task failed. Returning error to the client"
         );
@@ -270,11 +327,20 @@ async fn sync_publish<
     }
 
     let datum_retrieve_start = Instant::now();
-    let result = match proxy_state.orchestrator.clone().retrieve_saved(&id).await {
+    let result = match proxy_state
+        .orchestrator
+        .clone()
+        .retrieve_saved(&id, RequestType::Sync)
+        .await
+    {
         Ok(result) => {
             serving_metrics()
                 .datum_retrive_duration
                 .observe(datum_retrieve_start.elapsed().as_micros() as f64);
+            debug!(
+                "Time taken to retrieve message from store={:?}",
+                datum_retrieve_start.elapsed().as_millis()
+            );
             result
         }
         Err(e) => {
@@ -313,7 +379,6 @@ async fn sync_publish<
 
     // concatenate as single result, should use header values to decompose based on the len
     let body = result.concat();
-
     Ok((header_map, body).into_response())
 }
 
@@ -334,19 +399,20 @@ async fn async_publish<
         );
     }
     msg_headers.insert(proxy_state.tid_header.clone(), id.clone());
+    msg_headers.insert(DEFAULT_POD_HASH_KEY, proxy_state.pod_hash.clone());
 
     let processing_start = Instant::now();
     // Register request in Redis
     let notify = match proxy_state
         .orchestrator
         .clone()
-        .process_request(id.as_str())
+        .process_request(id.as_str(), RequestType::Async)
         .await
     {
         Ok(result) => result,
         Err(e) => {
             error!(error = ?e, "Registering request in data store");
-            if let StoreError::DuplicateRequest(id) = e {
+            if let Error::Duplicate(id) = e {
                 return Err(ApiError::Conflict(format!(
                     "Specified request id {id} already exists in the store"
                 )));
@@ -356,64 +422,62 @@ async fn async_publish<
             ));
         }
     };
-
+    debug!(
+        "Time taken to register message={:?}",
+        processing_start.elapsed().as_millis()
+    );
     {
         let span = tracing::Span::current();
         let mut orchestrator = proxy_state.orchestrator.clone();
         let id = id.clone();
+        let cln_token = proxy_state.cancellation_token.clone();
         // A send operation happens at the sender end of the notify channel when all callbacks are received.
         // We keep the receiver alive to avoid send failure.
         let wait_for_callbacks = async move {
-            match notify.await {
-                Ok(val) => {
+            tokio::select! {
+                val = notify => {
+                    let val = val.expect("Failed to get callback notify that its done");
                     match val {
                         Ok(_) => {}
                         Err(e) => {
-                            _ = orchestrator
-                                .mark_as_failed(id.as_str(), e.to_string().as_str())
-                                .await
-                                .inspect_err(|e| error!(error=?e, "Failed to mark as failed"));
+                            error!(error = ?e, "Error while waiting for callback notify to exit");
                         }
                     }
                     serving_metrics()
                         .processing_time
                         .observe(processing_start.elapsed().as_micros() as f64);
                 }
-                Err(e) => {
-                    error!(error = ?e, "Waiting for the pipeline output");
+                _ = cln_token.cancelled() => {
+                        _ = orchestrator
+                                .mark_as_failed(id.as_str(), "Cancelled")
+                                .await
+                                .inspect_err(|e| error!(error=?e, "Failed to mark as failed"));
                 }
-            };
+            }
         };
         tokio::spawn(wait_for_callbacks.instrument(span));
     }
 
-    let save_start = Instant::now();
-    proxy_state
-        .js_context
-        .publish_with_headers(proxy_state.stream.clone(), msg_headers, body)
-        .await
-        .inspect_err(|err| {
-            tracing::error!(
-                ?err,
-                id,
-                stream = proxy_state.stream,
-                "Saving message to Jetstream"
-            )
-        })
-        .map_err(|_| ApiError::BadGateway("Failed to write message to Jetstream".to_string()))?
-        .await
-        .inspect_err(|err| {
-            tracing::error!(
-                ?err,
-                id,
-                stream = proxy_state.stream,
-                "Waiting for message save confirmation from Jetstream"
-            )
-        })
-        .map_err(|_| ApiError::BadGateway("Failed to write message to Jetstream".to_string()))?;
-    serving_metrics()
-        .payload_save_duration
-        .observe(save_start.elapsed().as_micros() as f64);
+    // if the write to jetstream fails before returning we have to discard the state that we have
+    // stored in the tracker because it's not published to the pipeline otherwise when the users
+    // retry again we will show them its duplicate.
+    if let Err(err) = publish_to_jetstream(
+        proxy_state.js_context.clone(),
+        &proxy_state.stream,
+        &id,
+        msg_headers,
+        body,
+    )
+    .await
+    {
+        proxy_state
+            .orchestrator
+            .clone()
+            .discard(&id)
+            .await
+            .expect("Failed to discard");
+        return Err(err);
+    }
 
     Ok(Json(ServeResponse::new(
         "Successfully published message".to_string(),
@@ -427,8 +491,9 @@ mod tests {
     use super::*;
     use crate::app::orchestrator::OrchestratorState;
     use crate::app::router_with_auth;
-    use crate::app::store::cbstore::jetstreamstore::JetStreamCallbackStore;
+    use crate::app::store::cbstore::jetstream_store::JetStreamCallbackStore;
     use crate::app::store::datastore::jetstream::JetStreamDataStore;
+    use crate::app::store::status::StatusTracker;
     use crate::app::tracker::MessageGraph;
     use crate::callback::{Callback, Response};
     use crate::config::DEFAULT_ID_HEADER;
@@ -440,7 +505,9 @@ mod tests {
     use axum::http::header::{CONTENT_LENGTH, CONTENT_TYPE};
     use serde_json::{json, Value};
     use std::sync::Arc;
+    use std::time::Instant;
     use tokio::net::TcpListener;
+    use tokio_util::sync::CancellationToken;
     use tower::ServiceExt;
 
     const PIPELINE_SPEC_ENCODED: &str = "eyJ2ZXJ0aWNlcyI6W3sibmFtZSI6ImluIiwic291cmNlIjp7InNlcnZpbmciOnsiYXV0aCI6bnVsbCwic2VydmljZSI6dHJ1ZSwibXNnSURIZWFkZXJLZXkiOiJYLU51bWFmbG93LUlkIiwic3RvcmUiOnsidXJsIjoicmVkaXM6Ly9yZWRpczo2Mzc5In19fSwiY29udGFpbmVyVGVtcGxhdGUiOnsicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIiLCJlbnYiOlt7Im5hbWUiOiJSVVNUX0xPRyIsInZhbHVlIjoiZGVidWcifV19LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6InBsYW5uZXIiLCJ1ZGYiOnsiY29udGFpbmVyIjp7ImltYWdlIjoiYXNjaWk6MC4xIiwiYXJncyI6WyJwbGFubmVyIl0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sImJ1aWx0aW4iOm51bGwsImdyb3VwQnkiOm51bGx9LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6InRpZ2VyIiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsidGlnZXIiXSwicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwiYnVpbHRpbiI6bnVsbCwiZ3JvdXBCeSI6bnVsbH0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiZG9nIiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsiZG9nIl0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sImJ1aWx0aW4iOm51bGwsImdyb3VwQnkiOm51bGx9LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6ImVsZXBoYW50IiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsiZWxlcGhhbnQiXSwicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwiYnVpbHRpbiI6bnVsbCwiZ3JvdXBCeSI6bnVsbH0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiYXNjaWlhcnQiLCJ1ZGYiOnsiY29udGFpbmVyIjp7ImltYWdlIjoiYXNjaWk6MC4xIiwiYXJncyI6WyJhc2NpaWFydCJdLCJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJidWlsdGluIjpudWxsLCJncm91cEJ5IjpudWxsfSwiY29udGFpbmVyVGVtcGxhdGUiOnsicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwic2NhbGUiOnsibWluIjoxfSwidXBkYXRlU3RyYXRlZ3kiOnsidHlwZSI6IlJvbGxpbmdVcGRhdGUiLCJyb2xsaW5nVXBkYXRlIjp7Im1heFVuYXZhaWxhYmxlIjoiMjUlIn19fSx7Im5hbWUiOiJzZXJ2ZS1zaW5rIiwic2luayI6eyJ1ZHNpbmsiOnsiY29udGFpbmVyIjp7ImltYWdlIjoic2VydmVzaW5rOjAuMSIsImVudiI6W3sibmFtZSI6Ik5VTUFGTE9XX0NBTExCQUNLX1VSTF9LRVkiLCJ2YWx1ZSI6IlgtTnVtYWZsb3ctQ2FsbGJhY2stVXJsIn0seyJuYW1lIjoiTlVNQUZMT1dfTVNHX0lEX0hFQURFUl9LRVkiLCJ2YWx1ZSI6IlgtTnVtYWZsb3ctSWQifV0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn19LCJyZXRyeVN0cmF0ZWd5Ijp7fX0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiZXJyb3Itc2luayIsInNpbmsiOnsidWRzaW5rIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6InNlcnZlc2luazowLjEiLCJlbnYiOlt7Im5hbWUiOiJOVU1BRkxPV19DQUxMQkFDS19VUkxfS0VZIiwidmFsdWUiOiJYLU51bWFmbG93LUNhbGxiYWNrLVVybCJ9LHsibmFtZSI6Ik5VTUFGTE9XX01TR19JRF9IRUFERVJfS0VZIiwidmFsdWUiOiJYLU51bWFmbG93LUlkIn1dLCJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9fSwicmV0cnlTdHJhdGVneSI6e319LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19XSwiZWRnZXMiOlt7ImZyb20iOiJpbiIsInRvIjoicGxhbm5lciIsImNvbmRpdGlvbnMiOm51bGx9LHsiZnJvbSI6InBsYW5uZXIiLCJ0byI6ImFzY2lpYXJ0IiwiY29uZGl0aW9ucyI6eyJ0YWdzIjp7Im9wZXJhdG9yIjoib3IiLCJ2YWx1ZXMiOlsiYXNjaWlhcnQiXX19fSx7ImZyb20iOiJwbGFubmVyIiwidG8iOiJ0aWdlciIsImNvbmRpdGlvbnMiOnsidGFncyI6eyJvcGVyYXRvciI6Im9yIiwidmFsdWVzIjpbInRpZ2VyIl19fX0seyJmcm9tIjoicGxhbm5lciIsInRvIjoiZG9nIiwiY29uZGl0aW9ucyI6eyJ0YWdzIjp7Im9wZXJhdG9yIjoib3IiLCJ2YWx1ZXMiOlsiZG9nIl19fX0seyJmcm9tIjoicGxhbm5lciIsInRvIjoiZWxlcGhhbnQiLCJjb25kaXRpb25zIjp7InRhZ3MiOnsib3BlcmF0b3IiOiJvciIsInZhbHVlcyI6WyJlbGVwaGFudCJdfX19LHsiZnJvbSI6InRpZ2VyIiwidG8iOiJzZXJ2ZS1zaW5rIiwiY29uZGl0aW9ucyI6bnVsbH0seyJmcm9tIjoiZG9nIiwidG8iOiJzZXJ2ZS1zaW5rIiwiY29uZGl0aW9ucyI6bnVsbH0seyJmcm9tIjoiZWxlcGhhbnQiLCJ0byI6InNlcnZlLXNpbmsiLCJjb25kaXRpb25zIjpudWxsfSx7ImZyb20iOiJhc2NpaWFydCIsInRvIjoic2VydmUtc2luayIsImNvbmRpdGlvbnMiOm51bGx9LHsiZnJvbSI6InBsYW5uZXIiLCJ0byI6ImVycm9yLXNpbmsiLCJjb25kaXRpb25zIjp7InRhZ3MiOnsib3BlcmF0b3IiOiJvciIsInZhbHVlcyI6WyJlcnJvciJdfX19XSwibGlmZWN5Y2xlIjp7fSwid2F0ZXJtYXJrIjp7fX0=";
@@ -451,6 +518,7 @@ mod tests {
         const ID_HEADER: &str = "X-Numaflow-Id";
         const ID_VALUE: &str = "foobar";
 
+        let pod_hash = "0";
         let store_name = "test_async_publish";
         let message_stream_name = "test_async_publish_messages";
         let js_url = "localhost:4222";
@@ -458,6 +526,8 @@ mod tests {
         let context = jetstream::new(client);
         let _ = context.delete_key_value(store_name).await;
         let _ = context.delete_stream(message_stream_name).await;
+        let cancel_token = CancellationToken::new();
+        let _drop_guard = cancel_token.clone().drop_guard();
 
         let _ = context
             .create_key_value(jetstream::kv::Config {
@@ -469,7 +539,7 @@ mod tests {
             .unwrap();
 
         context
-            .create_stream(async_nats::jetstream::stream::Config {
+            .create_stream(jetstream::stream::Config {
                 name: message_stream_name.into(),
                 ..Default::default()
             })
@@ -483,26 +553,48 @@ mod tests {
             ..Default::default()
         };
 
-        let callback_store = JetStreamCallbackStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create callback store");
+        let callback_store = JetStreamCallbackStore::new(
+            context.clone(),
+            pod_hash,
+            store_name,
+            cancel_token.clone(),
+        )
+        .await
+        .expect("Failed to create callback store");
 
-        let datum_store = JetStreamDataStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create datum store");
+        let datum_store =
+            JetStreamDataStore::new(context.clone(), store_name, pod_hash, cancel_token.clone())
+                .await
+                .expect("Failed to create datum store");
+
+        let status_tracker = StatusTracker::new(
+            context.clone(),
+            store_name,
+            pod_hash.to_string(),
+            Some(store_name.to_string()),
+        )
+        .await
+        .unwrap();
 
         let pipeline_spec = PIPELINE_SPEC_ENCODED.parse().unwrap();
         let msg_graph = MessageGraph::from_pipeline(&pipeline_spec)?;
-        let orchestrator_state =
-            OrchestratorState::new(msg_graph, datum_store, callback_store).await?;
+        let orchestrator_state = OrchestratorState::new(
+            pod_hash,
+            msg_graph,
+            datum_store,
+            callback_store,
+            status_tracker,
+        )
+        .await?;
 
         let app_state = AppState {
             js_context: context,
             settings: Arc::new(settings),
             orchestrator_state,
+            cancellation_token: cancel_token.clone(),
         };
 
-        let app = jetstream_proxy(app_state).await?;
+        let app = serving_proxy(app_state).await?;
         let mut req = Request::builder()
             .method("POST")
             .uri("/async")
@@ -574,8 +666,9 @@ mod tests {
     #[cfg(feature = "nats-tests")]
     #[tokio::test]
     async fn test_sync_publish() {
-        const ID_HEADER: &str = "X-Numaflow-ID";
-        const ID_VALUE: &str = "foobar";
+        const ID_HEADER: &str = "X-Numaflow-Id";
+        const ID_VALUE: &str = "publish";
+        const POD_HASH: &str = "0";
 
         let store_name = "test_sync_publish";
         let messages_stream_name = "test_sync_publish_messages";
@@ -583,6 +676,8 @@ mod tests {
         let client = async_nats::connect(js_url).await.unwrap();
         let context = jetstream::new(client);
         let _ = context.delete_key_value(store_name).await;
+        let cancel_token = CancellationToken::new();
+        let _token_guard = cancel_token.clone().drop_guard();
 
         let kv_store = context
             .create_key_value(jetstream::kv::Config {
@@ -593,7 +688,7 @@ mod tests {
             .await
             .unwrap();
         context
-            .create_stream(async_nats::jetstream::stream::Config {
+            .create_stream(jetstream::stream::Config {
                 name: messages_stream_name.into(),
                 ..Default::default()
             })
@@ -602,43 +697,80 @@ mod tests {
 
         let settings = Settings {
             tid_header: ID_HEADER.into(),
+            pod_hash: POD_HASH.into(),
             js_callback_store: store_name.to_string(),
+            js_status_store: store_name.to_string(),
+            js_response_store: store_name.to_string(),
             js_message_stream: messages_stream_name.to_string(),
             ..Default::default()
         };
 
-        let callback_store = JetStreamCallbackStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create callback store");
+        let callback_store = JetStreamCallbackStore::new(
+            context.clone(),
+            POD_HASH,
+            store_name,
+            cancel_token.clone(),
+        )
+        .await
+        .expect("Failed to create callback store");
 
-        let datum_store = JetStreamDataStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create datum store");
+        let datum_store =
+            JetStreamDataStore::new(context.clone(), store_name, POD_HASH, cancel_token.clone())
+                .await
+                .expect("Failed to create datum store");
 
         let pipeline_spec: PipelineDCG = PIPELINE_SPEC_ENCODED.parse().unwrap();
         let msg_graph = MessageGraph::from_pipeline(&pipeline_spec).unwrap();
-
-        let orchestrator_state = OrchestratorState::new(msg_graph, datum_store, callback_store)
-            .await
-            .unwrap();
+        let status_tracker = StatusTracker::new(
+            context.clone(),
+            store_name,
+            POD_HASH.to_string(),
+            Some(store_name.to_string()),
+        )
+        .await
+        .unwrap();
+
+        let orchestrator_state = OrchestratorState::new(
+            POD_HASH,
+            msg_graph,
+            datum_store,
+            callback_store,
+            status_tracker,
+        )
+        .await
+        .unwrap();
 
         let app_state = AppState {
             js_context: context,
             settings: Arc::new(settings),
             orchestrator_state,
+            cancellation_token: cancel_token.clone(),
         };
 
-        let app = jetstream_proxy(app_state).await.unwrap();
+        let app = serving_proxy(app_state).await.unwrap();
 
         tokio::spawn(async move {
-            let rs_key = format!("rs.{ID_VALUE}.temp.12345");
+            // once the start processing rs is present then only write the response
+            let rs_key = format!("rs.{POD_HASH}.{ID_VALUE}.start.processing");
+            let start_time = Instant::now();
+            loop {
+                if let Some(_) = kv_store.get(&rs_key).await.unwrap() {
+                    break;
+                }
+
+                if start_time.elapsed().as_millis() > 1000 {
+                    panic!("Timed out waiting for start processing key");
+                }
+            }
+            let rs_key = format!("rs.{POD_HASH}.{ID_VALUE}.temp.12345");
             kv_store
                 .put(&rs_key, Bytes::from("test-output"))
                 .await
                 .unwrap();
+
             let cbs = create_default_callbacks(ID_VALUE);
             for cb in cbs {
-                let cb_key = format!("cb.{ID_VALUE}.temp.{}", cb.cb_time);
+                let cb_key = format!("cb.{POD_HASH}.{ID_VALUE}.temp.{}", cb.cb_time);
                 kv_store.put(&cb_key, cb.try_into().unwrap()).await.unwrap();
             }
         });
@@ -654,7 +786,6 @@ mod tests {
         req.extensions_mut().insert(Tid(ID_VALUE.to_string()));
 
         let response = app.clone().oneshot(req).await.unwrap();
-        // assert_eq!(message.id, ID_VALUE);
         assert_eq!(response.status(), StatusCode::OK);
 
         let result = to_bytes(response.into_body(), 10 * 1024).await.unwrap();
@@ -664,7 +795,8 @@ mod tests {
     #[cfg(feature = "nats-tests")]
     #[tokio::test]
     async fn test_sync_publish_serve() {
-        const ID_VALUE: &str = "foobar";
+        const ID_VALUE: &str = "publish-serve";
+        const POD_HASH: &str = "1";
 
         let store_name = "test_sync_publish_serve";
         let messages_stream_name = "test_sync_publish_serve_messages";
@@ -672,6 +804,8 @@ mod tests {
         let client = async_nats::connect(js_url).await.unwrap();
         let context = jetstream::new(client);
         let _ = context.delete_key_value(store_name).await;
+        let cancel_token = CancellationToken::new();
+        let _token_guard = cancel_token.clone().drop_guard();
 
         let kv_store = context
             .create_key_value(jetstream::kv::Config {
@@ -683,7 +817,7 @@ mod tests {
             .unwrap();
 
         context
-            .create_stream(async_nats::jetstream::stream::Config {
+            .create_stream(jetstream::stream::Config {
                 name: messages_stream_name.into(),
                 ..Default::default()
             })
@@ -692,42 +826,79 @@ mod tests {
 
         let settings = Settings {
             js_callback_store: store_name.to_string(),
+            js_status_store: store_name.to_string(),
+            js_response_store: store_name.to_string(),
             js_message_stream: messages_stream_name.to_string(),
+            pod_hash: POD_HASH.into(),
             ..Default::default()
         };
 
-        let callback_store = JetStreamCallbackStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create callback store");
+        let callback_store = JetStreamCallbackStore::new(
+            context.clone(),
+            POD_HASH,
+            store_name,
+            cancel_token.clone(),
+        )
+        .await
+        .expect("Failed to create callback store");
 
-        let datum_store = JetStreamDataStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create datum store");
+        let datum_store =
+            JetStreamDataStore::new(context.clone(), store_name, POD_HASH, cancel_token.clone())
+                .await
+                .expect("Failed to create datum store");
 
         let pipeline_spec: PipelineDCG = PIPELINE_SPEC_ENCODED.parse().unwrap();
         let msg_graph = MessageGraph::from_pipeline(&pipeline_spec).unwrap();
-
-        let orchestrator_state = OrchestratorState::new(msg_graph, datum_store, callback_store)
-            .await
-            .unwrap();
+        let status_tracker = StatusTracker::new(
+            context.clone(),
+            store_name,
+            POD_HASH.to_string(),
+            Some(store_name.to_string()),
+        )
+        .await
+        .unwrap();
+
+        let orchestrator_state = OrchestratorState::new(
+            POD_HASH,
+            msg_graph,
+            datum_store,
+            callback_store,
+            status_tracker,
+        )
+        .await
+        .unwrap();
 
         let app_state = AppState {
             js_context: context,
             settings: Arc::new(settings),
             orchestrator_state,
+            cancellation_token: cancel_token.clone(),
         };
 
-        let app = jetstream_proxy(app_state).await.unwrap();
+        let app = serving_proxy(app_state).await.unwrap();
 
         // pipeline is in -> cat -> out, so we will have 3 callback requests
         // spawn a tokio task which will insert the callback requests to the callback state
         tokio::spawn(async move {
-            let rs_key_one = format!("rs.{ID_VALUE}.temp.12345");
+            // once the start processing rs is present then only write the response
+            let rs_key = format!("rs.{POD_HASH}.{ID_VALUE}.start.processing");
+            let start_time = Instant::now();
+            loop {
+                if let Some(_) = kv_store.get(&rs_key).await.unwrap() {
+                    break;
+                }
+
+                if start_time.elapsed().as_millis() > 1000 {
+                    panic!("Timed out waiting for start processing key");
+                }
+            }
+
+            let rs_key_one = format!("rs.{POD_HASH}.{ID_VALUE}.temp.12345");
             kv_store
                 .put(&rs_key_one, Bytes::from("Test Message 1"))
                 .await
                 .unwrap();
-            let rs_key_two = format!("rs.{ID_VALUE}.temp.12346");
+            let rs_key_two = format!("rs.{POD_HASH}.{ID_VALUE}.temp.12346");
             kv_store
                 .put(&rs_key_two, Bytes::from("Another Test Message 2"))
                 .await
@@ -735,7 +906,7 @@ mod tests {
 
             let cbs = create_default_callbacks(ID_VALUE);
             for cb in cbs {
-                let cb_key = format!("cb.{ID_VALUE}.temp.{}", cb.cb_time);
+                let cb_key = format!("cb.{POD_HASH}.{ID_VALUE}.temp.{}", cb.cb_time);
                 kv_store.put(&cb_key, cb.try_into().unwrap()).await.unwrap();
             }
         });
@@ -785,7 +956,8 @@ mod tests {
     #[tokio::test]
     async fn test_sse_handler() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         const ID_HEADER: &str = "X-Numaflow-Id";
-        const ID_VALUE: &str = "foobar";
+        const POD_HASH: &str = "2";
+        const ID_VALUE: &str = "sse";
 
         let store_name = "test_sse_handler";
         let messages_stream_name = "test_sse_handler_messages";
@@ -793,6 +965,8 @@ mod tests {
         let client = async_nats::connect(js_url).await.unwrap();
         let context = jetstream::new(client);
         let _ = context.delete_key_value(store_name).await;
+        let cancel_token = CancellationToken::new();
+        let _drop_guard = cancel_token.clone().drop_guard();
 
         let kv_store = context
             .create_key_value(jetstream::kv::Config {
@@ -804,7 +978,7 @@ mod tests {
             .unwrap();
 
         context
-            .create_stream(async_nats::jetstream::stream::Config {
+            .create_stream(jetstream::stream::Config {
                 name: messages_stream_name.into(),
                 ..Default::default()
             })
@@ -813,28 +987,53 @@ mod tests {
 
         let settings = Settings {
             tid_header: ID_HEADER.into(),
+            pod_hash: POD_HASH.into(),
             js_callback_store: store_name.to_string(),
+            js_response_store: store_name.to_string(),
+            js_status_store: store_name.to_string(),
             js_message_stream: messages_stream_name.to_string(),
             ..Default::default()
         };
 
-        let callback_store = JetStreamCallbackStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create callback store");
+        let callback_store = JetStreamCallbackStore::new(
+            context.clone(),
+            POD_HASH,
+            store_name,
+            cancel_token.clone(),
+        )
+        .await
+        .expect("Failed to create callback store");
 
-        let datum_store = JetStreamDataStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create datum store");
+        let datum_store =
+            JetStreamDataStore::new(context.clone(), store_name, POD_HASH, cancel_token.clone())
+                .await
+                .expect("Failed to create datum store");
 
         let pipeline_spec = PIPELINE_SPEC_ENCODED.parse().unwrap();
         let msg_graph = MessageGraph::from_pipeline(&pipeline_spec)?;
-        let orchestrator_state =
-            OrchestratorState::new(msg_graph, datum_store, callback_store).await?;
+        let status_tracker = StatusTracker::new(
+            context.clone(),
+            store_name,
+            POD_HASH.to_string(),
+            Some(store_name.to_string()),
+        )
+        .await
+        .unwrap();
+
+        let orchestrator_state = OrchestratorState::new(
+            POD_HASH,
+            msg_graph,
+            datum_store,
+            callback_store,
+            status_tracker,
+        )
+        .await?;
 
         let app_state = AppState {
             js_context: context,
             settings: Arc::new(settings),
             orchestrator_state,
+            cancellation_token: cancel_token.clone(),
         };
 
         let app = router_with_auth(app_state).await?;
@@ -851,15 +1050,37 @@ mod tests {
         // Returns address (e.g. http://127.0.0.1{random_port})
         let listening_url = format!("http://{}:{}", host, port);
 
+        let mut event_stream = reqwest::Client::new()
+            .get(format!("{}/v1/process/sse", listening_url))
+            .header("Content-Type", "text/plain")
+            .header(DEFAULT_ID_HEADER, ID_VALUE)
+            .send()
+            .await
+            .unwrap()
+            .bytes_stream();
+
         // pipeline is in -> cat -> out, so we will have 3 callback requests
         // spawn a tokio task which will insert the callback requests to the callback state
         tokio::spawn(async move {
-            let rs_key_one = format!("rs.{ID_VALUE}.test.12345");
+            // once the start processing rs is present then only write the response
+            let rs_key = format!("rs.{POD_HASH}.{ID_VALUE}.start.processing");
+            let start_time = Instant::now();
+            loop {
+                if let Some(_) = kv_store.get(&rs_key).await.unwrap() {
+                    break;
+                }
+
+                if start_time.elapsed().as_millis() > 1000 {
+                    panic!("Timed out waiting for start processing key");
+                }
+            }
+
+            let rs_key_one = format!("rs.{POD_HASH}.{ID_VALUE}.test.12345");
             kv_store
                 .put(&rs_key_one, Bytes::from("Test Message 1"))
                 .await
                 .unwrap();
-            let rs_key_two = format!("rs.{ID_VALUE}.test.12346");
+            let rs_key_two = format!("rs.{POD_HASH}.{ID_VALUE}.test.12346");
             kv_store
                 .put(&rs_key_two, Bytes::from("Another Test Message 2"))
                 .await
@@ -867,20 +1088,11 @@ mod tests {
 
             let cbs = create_default_callbacks(ID_VALUE);
             for cb in cbs {
-                let cb_key = format!("cb.{ID_VALUE}.test.{}", cb.cb_time);
+                let cb_key = format!("cb.{POD_HASH}.{ID_VALUE}.test.{}", cb.cb_time);
                 kv_store.put(&cb_key, cb.try_into().unwrap()).await.unwrap();
             }
         });
 
-        let mut event_stream = reqwest::Client::new()
-            .get(format!("{}/v1/process/sse", listening_url))
-            .header("Content-Type", "text/plain")
-            .header(DEFAULT_ID_HEADER, ID_VALUE)
-            .send()
-            .await
-            .unwrap()
-            .bytes_stream();
-
         let mut event_data: Vec<Bytes> = vec![];
         while let Some(event) = event_stream.next().await {
             match event {
@@ -899,6 +1111,7 @@ mod tests {
             event_data[1],
             Bytes::from("data: Another Test Message 2\n\n")
         );
+
         Ok(())
     }
 }
diff --git a/rust/serving/src/app/message_path.rs b/rust/serving/src/app/message_path.rs
index d512f1cf..80b22647 100644
--- a/rust/serving/src/app/message_path.rs
+++ b/rust/serving/src/app/message_path.rs
@@ -47,8 +47,10 @@ mod tests {
     use super::*;
     use crate::app::store::cbstore::memstore::InMemoryCallbackStore;
     use crate::app::store::datastore::inmemory::InMemoryDataStore;
+    use crate::app::store::status::StatusTracker;
     use crate::app::tracker::MessageGraph;
     use crate::pipeline::PipelineDCG;
+    use async_nats::jetstream;
     use axum::body::Body;
     use axum::extract::Request;
     use axum::http::header::CONTENT_TYPE;
@@ -57,14 +59,34 @@ mod tests {
 
     const PIPELINE_SPEC_ENCODED: &str = "eyJ2ZXJ0aWNlcyI6W3sibmFtZSI6ImluIiwic291cmNlIjp7InNlcnZpbmciOnsiYXV0aCI6bnVsbCwic2VydmljZSI6dHJ1ZSwibXNnSURIZWFkZXJLZXkiOiJYLU51bWFmbG93LUlkIiwic3RvcmUiOnsidXJsIjoicmVkaXM6Ly9yZWRpczo2Mzc5In19fSwiY29udGFpbmVyVGVtcGxhdGUiOnsicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIiLCJlbnYiOlt7Im5hbWUiOiJSVVNUX0xPRyIsInZhbHVlIjoiZGVidWcifV19LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6InBsYW5uZXIiLCJ1ZGYiOnsiY29udGFpbmVyIjp7ImltYWdlIjoiYXNjaWk6MC4xIiwiYXJncyI6WyJwbGFubmVyIl0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sImJ1aWx0aW4iOm51bGwsImdyb3VwQnkiOm51bGx9LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6InRpZ2VyIiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsidGlnZXIiXSwicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwiYnVpbHRpbiI6bnVsbCwiZ3JvdXBCeSI6bnVsbH0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiZG9nIiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsiZG9nIl0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sImJ1aWx0aW4iOm51bGwsImdyb3VwQnkiOm51bGx9LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6ImVsZXBoYW50IiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsiZWxlcGhhbnQiXSwicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwiYnVpbHRpbiI6bnVsbCwiZ3JvdXBCeSI6bnVsbH0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiYXNjaWlhcnQiLCJ1ZGYiOnsiY29udGFpbmVyIjp7ImltYWdlIjoiYXNjaWk6MC4xIiwiYXJncyI6WyJhc2NpaWFydCJdLCJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJidWlsdGluIjpudWxsLCJncm91cEJ5IjpudWxsfSwiY29udGFpbmVyVGVtcGxhdGUiOnsicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwic2NhbGUiOnsibWluIjoxfSwidXBkYXRlU3RyYXRlZ3kiOnsidHlwZSI6IlJvbGxpbmdVcGRhdGUiLCJyb2xsaW5nVXBkYXRlIjp7Im1heFVuYXZhaWxhYmxlIjoiMjUlIn19fSx7Im5hbWUiOiJzZXJ2ZS1zaW5rIiwic2luayI6eyJ1ZHNpbmsiOnsiY29udGFpbmVyIjp7ImltYWdlIjoic2VydmVzaW5rOjAuMSIsImVudiI6W3sibmFtZSI6Ik5VTUFGTE9XX0NBTExCQUNLX1VSTF9LRVkiLCJ2YWx1ZSI6IlgtTnVtYWZsb3ctQ2FsbGJhY2stVXJsIn0seyJuYW1lIjoiTlVNQUZMT1dfTVNHX0lEX0hFQURFUl9LRVkiLCJ2YWx1ZSI6IlgtTnVtYWZsb3ctSWQifV0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn19LCJyZXRyeVN0cmF0ZWd5Ijp7fX0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiZXJyb3Itc2luayIsInNpbmsiOnsidWRzaW5rIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6InNlcnZlc2luazowLjEiLCJlbnYiOlt7Im5hbWUiOiJOVU1BRkxPV19DQUxMQkFDS19VUkxfS0VZIiwidmFsdWUiOiJYLU51bWFmbG93LUNhbGxiYWNrLVVybCJ9LHsibmFtZSI6Ik5VTUFGTE9XX01TR19JRF9IRUFERVJfS0VZIiwidmFsdWUiOiJYLU51bWFmbG93LUlkIn1dLCJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9fSwicmV0cnlTdHJhdGVneSI6e319LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19XSwiZWRnZXMiOlt7ImZyb20iOiJpbiIsInRvIjoicGxhbm5lciIsImNvbmRpdGlvbnMiOm51bGx9LHsiZnJvbSI6InBsYW5uZXIiLCJ0byI6ImFzY2lpYXJ0IiwiY29uZGl0aW9ucyI6eyJ0YWdzIjp7Im9wZXJhdG9yIjoib3IiLCJ2YWx1ZXMiOlsiYXNjaWlhcnQiXX19fSx7ImZyb20iOiJwbGFubmVyIiwidG8iOiJ0aWdlciIsImNvbmRpdGlvbnMiOnsidGFncyI6eyJvcGVyYXRvciI6Im9yIiwidmFsdWVzIjpbInRpZ2VyIl19fX0seyJmcm9tIjoicGxhbm5lciIsInRvIjoiZG9nIiwiY29uZGl0aW9ucyI6eyJ0YWdzIjp7Im9wZXJhdG9yIjoib3IiLCJ2YWx1ZXMiOlsiZG9nIl19fX0seyJmcm9tIjoicGxhbm5lciIsInRvIjoiZWxlcGhhbnQiLCJjb25kaXRpb25zIjp7InRhZ3MiOnsib3BlcmF0b3IiOiJvciIsInZhbHVlcyI6WyJlbGVwaGFudCJdfX19LHsiZnJvbSI6InRpZ2VyIiwidG8iOiJzZXJ2ZS1zaW5rIiwiY29uZGl0aW9ucyI6bnVsbH0seyJmcm9tIjoiZG9nIiwidG8iOiJzZXJ2ZS1zaW5rIiwiY29uZGl0aW9ucyI6bnVsbH0seyJmcm9tIjoiZWxlcGhhbnQiLCJ0byI6InNlcnZlLXNpbmsiLCJjb25kaXRpb25zIjpudWxsfSx7ImZyb20iOiJhc2NpaWFydCIsInRvIjoic2VydmUtc2luayIsImNvbmRpdGlvbnMiOm51bGx9LHsiZnJvbSI6InBsYW5uZXIiLCJ0byI6ImVycm9yLXNpbmsiLCJjb25kaXRpb25zIjp7InRhZ3MiOnsib3BlcmF0b3IiOiJvciIsInZhbHVlcyI6WyJlcnJvciJdfX19XSwibGlmZWN5Y2xlIjp7fSwid2F0ZXJtYXJrIjp7fX0=";
 
+    #[cfg(feature = "nats-tests")]
     #[tokio::test]
     async fn test_message_path_not_present() {
         let data_store = InMemoryDataStore::new(None);
         let cb_store = InMemoryCallbackStore::new(None);
+        let store_name = "test_async_publish";
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let _ = context.delete_key_value(store_name).await;
+
+        let _ = context
+            .create_key_value(jetstream::kv::Config {
+                bucket: store_name.to_string(),
+                history: 5,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let status_tracker =
+            StatusTracker::new(context.clone(), store_name, "xgda".to_string(), None)
+                .await
+                .unwrap();
 
         let pipeline_spec: PipelineDCG = PIPELINE_SPEC_ENCODED.parse().unwrap();
         let msg_graph = MessageGraph::from_pipeline(&pipeline_spec).unwrap();
-        let state = CallbackState::new(msg_graph, data_store, cb_store)
+        let state = CallbackState::new("xbda", msg_graph, data_store, cb_store, status_tracker)
             .await
             .unwrap();
         let app = get_message_path(state);
diff --git a/rust/serving/src/app/orchestrator.rs b/rust/serving/src/app/orchestrator.rs
index 9bd5fa4a..df10c10f 100644
--- a/rust/serving/src/app/orchestrator.rs
+++ b/rust/serving/src/app/orchestrator.rs
@@ -1,31 +1,30 @@
 //! To process a message in Serving, we need to track the progress of the processing state of the message
 //! and then fetch the messages once the results of the processing is stored in [crate::app::store::datastore] by
-//! different Sink. This subtly means that we even support streaming the results. The processing state
-//! is stored in [crate::app::store::cbstore].
-use std::sync::Arc;
+//! different Sink. The processing state is stored in [crate::app::store::cbstore].
 
 use bytes::Bytes;
+use std::sync::Arc;
 use tokio::sync::{mpsc, oneshot};
-use tokio::time::Instant;
 use tokio_stream::wrappers::ReceiverStream;
 use tokio_stream::StreamExt;
-use tracing::{error, info, Instrument};
+use tracing::{debug, error, info, trace, warn, Instrument};
 
-use crate::app::store::cbstore::{CallbackStore, ProcessingStatus};
+use crate::app::store::cbstore::CallbackStore;
 use crate::app::store::datastore::DataStore;
-use crate::app::store::datastore::Error as StoreError;
-use crate::app::store::datastore::Result as StoreResult;
+use crate::app::store::status;
+use crate::app::store::status::{ProcessingStatus, StatusTracker};
 use crate::app::tracker::MessageGraph;
+use crate::config::RequestType;
 use crate::metrics::serving_metrics;
 use crate::Error;
 
 #[derive(Clone)]
 pub(crate) struct OrchestratorState<T, U> {
-    // generator to generate subgraph
     msg_graph_generator: Arc<MessageGraph>,
-    // conn is to be used while reading and writing to redis.
     datum_store: T,
     callback_store: U,
+    status_tracker: StatusTracker,
+    pod_hash: String,
 }
 
 impl<T, U> OrchestratorState<T, U>
@@ -35,122 +34,133 @@ where
 {
     /// Create a new State to track connections and callback data
     pub(crate) async fn new(
+        pod_hash: &str,
         msg_graph: MessageGraph,
         datum_store: T,
         callback_store: U,
+        status_tracker: StatusTracker,
     ) -> crate::Result<Self> {
         Ok(Self {
             msg_graph_generator: Arc::new(msg_graph),
             datum_store,
             callback_store,
+            status_tracker,
+            pod_hash: pod_hash.to_string(),
         })
     }
 
-    /// register a new connection and spwans a watcher to watch on the callback stream. An oneshot rx is
+    /// register a new connection and spawns a watcher to watch on the callback stream. An oneshot rx is
     /// returned. This oneshot rx will be notified when all callbacks for this connection is received
     /// from the numaflow pipeline.
     pub(crate) async fn process_request(
         &mut self,
         id: &str,
-    ) -> StoreResult<oneshot::Receiver<Result<String, Error>>> {
+        request_type: RequestType,
+    ) -> crate::Result<oneshot::Receiver<Result<String, Error>>> {
         let (tx, rx) = oneshot::channel();
         let sub_graph_generator = Arc::clone(&self.msg_graph_generator);
         let msg_id = id.to_string();
-        let mut subgraph = None;
 
-        let start = Instant::now();
-        // register the request in the store
-        if let Err(e) = self.callback_store.register(id).await {
-            if let StoreError::DuplicateRequest(_) = e {
-                serving_metrics().cb_store_register_duplicate_count.inc();
-            }
-            serving_metrics().cb_store_register_fail_count.inc();
-            return Err(e);
-        }
-        serving_metrics()
-            .cb_store_register_duration
-            .observe(start.elapsed().as_micros() as f64);
+        // register the request id in the status tracker, if the request is already processed
+        // or if its getting processed the tracker will return duplicate error. If the request
+        // was cancelled before sending the response then the previous pod hash will be returned
+        // with cancelled error, and we will fetch the callbacks and responses again.
+        let pod_hash = match self.status_tracker.register(id, request_type).await {
+            Ok(_) => self.pod_hash.clone(),
+            Err(e) => match e {
+                status::Error::Cancelled {
+                    err,
+                    previous_pod_hash,
+                } => {
+                    warn!(
+                        ?err,
+                        "Request was cancelled, fetching callbacks and responses again"
+                    );
+                    previous_pod_hash
+                }
+                status::Error::Duplicate(msg) => {
+                    warn!(error = ?msg, "Request already exists in the store");
+                    serving_metrics().request_register_duplicate_count.inc();
+                    return Err(Error::Duplicate(msg));
+                }
+                _ => {
+                    error!(?e, "Failed to register request id in status store");
+                    serving_metrics().request_register_fail_count.inc();
+                    return Err(e.into());
+                }
+            },
+        };
 
         // start watching for callbacks
-        let (mut callbacks_stream, watch_handle) = self.callback_store.watch_callbacks(id).await?;
+        let mut callbacks_stream = self
+            .callback_store
+            .register_and_watch(id, &pod_hash)
+            .await?;
 
+        let mut status_tracker = self.status_tracker.clone();
         let mut cb_store = self.callback_store.clone();
-
         let span = tracing::Span::current();
-
         let callback_watcher = async move {
-            let _handle = watch_handle;
             let mut callbacks = Vec::new();
-
             while let Some(cb) = callbacks_stream.next().await {
-                info!(?cb, ?msg_id, "Received callback");
+                trace!(?cb, ?msg_id, "Received callback");
                 callbacks.push(cb);
-                subgraph = match sub_graph_generator
+                let subgraph = sub_graph_generator
                     .generate_subgraph_from_callbacks(msg_id.clone(), callbacks.clone())
-                {
-                    Ok(subgraph) => subgraph,
-                    Err(e) => {
-                        error!(?e, "Failed to generate subgraph");
-                        break;
+                    .expect("Failed to generate subgraph");
+
+                if let Some(graph) = subgraph {
+                    status_tracker
+                        .deregister(&msg_id, &graph)
+                        .await
+                        .expect("Failed to deregister");
+                    cb_store
+                        .deregister(&msg_id)
+                        .await
+                        .expect("Failed to deregister");
+
+                    // send can only fail if the request was cancelled by the client or the handler task
+                    // was terminated.
+                    if tx.send(Ok(graph)).is_err() {
+                        status_tracker
+                            .mark_as_failed(&msg_id, "Cancelled")
+                            .await
+                            .expect("Failed to mark the request as failed");
                     }
-                };
-                if subgraph.is_some() {
                     break;
                 }
             }
-
-            if let Some(subgraph) = subgraph {
-                cb_store
-                    .deregister(&msg_id, &subgraph)
-                    .await
-                    .expect("Failed to deregister");
-                tx.send(Ok(subgraph.clone()))
-                    .expect("Failed to send subgraph");
-            } else {
-                error!("Subgraph could not be generated for the given ID");
-                tx.send(Err(Error::SubGraphNotFound(
-                    "Subgraph could not be generated for the given ID",
-                )))
-                .inspect_err(|err| {
-                    tracing::error!(
-                        ?err,
-                        "Failed to send subgraph generation error to main task over channel"
-                    )
-                })
-                .expect("Failed to send subgraph");
-
-                cb_store
-                    .mark_as_failed(&msg_id, "Subgraph could not be generated")
-                    .await
-                    .inspect_err(|err| {
-                        tracing::error!(
-                            ?err,
-                            "Failed to mark subgraph generation failure in callback store"
-                        )
-                    })
-                    .expect("Failed to mark as failed");
-            }
         };
 
         // https://github.com/tokio-rs/tracing/blob/b4868674ba73f3963b125ec38b57efaadc714d90/examples/examples/tokio-spawny-thing.rs#L25
         tokio::spawn(callback_watcher.instrument(span));
-
         Ok(rx)
     }
 
     /// Retrieves the output of the processed request after checking whether the processing is complete.
-    pub(crate) async fn retrieve_saved(&mut self, id: &str) -> Result<Option<Vec<Vec<u8>>>, Error> {
+    pub(crate) async fn retrieve_saved(
+        &mut self,
+        id: &str,
+        request_type: RequestType,
+    ) -> Result<Option<Vec<Vec<u8>>>, Error> {
         // check the status of the request, if its completed, then retrieve the data
-        let status = self.callback_store.status(id).await?;
+        let status = self.status_tracker.status(id).await?;
         match status {
-            ProcessingStatus::InProgress => Ok(None),
-            ProcessingStatus::Completed(_) => {
-                let data = self.datum_store.retrieve_data(id).await?;
+            ProcessingStatus::InProgress { .. } => {
+                info!(?id, "Request is still in progress");
+                Ok(None)
+            }
+            ProcessingStatus::Completed { .. } => {
+                let data = self.datum_store.retrieve_data(id, None).await?;
                 Ok(Some(data))
             }
-            ProcessingStatus::Failed(error) => {
-                error!(?error, "Request failed");
-                Err(Error::Other(format!("Request failed: {}", error)))
+            ProcessingStatus::Failed { error, pod_hash } => {
+                warn!("Request was failed because of {error}, processing and fetching the responses again");
+                let notify = self.process_request(id, request_type).await?;
+                notify.await.expect("sender was dropped")?;
+                Ok(Some(
+                    self.datum_store.retrieve_data(id, Some(&pod_hash)).await?,
+                ))
             }
         }
     }
@@ -161,27 +171,55 @@ where
     pub(crate) async fn stream_response(
         &mut self,
         id: &str,
+        request_type: RequestType,
     ) -> Result<ReceiverStream<Arc<Bytes>>, Error> {
         let (tx, rx) = mpsc::channel(10);
         let sub_graph_generator = Arc::clone(&self.msg_graph_generator);
         let msg_id = id.to_string();
-        let mut subgraph = None;
 
-        // register the request in the store
-        self.callback_store.register(id).await?;
+        // register the request id in the status tracker, if the request is already processed
+        // or if its getting processed the tracker will return duplicate error. If the request
+        // was cancelled before sending the response then the previous pod hash will be returned
+        // with cancelled error, and we will fetch the callbacks and responses again.
+        let pod_hash = match self.status_tracker.register(id, request_type).await {
+            Ok(_) => self.pod_hash.clone(),
+            Err(e) => match e {
+                status::Error::Cancelled {
+                    err,
+                    previous_pod_hash,
+                } => {
+                    warn!(
+                        ?err,
+                        "Request was cancelled, fetching callbacks and responses again"
+                    );
+                    previous_pod_hash
+                }
+                status::Error::Duplicate(msg) => {
+                    warn!(error = ?msg, "Request already exists in the tracker");
+                    serving_metrics().request_register_duplicate_count.inc();
+                    return Err(Error::Duplicate(msg));
+                }
+                _ => {
+                    error!(?e, "Failed to register request id in status tracker");
+                    serving_metrics().request_register_fail_count.inc();
+                    return Err(e.into());
+                }
+            },
+        };
 
         // start watching for callbacks
-        let (mut callbacks_stream, watch_handle) = self.callback_store.watch_callbacks(id).await?;
+        let mut callbacks_stream = self
+            .callback_store
+            .register_and_watch(id, &pod_hash)
+            .await?;
 
         let mut cb_store = self.callback_store.clone();
+        let status_tracker = self.status_tracker.clone();
         tokio::spawn(async move {
-            let _handle = watch_handle;
             let mut callbacks = Vec::new();
-
             while let Some(cb) = callbacks_stream.next().await {
                 callbacks.push(cb);
-
-                subgraph = match sub_graph_generator
+                let subgraph = match sub_graph_generator
                     .generate_subgraph_from_callbacks(msg_id.clone(), callbacks.clone())
                 {
                     Ok(subgraph) => subgraph,
@@ -191,31 +229,33 @@ where
                     }
                 };
 
-                if subgraph.is_some() {
+                if let Some(graph) = subgraph {
+                    status_tracker
+                        .deregister(&msg_id, &graph)
+                        .await
+                        .expect("Failed to deregister in status tracker");
+                    cb_store
+                        .deregister(&msg_id)
+                        .await
+                        .expect("Failed to deregister in callback store");
                     break;
                 }
             }
-
-            if let Some(subgraph) = subgraph {
-                cb_store
-                    .deregister(&msg_id, &subgraph)
-                    .await
-                    .expect("Failed to deregister");
-            } else {
-                error!("Subgraph could not be generated for the given ID");
-
-                cb_store
-                    .mark_as_failed(&msg_id, "Subgraph could not be generated")
-                    .await
-                    .expect("Failed to mark as failed");
-            }
         });
 
         // watch for data stored in the datastore
-        let (mut response_stream, _handle) = self.datum_store.stream_data(id).await?;
+        let request_id = id.to_string();
+        let mut status_tracker = self.status_tracker.clone();
+        let mut response_stream = self.datum_store.stream_data(id, &pod_hash).await?;
         tokio::spawn(async move {
             while let Some(response) = response_stream.next().await {
-                tx.send(response).await.expect("Failed to send response");
+                if tx.send(response).await.is_err() {
+                    error!("Failed to send response");
+                    status_tracker
+                        .mark_as_failed(&request_id, "Cancelled")
+                        .await
+                        .expect("Failed to mark the request has failed");
+                }
             }
         });
 
@@ -228,21 +268,28 @@ where
         &mut self,
         id: &str,
     ) -> Result<String, Error> {
-        let status = self.callback_store.status(id).await?;
+        let status = self.status_tracker.status(id).await?;
         match status {
-            ProcessingStatus::InProgress => Ok("Request In Progress".to_string()),
-            ProcessingStatus::Completed(sub_graph) => Ok(sub_graph),
-            ProcessingStatus::Failed(error) => {
+            ProcessingStatus::InProgress { .. } => Ok("Request In Progress".to_string()),
+            ProcessingStatus::Completed { subgraph, .. } => Ok(subgraph),
+            ProcessingStatus::Failed { error, .. } => {
                 error!(?error, "Request failed");
                 Err(Error::SubGraphGenerator(error))
             }
         }
     }
 
-    /// marks the request as failed, FIXME: invoke this when the subgraph generation fails/critical errors
-    #[allow(dead_code)]
+    /// marks the request as failed
     pub(crate) async fn mark_as_failed(&mut self, id: &str, error: &str) -> Result<(), Error> {
-        self.callback_store.mark_as_failed(id, error).await?;
+        self.status_tracker.mark_as_failed(id, error).await?;
+        Ok(())
+    }
+
+    /// discards the request from the status tracker and callback store. This is used when the
+    /// request has not been accepted due to back pressure.
+    pub(crate) async fn discard(&mut self, id: &str) -> Result<(), Error> {
+        self.status_tracker.discard(id).await?;
+        self.callback_store.deregister(id).await?;
         Ok(())
     }
 
@@ -254,15 +301,16 @@ where
 
 #[cfg(test)]
 mod tests {
-    use async_nats::jetstream;
-    use axum::body::Bytes;
-    use chrono::Utc;
-
     use super::*;
-    use crate::app::store::cbstore::jetstreamstore::JetStreamCallbackStore;
+    use crate::app::store::cbstore::jetstream_store::JetStreamCallbackStore;
     use crate::app::store::datastore::jetstream::JetStreamDataStore;
     use crate::callback::{Callback, Response};
     use crate::pipeline::PipelineDCG;
+    use async_nats::jetstream;
+    use axum::body::Bytes;
+    use chrono::Utc;
+    use std::time::Instant;
+    use tokio_util::sync::CancellationToken;
 
     const PIPELINE_SPEC_ENCODED: &str = "eyJ2ZXJ0aWNlcyI6W3sibmFtZSI6ImluIiwic291cmNlIjp7InNlcnZpbmciOnsiYXV0aCI6bnVsbCwic2VydmljZSI6dHJ1ZSwibXNnSURIZWFkZXJLZXkiOiJYLU51bWFmbG93LUlkIiwic3RvcmUiOnsidXJsIjoicmVkaXM6Ly9yZWRpczo2Mzc5In19fSwiY29udGFpbmVyVGVtcGxhdGUiOnsicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIiLCJlbnYiOlt7Im5hbWUiOiJSVVNUX0xPRyIsInZhbHVlIjoiZGVidWcifV19LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6InBsYW5uZXIiLCJ1ZGYiOnsiY29udGFpbmVyIjp7ImltYWdlIjoiYXNjaWk6MC4xIiwiYXJncyI6WyJwbGFubmVyIl0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sImJ1aWx0aW4iOm51bGwsImdyb3VwQnkiOm51bGx9LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6InRpZ2VyIiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsidGlnZXIiXSwicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwiYnVpbHRpbiI6bnVsbCwiZ3JvdXBCeSI6bnVsbH0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiZG9nIiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsiZG9nIl0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sImJ1aWx0aW4iOm51bGwsImdyb3VwQnkiOm51bGx9LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19LHsibmFtZSI6ImVsZXBoYW50IiwidWRmIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6ImFzY2lpOjAuMSIsImFyZ3MiOlsiZWxlcGhhbnQiXSwicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwiYnVpbHRpbiI6bnVsbCwiZ3JvdXBCeSI6bnVsbH0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiYXNjaWlhcnQiLCJ1ZGYiOnsiY29udGFpbmVyIjp7ImltYWdlIjoiYXNjaWk6MC4xIiwiYXJncyI6WyJhc2NpaWFydCJdLCJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJidWlsdGluIjpudWxsLCJncm91cEJ5IjpudWxsfSwiY29udGFpbmVyVGVtcGxhdGUiOnsicmVzb3VyY2VzIjp7fSwiaW1hZ2VQdWxsUG9saWN5IjoiTmV2ZXIifSwic2NhbGUiOnsibWluIjoxfSwidXBkYXRlU3RyYXRlZ3kiOnsidHlwZSI6IlJvbGxpbmdVcGRhdGUiLCJyb2xsaW5nVXBkYXRlIjp7Im1heFVuYXZhaWxhYmxlIjoiMjUlIn19fSx7Im5hbWUiOiJzZXJ2ZS1zaW5rIiwic2luayI6eyJ1ZHNpbmsiOnsiY29udGFpbmVyIjp7ImltYWdlIjoic2VydmVzaW5rOjAuMSIsImVudiI6W3sibmFtZSI6Ik5VTUFGTE9XX0NBTExCQUNLX1VSTF9LRVkiLCJ2YWx1ZSI6IlgtTnVtYWZsb3ctQ2FsbGJhY2stVXJsIn0seyJuYW1lIjoiTlVNQUZMT1dfTVNHX0lEX0hFQURFUl9LRVkiLCJ2YWx1ZSI6IlgtTnVtYWZsb3ctSWQifV0sInJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn19LCJyZXRyeVN0cmF0ZWd5Ijp7fX0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImltYWdlUHVsbFBvbGljeSI6Ik5ldmVyIn0sInNjYWxlIjp7Im1pbiI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX0seyJuYW1lIjoiZXJyb3Itc2luayIsInNpbmsiOnsidWRzaW5rIjp7ImNvbnRhaW5lciI6eyJpbWFnZSI6InNlcnZlc2luazowLjEiLCJlbnYiOlt7Im5hbWUiOiJOVU1BRkxPV19DQUxMQkFDS19VUkxfS0VZIiwidmFsdWUiOiJYLU51bWFmbG93LUNhbGxiYWNrLVVybCJ9LHsibmFtZSI6Ik5VTUFGTE9XX01TR19JRF9IRUFERVJfS0VZIiwidmFsdWUiOiJYLU51bWFmbG93LUlkIn1dLCJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9fSwicmV0cnlTdHJhdGVneSI6e319LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJpbWFnZVB1bGxQb2xpY3kiOiJOZXZlciJ9LCJzY2FsZSI6eyJtaW4iOjF9LCJ1cGRhdGVTdHJhdGVneSI6eyJ0eXBlIjoiUm9sbGluZ1VwZGF0ZSIsInJvbGxpbmdVcGRhdGUiOnsibWF4VW5hdmFpbGFibGUiOiIyNSUifX19XSwiZWRnZXMiOlt7ImZyb20iOiJpbiIsInRvIjoicGxhbm5lciIsImNvbmRpdGlvbnMiOm51bGx9LHsiZnJvbSI6InBsYW5uZXIiLCJ0byI6ImFzY2lpYXJ0IiwiY29uZGl0aW9ucyI6eyJ0YWdzIjp7Im9wZXJhdG9yIjoib3IiLCJ2YWx1ZXMiOlsiYXNjaWlhcnQiXX19fSx7ImZyb20iOiJwbGFubmVyIiwidG8iOiJ0aWdlciIsImNvbmRpdGlvbnMiOnsidGFncyI6eyJvcGVyYXRvciI6Im9yIiwidmFsdWVzIjpbInRpZ2VyIl19fX0seyJmcm9tIjoicGxhbm5lciIsInRvIjoiZG9nIiwiY29uZGl0aW9ucyI6eyJ0YWdzIjp7Im9wZXJhdG9yIjoib3IiLCJ2YWx1ZXMiOlsiZG9nIl19fX0seyJmcm9tIjoicGxhbm5lciIsInRvIjoiZWxlcGhhbnQiLCJjb25kaXRpb25zIjp7InRhZ3MiOnsib3BlcmF0b3IiOiJvciIsInZhbHVlcyI6WyJlbGVwaGFudCJdfX19LHsiZnJvbSI6InRpZ2VyIiwidG8iOiJzZXJ2ZS1zaW5rIiwiY29uZGl0aW9ucyI6bnVsbH0seyJmcm9tIjoiZG9nIiwidG8iOiJzZXJ2ZS1zaW5rIiwiY29uZGl0aW9ucyI6bnVsbH0seyJmcm9tIjoiZWxlcGhhbnQiLCJ0byI6InNlcnZlLXNpbmsiLCJjb25kaXRpb25zIjpudWxsfSx7ImZyb20iOiJhc2NpaWFydCIsInRvIjoic2VydmUtc2luayIsImNvbmRpdGlvbnMiOm51bGx9LHsiZnJvbSI6InBsYW5uZXIiLCJ0byI6ImVycm9yLXNpbmsiLCJjb25kaXRpb25zIjp7InRhZ3MiOnsib3BlcmF0b3IiOiJvciIsInZhbHVlcyI6WyJlcnJvciJdfX19XSwibGlmZWN5Y2xlIjp7fSwid2F0ZXJtYXJrIjp7fX0=";
 
@@ -272,6 +320,7 @@ mod tests {
         let js_url = "localhost:4222";
         let client = async_nats::connect(js_url).await.unwrap();
         let context = jetstream::new(client);
+        let pod_hash = "0";
 
         let pipeline_spec: PipelineDCG = PIPELINE_SPEC_ENCODED.parse().unwrap();
         let msg_graph = MessageGraph::from_pipeline(&pipeline_spec).unwrap();
@@ -288,17 +337,42 @@ mod tests {
             .await
             .unwrap();
 
-        let callback_store = JetStreamCallbackStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create callback store");
-
-        let datum_store = JetStreamDataStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create datum store");
-
-        let mut state = OrchestratorState::new(msg_graph, datum_store, callback_store)
-            .await
-            .unwrap();
+        let callback_store = JetStreamCallbackStore::new(
+            context.clone(),
+            pod_hash,
+            store_name,
+            CancellationToken::new(),
+        )
+        .await
+        .expect("Failed to create callback store");
+
+        let datum_store = JetStreamDataStore::new(
+            context.clone(),
+            store_name,
+            pod_hash,
+            CancellationToken::new(),
+        )
+        .await
+        .expect("Failed to create datum store");
+
+        let status_tracker = StatusTracker::new(
+            context.clone(),
+            store_name,
+            pod_hash.to_string(),
+            Some(store_name.to_string()),
+        )
+        .await
+        .unwrap();
+
+        let mut state = OrchestratorState::new(
+            pod_hash,
+            msg_graph,
+            datum_store,
+            callback_store,
+            status_tracker,
+        )
+        .await
+        .unwrap();
 
         let id = "test_id".to_string();
 
@@ -350,41 +424,68 @@ mod tests {
             },
         ];
 
-        // put the callbacks into the store
-        for (i, cb) in cbs.drain(..).enumerate() {
-            kv_store
-                .put(
-                    format!("cb.{}.{}.{}", id, i, Utc::now().timestamp()),
-                    cb.try_into().unwrap(),
-                )
-                .await
-                .unwrap();
-        }
+        let request_id = id.clone();
+        tokio::spawn(async move {
+            // once the start processing rs is present then only write the response
+            let rs_key = format!("rs.{pod_hash}.{request_id}.start.processing");
+            let start_time = Instant::now();
+            loop {
+                if let Some(_) = kv_store.get(&rs_key).await.unwrap() {
+                    break;
+                }
 
-        // put responses to datum store
-        let responses = vec![
-            Bytes::from_static(b"response1"),
-            Bytes::from_static(b"response2"),
-        ];
+                if start_time.elapsed().as_millis() > 1000 {
+                    panic!("Timed out waiting for start processing key");
+                }
+            }
+            // put responses to datum store
+            let responses = vec![
+                Bytes::from_static(b"response1"),
+                Bytes::from_static(b"response2"),
+            ];
+
+            for (i, response) in responses.iter().enumerate() {
+                kv_store
+                    .put(
+                        format!(
+                            "rs.{pod_hash}.{request_id}.{}.{}",
+                            i,
+                            Utc::now().timestamp()
+                        ),
+                        response.clone(),
+                    )
+                    .await
+                    .unwrap();
+            }
 
-        for (i, response) in responses.iter().enumerate() {
-            kv_store
-                .put(
-                    format!("rs.{}.{}.{}", id, i, Utc::now().timestamp()),
-                    response.clone(),
-                )
-                .await
-                .unwrap();
-        }
+            // put the callbacks into the store
+            for (i, cb) in cbs.drain(..).enumerate() {
+                kv_store
+                    .put(
+                        format!(
+                            "cb.{pod_hash}.{request_id}.{}.{}",
+                            i,
+                            Utc::now().timestamp_millis()
+                        ),
+                        cb.try_into().unwrap(),
+                    )
+                    .await
+                    .unwrap();
+            }
+        });
 
         // Test process_request
-        let result = state.process_request(&id).await.unwrap();
+        let result = state.process_request(&id, RequestType::Sync).await.unwrap();
         let subgraph = result.await.unwrap();
         assert!(subgraph.is_ok());
         assert!(!subgraph.unwrap().is_empty());
 
         // test retrieve_saved
-        let result = state.retrieve_saved(&id).await.unwrap().unwrap();
+        let result = state
+            .retrieve_saved(&id, RequestType::Sync)
+            .await
+            .unwrap()
+            .unwrap();
         assert_eq!(result.len(), 2);
         assert_eq!(result[0], b"response1");
         assert_eq!(result[1], b"response2");
@@ -396,6 +497,7 @@ mod tests {
         let js_url = "localhost:4222";
         let client = async_nats::connect(js_url).await.unwrap();
         let context = jetstream::new(client);
+        let pod_hash = "0";
 
         let store_name = "test_stream_response";
         let _ = context.delete_key_value(store_name).await;
@@ -409,25 +511,49 @@ mod tests {
             .await
             .unwrap();
 
-        let callback_store = JetStreamCallbackStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create callback store");
-
-        let datum_store = JetStreamDataStore::new(context.clone(), store_name)
-            .await
-            .expect("Failed to create datum store");
+        let callback_store = JetStreamCallbackStore::new(
+            context.clone(),
+            pod_hash,
+            store_name,
+            CancellationToken::new(),
+        )
+        .await
+        .expect("Failed to create callback store");
+
+        let datum_store = JetStreamDataStore::new(
+            context.clone(),
+            store_name,
+            pod_hash,
+            CancellationToken::new(),
+        )
+        .await
+        .expect("Failed to create datum store");
+
+        let status_tracker = StatusTracker::new(
+            context.clone(),
+            store_name,
+            pod_hash.to_string(),
+            Some(store_name.to_string()),
+        )
+        .await
+        .unwrap();
 
         let pipeline_spec: PipelineDCG = PIPELINE_SPEC_ENCODED.parse().unwrap();
 
         let msg_graph = MessageGraph::from_pipeline(&pipeline_spec).unwrap();
 
-        let mut state = OrchestratorState::new(msg_graph, datum_store, callback_store)
-            .await
-            .unwrap();
+        let mut state = OrchestratorState::new(
+            pod_hash,
+            msg_graph,
+            datum_store,
+            callback_store,
+            status_tracker,
+        )
+        .await
+        .unwrap();
 
         let id = "test_id".to_string();
 
-        // Test insert_callback_requests
         let mut cbs = vec![
             Callback {
                 id: id.clone(),
@@ -475,38 +601,61 @@ mod tests {
             },
         ];
 
-        // put the callbacks into the store
-        for (i, cb) in cbs.drain(..).enumerate() {
-            kv_store
-                .put(
-                    format!("cb.{}.{}.{}", id, i, Utc::now().timestamp_millis()),
-                    cb.try_into().unwrap(),
-                )
-                .await
-                .unwrap();
-        }
+        let request_id = id.clone();
+        tokio::spawn(async move {
+            // once the start processing rs is present then only write the response
+            let rs_key = format!("rs.{pod_hash}.{request_id}.start.processing");
+            let start_time = Instant::now();
+            loop {
+                if let Some(_) = kv_store.get(&rs_key).await.unwrap() {
+                    break;
+                }
 
-        // put responses to datum store
-        let responses = vec![
-            Bytes::from_static(b"response1"),
-            Bytes::from_static(b"response2"),
-            Bytes::from_static(b"response3"),
-            Bytes::from_static(b"response4"),
-            Bytes::from_static(b"response5"),
-        ];
+                if start_time.elapsed().as_millis() > 1000 {
+                    panic!("Timed out waiting for start processing key");
+                }
+            }
+            // put responses to datum store
+            let responses = vec![
+                Bytes::from_static(b"response1"),
+                Bytes::from_static(b"response2"),
+                Bytes::from_static(b"response3"),
+                Bytes::from_static(b"response4"),
+                Bytes::from_static(b"response5"),
+            ];
+
+            for (i, response) in responses.iter().enumerate() {
+                kv_store
+                    .put(
+                        format!(
+                            "rs.{pod_hash}.{request_id}.{}.{}",
+                            i,
+                            Utc::now().timestamp()
+                        ),
+                        response.clone(),
+                    )
+                    .await
+                    .unwrap();
+            }
 
-        for (i, response) in responses.iter().enumerate() {
-            kv_store
-                .put(
-                    format!("rs.{}.{}.{}", id, i, Utc::now().timestamp()),
-                    response.clone(),
-                )
-                .await
-                .unwrap();
-        }
+            // put the callbacks into the store
+            for (i, cb) in cbs.drain(..).enumerate() {
+                kv_store
+                    .put(
+                        format!(
+                            "cb.{pod_hash}.{request_id}.{}.{}",
+                            i,
+                            Utc::now().timestamp_millis()
+                        ),
+                        cb.try_into().unwrap(),
+                    )
+                    .await
+                    .unwrap();
+            }
+        });
 
         // stream response
-        let response_stream = state.stream_response(&id).await.unwrap();
+        let response_stream = state.stream_response(&id, RequestType::Sse).await.unwrap();
         let output_responses: Vec<_> = response_stream.collect().await;
         assert_eq!(output_responses.len(), 5);
     }
diff --git a/rust/serving/src/app/response.rs b/rust/serving/src/app/response.rs
index 1824d763..ffa75dc0 100644
--- a/rust/serving/src/app/response.rs
+++ b/rust/serving/src/app/response.rs
@@ -29,8 +29,8 @@ impl ServeResponse {
 pub enum ApiError {
     BadRequest(String),
     InternalServerError(String),
-    BadGateway(String),
     Conflict(String),
+    TooManyRequests(String),
 }
 
 impl IntoResponse for ApiError {
@@ -45,8 +45,8 @@ impl IntoResponse for ApiError {
         let (status, message) = match self {
             ApiError::BadRequest(message) => (StatusCode::BAD_REQUEST, message),
             ApiError::InternalServerError(message) => (StatusCode::INTERNAL_SERVER_ERROR, message),
-            ApiError::BadGateway(message) => (StatusCode::BAD_GATEWAY, message),
             ApiError::Conflict(message) => (StatusCode::CONFLICT, message),
+            ApiError::TooManyRequests(message) => (StatusCode::TOO_MANY_REQUESTS, message),
         };
 
         (
diff --git a/rust/serving/src/app/store.rs b/rust/serving/src/app/store.rs
index ffb1180e..a2cd1479 100644
--- a/rust/serving/src/app/store.rs
+++ b/rust/serving/src/app/store.rs
@@ -1,3 +1,6 @@
+/// callback store for storing the callback state
 pub(crate) mod cbstore;
-/// datastore for storing the state
+/// datastore for storing the response state
 pub(crate) mod datastore;
+/// status for tracking the request lifecycle
+pub(crate) mod status;
diff --git a/rust/serving/src/app/store/cbstore.rs b/rust/serving/src/app/store/cbstore.rs
index 9a0cc56b..90062753 100644
--- a/rust/serving/src/app/store/cbstore.rs
+++ b/rust/serving/src/app/store/cbstore.rs
@@ -3,81 +3,30 @@
 //! The progress as the message moves from one vertex to another is **streamed** and the serving
 //! host that accepted that message from the client will be **watching** that stream of callbacks.
 use std::sync::Arc;
-
-use bytes::Bytes;
-use tokio::task::JoinHandle;
 use tokio_stream::wrappers::ReceiverStream;
 
 use crate::app::store::datastore::Result as StoreResult;
 use crate::callback::Callback;
 
-/// jetstream based callback store
-pub(crate) mod jetstreamstore;
+/// JetStream based callback store
+pub(crate) mod jetstream_store;
 
 /// In-memory based callback store
 pub(crate) mod memstore;
 
-/// Represents the current processing status of a request id in the `Store`.
-#[derive(Debug, PartialEq)]
-pub(crate) enum ProcessingStatus {
-    InProgress,
-    Completed(String), // subgraph of the completed request
-    Failed(String),    // error message of the failed request
-}
-
-const IN_PROGRESS: &str = "in-progress";
-const COMPLETED_PREFIX: &str = "completed:"; // prefix because we encode extra data after `:`
-const FAILED_PREFIX: &str = "failed:";
-
-impl From<Bytes> for ProcessingStatus {
-    fn from(value: Bytes) -> Self {
-        if value == IN_PROGRESS {
-            ProcessingStatus::InProgress
-        } else if value.starts_with(COMPLETED_PREFIX.as_ref()) {
-            let subgraph = String::from_utf8(value[COMPLETED_PREFIX.len()..].to_vec())
-                .expect("Invalid UTF-8 sequence");
-            ProcessingStatus::Completed(subgraph)
-        } else if value.starts_with(FAILED_PREFIX.as_ref()) {
-            let error = String::from_utf8(value[FAILED_PREFIX.len()..].to_vec())
-                .expect("Invalid UTF-8 sequence");
-            ProcessingStatus::Failed(error)
-        } else {
-            panic!("Invalid processing status")
-        }
-    }
-}
-
-impl From<ProcessingStatus> for Bytes {
-    fn from(value: ProcessingStatus) -> Self {
-        match value {
-            ProcessingStatus::InProgress => Bytes::from_static(IN_PROGRESS.as_ref()),
-            ProcessingStatus::Completed(subgraph) => {
-                Bytes::from(format!("{}{}", COMPLETED_PREFIX, subgraph))
-            }
-            ProcessingStatus::Failed(error) => Bytes::from(format!("{}{}", FAILED_PREFIX, error)),
-        }
-    }
-}
-
 /// Store trait to store the callback information.
 #[trait_variant::make(CallbackStore: Send)]
 #[allow(dead_code)]
 pub(crate) trait LocalCallbackStore {
-    /// Register a request id in the store. If the `id` already exists in the store,
-    /// `StoreError::DuplicateRequest` error is returned.
-    async fn register(&mut self, id: &str) -> StoreResult<()>;
-    /// This method will be called when processing is completed for a request id.
-    async fn deregister(&mut self, id: &str, sub_graph: &str) -> StoreResult<()>;
-    /// This method will be called when processing is failed for a request id.
-    async fn mark_as_failed(&mut self, id: &str, error: &str) -> StoreResult<()>;
-    /// watch the callback payloads for a given request id to determine whether the
-    /// request is complete.
-    async fn watch_callbacks(
+    /// register and watch the callbacks for a given request id to determine whether the request is
+    /// complete.
+    async fn register_and_watch(
         &mut self,
         id: &str,
-    ) -> StoreResult<(ReceiverStream<Arc<Callback>>, JoinHandle<()>)>;
-    /// retrieve the processing status of a request id.
-    async fn status(&mut self, id: &str) -> StoreResult<ProcessingStatus>;
+        pod_hash: &str,
+    ) -> StoreResult<ReceiverStream<Arc<Callback>>>;
+    /// This method will be called when processing is completed for a request id.
+    async fn deregister(&mut self, id: &str) -> StoreResult<()>;
     /// check if the store is ready
     async fn ready(&mut self) -> bool;
 }
diff --git a/rust/serving/src/app/store/cbstore/jetstream_store.rs b/rust/serving/src/app/store/cbstore/jetstream_store.rs
new file mode 100644
index 00000000..8ed01cb9
--- /dev/null
+++ b/rust/serving/src/app/store/cbstore/jetstream_store.rs
@@ -0,0 +1,524 @@
+//! Stores the watches the callbacks published for the request when it travels through the vertices
+//! A central watcher task per pod monitors callbacks for that pod and sends them to the appropriate
+//! sender.
+//!
+//! JetStream Callback Entry Format is as follows:
+//!  - Callback key - `cb.{pod-hash}.{id}.{vertex_name}.{timestamp}`
+//!  - Callback value - JSON serialized Callback struct
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+
+use crate::app::store::datastore::{Error as StoreError, Result as StoreResult};
+use crate::callback::Callback;
+use async_nats::jetstream::kv::{Store, Watch};
+use async_nats::jetstream::Context;
+use tokio::sync::{mpsc, Mutex};
+use tokio_stream::wrappers::ReceiverStream;
+use tokio_stream::StreamExt;
+use tokio_util::sync::CancellationToken;
+use tracing::{debug, error, info, warn, Instrument};
+
+const CALLBACK_KEY_PREFIX: &str = "cb";
+
+/// JetStream implementation of the callback store.
+#[derive(Clone)]
+pub(crate) struct JetStreamCallbackStore {
+    pod_hash: String,
+    callback_kv: Store,
+    callback_senders: Arc<Mutex<HashMap<String, mpsc::Sender<Arc<Callback>>>>>,
+    cln_token: CancellationToken,
+}
+
+impl JetStreamCallbackStore {
+    pub(crate) async fn new(
+        js_context: Context,
+        pod_hash: &str,
+        callback_bucket: &str,
+        cln_token: CancellationToken,
+    ) -> StoreResult<Self> {
+        info!(
+            pod_hash,
+            callback_bucket, "Initializing JetStreamCallbackStore"
+        );
+
+        let callback_kv = js_context
+            .get_key_value(callback_bucket)
+            .await
+            .map_err(|e| {
+                StoreError::Connection(format!(
+                    "Failed to get callback kv store '{callback_bucket}': {e:?}"
+                ))
+            })?;
+
+        let callback_senders = Arc::new(Mutex::new(HashMap::new()));
+
+        // start the central watcher for this pod
+        Self::spawn_central_watcher(
+            pod_hash.to_string(),
+            callback_kv.clone(),
+            Arc::clone(&callback_senders),
+            cln_token.clone(),
+        )
+        .await;
+
+        Ok(Self {
+            pod_hash: pod_hash.to_string(),
+            callback_kv,
+            callback_senders,
+            cln_token,
+        })
+    }
+
+    /// spawns the central task that watches callback keys and manages callback collection for this
+    /// pod.
+    async fn spawn_central_watcher(
+        pod_hash: String,
+        kv_store: Store,
+        senders: Arc<Mutex<HashMap<String, mpsc::Sender<Arc<Callback>>>>>,
+        cln_token: CancellationToken,
+    ) {
+        let span = tracing::info_span!("central_callback_watcher", pod_hash = ?pod_hash);
+
+        // we should start from the first revision when the pod comes up, it will watch for all the
+        // callbacks for the requests that are originated from this pod.
+        let mut latest_revision = 0;
+
+        // callback key format is cb.{pod_hash}.{request_id}.{vertex_name}.{timestamp}, so we should
+        // watch for all the callbacks with prefix cb.{pod_hash}.*
+        let watch_key_pattern = format!("{CALLBACK_KEY_PREFIX}.{pod_hash}.*.*.*");
+        let mut watcher = Self::create_watcher(
+            &kv_store,
+            &watch_key_pattern,
+            latest_revision + 1,
+            cln_token.clone(),
+        )
+        .await
+        .expect("Failed to create central watcher");
+
+        tokio::spawn(async move {
+            // Main watch loop
+            loop {
+                tokio::select! {
+                // Stop watching if the cancellation token is triggered, its the callee's responsibility
+                // to cancel the token after all the inflight messages are processed (after server shutdown)
+                _ = cln_token.cancelled() => {
+                    info!("Cancellation token triggered. Stopping central watcher.");
+                    break;
+                }
+
+                // Process the next entry from the watcher
+                maybe_entry = watcher.next() => {
+                    match maybe_entry {
+                        Some(Ok(entry)) => {
+                            latest_revision = entry.revision;
+                            debug!(key = ?entry.key, operation = ?entry.operation, revision = entry.revision, "Callback central watcher received entry");
+
+                            // Skip deletions/purges, only process Put operations for new callbacks
+                            if entry.operation != async_nats::jetstream::kv::Operation::Put {
+                                continue;
+                            }
+
+                            // Parse key: cb.{pod_hash}.{id}.{vertex_name}.{timestamp}
+                            let parts: Vec<&str> = entry.key.splitn(5, '.').collect();
+                            let cb_pod_hash = parts[1];
+                            let request_id = parts[2];
+
+                            if parts.len() < 5 || parts[0] != CALLBACK_KEY_PREFIX || cb_pod_hash != pod_hash {
+                                error!(?entry, "Received unexpected key format in central watcher");
+                                continue;
+                            }
+
+                            let callback: Arc<Callback> = Arc::new(entry.value.try_into().expect("Failed to deserialize callback"));
+                            // we need to send the callback to appropriate sender based on the id
+                            let senders_guard = senders.lock().await;
+                            if let Some(cb_sender) = senders_guard.get(request_id) {
+                                cb_sender.send(callback).await.expect("Failed to send callback");
+                            } else {
+                                error!(id = ?callback.id, "No active sender found for request id. Callback not sent.");
+                                continue;
+                            }
+                        }
+                        Some(Err(e)) => {
+                            error!(error = ?e, "Error from central watcher stream. Re-establishing watch...");
+                            // when we are recreating the watcher we should start watching from the latest revision+1
+                            // because we would have successfully read till latest revision, duplicate callbacks
+                            // will result in tracker not able to decide if the request has been completed.
+                            watcher = match Self::create_watcher(&kv_store, &watch_key_pattern, latest_revision + 1, cln_token.clone()).await {
+                                Ok(stream) => stream,
+                                Err(e) => {
+                                    error!(error = ?e, "Failed to recreate watcher. Stopping scan.");
+                                    break;
+                                }
+                            };
+                        }
+                        None => {
+                            error!("Central watcher stream ended unexpectedly. Attempting to restart.");
+                            // when we are recreating the watcher we should start watching from the latest revision+1
+                            // because we would have successfully read till latest revision, duplicate callbacks
+                            // will result in tracker not able to decide if the request has been completed.
+                            watcher = match Self::create_watcher(&kv_store, &watch_key_pattern, latest_revision + 1, cln_token.clone()).await {
+                                Ok(stream) => stream,
+                                Err(e) => {
+                                    error!(?e, "Failed to recreate watcher. Stopping scan.");
+                                    break;
+                                }
+                            };
+                        }
+                    }
+                }
+            }
+            }
+        }.instrument(span));
+    }
+
+    /// Watch for all the callbacks for a given request id that was getting processed by a different pod.
+    async fn watch_historical_callbacks(
+        callback_kv: Store,
+        previous_pod_hash: String,
+        id: String,
+        tx: mpsc::Sender<Arc<Callback>>,
+        cln_token: CancellationToken,
+    ) {
+        let key_pattern = format!("{}.{}.{}.*.*", CALLBACK_KEY_PREFIX, previous_pod_hash, id);
+        let mut latest_revision = 0;
+
+        info!(request_id = ?id, ?previous_pod_hash, key_pattern, "Historical scan task started");
+
+        let mut entries_stream = match Self::create_watcher(
+            &callback_kv,
+            &key_pattern,
+            latest_revision + 1,
+            cln_token.clone(),
+        )
+        .await
+        {
+            Ok(stream) => stream,
+            Err(e) => {
+                error!(request_id = ?id, error = ?e, "Failed to initiate historical callback scan");
+                return;
+            }
+        };
+
+        loop {
+            match entries_stream.next().await {
+                Some(Ok(entry)) => {
+                    latest_revision = entry.revision;
+                    if entry.operation == async_nats::jetstream::kv::Operation::Delete {
+                        return;
+                    }
+                    match entry.value.try_into() {
+                        Ok(callback_obj) => {
+                            let callback = Arc::new(callback_obj);
+                            if tx.send(callback).await.is_err() {
+                                warn!(request_id = ?id, "Receiver dropped during historical send. Stopping scan stream.");
+                                break;
+                            }
+                        }
+                        Err(e) => {
+                            error!(request_id = ?id, key = ?entry.key, error = ?e, "Failed to deserialize historical callback during scan");
+                        }
+                    }
+                }
+                Some(Err(e)) => {
+                    error!(request_id = ?id, error = ?e, "Error iterating historical callback scan stream. Recreating watcher...");
+                    entries_stream = match Self::create_watcher(
+                        &callback_kv,
+                        &key_pattern,
+                        latest_revision + 1,
+                        cln_token.clone(),
+                    )
+                    .await
+                    {
+                        Ok(stream) => stream,
+                        Err(re_e) => {
+                            error!(request_id = ?id, error = ?re_e, "Failed to recreate watcher. Stopping scan.");
+                            break;
+                        }
+                    };
+                }
+                None => {
+                    error!(request_id = ?id, "Historical callback stream ended unexpectedly. Recreating watcher...");
+                    entries_stream = match Self::create_watcher(
+                        &callback_kv,
+                        &key_pattern,
+                        latest_revision + 1,
+                        cln_token.clone(),
+                    )
+                    .await
+                    {
+                        Ok(stream) => stream,
+                        Err(re_e) => {
+                            error!(request_id = ?id, error = ?re_e, "Failed to recreate watcher. Stopping scan.");
+                            break;
+                        }
+                    };
+                }
+            }
+        }
+        info!(?id, "Finished streaming historical callbacks.");
+    }
+
+    /// creates a kv watcher for the given key pattern and revision, and it keeps retrying until it's
+    /// successful.
+    async fn create_watcher(
+        callback_kv: &Store,
+        key_pattern: &str,
+        start_revision: u64,
+        cln_token: CancellationToken,
+    ) -> StoreResult<Watch> {
+        loop {
+            match callback_kv
+                .watch_from_revision(key_pattern, start_revision)
+                .await
+            {
+                Ok(watcher) => {
+                    info!(key_pattern, start_revision, "Watcher created successfully");
+                    return Ok(watcher);
+                }
+                Err(e) => {
+                    error!(error = ?e, "Failed to create watcher. Retrying...");
+                    tokio::time::sleep(Duration::from_millis(100)).await;
+                }
+            }
+            if cln_token.is_cancelled() {
+                error!("Cancellation token triggered, exiting watcher creation retry loop.");
+                return Err(StoreError::Connection(
+                    "Cancellation token triggered".to_string(),
+                ));
+            }
+        }
+    }
+}
+
+impl super::CallbackStore for JetStreamCallbackStore {
+    /// Register the request and watch for callbacks.
+    async fn register_and_watch(
+        &mut self,
+        id: &str,
+        pod_hash: &str,
+    ) -> StoreResult<ReceiverStream<Arc<Callback>>> {
+        let (tx, rx) = mpsc::channel(10);
+
+        // if the pod_hash is same as the current pod, we can directly register the sender because
+        // central watcher is already watching the callbacks for this pod. If not we will have to
+        // create a new watcher for watching callbacks of the different pod.
+        if pod_hash.eq(&self.pod_hash) {
+            let mut senders_guard = self.callback_senders.lock().await;
+            senders_guard.insert(id.to_string(), tx);
+            Ok(ReceiverStream::new(rx))
+        } else {
+            let callback_kv_clone = self.callback_kv.clone();
+            let id_clone = id.to_string();
+            let tx_clone = tx.clone();
+            let previous_pod_hash = pod_hash.to_string();
+            let cln_token = self.cln_token.clone();
+
+            // Spawn a separate task for historical callbacks
+            tokio::spawn(async move {
+                JetStreamCallbackStore::watch_historical_callbacks(
+                    callback_kv_clone,
+                    previous_pod_hash,
+                    id_clone,
+                    tx_clone,
+                    cln_token,
+                )
+                .await;
+            });
+            Ok(ReceiverStream::new(rx))
+        }
+    }
+
+    /// deregister the request and remove the callback sender from the map.
+    async fn deregister(&mut self, id: &str) -> StoreResult<()> {
+        let mut senders_guard = self.callback_senders.lock().await;
+        if senders_guard.remove(id).is_none() {
+            error!(?id, "No active sender found during deregistration");
+        }
+        Ok(())
+    }
+
+    async fn ready(&mut self) -> bool {
+        true
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::app::store::cbstore::CallbackStore;
+    use async_nats::jetstream;
+    use async_nats::jetstream::kv::Config;
+    use bytes::Bytes;
+    use chrono::Utc;
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_register() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let serving_store = "test_serving_store";
+
+        context
+            .create_key_value(Config {
+                bucket: serving_store.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let mut store = JetStreamCallbackStore::new(
+            context.clone(),
+            "0",
+            serving_store,
+            CancellationToken::new(),
+        )
+        .await
+        .unwrap();
+
+        let id = "AFA7E0A1-3F0A-4C1B-AB94-BDA57694648D";
+        let result = store.register_and_watch(id, "xbac").await;
+        assert!(result.is_ok());
+
+        // delete store
+        context.delete_key_value(serving_store).await.unwrap();
+    }
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_watch_callbacks() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let serving_store = "test_watch_callbacks";
+        let pod_hash = "0";
+
+        // Delete bucket so that re-running the test won't fail
+        let _ = context.delete_key_value(serving_store).await;
+
+        context
+            .create_key_value(Config {
+                bucket: serving_store.to_string(),
+                description: "test_description".to_string(),
+                history: 15,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let mut store = JetStreamCallbackStore::new(
+            context.clone(),
+            pod_hash,
+            serving_store,
+            CancellationToken::new(),
+        )
+        .await
+        .unwrap();
+
+        let id = "test_watch_id_two";
+        let mut stream = store.register_and_watch(id, pod_hash).await.unwrap();
+
+        // Simulate a callback being added to the store
+        let callback = Callback {
+            id: id.to_string(),
+            vertex: "test_vertex".to_string(),
+            cb_time: 12345,
+            from_vertex: "test_from_vertex".to_string(),
+            responses: vec![],
+        };
+        let key = format!("cb.{pod_hash}.{id}.input.{}", Utc::now().timestamp());
+        store
+            .callback_kv
+            .put(key, Bytes::from(serde_json::to_vec(&callback).unwrap()))
+            .await
+            .unwrap();
+
+        // Verify that the callback is received
+        let received_callback = stream.next().await.unwrap();
+        assert_eq!(received_callback.id, callback.id);
+        assert_eq!(received_callback.vertex, callback.vertex);
+        assert_eq!(received_callback.cb_time, callback.cb_time);
+        assert_eq!(received_callback.from_vertex, callback.from_vertex);
+        assert_eq!(received_callback.responses.len(), callback.responses.len());
+
+        // delete store
+        context.delete_key_value(serving_store).await.unwrap();
+    }
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_watch_historical_callbacks() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let serving_store = "test_watch_historical_callbacks";
+        let current_pod_hash = "xxba";
+        let previous_pod_hash = "xbzb";
+
+        // Delete bucket so that re-running the test won't fail
+        let _ = context.delete_key_value(serving_store).await;
+
+        context
+            .create_key_value(Config {
+                bucket: serving_store.to_string(),
+                description: "test_description".to_string(),
+                history: 15,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let mut store = JetStreamCallbackStore::new(
+            context.clone(),
+            current_pod_hash,
+            serving_store,
+            CancellationToken::new(),
+        )
+        .await
+        .unwrap();
+
+        let id = "test_historical_id";
+
+        // Simulate historical callbacks in the KV store
+        let callback = Callback {
+            id: id.to_string(),
+            vertex: "test_vertex".to_string(),
+            cb_time: 12345,
+            from_vertex: "test_from_vertex".to_string(),
+            responses: vec![],
+        };
+        let callback_key = format!(
+            "cb.{}.{}.input.{}",
+            previous_pod_hash,
+            id,
+            Utc::now().timestamp()
+        );
+        store
+            .callback_kv
+            .put(
+                callback_key,
+                Bytes::from(serde_json::to_vec(&callback).unwrap()),
+            )
+            .await
+            .unwrap();
+
+        // Register and watch the callbacks
+        let mut stream = store
+            .register_and_watch(id, previous_pod_hash)
+            .await
+            .unwrap();
+
+        // Verify that the historical callback is received
+        let received_callback = stream.next().await.unwrap();
+        assert_eq!(received_callback.id, callback.id);
+        assert_eq!(received_callback.vertex, callback.vertex);
+        assert_eq!(received_callback.cb_time, callback.cb_time);
+        assert_eq!(received_callback.from_vertex, callback.from_vertex);
+        assert_eq!(received_callback.responses.len(), callback.responses.len());
+
+        // Clean up the KV store
+        context.delete_key_value(serving_store).await.unwrap();
+    }
+}
diff --git a/rust/serving/src/app/store/cbstore/jetstreamstore.rs b/rust/serving/src/app/store/cbstore/jetstreamstore.rs
deleted file mode 100644
index bf8525aa..00000000
--- a/rust/serving/src/app/store/cbstore/jetstreamstore.rs
+++ /dev/null
@@ -1,465 +0,0 @@
-//! Stores the callback and status information in JetStream KV store. Each callback from the vertex
-//! will be stored as a new entry in the store. The status of the request will be updated as the
-//! processing progresses.
-//!
-//! **JetStream Callback Entry Format**
-//!
-//! Callback key - cb.{id}.{vertex_name}.{timestamp}
-//!
-//! Callback value - JSON serialized Callback struct
-//!
-//! NOTE: cb.{id}.start.processing and cb.{id}.done.processing keys are used to signal the start and
-//! end of processing so that the watchers can stop watching for new callbacks.
-//!
-//! **Status Entry Format**
-//!
-//! Status Key - status.{id}
-//!
-//! Status Value - JSON serialized [ProcessingStatus] enum
-
-use std::sync::Arc;
-use std::time::Duration;
-
-use async_nats::jetstream::kv::{CreateErrorKind, Store};
-use async_nats::jetstream::Context;
-use bytes::Bytes;
-use chrono::Utc;
-use tokio::task::JoinHandle;
-use tokio_stream::wrappers::ReceiverStream;
-use tokio_stream::StreamExt;
-use tracing::{info, Instrument};
-
-use crate::app::store::cbstore::ProcessingStatus;
-use crate::app::store::datastore::{Error as StoreError, Result as StoreResult};
-use crate::callback::Callback;
-
-/// JetStream implementation of the callback store. JetStream KV store is used to store the
-/// callback and status information. We use the Watch feature of JetStream to monitor the status
-/// of the processing.
-#[derive(Clone)]
-pub(crate) struct JetStreamCallbackStore {
-    kv_store: Store,
-}
-
-impl JetStreamCallbackStore {
-    pub(crate) async fn new(js_context: Context, bucket_name: &str) -> StoreResult<Self> {
-        let kv_store = js_context.get_key_value(bucket_name).await.map_err(|e| {
-            StoreError::Connection(format!("Failed to get kv store '{bucket_name}': {e:?}"))
-        })?;
-        Ok(Self { kv_store })
-    }
-}
-
-impl super::CallbackStore for JetStreamCallbackStore {
-    /// registers a request id in the store. If the `id` already exists in the store,
-    /// `StoreError::DuplicateRequest` error is returned.
-    async fn register(&mut self, id: &str) -> StoreResult<()> {
-        // status.{id} is the key for the status of the request, we set it to InProgress
-        let status_key = format!("status.{id}");
-        info!(id, "Registering key in Jetstream KV store");
-
-        self.kv_store
-            .create(&status_key, ProcessingStatus::InProgress.into())
-            .await
-            .map_err(|e| {
-                if e.kind() == CreateErrorKind::AlreadyExists {
-                    StoreError::DuplicateRequest(id.to_string())
-                } else {
-                    StoreError::StoreWrite(format!(
-                        "Failed to register request id {status_key} in kv store: {e:?}"
-                    ))
-                }
-            })?;
-
-        let current_timestamp = Utc::now().timestamp().to_string();
-        let response_key = format!("rs.{id}.start.processing");
-        self.kv_store
-            // we start with some key so that JetStream won't return early thinking there is
-            // nothing to watch.
-            .create(&response_key, Bytes::from(current_timestamp.clone()))
-            .await
-            .map_err(|e| {
-                if e.kind() == CreateErrorKind::AlreadyExists {
-                    StoreError::DuplicateRequest(id.to_string())
-                } else {
-                    StoreError::StoreWrite(format!(
-                        "Failed to register request id {response_key} in kv store: {e:?}"
-                    ))
-                }
-            })?;
-
-        let callbacks_key = format!("cb.{id}.start.processing");
-        self.kv_store
-            // same as above
-            .create(&callbacks_key, Bytes::from(current_timestamp))
-            .await
-            .map_err(|e| {
-                if e.kind() == CreateErrorKind::AlreadyExists {
-                    StoreError::DuplicateRequest(id.to_string())
-                } else {
-                    StoreError::StoreWrite(format!(
-                        "Failed to register request id {callbacks_key} in kv store: {e:?}"
-                    ))
-                }
-            })?;
-
-        Ok(())
-    }
-    /// de-registers a request id from the store. Updates the status of the request to completed and
-    /// stores the subgraph of the request.
-    async fn deregister(&mut self, id: &str, sub_graph: &str) -> StoreResult<()> {
-        let key = format!("status.{id}");
-        self.kv_store
-            .put(
-                key,
-                ProcessingStatus::Completed(sub_graph.to_string()).into(),
-            )
-            .await
-            .map_err(|e| {
-                StoreError::StoreWrite(format!("Failed to mark request as done in kv store: {e:?}"))
-            })?;
-
-        // we need to update the callbacks and response keys to signal the end of processing so that
-        // the watchers can stop watching for new callbacks and responses.
-        let current_timestamp = Utc::now().timestamp().to_string();
-        let callbacks_key = format!("cb.{id}.done.processing");
-        self.kv_store
-            .put(callbacks_key, Bytes::from(current_timestamp.clone()))
-            .await
-            .map_err(|e| {
-                StoreError::StoreWrite(format!("Failed to mark request as done in kv store: {e:?}"))
-            })?;
-
-        let response_key = format!("rs.{id}.done.processing");
-        self.kv_store
-            .put(response_key, Bytes::from(current_timestamp.clone()))
-            .await
-            .map_err(|e| {
-                StoreError::StoreWrite(format!("Failed to mark request as done in kv store: {e:?}"))
-            })?;
-
-        Ok(())
-    }
-
-    /// updates the status of a request in the store to failed and stores the error message.
-    async fn mark_as_failed(&mut self, id: &str, error: &str) -> StoreResult<()> {
-        let key = format!("status.{id}");
-        self.kv_store
-            .put(key, ProcessingStatus::Failed(error.to_string()).into())
-            .await
-            .map_err(|e| {
-                StoreError::StoreWrite(format!(
-                    "Failed to mark request as failed in kv store: {e:?}"
-                ))
-            })?;
-        Ok(())
-    }
-
-    /// watches for callbacks in the store for a given request id.
-    async fn watch_callbacks(
-        &mut self,
-        id: &str,
-    ) -> StoreResult<(ReceiverStream<Arc<Callback>>, JoinHandle<()>)> {
-        // the callbacks are stored in the format cb.{id}.{vertex_name}.{timestamp}
-        // so we can watch for all keys that start with cb.{id}.*
-        let callbacks_key = format!("cb.{id}.*.*");
-        let mut watcher = self
-            .kv_store
-            .watch_from_revision(&callbacks_key, 1)
-            .await
-            .map_err(|e| {
-                StoreError::StoreRead(format!("Failed to watch request id in kv store: {e:?}"))
-            })?;
-
-        let kv_store = self.kv_store.clone();
-
-        let (tx, rx) = tokio::sync::mpsc::channel(10);
-
-        let callbacks_init_key = format!("cb.{id}.start.processing");
-        let callbacks_done_key = format!("cb.{id}.done.processing");
-
-        let span = tracing::Span::current();
-
-        let callback_watcher = async move {
-            let mut received_events = false;
-            let mut attempts = 0;
-            while !received_events && attempts < 5 {
-                attempts += 1;
-                while let Some(watch_event) = watcher.next().await {
-                    received_events = true;
-                    let entry = match watch_event {
-                        Ok(event) => event,
-                        Err(e) => {
-                            tracing::error!(?e, "Received error from Jetstream KV watcher");
-                            // Recreate the watcher and start processing the events again
-                            received_events = false;
-                            break;
-                        }
-                    };
-
-                    // init key is used to signal the start of processing, the value will be the timestamp.
-                    if entry.key == callbacks_init_key {
-                        tracing::debug!(
-                            callbacks_init_key,
-                            "Received event for the init key. Ignoring it"
-                        );
-                        continue;
-                    }
-
-                    // done key is used to signal the end of processing, we can break the loop
-                    if entry.key == callbacks_done_key {
-                        tracing::debug!(
-                            callbacks_done_key,
-                            "Received event for the callback_done_key. Stopping the watcher task"
-                        );
-                        break;
-                    }
-
-                    let cbr: Callback = entry
-                        .value
-                        .try_into()
-                        .inspect_err(|err| tracing::error!(?err, "Failed to deserialize callback"))
-                        .expect("Failed to deserialize callback");
-
-                    tx.send(Arc::new(cbr))
-                        .await
-                        .inspect_err(|err| {
-                            tracing::error!(
-                                ?err,
-                                "Failed to send callback details from watcher task to main task"
-                            )
-                        })
-                        .expect("Failed to send callback");
-                }
-                if !received_events && attempts < 5 {
-                    tracing::warn!(
-                        callbacks_key,
-                        "Watcher for Jetstream key didn't return any events. Will recreate the watcher in 20ms"
-                    );
-                    tokio::time::sleep(Duration::from_millis(20)).await;
-                    watcher = kv_store
-                        .watch_from_revision(&callbacks_key, 1)
-                        .await
-                        .map_err(|e| {
-                            StoreError::StoreRead(format!(
-                                "Failed to watch request id in kv store: {e:?}"
-                            ))
-                        })
-                        .inspect_err(|err| {
-                            tracing::error!(?err, callbacks_key, "Failed to recreate the watcher")
-                        })
-                        .expect("Failed to recreate the watcher");
-                }
-            }
-            if !received_events {
-                tracing::error!(
-                    callbacks_key,
-                    "Watcher for Jetstream key didn't return any events even after retries"
-                );
-            }
-        };
-
-        let handle = tokio::spawn(callback_watcher.instrument(span));
-
-        Ok((ReceiverStream::new(rx), handle))
-    }
-
-    /// returns the status of a request in the store.
-    async fn status(&mut self, id: &str) -> StoreResult<ProcessingStatus> {
-        let key = format!("status.{id}");
-        let status = self.kv_store.get(&key).await.map_err(|e| {
-            StoreError::StoreRead(format!("Failed to get status for request id: {e:?}"))
-        })?;
-        let Some(status) = status else {
-            return Err(StoreError::InvalidRequestId(id.to_string()));
-        };
-        Ok(status.into())
-    }
-
-    async fn ready(&mut self) -> bool {
-        // we do not need to implement a health check for the JetStream connection because it is the
-        // ISB and it has to be up.
-        true
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use async_nats::jetstream;
-    use async_nats::jetstream::kv::Config;
-    use bytes::Bytes;
-    use chrono::Utc;
-
-    use super::*;
-    use crate::app::store::cbstore::CallbackStore;
-
-    #[cfg(feature = "nats-tests")]
-    #[tokio::test]
-    async fn test_register() {
-        let js_url = "localhost:4222";
-        let client = async_nats::connect(js_url).await.unwrap();
-        let context = jetstream::new(client);
-        let serving_store = "test_serving_store";
-
-        context
-            .create_key_value(Config {
-                bucket: serving_store.to_string(),
-                history: 1,
-                ..Default::default()
-            })
-            .await
-            .unwrap();
-
-        let mut store = JetStreamCallbackStore::new(context.clone(), serving_store)
-            .await
-            .unwrap();
-
-        let id = "AFA7E0A1-3F0A-4C1B-AB94-BDA57694648D";
-        let result = store.register(id).await;
-        assert!(result.is_ok());
-
-        // delete store
-        context.delete_key_value(serving_store).await.unwrap();
-    }
-
-    #[cfg(feature = "nats-tests")]
-    #[tokio::test]
-    async fn test_watch_callbacks() {
-        let js_url = "localhost:4222";
-        let client = async_nats::connect(js_url).await.unwrap();
-        let context = jetstream::new(client);
-        let serving_store = "test_watch_callbacks";
-
-        // Delete bucket so that re-running the test won't fail
-        let _ = context.delete_key_value(serving_store).await;
-
-        context
-            .create_key_value(Config {
-                bucket: serving_store.to_string(),
-                description: "test_description".to_string(),
-                history: 15,
-                ..Default::default()
-            })
-            .await
-            .unwrap();
-
-        let mut store = JetStreamCallbackStore::new(context.clone(), serving_store)
-            .await
-            .unwrap();
-
-        let id = "test_watch_id_two";
-        store.register(id).await.unwrap();
-
-        let (mut rx, handle) = store.watch_callbacks(id).await.unwrap();
-
-        // Simulate a callback being added to the store
-        let callback = Callback {
-            id: id.to_string(),
-            vertex: "test_vertex".to_string(),
-            cb_time: 12345,
-            from_vertex: "test_from_vertex".to_string(),
-            responses: vec![],
-        };
-        let key = format!("cb.{id}.0.{}", Utc::now().timestamp());
-        store
-            .kv_store
-            .put(key, Bytes::from(serde_json::to_vec(&callback).unwrap()))
-            .await
-            .unwrap();
-
-        // Verify that the callback is received
-        let received_callback = rx.next().await.unwrap();
-        assert_eq!(received_callback.id, callback.id);
-        assert_eq!(received_callback.vertex, callback.vertex);
-        assert_eq!(received_callback.cb_time, callback.cb_time);
-        assert_eq!(received_callback.from_vertex, callback.from_vertex);
-        assert_eq!(received_callback.responses.len(), callback.responses.len());
-
-        handle.abort();
-
-        // delete store
-        context.delete_key_value(serving_store).await.unwrap();
-    }
-
-    #[cfg(feature = "nats-tests")]
-    #[tokio::test]
-    async fn test_deregister() {
-        let js_url = "localhost:4222";
-        let client = async_nats::connect(js_url).await.unwrap();
-        let context = jetstream::new(client);
-        let serving_store = "test_deregister_store";
-
-        // Delete bucket so that re-running the test won't fail
-        let _ = context.delete_key_value(serving_store).await;
-
-        context
-            .create_key_value(Config {
-                bucket: serving_store.to_string(),
-                description: "test_description".to_string(),
-                history: 15,
-                ..Default::default()
-            })
-            .await
-            .unwrap();
-
-        let mut store = JetStreamCallbackStore::new(context.clone(), serving_store)
-            .await
-            .unwrap();
-
-        let id = "test_deregister_id";
-        store.register(id).await.unwrap();
-
-        let sub_graph = "test_sub_graph";
-        let result = store.deregister(id, sub_graph).await;
-        assert!(result.is_ok());
-
-        // Verify that the status is marked as completed
-        let status = store.status(id).await.unwrap();
-        assert!(matches!(status, ProcessingStatus::Completed(_)));
-
-        // delete store
-        context.delete_key_value(serving_store).await.unwrap();
-    }
-
-    #[cfg(feature = "nats-tests")]
-    #[tokio::test]
-    async fn test_mark_as_failed() {
-        let js_url = "localhost:4222";
-        let client = async_nats::connect(js_url).await.unwrap();
-        let context = jetstream::new(client);
-        let serving_store = "test_mark_as_failed_store";
-
-        // Delete bucket so that re-running the test won't fail
-        let _ = context.delete_key_value(serving_store).await;
-
-        context
-            .create_key_value(Config {
-                bucket: serving_store.to_string(),
-                description: "test_description".to_string(),
-                history: 15,
-                ..Default::default()
-            })
-            .await
-            .unwrap();
-
-        let mut store = JetStreamCallbackStore::new(context.clone(), serving_store)
-            .await
-            .unwrap();
-
-        assert!(store.ready().await);
-
-        let id = "test_mark_as_failed_id";
-        store.register(id).await.unwrap();
-
-        let error_message = "test_error_message";
-        let result = store.mark_as_failed(id, error_message).await;
-        assert!(result.is_ok());
-
-        // Verify that the status is marked as failed
-        let status = store.status(id).await.unwrap();
-        assert!(matches!(status, ProcessingStatus::Failed(_)));
-
-        // delete store
-        context.delete_key_value(serving_store).await.unwrap();
-    }
-}
diff --git a/rust/serving/src/app/store/cbstore/memstore.rs b/rust/serving/src/app/store/cbstore/memstore.rs
index db091eae..408cbffe 100644
--- a/rust/serving/src/app/store/cbstore/memstore.rs
+++ b/rust/serving/src/app/store/cbstore/memstore.rs
@@ -1,14 +1,12 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
+use crate::app::store::datastore::{Error as StoreError, Result as StoreResult};
+use crate::callback::Callback;
 use tokio::sync::mpsc;
 use tokio::sync::Mutex;
-use tokio::task::JoinHandle;
 use tokio_stream::wrappers::ReceiverStream;
 
-use crate::app::store::datastore::{Error as StoreError, Result as StoreResult};
-use crate::callback::Callback;
-
 /// An in-memory implementation of the callback store. Only used for testing.
 #[derive(Clone)]
 #[allow(dead_code)]
@@ -26,20 +24,9 @@ impl InMemoryCallbackStore {
 }
 
 impl super::CallbackStore for InMemoryCallbackStore {
-    /// Register a request id in the store. If the `id` already exists in the store,
-    /// `StoreError::DuplicateRequest` error is returned.
-    async fn register(&mut self, id: &str) -> StoreResult<()> {
-        let mut data = self.data.lock().await;
-        if data.contains_key(id) {
-            return Err(StoreError::DuplicateRequest(id.to_string()));
-        }
-        data.insert(id.to_string(), Vec::new());
-        Ok(())
-    }
-
     /// De-register a request id from the store. If the `id` does not exist in the store,
     /// `StoreError::InvalidRequestId` error is returned.
-    async fn deregister(&mut self, id: &str, _sub_graph: &str) -> StoreResult<()> {
+    async fn deregister(&mut self, id: &str) -> StoreResult<()> {
         let mut data = self.data.lock().await;
         if data.remove(id).is_none() {
             return Err(StoreError::InvalidRequestId(id.to_string()));
@@ -47,16 +34,16 @@ impl super::CallbackStore for InMemoryCallbackStore {
         Ok(())
     }
 
-    async fn mark_as_failed(&mut self, _id: &str, _error: &str) -> StoreResult<()> {
-        Ok(())
-    }
-
-    async fn watch_callbacks(
+    async fn register_and_watch(
         &mut self,
         id: &str,
-    ) -> StoreResult<(ReceiverStream<Arc<Callback>>, JoinHandle<()>)> {
+        _pod_hash: &str,
+    ) -> StoreResult<ReceiverStream<Arc<Callback>>> {
+        let mut data = self.data.lock().await;
+        if !data.contains_key(id) {
+            data.insert(id.to_string(), Vec::new());
+        }
         let (tx, rx) = mpsc::channel(10);
-        let data = self.data.lock().await;
         if let Some(callbacks) = data.get(id) {
             for callback in callbacks {
                 tx.send(Arc::clone(callback))
@@ -66,15 +53,7 @@ impl super::CallbackStore for InMemoryCallbackStore {
         } else {
             return Err(StoreError::InvalidRequestId(id.to_string()));
         }
-        Ok((ReceiverStream::new(rx), tokio::spawn(async {})))
-    }
-
-    async fn status(&mut self, _id: &str) -> StoreResult<super::ProcessingStatus> {
-        let data = self.data.lock().await;
-        if data.get(_id).is_none() {
-            return Err(StoreError::InvalidRequestId(_id.to_string()));
-        }
-        Ok(super::ProcessingStatus::InProgress)
+        Ok(ReceiverStream::new(rx))
     }
 
     async fn ready(&mut self) -> bool {
@@ -90,7 +69,7 @@ mod tests {
     use tokio_stream::StreamExt;
 
     use super::*;
-    use crate::app::store::cbstore::{CallbackStore, ProcessingStatus};
+    use crate::app::store::cbstore::CallbackStore;
     use crate::callback::Callback;
 
     fn create_test_store() -> InMemoryCallbackStore {
@@ -106,28 +85,16 @@ mod tests {
         InMemoryCallbackStore::new(Some(callback_map))
     }
 
-    #[tokio::test]
-    async fn test_register() {
-        let mut store = create_test_store();
-        let id = "new_test_id";
-        let result = store.register(id).await;
-        assert!(result.is_ok());
-
-        // Try to register the same id again, should return an error
-        let result = store.register(id).await;
-        assert!(matches!(result, Err(StoreError::DuplicateRequest(_))));
-    }
-
     #[tokio::test]
     async fn test_deregister() {
         let mut store = create_test_store();
         let id = "test_id";
 
-        let result = store.deregister(id, "sub_graph").await;
+        let result = store.deregister(id).await;
         assert!(result.is_ok());
 
         // Try to deregister the same id again, should return an error
-        let result = store.deregister(id, "sub_graph").await;
+        let result = store.deregister(id).await;
         assert!(matches!(result, Err(StoreError::InvalidRequestId(_))));
     }
 
@@ -136,20 +103,11 @@ mod tests {
         let mut store = create_test_store();
         let id = "test_id";
 
-        let (mut rx, _handle) = store.watch_callbacks(id).await.unwrap();
+        let mut rx = store.register_and_watch(id, "0").await.unwrap();
         let received_callback = rx.next().await.unwrap();
         assert_eq!(received_callback.id, "test_id");
         assert_eq!(received_callback.vertex, "vertex");
         assert_eq!(received_callback.cb_time, 12345);
         assert_eq!(received_callback.from_vertex, "from_vertex");
     }
-
-    #[tokio::test]
-    async fn test_status() {
-        let mut store = create_test_store();
-        let id = "test_id";
-
-        let status = store.status(id).await.unwrap();
-        assert!(matches!(status, ProcessingStatus::InProgress));
-    }
 }
diff --git a/rust/serving/src/app/store/datastore.rs b/rust/serving/src/app/store/datastore.rs
index 931e5f5d..980de22d 100644
--- a/rust/serving/src/app/store/datastore.rs
+++ b/rust/serving/src/app/store/datastore.rs
@@ -3,7 +3,6 @@ use std::sync::Arc;
 
 use bytes::Bytes;
 use thiserror::Error;
-use tokio::task::JoinHandle;
 use tokio_stream::wrappers::ReceiverStream;
 
 /// JetStream based datum store
@@ -23,14 +22,12 @@ pub(crate) enum Error {
     #[error("Request id {0} doesn't exist in store")]
     InvalidRequestId(String),
 
+    #[allow(dead_code)]
     #[error("Request id {0} already exists in the store")]
     DuplicateRequest(String),
 
     #[error("Reading from the store: {0}")]
     StoreRead(String),
-
-    #[error("Writing payload to the store: {0}")]
-    StoreWrite(String),
 }
 
 impl From<Error> for crate::Error {
@@ -45,13 +42,12 @@ pub(crate) type Result<T> = std::result::Result<T, Error>;
 #[trait_variant::make(DataStore: Send)]
 #[allow(dead_code)]
 pub(crate) trait LocalDataStore {
-    /// retrieve the data from the store
-    async fn retrieve_data(&mut self, id: &str) -> Result<Vec<Vec<u8>>>;
+    /// retrieve the data from the store.
+    /// FIXME: make pod_hash optional
+    async fn retrieve_data(&mut self, id: &str, pod_hash: Option<&str>) -> Result<Vec<Vec<u8>>>;
     /// streams the data from the store
-    async fn stream_data(
-        &mut self,
-        id: &str,
-    ) -> Result<(ReceiverStream<Arc<Bytes>>, JoinHandle<()>)>;
+    async fn stream_data(&mut self, id: &str, pod_hash: &str)
+        -> Result<ReceiverStream<Arc<Bytes>>>;
     /// check if the store is ready
     async fn ready(&mut self) -> bool;
 }
diff --git a/rust/serving/src/app/store/datastore/inmemory.rs b/rust/serving/src/app/store/datastore/inmemory.rs
index 94e90b49..bc72e63c 100644
--- a/rust/serving/src/app/store/datastore/inmemory.rs
+++ b/rust/serving/src/app/store/datastore/inmemory.rs
@@ -2,7 +2,6 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use bytes::Bytes;
-use tokio::task::JoinHandle;
 use tokio_stream::wrappers::ReceiverStream;
 
 use crate::app::store::datastore::{Error as StoreError, Result as StoreResult};
@@ -30,7 +29,11 @@ impl InMemoryDataStore {
 impl super::DataStore for InMemoryDataStore {
     /// Retrieves data for a given id from the `HashMap`.
     /// Each piece of data is deserialized from bytes into a `String`.
-    async fn retrieve_data(&mut self, id: &str) -> StoreResult<Vec<Vec<u8>>> {
+    async fn retrieve_data(
+        &mut self,
+        id: &str,
+        _pod_hash: Option<&str>,
+    ) -> StoreResult<Vec<Vec<u8>>> {
         let id = format!("{id}_{STORE_KEY_SUFFIX}");
         let data = self.data.lock().await;
         match data.get(&id) {
@@ -46,7 +49,8 @@ impl super::DataStore for InMemoryDataStore {
     async fn stream_data(
         &mut self,
         id: &str,
-    ) -> StoreResult<(ReceiverStream<Arc<Bytes>>, JoinHandle<()>)> {
+        _pod_hash: &str,
+    ) -> StoreResult<ReceiverStream<Arc<Bytes>>> {
         let (tx, rx) = tokio::sync::mpsc::channel(10);
         let data = self.data.lock().await;
         if let Some(response) = data.get(id) {
@@ -56,7 +60,7 @@ impl super::DataStore for InMemoryDataStore {
                     .map_err(|_| StoreError::StoreRead("Failed to send datum".to_string()))?;
             }
         }
-        Ok((ReceiverStream::new(rx), tokio::task::spawn(async {})))
+        Ok(ReceiverStream::new(rx))
     }
 
     async fn ready(&mut self) -> bool {
@@ -85,7 +89,7 @@ mod tests {
     async fn test_retrieve_datum() {
         let mut store = create_test_store();
         let id = "test_id";
-        let result = store.retrieve_data(id).await.unwrap();
+        let result = store.retrieve_data(id, None).await.unwrap();
         assert!(result.len() > 0);
         assert_eq!(result[0], b"test_payload");
     }
@@ -94,7 +98,7 @@ mod tests {
     async fn test_retrieve_datum_not_found() {
         let mut store = create_test_store();
         let id = "non_existent_id";
-        let result = store.retrieve_data(id).await;
+        let result = store.retrieve_data(id, None).await;
         assert!(matches!(result, Err(StoreError::InvalidRequestId(_))));
     }
 
@@ -102,7 +106,7 @@ mod tests {
     async fn test_stream_response() {
         let mut store = create_test_store();
         let id = "test_id_saved";
-        let (mut rx, _handle) = store.stream_data(id).await.unwrap();
+        let mut rx = store.stream_data(id, "0").await.unwrap();
         let received_response = rx.next().await.unwrap();
         assert_eq!(
             received_response,
diff --git a/rust/serving/src/app/store/datastore/jetstream.rs b/rust/serving/src/app/store/datastore/jetstream.rs
index 76526a50..aa234a4f 100644
--- a/rust/serving/src/app/store/datastore/jetstream.rs
+++ b/rust/serving/src/app/store/datastore/jetstream.rs
@@ -1,135 +1,541 @@
 //! Stores the response from the processing in a JetStream key-value store. Each response is stored
 //! as a new key-value pair in the store, so that no responses are lost. We cannot use the same key
 //! and overwrite the value, since the history of the kv store is not per key, but for the entire
-//! store.
+//! store. Store is capable of storing multiple results for the same request so it is stored in a
+//! concatenated format (length-prefixed format).
 //!
-//! **JetStream Response Entry Format**
-//!
-//! Response Key - rs.{id}.{vertex_name}.{timestamp}
-//!
-//! Response Value - response_payload
-use std::sync::Arc;
+//! JetStream Response Entry Format is as follows:
+//!  - Response Key - rs.{pod_hash}.{id}.{vertex_name}.{timestamp}
+//!  - Response Value - Actual response payload as Bytes
 
-use async_nats::jetstream::kv::Store;
+use crate::app::store::datastore::{DataStore, Error as StoreError, Result as StoreResult};
+use crate::config::RequestType;
+use async_nats::jetstream::kv::{Store, Watch};
 use async_nats::jetstream::Context;
-use bytes::Bytes;
+use bytes::{Buf, BufMut, Bytes};
+use std::collections::hash_map::Entry;
+use std::collections::HashMap;
+use std::io::Cursor;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::mpsc;
+use tokio::sync::mpsc::Sender;
+use tokio::sync::Mutex;
 use tokio::task::JoinHandle;
+use tokio::time::sleep;
 use tokio_stream::wrappers::ReceiverStream;
 use tokio_stream::StreamExt;
+use tokio_util::sync::CancellationToken;
+use tracing::{debug, error, info, warn};
+
+const RESPONSE_KEY_PREFIX: &str = "rs";
+const FINAL_RESULT_KEY_SUFFIX: &str = "final.result.processed";
+const START_PROCESSING_MARKER: &str = "start.processing";
+const DONE_PROCESSING_MARKER: &str = "done.processing";
+
+/// Merges a list of Bytes using u64 length prefixing.
+fn merge_bytes_list(list: &[Bytes]) -> Bytes {
+    let mut accumulator = Vec::new();
+    for b in list {
+        let len = b.len() as u64;
+        accumulator.put_u64_le(len);
+        accumulator.put_slice(b);
+    }
+    Bytes::from(accumulator)
+}
 
-use crate::app::store::datastore::{DataStore, Error as StoreError, Result as StoreResult};
+/// Extracts a list of Bytes merged with u64 length prefixing.
+fn extract_bytes_list(merged: &Bytes) -> Result<Vec<Bytes>, String> {
+    let mut extracted = Vec::new();
+    let mut cursor = Cursor::new(merged);
+
+    while cursor.has_remaining() {
+        // get size/length first
+        if cursor.remaining() < size_of::<u64>() {
+            if cursor.remaining() == 0 {
+                break;
+            } else {
+                return Err(format!(
+                    "Truncated data: not enough bytes for length prefix. Remaining: {}",
+                    cursor.remaining()
+                ));
+            }
+        }
+        // get the data for the given len
+        let len = cursor.get_u64_le();
+        if cursor.remaining() < len as usize {
+            return Err(format!(
+                "Truncated data: declared payload length {} exceeds remaining bytes {}",
+                len,
+                cursor.remaining()
+            ));
+        }
+        // extract the data
+        let current_pos = cursor.position() as usize;
+        let end_pos = current_pos + (len as usize);
+        extracted.push(merged.slice(current_pos..end_pos));
+
+        cursor.advance(len as usize);
+    }
+
+    Ok(extracted)
+}
 
 /// A JetStream implementation of the data store.
 #[derive(Clone)]
 pub(crate) struct JetStreamDataStore {
     kv_store: Store,
+    pod_hash: String,
+    responses_map: Arc<Mutex<HashMap<String, ResponseMode>>>,
+    cln_token: CancellationToken,
 }
 
 impl JetStreamDataStore {
-    pub(crate) async fn new(js_context: Context, datum_store_name: &str) -> StoreResult<Self> {
+    pub(crate) async fn new(
+        js_context: Context,
+        datum_store_name: &str,
+        pod_hash: &str,
+        cln_token: CancellationToken,
+    ) -> StoreResult<Self> {
         let kv_store = js_context
             .get_key_value(datum_store_name)
             .await
             .map_err(|e| StoreError::Connection(format!("Failed to get datum kv store: {e:?}")))?;
-        Ok(Self { kv_store })
-    }
-}
 
-impl DataStore for JetStreamDataStore {
-    /// Retrieve a data from the store. If the datum is not found, `None` is returned.
-    async fn retrieve_data(&mut self, id: &str) -> StoreResult<Vec<Vec<u8>>> {
-        // the responses in kv bucket are stored in the format rs.{id}.{vertex_name}.{timestamp}
-        // so we should watch for all keys that start with rs.{id}
-        let watch_key = format!("rs.{id}.*.*");
-        let mut watcher = self
-            .kv_store
-            .watch_from_revision(&watch_key, 1)
-            .await
-            .map_err(|e| {
-                StoreError::StoreRead(format!("Failed to watch request id in kv store: {e:?}"))
-            })?;
+        let responses_map = Arc::new(Mutex::new(HashMap::new()));
+        // start the central watcher
+        Self::spawn_central_watcher(
+            pod_hash.to_string(),
+            kv_store.clone(),
+            Arc::clone(&responses_map),
+            cln_token.clone(),
+        )
+        .await;
+        Ok(Self {
+            kv_store,
+            pod_hash: pod_hash.to_string(),
+            responses_map,
+            cln_token,
+        })
+    }
 
-        let response_init_key = format!("rs.{id}.start.processing");
-        let response_done_key = format!("rs.{id}.done.processing");
+    /// Spawns the central task that watches response keys and manages data collection.
+    async fn spawn_central_watcher(
+        pod_hash: String,
+        kv_store: Store,
+        responses_map: Arc<Mutex<HashMap<String, ResponseMode>>>,
+        cln_token: CancellationToken,
+    ) -> JoinHandle<()> {
+        // watch for all the responses with prefix rs.{pod_hash}.* (ie, this pod)
+        let watch_pattern = format!("{}.{}.*.*.*", RESPONSE_KEY_PREFIX, pod_hash);
+        // monotonic revision which is reset when pod is restarted
+        let mut latest_revision = 0;
+
+        let mut watcher = Self::create_watcher(
+            &kv_store,
+            &watch_pattern,
+            latest_revision + 1,
+            cln_token.clone(),
+        )
+        .await
+        .expect("Failed to create central result watcher");
+
+        tokio::spawn(async move {
+            loop {
+                tokio::select! {
+                    // Stop watching if the cancellation token is triggered
+                    _ = cln_token.cancelled() => {
+                        info!("Cancellation token triggered. Stopping central watcher.");
+                        break;
+                    }
+                    // Process the next entry from the watcher
+                    maybe_entry = watcher.next() => {
+                        match maybe_entry {
+                            Some(Ok(entry)) => {
+                                debug!(key = ?entry.key, operation = ?entry.operation, revision = entry.revision, "Response central watcher received entry");
+                                latest_revision = entry.revision;
+                                if entry.operation != async_nats::jetstream::kv::Operation::Put {
+                                    continue;
+                                }
+
+                                if let Err(e) = Self::process_entry(&entry, &responses_map, &pod_hash, &kv_store).await {
+                                    error!(error = ?e, "Failed to process entry");
+                                }
+                            }
+                            Some(Err(e)) => {
+                                error!(error = ?e, "Error from watcher stream. Re-establishing...");
+                                watcher = match Self::create_watcher(&kv_store, &watch_pattern, latest_revision + 1, cln_token.clone()).await {
+                                    Ok(w) => w,
+                                    Err(re_e) => {
+                                        error!(error = ?re_e, "Failed to re-establish watcher. Exiting...");
+                                        break;
+                                    }
+                                };
+                            }
+                            None => {
+                                error!("Watcher stream ended unexpectedly. Re-establishing...");
+                                watcher = match Self::create_watcher(&kv_store, &watch_pattern, latest_revision + 1, cln_token.clone()).await {
+                                    Ok(w) => w,
+                                    Err(re_e) => {
+                                        error!(error = ?re_e, "Failed to re-establish watcher. Exiting...");
+                                        break;
+                                    }
+                                };
+                            }
+                        }
+                    }
+                }
+            }
+        })
+    }
 
-        let mut results = Vec::new();
-        while let Some(watch_event) = watcher.next().await {
-            let entry = match watch_event {
-                Ok(event) => event,
+    /// Creates a new kv watcher for the store with the given watch pattern and revision.
+    async fn create_watcher(
+        kv_store: &Store,
+        watch_pattern: &str,
+        start_revision: u64,
+        cln_token: CancellationToken,
+    ) -> StoreResult<Watch> {
+        loop {
+            match kv_store
+                .watch_from_revision(watch_pattern, start_revision)
+                .await
+            {
+                Ok(watcher) => {
+                    info!(
+                        watch_pattern,
+                        start_revision, "Watcher created successfully"
+                    );
+                    return Ok(watcher);
+                }
                 Err(e) => {
-                    tracing::error!(?e, "Received error from Jetstream KV watcher");
-                    continue;
+                    error!(error = ?e, "Failed to create watcher. Retrying...");
+                    tokio::time::sleep(Duration::from_millis(100)).await;
                 }
-            };
-
-            // init key is used to signal the start of processing, the value will be empty
-            // we can skip the init key
-            if entry.key == response_init_key {
-                continue;
             }
-
-            // done key is used to signal the end of processing, we can break the loop
-            if entry.key == response_done_key {
-                break;
+            if cln_token.is_cancelled() {
+                info!("Cancellation token triggered, exiting response watcher creation retry loop");
+                return Err(StoreError::Connection(
+                    "Cancellation token triggered".to_string(),
+                ));
             }
+        }
+    }
+
+    /// process every entry from the response store central watcher.
+    async fn process_entry(
+        entry: &async_nats::jetstream::kv::Entry,
+        responses_map: &Arc<Mutex<HashMap<String, ResponseMode>>>,
+        pod_hash: &str,
+        kv_store: &Store,
+    ) -> Result<(), String> {
+        let parts: Vec<&str> = entry.key.splitn(5, '.').collect();
+        let key_prefix = parts[0];
+        let request_id = parts[2];
+        let key_suffix = format!("{}.{}", parts[3], parts[4]);
+
+        // key format is rs.{pod_hash}.{request_id}.{key_suffix}, anything else is not valid
+        if parts.len() != 5 || key_prefix != RESPONSE_KEY_PREFIX || parts[1] != pod_hash {
+            return Err(format!("Unexpected key format: {}", entry.key));
+        }
 
-            results.push(entry.value.to_vec());
+        // If it's start of the request, we need to add it to the map for tracking the responses
+        // for the given id. The value of start processing marker will be the type of the request
+        // based on that we can add the appropriate entry in the map.
+        if key_suffix == START_PROCESSING_MARKER {
+            let request_type = RequestType::try_from(entry.value.clone())
+                .expect("Failed to convert entry value to RequestType");
+
+            let mut response_map = responses_map.lock().await;
+            let entry = response_map.entry(request_id.to_string());
+            // If the entry is vacant, we can insert the appropriate entry
+            // else (if map contains the entry) we need to return an error
+            if let Entry::Vacant(vacant_entry) = entry {
+                match request_type {
+                    // For SSE, create a channel and insert it
+                    RequestType::Sse => {
+                        let (tx, rx) = mpsc::channel(10);
+                        vacant_entry.insert(ResponseMode::Stream {
+                            tx,
+                            rx: Some(ReceiverStream::new(rx)),
+                        });
+                    }
+                    // For Sync and Async, insert an empty list
+                    RequestType::Sync | RequestType::Async => {
+                        vacant_entry.insert(ResponseMode::Unary(Vec::new()));
+                    }
+                }
+            } else {
+                error!(
+                    ?request_id,
+                    "Start processing marker already exists for this request ID"
+                );
+                return Err(format!(
+                    "Start processing marker already exists for request ID: {}",
+                    request_id
+                ));
+            }
+        } else if key_suffix == DONE_PROCESSING_MARKER {
+            // when we get the done processing marker, it means we won't get any more responses for
+            // this request id. So we need to merge the responses and put them in the response store, so
+            // that it can be retrieved when needed. We only need to store the merged response for
+            // sync and async because for sse it will be streaming. It's safe to store for sync request
+            // because during failure of sync request it will be converted as an async request.
+            let mut response_map = responses_map.lock().await;
+            if let Some(ResponseMode::Unary(responses)) = response_map.remove(request_id) {
+                // we use rs.{pod_hash}.{request_id}.final.result.processed as the key for the final result
+                // and the value will be the merged response.
+                let merged_response = merge_bytes_list(&responses);
+                let final_result_key =
+                    format!("{}.{}", RESPONSE_KEY_PREFIX, FINAL_RESULT_KEY_SUFFIX);
+                kv_store
+                    .put(final_result_key, merged_response)
+                    .await
+                    .map_err(|e| format!("Failed to put final result: {:?}", e))?;
+            }
+        } else {
+            // we can append the response for the list for sync and async, for sse we can write the
+            // response to the response channel
+            let mut response_map = responses_map.lock().await;
+            if let Some(response_mode) = response_map.get_mut(request_id) {
+                match response_mode {
+                    ResponseMode::Stream { tx, .. } => {
+                        tx.send(Arc::new(entry.value.clone()))
+                            .await
+                            .map_err(|e| format!("Failed to send SSE response: {:?}", e))?;
+                    }
+                    ResponseMode::Unary(responses) => {
+                        responses.push(entry.value.clone());
+                    }
+                }
+            }
         }
+        Ok(())
+    }
 
-        Ok(results)
+    /// Retrieves the final result from the KV store for a given request ID using the final key, it
+    /// keeps retrying until the final result is available.
+    async fn get_data_from_response_store(&self, id: &str) -> StoreResult<Vec<Vec<u8>>> {
+        let final_result_key = format!("{}.{}", RESPONSE_KEY_PREFIX, FINAL_RESULT_KEY_SUFFIX);
+
+        loop {
+            match self.kv_store.get(&final_result_key).await {
+                Ok(Some(result)) => {
+                    return extract_bytes_list(&result)
+                        .map(|extracted| extracted.iter().map(|b| b.to_vec()).collect())
+                        .map_err(|e| {
+                            StoreError::StoreRead(format!("Failed to extract bytes list: {e:?}"))
+                        });
+                }
+                Ok(None) => {
+                    warn!(?id, "No final result key found in KV store, retrying...");
+                    tokio::time::sleep(Duration::from_millis(5)).await;
+                }
+                Err(e) => {
+                    return Err(StoreError::StoreRead(format!(
+                        "Failed to get final result key from KV store: {e:?}"
+                    )));
+                }
+            }
+        }
     }
 
-    /// Stream the response from the store. The response is streamed as it is added to the store.
-    async fn stream_data(
-        &mut self,
+    /// Retrieves the historic response for a given request ID and pod hash. It uses a watcher to
+    /// stream the responses from the KV store. The responses are collected until the done processing
+    /// marker is received.
+    /// This is used during a failover case (pod processing the request is different the one that
+    /// originated the request).
+    pub(crate) async fn get_historic_response(
+        store: Store,
         id: &str,
-    ) -> StoreResult<(ReceiverStream<Arc<Bytes>>, JoinHandle<()>)> {
-        // the responses in kv bucket are stored in the format rs.{id}.{vertex_name}.{timestamp}
-        // so we should watch for all keys that start with rs.{id}
-        let watch_key = format!("rs.{id}.*.*");
-
-        let mut watcher = self
-            .kv_store
-            .watch_from_revision(&watch_key, 1)
-            .await
-            .map_err(|e| {
-                StoreError::StoreRead(format!("Failed to watch request id in kv store: {e:?}"))
-            })?;
-
-        let (tx, rx) = tokio::sync::mpsc::channel(10);
-        let response_init_key = format!("rs.{id}.start.processing");
-        let response_done_key = format!("rs.{id}.done.processing");
-
-        let handle = tokio::spawn(async move {
-            while let Some(watch_event) = watcher.next().await {
-                let entry = match watch_event {
-                    Ok(event) => event,
-                    Err(e) => {
-                        tracing::error!(?e, "Received error from Jetstream KV watcher");
+        pod_hash: &str,
+        cln_token: CancellationToken,
+    ) -> StoreResult<Vec<Vec<u8>>> {
+        let mut latest_revision = 0;
+        let watch_pattern = format!("{RESPONSE_KEY_PREFIX}.{pod_hash}.{id}.*.*");
+        let mut watcher = Self::create_watcher(
+            &store,
+            &watch_pattern,
+            latest_revision + 1,
+            cln_token.clone(),
+        )
+        .await?;
+        let mut responses = vec![];
+        loop {
+            match watcher.next().await {
+                None => {
+                    error!(request_id = ?id, "Historical watcher stream ended unexpectedly. Recreating watcher...");
+                    watcher = Self::create_watcher(
+                        &store,
+                        &watch_pattern,
+                        latest_revision + 1,
+                        cln_token.clone(),
+                    )
+                    .await?;
+                }
+                Some(Ok(entry)) => {
+                    latest_revision = entry.revision;
+                    if entry.key.contains(START_PROCESSING_MARKER) {
                         continue;
                     }
-                };
 
-                // init key is used to signal the start of processing, the value will be empty
-                // we can skip the init key
-                if entry.key == response_init_key {
-                    continue;
+                    if entry.key.contains(DONE_PROCESSING_MARKER) {
+                        break;
+                    }
+                    responses.push(entry.value.to_vec());
+                }
+                Some(Err(e)) => {
+                    error!(request_id = ?id, "Error while receiving response from watcher: {e:?}, Recreating watcher...");
+                    watcher = Self::create_watcher(
+                        &store,
+                        &watch_pattern,
+                        latest_revision + 1,
+                        cln_token.clone(),
+                    )
+                    .await?;
                 }
+            }
+        }
+        Ok(responses)
+    }
 
-                // done key is used to signal the end of processing, we can break the loop
-                if entry.key == response_done_key {
-                    break;
+    /// Streams the historic responses for a given request ID and pod hash. It uses a watcher to
+    /// stream the responses from the KV store. The responses are sent to the provided channel
+    /// until the done processing marker is received.
+    /// This is like [get_historic_response] except that this is for SSE.
+    pub(crate) async fn stream_historic_responses(
+        store: Store,
+        id: &str,
+        pod_hash: &str,
+        tx: Sender<Arc<Bytes>>,
+        cln_token: CancellationToken,
+    ) {
+        let mut latest_revision = 0;
+        let watch_pattern = format!("{RESPONSE_KEY_PREFIX}.{pod_hash}.{id}.*.*");
+        let mut watcher = Self::create_watcher(
+            &store,
+            &watch_pattern,
+            latest_revision + 1,
+            cln_token.clone(),
+        )
+        .await
+        .expect("Failed to create watcher");
+        loop {
+            match watcher.next().await {
+                None => {
+                    error!(request_id = ?id, "Historical watcher stream ended unexpectedly. Recreating watcher...");
+                    watcher = Self::create_watcher(
+                        &store,
+                        &watch_pattern,
+                        latest_revision + 1,
+                        cln_token.clone(),
+                    )
+                    .await
+                    .expect("Failed to create watcher");
                 }
+                Some(Ok(entry)) => {
+                    latest_revision = entry.revision;
+                    if entry.key.contains(START_PROCESSING_MARKER) {
+                        continue;
+                    }
 
-                if !entry.value.is_empty() {
+                    if entry.key.contains(DONE_PROCESSING_MARKER) {
+                        break;
+                    }
                     tx.send(Arc::new(entry.value))
                         .await
                         .expect("Failed to send response");
                 }
+                Some(Err(e)) => {
+                    error!(request_id = ?id, "Error while receiving response from watcher: {e:?}, Recreating watcher...");
+                    watcher = Self::create_watcher(
+                        &store,
+                        &watch_pattern,
+                        latest_revision + 1,
+                        cln_token.clone(),
+                    )
+                    .await
+                    .expect("Failed to create watcher");
+                }
             }
-        });
+        }
+    }
+}
 
-        Ok((ReceiverStream::new(rx), handle))
+/// Represents the response mode for the request. For sse it will be streaming but for sync and async
+/// it will be unary.
+enum ResponseMode {
+    Stream {
+        tx: Sender<Arc<Bytes>>,
+        rx: Option<ReceiverStream<Arc<Bytes>>>,
+    },
+    Unary(Vec<Bytes>),
+}
+
+impl DataStore for JetStreamDataStore {
+    /// Retrieve a data from the store. If the datum is not found, `None` is returned.
+    async fn retrieve_data(
+        &mut self,
+        id: &str,
+        pod_hash: Option<&str>,
+    ) -> StoreResult<Vec<Vec<u8>>> {
+        // if the pod_hash is same as the current pod hash then we can rely on the central watcher for
+        // responses else we will have to create a new watcher for the old pod hash to get all the
+        // responses.
+        if let Some(pod_hash_value) = pod_hash {
+            Self::get_historic_response(
+                self.kv_store.clone(),
+                id,
+                pod_hash_value,
+                self.cln_token.clone(),
+            )
+            .await
+        } else {
+            self.get_data_from_response_store(id).await
+        }
+    }
+
+    /// Stream the response from the store. The response is streamed as it is added to the store.
+    async fn stream_data(
+        &mut self,
+        id: &str,
+        pod_hash: &str,
+    ) -> StoreResult<ReceiverStream<Arc<Bytes>>> {
+        // if the pod_hash is same as the current pod hash then we can rely on the central watcher for
+        // responses else we will have to create a new watcher for the old pod hash.
+        if pod_hash.eq(&self.pod_hash) {
+            // we need to wait for the start processing marker entry to be processed by the central watcher
+            loop {
+                {
+                    let mut response_map = self.responses_map.lock().await;
+                    if let Some(ResponseMode::Stream { rx, .. }) = response_map.get_mut(id) {
+                        if let Some(receiver_stream) = rx.take() {
+                            return Ok(receiver_stream);
+                        }
+                    }
+                }
+                sleep(Duration::from_millis(5)).await;
+            }
+        } else {
+            let (tx, rx) = mpsc::channel(10);
+            tokio::spawn({
+                let kv_store = self.kv_store.clone();
+                let request_id = id.to_string();
+                let pod_hash = pod_hash.to_string();
+                let cln_token = self.cln_token.clone();
+                async move {
+                    Self::stream_historic_responses(
+                        kv_store,
+                        &request_id,
+                        &pod_hash,
+                        tx,
+                        cln_token,
+                    )
+                    .await;
+                }
+            });
+            Ok(ReceiverStream::new(rx))
+        }
     }
 
     async fn ready(&mut self) -> bool {
@@ -139,15 +545,14 @@ impl DataStore for JetStreamDataStore {
 
 #[cfg(test)]
 mod tests {
+    use super::*;
+    use crate::app::store::datastore::DataStore;
     use async_nats::jetstream;
     use async_nats::jetstream::kv::Config;
     use bytes::Bytes;
     use chrono::Utc;
     use tokio_stream::StreamExt;
 
-    use super::*;
-    use crate::app::store::datastore::DataStore;
-
     #[cfg(feature = "nats-tests")]
     #[tokio::test]
     async fn test_retrieve_datum() {
@@ -155,6 +560,7 @@ mod tests {
         let client = async_nats::connect(js_url).await.unwrap();
         let context = jetstream::new(client);
         let datum_store_name = "test_retrieve_datum_store";
+        let pod_hash = "0";
 
         let _ = context.delete_key_value(datum_store_name).await;
 
@@ -167,22 +573,35 @@ mod tests {
             .await
             .unwrap();
 
-        let mut store = JetStreamDataStore::new(context.clone(), datum_store_name)
+        let mut store = JetStreamDataStore::new(
+            context.clone(),
+            datum_store_name,
+            pod_hash,
+            CancellationToken::new(),
+        )
+        .await
+        .unwrap();
+
+        let id = "test-id";
+        let start_key = format!("rs.{pod_hash}.{id}.start.processing");
+        let start_payload: Bytes = RequestType::Sync.try_into().unwrap();
+        store
+            .kv_store
+            .put(start_key, start_payload.clone())
             .await
             .unwrap();
 
-        let id = "test-id";
-        let key = format!("rs.{id}.0.{}", Utc::now().timestamp());
+        let key = format!("rs.{pod_hash}.{id}.sink.{}", Utc::now().timestamp());
         store
             .kv_store
             .put(key, Bytes::from_static(b"test_payload"))
             .await
             .unwrap();
 
-        let done_key = format!("rs.{id}.done.processing");
+        let done_key = format!("rs.{pod_hash}.{id}.done.processing");
         store.kv_store.put(done_key, Bytes::new()).await.unwrap();
 
-        let result = store.retrieve_data(id).await.unwrap();
+        let result = store.retrieve_data(id, None).await.unwrap();
         assert!(result.len() > 0);
         assert_eq!(result[0], b"test_payload");
 
@@ -197,6 +616,8 @@ mod tests {
         let client = async_nats::connect(js_url).await.unwrap();
         let context = jetstream::new(client);
         let datum_store_name = "test_stream_response_store";
+        let pod_hash = "0";
+        let _ = context.delete_key_value(datum_store_name).await;
 
         context
             .create_key_value(Config {
@@ -207,23 +628,151 @@ mod tests {
             .await
             .unwrap();
 
-        let mut store = JetStreamDataStore::new(context.clone(), datum_store_name)
+        let mut store = JetStreamDataStore::new(
+            context.clone(),
+            datum_store_name,
+            pod_hash,
+            CancellationToken::new(),
+        )
+        .await
+        .unwrap();
+
+        let id = "test_stream_id";
+        let start_key = format!("rs.{pod_hash}.{id}.start.processing");
+        let start_payload: Bytes = RequestType::Sse.try_into().unwrap();
+        store
+            .kv_store
+            .put(start_key, start_payload.clone())
             .await
             .unwrap();
 
-        let id = "test_stream_id";
         // Simulate a response being added to the store
-        let key = format!("rs.{id}.0.1234");
+        let key = format!("rs.{pod_hash}.{id}.0.1234");
         let payload = Bytes::from_static(b"test_payload");
         store.kv_store.put(key, payload.clone()).await.unwrap();
 
-        let (mut rx, handle) = store.stream_data(id).await.unwrap();
+        let mut stream = store.stream_data(id, pod_hash).await.unwrap();
 
         // Verify that the response is received
-        let received_response = rx.next().await.unwrap();
+        let received_response = stream.next().await.unwrap();
         assert_eq!(received_response, Arc::new(payload));
 
-        handle.abort();
+        // delete store
+        context.delete_key_value(datum_store_name).await.unwrap();
+    }
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_retrieve_data_with_different_pod_hash() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let datum_store_name = "test_retrieve_data_different_pod_hash_store";
+        let current_pod_hash = "abcd";
+        let other_pod_hash = "efgh";
+
+        let _ = context.delete_key_value(datum_store_name).await;
+
+        context
+            .create_key_value(Config {
+                bucket: datum_store_name.to_string(),
+                history: 5,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let mut store = JetStreamDataStore::new(
+            context.clone(),
+            datum_store_name,
+            current_pod_hash,
+            CancellationToken::new(),
+        )
+        .await
+        .unwrap();
+
+        let id = "test-unary-id";
+        let start_key = format!("rs.{other_pod_hash}.{id}.start.processing");
+        let start_payload: Bytes = RequestType::Sync.try_into().unwrap();
+        store
+            .kv_store
+            .put(start_key, start_payload.clone())
+            .await
+            .unwrap();
+
+        let key = format!("rs.{other_pod_hash}.{id}.sink.1234");
+        store
+            .kv_store
+            .put(key, Bytes::from_static(b"test_payload"))
+            .await
+            .unwrap();
+
+        let done_key = format!("rs.{other_pod_hash}.{id}.done.processing");
+        store.kv_store.put(done_key, Bytes::new()).await.unwrap();
+
+        let result = store.retrieve_data(id, Some(other_pod_hash)).await.unwrap();
+        assert!(result.len() > 0);
+        assert_eq!(result[0], b"test_payload");
+
+        // delete store
+        context.delete_key_value(datum_store_name).await.unwrap();
+    }
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_stream_data_with_different_pod_hash() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let datum_store_name = "test_stream_data_different_pod_hash_store";
+        let current_pod_hash = "hjck";
+        let other_pod_hash = "kclm";
+
+        let _ = context.delete_key_value(datum_store_name).await;
+
+        context
+            .create_key_value(Config {
+                bucket: datum_store_name.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let mut store = JetStreamDataStore::new(
+            context.clone(),
+            datum_store_name,
+            current_pod_hash,
+            CancellationToken::new(),
+        )
+        .await
+        .unwrap();
+
+        let id = "test-stream-id";
+        let start_key = format!("rs.{other_pod_hash}.{id}.start.processing");
+        let start_payload: Bytes = RequestType::Sse.try_into().unwrap();
+        store
+            .kv_store
+            .put(start_key, start_payload.clone())
+            .await
+            .unwrap();
+        let key = format!("rs.{other_pod_hash}.{id}.sink.1234");
+        store
+            .kv_store
+            .put(key, Bytes::from_static(b"test_payload"))
+            .await
+            .unwrap();
+
+        let done_key = format!("rs.{other_pod_hash}.{id}.done.processing");
+        store.kv_store.put(done_key, Bytes::new()).await.unwrap();
+
+        let mut stream = store.stream_data(id, other_pod_hash).await.unwrap();
+
+        let received_response = stream.next().await.unwrap();
+        assert_eq!(
+            received_response,
+            Arc::new(Bytes::from_static(b"test_payload"))
+        );
 
         // delete store
         context.delete_key_value(datum_store_name).await.unwrap();
diff --git a/rust/serving/src/app/store/datastore/user_defined.rs b/rust/serving/src/app/store/datastore/user_defined.rs
index 94291522..aad10698 100644
--- a/rust/serving/src/app/store/datastore/user_defined.rs
+++ b/rust/serving/src/app/store/datastore/user_defined.rs
@@ -8,7 +8,6 @@ use http::Uri;
 use numaflow_pb::clients::serving::serving_store_client::ServingStoreClient;
 use numaflow_pb::clients::serving::GetRequest;
 use tokio::net::UnixStream;
-use tokio::task::JoinHandle;
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::transport::{Channel, Endpoint};
 use tower::service_fn;
@@ -35,7 +34,11 @@ impl UserDefinedStore {
 
 impl DataStore for UserDefinedStore {
     // FIXME(serving): we need to return the origin details along with the payload
-    async fn retrieve_data(&mut self, id: &str) -> StoreResult<Vec<Vec<u8>>> {
+    async fn retrieve_data(
+        &mut self,
+        id: &str,
+        _pod_hash: Option<&str>,
+    ) -> StoreResult<Vec<Vec<u8>>> {
         let request = GetRequest { id: id.to_string() };
         let response = self
             .client
@@ -49,7 +52,8 @@ impl DataStore for UserDefinedStore {
     async fn stream_data(
         &mut self,
         _id: &str,
-    ) -> StoreResult<(ReceiverStream<Arc<Bytes>>, JoinHandle<()>)> {
+        _pod_hash: &str,
+    ) -> StoreResult<ReceiverStream<Arc<Bytes>>> {
         unimplemented!("stream_response is not supported for UserDefinedStore")
     }
 
@@ -145,7 +149,7 @@ mod tests {
             store: Arc::new(Mutex::new(HashMap::new())),
         };
         let id = "test_id".to_string();
-        let payload = vec![numaflow::serving_store::Payload {
+        let payload = vec![serving_store::Payload {
             origin: "test_origin".to_string(),
             value: vec![1, 2, 3],
         }];
@@ -177,7 +181,7 @@ mod tests {
         assert!(store.ready().await);
 
         // Test retrieve_data
-        let retrieved_data = store.retrieve_data(&id).await.unwrap();
+        let retrieved_data = store.retrieve_data(&id, None).await.unwrap();
         assert_eq!(retrieved_data, vec![payload[0].value.clone()]);
 
         drop(store);
diff --git a/rust/serving/src/app/store/status.rs b/rust/serving/src/app/store/status.rs
new file mode 100644
index 00000000..8213a756
--- /dev/null
+++ b/rust/serving/src/app/store/status.rs
@@ -0,0 +1,720 @@
+//! Status tracking for requests using jetstream key-value store.
+//!JetStream Status Entry uses the below format:
+//! - Status key - `status.{id}`
+//! - Status value - JSON serialized [ProcessingStatus]
+
+use crate::config::RequestType;
+use async_nats::jetstream::kv::{CreateErrorKind, Store};
+use async_nats::jetstream::Context;
+use bytes::Bytes;
+use serde::{Deserialize, Serialize};
+use std::result;
+use thiserror::Error;
+use tracing::{error, info, warn};
+
+const STATUS_KEY_PREFIX: &str = "status";
+const RESPONSE_KEY_PREFIX: &str = "rs";
+const START_PROCESSING_MARKER: &str = "start.processing";
+const DONE_PROCESSING_MARKER: &str = "done.processing";
+
+#[derive(Error, Debug, Clone)]
+pub(crate) enum Error {
+    #[error("Invalid status : {0}")]
+    InvalidStatus(String),
+    #[error("Client cancelled")]
+    Cancelled {
+        err: String,
+        previous_pod_hash: String,
+    },
+    #[error("Duplicate request : {0}")]
+    Duplicate(String),
+    #[error("Other : {0}")]
+    Other(String),
+}
+
+pub(crate) type Result<T> = result::Result<T, Error>;
+
+impl From<Error> for crate::Error {
+    fn from(value: Error) -> Self {
+        crate::Error::Status(value.to_string())
+    }
+}
+
+/// Represents the current processing status of a request id in the Jetstream KV Store. It also has the
+/// information about the pod hash which has accepted the request.
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
+pub(crate) enum ProcessingStatus {
+    /// InProgress indicates that the request is currently being processed.
+    InProgress { pod_hash: String },
+    /// Completed indicates that the request has been completed successfully.
+    Completed {
+        subgraph: String, // subgraph of the completed request
+        pod_hash: String,
+    },
+    /// Failed indicates that the request has failed.
+    Failed {
+        error: String, // error message of the failed request
+        pod_hash: String,
+    },
+}
+
+impl TryFrom<Bytes> for ProcessingStatus {
+    type Error = serde_json::Error;
+
+    fn try_from(value: Bytes) -> result::Result<Self, Self::Error> {
+        serde_json::from_slice(&value)
+    }
+}
+
+impl TryFrom<ProcessingStatus> for Bytes {
+    type Error = serde_json::Error;
+    fn try_from(status: ProcessingStatus) -> result::Result<Bytes, serde_json::Error> {
+        Ok(Bytes::from(serde_json::to_vec(&status)?))
+    }
+}
+
+/// StatusTracker is responsible for tracking the status of requests in the KV store.
+#[derive(Clone)]
+pub(crate) struct StatusTracker {
+    pod_hash: String,
+    status_kv: Store,
+    /// Optional response kv store to write start and done processing marker. For user defined store,
+    /// this is not used.
+    response_kv: Option<Store>,
+}
+
+impl StatusTracker {
+    pub(crate) async fn new(
+        js_context: Context,
+        status_bucket: &str,
+        pod_hash: String,
+        response_bucket: Option<String>,
+    ) -> Result<Self> {
+        let status_kv = js_context.get_key_value(status_bucket).await.map_err(|e| {
+            Error::Other(format!(
+                "Failed to get status kv store '{status_bucket}': {e:?}"
+            ))
+        })?;
+
+        // we need to optional response kv store to write start and done processing marker for the
+        // response.
+        let response_kv = match &response_bucket {
+            None => None,
+            Some(bucket_name) => {
+                let response_kv = js_context.get_key_value(bucket_name).await.map_err(|e| {
+                    Error::Other(format!(
+                        "Failed to get response kv store '{bucket_name}': {e:?}"
+                    ))
+                })?;
+                Some(response_kv)
+            }
+        };
+
+        Ok(Self {
+            status_kv,
+            pod_hash,
+            response_kv,
+        })
+    }
+
+    /// register a request id in the status kv store.
+    pub(crate) async fn register(&self, id: &str, request_type: RequestType) -> Result<()> {
+        let status_key = format!("{}.{}", STATUS_KEY_PREFIX, id);
+        let initial_status = ProcessingStatus::InProgress {
+            pod_hash: self.pod_hash.clone(),
+        };
+
+        match self
+            .status_kv
+            .create(
+                &status_key,
+                initial_status
+                    .clone()
+                    .try_into()
+                    .expect("status to bytes failed"),
+            )
+            .await
+        {
+            Ok(_) => {
+                // if the request is registered successfully, we need to write the start processing marker
+                // to signal the response watcher to wait for the responses. We need to do it here, since
+                // only the status tracker will know the lifecycle of the request.
+                // TODO: This makes `discard` method trickier.
+                if let Some(response_kv) = &self.response_kv {
+                    let response_key = format!(
+                        "{}.{}.{}.{}",
+                        RESPONSE_KEY_PREFIX, self.pod_hash, id, START_PROCESSING_MARKER
+                    );
+                    let request_type_bytes: Bytes = request_type
+                        .try_into()
+                        .expect("Failed to convert request type");
+                    response_kv
+                        .put(&response_key, request_type_bytes)
+                        .await
+                        .map_err(|e| {
+                            Error::Other(format!("Failed to create response entry for {id}: {e:?}"))
+                        })?;
+                }
+                Ok(())
+            }
+            Err(e) if e.kind() == CreateErrorKind::AlreadyExists => {
+                warn!(id, status_key, "Request ID already exists.");
+                let existing_status = self.get_status(&status_key).await?;
+                self.handle_existing_status(id, &status_key, existing_status)
+            }
+            Err(e) => {
+                error!(id, status_key, error = ?e, "Failed to create status entry.");
+                Err(Error::Other(format!(
+                    "Failed to register request id {status_key}: {e:?}"
+                )))
+            }
+        }
+    }
+
+    /// Helper to fetch and deserialize the status from the KV store.
+    async fn get_status(&self, status_key: &str) -> Result<ProcessingStatus> {
+        let status_entry =
+            self.status_kv.get(status_key).await.map_err(|e| {
+                Error::Other(format!("Failed to get status for {status_key}: {e:?}"))
+            })?;
+        let status_bytes = status_entry
+            .ok_or_else(|| Error::Other(format!("Status entry not found for {status_key}")))?;
+        ProcessingStatus::try_from(status_bytes).map_err(|e| {
+            Error::Other(format!(
+                "Failed to deserialize status for {status_key}: {e:?}"
+            ))
+        })
+    }
+
+    /// Handles scenarios where the status for a request ID already exists in the KV store.
+    ///
+    /// * [ProcessingStatus::Completed]: Indicates the request has already been processed. In this
+    ///   case, a duplicate error is returned.
+    /// * [ProcessingStatus::InProgress]: If the request is currently being processed and the pod
+    ///   hash matches the current pod, it is treated as a duplicate, and a duplicate error is
+    ///   returned. If the pod hash does not match, it means that pod crashed before sending the
+    ///   response. A 500 status code is returned and the request cannot be retried.
+    /// * [ProcessingStatus::Failed]: Indicates the request failed due to client cancellation or a
+    ///   failure in the serving pod (pod that accepted the request). In this case, the request can
+    ///   be retried.
+    fn handle_existing_status(
+        &self,
+        id: &str,
+        status_key: &str,
+        existing_status: ProcessingStatus,
+    ) -> Result<()> {
+        match existing_status {
+            ProcessingStatus::InProgress { pod_hash } => {
+                if pod_hash == self.pod_hash {
+                    warn!(id, status_key, "Request is getting processed");
+                    Err(Error::Duplicate(id.to_string()))
+                } else {
+                    // TODO: There can be a case where the pod terminates without updating the request
+                    //  status(SIGKILL). Return with a body asking for retrying with a different ID.
+                    warn!(
+                        id,
+                        status_key, "Request is getting processed in another pod"
+                    );
+                    Err(Error::Other(
+                        "Request is getting processed in another pod".to_string(),
+                    ))
+                }
+            }
+            ProcessingStatus::Completed { .. } => {
+                warn!(
+                    id,
+                    status_key, "Request already completed, returning as duplicate."
+                );
+                Err(Error::Duplicate(id.to_string()))
+            }
+            ProcessingStatus::Failed { error, pod_hash } => {
+                warn!(id, status_key, error = ?error, "Request had failed, retrying...");
+                Err(Error::Cancelled {
+                    err: error,
+                    previous_pod_hash: pod_hash,
+                })
+            }
+        }
+    }
+
+    /// returns the status of the request by checking the status kv store.
+    pub(crate) async fn status(&mut self, id: &str) -> Result<ProcessingStatus> {
+        let key = format!("{}.{}", STATUS_KEY_PREFIX, id);
+        let entry = self.status_kv.get(&key).await.map_err(|e| {
+            Error::Other(format!("Failed to get status for request id {id}: {e:?}"))
+        })?;
+        match entry {
+            Some(status) => ProcessingStatus::try_from(status)
+                .map_err(|e| Error::Other(format!("Failed to deserialize status for {id}: {e:?}"))),
+            None => Err(Error::InvalidStatus(id.to_string())),
+        }
+    }
+
+    /// Mark the request as failed in the status kv store.
+    pub(crate) async fn mark_as_failed(&mut self, id: &str, error: &str) -> Result<()> {
+        info!(id, replica_id = ?self.pod_hash, error, "Marking request as failed");
+        let failed_status = ProcessingStatus::Failed {
+            error: error.to_string(),
+            pod_hash: self.pod_hash.clone(),
+        };
+        self.update_status(id, failed_status).await
+    }
+
+    /// Update the status of a request in the status kv store.
+    async fn update_status(&self, id: &str, status: ProcessingStatus) -> Result<()> {
+        let key = format!("{}.{}", STATUS_KEY_PREFIX, id);
+
+        let status_bytes: Bytes = status.try_into().expect("Failed to convert into bytes");
+
+        self.status_kv
+            .put(key, status_bytes)
+            .await
+            .map_err(|e| Error::Other(format!("Failed to update status for {id}: {e:?}")))?;
+        Ok(())
+    }
+
+    /// Deregister marks the request as completed in the status kv store, it also (unlike registed)
+    /// writes a done processing marker to the response kv store (if nats is used) to signal the
+    /// response watcher.
+    pub(crate) async fn deregister(&self, id: &str, subgraph: &str) -> Result<()> {
+        let completed_status = ProcessingStatus::Completed {
+            subgraph: subgraph.to_string(),
+            pod_hash: self.pod_hash.clone(),
+        };
+
+        self.update_status(id, completed_status).await?;
+
+        let Some(response_kv) = &self.response_kv else {
+            return Ok(());
+        };
+
+        // if nats is used, we need to write the done processing marker to the response kv store to
+        // let the response watcher know that the processing is done. Since only tracker knows the
+        // lifecycle of the request, we need to do it here.
+        let done_key = format!(
+            "{}.{}.{}.{}",
+            RESPONSE_KEY_PREFIX, self.pod_hash, id, DONE_PROCESSING_MARKER
+        );
+        response_kv
+            .put(done_key.clone(), Bytes::new())
+            .await
+            .map_err(|e| {
+                Error::Other(format!(
+                    "Failed to write done processing marker {done_key}: {e:?}"
+                ))
+            })?;
+        Ok(())
+    }
+
+    /// Discard the request from the status kv store. This is used when the request is not accepted
+    /// because of back pressure.
+    /// It deletes the entry from the status kv store (so ID is removed and can be retried) and also
+    /// sends a done processing marker to the response kv store (if nats is used) to signal the end.
+    /// This is because we just start processing marker in the register before inserting into the pipeline.
+    pub(crate) async fn discard(&self, id: &str) -> Result<()> {
+        let key = format!("{}.{}", STATUS_KEY_PREFIX, id);
+        self.status_kv
+            .delete(&key)
+            .await
+            .map_err(|e| Error::Other(format!("Failed to delete status for {id}: {e:?}")))?;
+
+        let Some(response_kv) = &self.response_kv else {
+            return Ok(());
+        };
+        let done_key = format!(
+            "{}.{}.{}.{}",
+            RESPONSE_KEY_PREFIX, self.pod_hash, id, DONE_PROCESSING_MARKER
+        );
+
+        response_kv
+            .put(done_key.clone(), Bytes::new())
+            .await
+            .map_err(|e| {
+                Error::Other(format!(
+                    "Failed to write done processing marker {done_key}: {e:?}"
+                ))
+            })?;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use async_nats::jetstream;
+    use async_nats::jetstream::kv::Config;
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_register_request() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let status_bucket = "test_register_status";
+        let response_bucket = "test_register_response";
+
+        // Clean up buckets
+        let _ = context.delete_key_value(status_bucket).await;
+        let _ = context.delete_key_value(response_bucket).await;
+
+        // Create buckets
+        context
+            .create_key_value(Config {
+                bucket: status_bucket.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        context
+            .create_key_value(Config {
+                bucket: response_bucket.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let mut tracker = StatusTracker::new(
+            context.clone(),
+            status_bucket,
+            "test_pod".to_string(),
+            Some(response_bucket.to_string()),
+        )
+        .await
+        .unwrap();
+
+        let request_id = "test_request";
+        tracker
+            .register(request_id, RequestType::Sync)
+            .await
+            .expect("Failed to register request");
+
+        // Verify status
+        let status = tracker.status(request_id).await.unwrap();
+        assert_eq!(
+            status,
+            ProcessingStatus::InProgress {
+                pod_hash: "test_pod".to_string()
+            }
+        );
+
+        // Verify start processing marker
+        let response_kv = context.get_key_value(response_bucket).await.unwrap();
+        let start_key = format!(
+            "{}.{}.{}.{}",
+            RESPONSE_KEY_PREFIX, "test_pod", request_id, START_PROCESSING_MARKER
+        );
+        let start_marker = response_kv.get(&start_key).await.unwrap();
+        assert!(start_marker.is_some());
+
+        // Clean up
+        context.delete_key_value(status_bucket).await.unwrap();
+        context.delete_key_value(response_bucket).await.unwrap();
+    }
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_handle_existing_status() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let status_bucket = "test_existing_status";
+
+        // Clean up bucket
+        let _ = context.delete_key_value(status_bucket).await;
+
+        // Create bucket
+        context
+            .create_key_value(Config {
+                bucket: status_bucket.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let tracker =
+            StatusTracker::new(context.clone(), status_bucket, "test_pod".to_string(), None)
+                .await
+                .unwrap();
+
+        let request_id = "test_request";
+        tracker
+            .register(request_id, RequestType::Sync)
+            .await
+            .expect("Failed to register request");
+
+        // Attempt to register the same request again
+        let result = tracker.register(request_id, RequestType::Sync).await;
+        assert!(matches!(result, Err(Error::Duplicate(_))));
+
+        // Clean up
+        context.delete_key_value(status_bucket).await.unwrap();
+    }
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_mark_as_failed() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let status_bucket = "test_mark_failed";
+
+        // Clean up bucket
+        let _ = context.delete_key_value(status_bucket).await;
+
+        // Create bucket
+        context
+            .create_key_value(Config {
+                bucket: status_bucket.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let mut tracker =
+            StatusTracker::new(context.clone(), status_bucket, "test_pod".to_string(), None)
+                .await
+                .unwrap();
+
+        let request_id = "test_request";
+        tracker
+            .register(request_id, RequestType::Sync)
+            .await
+            .expect("Failed to register request");
+
+        tracker
+            .mark_as_failed(request_id, "Test error")
+            .await
+            .expect("Failed to mark request as failed");
+
+        // Verify status
+        let status = tracker.status(request_id).await.unwrap();
+        assert_eq!(
+            status,
+            ProcessingStatus::Failed {
+                error: "Test error".to_string(),
+                pod_hash: "test_pod".to_string()
+            }
+        );
+
+        // Clean up
+        context.delete_key_value(status_bucket).await.unwrap();
+    }
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_deregister_request() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let status_bucket = "test_deregister_status";
+        let response_bucket = "test_deregister_response";
+
+        // Clean up buckets
+        let _ = context.delete_key_value(status_bucket).await;
+        let _ = context.delete_key_value(response_bucket).await;
+
+        // Create buckets
+        context
+            .create_key_value(Config {
+                bucket: status_bucket.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        context
+            .create_key_value(Config {
+                bucket: response_bucket.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let mut tracker = StatusTracker::new(
+            context.clone(),
+            status_bucket,
+            "test_pod".to_string(),
+            Some(response_bucket.to_string()),
+        )
+        .await
+        .unwrap();
+
+        let request_id = "test_request";
+        tracker
+            .register(request_id, RequestType::Sync)
+            .await
+            .expect("Failed to register request");
+
+        tracker
+            .deregister(request_id, "test_subgraph")
+            .await
+            .expect("Failed to deregister request");
+
+        // Verify status
+        let status = tracker.status(request_id).await;
+        assert!(matches!(status, Ok(ProcessingStatus::Completed { .. })));
+
+        // Verify done processing marker
+        let response_kv = context.get_key_value(response_bucket).await.unwrap();
+        let done_key = format!(
+            "{}.{}.{}.{}",
+            RESPONSE_KEY_PREFIX, "test_pod", request_id, DONE_PROCESSING_MARKER
+        );
+        let done_marker = response_kv.get(&done_key).await.unwrap();
+        assert!(done_marker.is_some());
+
+        // Clean up
+        context.delete_key_value(status_bucket).await.unwrap();
+        context.delete_key_value(response_bucket).await.unwrap();
+    }
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_discard_request() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let status_bucket = "test_discard_status";
+        let response_bucket = "test_discard_response";
+
+        // Clean up buckets
+        let _ = context.delete_key_value(status_bucket).await;
+        let _ = context.delete_key_value(response_bucket).await;
+
+        // Create buckets
+        context
+            .create_key_value(Config {
+                bucket: status_bucket.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        context
+            .create_key_value(Config {
+                bucket: response_bucket.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let mut tracker = StatusTracker::new(
+            context.clone(),
+            status_bucket,
+            "test_pod".to_string(),
+            Some(response_bucket.to_string()),
+        )
+        .await
+        .unwrap();
+
+        let request_id = "test_request";
+        tracker
+            .register(request_id, RequestType::Sync)
+            .await
+            .expect("Failed to register request");
+
+        tracker
+            .discard(request_id)
+            .await
+            .expect("Failed to discard request");
+
+        // Verify status is removed
+        let status = tracker.status(request_id).await;
+        assert!(status.is_err());
+
+        // Verify done processing marker is sent so that response watcher can clean up
+        let response_kv = context.get_key_value(response_bucket).await.unwrap();
+        let done_key = format!(
+            "{}.{}.{}.{}",
+            RESPONSE_KEY_PREFIX, "test_pod", request_id, DONE_PROCESSING_MARKER
+        );
+        let done_marker = response_kv.get(&done_key).await.unwrap();
+        assert!(done_marker.is_some());
+
+        // Clean up
+        context.delete_key_value(status_bucket).await.unwrap();
+        context.delete_key_value(response_bucket).await.unwrap();
+    }
+
+    #[cfg(feature = "nats-tests")]
+    #[tokio::test]
+    async fn test_handle_existing_status_cases() {
+        let js_url = "localhost:4222";
+        let client = async_nats::connect(js_url).await.unwrap();
+        let context = jetstream::new(client);
+        let status_bucket = "test_handle_existing_status_cases";
+
+        // Clean up bucket
+        let _ = context.delete_key_value(status_bucket).await;
+
+        // Create bucket
+        context
+            .create_key_value(Config {
+                bucket: status_bucket.to_string(),
+                history: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        let tracker =
+            StatusTracker::new(context.clone(), status_bucket, "test_pod".to_string(), None)
+                .await
+                .unwrap();
+
+        let request_id = "test_request";
+
+        // Case 1: InProgress (different pod)
+        tracker
+            .update_status(
+                request_id,
+                ProcessingStatus::InProgress {
+                    pod_hash: "other_pod".to_string(),
+                },
+            )
+            .await
+            .unwrap();
+        let result = tracker.register(request_id, RequestType::Sync).await;
+        assert!(matches!(result, Err(Error::Other(_))));
+
+        // Case 2: Completed
+        tracker
+            .update_status(
+                request_id,
+                ProcessingStatus::Completed {
+                    subgraph: "test_subgraph".to_string(),
+                    pod_hash: "test_pod".to_string(),
+                },
+            )
+            .await
+            .unwrap();
+        let result = tracker.register(request_id, RequestType::Sync).await;
+        assert!(matches!(result, Err(Error::Duplicate(_))));
+
+        // Case 3: Failed
+        tracker
+            .update_status(
+                request_id,
+                ProcessingStatus::Failed {
+                    error: "Test error".to_string(),
+                    pod_hash: "test_pod".to_string(),
+                },
+            )
+            .await
+            .unwrap();
+        let result = tracker.register(request_id, RequestType::Sync).await;
+        assert!(matches!(result, Err(Error::Cancelled { .. })));
+
+        // Clean up
+        context.delete_key_value(status_bucket).await.unwrap();
+    }
+}
diff --git a/rust/serving/src/app/tracker.rs b/rust/serving/src/app/tracker.rs
index 8793f881..9be3cbae 100644
--- a/rust/serving/src/app/tracker.rs
+++ b/rust/serving/src/app/tracker.rs
@@ -8,6 +8,8 @@ use crate::callback::Callback;
 use crate::pipeline::{Edge, OperatorType, PipelineDCG};
 use crate::Error;
 
+const DROP: &str = "U+005C__DROP__";
+
 fn compare_slice(operator: &OperatorType, a: &[String], b: &[String]) -> bool {
     match operator {
         OperatorType::And => a.iter().all(|val| b.contains(val)),
@@ -188,7 +190,19 @@ impl MessageGraph {
             cb_time: current_callback.cb_time,
         });
 
-        // if there are no responses, means the message is dropped, we can return true
+        // if there are no responses or if any of the response contains drop tag means the message
+        // was dropped we can return true.
+        if current_callback.responses.is_empty()
+            || current_callback.responses.iter().any(|response| {
+                response
+                    .tags
+                    .as_ref()
+                    .is_some_and(|tags| tags.contains(&DROP.to_string()))
+            })
+        {
+            return true;
+        }
+
         if current_callback.responses.is_empty() {
             return true;
         }
diff --git a/rust/serving/src/callback.rs b/rust/serving/src/callback.rs
index ef3e3ded..2d92446d 100644
--- a/rust/serving/src/callback.rs
+++ b/rust/serving/src/callback.rs
@@ -1,7 +1,4 @@
-use std::{
-    sync::Arc,
-    time::{SystemTime, UNIX_EPOCH},
-};
+use std::sync::Arc;
 
 use async_nats::jetstream::kv::Store;
 use async_nats::jetstream::Context;
@@ -85,13 +82,11 @@ impl CallbackHandler {
     pub async fn callback(
         &self,
         id: String,
+        pod_hash: String,
         previous_vertex: String,
         responses: Vec<Option<Vec<String>>>,
     ) -> crate::Result<JoinHandle<()>> {
-        let cb_time = SystemTime::now()
-            .duration_since(UNIX_EPOCH)
-            .expect("System time is older than Unix epoch time")
-            .as_millis() as u64;
+        let cb_time = Utc::now().timestamp_millis() as u64;
 
         let responses = responses
             .into_iter()
@@ -111,7 +106,7 @@ impl CallbackHandler {
         let store = self.store.clone();
         let handle = tokio::spawn(async move {
             let timestamp = Utc::now().timestamp_nanos_opt().unwrap();
-            let callbacks_key = format!("cb.{id}.{vertex_name}.{timestamp}");
+            let callbacks_key = format!("cb.{pod_hash}.{id}.{vertex_name}.{timestamp}");
             let interval = fixed::Interval::from_millis(1000).take(2);
 
             let _permit = permit;
@@ -184,7 +179,12 @@ mod tests {
             CallbackHandler::new("test_vertex".to_string(), context.clone(), store_name, 1).await;
 
         callback_handler
-            .callback(id.to_string(), "test_from_vertex".to_string(), vec![])
+            .callback(
+                id.to_string(),
+                "0".to_string(),
+                "test_from_vertex".to_string(),
+                vec![],
+            )
             .await
             .unwrap()
             .await
diff --git a/rust/serving/src/config.rs b/rust/serving/src/config.rs
index 7f410e68..921a1a26 100644
--- a/rust/serving/src/config.rs
+++ b/rust/serving/src/config.rs
@@ -1,25 +1,30 @@
-use std::collections::HashMap;
-use std::fmt::Debug;
-
-use base64::prelude::BASE64_STANDARD;
-use base64::Engine;
-use rcgen::{generate_simple_self_signed, Certificate, CertifiedKey, KeyPair};
-use serde::{Deserialize, Serialize};
-
 use crate::{
     pipeline::PipelineDCG,
     Error::{self, ParseConfig},
 };
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use bytes::Bytes;
+use rcgen::{generate_simple_self_signed, Certificate, CertifiedKey, KeyPair};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::fmt::Debug;
+use tracing::info;
 
 const ENV_NUMAFLOW_SERVING_APP_PORT: &str = "NUMAFLOW_SERVING_APP_LISTEN_PORT";
 pub const ENV_MIN_PIPELINE_SPEC: &str = "NUMAFLOW_SERVING_MIN_PIPELINE_SPEC";
 pub const DEFAULT_ID_HEADER: &str = "X-Numaflow-Id";
+pub const DEFAULT_POD_HASH_KEY: &str = "X-Numaflow-Pod-Hash";
 pub const DEFAULT_CALLBACK_URL_HEADER_KEY: &str = "X-Numaflow-Callback-Url";
 const DEFAULT_GRPC_MAX_MESSAGE_SIZE: usize = 64 * 1024 * 1024; // 64 MB
 const DEFAULT_SERVING_STORE_SOCKET: &str = "/var/run/numaflow/serving.sock";
 const DEFAULT_SERVING_STORE_SERVER_INFO_FILE: &str = "/var/run/numaflow/serving-server-info";
 const ENV_NUMAFLOW_SERVING_JETSTREAM_USER: &str = "NUMAFLOW_ISBSVC_JETSTREAM_USER";
 const ENV_NUMAFLOW_SERVING_JETSTREAM_PASSWORD: &str = "NUMAFLOW_ISBSVC_JETSTREAM_PASSWORD";
+const ENV_NUMAFLOW_SERVING_CALLBACK_STORE: &str = "NUMAFLOW_SERVING_CALLBACK_STORE";
+const ENV_NUMAFLOW_SERVING_RESPONSE_STORE: &str = "NUMAFLOW_SERVING_RESPONSE_STORE";
+const ENV_NUMAFLOW_SERVING_STATUS_STORE: &str = "NUMAFLOW_SERVING_STATUS_STORE";
+const ENV_NUMAFLOW_POD: &str = "NUMAFLOW_POD";
 
 pub fn generate_certs() -> Result<(Certificate, KeyPair), String> {
     let CertifiedKey { cert, key_pair } = generate_simple_self_signed(vec!["localhost".into()])
@@ -27,6 +32,29 @@ pub fn generate_certs() -> Result<(Certificate, KeyPair), String> {
     Ok((cert, key_pair))
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
+pub(crate) enum RequestType {
+    Sse,
+    Sync,
+    Async,
+}
+
+impl TryInto<Bytes> for RequestType {
+    type Error = serde_json::Error;
+
+    fn try_into(self) -> Result<Bytes, Self::Error> {
+        Ok(Bytes::from(serde_json::to_vec(&self)?))
+    }
+}
+
+impl TryFrom<Bytes> for RequestType {
+    type Error = serde_json::Error;
+
+    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
+        serde_json::from_slice(&value)
+    }
+}
+
 #[derive(Debug, Deserialize, Clone, PartialEq)]
 pub struct UserDefinedStoreConfig {
     pub grpc_max_message_size: usize,
@@ -49,12 +77,15 @@ pub struct Settings {
     /// The HTTP header used to communicate to the client about the unique id assigned for a request in the store
     /// The client may also set the value of this header when sending the payload.
     pub tid_header: String,
+    pub pod_hash: String,
     pub app_listen_port: u16,
     pub metrics_server_listen_port: u16,
     pub upstream_addr: String,
     pub drain_timeout_secs: u64,
     pub store_type: StoreType,
     pub js_callback_store: String,
+    pub js_response_store: String,
+    pub js_status_store: String,
     pub api_auth_token: Option<String>,
     pub pipeline_spec: PipelineDCG,
     pub nats_basic_auth: Option<(String, String)>,
@@ -73,12 +104,15 @@ impl Default for Settings {
     fn default() -> Self {
         Self {
             tid_header: DEFAULT_ID_HEADER.to_owned(),
+            pod_hash: "0".to_owned(),
             app_listen_port: 3000,
             metrics_server_listen_port: 3001,
             upstream_addr: "localhost:8888".to_owned(),
             drain_timeout_secs: 600,
             store_type: StoreType::default(),
-            js_callback_store: "kv".to_owned(),
+            js_callback_store: "callback-kv".to_owned(),
+            js_response_store: "response-kv".to_owned(),
+            js_status_store: "status-kv".to_owned(),
             api_auth_token: None,
             pipeline_spec: Default::default(),
             nats_basic_auth: None,
@@ -184,16 +218,32 @@ impl TryFrom<HashMap<String, String>> for Settings {
             None => None,
         };
 
-        let js_store =  env_vars.get("NUMAFLOW_SERVING_KV_STORE").ok_or_else(|| {
-            ParseConfig("Serving store is default, but environment variable NUMAFLOW_SERVING_KV_STORE is not set".into())
-        })?;
+        let js_cb_store = env_vars
+            .get(ENV_NUMAFLOW_SERVING_CALLBACK_STORE)
+            .ok_or_else(|| {
+                ParseConfig("Serving store is default, but environment variable is not set".into())
+            })?;
+
+        let js_response_store = env_vars
+            .get(ENV_NUMAFLOW_SERVING_RESPONSE_STORE)
+            .ok_or_else(|| {
+                ParseConfig("Serving store is default, but environment variable is not set".into())
+            })?;
+
+        let js_status_store = env_vars
+            .get(ENV_NUMAFLOW_SERVING_STATUS_STORE)
+            .ok_or_else(|| {
+                ParseConfig("Serving store is default, but environment variable is not set".into())
+            })?;
 
         let mut settings = Settings {
             pipeline_spec,
-            js_callback_store: js_store.into(),
+            js_callback_store: js_cb_store.into(),
             nats_basic_auth,
             js_message_stream: js_source_spec.stream,
             jetstream_url: js_source_spec.url,
+            js_response_store: js_response_store.into(),
+            js_status_store: js_status_store.into(),
             ..Default::default()
         };
 
@@ -230,11 +280,18 @@ impl TryFrom<HashMap<String, String>> for Settings {
         };
 
         settings.tid_header = serving_server_settings.msg_id_header_key;
+        settings.pod_hash = env_vars
+            .get(ENV_NUMAFLOW_POD)
+            .unwrap()
+            .split('-')
+            .last()
+            .unwrap_or("0")
+            .to_string();
 
         settings.drain_timeout_secs = serving_server_settings
             .request_timeout_seconds
             .unwrap_or(120)
-            .max(1) as u64; // Ensure timeout is atleast 1 second
+            .max(1) as u64; // Ensure timeout is at least 1 second
 
         if let Some(auth) = serving_server_settings.auth {
             let token = auth.token.unwrap();
@@ -243,6 +300,7 @@ impl TryFrom<HashMap<String, String>> for Settings {
             settings.api_auth_token = Some(auth_token);
         }
 
+        info!("Settings {:?}", settings);
         Ok(settings)
     }
 }
@@ -266,6 +324,7 @@ mod tests {
         let settings = Settings::default();
 
         assert_eq!(settings.tid_header, "X-Numaflow-Id");
+        assert_eq!(settings.pod_hash, "0");
         assert_eq!(settings.app_listen_port, 3000);
         assert_eq!(settings.metrics_server_listen_port, 3001);
         assert_eq!(settings.upstream_addr, "localhost:8888");
@@ -278,9 +337,12 @@ mod tests {
         // Set up the environment variables
         let env_vars = [
             (ENV_NUMAFLOW_SERVING_APP_PORT, "8443"),
+            (ENV_NUMAFLOW_POD, "serving-server-kddc"),
             ("NUMAFLOW_ISBSVC_JETSTREAM_USER", "testuser"),
             ("NUMAFLOW_ISBSVC_JETSTREAM_PASSWORD", "testpasswd"),
-            ("NUMAFLOW_SERVING_KV_STORE", "test-kv-store"),
+            ("NUMAFLOW_SERVING_CALLBACK_STORE", "test-kv-store"),
+            ("NUMAFLOW_SERVING_STATUS_STORE", "test-kv-store"),
+            ("NUMAFLOW_SERVING_RESPONSE_STORE", "test-kv-store"),
             ("NUMAFLOW_SERVING_SOURCE_SETTINGS", "eyJhdXRoIjpudWxsLCJzZXJ2aWNlIjp0cnVlLCJtc2dJREhlYWRlcktleSI6IlgtTnVtYWZsb3ctSWQifQ=="),
             (ENV_MIN_PIPELINE_SPEC, "eyJ2ZXJ0aWNlcyI6W3sibmFtZSI6ImluIiwic291cmNlIjp7ImpldHN0cmVhbSI6eyJ1cmwiOiJuYXRzOi8vaXNic3ZjLWRlZmF1bHQtanMtc3ZjLmRlZmF1bHQuc3ZjOjQyMjIiLCJzdHJlYW0iOiJzZXJ2aW5nLXNvdXJjZS1zaW1wbGUtcGlwZWxpbmUiLCJ0bHMiOm51bGwsImF1dGgiOnsiYmFzaWMiOnsidXNlciI6eyJuYW1lIjoiaXNic3ZjLWRlZmF1bHQtanMtY2xpZW50LWF1dGgiLCJrZXkiOiJjbGllbnQtYXV0aC11c2VyIn0sInBhc3N3b3JkIjp7Im5hbWUiOiJpc2JzdmMtZGVmYXVsdC1qcy1jbGllbnQtYXV0aCIsImtleSI6ImNsaWVudC1hdXRoLXBhc3N3b3JkIn19fX19LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJlbnYiOlt7Im5hbWUiOiJOVU1BRkxPV19DQUxMQkFDS19FTkFCTEVEIiwidmFsdWUiOiJ0cnVlIn0seyJuYW1lIjoiTlVNQUZMT1dfU0VSVklOR19TT1VSQ0VfU0VUVElOR1MiLCJ2YWx1ZSI6ImV5SmhkWFJvSWpwdWRXeHNMQ0p6WlhKMmFXTmxJanAwY25WbExDSnRjMmRKUkVobFlXUmxja3RsZVNJNklsZ3RUblZ0WVdac2IzY3RTV1FpZlE9PSJ9LHsibmFtZSI6Ik5VTUFGTE9XX1NFUlZJTkdfS1ZfU1RPUkUiLCJ2YWx1ZSI6InNlcnZpbmctc3RvcmUtc2ltcGxlLXBpcGVsaW5lX1NFUlZJTkdfS1ZfU1RPUkUifV19LCJzY2FsZSI6eyJtaW4iOjEsIm1heCI6MX0sImluaXRDb250YWluZXJzIjpbeyJuYW1lIjoidmFsaWRhdGUtc3RyZWFtLWluaXQiLCJpbWFnZSI6InF1YXkuaW8vbnVtYXByb2ovbnVtYWZsb3c6ODA4MkU2NTYtQTAxOS00QjE1LUE2ODQtNTg2RkM3RDJDQTFGIiwiYXJncyI6WyJpc2JzdmMtdmFsaWRhdGUiLCItLWlzYnN2Yy10eXBlPWpldHN0cmVhbSIsIi0tYnVmZmVycz1zZXJ2aW5nLXNvdXJjZS1zaW1wbGUtcGlwZWxpbmUiXSwiZW52IjpbeyJuYW1lIjoiTlVNQUZMT1dfUElQRUxJTkVfTkFNRSIsInZhbHVlIjoicy1zaW1wbGUtcGlwZWxpbmUifSx7Im5hbWUiOiJHT0RFQlVHIn0seyJuYW1lIjoiTlVNQUZMT1dfSVNCU1ZDX0NPTkZJRyIsInZhbHVlIjoiZXlKcVpYUnpkSEpsWVcwaU9uc2lkWEpzSWpvaWJtRjBjem92TDJselluTjJZeTFrWldaaGRXeDBMV3B6TFhOMll5NWtaV1poZFd4MExuTjJZem8wTWpJeUlpd2lZWFYwYUNJNmV5SmlZWE5wWXlJNmV5SjFjMlZ5SWpwN0ltNWhiV1VpT2lKcGMySnpkbU10WkdWbVlYVnNkQzFxY3kxamJHbGxiblF0WVhWMGFDSXNJbXRsZVNJNkltTnNhV1Z1ZEMxaGRYUm9MWFZ6WlhJaWZTd2ljR0Z6YzNkdmNtUWlPbnNpYm1GdFpTSTZJbWx6WW5OMll5MWtaV1poZFd4MExXcHpMV05zYVdWdWRDMWhkWFJvSWl3aWEyVjVJam9pWTJ4cFpXNTBMV0YxZEdndGNHRnpjM2R2Y21RaWZYMTlMQ0p6ZEhKbFlXMURiMjVtYVdjaU9pSmpiMjV6ZFcxbGNqcGNiaUFnWVdOcmQyRnBkRG9nTmpCelhHNGdJRzFoZUdGamEzQmxibVJwYm1jNklESTFNREF3WEc1dmRHSjFZMnRsZERwY2JpQWdhR2x6ZEc5eWVUb2dNVnh1SUNCdFlYaGllWFJsY3pvZ01GeHVJQ0J0WVhoMllXeDFaWE5wZW1VNklEQmNiaUFnY21Wd2JHbGpZWE02SUROY2JpQWdjM1J2Y21GblpUb2dNRnh1SUNCMGRHdzZJRE5vWEc1d2NtOWpZblZqYTJWME9seHVJQ0JvYVhOMGIzSjVPaUF4WEc0Z0lHMWhlR0o1ZEdWek9pQXdYRzRnSUcxaGVIWmhiSFZsYzJsNlpUb2dNRnh1SUNCeVpYQnNhV05oY3pvZ00xeHVJQ0J6ZEc5eVlXZGxPaUF3WEc0Z0lIUjBiRG9nTnpKb1hHNXpkSEpsWVcwNlhHNGdJR1IxY0d4cFkyRjBaWE02SURZd2MxeHVJQ0J0WVhoaFoyVTZJRGN5YUZ4dUlDQnRZWGhpZVhSbGN6b2dMVEZjYmlBZ2JXRjRiWE5uY3pvZ01UQXdNREF3WEc0Z0lISmxjR3hwWTJGek9pQXpYRzRnSUhKbGRHVnVkR2x2YmpvZ01GeHVJQ0J6ZEc5eVlXZGxPaUF3WEc0aWZYMD0ifSx7Im5hbWUiOiJOVU1BRkxPV19JU0JTVkNfSkVUU1RSRUFNX1VSTCIsInZhbHVlIjoibmF0czovL2lzYnN2Yy1kZWZhdWx0LWpzLXN2Yy5kZWZhdWx0LnN2Yzo0MjIyIn0seyJuYW1lIjoiTlVNQUZMT1dfSVNCU1ZDX0pFVFNUUkVBTV9UTFNfRU5BQkxFRCIsInZhbHVlIjoiZmFsc2UifSx7Im5hbWUiOiJOVU1BRkxPV19JU0JTVkNfSkVUU1RSRUFNX1VTRVIiLCJ2YWx1ZUZyb20iOnsic2VjcmV0S2V5UmVmIjp7Im5hbWUiOiJpc2JzdmMtZGVmYXVsdC1qcy1jbGllbnQtYXV0aCIsImtleSI6ImNsaWVudC1hdXRoLXVzZXIifX19LHsibmFtZSI6Ik5VTUFGTE9XX0lTQlNWQ19KRVRTVFJFQU1fUEFTU1dPUkQiLCJ2YWx1ZUZyb20iOnsic2VjcmV0S2V5UmVmIjp7Im5hbWUiOiJpc2JzdmMtZGVmYXVsdC1qcy1jbGllbnQtYXV0aCIsImtleSI6ImNsaWVudC1hdXRoLXBhc3N3b3JkIn19fV0sInJlc291cmNlcyI6eyJyZXF1ZXN0cyI6eyJjcHUiOiIxMDBtIiwibWVtb3J5IjoiMTI4TWkifX0sImltYWdlUHVsbFBvbGljeSI6IklmTm90UHJlc2VudCJ9XSwidXBkYXRlU3RyYXRlZ3kiOnsidHlwZSI6IlJvbGxpbmdVcGRhdGUiLCJyb2xsaW5nVXBkYXRlIjp7Im1heFVuYXZhaWxhYmxlIjoiMjUlIn19fSx7Im5hbWUiOiJjYXQiLCJ1ZGYiOnsiY29udGFpbmVyIjp7ImltYWdlIjoicXVheS5pby9udW1haW8vbnVtYWZsb3ctZ28vbWFwLWZvcndhcmQtbWVzc2FnZTpzdGFibGUiLCJyZXNvdXJjZXMiOnt9fSwiYnVpbHRpbiI6bnVsbCwiZ3JvdXBCeSI6bnVsbH0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImVudiI6W3sibmFtZSI6Ik5VTUFGTE9XX0NBTExCQUNLX0VOQUJMRUQiLCJ2YWx1ZSI6InRydWUifSx7Im5hbWUiOiJOVU1BRkxPV19TRVJWSU5HX1NPVVJDRV9TRVRUSU5HUyIsInZhbHVlIjoiZXlKaGRYUm9JanB1ZFd4c0xDSnpaWEoyYVdObElqcDBjblZsTENKdGMyZEpSRWhsWVdSbGNrdGxlU0k2SWxndFRuVnRZV1pzYjNjdFNXUWlmUT09In0seyJuYW1lIjoiTlVNQUZMT1dfU0VSVklOR19LVl9TVE9SRSIsInZhbHVlIjoic2VydmluZy1zdG9yZS1zaW1wbGUtcGlwZWxpbmVfU0VSVklOR19LVl9TVE9SRSJ9XX0sInNjYWxlIjp7Im1pbiI6MSwibWF4IjoxfSwidXBkYXRlU3RyYXRlZ3kiOnsidHlwZSI6IlJvbGxpbmdVcGRhdGUiLCJyb2xsaW5nVXBkYXRlIjp7Im1heFVuYXZhaWxhYmxlIjoiMjUlIn19fSx7Im5hbWUiOiJvdXQiLCJzaW5rIjp7InVkc2luayI6eyJjb250YWluZXIiOnsiaW1hZ2UiOiJxdWF5LmlvL251bWFpby9udW1hZmxvdy1nby9zaW5rLXNlcnZlOnN0YWJsZSIsInJlc291cmNlcyI6e319fSwicmV0cnlTdHJhdGVneSI6e319LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJlbnYiOlt7Im5hbWUiOiJOVU1BRkxPV19DQUxMQkFDS19FTkFCTEVEIiwidmFsdWUiOiJ0cnVlIn0seyJuYW1lIjoiTlVNQUZMT1dfU0VSVklOR19TT1VSQ0VfU0VUVElOR1MiLCJ2YWx1ZSI6ImV5SmhkWFJvSWpwdWRXeHNMQ0p6WlhKMmFXTmxJanAwY25WbExDSnRjMmRKUkVobFlXUmxja3RsZVNJNklsZ3RUblZ0WVdac2IzY3RTV1FpZlE9PSJ9LHsibmFtZSI6Ik5VTUFGTE9XX1NFUlZJTkdfS1ZfU1RPUkUiLCJ2YWx1ZSI6InNlcnZpbmctc3RvcmUtc2ltcGxlLXBpcGVsaW5lX1NFUlZJTkdfS1ZfU1RPUkUifV19LCJzY2FsZSI6eyJtaW4iOjEsIm1heCI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX1dLCJlZGdlcyI6W3siZnJvbSI6ImluIiwidG8iOiJjYXQiLCJjb25kaXRpb25zIjpudWxsfSx7ImZyb20iOiJjYXQiLCJ0byI6Im91dCIsImNvbmRpdGlvbnMiOm51bGx9XSwibGlmZWN5Y2xlIjp7fSwid2F0ZXJtYXJrIjp7fX0="),
         ];
@@ -295,6 +357,7 @@ mod tests {
 
         let expected_config = Settings {
             tid_header: "X-Numaflow-Id".into(),
+            pod_hash: "kddc".into(),
             app_listen_port: 8443,
             metrics_server_listen_port: 3001,
             upstream_addr: "localhost:8888".into(),
@@ -302,6 +365,8 @@ mod tests {
             store_type: StoreType::Nats,
             api_auth_token: None,
             js_callback_store: "test-kv-store".into(),
+            js_status_store: "test-kv-store".into(),
+            js_response_store: "test-kv-store".into(),
             nats_basic_auth: Some(("testuser".into(), "testpasswd".into())),
             js_message_stream: "serving-source-simple-pipeline".into(),
             jetstream_url: "nats://isbsvc-default-js-svc.default.svc:4222".into(),
diff --git a/rust/serving/src/error.rs b/rust/serving/src/error.rs
index 1129ba6d..6774b043 100644
--- a/rust/serving/src/error.rs
+++ b/rust/serving/src/error.rs
@@ -23,6 +23,9 @@ pub enum Error {
     // Store operation errors
     Store(String),
 
+    #[error("Status - {0}")]
+    Status(String),
+
     #[error("SubGraphNotFound Error - {0}")]
     // Sub Graph Not Found Error
     SubGraphNotFound(&'static str),
@@ -47,6 +50,9 @@ pub enum Error {
     #[error("Serving source error - {0}")]
     Source(String),
 
+    #[error("Duplicate request - {0}")]
+    Duplicate(String),
+
     #[error("Other Error - {0}")]
     // catch-all variant for now
     Other(String),
diff --git a/rust/serving/src/lib.rs b/rust/serving/src/lib.rs
index 809a5cd2..1414d763 100644
--- a/rust/serving/src/lib.rs
+++ b/rust/serving/src/lib.rs
@@ -4,12 +4,13 @@ use axum_server::tls_rustls::RustlsConfig;
 use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Duration;
+use tokio_util::sync::CancellationToken;
 use tracing::info;
 
 pub use self::error::{Error, Result};
 use crate::app::orchestrator::OrchestratorState as CallbackState;
 use crate::app::start_main_server;
-use crate::app::store::cbstore::jetstreamstore::JetStreamCallbackStore;
+use crate::app::store::cbstore::jetstream_store::JetStreamCallbackStore;
 use crate::app::store::datastore::jetstream::JetStreamDataStore;
 use crate::app::store::datastore::user_defined::UserDefinedStore;
 use crate::app::tracker::MessageGraph;
@@ -23,7 +24,7 @@ mod app;
 mod config;
 pub use {
     config::Settings, config::DEFAULT_CALLBACK_URL_HEADER_KEY, config::DEFAULT_ID_HEADER,
-    config::ENV_MIN_PIPELINE_SPEC,
+    config::DEFAULT_POD_HASH_KEY, config::ENV_MIN_PIPELINE_SPEC,
 };
 
 mod error;
@@ -32,6 +33,7 @@ mod pipeline;
 
 use crate::app::store::cbstore::CallbackStore;
 use crate::app::store::datastore::DataStore;
+use crate::app::store::status::StatusTracker;
 
 pub mod callback;
 
@@ -40,6 +42,7 @@ pub(crate) struct AppState<T, U> {
     pub(crate) js_context: Context,
     pub(crate) settings: Arc<Settings>,
     pub(crate) orchestrator_state: OrchestratorState<T, U>,
+    pub(crate) cancellation_token: CancellationToken,
 }
 
 /// Sets up and starts the main application server and the metrics server
@@ -47,6 +50,7 @@ pub(crate) struct AppState<T, U> {
 /// handles any errors that may occur during their execution.
 async fn serve<T, U>(
     app: AppState<T, U>,
+    cln_token: CancellationToken,
 ) -> std::result::Result<(), Box<dyn std::error::Error + Send + Sync + 'static>>
 where
     T: Clone + Send + Sync + DataStore + 'static,
@@ -64,11 +68,14 @@ where
     let metrics_addr: SocketAddr =
         format!("0.0.0.0:{}", &app.settings.metrics_server_listen_port).parse()?;
 
-    let metrics_server_handle =
-        tokio::spawn(start_https_metrics_server(metrics_addr, tls_config.clone()));
+    let metrics_server_handle = tokio::spawn(start_https_metrics_server(
+        metrics_addr,
+        tls_config.clone(),
+        cln_token.clone(),
+    ));
 
     // Start the main server, which serves the application.
-    let app_server_handle = tokio::spawn(start_main_server(app, tls_config));
+    let app_server_handle = tokio::spawn(start_main_server(app, tls_config, cln_token));
 
     // TODO: is try_join the best? we need to short-circuit at the first failure
     tokio::try_join!(flatten(app_server_handle), flatten(metrics_server_handle))?;
@@ -102,15 +109,21 @@ pub async fn run(config: Settings) -> Result<()> {
     start(js_context, Arc::new(config))
         .await
         .map_err(|e| Error::Source(e.to_string()))?;
-
     Ok(())
 }
 
 /// Starts the serving code after setting up the callback store, message DCG processor, and data store.
 async fn start(js_context: Context, settings: Arc<Settings>) -> Result<()> {
+    let cancel_token = CancellationToken::new();
+
     // create a callback store for tracking
-    let callback_store =
-        JetStreamCallbackStore::new(js_context.clone(), &settings.js_callback_store).await?;
+    let callback_store = JetStreamCallbackStore::new(
+        js_context.clone(),
+        &settings.pod_hash,
+        &settings.js_callback_store,
+        cancel_token.clone(),
+    )
+    .await?;
 
     // Create the message graph from the pipeline spec and the redis store
     let msg_graph = MessageGraph::from_pipeline(&settings.pipeline_spec).map_err(|e| {
@@ -123,28 +136,68 @@ async fn start(js_context: Context, settings: Arc<Settings>) -> Result<()> {
     // Create a store (builtin or user-defined) to store the callbacks and the custom responses
     match &settings.store_type {
         StoreType::Nats => {
-            let nats_store =
-                JetStreamDataStore::new(js_context.clone(), &settings.js_callback_store).await?;
-            let callback_state = CallbackState::new(msg_graph, nats_store, callback_store).await?;
+            let nats_store = JetStreamDataStore::new(
+                js_context.clone(),
+                &settings.js_response_store,
+                &settings.pod_hash,
+                cancel_token.clone(),
+            )
+            .await?;
+            let status_tracker = StatusTracker::new(
+                js_context.clone(),
+                &settings.js_status_store,
+                settings.pod_hash.clone(),
+                Some(settings.js_response_store.to_string()),
+            )
+            .await?;
+            let callback_state = CallbackState::new(
+                &settings.pod_hash,
+                msg_graph,
+                nats_store,
+                callback_store,
+                status_tracker,
+            )
+            .await?;
             let app = AppState {
                 js_context,
                 settings,
                 orchestrator_state: callback_state,
+                cancellation_token: cancel_token.clone(),
             };
-            serve(app).await.map_err(|e| Error::Store(e.to_string()))?
+            info!("Starting app");
+            serve(app, cancel_token.clone())
+                .await
+                .map_err(|e| Error::Store(e.to_string()))?
         }
         StoreType::UserDefined(ud_config) => {
+            let status_tracker = StatusTracker::new(
+                js_context.clone(),
+                &settings.js_status_store,
+                settings.pod_hash.clone(),
+                None,
+            )
+            .await?;
             let ud_store = UserDefinedStore::new(ud_config.clone()).await?;
-            let callback_state = CallbackState::new(msg_graph, ud_store, callback_store).await?;
+            let callback_state = CallbackState::new(
+                &settings.pod_hash,
+                msg_graph,
+                ud_store,
+                callback_store,
+                status_tracker,
+            )
+            .await?;
             let app = AppState {
                 js_context,
                 settings,
                 orchestrator_state: callback_state,
+                cancellation_token: cancel_token.clone(),
             };
-            serve(app).await.map_err(|e| Error::Store(e.to_string()))?
+            serve(app, cancel_token.clone())
+                .await
+                .map_err(|e| Error::Store(e.to_string()))?
         }
     }
-
+    info!("Done app");
     Ok(())
 }
 
@@ -165,8 +218,11 @@ mod tests {
         let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
 
         let env_vars = [
+            ("NUMAFLOW_POD", "serving-serve-cbdf"),
             ("NUMAFLOW_SERVING_APP_LISTEN_PORT", "32443"),
-            ("NUMAFLOW_SERVING_KV_STORE", "serving-test-run-kv-store"),
+            ("NUMAFLOW_SERVING_CALLBACK_STORE", "serving-test-run-kv-store"),
+            ("NUMAFLOW_SERVING_RESPONSE_STORE", "serving-test-run-kv-store"),
+            ("NUMAFLOW_SERVING_STATUS_STORE", "serving-test-run-kv-store"),
             ("NUMAFLOW_SERVING_SOURCE_SETTINGS", "eyJhdXRoIjpudWxsLCJzZXJ2aWNlIjp0cnVlLCJtc2dJREhlYWRlcktleSI6IlgtTnVtYWZsb3ctSWQifQ=="),
             ("NUMAFLOW_SERVING_MIN_PIPELINE_SPEC", "eyJ2ZXJ0aWNlcyI6W3sibmFtZSI6ImluIiwic291cmNlIjp7ImpldHN0cmVhbSI6eyJ1cmwiOiJuYXRzOi8vbG9jYWxob3N0OjQyMjIiLCJzdHJlYW0iOiJzZXJ2aW5nLXNvdXJjZS1zaW1wbGUtcGlwZWxpbmUiLCJ0bHMiOm51bGwsImF1dGgiOnsiYmFzaWMiOnsidXNlciI6eyJuYW1lIjoiaXNic3ZjLWRlZmF1bHQtanMtY2xpZW50LWF1dGgiLCJrZXkiOiJjbGllbnQtYXV0aC11c2VyIn0sInBhc3N3b3JkIjp7Im5hbWUiOiJpc2JzdmMtZGVmYXVsdC1qcy1jbGllbnQtYXV0aCIsImtleSI6ImNsaWVudC1hdXRoLXBhc3N3b3JkIn19fX19LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJlbnYiOlt7Im5hbWUiOiJOVU1BRkxPV19DQUxMQkFDS19FTkFCTEVEIiwidmFsdWUiOiJ0cnVlIn0seyJuYW1lIjoiTlVNQUZMT1dfU0VSVklOR19TT1VSQ0VfU0VUVElOR1MiLCJ2YWx1ZSI6ImV5SmhkWFJvSWpwdWRXeHNMQ0p6WlhKMmFXTmxJanAwY25WbExDSnRjMmRKUkVobFlXUmxja3RsZVNJNklsZ3RUblZ0WVdac2IzY3RTV1FpZlE9PSJ9LHsibmFtZSI6Ik5VTUFGTE9XX1NFUlZJTkdfS1ZfU1RPUkUiLCJ2YWx1ZSI6InNlcnZpbmctc3RvcmUtc2ltcGxlLXBpcGVsaW5lX1NFUlZJTkdfS1ZfU1RPUkUifV19LCJzY2FsZSI6eyJtaW4iOjEsIm1heCI6MX0sImluaXRDb250YWluZXJzIjpbeyJuYW1lIjoidmFsaWRhdGUtc3RyZWFtLWluaXQiLCJpbWFnZSI6InF1YXkuaW8vbnVtYXByb2ovbnVtYWZsb3c6ODA4MkU2NTYtQTAxOS00QjE1LUE2ODQtNTg2RkM3RDJDQTFGIiwiYXJncyI6WyJpc2JzdmMtdmFsaWRhdGUiLCItLWlzYnN2Yy10eXBlPWpldHN0cmVhbSIsIi0tYnVmZmVycz1zZXJ2aW5nLXNvdXJjZS1zaW1wbGUtcGlwZWxpbmUiXSwiZW52IjpbeyJuYW1lIjoiTlVNQUZMT1dfUElQRUxJTkVfTkFNRSIsInZhbHVlIjoicy1zaW1wbGUtcGlwZWxpbmUifSx7Im5hbWUiOiJHT0RFQlVHIn0seyJuYW1lIjoiTlVNQUZMT1dfSVNCU1ZDX0NPTkZJRyIsInZhbHVlIjoiZXlKcVpYUnpkSEpsWVcwaU9uc2lkWEpzSWpvaWJtRjBjem92TDJselluTjJZeTFrWldaaGRXeDBMV3B6TFhOMll5NWtaV1poZFd4MExuTjJZem8wTWpJeUlpd2lZWFYwYUNJNmV5SmlZWE5wWXlJNmV5SjFjMlZ5SWpwN0ltNWhiV1VpT2lKcGMySnpkbU10WkdWbVlYVnNkQzFxY3kxamJHbGxiblF0WVhWMGFDSXNJbXRsZVNJNkltTnNhV1Z1ZEMxaGRYUm9MWFZ6WlhJaWZTd2ljR0Z6YzNkdmNtUWlPbnNpYm1GdFpTSTZJbWx6WW5OMll5MWtaV1poZFd4MExXcHpMV05zYVdWdWRDMWhkWFJvSWl3aWEyVjVJam9pWTJ4cFpXNTBMV0YxZEdndGNHRnpjM2R2Y21RaWZYMTlMQ0p6ZEhKbFlXMURiMjVtYVdjaU9pSmpiMjV6ZFcxbGNqcGNiaUFnWVdOcmQyRnBkRG9nTmpCelhHNGdJRzFoZUdGamEzQmxibVJwYm1jNklESTFNREF3WEc1dmRHSjFZMnRsZERwY2JpQWdhR2x6ZEc5eWVUb2dNVnh1SUNCdFlYaGllWFJsY3pvZ01GeHVJQ0J0WVhoMllXeDFaWE5wZW1VNklEQmNiaUFnY21Wd2JHbGpZWE02SUROY2JpQWdjM1J2Y21GblpUb2dNRnh1SUNCMGRHdzZJRE5vWEc1d2NtOWpZblZqYTJWME9seHVJQ0JvYVhOMGIzSjVPaUF4WEc0Z0lHMWhlR0o1ZEdWek9pQXdYRzRnSUcxaGVIWmhiSFZsYzJsNlpUb2dNRnh1SUNCeVpYQnNhV05oY3pvZ00xeHVJQ0J6ZEc5eVlXZGxPaUF3WEc0Z0lIUjBiRG9nTnpKb1hHNXpkSEpsWVcwNlhHNGdJR1IxY0d4cFkyRjBaWE02SURZd2MxeHVJQ0J0WVhoaFoyVTZJRGN5YUZ4dUlDQnRZWGhpZVhSbGN6b2dMVEZjYmlBZ2JXRjRiWE5uY3pvZ01UQXdNREF3WEc0Z0lISmxjR3hwWTJGek9pQXpYRzRnSUhKbGRHVnVkR2x2YmpvZ01GeHVJQ0J6ZEc5eVlXZGxPaUF3WEc0aWZYMD0ifSx7Im5hbWUiOiJOVU1BRkxPV19JU0JTVkNfSkVUU1RSRUFNX1VSTCIsInZhbHVlIjoibmF0czovL2lzYnN2Yy1kZWZhdWx0LWpzLXN2Yy5kZWZhdWx0LnN2Yzo0MjIyIn0seyJuYW1lIjoiTlVNQUZMT1dfSVNCU1ZDX0pFVFNUUkVBTV9UTFNfRU5BQkxFRCIsInZhbHVlIjoiZmFsc2UifSx7Im5hbWUiOiJOVU1BRkxPV19JU0JTVkNfSkVUU1RSRUFNX1VTRVIiLCJ2YWx1ZUZyb20iOnsic2VjcmV0S2V5UmVmIjp7Im5hbWUiOiJpc2JzdmMtZGVmYXVsdC1qcy1jbGllbnQtYXV0aCIsImtleSI6ImNsaWVudC1hdXRoLXVzZXIifX19LHsibmFtZSI6Ik5VTUFGTE9XX0lTQlNWQ19KRVRTVFJFQU1fUEFTU1dPUkQiLCJ2YWx1ZUZyb20iOnsic2VjcmV0S2V5UmVmIjp7Im5hbWUiOiJpc2JzdmMtZGVmYXVsdC1qcy1jbGllbnQtYXV0aCIsImtleSI6ImNsaWVudC1hdXRoLXBhc3N3b3JkIn19fV0sInJlc291cmNlcyI6eyJyZXF1ZXN0cyI6eyJjcHUiOiIxMDBtIiwibWVtb3J5IjoiMTI4TWkifX0sImltYWdlUHVsbFBvbGljeSI6IklmTm90UHJlc2VudCJ9XSwidXBkYXRlU3RyYXRlZ3kiOnsidHlwZSI6IlJvbGxpbmdVcGRhdGUiLCJyb2xsaW5nVXBkYXRlIjp7Im1heFVuYXZhaWxhYmxlIjoiMjUlIn19fSx7Im5hbWUiOiJjYXQiLCJ1ZGYiOnsiY29udGFpbmVyIjp7ImltYWdlIjoicXVheS5pby9udW1haW8vbnVtYWZsb3ctZ28vbWFwLWZvcndhcmQtbWVzc2FnZTpzdGFibGUiLCJyZXNvdXJjZXMiOnt9fSwiYnVpbHRpbiI6bnVsbCwiZ3JvdXBCeSI6bnVsbH0sImNvbnRhaW5lclRlbXBsYXRlIjp7InJlc291cmNlcyI6e30sImVudiI6W3sibmFtZSI6Ik5VTUFGTE9XX0NBTExCQUNLX0VOQUJMRUQiLCJ2YWx1ZSI6InRydWUifSx7Im5hbWUiOiJOVU1BRkxPV19TRVJWSU5HX1NPVVJDRV9TRVRUSU5HUyIsInZhbHVlIjoiZXlKaGRYUm9JanB1ZFd4c0xDSnpaWEoyYVdObElqcDBjblZsTENKdGMyZEpSRWhsWVdSbGNrdGxlU0k2SWxndFRuVnRZV1pzYjNjdFNXUWlmUT09In0seyJuYW1lIjoiTlVNQUZMT1dfU0VSVklOR19LVl9TVE9SRSIsInZhbHVlIjoic2VydmluZy1zdG9yZS1zaW1wbGUtcGlwZWxpbmVfU0VSVklOR19LVl9TVE9SRSJ9XX0sInNjYWxlIjp7Im1pbiI6MSwibWF4IjoxfSwidXBkYXRlU3RyYXRlZ3kiOnsidHlwZSI6IlJvbGxpbmdVcGRhdGUiLCJyb2xsaW5nVXBkYXRlIjp7Im1heFVuYXZhaWxhYmxlIjoiMjUlIn19fSx7Im5hbWUiOiJvdXQiLCJzaW5rIjp7InVkc2luayI6eyJjb250YWluZXIiOnsiaW1hZ2UiOiJxdWF5LmlvL251bWFpby9udW1hZmxvdy1nby9zaW5rLXNlcnZlOnN0YWJsZSIsInJlc291cmNlcyI6e319fSwicmV0cnlTdHJhdGVneSI6e319LCJjb250YWluZXJUZW1wbGF0ZSI6eyJyZXNvdXJjZXMiOnt9LCJlbnYiOlt7Im5hbWUiOiJOVU1BRkxPV19DQUxMQkFDS19FTkFCTEVEIiwidmFsdWUiOiJ0cnVlIn0seyJuYW1lIjoiTlVNQUZMT1dfU0VSVklOR19TT1VSQ0VfU0VUVElOR1MiLCJ2YWx1ZSI6ImV5SmhkWFJvSWpwdWRXeHNMQ0p6WlhKMmFXTmxJanAwY25WbExDSnRjMmRKUkVobFlXUmxja3RsZVNJNklsZ3RUblZ0WVdac2IzY3RTV1FpZlE9PSJ9LHsibmFtZSI6Ik5VTUFGTE9XX1NFUlZJTkdfS1ZfU1RPUkUiLCJ2YWx1ZSI6InNlcnZpbmctc3RvcmUtc2ltcGxlLXBpcGVsaW5lX1NFUlZJTkdfS1ZfU1RPUkUifV19LCJzY2FsZSI6eyJtaW4iOjEsIm1heCI6MX0sInVwZGF0ZVN0cmF0ZWd5Ijp7InR5cGUiOiJSb2xsaW5nVXBkYXRlIiwicm9sbGluZ1VwZGF0ZSI6eyJtYXhVbmF2YWlsYWJsZSI6IjI1JSJ9fX1dLCJlZGdlcyI6W3siZnJvbSI6ImluIiwidG8iOiJjYXQiLCJjb25kaXRpb25zIjpudWxsfSx7ImZyb20iOiJjYXQiLCJ0byI6Im91dCIsImNvbmRpdGlvbnMiOm51bGx9XSwibGlmZWN5Y2xlIjp7fSwid2F0ZXJtYXJrIjp7fX0="),
         ];
@@ -194,7 +250,7 @@ mod tests {
             .unwrap();
 
         context
-            .create_stream(async_nats::jetstream::stream::Config {
+            .create_stream(jetstream::stream::Config {
                 name: message_stream_name.into(),
                 ..Default::default()
             })
diff --git a/rust/serving/src/metrics.rs b/rust/serving/src/metrics.rs
index a3cd9c53..d37b4c3e 100644
--- a/rust/serving/src/metrics.rs
+++ b/rust/serving/src/metrics.rs
@@ -13,12 +13,14 @@ use axum::{
     Router,
 };
 use axum_server::tls_rustls::RustlsConfig;
+use axum_server::Handle;
 use prometheus_client::encoding::text::encode;
 use prometheus_client::metrics::counter::Counter;
 use prometheus_client::metrics::family::Family;
 use prometheus_client::metrics::histogram::{exponential_buckets, Histogram};
 use prometheus_client::registry::Registry;
-use tracing::debug;
+use tokio_util::sync::CancellationToken;
+use tracing::{debug, info};
 
 use crate::Error::MetricsServer;
 
@@ -54,10 +56,10 @@ pub(crate) struct ServingMetrics {
     pub(crate) http_requests_count: Family<Vec<(String, String)>, Counter>,
     pub(crate) http_requests_duration: Family<Vec<(String, String)>, Histogram>,
 
-    pub(crate) cb_store_register_count: Counter,
-    pub(crate) cb_store_register_fail_count: Counter,
-    pub(crate) cb_store_register_duplicate_count: Counter,
-    pub(crate) cb_store_register_duration: Histogram,
+    pub(crate) request_register_count: Counter,
+    pub(crate) request_register_fail_count: Counter,
+    pub(crate) request_register_duplicate_count: Counter,
+    pub(crate) request_register_duration: Histogram,
 
     pub(crate) payload_save_duration: Histogram,
     pub(crate) datum_retrive_duration: Histogram,
@@ -86,10 +88,10 @@ impl ServingMetrics {
         let metrics = Self {
             http_requests_count: http_requests_total,
             http_requests_duration,
-            cb_store_register_count,
-            cb_store_register_fail_count,
-            cb_store_register_duplicate_count,
-            cb_store_register_duration,
+            request_register_count: cb_store_register_count,
+            request_register_fail_count: cb_store_register_fail_count,
+            request_register_duplicate_count: cb_store_register_duplicate_count,
+            request_register_duration: cb_store_register_duration,
             payload_save_duration,
             datum_retrive_duration,
             processing_time,
@@ -109,27 +111,27 @@ impl ServingMetrics {
         );
 
         registry.register(
-            "CALLBACK_STORE_REGISTER",
-            "A Counter to keep track of the number of callback store register requests",
-            metrics.cb_store_register_count.clone(),
+            "REQUEST_REGISTER",
+            "A Counter to keep track of the number of register requests",
+            metrics.request_register_count.clone(),
         );
 
         registry.register(
-            "CALLBACK_STORE_REGISTER_FAIL",
-            "A Counter to keep track of the number of failed callback store register requests",
-            metrics.cb_store_register_fail_count.clone(),
+            "REQUEST_REGISTER_FAIL",
+            "A Counter to keep track of the number of failed register requests",
+            metrics.request_register_fail_count.clone(),
         );
 
         registry.register(
-            "CALLBACK_STORE_REGISTER_DUPLICATES",
-            "A Counter to keep track of the number of failed callback store register requests due to duplicate request id",
-            metrics.cb_store_register_fail_count.clone(),
+            "REQUEST_REGISTER_DUPLICATES",
+            "A Counter to keep track of the number of failed register requests due to duplicate request id",
+            metrics.request_register_fail_count.clone(),
         );
 
         registry.register(
-            "CALLBACK_STORE_REGISTER_DURATION",
-            "A Histogram to keep track of the duration of the successful callback store register requests",
-            metrics.cb_store_register_duration.clone(),
+            "REQUEST_REGISTER_DURATION",
+            "A Histogram to keep track of the duration of the successful register of requests",
+            metrics.request_register_duration.clone(),
         );
 
         registry.register(
@@ -163,15 +165,24 @@ pub(crate) fn serving_metrics() -> &'static ServingMetrics {
 pub(crate) async fn start_https_metrics_server(
     addr: SocketAddr,
     tls_config: RustlsConfig,
+    cln_token: CancellationToken,
 ) -> crate::Result<()> {
     let metrics_app = Router::new().route("/metrics", get(metrics_handler));
 
-    tracing::info!(?addr, "Starting metrics server");
+    let handle = Handle::new();
+    let shutdown_handle = handle.clone();
+    tokio::spawn(async move {
+        cln_token.cancelled().await;
+        shutdown_handle.shutdown();
+    });
+
     axum_server::bind_rustls(addr, tls_config)
+        .handle(handle)
         .serve(metrics_app.into_make_service())
         .await
         .map_err(|e| MetricsServer(format!("Starting web server for metrics: {}", e)))?;
 
+    info!(?addr, "Metrics server stopped");
     Ok(())
 }
 
@@ -254,7 +265,8 @@ mod tests {
 
         let addr = SocketAddr::from(([127, 0, 0, 1], 0));
         let server = tokio::spawn(async move {
-            let result = start_https_metrics_server(addr, tls_config).await;
+            let result =
+                start_https_metrics_server(addr, tls_config, CancellationToken::new()).await;
             assert!(result.is_ok())
         });
 



Test patch:
diff --git a/pkg/apis/numaflow/v1alpha1/serving_pipeline_types_test.go b/pkg/apis/numaflow/v1alpha1/serving_pipeline_types_test.go
index bfdeccc8..68b5465d 100644
--- a/pkg/apis/numaflow/v1alpha1/serving_pipeline_types_test.go
+++ b/pkg/apis/numaflow/v1alpha1/serving_pipeline_types_test.go
@@ -272,7 +272,9 @@ func Test_GetPipelineObj(t *testing.T) {
 	// Validate environment variables
 	envVars := sourceVertex.ContainerTemplate.Env
 	assert.Contains(t, envVars, corev1.EnvVar{Name: "NUMAFLOW_SERVING_SOURCE_SETTINGS", Value: "eyJhdXRoIjpudWxsLCJzZXJ2aWNlIjp0cnVlLCJtc2dJREhlYWRlcktleSI6bnVsbCwiY29udGFpbmVyVGVtcGxhdGUiOnsicmVzb3VyY2VzIjp7InJlcXVlc3RzIjp7Im1lbW9yeSI6IjE5NDJNaSJ9fSwiZW52IjpbeyJuYW1lIjoiVEVTVF9FTlYiLCJ2YWx1ZSI6InRlc3QtdmFsdWUifV19fQ=="})
-	assert.Contains(t, envVars, corev1.EnvVar{Name: "NUMAFLOW_SERVING_KV_STORE", Value: "serving-store-test-serving-pipeline_SERVING_KV_STORE"})
+	assert.Contains(t, envVars, corev1.EnvVar{Name: "NUMAFLOW_SERVING_CALLBACK_STORE", Value: "serving-store-test-serving-pipeline_SERVING_CALLBACK_STORE"})
+	assert.Contains(t, envVars, corev1.EnvVar{Name: "NUMAFLOW_SERVING_RESPONSE_STORE", Value: "serving-store-test-serving-pipeline_SERVING_RESPONSE_STORE"})
+	assert.Contains(t, envVars, corev1.EnvVar{Name: "NUMAFLOW_SERVING_STATUS_STORE", Value: "serving-store-test-serving-pipeline_SERVING_STATUS_STORE"})
 
 	servingDeployReq := GetServingPipelineResourceReq{
 		ISBSvcConfig: BufferServiceConfig{JetStream: &JetStreamConfig{URL: "nats://test-url"}},