Close WebSocket cleanly before domain reload

[liuzqk]

Description

Replace immediate WebSocket abort during Unity domain reload with a bounded close-output attempt that is serialized against normal sends and falls back to abort without blocking the Editor.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation update
• Refactoring (no functional changes)
• Test update

Changes Made

• Send EndpointUnavailable close output for Open and CloseReceived sockets with a 100 ms budget.
• Acquire the normal WebSocket send lock non-blockingly before the close output; abort immediately if a send is already active.
• Abort on timeout, exception, or unfinished states such as Connecting and CloseSent.
• Leave Closed and Aborted sockets untouched.

Compatibility / Package Source

• Unity version(s) tested: 2022.3.62f3
• Package source used (#beta, #main, tag, branch, or file:): branch derived directly from official beta (bd72241ab91c664176091789c9408d676783abec); local file: source for validation only
• Resolved commit hash from Packages/packages-lock.json (if using a Git package URL): N/A

The 100 ms close budget is bounded; all failure paths retain the previous abort behavior.

Testing/Screenshots/Recordings

• Python tests (cd Server && uv run pytest tests/ -v)
• Unity EditMode tests
• Unity PlayMode tests
• Package import/compile check
• Not applicable (explain why in Additional Notes)

Unity 2022.3.62f3 targeted reload suite: 21 passed, 0 failed. Tests cover close output, CloseReceived, timeout, exception, unfinished states, and terminal states; the send-lock serialization path is compiled and source-reviewed.

Documentation Updates

• I have added/removed/modified tools or resources
• If yes, I have updated all documentation files using:
  • The LLM prompt at tools/UPDATE_DOCS_PROMPT.md (recommended)
  • Manual review of the generated changes

Related Issues

Relates to #1207.

Additional Notes

A rebuilt stable branch containing this PR plus #1285 and #1286 passed 162 targeted EditMode tests: 161 passed, 0 failed, 1 intentional stress-only skip. git diff --check passes.

[coderabbitai]

📝 Walkthrough

Walkthrough

ForceStop() now performs a bounded graceful WebSocket close during domain reload, aborting when necessary before disposing resources. Tests cover open, close-received, connecting, terminal, failure, and timeout states.

Changes

WebSocket domain-reload shutdown

Layer / File(s)	Summary
Bounded shutdown helper and ForceStop() wiring MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs	ForceStop() calls CloseForDomainReload with a 100ms timeout. Open and close-received sockets send an EndpointUnavailable close frame; failures and unsupported states abort.
Shutdown behavior tests TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/WebSocketTransportClientTests.cs	Tests and a RecordingWebSocket stub verify close output, timeout cancellation, exceptions, close-received behavior, connecting-state aborts, and terminal-state no-op behavior.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ForceStop
  participant CloseForDomainReload
  participant WebSocket
  ForceStop->>CloseForDomainReload: request bounded close
  CloseForDomainReload->>WebSocket: send EndpointUnavailable output
  CloseForDomainReload->>WebSocket: abort on failure or timeout
  ForceStop->>WebSocket: dispose socket resources

Loading

Possibly related PRs

• CoplayDev/unity-mcp#688: Updates WebSocketTransportClient reload teardown behavior involving ForceStop().

Suggested labels: bug

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title is concise and accurately summarizes the main change: a cleaner WebSocket shutdown before domain reload.
Description check	✅ Passed	The description follows the template and includes the required sections, change summary, compatibility, testing, related issue, and notes.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/WebSocketTransportClientTests.cs (1)

88-144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add terminal-state regressions.

Add Closed and Aborted cases asserting that neither CloseOutputAsync nor Abort is called. Those branches are part of the stated shutdown contract but are not covered here.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/WebSocketTransportClientTests.cs`
around lines 88 - 144, Add terminal-state regression tests alongside the
existing CloseForDomainReload tests using RecordingWebSocket configured with
WebSocketState.Closed and WebSocketState.Aborted. Assert that both
CloseOutputCalled and AbortCalled remain false in each case, preserving the
existing ConnectingSocket test pattern.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In
`@MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs`:
- Around line 181-183: Update CloseForDomainReload to attempt acquiring
_sendLock without blocking before calling CloseOutputAsync; perform the
domain-reload close only when the lock is acquired, release it afterward, and
abort immediately when a normal SendJsonAsync already holds or is queued for the
lock.

---

Nitpick comments:
In
`@TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/WebSocketTransportClientTests.cs`:
- Around line 88-144: Add terminal-state regression tests alongside the existing
CloseForDomainReload tests using RecordingWebSocket configured with
WebSocketState.Closed and WebSocketState.Aborted. Assert that both
CloseOutputCalled and AbortCalled remain false in each case, preserving the
existing ConnectingSocket test pattern.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 725ad9b5-aabf-4694-b456-5404fb86fea1

📥 Commits

Reviewing files that changed from the base of the PR and between bd72241 and 5d4e473.

📒 Files selected for processing (2)

• MCPForUnity/Editor/Services/Transport/Transports/WebSocketTransportClient.cs
• TestProjects/UnityMCPTests/Assets/Tests/EditMode/Services/WebSocketTransportClientTests.cs

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5d4e4731dd

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".