Function wrapObjectStreamSeq

coderd/x/chatd/chatdebug/model.go:578–720

(
	ctx context.Context,
	handle *stepHandle,
	seq iter.Seq[fantasy.ObjectStreamPart],
)

}

func wrapObjectStreamSeq(
	ctx context.Context,
	handle *stepHandle,
	seq iter.Seq[fantasy.ObjectStreamPart],
) fantasy.ObjectStreamResponse {
	// Same safety-net pattern as wrapStreamSeq: a mutex rather
	// than sync.Once lets the AfterFunc yield to the normal
	// finalization path when the stream has already completed.
	var (
		mu             sync.Mutex
		finalized      bool
		streamComplete atomic.Bool
	)

	heartbeatDone := make(chan struct{})

	stop := context.AfterFunc(ctx, func() {
		mu.Lock()
		defer mu.Unlock()
		if finalized || streamComplete.Load() {
			return
		}
		finalized = true
		close(heartbeatDone)
		handle.finish(ctx, StatusInterrupted, nil, nil, nil, nil)
	})

	// Deferred heartbeat: start the heartbeat goroutine only when the
	// caller begins consuming the stream.
	startHeartbeat := sync.OnceFunc(func() {
		launchHeartbeat(ctx, handle.svc, handle.stepCtx.StepID, handle.stepCtx.RunID, handle.stepCtx.ChatID, heartbeatDone)
	})

	return func(yield func(fantasy.ObjectStreamPart) bool) {
		startHeartbeat()
		var (
			summary       = objectStreamSummary{StructuredOutput: true}
			latestUsage   fantasy.Usage
			usageSeen     bool
			finishSeen    bool
			finishReason  fantasy.FinishReason
			rawTextLength int
			warnings      []normalizedWarning
			streamError   any
			streamStatus  = StatusCompleted
		)

		finalize := func(status Status) {
			if stop != nil {
				stop()
			}
			mu.Lock()
			defer mu.Unlock()
			if finalized {
				return
			}
			finalized = true
			close(heartbeatDone)