@@ -4,7 +4,6 @@ import { type IPty } from "bun-pty"
44import z from "zod"
55import { Identifier } from "../id/id"
66import { Log } from "../util/log"
7- import type { WSContext } from "hono/ws"
87import { Instance } from "../project/instance"
98import { lazy } from "@opencode-ai/util/lazy"
109import { Shell } from "@/shell/shell"
@@ -17,6 +16,22 @@ export namespace Pty {
1716 const BUFFER_CHUNK = 64 * 1024
1817 const encoder = new TextEncoder ( )
1918
19+ type Socket = {
20+ readyState : number
21+ send : ( data : string | Uint8Array < ArrayBuffer > | ArrayBuffer ) => void
22+ close : ( code ?: number , reason ?: string ) => void
23+ }
24+
25+ const sockets = new WeakMap < object , number > ( )
26+ let socketCounter = 0
27+
28+ const tagSocket = ( ws : Socket ) => {
29+ if ( ! ws || typeof ws !== "object" ) return
30+ const next = ( socketCounter = ( socketCounter + 1 ) % Number . MAX_SAFE_INTEGER )
31+ sockets . set ( ws , next )
32+ return next
33+ }
34+
2035 // WebSocket control frame: 0x00 + UTF-8 JSON (currently { cursor }).
2136 const meta = ( cursor : number ) => {
2237 const json = JSON . stringify ( { cursor } )
@@ -81,7 +96,7 @@ export namespace Pty {
8196 buffer : string
8297 bufferCursor : number
8398 cursor : number
84- subscribers : Set < WSContext >
99+ subscribers : Map < Socket , number >
85100 }
86101
87102 const state = Instance . state (
@@ -91,8 +106,12 @@ export namespace Pty {
91106 try {
92107 session . process . kill ( )
93108 } catch { }
94- for ( const ws of session . subscribers ) {
95- ws . close ( )
109+ for ( const ws of session . subscribers . keys ( ) ) {
110+ try {
111+ ws . close ( )
112+ } catch {
113+ // ignore
114+ }
96115 }
97116 }
98117 sessions . clear ( )
@@ -154,18 +173,26 @@ export namespace Pty {
154173 buffer : "" ,
155174 bufferCursor : 0 ,
156175 cursor : 0 ,
157- subscribers : new Set ( ) ,
176+ subscribers : new Map ( ) ,
158177 }
159178 state ( ) . set ( id , session )
160179 ptyProcess . onData ( ( data ) => {
161180 session . cursor += data . length
162181
163- for ( const ws of session . subscribers ) {
182+ for ( const [ ws , id ] of session . subscribers ) {
164183 if ( ws . readyState !== 1 ) {
165184 session . subscribers . delete ( ws )
166185 continue
167186 }
168- ws . send ( data )
187+ if ( typeof ws === "object" && sockets . get ( ws ) !== id ) {
188+ session . subscribers . delete ( ws )
189+ continue
190+ }
191+ try {
192+ ws . send ( data )
193+ } catch {
194+ session . subscribers . delete ( ws )
195+ }
169196 }
170197
171198 session . buffer += data
@@ -177,14 +204,15 @@ export namespace Pty {
177204 ptyProcess . onExit ( ( { exitCode } ) => {
178205 log . info ( "session exited" , { id, exitCode } )
179206 session . info . status = "exited"
180- for ( const ws of session . subscribers ) {
181- ws . close ( )
207+ for ( const ws of session . subscribers . keys ( ) ) {
208+ try {
209+ ws . close ( )
210+ } catch {
211+ // ignore
212+ }
182213 }
183214 session . subscribers . clear ( )
184215 Bus . publish ( Event . Exited , { id, exitCode } )
185- for ( const ws of session . subscribers ) {
186- ws . close ( )
187- }
188216 state ( ) . delete ( id )
189217 } )
190218 Bus . publish ( Event . Created , { info } )
@@ -211,9 +239,14 @@ export namespace Pty {
211239 try {
212240 session . process . kill ( )
213241 } catch { }
214- for ( const ws of session . subscribers ) {
215- ws . close ( )
242+ for ( const ws of session . subscribers . keys ( ) ) {
243+ try {
244+ ws . close ( )
245+ } catch {
246+ // ignore
247+ }
216248 }
249+ session . subscribers . clear ( )
217250 state ( ) . delete ( id )
218251 Bus . publish ( Event . Deleted , { id } )
219252 }
@@ -232,7 +265,7 @@ export namespace Pty {
232265 }
233266 }
234267
235- export function connect ( id : string , ws : WSContext , cursor ?: number ) {
268+ export function connect ( id : string , ws : Socket , cursor ?: number ) {
236269 const session = state ( ) . get ( id )
237270 if ( ! session ) {
238271 ws . close ( )
@@ -272,7 +305,8 @@ export namespace Pty {
272305 return
273306 }
274307
275- session . subscribers . add ( ws )
308+ const socketId = tagSocket ( ws )
309+ if ( typeof socketId === "number" ) session . subscribers . set ( ws , socketId )
276310 return {
277311 onMessage : ( message : string | ArrayBuffer ) => {
278312 session . process . write ( String ( message ) )
0 commit comments