| 402 | } |
| 403 | |
| 404 | func wrapStreamSeq( |
| 405 | ctx context.Context, |
| 406 | handle *stepHandle, |
| 407 | seq iter.Seq[fantasy.StreamPart], |
| 408 | ) fantasy.StreamResponse { |
| 409 | // mu and finalized guard both the normal finalization path |
| 410 | // inside the iterator and the safety-net AfterFunc below. |
| 411 | // This ensures handle.finish is called exactly once regardless |
| 412 | // of whether the caller iterates, drops the stream, or the |
| 413 | // context is canceled mid-flight. We use a mutex rather than |
| 414 | // sync.Once so the AfterFunc can yield to the normal path |
| 415 | // when the stream already received its terminal chunk |
| 416 | // (streamComplete), preventing the AfterFunc from clobbering |
| 417 | // completed stream data with nil. |
| 418 | var ( |
| 419 | mu sync.Mutex |
| 420 | finalized bool |
| 421 | streamComplete atomic.Bool |
| 422 | ) |
| 423 | |
| 424 | // heartbeatDone is closed when the stream finalizes (either |
| 425 | // normally or via the safety net) to stop the heartbeat goroutine. |
| 426 | heartbeatDone := make(chan struct{}) |
| 427 | |
| 428 | // Safety net: if the caller drops the returned iterator without |
| 429 | // consuming it (or abandons mid-stream and the context is |
| 430 | // canceled), finalize the step so it does not remain permanently |
| 431 | // in_progress once persistence lands in later branches. |
| 432 | stop := context.AfterFunc(ctx, func() { |
| 433 | mu.Lock() |
| 434 | defer mu.Unlock() |
| 435 | // If the stream already received a finish chunk, let |
| 436 | // finalize handle it; it has the real response payload |
| 437 | // and usage data that we would otherwise clobber. |
| 438 | if finalized || streamComplete.Load() { |
| 439 | return |
| 440 | } |
| 441 | finalized = true |
| 442 | close(heartbeatDone) |
| 443 | handle.finish(ctx, StatusInterrupted, nil, nil, nil, nil) |
| 444 | }) |
| 445 | |
| 446 | // startHeartbeat launches the heartbeat goroutine on first call. |
| 447 | // Deferring the start until the caller begins consuming the stream |
| 448 | // prevents leaked goroutines when the iterator is dropped without |
| 449 | // being iterated. |
| 450 | startHeartbeat := sync.OnceFunc(func() { |
| 451 | launchHeartbeat(ctx, handle.svc, handle.stepCtx.StepID, handle.stepCtx.RunID, handle.stepCtx.ChatID, heartbeatDone) |
| 452 | }) |
| 453 | |
| 454 | return func(yield func(fantasy.StreamPart) bool) { |
| 455 | startHeartbeat() |
| 456 | var ( |
| 457 | summary streamSummary |
| 458 | latestUsage fantasy.Usage |
| 459 | usageSeen bool |
| 460 | finishSeen bool |
| 461 | finishReason fantasy.FinishReason |