Make cancelled stream transcript retention explicit - #266
Conversation
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Concurrent responses can erase unrelated transcript entries or cause the completion wait to observe the wrong relay.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 1
Open (2)
What changed in this PR
Makes transcript retention explicit when generation fails or streaming is cancelled.
Changes:
- Adds configurable preserve/revert transcript policies.
- Prevents cancelled partial responses from being committed.
- Adds relay-completion synchronization and deterministic tests.
| File | Description |
|---|---|
LanguageModelSession.swift |
Implements retention policies, cancellation checks, and relay waiting. |
TranscriptErrorHandlingTests.swift |
Tests cancellation, failures, checkpoints, rollback, and success. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
Hi @qoli. Thanks for splitting this out so quickly, and it's great that it lines up with Foundation Models'
|
|
Thanks @mattt — addressed all three points in 8475e49.
I also replaced the single relay handle with a registry: Added gated concurrency tests covering stream/nonstream overlap with both start and finish orders, multimodal rollback, older-relay cancellation after a newer response completes, cleanup waiting, and Observation notifications. All 598 offline tests passed, as did iOS and watchOS Simulator test-target builds and formatting checks. The nil/default policy and the nonstreaming failure limitations remain explicit in the PR description; I am not claiming parity with Apple's default behavior. This refactor is now adopted alongside #264 in my maintained fork ( |
|
Hi @qoli. Thanks for turning these around so quickly! Everything looks good now except one case: when the upstream stream finishes without yielding a snapshot, the relay finishes normally, and |
|
Thanks @mattt — fixed in 7338a62. The relay now throws the existing Added a regression test covering all three policy settings with both consumption styles (six cases). It reproduced the missing iterator error and retained prompt before the fix, and now verifies that policy cleanup and Validation: all 599 offline tests passed; iOS Simulator and watchOS Simulator test-target builds, strict formatting, and whitespace checks passed. No new public API or error type was added. |


Split independently from #264. This revision addresses the struct, concurrent rollback, and Observation requests in #266 (comment).
Cancelling a streaming consumer can make
AsyncThrowingStreamend normally inside the session relay. Without a cancellation check, the session commits its last partial answer as a completed assistant response. This change checks cancellation before that commit and makes failed-request transcript retention explicit.TranscriptErrorHandlingPolicyis a struct with staticpreserveTranscriptandrevertTranscriptvalues, matching the Foundation Models 27 API shape. The session's optional property usesaccessandwithMutationso Observation clients receive policy changes.nilretains the prompt without committing streaming checkpoint entries on failure. This preserves existing thrown-error behavior while fixing cancelled-stream partial-answer commits; it does not claim parity with Apple's default policy..preserveTranscriptcommits only the latest cumulativeSnapshot.transcriptEntrieson streaming failure, preserving completed tool records without synthesizing a completed answer or duplicating repeated checkpoints..revertTranscriptremoves only the failing request's unique prompt entry. The prompt is the only entry committed before success; checkpoints and the final response commit atomically at completion. Rollback does not restore an old transcript or truncate later entries, so overlapping requests retain their own prompts, tools, and responses. External tool side effects are not undone.Nonstreaming failures have no checkpoint channel: preserve retains the prompt, and revert removes that request's prompt. A nonstreaming provider that ignores cancellation and returns successfully still follows the existing success path; this PR does not add a new nonstreaming cancellation contract.
waitForResponseCompletion()is an AnyLanguageModel extension that snapshots and waits for all streaming relays registered when the call begins. Later-started relays are excluded; nonstreaming operations must be awaited separately. Cancelling the wait does not cancel generation. Applications can call it after cancelling consumers before persisting the transcript. It must not be called from a tool running inside one of the included responses. Registration and completed-relay removal share one lock, preventing fast completion from leaving stale task handles. Requests remain free to overlap.Tests cover existing cancellation/error policies, a gated eight-case stream/nonstream overlap matrix with both start and finish orders, multimodal rollback, cancellation of an older relay after a newer request completes, waiting for an older active relay, and Observation notifications. No reasoning APIs or provider changes are included.
Validation against current upstream
main(0626c0f):swift testpassed 599 tests in 60 suites withCI=1, live-provider credentials removed, and local-model opt-ins disabled. iOS Simulator and watchOS Simulatorxcodebuild build-for-testingboth passed. Changed Swift files passswift format lint;git diff --checkpasses.Empty-stream follow-up
Addresses #266 (comment). A stream that finishes without a snapshot now throws the existing
noSnapshotserror inside the relay, after the cancellation check. The normal policy/error cleanup runs before the error reaches either direct iteration orcollect().The regression test covers nil, preserve, and revert with both consumption styles (six cases), checking retained prior history, the request prompt, absence of a fabricated answer, and
isRespondingcleanup. It reproduced four failures before the fix and passes afterward. No new error type or public API was added.