fix: expose stable Langfuse trace IO

Author: erichare

src/backend/base/langflow/services/tracing/langfuse.py

@@ -28,6 +28,59 @@
 LANGFUSE_FEEDBACK_SCORE_NAME = "user-feedback"
 
 
+def _normalize_boundary_messages(value: Any) -> Any:
+    """Replace Langflow messages with their text while preserving container shape."""
+    from lfx.schema.message import Message
+
+    if isinstance(value, Message):
+        return value.get_text()
+    if isinstance(value, dict):
+        return {key: _normalize_boundary_messages(item) for key, item in value.items()}
+    if isinstance(value, list | tuple):
+        return [_normalize_boundary_messages(item) for item in value]
+    return value
+
+
+def _serialize_component_boundary(component_output: Any) -> Any:
+    """Collapse a sole component output and normalize Langflow messages to text."""
+    value = (
+        next(iter(component_output.values()))
+        if isinstance(component_output, dict) and len(component_output) == 1
+        else component_output
+    )
+    return serialize(_normalize_boundary_messages(value))
+
+
+def _trace_boundary_value(
+    component_values: dict[str, Any],
+    boundary_traces: dict[str, str],
+    *,
+    fallback_component_values: dict[str, Any] | None = None,
+    prefer_fallback_trace_ids: set[str] | None = None,
+) -> tuple[bool, Any]:
+    """Return marked graph-boundary outputs in deterministic component-id order."""
+    values = []
+    prefer_fallback_trace_ids = prefer_fallback_trace_ids or set()
+    for trace_id, trace_name in sorted(boundary_traces.items()):
+        sources = (
+            (fallback_component_values, component_values)
+            if trace_id in prefer_fallback_trace_ids
+            else (component_values, fallback_component_values)
+        )
+        for source in sources:
+            if source is None or trace_name not in source:
+                continue
+            component_value = source[trace_name]
+            if isinstance(component_value, dict) and not component_value:
+                continue
+            values.append(_serialize_component_boundary(component_value))
+            break
+
+    if not values:
+        return False, None
+    return True, values[0] if len(boundary_traces) == 1 else values
+
+
 class _SharedClient:
     """Process-wide cached Langfuse client.
 
@@ -299,6 +352,8 @@ def __init__(
         self.session_id = session_id
         self.flow_id = trace_name.split(" - ")[-1]
         self.spans: dict[str, LangfuseSpan] = OrderedDict()
+        self._input_trace_names: dict[str, str] = {}
+        self._output_trace_names: dict[str, str] = {}
         self.langfuse_trace_id = None
 
         config = self._get_config()
@@ -397,6 +452,12 @@ def add_trace(
 
         name = trace_name.removesuffix(f" ({trace_id})")
 
+        if vertex is not None:
+            if vertex.is_input:
+                self._input_trace_names[trace_id] = trace_name
+            if vertex.is_output:
+                self._output_trace_names[trace_id] = trace_name
+
         # Create child span under the root span
         span = self._root_span.start_span(
             name=name,
@@ -440,11 +501,29 @@ def end(
         if not self._ready:
             return
 
-        # Serialize once and reuse to avoid duplicate work
+        # Keep the complete component aggregates on the root observation.
         inputs_ser = serialize(inputs)
         outputs_ser = serialize(outputs)
         metadata_ser = serialize(metadata) if metadata else None
 
+        # Input components emit the normalized external request as their output;
+        # output components emit the final graph result. If a custom graph has no
+        # boundary marker, retain the full aggregate rather than guessing from
+        # concurrent component completion order.
+        dual_role_trace_ids = self._input_trace_names.keys() & self._output_trace_names.keys()
+        input_found, trace_input = _trace_boundary_value(
+            outputs,
+            self._input_trace_names,
+            fallback_component_values=inputs,
+            prefer_fallback_trace_ids=dual_role_trace_ids,
+        )
+        if not input_found:
+            trace_input = inputs_ser
+
+        output_found, trace_output = _trace_boundary_value(outputs, self._output_trace_names)
+        if not output_found:
+            trace_output = outputs_ser
+
         # Update the root span with final input/output
         self._root_span.update(
             input=inputs_ser,
@@ -454,8 +533,8 @@ def end(
 
         # Update trace-level data
         self._root_span.update_trace(
-            input=inputs_ser,
-            output=outputs_ser,
+            input={"input": trace_input},
+            output={"output": trace_output},
             metadata=metadata_ser,
         )
 

src/backend/tests/unit/services/tracing/test_langfuse_v3_compatibility.py

@@ -317,9 +317,11 @@ def test_end_trace_updates_and_ends_span(self, mock_langfuse):
         mock_langfuse["child_span"].update.assert_called()
         mock_langfuse["child_span"].end.assert_called()
 
-    def test_end_updates_root_span_and_trace(self, mock_langfuse):
-        """Test that end() updates both root span and trace, then ends."""
+    def test_end_exposes_stable_evaluator_trace_input_and_output(self, mock_langfuse):
+        """Trace input/output should expose stable JSONPath keys with the flow boundary messages."""
+        from langflow.serialization.serialization import serialize
         from langflow.services.tracing.langfuse import LangFuseTracer
+        from lfx.schema.message import Message
 
         tracer = LangFuseTracer(
             trace_name="test-flow - flow-123",
@@ -328,18 +330,197 @@ def test_end_updates_root_span_and_trace(self, mock_langfuse):
             trace_id=uuid.uuid4(),
         )
 
+        component_inputs = {
+            "Config (config-id)": {"model": "test-model"},
+            "Chat Input (chat-input-id)": {"input_value": "What is Langflow?", "sender": "User"},
+            "Prompt (prompt-id)": {"template": "Answer the user"},
+        }
+        component_outputs = {
+            "Chat Input (chat-input-id)": {"message": Message(text="What is Langflow?")},
+            "Agent (agent-id)": {"response": Message(text="An intermediate response")},
+            "Chat Output (chat-output-id)": {"message": Message(text="Langflow is a visual workflow builder.")},
+            "Audit Sink (audit-id)": {"record": {"status": "stored"}},
+        }
+        tracer.add_trace(
+            trace_id="chat-input-id",
+            trace_name="Chat Input (chat-input-id)",
+            trace_type="chain",
+            inputs=component_inputs["Chat Input (chat-input-id)"],
+            vertex=MagicMock(is_input=True, is_output=False),
+        )
+        tracer.add_trace(
+            trace_id="chat-output-id",
+            trace_name="Chat Output (chat-output-id)",
+            trace_type="chain",
+            inputs={"input_value": "Langflow is a visual workflow builder."},
+            vertex=MagicMock(is_input=False, is_output=True),
+        )
         tracer.end(
-            inputs={"flow_input": "test"},
-            outputs={"flow_output": "result"},
+            inputs=component_inputs,
+            outputs=component_outputs,
             metadata={"final": True},
         )
 
-        # Should update root span
-        mock_langfuse["root_span"].update.assert_called()
-        # Should update trace metadata
-        assert mock_langfuse["root_span"].update_trace.call_count >= 2  # init + end
-        # Should end root span
-        mock_langfuse["root_span"].end.assert_called()
+        root_update = mock_langfuse["root_span"].update.call_args.kwargs
+        assert root_update["input"] == serialize(component_inputs)
+        assert root_update["output"] == serialize(component_outputs)
+
+        trace_update = mock_langfuse["root_span"].update_trace.call_args.kwargs
+        assert trace_update["input"] == {"input": "What is Langflow?"}
+        assert trace_update["output"] == {"output": "Langflow is a visual workflow builder."}
+        assert trace_update["metadata"] == {"final": True}
+        mock_langfuse["root_span"].end.assert_called_once()
+
+    def test_end_exposes_stable_trace_io_for_marked_arbitrary_graphs(self, mock_langfuse):
+        """Marked non-chat boundaries should expose their structured input and output values."""
+        from langflow.services.tracing.langfuse import LangFuseTracer
+
+        tracer = LangFuseTracer(
+            trace_name="test-flow - flow-123",
+            trace_type="chain",
+            project_name="test-project",
+            trace_id=uuid.uuid4(),
+        )
+
+        component_inputs = {
+            "Webhook (webhook-id)": {"payload": {"order_id": 42}, "method": "POST"},
+            "Transform (transform-id)": {"mapping": "order"},
+        }
+        component_outputs = {
+            "Webhook (webhook-id)": {},
+            "Transform (transform-id)": {"record": {"order_id": 42, "status": "ready"}},
+            "Data Output (data-output-id)": {"data": [{"order_id": 42, "status": "ready"}]},
+        }
+        tracer.add_trace(
+            trace_id="webhook-id",
+            trace_name="Webhook (webhook-id)",
+            trace_type="chain",
+            inputs=component_inputs["Webhook (webhook-id)"],
+            vertex=MagicMock(is_input=True, is_output=False),
+        )
+        tracer.add_trace(
+            trace_id="data-output-id",
+            trace_name="Data Output (data-output-id)",
+            trace_type="chain",
+            inputs={},
+            vertex=MagicMock(is_input=False, is_output=True),
+        )
+        tracer.end(inputs=component_inputs, outputs=component_outputs)
+
+        trace_update = mock_langfuse["root_span"].update_trace.call_args.kwargs
+        assert trace_update["input"] == {"input": component_inputs["Webhook (webhook-id)"]}
+        assert trace_update["output"] == {"output": [{"order_id": 42, "status": "ready"}]}
+
+    def test_end_falls_back_to_stable_aggregate_when_boundaries_are_unmarked(self, mock_langfuse):
+        """Custom graphs without boundary markers should retain their complete structured aggregates."""
+        from langflow.services.tracing.langfuse import LangFuseTracer
+
+        tracer = LangFuseTracer(
+            trace_name="test-flow - flow-123",
+            trace_type="chain",
+            project_name="test-project",
+            trace_id=uuid.uuid4(),
+        )
+        component_inputs = {"Custom Source (source-id)": {"payload": {"order_id": 42}}}
+        component_outputs = {"Custom Sink (sink-id)": {"record": {"order_id": 42, "status": "ready"}}}
+
+        tracer.end(inputs=component_inputs, outputs=component_outputs)
+
+        trace_update = mock_langfuse["root_span"].update_trace.call_args.kwargs
+        assert trace_update["input"] == {"input": component_inputs}
+        assert trace_update["output"] == {"output": component_outputs}
+
+    def test_end_orders_multiple_boundary_outputs_by_component_id(self, mock_langfuse):
+        """Multiple graph outputs should be deterministic even when their completion order is not."""
+        from langflow.services.tracing.langfuse import LangFuseTracer
+        from lfx.schema.message import Message
+
+        tracer = LangFuseTracer(
+            trace_name="test-flow - flow-123",
+            trace_type="chain",
+            project_name="test-project",
+            trace_id=uuid.uuid4(),
+        )
+        component_outputs = {
+            "Data Output (z-output-id)": {"data": {"order_id": 42}},
+            "Text Output (a-output-id)": {"text": Message(text="ready")},
+        }
+        tracer.add_trace(
+            trace_id="z-output-id",
+            trace_name="Data Output (z-output-id)",
+            trace_type="chain",
+            inputs={},
+            vertex=MagicMock(is_input=False, is_output=True),
+        )
+        tracer.add_trace(
+            trace_id="a-output-id",
+            trace_name="Text Output (a-output-id)",
+            trace_type="chain",
+            inputs={},
+            vertex=MagicMock(is_input=False, is_output=True),
+        )
+
+        tracer.end(inputs={}, outputs=component_outputs)
+
+        trace_update = mock_langfuse["root_span"].update_trace.call_args.kwargs
+        assert trace_update["output"] == {"output": ["ready", {"order_id": 42}]}
+
+    def test_end_uses_component_input_for_dual_role_boundary(self, mock_langfuse):
+        """A component marked as both input and output should preserve the original request."""
+        from langflow.services.tracing.langfuse import LangFuseTracer
+        from lfx.schema.message import Message
+
+        tracer = LangFuseTracer(
+            trace_name="test-flow - flow-123",
+            trace_type="chain",
+            project_name="test-project",
+            trace_id=uuid.uuid4(),
+        )
+        trace_name = "Bidirectional (component-id)"
+        tracer.add_trace(
+            trace_id="component-id",
+            trace_name=trace_name,
+            trace_type="chain",
+            inputs={"request": "ping"},
+            vertex=MagicMock(is_input=True, is_output=True),
+        )
+
+        tracer.end(
+            inputs={trace_name: {"request": "ping"}},
+            outputs={trace_name: {"response": Message(text="pong")}},
+        )
+
+        trace_update = mock_langfuse["root_span"].update_trace.call_args.kwargs
+        assert trace_update["input"] == {"input": "ping"}
+        assert trace_update["output"] == {"output": "pong"}
+
+    def test_end_normalizes_messages_nested_in_multi_output_boundary(self, mock_langfuse):
+        """Message values should become evaluator-ready text without dropping sibling outputs."""
+        from langflow.services.tracing.langfuse import LangFuseTracer
+        from lfx.schema.message import Message
+
+        tracer = LangFuseTracer(
+            trace_name="test-flow - flow-123",
+            trace_type="chain",
+            project_name="test-project",
+            trace_id=uuid.uuid4(),
+        )
+        trace_name = "Composite Output (output-id)"
+        tracer.add_trace(
+            trace_id="output-id",
+            trace_name=trace_name,
+            trace_type="chain",
+            inputs={},
+            vertex=MagicMock(is_input=False, is_output=True),
+        )
+
+        tracer.end(
+            inputs={},
+            outputs={trace_name: {"message": Message(text="ready"), "usage": {"tokens": 3}}},
+        )
+
+        trace_update = mock_langfuse["root_span"].update_trace.call_args.kwargs
+        assert trace_update["output"] == {"output": {"message": "ready", "usage": {"tokens": 3}}}
 
     def test_get_langchain_callback_uses_trace_context(self, mock_langfuse):
         """Test that get_langchain_callback creates handler with trace context."""