feature for agent outputs in events & streaming fixes

flo_ai/flo_ai/arium/arium.py

@@ -507,6 +507,7 @@ async def _execute_node(
                         node_name=node.name,
                         node_type=node_type,
                         execution_time=execution_time,
+                        node_output=self._serialize_node_output(result),
                     )
 
                     return result
@@ -561,6 +562,7 @@ async def _execute_node(
                     node_name=node.name,
                     node_type=node_type,
                     execution_time=execution_time,
+                    node_output=self._serialize_node_output(result),
                 )
 
                 return result
@@ -638,3 +640,25 @@ def _add_to_memory(self, message: MessageMemoryItem):
         Store message in memory
         """
         self.memory.add(message)
+
+    def _serialize_node_output(self, result: Any) -> Optional[str]:
+        if result is None:
+            return None
+        if isinstance(result, str):
+            return result
+        if isinstance(result, list):
+            parts = [self._serialize_node_output(item) for item in result]
+            return '\n'.join(p for p in parts if p) or None
+        if hasattr(result, 'content'):
+            return self._serialize_node_output(result.content)
+        if hasattr(result, 'text'):
+            return result.text
+        # DocumentMessageContent / ImageMessageContent — show url or type label
+        media_type = getattr(result, 'type', None)
+        if media_type in ('document', 'image'):
+            url = getattr(result, 'url', None)
+            mime = getattr(result, 'mime_type', None)
+            if url:
+                return f'[{media_type}: {url}]'
+            return f'[{media_type}{f": {mime}" if mime else ""}]'
+        return str(result)

flo_ai/flo_ai/arium/events.py

@@ -48,6 +48,7 @@ class AriumEvent:
     execution_time: Optional[float] = None
     error: Optional[str] = None
     router_choice: Optional[str] = None
+    node_output: Optional[str] = None
     metadata: Optional[dict] = None
 
 

wavefront/client/src/components/InferencePopup.tsx

@@ -254,6 +254,7 @@ const InferencePopup: React.FC<InferencePopupProps> = ({ onClose, renderModal =
             document_type: doc.documentType,
             document_base64: doc.base64Content,
             mime_type: doc.mimeType,
+            file_name: doc.file.name,
             metadata: {
               filename: doc.file.name,
               size: doc.file.size,

wavefront/client/src/components/Stream.tsx

@@ -10,6 +10,7 @@ interface StreamProps {
     execution_time?: number;
     error?: string;
     router_choice?: string;
+    node_output?: string;
   }>;
   isStreaming?: boolean;
   eventsContainerRef?: RefObject<HTMLDivElement | null>;
@@ -74,6 +75,12 @@ const Stream: React.FC<StreamProps> = ({ listenEventsEnabled, streamingEvents, i
                   {'router_choice' in event && event.router_choice && (
                     <div className="mt-1 text-blue-600">Router choice: {event.router_choice}</div>
                   )}
+                  {'node_output' in event && event.node_output && (
+                    <div className="mt-1 rounded border border-gray-200 bg-white p-2">
+                      <span className="font-medium text-gray-500">Output: </span>
+                      <span className="break-words whitespace-pre-wrap text-gray-700">{event.node_output}</span>
+                    </div>
+                  )}
                 </div>
               ))}
             </div>

wavefront/client/src/pages/apps/[appId]/agents/[id].tsx

@@ -432,6 +432,7 @@ const AgentDetail: React.FC = () => {
             document_type: doc.documentType,
             document_base64: doc.base64Content,
             mime_type: doc.mimeType,
+            file_name: doc.file.name,
             metadata: {
               filename: doc.file.name,
               size: doc.file.size,

wavefront/client/src/pages/apps/[appId]/workflows/[id].tsx

@@ -391,6 +391,7 @@ const WorkflowDetail: React.FC = () => {
             document_type: doc.documentType,
             document_base64: doc.base64Content,
             mime_type: doc.mimeType,
+            file_name: doc.file.name,
             metadata: {
               filename: doc.file.name,
               size: doc.file.size,

wavefront/client/src/types/chat-message.ts

@@ -8,6 +8,7 @@ export interface DocumentContent {
   document_type: string;
   document_base64?: string;
   mime_type?: string;
+  file_name?: string;
   metadata?: {
     filename?: string;
     size?: number;

wavefront/client/src/types/workflow.ts

@@ -79,6 +79,7 @@ export interface WorkflowEventBase {
   execution_time?: number;
   error?: string;
   router_choice?: string;
+  node_output?: string;
   metadata?: Record<string, unknown>;
 }
 
@@ -105,6 +106,7 @@ export interface NodeCompletedEvent extends WorkflowEventBase {
   event_type: 'node_completed';
   node_name: string;
   execution_time: number;
+  node_output?: string;
 }
 
 export interface NodeFailedEvent extends WorkflowEventBase {

wavefront/server/apps/floconsole/floconsole/di/application_container.py

@@ -66,8 +66,12 @@ class ApplicationContainer(containers.DeclarativeContainer):
         app_user_repository=app_user_repository,
     )
 
-    kms_service = providers.Singleton(
-        FloKmsService, cloud_provider=config.cloud_config.cloud_provider
+    kms_service = providers.Selector(
+        config.jwt_token.enable_cloud_kms,
+        true=providers.Singleton(
+            FloKmsService, cloud_provider=config.cloud_config.cloud_provider
+        ),
+        false=providers.Object(None),
     )
 
     token_service = providers.Singleton(

wavefront/server/modules/agents_module/agents_module/controllers/workflow_controller.py

@@ -6,6 +6,8 @@
 from dependency_injector.wiring import inject, Provide
 import json
 import asyncio
+import uuid
+import time
 
 from common_module.log.logger import logger
 from common_module.response_formatter import ResponseFormatter
@@ -75,9 +77,6 @@ async def workflow_inference(
         f'Starting inference for namespace: {namespace}, workflow_id: {workflow_id}, listen_events: {listen_events}'
     )
 
-    # Extract user_id from authenticated session
-    user_id = request.state.session.user_id
-
     # Extract authentication credentials
     access_token, app_key = extract_auth_credentials(request)
 
@@ -89,27 +88,26 @@ async def workflow_inference(
     events_filter = None
 
     if listen_events or request_body.listen_events:
-        event_callback = create_workflow_event_callback(user_id, namespace, workflow_id)
+        execution_id = str(uuid.uuid4())
+        event_callback = create_workflow_event_callback(
+            execution_id, namespace, workflow_id
+        )
         events_filter = DEFAULT_EVENTS_FILTER
         logger.info(
-            f'Event streaming enabled for user {user_id}, workflow {namespace}/{workflow_id}'
+            f'Event streaming enabled for execution {execution_id}, workflow {namespace}/{workflow_id}'
         )
 
     # Check if streaming is requested
     if listen_events or request_body.listen_events:
         logger.info(
-            f'Streaming inference for user {user_id}, workflow {namespace}/{workflow_id}'
+            f'Streaming inference for execution {execution_id}, workflow {namespace}/{workflow_id}'
         )
 
-        # Get or create event queue for this user-workflow
-        event_queue = event_streamer.get_or_create_queue(
-            user_id, namespace, workflow_id
-        )
+        event_queue = event_streamer.create_queue(execution_id)
 
         async def generate_inference_stream():
             """Generate streaming inference with events and final output"""
             try:
-                # Start inference in background task
                 inference_task = asyncio.create_task(
                     workflow_inference_service.perform_inference(
                         workflow_name=workflow_id,
@@ -126,60 +124,54 @@ async def generate_inference_stream():
                     )
                 )
 
-                # Stream events while workflow is running
-                workflow_completed = False
-                while not workflow_completed and not inference_task.done():
+                # Stream events until inference completes
+                while not inference_task.done():
                     try:
-                        # Wait for event with timeout
                         event_data = await asyncio.wait_for(
                             event_queue.get(), timeout=1.0
                         )
                         yield f'data: {json.dumps(event_data)}\n\n'
-                        await asyncio.sleep(0.1)  # remove it later
-
-                        # Check if workflow ended
-                        if event_data.get('event_type') in [
-                            'workflow_completed',
-                            'workflow_failed',
-                        ]:
-                            workflow_completed = True
-
                     except asyncio.TimeoutError:
-                        # Continue waiting if no events
                         continue
 
-                # Wait for inference to complete and get result
+                # Yield to the event loop so any ensure_future(add_event(...))
+                # callbacks scheduled inside the inference task have a chance
+                # to run and enqueue their events before we drain.
+                await asyncio.sleep(0)
+
+                # Drain any remaining events queued after task completion
+                while not event_queue.empty():
+                    event_data = event_queue.get_nowait()
+                    yield f'data: {json.dumps(event_data)}\n\n'
+
                 result, execution_time = await inference_task
 
-                # Send final output event
                 output_event = {
                     'event_type': 'output',
                     'result': result,
                     'workflow_id': workflow_id,
                     'namespace': namespace,
                     'execution_time': execution_time,
-                    'timestamp': asyncio.get_event_loop().time(),
+                    'timestamp': time.time(),
                 }
                 yield f'data: {json.dumps(output_event)}\n\n'
-                await asyncio.sleep(0.1)  # remove it later
 
                 logger.info(
-                    f'Streaming inference completed for user {user_id}, workflow {namespace}/{workflow_id}'
+                    f'Streaming inference completed for execution {execution_id}, workflow {namespace}/{workflow_id}'
                 )
 
             except Exception as e:
                 logger.error(
-                    f'Error in streaming inference for user {user_id}, workflow {namespace}/{workflow_id}: {e}'
+                    f'Error in streaming inference for execution {execution_id}, workflow {namespace}/{workflow_id}: {e}'
                 )
                 error_event = {
                     'event_type': 'error',
                     'error': str(e),
-                    'timestamp': asyncio.get_event_loop().time(),
+                    'timestamp': time.time(),
                 }
                 yield f'data: {json.dumps(error_event)}\n\n'
             finally:
-                # Clean up queue
-                event_streamer.cleanup_queue(user_id, namespace, workflow_id)
+                event_streamer.cleanup_queue(execution_id)
 
         return StreamingResponse(
             generate_inference_stream(),
@@ -278,10 +270,6 @@ async def workflow_inference_v2(
     logger.info(
         f'Starting v2 inference for workflow_id: {workflow_id}, listen_events: {listen_events}'
     )
-
-    # Extract user_id from authenticated session
-    user_id = request.state.session.user_id
-
     # Extract authentication credentials
     access_token, app_key = extract_auth_credentials(request)
 
@@ -298,30 +286,26 @@ async def workflow_inference_v2(
     events_filter = None
 
     if listen_events or request_body.listen_events:
-        # Use real namespace and workflow name for event streaming
+        execution_id = str(uuid.uuid4())
         event_callback = create_workflow_event_callback(
-            user_id, namespace, workflow_name
+            execution_id, namespace, workflow_name
         )
         events_filter = DEFAULT_EVENTS_FILTER
         logger.info(
-            f'Event streaming enabled for user {user_id}, workflow {namespace}/{workflow_name}'
+            f'Event streaming enabled for execution {execution_id}, workflow {namespace}/{workflow_name}'
         )
 
     # Check if streaming is requested
     if listen_events or request_body.listen_events:
         logger.info(
-            f'Streaming inference for user {user_id}, workflow {namespace}/{workflow_name}'
+            f'Streaming inference for execution {execution_id}, workflow {namespace}/{workflow_name}'
         )
 
-        # Get or create event queue for this user-workflow
-        event_queue = event_streamer.get_or_create_queue(
-            user_id, namespace, workflow_name
-        )
+        event_queue = event_streamer.create_queue(execution_id)
 
         async def generate_inference_stream():
             """Generate streaming inference with events and final output"""
             try:
-                # Start inference in background task
                 inference_task = asyncio.create_task(
                     workflow_inference_service.perform_inference_v2(
                         workflow_data=workflow_data,
@@ -337,66 +321,54 @@ async def generate_inference_stream():
                     )
                 )
 
-                # Stream events while workflow is running
-                workflow_completed = False
-                while not workflow_completed and not inference_task.done():
+                # Stream events until inference completes
+                while not inference_task.done():
                     try:
-                        # Wait for event with timeout
                         event_data = await asyncio.wait_for(
                             event_queue.get(), timeout=1.0
                         )
                         yield f'data: {json.dumps(event_data)}\n\n'
-                        await asyncio.sleep(0.1)  # remove it later
-
-                        # Check if workflow ended
-                        if event_data.get('event_type') in [
-                            'workflow_completed',
-                            'workflow_failed',
-                        ]:
-                            workflow_completed = True
-
                     except asyncio.TimeoutError:
-                        # Continue waiting if no events
                         continue
 
-                # Wait for inference to complete and get result
+                # Yield to the event loop so any ensure_future(add_event(...))
+                # callbacks scheduled inside the inference task have a chance
+                # to run and enqueue their events before we drain.
+                await asyncio.sleep(0)
+
+                # Drain any remaining events queued after task completion
+                while not event_queue.empty():
+                    event_data = event_queue.get_nowait()
+                    yield f'data: {json.dumps(event_data)}\n\n'
+
                 result, execution_time = await inference_task
 
-                # Send final output event
                 output_event = {
                     'event_type': 'output',
                     'result': result,
                     'workflow_id': workflow_name,
                     'namespace': namespace,
                     'execution_time': execution_time,
-                    'timestamp': asyncio.get_event_loop().time(),
+                    'timestamp': time.time(),
                 }
                 yield f'data: {json.dumps(output_event)}\n\n'
-                await asyncio.sleep(0.1)  # remove it later
 
                 logger.info(
-                    f'Streaming inference completed for user {user_id}, workflow {namespace}/{workflow_name}'
+                    f'Streaming inference completed for execution {execution_id}, workflow {namespace}/{workflow_name}'
                 )
 
-            except ValueError as e:
-                logger.error(f'Error in streaming inference: {e}')
-                error_event = {
-                    'event_type': 'error',
-                    'error': str(e),
-                    'timestamp': asyncio.get_event_loop().time(),
-                }
-                yield f'data: {json.dumps(error_event)}\n\n'
             except Exception as e:
-                logger.error(f'Error in streaming inference: {e}')
+                logger.error(
+                    f'Error in streaming inference for execution {execution_id}: {e}'
+                )
                 error_event = {
                     'event_type': 'error',
                     'error': str(e),
-                    'timestamp': asyncio.get_event_loop().time(),
+                    'timestamp': time.time(),
                 }
                 yield f'data: {json.dumps(error_event)}\n\n'
             finally:
-                # Clean up queue
-                event_streamer.cleanup_queue(user_id, namespace, workflow_name)
+                event_streamer.cleanup_queue(execution_id)
 
         return StreamingResponse(
             generate_inference_stream(),