Handle stream cancellation and sanitize tool call parsing

[christopherkarani]

Summary

• propagate chat stream cancellation to the provider so cancelGeneration() is invoked when the session is cancelled
• clamp or default unsafe PartialToolCall data and add tests for sanitization and stream cancellation behavior
• guard against OpenAI tool call indexes exceeding maxToolCallIndex and track skipped accumulators while parsing

Testing

• Not run (not requested)

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

[christopherkarani]

@claude Review

[claude]

PR Review

Overall this is a solid set of improvements — replacing crash-on-invalid-input with graceful sanitization, propagating cancellation correctly, and fixing the CRLF boundary edge case are all good defensive changes. A few issues worth addressing before merge:

---

Issues

1. Stale doc comment on PartialToolCall.index property (GenerationChunk.swift:92)

The init doc comment was correctly updated, but the property-level doc still says:

/// - Precondition: Must be in range `0...maxToolCallIndex` (0...100).

and line 84:

/// The index must be in the range `0...100`

These contradict the new clamping behavior and will mislead API consumers reading the generated docs. Both should be updated to reflect that values outside the range are clamped rather than rejected.

---

2. Duplicate if case .cancelled = termination pattern match (ChatSession.swift:745–752)

The closure matches on termination twice in sequence:

if case .cancelled = termination {
    task.cancel()
}
guard let strongSelf = self else { return }
if case .cancelled = termination {   // redundant second match
    Task { await strongSelf.provider.cancelGeneration() }
}

Both branches only care about the .cancelled case. Combining them into a single check reduces the noise and makes the intent clearer:

if case .cancelled = termination {
    task.cancel()
}
guard let strongSelf = self else { return }
strongSelf.withLock { strongSelf.generationTask = nil }
if case .cancelled = termination {
    Task { await strongSelf.provider.cancelGeneration() }
}

Or group the two cancellation side-effects together before the guard.

---

3. Fire-and-forget Task for cancelGeneration() creates a potential test race (ChatSession.swift:750–752 / ChatSessionTests.swift:822–825)

The cancellation propagation spawns an unstructured Task:

Task {
    await strongSelf.provider.cancelGeneration()
}

In the new test, after consumer.cancel() the assertion runs immediately once consumer.result completes:

consumer.cancel()
_ = await consumer.result
let cancelCount = await provider.cancelCallCount
#expect(cancelCount >= 1)

consumer.result completing doesn't guarantee the fire-and-forget Task has run yet — they're racing on Swift's cooperative scheduler. In practice it usually works because awaiting consumer.result yields to the runtime, but it's fragile. A small Task.yield() or Task.sleep after the await would make the test deterministic, or the design could track the cancellation task and await it explicitly.

---

Observations (no action required)

• CRLF fix (URLSessionAsyncBytes.swift): The peek-one-extra-byte approach is correct. When \r lands at the last position of the buffer, peeking the next byte before deciding whether it's a CRLF pair is the right strategy. The fall-through path when next() returns nil is also handled correctly by the existing lone-\r logic below.

• skippedToolAccumulatorIDs (OpenAIProvider+Streaming.swift): The nextToolIndex > maxToolCallIndex guard is correct and the skip-set approach prevents re-processing events for already-rejected IDs. The set will grow if a malicious server sends a very large number of distinct IDs, but for realistic workloads this is fine.

• maxToolCallIndex constant usage: Good cleanup — removing the hardcoded 100 and using the named constant makes the bound easier to adjust in future and avoids the earlier inconsistency between PartialToolCall.init and the provider code.

• Test additions: The new sanitization tests and the streamCancellationPropagatesToProvider test cover the right behaviors. The MockTextProvider changes are clean and correctly reset state in reset().

[christopherkarani]

@claude review and fix all identified issues

[claude]

Claude finished @christopherkarani's task in 2m 13s —— View job

---

Review Issues Fixed

All three issues from the previous review have been addressed in commit 9c32940.

• Read repo CLAUDE.md and relevant source files
• Fix stale doc comment on PartialToolCall.index (GenerationChunk.swift)
• Deduplicate if case .cancelled pattern match (ChatSession.swift)
• Fix fire-and-forget test race for cancelGeneration() (ChatSessionTests.swift)
• Commit and push fixes

---

Changes Made

1. Stale doc comment — GenerationChunk.swift lines 84 & 92

Replaced the - Precondition: tag (implying a crash on violation) and the "must be in range" wording with accurate descriptions of the clamping behavior:

-/// The index must be in the range `0...100` (see ``maxToolCallIndex``).
+/// The index is clamped to `0...maxToolCallIndex` (see ``maxToolCallIndex``).
 ...
-/// - Precondition: Must be in range `0...maxToolCallIndex` (0...100).
+/// - Note: Values outside `0...maxToolCallIndex` are clamped rather than rejected.

2. Duplicate if case .cancelled check — ChatSession.swift onTermination closure

Replaced the two separate if case .cancelled checks with a single guard case .cancelled that groups both cancellation side-effects together. The non-cancellation path now uses optional chaining to clear generationTask before returning early:

// Before: two separate if-case checks around a guard
if case .cancelled = termination { task.cancel() }
guard let strongSelf = self else { return }
if case .cancelled = termination { Task { await strongSelf.provider.cancelGeneration() } }
strongSelf.withLock { strongSelf.generationTask = nil }

// After: single guard-case, cancellation effects grouped
guard case .cancelled = termination else {
    self?.withLock { self?.generationTask = nil }
    return
}
task.cancel()
guard let strongSelf = self else { return }
strongSelf.withLock { strongSelf.generationTask = nil }
Task { await strongSelf.provider.cancelGeneration() }

3. Test race condition — ChatSessionTests.swift:streamCancellationPropagatesToProvider

Added await Task.yield() between awaiting consumer.result and reading cancelCallCount. This yields to the cooperative scheduler, giving the fire-and-forget Task { await strongSelf.provider.cancelGeneration() } a chance to execute before the assertion runs:

consumer.cancel()
_ = await consumer.result
// Yield to the cooperative scheduler so the fire-and-forget
// cancelGeneration() Task has a chance to run before we assert.
await Task.yield()
let cancelCount = await provider.cancelCallCount
#expect(cancelCount >= 1)

| Branch