From 9e4074474bc9bb7150d827bfa24a90aae7748992 Mon Sep 17 00:00:00 2001
From: john <cynell@139.com>
Date: Fri, 3 Jul 2026 16:34:05 +0800
Subject: [PATCH] fix: harden session event proxy errors

---
 tkmind-proxy.mjs      |  21 +++++++--
 tkmind-proxy.test.mjs | 103 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 121 insertions(+), 3 deletions(-)

diff --git a/tkmind-proxy.mjs b/tkmind-proxy.mjs
index 1fc6798..07caf80 100644
--- a/tkmind-proxy.mjs
+++ b/tkmind-proxy.mjs
@@ -390,6 +390,20 @@ function sendProxyResponse(res, upstream) {
   Readable.fromWeb(upstream.body).pipe(res);
 }
 
+function writeSseProxyErrorAndEnd(res, message) {
+  if (res.writableEnded) return;
+  if (res.headersSent) {
+    try {
+      res.write(`event: error\ndata: ${JSON.stringify({ message })}\n\n`);
+    } catch {
+      // The client may already be gone. Ending below is still harmless.
+    }
+    res.end();
+    return;
+  }
+  res.status(502).json({ message });
+}
+
 function extractSessionId(req, body) {
   const fromParams = req.params?.sessionId ?? req.params?.id;
   if (fromParams) return fromParams;
@@ -1756,12 +1770,13 @@ export function createTkmindProxy({
         if (!res.writableEnded) res.end();
         return;
       }
-      res.status(502).json({
-        message: sanitizeUserFacingProxyMessage(
+      writeSseProxyErrorAndEnd(
+        res,
+        sanitizeUserFacingProxyMessage(
           err instanceof Error ? err.message : 'SSE 代理失败',
           'SSE 代理失败',
         ),
-      });
+      );
     }
   };
 
diff --git a/tkmind-proxy.test.mjs b/tkmind-proxy.test.mjs
index c7ce91a..f0e6e16 100644
--- a/tkmind-proxy.test.mjs
+++ b/tkmind-proxy.test.mjs
@@ -1,4 +1,5 @@
 import assert from 'node:assert/strict';
+import { EventEmitter } from 'node:events';
 import fs from 'node:fs';
 import { createServer } from 'node:http';
 import os from 'node:os';
@@ -108,9 +109,39 @@ function createMemoryTestUserAuth(workingDir) {
     async getSessionTarget() {
       return { target: null, node: 0 };
     },
+    async billSessionUsage() {
+      return { ok: true, costCents: 0, balanceCents: 0 };
+    },
   };
 }
 
+function attachExpressResponseHelpers(res) {
+  res.status = (code) => {
+    res.statusCode = code;
+    return res;
+  };
+  res.json = (payload) => {
+    res.setHeader('Content-Type', 'application/json; charset=utf-8');
+    res.end(JSON.stringify(payload));
+    return res;
+  };
+  res.send = (payload = '') => {
+    res.end(payload);
+    return res;
+  };
+  return res;
+}
+
+async function listen(server) {
+  await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve));
+  return server.address().port;
+}
+
+async function closeServer(server) {
+  if (!server?.listening) return;
+  await new Promise((resolve) => server.close(resolve));
+}
+
 test('sanitizePublicHtmlLinksInText downgrades missing own markdown public html links', () => {
   const owner = `test-user-${Date.now()}-markdown-missing`;
   const publicRoot = path.join(process.cwd(), 'MindSpace', owner, 'public');
@@ -251,6 +282,78 @@ test('buildVisionPayload injects public standard image urls for page generation'
   assert.match(text, /不得改写图片里人物的年龄、性别、人数或主体关系/);
 });
 
+test('proxySessionEvents does not send JSON after SSE headers were sent', async () => {
+  let upstream;
+
+  try {
+    upstream = createServer((req, res) => {
+      if (req.method === 'GET' && req.url === '/sessions/session-1/events') {
+        res.writeHead(200, { 'Content-Type': 'text/event-stream' });
+        res.write('data: {"type":"message","content":"hello"}\n\n');
+        res.end('data: {"type":"message","content":"second"}\n\n');
+        return;
+      }
+      res.writeHead(404, { 'Content-Type': 'application/json' });
+      res.end(JSON.stringify({ message: `unexpected ${req.method} ${req.url}` }));
+    });
+    const upstreamPort = await listen(upstream);
+    const proxy = createTkmindProxy({
+      apiTarget: `http://127.0.0.1:${upstreamPort}`,
+      apiSecret: 'test-secret',
+      userAuth: createMemoryTestUserAuth(process.cwd()),
+    });
+
+    const req = new EventEmitter();
+    req.currentUser = { id: 'user-1', username: 'john' };
+    req.get = () => '';
+    req.once = req.once.bind(req);
+    req.off = req.off.bind(req);
+
+    let writes = 0;
+    let jsonCalled = false;
+    const chunks = [];
+    const res = new EventEmitter();
+    res.headersSent = false;
+    res.writableEnded = false;
+    res.statusCode = 200;
+    res.status = (code) => {
+      res.statusCode = code;
+      return res;
+    };
+    res.setHeader = () => {};
+    res.flushHeaders = () => {
+      res.headersSent = true;
+    };
+    res.write = (chunk) => {
+      res.headersSent = true;
+      writes += 1;
+      if (writes > 1) throw new Error('simulated client write failure');
+      chunks.push(String(chunk));
+      return true;
+    };
+    res.end = () => {
+      res.writableEnded = true;
+    };
+    res.json = () => {
+      jsonCalled = true;
+      throw new Error('json must not be called after SSE started');
+    };
+    res.send = () => {
+      throw new Error('send must not be called after SSE started');
+    };
+
+    await proxy.proxySessionEvents(req, res, 'session-1');
+
+    assert.equal(res.statusCode, 200);
+    assert.equal(res.headersSent, true);
+    assert.equal(res.writableEnded, true);
+    assert.equal(jsonCalled, false);
+    assert.match(chunks.join(''), /"content":"hello"/);
+  } finally {
+    await closeServer(upstream);
+  }
+});
+
 test('startSessionForUser resolves memories through Memory V2 facade', async () => {
   let resolveInput = null;
   await withFakeGoosedSession(async ({ apiTarget, workingDir, harnessEntries }) => {