Skip to content

Make cancelled stream transcript retention explicit - #266

Merged
mattt merged 4 commits into
huggingface:mainfrom
qoli:codex/stream-cancellation-policy
Sep 27, 2026
Merged

mattt merged 4 commits into
huggingface:mainfrom
qoli:codex/stream-cancellation-policy

Conversation

@qoli

@qoli qoli commented Sep 24, 2026 •

Copy link
Copy Markdown
Contributor

Split independently from #264. This revision addresses the struct, concurrent rollback, and Observation requests in #266 (comment).

Cancelling a streaming consumer can make AsyncThrowingStream end 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.

TranscriptErrorHandlingPolicy is a struct with static preserveTranscript and revertTranscript values, matching the Foundation Models 27 API shape. The session's optional property uses access and withMutation so Observation clients receive policy changes.

  • nil retains 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.
  • .preserveTranscript commits only the latest cumulative Snapshot.transcriptEntries on streaming failure, preserving completed tool records without synthesizing a completed answer or duplicating repeated checkpoints.
  • .revertTranscript removes 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 test passed 599 tests in 60 suites with CI=1, live-provider credentials removed, and local-model opt-ins disabled. iOS Simulator and watchOS Simulator xcodebuild build-for-testing both passed. Changed Swift files pass swift format lint; git diff --check passes.

Empty-stream follow-up

Addresses #266 (comment). A stream that finishes without a snapshot now throws the existing noSnapshots error inside the relay, after the cancellation check. The normal policy/error cleanup runs before the error reaches either direct iteration or collect().

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 isResponding cleanup. It reproduced four failures before the fix and passes afterward. No new error type or public API was added.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 High severity · 1 Medium severity

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.

Comment thread Sources/AnyLanguageModel/LanguageModelSession.swift Outdated
Comment thread Sources/AnyLanguageModel/LanguageModelSession.swift
@mattt

mattt commented Sep 25, 2026

Copy link
Copy Markdown
Collaborator

Hi @qoli. Thanks for splitting this out so quickly, and it's great that it lines up with Foundation Models' transcriptErrorHandlingPolicy. A few things before we merge:

  • Foundation Models declares TranscriptErrorHandlingPolicy as a struct with static preserveTranscript and revertTranscript, not an enum. That leaves room for more policies without breaking anyone's exhaustive switch. Could you match that?
  • Reverting restores a saved copy of the whole transcript, or cuts it at this request's prompt. If two requests overlap on one session, a failure in one can remove the other's entries. Removing only the entries this request added would avoid that.
  • transcriptErrorHandlingPolicy should go through the same access/withMutation calls as transcript and usage, so observers see changes.

@qoli

qoli commented Sep 26, 2026

Copy link
Copy Markdown
Contributor Author

Thanks @mattt — addressed all three points in 8475e49.

  • TranscriptErrorHandlingPolicy is now a struct with static preserveTranscript and revertTranscript values.
  • Revert removes only the failing request's own prompt ID. The prompt is the only entry committed before success; checkpoints and the final response are committed atomically. It no longer restores a saved whole transcript or truncates entries belonging to overlapping requests.
  • The policy getter/setter now use access and withMutation, with an Observation regression test.

I also replaced the single relay handle with a registry: waitForResponseCompletion() waits for all streaming relays registered when the call begins, so a newer response cannot hide an older response's cleanup. Later-started relays are excluded, and nonstreaming operations must be awaited separately. Requests are not serialized or rejected.

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 (4d088af) and SwiftChat Alpha 14. The combined fork and downstream integrations were verified before publication; this PR still contains no reasoning APIs or provider changes.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🟡 Changes recommended

Empty streams bypass the rollback policy because their error is produced after relay cleanup.

Review effort: Balanced
Findings: 1 Medium severity

Open (1)
Resolved since last review (2)

Comment thread Sources/AnyLanguageModel/LanguageModelSession.swift Outdated
@mattt

mattt commented Sep 26, 2026

Copy link
Copy Markdown
Collaborator

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 collect() throws noSnapshots after that. So with .revertTranscript, the prompt stays in the transcript. Could the relay treat a missing snapshot as an error, so the policy applies? Once that's in, we can include this in the next release.

@qoli

qoli commented Sep 27, 2026

Copy link
Copy Markdown
Contributor Author

Thanks @mattt — fixed in 7338a62.

The relay now throws the existing ResponseStreamError.noSnapshots when the upstream finishes without yielding a snapshot, after the cancellation check. This goes through the relay's existing error-policy cleanup before reporting the error to either a direct iterator or collect(). With .revertTranscript, only the failed request's prompt is removed; nil and .preserveTranscript retain it. No completed response is fabricated.

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 isResponding cleanup have completed when the caller receives the error, while preserving prior history.

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.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🟢 Approval recommended

The implementation matches the documented semantics and provides strong regression coverage for the affected concurrency paths.

Review effort: Balanced
Findings: None

Resolved since last review (1)

@mattt
mattt merged commit 6da8838 into huggingface:main Sep 27, 2026
12 checks passed
@mattt

mattt commented Sep 27, 2026

Copy link
Copy Markdown
Collaborator

Thanks again, @qoli! This is out now in 0.14.1.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants