diff --git a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java
index 0980dc92803..9a13668d6ec 100644
--- a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java
+++ b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java
@@ -297,14 +297,10 @@ public Flux<ChatResponse> stream(Prompt prompt) {
 					return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
 				}
 				return Flux.just(response);
-			}).doOnError(observation::error).doFinally(s -> {
-				// TODO: Consider a custom ObservationContext and
-				// include additional metadata
-				// if (s == SignalType.CANCEL) {
-				// observationContext.setAborted(true);
-				// }
-				observation.stop();
-			}).contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
+			})
+			.doOnError(observation::error)
+			.doFinally(s -> observation.stop())
+			.contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
 			// @formatter:on
 
 			return new MessageAggregator().aggregate(flux, observationContext::setResponse);