fix(openai): lease WebSocket connections exclusively - #4346
Merged
Merged
Conversation
aheritier
added
area/providers
This was referenced Sep 17, 2026
dgageot
enabled auto-merge
September 18, 2026 12:03
gtardif
previously approved these changes
Sep 18, 2026
github-merge-queue
Bot
removed this pull request from the merge queue due to a conflict with the base branch
Sep 18, 2026
Adds ExclusiveStreamLease, a program-level cop that flags any function returning a WebSocket stream that aliases a receiver field — a pattern that lets the caller and pool share the same connection concurrently. Red-first evidence (go run ./lint . before production fix): ws_pool.go:103 Lint/ExclusiveStreamLease: returned WebSocket stream aliases receiver field conn ws_pool.go:127 Lint/ExclusiveStreamLease: returned WebSocket stream aliases receiver field conn 2 offense(s) detected Signed-off-by: David Gageot <david.gageot@docker.com> Assisted-By: claude-opus-4-5
Pool is idle-only: the idle slot is cleared the moment a stream takes ownership, so overlapping requests each dial their own connection. Active connections are tracked separately for shutdown only. Key points: - terminal-event completion sets a flag; Close is idempotent via sync.Once - early Close discards the lease and closes the connection immediately - active map lets Close() shut down in-flight connections cleanly - cancelled requests (ctx.Err at entry) don't consume the idle slot Signed-off-by: David Gageot <david@gageot.net> Assisted-By: Claude
Assisted-By: Claude
Use atomic.Bool for done, sync.Once for Close, and WriteControl (concurrency-safe, bounded timeout) for the close handshake so concurrent Next/Close calls cannot race or double-close the socket. Assisted-By: Claude Signed-off-by: David Gageot <david.gageot@docker.com>
Protect Close read vs retry pointer write with mutex; close flag ensures replacement returned after closure is discarded. Both adapters delegate to synchronized close. Tests cover standard/beta Close before/after retry and successful retry; race detector confirmed issues before and clean after. Assisted-By: claude-sonnet-4-5 Signed-off-by: David Gageot <david.gageot@docker.com>
fix: detect and repair unsafe stream closure
aheritier
added
the
status/needs-rebase
Collaborator
|
👋 This PR has merge conflicts with the base branch. Please rebase or merge the latest base branch and resolve them. I've moved it to draft and added |
dgageot
force-pushed
the
fix/websocket-stream-lease
branch
from
September 18, 2026 12:31
22a7dac to
c713e64
Compare
trungutt
approved these changes
Sep 18, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The OpenAI WebSocket pool allowed a single idle connection to be shared across concurrent streaming requests. Because each WebSocket carries exactly one in-flight exchange at a time, two concurrent callers could interleave reads on the same connection, silently corrupting both responses.
A new
Lint/ExclusiveStreamLeaseanalyzer was written first and run against the codebase to confirm the problem. It uses resolved-type, intra-procedural CFG and origin tracking via the existingprog.Copframework to flag WebSocket wrapper values that alias a retained receiver field and are returned to callers — covering branch paths and named-return cases. The analyzer flagged exactly two violations atws_pool.go:103andws_pool.go:127. It is intra-procedural only and does not prove whole-program concurrency safety; helper side-effects, receiver aliases in callers, container tracking, and general lease protocols are out of scope.The fix gives the idle slot take-and-clear semantics: consuming it transfers exclusive connection ownership to the caller for the lifetime of its streaming request. A second concurrent request dials its own connection rather than blocking or sharing. Active connections are tracked so pool shutdown closes both idle and in-flight connections.
Closeis idempotent; connections that have not finished are closed to unblock the reader; a request already cancelled at entry leaves the idle slot unchanged.task buildandtask lintpass. The targeted race tests pass with-race -count=3.task testfails only onTestLoadExamplesentries that require Docker Model Runner to be running locally; this failure reproduces identically on the pre-fix commit and is unrelated to this change.The rest of the suite passes with
go test ./... -skip '^TestLoadExamples$'. Pool-widelastResponseIDsemantics and holding the pool mutex during dialing predate this PR and remain out of scope.