fix: propagate full socket close to guest peers

Author: NathanFlurry

crates/native-sidecar/src/execution/mod.rs

@@ -292,6 +292,7 @@ use url::Url;
 
 const DEFAULT_KERNEL_STDIN_READ_MAX_BYTES: usize = 64 * 1024;
 const DEFAULT_KERNEL_STDIN_READ_TIMEOUT_MS: u64 = 100;
+const JAVASCRIPT_NET_CLOSE_SENTINEL: &str = "__agentos_net_close__";
 const JAVASCRIPT_NET_TIMEOUT_SENTINEL: &str = "__agentos_net_timeout__";
 const PYTHON_PYODIDE_GUEST_ROOT: &str = "/__agentos_pyodide";
 const PYTHON_PYODIDE_CACHE_GUEST_ROOT: &str = "/__agentos_pyodide_cache";

crates/native-sidecar/src/execution/network/tcp.rs

@@ -609,7 +609,14 @@ impl ActiveTcpSocket {
                                 .fetch_add(1, Ordering::Relaxed);
                         }
                         self.saw_remote_end.store(true, Ordering::SeqCst);
-                        Ok(Some(JavascriptTcpSocketEvent::End))
+                        if kernel
+                            .socket_get(socket_id)
+                            .is_some_and(|record| record.peer_socket_id().is_none())
+                        {
+                            Ok(Some(JavascriptTcpSocketEvent::Close { had_error: false }))
+                        } else {
+                            Ok(Some(JavascriptTcpSocketEvent::End))
+                        }
                     }
                     Err(error) if error.code() == "EAGAIN" => {
                         if trace_enabled {
@@ -634,7 +641,16 @@ impl ActiveTcpSocket {
             }
             if revents.intersects(POLLHUP) {
                 self.saw_remote_end.store(true, Ordering::SeqCst);
-                return Ok(Some(JavascriptTcpSocketEvent::End));
+                return Ok(Some(
+                    if kernel
+                        .socket_get(socket_id)
+                        .is_some_and(|record| record.peer_socket_id().is_none())
+                    {
+                        JavascriptTcpSocketEvent::Close { had_error: false }
+                    } else {
+                        JavascriptTcpSocketEvent::End
+                    },
+                ));
             }
             if revents.intersects(POLLERR) {
                 return Ok(Some(JavascriptTcpSocketEvent::Error {
@@ -2996,8 +3012,9 @@ pub(in crate::execution) fn javascript_net_read_value(
         Some(JavascriptTcpSocketEvent::Data { bytes, .. }) => Ok(Value::String(
             base64::engine::general_purpose::STANDARD.encode(bytes),
         )),
-        Some(JavascriptTcpSocketEvent::End | JavascriptTcpSocketEvent::Close { .. }) => {
-            Ok(Value::Null)
+        Some(JavascriptTcpSocketEvent::End) => Ok(Value::Null),
+        Some(JavascriptTcpSocketEvent::Close { .. }) => {
+            Ok(Value::String(String::from(JAVASCRIPT_NET_CLOSE_SENTINEL)))
         }
         Some(JavascriptTcpSocketEvent::Error { code, message }) => {
             let detail = code.unwrap_or_else(|| String::from("socket read"));

packages/build-tools/bridge-src/builtins/net.ts

@@ -1282,10 +1282,11 @@ function createAcceptedClientHandle(socketId, info) {
 	};
 }
 
-// Must match JAVASCRIPT_NET_TIMEOUT_SENTINEL in crates/native-sidecar/src/execution/mod.rs.
+// Must match the sentinels in crates/native-sidecar/src/execution/mod.rs.
 // A mismatched sentinel is NOT a soft failure: every no-data poll response then
 // falls through to base64 decoding and injects the decoded sentinel bytes into
 // the socket stream as phantom data.
+var NET_BRIDGE_CLOSE_SENTINEL = "__agentos_net_close__";
 var NET_BRIDGE_TIMEOUT_SENTINEL = "__agentos_net_timeout__";
 
 function isNetBridgeTraceEnabled() {
@@ -2897,6 +2898,13 @@ var NetSocket = class _NetSocket extends CanonicalDuplex {
 					countNetBridgeMetric("readWaitsForWake");
 					return;
 				}
+				if (chunk === NET_BRIDGE_CLOSE_SENTINEL) {
+					countNetBridgeMetric("readCloseEvents");
+					this._pendingBridgeWake = false;
+					this._pendingBridgeWakeRetries = 0;
+					this.destroy();
+					return;
+				}
 				if (chunk === null) {
 					if (firstPumpRun && !firstPumpResultRecorded) {
 						firstPumpResultRecorded = true;
@@ -3830,6 +3838,7 @@ export {
 	isValidIPv6Zone,
 	isValidTcpPort,
 	maxNetBridgeMetric,
+	NET_BRIDGE_CLOSE_SENTINEL,
 	NET_BRIDGE_MAX_RAW_WRITE_BYTES,
 	NET_BRIDGE_TIMEOUT_SENTINEL,
 	NET_SERVER_HANDLE_PREFIX,

packages/core/tests/network-http-request.test.ts

@@ -12,7 +12,7 @@ async function runSpawnedProcess(
 ): Promise<{ exitCode: number; stdout: string; stderr: string }> {
 	const stdoutChunks: string[] = [];
 	const stderrChunks: string[] = [];
-	const { pid } = vm.spawn(command, args, {
+	const { pid } = await vm.spawn(command, args, {
 		onStdout: (chunk) => {
 			stdoutChunks.push(textDecoder.decode(chunk));
 		},
@@ -22,8 +22,9 @@ async function runSpawnedProcess(
 		},
 	});
 
+	const exit = await vm.process.wait(pid);
 	return {
-		exitCode: await vm.waitProcess(pid),
+		exitCode: exit.exitCode ?? -1,
 		stdout: stdoutChunks.join(""),
 		stderr: stderrChunks.join(""),
 	};
@@ -248,7 +249,7 @@ describe("guest http.request transport", () => {
 		});
 	});
 
-	test("streams a guest response after the handler opens an outbound websocket", async () => {
+	test("reclaims cancelled guest response streams while an outbound websocket stays open", async () => {
 		const upstream = createServer();
 		const upstreamWebSocket = new WebSocketServer({ server: upstream });
 		upstreamWebSocket.on("connection", (socket) => {
@@ -282,20 +283,21 @@ describe("guest http.request transport", () => {
 			});
 			const script = [
 				'const http = require("node:http");',
-				"const server = http.createServer(async (_request, response) => {",
+				"void (async () => {",
 				`  const socket = new WebSocket("ws://127.0.0.1:${upstreamAddress.port}/", ["rivet", "rivet_token.token"]);`,
 				'  socket.binaryType = "arraybuffer";',
 				"  const binaryLength = await new Promise((resolve, reject) => {",
 				"    socket.onopen = () => queueMicrotask(() => socket.send(new Uint8Array([4, 5, 6])));",
 				"    socket.onmessage = (event) => resolve(event.data.byteLength);",
 				"    socket.onerror = reject;",
 				"  });",
-				"  socket.close();",
-				'  response.writeHead(200, { "Content-Type": "text/event-stream" });',
-				"  response.flushHeaders();",
-				"  response.write(`data: websocket-${binaryLength}\\n\\n`);",
-				"});",
-				'server.listen(3000, "0.0.0.0", () => console.log("READY"));',
+				"  const server = http.createServer((_request, response) => {",
+				'    response.writeHead(200, { "Content-Type": "text/event-stream" });',
+				"    response.flushHeaders();",
+				"    response.write(`data: websocket-${binaryLength}\\n\\n`);",
+				"  });",
+				'  server.listen(3000, "0.0.0.0", () => console.log("READY"));',
+				"})().catch((error) => { console.error(error); process.exitCode = 1; });",
 			].join("\n");
 			const child = await vm.spawn("node", ["-e", script], {
 				onStdout: (chunk) => {
@@ -320,15 +322,24 @@ describe("guest http.request transport", () => {
 				),
 			]);
 
-			for (let requestIndex = 0; requestIndex < 10; requestIndex++) {
+			const requestCount = Number.parseInt(
+				process.env.AGENTOS_VM_FETCH_STREAM_REQUESTS ?? "300",
+				10,
+			);
+			for (let requestIndex = 0; requestIndex < requestCount; requestIndex++) {
 				const head = await Promise.race([
 					vm.fetchStreamStart(
 						3000,
 						new Request(`http://guest/events-${requestIndex}`),
 					),
 					new Promise<never>((_, reject) =>
 						setTimeout(
-							() => reject(new Error("stream response head timed out")),
+							() =>
+								reject(
+									new Error(
+										`stream response head timed out at request ${requestIndex + 1}/${requestCount}`,
+									),
+								),
 							5_000,
 						),
 					),