Skip to content

Commit 0443dda

Browse files
authored
Merge pull request #352 from esokullu/main
Harden managed cloud bridge lifecycle
2 parents 1d8d93c + b548e41 commit 0443dda

3 files changed

Lines changed: 270 additions & 29 deletions

File tree

src/chrome/src/cloud-runs.js

Lines changed: 68 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@ const DEFAULT_CLOUD_BRIDGE_URL = 'ws://127.0.0.1:17373/extension';
22
const CLOUD_RUN_STORAGE_KEY = 'webbrainCloudRunSnapshots';
33
const CLOUD_UPDATE_LIMIT = 200;
44
const CLOUD_RUN_LIMIT = 50;
5+
const CLOUD_STRING_LIMIT = 16 * 1024;
6+
const CLOUD_RUN_PERSIST_BYTES_LIMIT = 256 * 1024;
7+
const CLOUD_PERSIST_BYTES_LIMIT = 4 * 1024 * 1024;
58
const TERMINAL_STATUSES = new Set(['completed', 'failed', 'aborted']);
69

710
export function normalizeCloudBridgeUrl(value = DEFAULT_CLOUD_BRIDGE_URL) {
@@ -22,22 +25,85 @@ function scrubCloudValue(value) {
2225
if ((key === '_attachImage' || key === 'screenshot') && typeof item === 'string' && item.length > 500) {
2326
return `[large payload omitted: ${item.length} chars]`;
2427
}
28+
if (typeof item === 'string' && item.length > CLOUD_STRING_LIMIT) {
29+
return `${item.slice(0, CLOUD_STRING_LIMIT)}\n[truncated ${item.length - CLOUD_STRING_LIMIT} chars for cloud persistence]`;
30+
}
2531
return item;
2632
}));
2733
} catch {
2834
return { unserializable: true };
2935
}
3036
}
3137

38+
function serializedBytes(value) {
39+
return new TextEncoder().encode(JSON.stringify(value)).length;
40+
}
41+
42+
function compactCloudRunForPersistence(run) {
43+
const row = scrubCloudValue(run);
44+
row.structured = row.structured ?? !!run?.outputSchema;
45+
if (serializedBytes(row) <= CLOUD_RUN_PERSIST_BYTES_LIMIT) return row;
46+
47+
const omittedUpdates = Array.isArray(row.updates) ? row.updates.length : 0;
48+
row.updates = [];
49+
row.persistenceTruncated = { omittedUpdates };
50+
if (serializedBytes(row) <= CLOUD_RUN_PERSIST_BYTES_LIMIT) return row;
51+
52+
row.content = '';
53+
row.outputSchema = null;
54+
row.persistenceTruncated.omittedContent = true;
55+
row.persistenceTruncated.omittedSchema = true;
56+
if (serializedBytes(row) <= CLOUD_RUN_PERSIST_BYTES_LIMIT) return row;
57+
58+
delete row.result;
59+
row.persistenceTruncated.omittedResult = true;
60+
if (serializedBytes(row) <= CLOUD_RUN_PERSIST_BYTES_LIMIT) return row;
61+
62+
return scrubCloudValue({
63+
runId: run?.runId,
64+
status: run?.status,
65+
tabId: run?.tabId,
66+
task: run?.task,
67+
structured: !!run?.outputSchema || run?.structured === true,
68+
summary: run?.summary,
69+
content: '',
70+
finalUrl: run?.finalUrl,
71+
error: run?.error,
72+
createdAt: run?.createdAt,
73+
updatedAt: run?.updatedAt,
74+
completedAt: run?.completedAt,
75+
updates: [],
76+
persistenceTruncated: { omittedUpdates, omittedResult: true, omittedSchema: true },
77+
});
78+
}
79+
80+
export function buildCloudPersistenceRows(runs) {
81+
const values = Array.isArray(runs) ? [...runs] : [...(runs?.values?.() || [])];
82+
const candidates = values
83+
.sort((a, b) => String(b?.createdAt || '').localeCompare(String(a?.createdAt || '')))
84+
.slice(0, CLOUD_RUN_LIMIT)
85+
.map(compactCloudRunForPersistence);
86+
const rows = [];
87+
let totalBytes = 2;
88+
for (const row of candidates) {
89+
const rowBytes = serializedBytes(row) + (rows.length ? 1 : 0);
90+
if (totalBytes + rowBytes > CLOUD_PERSIST_BYTES_LIMIT) continue;
91+
rows.push(row);
92+
totalBytes += rowBytes;
93+
}
94+
return rows;
95+
}
96+
3297
function cloudSnapshot(run, { includeUpdates = true } = {}) {
3398
if (!run) return null;
3499
return {
35100
runId: run.runId,
36101
status: run.status,
37102
tabId: run.tabId,
38103
task: run.task,
39-
structured: !!run.outputSchema,
104+
structured: run.structured ?? !!run.outputSchema,
40105
result: run.result,
106+
persistenceTruncated: run.persistenceTruncated,
41107
summary: run.summary,
42108
content: run.content,
43109
finalUrl: run.finalUrl,
@@ -79,10 +145,7 @@ export function createCloudRunController({
79145

80146
async function persist() {
81147
if (!api.storage?.session?.set) return;
82-
const rows = [...runs.values()]
83-
.sort((a, b) => String(b.createdAt).localeCompare(String(a.createdAt)))
84-
.slice(0, CLOUD_RUN_LIMIT)
85-
.map(run => scrubCloudValue(run));
148+
const rows = buildCloudPersistenceRows(runs);
86149
persistQueue = persistQueue
87150
.catch(() => {})
88151
.then(() => api.storage.session.set({ [CLOUD_RUN_STORAGE_KEY]: rows }));

src/chrome/src/offscreen/cloud-bridge.js

Lines changed: 33 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
*/
88

99
(() => {
10+
const ALLOWED_BRIDGE_ACTIONS = new Set(['cloud_run', 'cloud_status', 'cloud_abort']);
1011
let socket = null;
1112
let bridgeUrl = null;
1213
let enabled = false;
@@ -34,10 +35,10 @@
3435
};
3536
}
3637

37-
function sendJson(obj) {
38-
if (!socket || socket.readyState !== WebSocket.OPEN) return;
38+
function sendJson(obj, target = socket) {
39+
if (!target || target.readyState !== WebSocket.OPEN) return;
3940
try {
40-
socket.send(JSON.stringify(obj));
41+
target.send(JSON.stringify(obj));
4142
} catch (e) {
4243
lastError = e.message || String(e);
4344
}
@@ -56,26 +57,33 @@
5657
if (!enabled || !bridgeUrl) return;
5758
if (socket && (socket.readyState === WebSocket.OPEN || socket.readyState === WebSocket.CONNECTING)) return;
5859
try {
59-
socket = new WebSocket(bridgeUrl);
60-
socket.addEventListener('open', () => {
60+
const nextSocket = new WebSocket(bridgeUrl);
61+
socket = nextSocket;
62+
nextSocket.addEventListener('open', () => {
63+
if (socket !== nextSocket) return;
6164
reconnectAttempt = 0;
6265
lastError = '';
63-
sendJson({ type: 'hello', client: 'webbrain-extension', status: status() });
66+
sendJson({ type: 'hello', client: 'webbrain-extension', status: status() }, nextSocket);
6467
});
65-
socket.addEventListener('message', async (event) => {
68+
nextSocket.addEventListener('message', async (event) => {
69+
if (socket !== nextSocket) return;
6670
let msg;
6771
try {
6872
msg = JSON.parse(event.data);
6973
} catch (e) {
70-
sendJson({ ok: false, error: `Invalid JSON message: ${e.message}` });
74+
sendJson({ ok: false, error: `Invalid JSON message: ${e.message}` }, nextSocket);
7175
return;
7276
}
7377

7478
const id = msg.id || null;
7579
const action = msg.action || msg.command;
7680
const payload = msg.payload || msg;
7781
if (!action) {
78-
sendJson({ id, ok: false, error: 'Missing action' });
82+
sendJson({ id, ok: false, error: 'Missing action' }, nextSocket);
83+
return;
84+
}
85+
if (!ALLOWED_BRIDGE_ACTIONS.has(action)) {
86+
sendJson({ id, ok: false, error: `Unsupported cloud bridge action: ${action}` }, nextSocket);
7987
return;
8088
}
8189

@@ -85,20 +93,25 @@
8593
target: 'background',
8694
action,
8795
});
88-
if (response && response.error) {
89-
sendJson({ id, ok: false, error: response.error });
96+
const isRunSnapshot = !!response
97+
&& (response.runId != null || response.run_id != null)
98+
&& typeof response.status === 'string';
99+
if (response?.error && !isRunSnapshot) {
100+
sendJson({ id, ok: false, error: response.error }, nextSocket);
90101
} else {
91-
sendJson({ id, ok: true, result: response });
102+
sendJson({ id, ok: true, result: response }, nextSocket);
92103
}
93104
} catch (e) {
94-
sendJson({ id, ok: false, error: e.message || String(e) });
105+
sendJson({ id, ok: false, error: e.message || String(e) }, nextSocket);
95106
}
96107
});
97-
socket.addEventListener('close', () => {
108+
nextSocket.addEventListener('close', () => {
109+
if (socket !== nextSocket) return;
98110
socket = null;
99111
scheduleReconnect();
100112
});
101-
socket.addEventListener('error', () => {
113+
nextSocket.addEventListener('error', () => {
114+
if (socket !== nextSocket) return;
102115
lastError = 'WebSocket error';
103116
});
104117
} catch (e) {
@@ -122,8 +135,9 @@
122135
enabled = true;
123136
bridgeUrl = nextUrl;
124137
if (changed && socket) {
125-
try { socket.close(); } catch {}
138+
const previousSocket = socket;
126139
socket = null;
140+
try { previousSocket.close(); } catch {}
127141
}
128142
connect();
129143
sendResponse(status());
@@ -135,9 +149,10 @@
135149
reconnectTimer = null;
136150
reconnectAttempt = 0;
137151
if (socket) {
138-
try { socket.close(); } catch {}
152+
const previousSocket = socket;
153+
socket = null;
154+
try { previousSocket.close(); } catch {}
139155
}
140-
socket = null;
141156
sendResponse(status());
142157
return false;
143158
}

0 commit comments

Comments
 (0)