Finalize streaming state and handle terminal events

Add a finalized flag to streamingState (markFinalized/isFinalized) and use it to guard finalization paths across the streaming lifecycle. Ensure streams are closed in runAgentLoopStreamStep and short-circuit when state is nil in finishStreamingWithFailure. Update responsesTurnAdapter/FinaleAgentLoop and chatCompletionsTurnAdapter to use isFinalized checks, make processResponseStreamEvent explicitly handle response.failed and response.incomplete (finalize metadata, log, and return loop-stop), and return on response.completed. Add early-finalize guard in completeStreamingSuccess. Include new tests to cover completed, failed, and finalize behavior. These changes prevent duplicate finalization/races and ensure terminal events stop the loop and close resources correctly.

Author: batuhan

bridges/ai/agent_loop_runtime.go

@@ -90,6 +90,9 @@ func runAgentLoopStreamStep[T any](
 	handleEvent func(T) (done bool, cle *ContextLengthError, err error),
 	handleErr func(error) (cle *ContextLengthError, err error),
 ) (bool, *ContextLengthError, error) {
+	if stream != nil {
+		defer stream.Close()
+	}
 	writer := state.writer()
 	writer.StepStart(ctx)
 	defer writer.StepFinish(ctx)

bridges/ai/streaming_chat_completions.go

@@ -203,7 +203,7 @@ func (a *chatCompletionsTurnAdapter) FinalizeAgentLoop(ctx context.Context) {
 	state := a.state
 	portal := a.portal
 	meta := a.meta
-	if state == nil || state.completedAtMs != 0 {
+	if state == nil || state.isFinalized() {
 		return
 	}
 

bridges/ai/streaming_error_handling.go

@@ -40,6 +40,12 @@ func (oc *AIClient) finishStreamingWithFailure(
 	reason string,
 	err error,
 ) error {
+	if state == nil {
+		return err
+	}
+	if !state.markFinalized() {
+		return streamFailureError(state, err)
+	}
 	if state != nil && state.stop.Load() != nil && reason == "cancelled" {
 		reason = "stop"
 	}

bridges/ai/streaming_lifecycle_cluster_test.go

@@ -162,3 +162,104 @@ func TestProcessResponseStreamEventUpdatesCompletedResponseStatus(t *testing.T)
 		t.Fatalf("expected writer metadata to be completed, got %#v", metadata["response_status"])
 	}
 }
+
+func TestProcessResponseStreamEventCompletedSignalsLoopStop(t *testing.T) {
+	state := newTestStreamingStateWithTurn()
+	oc := &AIClient{}
+
+	rsc := &responseStreamContext{
+		base: &agentLoopProviderBase{
+			oc:    oc,
+			log:   zerolog.Nop(),
+			state: state,
+		},
+	}
+
+	done, cle, err := oc.processResponseStreamEvent(context.Background(), rsc, responses.ResponseStreamEventUnion{
+		Type: "response.completed",
+		Response: responses.Response{
+			ID:     "resp_done",
+			Status: "completed",
+		},
+	}, false)
+	if !done {
+		t.Fatal("expected completed response event to stop the stream loop")
+	}
+	if cle != nil {
+		t.Fatalf("did not expect context-length error, got %#v", cle)
+	}
+	if err != nil {
+		t.Fatalf("did not expect error, got %v", err)
+	}
+}
+
+func TestResponsesTurnAdapterFinalizeAgentLoopDoesNotSkipTerminalLifecycle(t *testing.T) {
+	state := newTestStreamingStateWithTurn()
+	state.turn.SetSuppressSend(true)
+	state.writer().TextDelta(context.Background(), "done")
+	state.completedAtMs = 123
+	state.finishReason = "stop"
+
+	adapter := &responsesTurnAdapter{
+		agentLoopProviderBase: agentLoopProviderBase{
+			oc:    &AIClient{},
+			log:   zerolog.Nop(),
+			state: state,
+		},
+	}
+
+	adapter.FinalizeAgentLoop(context.Background())
+
+	if !state.isFinalized() {
+		t.Fatal("expected finalize agent loop to finalize terminal response state")
+	}
+
+	message := streamui.SnapshotUIMessage(state.turn.UIState())
+	metadata, _ := message["metadata"].(map[string]any)
+	if metadata["finish_reason"] != "stop" {
+		t.Fatalf("expected finalized UI message finish_reason stop, got %#v", metadata["finish_reason"])
+	}
+}
+
+func TestProcessResponseStreamEventFailedFinalizesAsError(t *testing.T) {
+	state := newTestStreamingStateWithTurn()
+	state.turn.SetSuppressSend(true)
+	state.writer().TextDelta(context.Background(), "hello")
+	oc := &AIClient{}
+
+	rsc := &responseStreamContext{
+		base: &agentLoopProviderBase{
+			oc:    oc,
+			log:   zerolog.Nop(),
+			state: state,
+		},
+	}
+
+	done, cle, err := oc.processResponseStreamEvent(context.Background(), rsc, responses.ResponseStreamEventUnion{
+		Type: "response.failed",
+		Response: responses.Response{
+			ID:     "resp_failed",
+			Status: "failed",
+			Error: responses.ResponseError{
+				Message: "boom",
+			},
+		},
+	}, false)
+	if !done {
+		t.Fatal("expected failed response event to stop the stream loop")
+	}
+	if cle != nil {
+		t.Fatalf("did not expect context-length error, got %#v", cle)
+	}
+	if err == nil {
+		t.Fatal("expected failed response event to return an error")
+	}
+	if !state.isFinalized() {
+		t.Fatal("expected failed response event to finalize the turn")
+	}
+	message := streamui.SnapshotUIMessage(state.turn.UIState())
+	metadata, _ := message["metadata"].(map[string]any)
+	if metadata["finish_reason"] != "error" {
+		t.Fatalf("expected error finish_reason, got %#v", metadata["finish_reason"])
+	}
+}

bridges/ai/streaming_responses_api.go

@@ -176,7 +176,7 @@ func (a *responsesTurnAdapter) RunAgentTurn(
 }
 
 func (a *responsesTurnAdapter) FinalizeAgentLoop(ctx context.Context) {
-	if a.state == nil || a.state.completedAtMs != 0 {
+	if a.state == nil || a.state.isFinalized() {
 		return
 	}
 	a.oc.finalizeResponsesStream(ctx, a.log, a.portal, a.state, a.meta)
@@ -217,9 +217,29 @@ func (oc *AIClient) processResponseStreamEvent(
 	)
 
 	switch streamEvent.Type {
-	case "response.created", "response.queued", "response.in_progress", "response.failed", "response.incomplete":
+	case "response.created", "response.queued", "response.in_progress":
 		oc.handleResponseLifecycleEvent(ctx, portal, state, meta, streamEvent.Type, streamEvent.Response)
 
+	case "response.failed":
+		oc.handleResponseLifecycleEvent(ctx, portal, state, meta, streamEvent.Type, streamEvent.Response)
+		state.completedAtMs = time.Now().UnixMilli()
+		errText := strings.TrimSpace(streamEvent.Response.Error.Message)
+		if errText == "" {
+			errText = "response failed"
+		}
+		return true, nil, oc.finishStreamingWithFailure(ctx, log, portal, state, meta, "error", errors.New(errText))
+
+	case "response.incomplete":
+		oc.handleResponseLifecycleEvent(ctx, portal, state, meta, streamEvent.Type, streamEvent.Response)
+		state.completedAtMs = time.Now().UnixMilli()
+		actions.finalizeMetadata()
+		log.Debug().
+			Str("reason", state.finishReason).
+			Str("response_id", state.responseID).
+			Str("response_status", state.responseStatus).
+			Msg("Response stream ended incomplete" + contSuffix)
+		return true, nil, nil
+
 	case "response.output_item.added":
 		actions.outputItemAdded(streamEvent.Item)
 
@@ -377,6 +397,7 @@ func (oc *AIClient) processResponseStreamEvent(
 		}
 		log.Debug().Str("reason", state.finishReason).Str("response_id", state.responseID).Int("images", len(state.pendingImages)).
 			Msg("Response stream completed" + contSuffix)
+		return true, nil, nil
 
 	case "error":
 		apiErr := fmt.Errorf("API error: %s", streamEvent.Message)

bridges/ai/streaming_state.go

@@ -71,7 +71,8 @@ type streamingState struct {
 	pendingMcpApprovals     []mcpApprovalRequest
 	pendingMcpApprovalsSeen map[string]bool
 
-	stop atomic.Pointer[assistantStopMetadata]
+	finalized atomic.Bool
+	stop      atomic.Pointer[assistantStopMetadata]
 }
 
 // sourceEventID returns the triggering user message event ID from the turn's source ref.
@@ -109,6 +110,20 @@ func (s *streamingState) writer() *sdk.Writer {
 	return s.turn.Writer()
 }
 
+func (s *streamingState) markFinalized() bool {
+	if s == nil {
+		return false
+	}
+	return s.finalized.CompareAndSwap(false, true)
+}
+
+func (s *streamingState) isFinalized() bool {
+	if s == nil {
+		return false
+	}
+	return s.finalized.Load()
+}
+
 func (s *streamingState) nextMessageTiming() agentremote.EventTiming {
 	if s == nil {
 		return agentremote.ResolveEventTiming(time.Time{}, 0)

bridges/ai/streaming_success.go

@@ -17,6 +17,9 @@ func (oc *AIClient) completeStreamingSuccess(
 	state *streamingState,
 	meta *PortalMetadata,
 ) {
+	if state == nil || !state.markFinalized() {
+		return
+	}
 	state.completedAtMs = time.Now().UnixMilli()
 	if state.finishReason == "" {
 		state.finishReason = "stop"