fix: address review comments — SDK compatibility, race condition, wording

1. Python SDK: Add STREAM_KEEPALIVE to ActionType enum so the SDK
   recognizes the new value instead of crashing on unknown enum.

2. Python SDK: Make convert_proto_enum_to_python return None for
   unknown enum values instead of raising KeyError. This makes the
   SDK forward-compatible with future ActionType additions.

3. Python SDK: Skip STREAM_KEEPALIVE (and any unknown) actions early
   in the action listener generator before attempting to parse payload.

4. Engine: Use subscribedWorker.sendMu mutex for keepalive sends to
   avoid racing with real AssignedAction dispatches from other goroutines.
   gRPC server streams are not safe for concurrent sends.

5. Engine: Update comment to say "periodic" instead of "idle" — the
   keepalive fires unconditionally at the interval, not only when idle.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

Author: shaunpatterson

internal/services/dispatcher/server.go

@@ -343,7 +343,8 @@ func (s *DispatcherImpl) ListenV2(request *contracts.WorkerListenRequest, stream
 
 	fin := make(chan bool)
 
-	s.workers.Add(workerId, sessionId, newSubscribedWorker(stream, fin, workerId, s.defaultMaxWorkerBacklogSize, s.pubBuffer))
+	sw := newSubscribedWorker(stream, fin, workerId, s.defaultMaxWorkerBacklogSize, s.pubBuffer)
+	s.workers.Add(workerId, sessionId, sw)
 
 	defer func() {
 		// non-blocking send
@@ -355,10 +356,12 @@ func (s *DispatcherImpl) ListenV2(request *contracts.WorkerListenRequest, stream
 		s.workers.DeleteForSession(workerId, sessionId)
 	}()
 
-	// Optionally send periodic keepalive messages on the stream to prevent
-	// L7 proxies (e.g., Envoy on Azure Container Apps) from killing the stream
-	// due to stream_idle_timeout. This does NOT affect worker liveness detection,
-	// which is handled by the separate Heartbeat RPC.
+	// Optionally send periodic keepalive messages on the ListenV2 stream to
+	// prevent L7 proxies (e.g., Envoy on Azure Container Apps) from killing the
+	// stream due to stream_idle_timeout. Keepalives are sent unconditionally at
+	// the configured interval regardless of other stream activity. This does NOT
+	// affect worker liveness detection, which is handled by the separate
+	// Heartbeat RPC.
 	var keepaliveTicker *time.Ticker
 	var keepaliveC <-chan time.Time
 
@@ -397,9 +400,13 @@ func (s *DispatcherImpl) ListenV2(request *contracts.WorkerListenRequest, stream
 
 			return nil
 		case <-keepaliveC:
-			if err := stream.Send(&contracts.AssignedAction{
+			sw.sendMu.Lock()
+			err := stream.Send(&contracts.AssignedAction{
 				ActionType: contracts.ActionType_STREAM_KEEPALIVE,
-			}); err != nil {
+			})
+			sw.sendMu.Unlock()
+
+			if err != nil {
 				s.l.Debug().Err(err).Msgf("failed to send stream keepalive for worker %s, closing stream", request.WorkerId)
 				return nil
 			}

sdks/python/hatchet_sdk/clients/dispatcher/action_listener.py

@@ -215,6 +215,17 @@ async def _generator(self) -> AsyncGenerator[Action | None, None]:
 
                     assigned_action = result.data
 
+                    # Skip keepalive messages — they only exist to prevent
+                    # L7 proxies from killing idle streams.
+                    action_type = convert_proto_enum_to_python(
+                        assigned_action.action_type,
+                        ActionType,
+                        ActionTypeProto,
+                    )
+
+                    if action_type is None or action_type == ActionType.STREAM_KEEPALIVE:
+                        continue
+
                     try:
                         action_payload = (
                             ActionPayload()
@@ -239,11 +250,7 @@ async def _generator(self) -> AsyncGenerator[Action | None, None]:
                         step_run_id=assigned_action.task_run_external_id,
                         action_id=assigned_action.action_id,
                         action_payload=action_payload,
-                        action_type=convert_proto_enum_to_python(
-                            assigned_action.action_type,
-                            ActionType,
-                            ActionTypeProto,
-                        ),
+                        action_type=action_type,
                         retry_count=assigned_action.retry_count,
                         additional_metadata=parse_additional_metadata(
                             assigned_action.additional_metadata

sdks/python/hatchet_sdk/runnables/action.py

@@ -49,6 +49,7 @@ def validate_filter_payload(self) -> "ActionPayload":
 class ActionType(str, Enum):
     START_STEP_RUN = "START_STEP_RUN"
     CANCEL_STEP_RUN = "CANCEL_STEP_RUN"
+    STREAM_KEEPALIVE = "STREAM_KEEPALIVE"
 
 
 class Action(BaseModel):

sdks/python/hatchet_sdk/utils/proto_enums.py

@@ -44,4 +44,12 @@ def convert_proto_enum_to_python(
     if value is None:
         return None
 
-    return python_enum_class[proto_enum.Name(value)]
+    name = proto_enum.Name(value)
+
+    try:
+        return python_enum_class[name]
+    except KeyError:
+        # Unknown enum values (e.g., from a newer server) are returned as None
+        # so callers can skip them gracefully. This prevents SDK crashes when the
+        # engine sends new ActionType values like STREAM_KEEPALIVE.
+        return None