Skip to content

Commit b62b42f

Browse files
jacobjmcCopilot
andauthored
Reduce memory pressure in heavy threads (#9)
* Reduce thread memory pressure * Guard base64 data URL image decoding against unbounded memory allocation (#10) * Initial plan * Add pre/post-decode size guards for data URL base64 images Co-authored-by: jacobjmc <111402762+jacobjmc@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: jacobjmc <111402762+jacobjmc@users.noreply.github.com> --------- Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
1 parent dd7f377 commit b62b42f

File tree

7 files changed

+524
-56
lines changed

7 files changed

+524
-56
lines changed

src-tauri/src/shared/codex_core.rs

Lines changed: 79 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
1+
use base64::Engine;
12
use serde_json::{json, Value};
3+
use std::collections::hash_map::DefaultHasher;
24
use std::collections::{HashMap, HashSet};
5+
use std::hash::{Hash, Hasher};
36
use std::net::IpAddr;
47
use std::path::PathBuf;
58
use std::sync::Arc;
@@ -576,7 +579,10 @@ pub(crate) async fn resume_thread_core<E: EventSink>(
576579
}
577580
"file" => {
578581
if let Some(url) = part.get("url").and_then(|v| v.as_str()) {
579-
content_parts.push(json!({ "type": "image", "value": url }));
582+
content_parts.push(json!({
583+
"type": "image",
584+
"value": frontend_image_value(url)
585+
}));
580586
}
581587
}
582588
_ => {}
@@ -1102,6 +1108,62 @@ pub(crate) async fn set_thread_name_core(
11021108
const URL_IMAGE_FETCH_TIMEOUT: Duration = Duration::from_secs(10);
11031109
const URL_IMAGE_MAX_BYTES: usize = 8 * 1024 * 1024;
11041110

1111+
fn extension_from_mime(mime: &str) -> &'static str {
1112+
match mime.trim().to_ascii_lowercase().as_str() {
1113+
"image/png" => "png",
1114+
"image/jpeg" => "jpg",
1115+
"image/gif" => "gif",
1116+
"image/webp" => "webp",
1117+
"image/svg+xml" => "svg",
1118+
"image/bmp" => "bmp",
1119+
"image/tiff" => "tiff",
1120+
_ => "bin",
1121+
}
1122+
}
1123+
1124+
fn persist_data_image_to_temp_file(data_url: &str) -> Option<String> {
1125+
let trimmed = data_url.trim();
1126+
let (metadata, encoded) = trimmed
1127+
.strip_prefix("data:")?
1128+
.split_once(";base64,")?;
1129+
if !metadata.starts_with("image/") {
1130+
return None;
1131+
}
1132+
let estimated_len = encoded.len().saturating_mul(3) / 4;
1133+
if estimated_len > URL_IMAGE_MAX_BYTES {
1134+
return None;
1135+
}
1136+
let bytes = base64::engine::general_purpose::STANDARD
1137+
.decode(encoded)
1138+
.ok()?;
1139+
if bytes.len() > URL_IMAGE_MAX_BYTES {
1140+
return None;
1141+
}
1142+
1143+
let mut hasher = DefaultHasher::new();
1144+
metadata.hash(&mut hasher);
1145+
bytes.hash(&mut hasher);
1146+
let digest = hasher.finish();
1147+
1148+
let cache_dir = std::env::temp_dir().join("opencode-monitor-image-cache");
1149+
std::fs::create_dir_all(&cache_dir).ok()?;
1150+
1151+
let extension = extension_from_mime(metadata);
1152+
let path = cache_dir.join(format!("{digest:016x}.{extension}"));
1153+
if !path.exists() {
1154+
std::fs::write(&path, &bytes).ok()?;
1155+
}
1156+
path.to_str().map(|value| value.to_string())
1157+
}
1158+
1159+
fn frontend_image_value(raw: &str) -> String {
1160+
let trimmed = raw.trim();
1161+
if trimmed.starts_with("data:image/") {
1162+
return persist_data_image_to_temp_file(trimmed).unwrap_or_else(|| trimmed.to_string());
1163+
}
1164+
trimmed.to_string()
1165+
}
1166+
11051167
/// Build REST prompt parts from frontend input.
11061168
///
11071169
/// REST uses `{ type: "file", mime, url: "data:...", filename }` for images.
@@ -1141,7 +1203,6 @@ async fn build_rest_prompt_parts(
11411203
// Local file path — read and base64-encode.
11421204
match std::fs::read(trimmed) {
11431205
Ok(bytes) => {
1144-
use base64::Engine;
11451206
let encoded = base64::engine::general_purpose::STANDARD.encode(&bytes);
11461207
let mime = mime_from_extension(trimmed);
11471208
let filename = std::path::Path::new(trimmed)
@@ -1629,7 +1690,10 @@ pub(crate) async fn send_user_message_core<E: EventSink>(
16291690
if trimmed.is_empty() {
16301691
continue;
16311692
}
1632-
content_parts.push(json!({ "type": "image", "value": trimmed }));
1693+
content_parts.push(json!({
1694+
"type": "image",
1695+
"value": frontend_image_value(trimmed)
1696+
}));
16331697
}
16341698
if !content_parts.is_empty() {
16351699
let user_item_id = {
@@ -2350,6 +2414,18 @@ mod tests {
23502414
});
23512415
}
23522416

2417+
#[test]
2418+
fn frontend_image_value_materializes_data_urls_to_temp_files() {
2419+
let data_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jXioAAAAASUVORK5CYII=";
2420+
2421+
let path = frontend_image_value(data_url);
2422+
let repeated = frontend_image_value(data_url);
2423+
2424+
assert!(!path.starts_with("data:"));
2425+
assert_eq!(path, repeated);
2426+
assert!(PathBuf::from(&path).exists());
2427+
}
2428+
23532429
#[test]
23542430
fn hidden_session_ids_are_read_from_workspace_settings() {
23552431
let runtime = Builder::new_current_thread()

src/features/messages/components/Messages.test.tsx

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -410,6 +410,92 @@ describe("Messages", () => {
410410
expect(useFileLinkOpenerMock).toHaveBeenCalledTimes(1);
411411
});
412412

413+
it("virtualizes large message lists instead of rendering every row", async () => {
414+
const items: ConversationItem[] = Array.from({ length: 60 }, (_, index) => ({
415+
id: `msg-virtual-${index}`,
416+
kind: "message",
417+
role: "assistant",
418+
text: `Virtualized message ${index}`,
419+
}));
420+
const offsetHeightDescriptor = Object.getOwnPropertyDescriptor(
421+
HTMLElement.prototype,
422+
"offsetHeight",
423+
);
424+
const offsetWidthDescriptor = Object.getOwnPropertyDescriptor(
425+
HTMLElement.prototype,
426+
"offsetWidth",
427+
);
428+
const resizeObserverPrototype = window.ResizeObserver?.prototype;
429+
const originalObserve = resizeObserverPrototype?.observe;
430+
const originalUnobserve = resizeObserverPrototype?.unobserve;
431+
const originalDisconnect = resizeObserverPrototype?.disconnect;
432+
433+
Object.defineProperty(HTMLElement.prototype, "offsetHeight", {
434+
configurable: true,
435+
get() {
436+
const element = this as HTMLElement;
437+
if (element.classList.contains("messages")) {
438+
return 720;
439+
}
440+
return 180;
441+
},
442+
});
443+
Object.defineProperty(HTMLElement.prototype, "offsetWidth", {
444+
configurable: true,
445+
get() {
446+
return 1024;
447+
},
448+
});
449+
if (resizeObserverPrototype) {
450+
resizeObserverPrototype.observe = () => undefined;
451+
resizeObserverPrototype.unobserve = () => undefined;
452+
resizeObserverPrototype.disconnect = () => undefined;
453+
}
454+
455+
let unmount: (() => void) | null = null;
456+
try {
457+
const view = render(
458+
<Messages
459+
items={items}
460+
threadId="thread-virtual"
461+
workspaceId="ws-1"
462+
isThinking={false}
463+
openTargets={[]}
464+
selectedOpenAppId=""
465+
/>,
466+
);
467+
unmount = view.unmount;
468+
const { container } = view;
469+
470+
await waitFor(() => {
471+
const renderedMessages = container.querySelectorAll(".message");
472+
expect(renderedMessages.length).toBeGreaterThan(0);
473+
expect(renderedMessages.length).toBeLessThan(items.length);
474+
});
475+
476+
expect(screen.getByText("Virtualized message 0")).toBeTruthy();
477+
expect(screen.queryByText("Virtualized message 20")).toBeNull();
478+
expect(screen.queryByText("Virtualized message 59")).toBeNull();
479+
} finally {
480+
unmount?.();
481+
if (offsetHeightDescriptor) {
482+
Object.defineProperty(HTMLElement.prototype, "offsetHeight", offsetHeightDescriptor);
483+
} else {
484+
delete (HTMLElement.prototype as { offsetHeight?: number }).offsetHeight;
485+
}
486+
if (offsetWidthDescriptor) {
487+
Object.defineProperty(HTMLElement.prototype, "offsetWidth", offsetWidthDescriptor);
488+
} else {
489+
delete (HTMLElement.prototype as { offsetWidth?: number }).offsetWidth;
490+
}
491+
if (resizeObserverPrototype) {
492+
resizeObserverPrototype.observe = originalObserve ?? (() => undefined);
493+
resizeObserverPrototype.unobserve = originalUnobserve ?? (() => undefined);
494+
resizeObserverPrototype.disconnect = originalDisconnect ?? (() => undefined);
495+
}
496+
}
497+
});
498+
413499
it("renders title-only reasoning rows and keeps the working indicator generic", () => {
414500
const items: ConversationItem[] = [
415501
{

0 commit comments

Comments
 (0)