feat: support nats streaming

Author: 0x3bfc

pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "httpx>=0.24.0,<1.0.0",         # For HTTP client
     "structlog>=25.4.0,<26.0.0",    # For structured logging
     "python-dotenv>=1.0.0,<2.0.0",  # For environment variables
+    "nats-py>=2.6.0,<3.0.0",        # For NATS streaming
     # Note: Langflow dependencies will be added via path or package reference
 ]
 

src/node/api.py

@@ -8,6 +8,7 @@
 import os
 import sys
 import time
+import uuid
 from typing import Any
 
 from fastapi import FastAPI, HTTPException
@@ -43,6 +44,29 @@
 
 app = FastAPI(title="Langflow Executor Node", version="0.1.0")
 
+# Initialize NATS client (lazy connection)
+_nats_client = None
+
+
+async def get_nats_client():
+    """Get or create NATS client instance."""
+    global _nats_client
+    if _nats_client is None:
+        logger.info("[NATS] Creating new NATS client instance...")
+        from node.nats import NATSClient
+        nats_url = os.getenv("NATS_URL", "nats://localhost:4222")
+        logger.info(f"[NATS] Connecting to NATS at {nats_url}")
+        _nats_client = NATSClient(nats_url=nats_url)
+        try:
+            await _nats_client.connect()
+            logger.info("[NATS] ✅ Successfully connected to NATS")
+        except Exception as e:
+            logger.warning(f"[NATS] ❌ Failed to connect to NATS (non-critical): {e}", exc_info=True)
+            _nats_client = None
+    else:
+        logger.debug("[NATS] Using existing NATS client instance")
+    return _nats_client
+
 
 class ComponentState(BaseModel):
     """Component state for execution."""
@@ -55,6 +79,7 @@ class ComponentState(BaseModel):
     config: dict[str, Any] | None = None
     display_name: str | None = None
     component_id: str | None = None
+    stream_topic: str | None = None  # NATS stream topic for publishing results
 
 
 class ExecutionRequest(BaseModel):
@@ -64,6 +89,7 @@ class ExecutionRequest(BaseModel):
     method_name: str
     is_async: bool = False
     timeout: int = 30
+    message_id: str | None = None  # Unique message ID from backend for tracking published messages
 
 
 class ExecutionResponse(BaseModel):
@@ -74,6 +100,7 @@ class ExecutionResponse(BaseModel):
     result_type: str
     execution_time: float
     error: str | None = None
+    message_id: str | None = None  # Unique ID for the published NATS message
 
 
 async def load_component_class(
@@ -439,12 +466,16 @@ async def execute_component(request: ExecutionRequest) -> ExecutionResponse:
 
     try:
         # Log what we received
-        logger.info(
+        stream_topic_value = request.component_state.stream_topic
+        log_msg = (
             f"Received execution request: "
             f"class={request.component_state.component_class}, "
             f"module={request.component_state.component_module}, "
-            f"code_length={len(request.component_state.component_code or '') if request.component_state.component_code else 0}"
+            f"code_length={len(request.component_state.component_code or '') if request.component_state.component_code else 0}, "
+            f"stream_topic={stream_topic_value}"
         )
+        logger.info(log_msg)
+        print(f"[EXECUTOR] {log_msg}")  # Also print to ensure visibility
 
         # Load component class dynamically
         component_class = await load_component_class(
@@ -517,11 +548,52 @@ async def execute_component(request: ExecutionRequest) -> ExecutionResponse:
             f"in {execution_time:.3f}s, result type: {type(result).__name__}"
         )
 
+        # Use message_id from request (generated by backend) or generate one if not provided
+        message_id = request.message_id or str(uuid.uuid4())
+        
+        # Publish result to NATS stream if topic is provided
+        if request.component_state.stream_topic:
+            topic = request.component_state.stream_topic
+            logger.info(f"[NATS] Attempting to publish to topic: {topic} with message_id: {message_id}")
+            print(f"[NATS] Attempting to publish to topic: {topic} with message_id: {message_id}")
+            try:
+                nats_client = await get_nats_client()
+                if nats_client:
+                    logger.info(f"[NATS] NATS client obtained, preparing publish data...")
+                    print(f"[NATS] NATS client obtained, preparing publish data...")
+                    # Publish result to NATS with message ID from backend
+                    publish_data = {
+                        "message_id": message_id,  # Use message_id from backend request
+                        "component_id": request.component_state.component_id,
+                        "component_class": request.component_state.component_class,
+                        "result": serialized_result,
+                        "result_type": type(result).__name__,
+                        "execution_time": execution_time,
+                    }
+                    logger.info(f"[NATS] Publishing to topic: {topic}, message_id: {message_id}, data keys: {list(publish_data.keys())}")
+                    print(f"[NATS] Publishing to topic: {topic}, message_id: {message_id}, data keys: {list(publish_data.keys())}")
+                    # Use the topic directly (already in format: droq.local.public.userid.workflowid.component.out)
+                    await nats_client.publish(topic, publish_data)
+                    logger.info(f"[NATS] ✅ Successfully published result to NATS topic: {topic} with message_id: {message_id}")
+                    print(f"[NATS] ✅ Successfully published result to NATS topic: {topic} with message_id: {message_id}")
+                else:
+                    logger.warning(f"[NATS] NATS client is None, cannot publish")
+                    print(f"[NATS] ⚠️  NATS client is None, cannot publish")
+            except Exception as e:
+                # Non-critical: log but don't fail execution
+                logger.warning(f"[NATS] ❌ Failed to publish to NATS (non-critical): {e}", exc_info=True)
+                print(f"[NATS] ❌ Failed to publish to NATS (non-critical): {e}")
+        else:
+            msg = f"[NATS] ⚠️  No stream_topic provided in request, skipping NATS publish. Component: {request.component_state.component_class}, ID: {request.component_state.component_id}"
+            logger.info(msg)
+            print(msg)
+
         return ExecutionResponse(
             result=serialized_result,
             success=True,
             result_type=type(result).__name__,
             execution_time=execution_time,
+            message_id=message_id,  # Return message ID (from request or generated) so backend can match it
         )
 
     except asyncio.TimeoutError:

src/node/nats.py

@@ -54,20 +54,39 @@ async def _ensure_stream(self) -> None:
         """Ensure the JetStream exists, create if it doesn't."""
         try:
             # Try to get stream info
-            await self.js.stream_info(self.stream_name)
+            stream_info = await self.js.stream_info(self.stream_name)
             logger.info(f"Stream '{self.stream_name}' already exists")
-        except Exception:
+            logger.info(f"Stream subjects: {stream_info.config.subjects}")
+            
+            # Check if 'droq.local.public.>' is in subjects, if not, update stream
+            required_subject = "droq.local.public.>"
+            if required_subject not in stream_info.config.subjects:
+                logger.warning(f"Stream '{self.stream_name}' missing required subject '{required_subject}', updating...")
+                subjects = list(stream_info.config.subjects) + [required_subject]
+                await self.js.update_stream(
+                    StreamConfig(
+                        name=self.stream_name,
+                        subjects=subjects,
+                        retention=stream_info.config.retention,
+                        storage=stream_info.config.storage,
+                    )
+                )
+                logger.info(f"Stream '{self.stream_name}' updated with subject '{required_subject}'")
+        except Exception as e:
             # Stream doesn't exist, create it
-            logger.info(f"Creating stream '{self.stream_name}'")
+            logger.info(f"Creating stream '{self.stream_name}' (error: {e})")
             await self.js.add_stream(
                 StreamConfig(
                     name=self.stream_name,
-                    subjects=[f"{self.stream_name}.>"],
+                    subjects=[
+                        f"{self.stream_name}.>",  # Backward compatibility
+                        "droq.local.public.>",    # Full topic path format
+                    ],
                     retention=RetentionPolicy.WORK_QUEUE,
                     storage=StorageType.FILE,
                 )
             )
-            logger.info(f"Stream '{self.stream_name}' created")
+            logger.info(f"Stream '{self.stream_name}' created with subjects: ['{self.stream_name}.>', 'droq.local.public.>']")
 
     async def publish(
         self,
@@ -79,27 +98,34 @@ async def publish(
         Publish a message to a NATS subject.
 
         Args:
-            subject: NATS subject to publish to
+            subject: NATS subject to publish to (can be full topic path or relative)
             data: Data to publish (will be JSON encoded)
             headers: Optional headers to include
         """
         if not self.js:
             raise RuntimeError("Not connected to NATS. Call connect() first.")
 
         try:
-            # Full subject with stream prefix
-            full_subject = f"{self.stream_name}.{subject}"
+            # If subject starts with "droq.", use it as full topic path
+            # Otherwise, prefix with stream name for backward compatibility
+            if subject.startswith("droq."):
+                full_subject = subject
+            else:
+                full_subject = f"{self.stream_name}.{subject}"
 
             # Encode data as JSON
             payload = json.dumps(data).encode()
+            payload_size = len(payload)
+            
+            logger.info(f"[NATS] Publishing to subject: {full_subject}, payload size: {payload_size} bytes")
 
             # Publish with headers if provided
             if headers:
-                await self.js.publish(full_subject, payload, headers=headers)
+                ack = await self.js.publish(full_subject, payload, headers=headers)
             else:
-                await self.js.publish(full_subject, payload)
+                ack = await self.js.publish(full_subject, payload)
 
-            logger.debug(f"Published message to {full_subject}")
+            logger.info(f"[NATS] ✅ Published message to {full_subject} (seq: {ack.seq if hasattr(ack, 'seq') else 'N/A'})")
         except Exception as e:
             logger.error(f"Failed to publish message: {e}")
             raise