fix(langfuse): close partially initialized spans (#605)

Author: zxuexingzhijie

data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/langfuse/NodeTracingLifecycleListener.java

@@ -102,11 +102,12 @@ public void before(String nodeId, Map<String, Object> state, RunnableConfig conf
 			return;
 		}
 
+		Span span = null;
 		try {
 			int attempt = attemptCounters.computeIfAbsent(counterKey(threadId, nodeId), k -> new AtomicInteger())
 				.incrementAndGet();
 
-			Span span = tracer.spanBuilder(nodeId)
+			span = tracer.spanBuilder(nodeId)
 				.setSpanKind(SpanKind.INTERNAL)
 				.setParent(io.opentelemetry.context.Context.current().with(rootSpan))
 				.startSpan();
@@ -122,9 +123,19 @@ public void before(String nodeId, Map<String, Object> state, RunnableConfig conf
 			LangfuseService.registerActiveAccumulator(threadId);
 
 			activeNodes.computeIfAbsent(threadId, k -> new ArrayDeque<>()).push(new ActiveNode(nodeId, span, snapshot));
+			// Ownership is now held by activeNodes and transferred to after/onError.
+			span = null;
 		}
 		catch (Exception e) {
 			log.warn("Failed to start Langfuse span for node {}", nodeId, e);
+			if (span != null) {
+				try {
+					span.end();
+				}
+				catch (Exception endException) {
+					log.warn("Failed to clean up partially started Langfuse span for node {}", nodeId, endException);
+				}
+			}
 		}
 	}
 

data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/controller/GraphControllerTest.java

@@ -29,8 +29,8 @@
 import org.springframework.http.codec.ServerSentEvent;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.http.HttpHeaders;
-import reactor.core.publisher.Flux;
 import reactor.core.publisher.Sinks;
+import reactor.test.StepVerifier;
 
 import java.time.Duration;
 
@@ -64,11 +64,16 @@ private void stubResponseHeaders() {
 	@Test
 	void streamSearch_validRequest_invokesGraphServiceAndReturnsFlux() {
 		stubResponseHeaders();
+		doAnswer(invocation -> {
+			Sinks.Many<ServerSentEvent<GraphNodeResponse>> sink = invocation.getArgument(0);
+			sink.tryEmitComplete();
+			return null;
+		}).when(graphService).graphStreamProcess(any(Sinks.Many.class), any(GraphRequest.class));
 
-		Flux<ServerSentEvent<GraphNodeResponse>> result = graphController.streamSearch("agent-1", "conversation-1",
-				"thread-1", "show me sales data", false, null, false, false, serverHttpResponse);
-
-		assertNotNull(result);
+		StepVerifier
+			.create(graphController.streamSearch("agent-1", "conversation-1", "thread-1", "show me sales data", false,
+					null, false, false, serverHttpResponse))
+			.verifyComplete();
 
 		ArgumentCaptor<GraphRequest> requestCaptor = ArgumentCaptor.forClass(GraphRequest.class);
 		verify(graphService).graphStreamProcess(any(Sinks.Many.class), requestCaptor.capture());
@@ -79,16 +84,17 @@ void streamSearch_validRequest_invokesGraphServiceAndReturnsFlux() {
 		assertEquals("thread-1", captured.getThreadId());
 		assertEquals("show me sales data", captured.getQuery());
 		assertFalse(captured.isHumanFeedback());
+		verify(httpHeaders).add("Cache-Control", "no-cache");
+		verify(httpHeaders).add("Connection", "keep-alive");
+		verify(httpHeaders).add("Access-Control-Allow-Origin", "*");
 	}
 
 	@Test
 	void streamSearch_humanFeedback_passesHumanFeedbackParams() {
 		stubResponseHeaders();
 
-		Flux<ServerSentEvent<GraphNodeResponse>> result = graphController.streamSearch("agent-1", "conversation-2",
-				"thread-2", "approve this plan", true, "looks good", false, false, serverHttpResponse);
-
-		assertNotNull(result);
+		graphController.streamSearch("agent-1", "conversation-2", "thread-2", "approve this plan", true, "looks good",
+				false, false, serverHttpResponse);
 
 		ArgumentCaptor<GraphRequest> requestCaptor = ArgumentCaptor.forClass(GraphRequest.class);
 		verify(graphService).graphStreamProcess(any(Sinks.Many.class), requestCaptor.capture());
@@ -103,10 +109,8 @@ void streamSearch_humanFeedback_passesHumanFeedbackParams() {
 	void streamSearch_nl2sqlOnly_setsNl2sqlOnlyFlag() {
 		stubResponseHeaders();
 
-		Flux<ServerSentEvent<GraphNodeResponse>> result = graphController.streamSearch("agent-1", "conversation-3",
-				null, "SELECT query", false, null, false, true, serverHttpResponse);
-
-		assertNotNull(result);
+		graphController.streamSearch("agent-1", "conversation-3", null, "SELECT query", false, null, false, true,
+				serverHttpResponse);
 
 		ArgumentCaptor<GraphRequest> requestCaptor = ArgumentCaptor.forClass(GraphRequest.class);
 		verify(graphService).graphStreamProcess(any(Sinks.Many.class), requestCaptor.capture());

data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/service/chat/SessionTitleServiceTest.java

@@ -25,11 +25,12 @@
 import org.springframework.ai.chat.model.ChatResponse;
 import reactor.core.publisher.Flux;
 
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
 
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.*;
 
@@ -77,7 +78,7 @@ void scheduleTitleGeneration_sessionNotFound_doesNotCallRename() throws Exceptio
 
 		service.scheduleTitleGeneration("session-1", "hello");
 		executorService.shutdown();
-		executorService.awaitTermination(5, TimeUnit.SECONDS);
+		assertTrue(executorService.awaitTermination(5, TimeUnit.SECONDS), "the title task must actually finish");
 
 		verify(chatSessionService).findBySessionId("session-1");
 		verify(chatSessionService, never()).renameSession(anyString(), anyString());
@@ -90,7 +91,7 @@ void scheduleTitleGeneration_sessionHasCustomTitle_skips() throws Exception {
 
 		service.scheduleTitleGeneration("session-1", "hello");
 		executorService.shutdown();
-		executorService.awaitTermination(5, TimeUnit.SECONDS);
+		assertTrue(executorService.awaitTermination(5, TimeUnit.SECONDS), "the title task must actually finish");
 
 		verify(chatSessionService, never()).renameSession(anyString(), anyString());
 	}
@@ -106,40 +107,44 @@ void scheduleTitleGeneration_generatesAndPersistsTitle() throws Exception {
 
 		service.scheduleTitleGeneration("session-1", "hello world");
 		executorService.shutdown();
-		executorService.awaitTermination(5, TimeUnit.SECONDS);
+		assertTrue(executorService.awaitTermination(5, TimeUnit.SECONDS), "the title task must actually finish");
 
 		verify(chatSessionService).renameSession(eq("session-1"), eq("Generated Title"));
 		verify(sessionEventPublisher).publishTitleUpdated(eq(1), eq("session-1"), eq("Generated Title"));
 	}
 
 	@Test
 	void scheduleTitleGeneration_duplicateSessionId_skipsSecondCall() throws Exception {
-		AtomicReference<ChatSession> sessionRef = new AtomicReference<>(
-				ChatSession.builder().id("session-1").agentId(1).title("\u65b0\u4f1a\u8bdd").build());
-		when(chatSessionService.findBySessionId("session-1")).thenAnswer(invocation -> sessionRef.get());
-
-		doAnswer(invocation -> {
-			String newTitle = invocation.getArgument(1);
-			ChatSession session = sessionRef.get();
-			ChatSession newSession = ChatSession.builder()
-				.title(newTitle)
-				.id(session.getId())
-				.agentId(session.getAgentId())
-				.build();
-			sessionRef.set(newSession);
-			return null;
-		}).when(chatSessionService).renameSession(anyString(), anyString());
+		ChatSession session = ChatSession.builder().id("session-1").agentId(1).title("\u65b0\u4f1a\u8bdd").build();
+		CountDownLatch firstTaskStarted = new CountDownLatch(1);
+		CountDownLatch allowFirstTaskToFinish = new CountDownLatch(1);
+		when(chatSessionService.findBySessionId("session-1")).thenAnswer(invocation -> {
+			firstTaskStarted.countDown();
+			if (!allowFirstTaskToFinish.await(5, TimeUnit.SECONDS)) {
+				throw new IllegalStateException("timed out waiting to finish the first title task");
+			}
+			return session;
+		});
 
 		Flux<ChatResponse> chatResponseFlux = Flux.empty();
 		when(llmService.call(anyString(), anyString())).thenReturn(chatResponseFlux);
 		when(llmService.toStringFlux(chatResponseFlux)).thenReturn(Flux.just("Title"));
 
 		service.scheduleTitleGeneration("session-1", "hello");
-		service.scheduleTitleGeneration("session-1", "hello again");
+		try {
+			assertTrue(firstTaskStarted.await(5, TimeUnit.SECONDS), "the first title task must actually start");
+			service.scheduleTitleGeneration("session-1", "hello again");
+		}
+		finally {
+			allowFirstTaskToFinish.countDown();
+		}
 		executorService.shutdown();
-		executorService.awaitTermination(5, TimeUnit.SECONDS);
+		assertTrue(executorService.awaitTermination(5, TimeUnit.SECONDS), "the title task must actually finish");
 
-		verify(chatSessionService, atMostOnce()).renameSession(anyString(), anyString());
+		verify(chatSessionService, times(1)).findBySessionId("session-1");
+		verify(llmService, times(1)).call(anyString(), anyString());
+		verify(chatSessionService, times(1)).renameSession("session-1", "Title");
+		verify(sessionEventPublisher, times(1)).publishTitleUpdated(1, "session-1", "Title");
 	}
 
 }

data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/service/langfuse/NodeTracingLifecycleListenerTest.java

@@ -550,8 +550,16 @@ void exceptionWhileRecordingErrorIsSwallowedAndSpanEnded() {
 
 	@Test
 	void discardThread_isSafeForUnknownAndNullThreadId() {
+		Map<String, Object> state = new HashMap<>();
+		listener.before(SQL_GENERATE_NODE, state, config, 1L);
+
 		assertDoesNotThrow(() -> listener.discardThread("never-seen"));
 		assertDoesNotThrow(() -> listener.discardThread(null));
+
+		state.put(SQL_GENERATE_OUTPUT, "SELECT 1");
+		listener.after(SQL_GENERATE_NODE, state, config, 2L);
+		assertEquals(StatusCode.OK, onlyNodeSpan().getStatus().getStatusCode(),
+				"discarding an unknown thread must not disturb another thread's active span");
 	}
 
 	/**
@@ -630,10 +638,11 @@ void onErrorWithoutBefore_isIgnored() {
 	}
 
 	/**
-	 * span 自身在结束阶段抛异常时，catch 分支必须兜住，且仍要 end()。
+	 * span 创建后、登记到活动栈之前初始化失败时，catch 分支必须兜住并 end()，否则该 span 永远不会被
+	 * after/onError/discardThread 找到。
 	 */
 	@Test
-	void exceptionWhileFinishingSpanIsSwallowed() {
+	void exceptionWhileInitializingSpanIsSwallowedAndSpanEnded() {
 		Tracer flakyTracer = mock(Tracer.class);
 		io.opentelemetry.api.trace.SpanBuilder builder = mock(io.opentelemetry.api.trace.SpanBuilder.class);
 		Span flakySpan = mock(Span.class);
@@ -650,6 +659,8 @@ void exceptionWhileFinishingSpanIsSwallowed() {
 			flakyListener.before(SQL_GENERATE_NODE, Map.of(), config, 1L);
 			flakyListener.after(SQL_GENERATE_NODE, Map.of(SQL_GENERATE_OUTPUT, "x"), config, 2L);
 		});
+
+		org.mockito.Mockito.verify(flakySpan).end();
 	}
 
 	/**
@@ -717,6 +728,9 @@ void exceptionFromTracerIsSwallowed() {
 
 		assertDoesNotThrow(() -> brokenListener.before(INTENT_RECOGNITION_NODE, Map.of(INPUT_KEY, "q"), config, 1L));
 		assertDoesNotThrow(() -> brokenListener.after(INTENT_RECOGNITION_NODE, Map.of(), config, 2L));
+
+		org.mockito.Mockito.verify(brokenTracer).spanBuilder(INTENT_RECOGNITION_NODE);
+		assertTrue(nodeSpans().isEmpty(), "a failed tracer must not leave a phantom node span");
 	}
 
 	// --- helpers ---